Merge pull request #263 from AI21Labs/rc-tools-dict

fix: change tool to a dict
AI21Labs · Jan 28, 2025 · 5d6e647 · 5d6e647
2 parents aedf04a + 5cb61b8
commit 5d6e647
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 11 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -1,7 +1,7 @@
 minimum_pre_commit_version: 2.20.0
 fail_fast: false
 default_stages:
-  - commit
+  - pre-commit
 exclude: (.idea|vscode)
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks

diff --git a/ai21/clients/common/beta/assistant/assistants.py b/ai21/clients/common/beta/assistant/assistants.py
@@ -24,7 +24,7 @@ def create(
         description: str | NotGiven = NOT_GIVEN,
         optimization: Optimization | NotGiven = NOT_GIVEN,
         models: List[str] | NotGiven = NOT_GIVEN,
-        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tools: List[Dict[str, Tool]] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
         **kwargs,
     ) -> AssistantResponse:
@@ -37,7 +37,7 @@ def _create_body(
         description: str | NotGiven,
         optimization: str | NotGiven,
         models: List[str] | NotGiven,
-        tools: List[str] | NotGiven,
+        tools: List[Dict[str, Tool]] | NotGiven,
         tool_resources: dict | NotGiven,
         **kwargs,
     ) -> Dict[str, Any]:
@@ -71,7 +71,7 @@ def modify(
         optimization: Optimization | NotGiven = NOT_GIVEN,
         is_archived: bool | NotGiven = NOT_GIVEN,
         models: List[str] | NotGiven = NOT_GIVEN,
-        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tools: List[Dict[str, Tool]] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
     ) -> AssistantResponse:
         pass

diff --git a/ai21/clients/studio/resources/beta/assistant/assistant.py b/ai21/clients/studio/resources/beta/assistant/assistant.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import List
+from typing import Dict, List
 
 from ai21.clients.common.beta.assistant.assistants import BaseAssistants
 from ai21.clients.studio.resources.beta.assistant.assistants_plans import AssistantPlans, AsyncAssistantPlans
@@ -30,7 +30,7 @@ def create(
         description: str | NotGiven = NOT_GIVEN,
         optimization: str | NotGiven = NOT_GIVEN,
         models: List[str] | NotGiven = NOT_GIVEN,
-        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tools: List[Dict[str, Tool]] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
         **kwargs,
     ) -> AssistantResponse:
@@ -61,7 +61,7 @@ def modify(
         optimization: str | NotGiven = NOT_GIVEN,
         is_archived: bool | NotGiven = NOT_GIVEN,
         models: List[str] | NotGiven = NOT_GIVEN,
-        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tools: List[Dict[str, Tool]] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
     ) -> AssistantResponse:
         body = self._create_body(
@@ -94,7 +94,7 @@ async def create(
         description: str | NotGiven = NOT_GIVEN,
         optimization: str | NotGiven = NOT_GIVEN,
         models: List[str] | NotGiven = NOT_GIVEN,
-        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tools: List[Dict[str, Tool]] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
         **kwargs,
     ) -> AssistantResponse:
@@ -125,7 +125,7 @@ async def modify(
         optimization: str | NotGiven = NOT_GIVEN,
         is_archived: bool | NotGiven = NOT_GIVEN,
         models: List[str] | NotGiven = NOT_GIVEN,
-        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tools: List[Dict[str, Tool]] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
     ) -> AssistantResponse:
         body = self._create_body(

diff --git a/ai21/models/responses/assistant_response.py b/ai21/models/responses/assistant_response.py
@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Optional, List, Literal
+from typing import Optional, List, Literal, Any, Dict
 
 from ai21.models.ai21_base_model import AI21BaseModel
 from ai21.models.assistant.assistant import ToolResources
@@ -18,7 +18,7 @@ class AssistantResponse(AI21BaseModel):
     avatar: Optional[str] = None
     is_archived: bool = False
     models: Optional[List[str]] = None
-    tools: Optional[List[str]] = None
+    tools: Optional[List[Dict[str, Any]]] = None
     tool_resources: Optional[ToolResources] = None