From e8725bcce1b0afdae4bb42f1f9413a3981e95305 Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 3 Dec 2024 10:39:06 +0200 Subject: [PATCH 01/13] feat: :sparkles: add support for Run resource --- ai21/clients/common/assistant/runs.py | 67 +++++++++++ .../resources/assistant/studio_thread.py | 3 + .../resources/assistant/studio_thread_run.py | 104 ++++++++++++++++++ ai21/models/assistant/run.py | 41 +++++++ ai21/models/responses/run_response.py | 19 ++++ 5 files changed, 234 insertions(+) create mode 100644 ai21/clients/common/assistant/runs.py create mode 100644 ai21/clients/studio/resources/assistant/studio_thread_run.py create mode 100644 ai21/models/assistant/run.py create mode 100644 ai21/models/responses/run_response.py diff --git a/ai21/clients/common/assistant/runs.py b/ai21/clients/common/assistant/runs.py new file mode 100644 index 00000000..656bf26e --- /dev/null +++ b/ai21/clients/common/assistant/runs.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import List + +from ai21.models.assistant.assistant import Optimization +from ai21.models.assistant.run import ToolOutput +from ai21.models.responses.run_response import RunResponse +from ai21.types import NOT_GIVEN, NotGiven +from ai21.utils.typing import remove_not_given + + +class Runs(ABC): + _module_name = "runs" + + @abstractmethod + def create( + self, + *, + thread_id: str, + assistant_id: str, + description: str | NotGiven = NOT_GIVEN, + optimization: Optimization | NotGiven = NOT_GIVEN, + **kwargs, + ) -> RunResponse: + pass + + def _create_body( + self, + *, + thread_id: str, + assistant_id: str, + description: str | NotGiven, + optimization: str | NotGiven, + **kwargs, + ) -> dict: + return remove_not_given( + { + "thread_id": thread_id, + "assistant_id": assistant_id, + "description": description, + "optimization": optimization, + **kwargs, + } + ) + + @abstractmethod + def retrieve( + self, + *, + thread_id: str, + run_id: str, + ) -> RunResponse: + pass + + @abstractmethod + def cancel( + self, + *, + thread_id: str, + run_id: str, + ) -> RunResponse: + pass + + @abstractmethod + def submit_tool_outputs(self, *, thread_id: str, run_id: str, tool_outputs: List[ToolOutput]) -> RunResponse: + pass diff --git a/ai21/clients/studio/resources/assistant/studio_thread.py b/ai21/clients/studio/resources/assistant/studio_thread.py index 60c88e22..2a33e595 100644 --- a/ai21/clients/studio/resources/assistant/studio_thread.py +++ b/ai21/clients/studio/resources/assistant/studio_thread.py @@ -4,6 +4,7 @@ from ai21.clients.common.assistant.threads import Threads from ai21.clients.studio.resources.assistant.studio_thread_message import StudioThreadMessage, AsyncStudioThreadMessage +from ai21.clients.studio.resources.assistant.studio_thread_run import AsyncStudioThreadRun, StudioThreadRun from ai21.clients.studio.resources.studio_resource import StudioResource, AsyncStudioResource from ai21.http_client.async_http_client import AsyncAI21HTTPClient from ai21.http_client.http_client import AI21HTTPClient @@ -16,6 +17,7 @@ def __init__(self, client: AI21HTTPClient): super().__init__(client) self.messages = StudioThreadMessage(client) + self.runs = StudioThreadRun(client) def create( self, @@ -35,6 +37,7 @@ def __init__(self, client: AsyncAI21HTTPClient): super().__init__(client) self.messages = AsyncStudioThreadMessage(client) + self.runs = AsyncStudioThreadRun(client) async def create( self, diff --git a/ai21/clients/studio/resources/assistant/studio_thread_run.py b/ai21/clients/studio/resources/assistant/studio_thread_run.py new file mode 100644 index 00000000..2df87247 --- /dev/null +++ b/ai21/clients/studio/resources/assistant/studio_thread_run.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +from typing import List + +from ai21.clients.common.assistant.runs import Runs +from ai21.clients.studio.resources.studio_resource import StudioResource, AsyncStudioResource +from ai21.models.assistant.assistant import Optimization +from ai21.models.assistant.run import ToolOutput +from ai21.models.responses.run_response import RunResponse +from ai21.types import NotGiven, NOT_GIVEN + + +class StudioThreadRun(StudioResource, Runs): + def create( + self, + *, + thread_id: str, + assistant_id: str, + description: str | NotGiven = NOT_GIVEN, + optimization: Optimization | NotGiven = NOT_GIVEN, + **kwargs, + ) -> RunResponse: + body = self._create_body( + thread_id=thread_id, + assistant_id=assistant_id, + description=description, + optimization=optimization, + **kwargs, + ) + + return self._post(path=f"/threads/{thread_id}/{self._module_name}", body=body, response_cls=RunResponse) + + def retrieve( + self, + *, + thread_id: str, + run_id: str, + ) -> RunResponse: + return self._get(path=f"/threads/{thread_id}/{self._module_name}/{run_id}", response_cls=RunResponse) + + def cancel( + self, + *, + thread_id: str, + run_id: str, + ) -> RunResponse: + return self._post(path=f"/threads/{thread_id}/{self._module_name}/{run_id}/cancel", response_cls=RunResponse) + + def submit_tool_outputs(self, *, thread_id: str, run_id: str, tool_outputs: List[ToolOutput]) -> RunResponse: + body = dict(tool_outputs=tool_outputs) + + return self._post( + path=f"/threads/{thread_id}/{self._module_name}/{run_id}/submit_tool_outputs", + body=body, + response_cls=RunResponse, + ) + + +class AsyncStudioThreadRun(AsyncStudioResource, Runs): + async def create( + self, + *, + thread_id: str, + assistant_id: str, + description: str | NotGiven = NOT_GIVEN, + optimization: Optimization | NotGiven = NOT_GIVEN, + **kwargs, + ) -> RunResponse: + body = self._create_body( + thread_id=thread_id, + assistant_id=assistant_id, + description=description, + optimization=optimization, + **kwargs, + ) + + return await self._post(path=f"/threads/{thread_id}/{self._module_name}", body=body, response_cls=RunResponse) + + async def retrieve( + self, + *, + thread_id: str, + run_id: str, + ) -> RunResponse: + return await self._get(path=f"/threads/{thread_id}/{self._module_name}/{run_id}", response_cls=RunResponse) + + async def cancel( + self, + *, + thread_id: str, + run_id: str, + ) -> RunResponse: + return await self._post( + path=f"/threads/{thread_id}/{self._module_name}/{run_id}/cancel", response_cls=RunResponse + ) + + async def submit_tool_outputs(self, *, thread_id: str, run_id: str, tool_outputs: List[ToolOutput]) -> RunResponse: + body = dict(tool_outputs=tool_outputs) + + return await self._post( + path=f"/threads/{thread_id}/{self._module_name}/{run_id}/submit_tool_outputs", + body=body, + response_cls=RunResponse, + ) diff --git a/ai21/models/assistant/run.py b/ai21/models/assistant/run.py new file mode 100644 index 00000000..9447828d --- /dev/null +++ b/ai21/models/assistant/run.py @@ -0,0 +1,41 @@ +from typing import Literal, Any + +from typing_extensions import TypedDict + + +RunStatus = Literal[ + "cancelled", + "cancelling", + "completed", + "expired", + "failed", + "incomplete", + "in_progress", + "queued", + "requires_action", +] + + +class ToolOutput(TypedDict): + tool_call_id: str + output: Any + + +class Function(TypedDict): + name: str + arguments: Any + + +class ToolCall(TypedDict): + type: Literal["function"] + id: str + function: Function + + +class SubmitToolCallOutputs(TypedDict): + tool_calls: list[ToolOutput] + + +class RequiredAction(TypedDict): + type: Literal["submit_tool_outputs"] + submit_tool_outputs: dict[str, Any] diff --git a/ai21/models/responses/run_response.py b/ai21/models/responses/run_response.py new file mode 100644 index 00000000..2bf45c8a --- /dev/null +++ b/ai21/models/responses/run_response.py @@ -0,0 +1,19 @@ +from datetime import datetime +from typing import Optional + +from ai21.models.ai21_base_model import AI21BaseModel +from ai21.models.assistant.assistant import Optimization +from ai21.models.assistant.run import RunStatus, RequiredAction + + +class RunResponse(AI21BaseModel): + id: str + created_at: datetime + updated_at: datetime + thread_id: str + assistant_id: str + description: str = None + status: RunStatus + optimization: Optimization + execution_id: Optional[str] = None + required_action: Optional[RequiredAction] = None From 810096e109492ddbb21ba354775910bed66cf68f Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 3 Dec 2024 11:06:14 +0200 Subject: [PATCH 02/13] fix: :label: typing backwards compatability --- ai21/models/assistant/run.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ai21/models/assistant/run.py b/ai21/models/assistant/run.py index 9447828d..6a727141 100644 --- a/ai21/models/assistant/run.py +++ b/ai21/models/assistant/run.py @@ -1,4 +1,4 @@ -from typing import Literal, Any +from typing import Literal, Any, List, Dict from typing_extensions import TypedDict @@ -33,9 +33,9 @@ class ToolCall(TypedDict): class SubmitToolCallOutputs(TypedDict): - tool_calls: list[ToolOutput] + tool_calls: List[ToolOutput] class RequiredAction(TypedDict): type: Literal["submit_tool_outputs"] - submit_tool_outputs: dict[str, Any] + submit_tool_outputs: Dict[str, Any] From 0cc133710b683f42ce74b3ed871099f8d4f9bfb1 Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 3 Dec 2024 11:19:35 +0200 Subject: [PATCH 03/13] refactor: :truck: rename classes & types and such --- ai21/clients/common/assistant/assistants.py | 8 ++--- ai21/clients/common/assistant/threads.py | 6 ++-- .../resources/assistant/studio_assistant.py | 30 ++++++++-------- .../resources/assistant/studio_thread.py | 34 +++++++++---------- .../assistant/studio_thread_message.py | 4 +-- .../resources/assistant/studio_thread_run.py | 4 +-- .../studio/resources/beta/async_beta.py | 8 ++--- ai21/clients/studio/resources/beta/beta.py | 8 ++--- ai21/models/responses/assistant_response.py | 4 +-- ai21/models/responses/thread_response.py | 4 +-- 10 files changed, 55 insertions(+), 55 deletions(-) diff --git a/ai21/clients/common/assistant/assistants.py b/ai21/clients/common/assistant/assistants.py index 1beb2b7f..07ecb51c 100644 --- a/ai21/clients/common/assistant/assistants.py +++ b/ai21/clients/common/assistant/assistants.py @@ -4,7 +4,7 @@ from typing import Any, Dict, List from ai21.models.assistant.assistant import Optimization, Tool, ToolResources -from ai21.models.responses.assistant_response import Assistant, ListAssistant +from ai21.models.responses.assistant_response import AssistantResponse, ListAssistant from ai21.types import NotGiven, NOT_GIVEN from ai21.utils.typing import remove_not_given @@ -24,7 +24,7 @@ def create( tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, **kwargs, - ) -> Assistant: + ) -> AssistantResponse: pass def _create_body( @@ -57,7 +57,7 @@ def list(self) -> ListAssistant: pass @abstractmethod - def retrieve(self, assistant_id: str) -> Assistant: + def retrieve(self, assistant_id: str) -> AssistantResponse: pass @abstractmethod @@ -73,5 +73,5 @@ def modify( models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, - ) -> Assistant: + ) -> AssistantResponse: pass diff --git a/ai21/clients/common/assistant/threads.py b/ai21/clients/common/assistant/threads.py index 9b64580d..e7359cc9 100644 --- a/ai21/clients/common/assistant/threads.py +++ b/ai21/clients/common/assistant/threads.py @@ -5,7 +5,7 @@ from ai21.clients.common.assistant.messages import Messages from ai21.models.assistant.message import Message -from ai21.models.responses.thread_response import Thread +from ai21.models.responses.thread_response import ThreadResponse class Threads(ABC): @@ -17,9 +17,9 @@ def create( self, messages: List[Message], **kwargs, - ) -> Thread: + ) -> ThreadResponse: pass @abstractmethod - def retrieve(self, thread_id: str) -> Thread: + def retrieve(self, thread_id: str) -> ThreadResponse: pass diff --git a/ai21/clients/studio/resources/assistant/studio_assistant.py b/ai21/clients/studio/resources/assistant/studio_assistant.py index f2a3c306..49cf3b34 100644 --- a/ai21/clients/studio/resources/assistant/studio_assistant.py +++ b/ai21/clients/studio/resources/assistant/studio_assistant.py @@ -8,11 +8,11 @@ StudioResource, ) from ai21.models.assistant.assistant import Tool, ToolResources -from ai21.models.responses.assistant_response import Assistant, ListAssistant +from ai21.models.responses.assistant_response import AssistantResponse, ListAssistant from ai21.types import NotGiven, NOT_GIVEN -class StudioAssistant(StudioResource, Assistants): +class Assistant(StudioResource, Assistants): def create( self, name: str, @@ -24,7 +24,7 @@ def create( tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, **kwargs, - ) -> Assistant: + ) -> AssistantResponse: body = self._create_body( name=name, description=description, @@ -36,10 +36,10 @@ def create( **kwargs, ) - return self._post(path=f"/{self._module_name}", body=body, response_cls=Assistant) + return self._post(path=f"/{self._module_name}", body=body, response_cls=AssistantResponse) - def retrieve(self, assistant_id: str) -> Assistant: - return self._get(path=f"/{self._module_name}/{assistant_id}", response_cls=Assistant) + def retrieve(self, assistant_id: str) -> AssistantResponse: + return self._get(path=f"/{self._module_name}/{assistant_id}", response_cls=AssistantResponse) def list(self) -> ListAssistant: return self._get(path=f"/{self._module_name}", response_cls=ListAssistant) @@ -56,7 +56,7 @@ def modify( models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, - ) -> Assistant: + ) -> AssistantResponse: body = self._create_body( name=name, description=description, @@ -68,10 +68,10 @@ def modify( tool_resources=tool_resources, ) - return self._patch(path=f"/{self._module_name}/{assistant_id}", body=body, response_cls=Assistant) + return self._patch(path=f"/{self._module_name}/{assistant_id}", body=body, response_cls=AssistantResponse) -class AsyncStudioAssistant(AsyncStudioResource, Assistants): +class AsyncAssistant(AsyncStudioResource, Assistants): async def create( self, name: str, @@ -83,7 +83,7 @@ async def create( tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, **kwargs, - ) -> Assistant: + ) -> AssistantResponse: body = self._create_body( name=name, description=description, @@ -95,10 +95,10 @@ async def create( **kwargs, ) - return await self._post(path=f"/{self._module_name}", body=body, response_cls=Assistant) + return await self._post(path=f"/{self._module_name}", body=body, response_cls=AssistantResponse) - async def retrieve(self, assistant_id: str) -> Assistant: - return await self._get(path=f"/{self._module_name}/{assistant_id}", response_cls=Assistant) + async def retrieve(self, assistant_id: str) -> AssistantResponse: + return await self._get(path=f"/{self._module_name}/{assistant_id}", response_cls=AssistantResponse) async def list(self) -> ListAssistant: return await self._get(path=f"/{self._module_name}", response_cls=ListAssistant) @@ -115,7 +115,7 @@ async def modify( models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, - ) -> Assistant: + ) -> AssistantResponse: body = self._create_body( name=name, description=description, @@ -127,4 +127,4 @@ async def modify( tool_resources=tool_resources, ) - return await self._patch(path=f"/{self._module_name}/{assistant_id}", body=body, response_cls=Assistant) + return await self._patch(path=f"/{self._module_name}/{assistant_id}", body=body, response_cls=AssistantResponse) diff --git a/ai21/clients/studio/resources/assistant/studio_thread.py b/ai21/clients/studio/resources/assistant/studio_thread.py index 2a33e595..e2ac565a 100644 --- a/ai21/clients/studio/resources/assistant/studio_thread.py +++ b/ai21/clients/studio/resources/assistant/studio_thread.py @@ -3,50 +3,50 @@ from typing import List from ai21.clients.common.assistant.threads import Threads -from ai21.clients.studio.resources.assistant.studio_thread_message import StudioThreadMessage, AsyncStudioThreadMessage -from ai21.clients.studio.resources.assistant.studio_thread_run import AsyncStudioThreadRun, StudioThreadRun +from ai21.clients.studio.resources.assistant.studio_thread_message import ThreadMessage, AsyncThreadMessage +from ai21.clients.studio.resources.assistant.studio_thread_run import AsyncThreadRun, ThreadRun from ai21.clients.studio.resources.studio_resource import StudioResource, AsyncStudioResource from ai21.http_client.async_http_client import AsyncAI21HTTPClient from ai21.http_client.http_client import AI21HTTPClient from ai21.models.assistant.message import Message -from ai21.models.responses.thread_response import Thread +from ai21.models.responses.thread_response import ThreadResponse -class StudioThread(StudioResource, Threads): +class Thread(StudioResource, Threads): def __init__(self, client: AI21HTTPClient): super().__init__(client) - self.messages = StudioThreadMessage(client) - self.runs = StudioThreadRun(client) + self.messages = ThreadMessage(client) + self.runs = ThreadRun(client) def create( self, messages: List[Message], **kwargs, - ) -> Thread: + ) -> ThreadResponse: body = dict(messages=messages) - return self._post(path=f"/{self._module_name}", body=body, response_cls=Thread) + return self._post(path=f"/{self._module_name}", body=body, response_cls=ThreadResponse) - def retrieve(self, thread_id: str) -> Thread: - return self._get(path=f"/{self._module_name}/{thread_id}", response_cls=Thread) + def retrieve(self, thread_id: str) -> ThreadResponse: + return self._get(path=f"/{self._module_name}/{thread_id}", response_cls=ThreadResponse) -class AsyncStudioThread(AsyncStudioResource, Threads): +class AsyncThread(AsyncStudioResource, Threads): def __init__(self, client: AsyncAI21HTTPClient): super().__init__(client) - self.messages = AsyncStudioThreadMessage(client) - self.runs = AsyncStudioThreadRun(client) + self.messages = AsyncThreadMessage(client) + self.runs = AsyncThreadRun(client) async def create( self, messages: List[Message], **kwargs, - ) -> Thread: + ) -> ThreadResponse: body = dict(messages=messages) - return await self._post(path=f"/{self._module_name}", body=body, response_cls=Thread) + return await self._post(path=f"/{self._module_name}", body=body, response_cls=ThreadResponse) - async def retrieve(self, thread_id: str) -> Thread: - return await self._get(path=f"/{self._module_name}/{thread_id}", response_cls=Thread) + async def retrieve(self, thread_id: str) -> ThreadResponse: + return await self._get(path=f"/{self._module_name}/{thread_id}", response_cls=ThreadResponse) diff --git a/ai21/clients/studio/resources/assistant/studio_thread_message.py b/ai21/clients/studio/resources/assistant/studio_thread_message.py index 95cb1b47..bc0494dd 100644 --- a/ai21/clients/studio/resources/assistant/studio_thread_message.py +++ b/ai21/clients/studio/resources/assistant/studio_thread_message.py @@ -6,7 +6,7 @@ from ai21.models.responses.message_response import MessageResponse, ListMessageResponse -class StudioThreadMessage(StudioResource, Messages): +class ThreadMessage(StudioResource, Messages): def create( self, thread_id: str, @@ -26,7 +26,7 @@ def list(self, thread_id: str) -> ListMessageResponse: return self._get(path=f"/threads/{thread_id}/{self._module_name}", response_cls=ListMessageResponse) -class AsyncStudioThreadMessage(AsyncStudioResource, Messages): +class AsyncThreadMessage(AsyncStudioResource, Messages): async def create( self, thread_id: str, diff --git a/ai21/clients/studio/resources/assistant/studio_thread_run.py b/ai21/clients/studio/resources/assistant/studio_thread_run.py index 2df87247..092ba242 100644 --- a/ai21/clients/studio/resources/assistant/studio_thread_run.py +++ b/ai21/clients/studio/resources/assistant/studio_thread_run.py @@ -10,7 +10,7 @@ from ai21.types import NotGiven, NOT_GIVEN -class StudioThreadRun(StudioResource, Runs): +class ThreadRun(StudioResource, Runs): def create( self, *, @@ -56,7 +56,7 @@ def submit_tool_outputs(self, *, thread_id: str, run_id: str, tool_outputs: List ) -class AsyncStudioThreadRun(AsyncStudioResource, Runs): +class AsyncThreadRun(AsyncStudioResource, Runs): async def create( self, *, diff --git a/ai21/clients/studio/resources/beta/async_beta.py b/ai21/clients/studio/resources/beta/async_beta.py index 1bb13bbe..9afbc4e0 100644 --- a/ai21/clients/studio/resources/beta/async_beta.py +++ b/ai21/clients/studio/resources/beta/async_beta.py @@ -1,5 +1,5 @@ -from ai21.clients.studio.resources.assistant.studio_assistant import AsyncStudioAssistant -from ai21.clients.studio.resources.assistant.studio_thread import AsyncStudioThread +from ai21.clients.studio.resources.assistant.studio_assistant import AsyncAssistant +from ai21.clients.studio.resources.assistant.studio_thread import AsyncThread from ai21.clients.studio.resources.studio_conversational_rag import AsyncStudioConversationalRag from ai21.clients.studio.resources.studio_resource import AsyncStudioResource from ai21.http_client.async_http_client import AsyncAI21HTTPClient @@ -9,6 +9,6 @@ class AsyncBeta(AsyncStudioResource): def __init__(self, client: AsyncAI21HTTPClient): super().__init__(client) - self.assistants = AsyncStudioAssistant(client) + self.assistants = AsyncAssistant(client) self.conversational_rag = AsyncStudioConversationalRag(client) - self.threads = AsyncStudioThread(client) + self.threads = AsyncThread(client) diff --git a/ai21/clients/studio/resources/beta/beta.py b/ai21/clients/studio/resources/beta/beta.py index affede10..18ed1c0c 100644 --- a/ai21/clients/studio/resources/beta/beta.py +++ b/ai21/clients/studio/resources/beta/beta.py @@ -1,5 +1,5 @@ -from ai21.clients.studio.resources.assistant.studio_assistant import StudioAssistant -from ai21.clients.studio.resources.assistant.studio_thread import StudioThread +from ai21.clients.studio.resources.assistant.studio_assistant import Assistant +from ai21.clients.studio.resources.assistant.studio_thread import Thread from ai21.clients.studio.resources.studio_conversational_rag import StudioConversationalRag from ai21.clients.studio.resources.studio_resource import StudioResource from ai21.http_client.http_client import AI21HTTPClient @@ -9,6 +9,6 @@ class Beta(StudioResource): def __init__(self, client: AI21HTTPClient): super().__init__(client) - self.assistants = StudioAssistant(client) + self.assistants = Assistant(client) self.conversational_rag = StudioConversationalRag(client) - self.threads = StudioThread(client) + self.threads = Thread(client) diff --git a/ai21/models/responses/assistant_response.py b/ai21/models/responses/assistant_response.py index 3263b5a8..0bf40c5f 100644 --- a/ai21/models/responses/assistant_response.py +++ b/ai21/models/responses/assistant_response.py @@ -5,7 +5,7 @@ from ai21.models.assistant.assistant import ToolResources -class Assistant(AI21BaseModel): +class AssistantResponse(AI21BaseModel): id: str created_at: datetime updated_at: datetime @@ -23,4 +23,4 @@ class Assistant(AI21BaseModel): class ListAssistant(AI21BaseModel): - results: List[Assistant] + results: List[AssistantResponse] diff --git a/ai21/models/responses/thread_response.py b/ai21/models/responses/thread_response.py index b2c9bbc7..f7206a40 100644 --- a/ai21/models/responses/thread_response.py +++ b/ai21/models/responses/thread_response.py @@ -4,7 +4,7 @@ from ai21.models.ai21_base_model import AI21BaseModel -class Thread(AI21BaseModel): +class ThreadResponse(AI21BaseModel): id: str created_at: datetime updated_at: datetime @@ -12,4 +12,4 @@ class Thread(AI21BaseModel): class ListThread(AI21BaseModel): - results: List[Thread] + results: List[ThreadResponse] From 3ab7bec56d52730197407d2c8019a24cc0e197a0 Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 3 Dec 2024 11:36:14 +0200 Subject: [PATCH 04/13] refactor: :fire: remove `avatar` option --- ai21/clients/common/assistant/assistants.py | 4 ---- .../studio/resources/assistant/studio_assistant.py | 8 -------- 2 files changed, 12 deletions(-) diff --git a/ai21/clients/common/assistant/assistants.py b/ai21/clients/common/assistant/assistants.py index 07ecb51c..65bbabf2 100644 --- a/ai21/clients/common/assistant/assistants.py +++ b/ai21/clients/common/assistant/assistants.py @@ -19,7 +19,6 @@ def create( *, description: str | NotGiven = NOT_GIVEN, optimization: Optimization | NotGiven = NOT_GIVEN, - avatar: str | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, @@ -33,7 +32,6 @@ def _create_body( name: str, description: str | NotGiven, optimization: str | NotGiven, - avatar: str | NotGiven, models: List[str] | NotGiven, tools: List[str] | NotGiven, tool_resources: dict | NotGiven, @@ -44,7 +42,6 @@ def _create_body( "name": name, "description": description, "optimization": optimization, - "avatar": avatar, "models": models, "tools": tools, "tool_resources": tool_resources, @@ -68,7 +65,6 @@ def modify( name: str | NotGiven = NOT_GIVEN, description: str | NotGiven = NOT_GIVEN, optimization: Optimization | NotGiven = NOT_GIVEN, - avatar: str | NotGiven = NOT_GIVEN, is_archived: bool | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, diff --git a/ai21/clients/studio/resources/assistant/studio_assistant.py b/ai21/clients/studio/resources/assistant/studio_assistant.py index 49cf3b34..259302bf 100644 --- a/ai21/clients/studio/resources/assistant/studio_assistant.py +++ b/ai21/clients/studio/resources/assistant/studio_assistant.py @@ -19,7 +19,6 @@ def create( *, description: str | NotGiven = NOT_GIVEN, optimization: str | NotGiven = NOT_GIVEN, - avatar: str | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, @@ -29,7 +28,6 @@ def create( name=name, description=description, optimization=optimization, - avatar=avatar, models=models, tools=tools, tool_resources=tool_resources, @@ -51,7 +49,6 @@ def modify( name: str | NotGiven = NOT_GIVEN, description: str | NotGiven = NOT_GIVEN, optimization: str | NotGiven = NOT_GIVEN, - avatar: str | NotGiven = NOT_GIVEN, is_archived: bool | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, @@ -61,7 +58,6 @@ def modify( name=name, description=description, optimization=optimization, - avatar=avatar, is_archived=is_archived, models=models, tools=tools, @@ -78,7 +74,6 @@ async def create( *, description: str | NotGiven = NOT_GIVEN, optimization: str | NotGiven = NOT_GIVEN, - avatar: str | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, @@ -88,7 +83,6 @@ async def create( name=name, description=description, optimization=optimization, - avatar=avatar, models=models, tools=tools, tool_resources=tool_resources, @@ -110,7 +104,6 @@ async def modify( name: str | NotGiven = NOT_GIVEN, description: str | NotGiven = NOT_GIVEN, optimization: str | NotGiven = NOT_GIVEN, - avatar: str | NotGiven = NOT_GIVEN, is_archived: bool | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, @@ -120,7 +113,6 @@ async def modify( name=name, description=description, optimization=optimization, - avatar=avatar, is_archived=is_archived, models=models, tools=tools, From f4da2cdbe97794dcb93edf35dca79fd4f6634189 Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 3 Dec 2024 11:59:37 +0200 Subject: [PATCH 05/13] fix: :label: typing yay --- ai21/models/assistant/run.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ai21/models/assistant/run.py b/ai21/models/assistant/run.py index 6a727141..ed118545 100644 --- a/ai21/models/assistant/run.py +++ b/ai21/models/assistant/run.py @@ -1,4 +1,4 @@ -from typing import Literal, Any, List, Dict +from typing import Literal, Any, List from typing_extensions import TypedDict @@ -38,4 +38,4 @@ class SubmitToolCallOutputs(TypedDict): class RequiredAction(TypedDict): type: Literal["submit_tool_outputs"] - submit_tool_outputs: Dict[str, Any] + submit_tool_outputs: SubmitToolCallOutputs From 4d2bb958414ec1af42484462d9564e7da2a177ea Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 3 Dec 2024 12:00:02 +0200 Subject: [PATCH 06/13] feat: :sparkles: add examples --- examples/studio/assistant/assistant.py | 38 ++++++++++++++++++ examples/studio/assistant/async_assistant.py | 42 ++++++++++++++++++++ 2 files changed, 80 insertions(+) create mode 100644 examples/studio/assistant/assistant.py create mode 100644 examples/studio/assistant/async_assistant.py diff --git a/examples/studio/assistant/assistant.py b/examples/studio/assistant/assistant.py new file mode 100644 index 00000000..b0357a74 --- /dev/null +++ b/examples/studio/assistant/assistant.py @@ -0,0 +1,38 @@ +from ai21 import AI21Client + + +def main(): + ai21_client = AI21Client() + + assistant = ai21_client.beta.assistants.create(name="My Assistant") + + thread = ai21_client.beta.threads.create( + messages=[ + { + "role": "user", + "content": { + "type": "text", + "text": "Hello", + }, + }, + ] + ) + + run = ai21_client.beta.threads.runs.create( + thread_id=thread.id, + assistant_id=assistant.id, + description="Get response", + optimization="latency", + ) + + while run.status == "in_progress": + run = ai21_client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id) + + if run.status == "completed": + messages = ai21_client.beta.threads.messages.list(thread_id=thread.id) + print("Messages:") + print("\n".join(f"{msg.role}: {msg.content['text']}" for msg in messages.results)) + + +if __name__ == "__main__": + main() diff --git a/examples/studio/assistant/async_assistant.py b/examples/studio/assistant/async_assistant.py new file mode 100644 index 00000000..396b73f7 --- /dev/null +++ b/examples/studio/assistant/async_assistant.py @@ -0,0 +1,42 @@ +import asyncio + +from ai21 import AsyncAI21Client + + +async def main(): + ai21_client = AsyncAI21Client( + api_key="413NQymWvgp83hNaqbA3EwYgqUjvREgn", api_host="https://api-stage.ai21.com/studio/v1" + ) + + assistant = await ai21_client.beta.assistants.create(name="My Assistant") + + thread = await ai21_client.beta.threads.create( + messages=[ + { + "role": "user", + "content": { + "type": "text", + "text": "Hello", + }, + }, + ] + ) + + run = await ai21_client.beta.threads.runs.create( + thread_id=thread.id, + assistant_id=assistant.id, + description="Get response", + optimization="latency", + ) + + while run.status == "in_progress": + run = await ai21_client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id) + + if run.status == "completed": + messages = await ai21_client.beta.threads.messages.list(thread_id=thread.id) + print("Messages:") + print("\n".join(f"{msg.role}: {msg.content['text']}" for msg in messages.results)) + + +if __name__ == "__main__": + asyncio.run(main()) From bdb4bc3eba0db597e7965f80afc057c1aa48ea0e Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 3 Dec 2024 13:49:32 +0200 Subject: [PATCH 07/13] fix: :label: make `description` optional --- ai21/models/responses/run_response.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ai21/models/responses/run_response.py b/ai21/models/responses/run_response.py index 2bf45c8a..0a260235 100644 --- a/ai21/models/responses/run_response.py +++ b/ai21/models/responses/run_response.py @@ -12,7 +12,7 @@ class RunResponse(AI21BaseModel): updated_at: datetime thread_id: str assistant_id: str - description: str = None + description: Optional[str] = None status: RunStatus optimization: Optimization execution_id: Optional[str] = None From d9a3209811ae668fa0317eb765373acfdd6a21d2 Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 3 Dec 2024 13:54:34 +0200 Subject: [PATCH 08/13] docs: :memo: add assistants info to README also, add examples --- README.md | 35 ++++++++++++++++++++ examples/studio/assistant/assistant.py | 9 +++-- examples/studio/assistant/async_assistant.py | 2 -- 3 files changed, 41 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index b3bbeb80..8658ad0e 100644 --- a/README.md +++ b/README.md @@ -388,6 +388,41 @@ For a more detailed example, see the chat [sync](examples/studio/conversational_ --- +### Assistants (Beta) + +Create assistants to help you with your tasks. + +```python +from time import sleep +from ai21 import AI21Client +from ai21.models.assistant.message import Message + +messages = [ + Message(content={"type": "text", "text": "Youre message here"}, role="user"), +] + +client = AI21Client() + +assistant = client.beta.assistants.create(name="My assistant") +thread = client.beta.threads.create(messages=messages) +run = client.beta.threads.runs.create(thread_id=thread.id, assistant_id=assistant.id) + +while run.status == "in_progress": + run = client.beta.threads.runs.get(thread_id=thread.id, run_id=run.id) + sleep(1) + +if run.status == "completed": + messages = client.beta.threads.messages.list(thread_id=thread.id) + print(messages) +else: + # handle error or required action + pass +``` + +For a more detailed example, see the chat [sync](examples/studio/assistant/assistant.py) and [async](examples/studio/assistant/async_assistant.py) examples. + +--- + ### File Upload ```python diff --git a/examples/studio/assistant/assistant.py b/examples/studio/assistant/assistant.py index b0357a74..923187c8 100644 --- a/examples/studio/assistant/assistant.py +++ b/examples/studio/assistant/assistant.py @@ -2,7 +2,9 @@ def main(): - ai21_client = AI21Client() + ai21_client = AI21Client( + api_key="413NQymWvgp83hNaqbA3EwYgqUjvREgn", api_host="https://api-stage.ai21.com/studio/v1" + ) assistant = ai21_client.beta.assistants.create(name="My Assistant") @@ -21,8 +23,6 @@ def main(): run = ai21_client.beta.threads.runs.create( thread_id=thread.id, assistant_id=assistant.id, - description="Get response", - optimization="latency", ) while run.status == "in_progress": @@ -32,6 +32,9 @@ def main(): messages = ai21_client.beta.threads.messages.list(thread_id=thread.id) print("Messages:") print("\n".join(f"{msg.role}: {msg.content['text']}" for msg in messages.results)) + else: + # handle error or required action + pass if __name__ == "__main__": diff --git a/examples/studio/assistant/async_assistant.py b/examples/studio/assistant/async_assistant.py index 396b73f7..b141e9d7 100644 --- a/examples/studio/assistant/async_assistant.py +++ b/examples/studio/assistant/async_assistant.py @@ -25,8 +25,6 @@ async def main(): run = await ai21_client.beta.threads.runs.create( thread_id=thread.id, assistant_id=assistant.id, - description="Get response", - optimization="latency", ) while run.status == "in_progress": From b65fba23c47980161d02fe080f2df1d9faea47a2 Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 3 Dec 2024 13:57:38 +0200 Subject: [PATCH 09/13] test: :white_check_mark: add assistants to tests --- examples/studio/assistant/assistant.py | 4 +--- examples/studio/assistant/async_assistant.py | 4 +--- tests/integration_tests/clients/test_studio.py | 4 ++++ 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/examples/studio/assistant/assistant.py b/examples/studio/assistant/assistant.py index 923187c8..edaebd93 100644 --- a/examples/studio/assistant/assistant.py +++ b/examples/studio/assistant/assistant.py @@ -2,9 +2,7 @@ def main(): - ai21_client = AI21Client( - api_key="413NQymWvgp83hNaqbA3EwYgqUjvREgn", api_host="https://api-stage.ai21.com/studio/v1" - ) + ai21_client = AI21Client() assistant = ai21_client.beta.assistants.create(name="My Assistant") diff --git a/examples/studio/assistant/async_assistant.py b/examples/studio/assistant/async_assistant.py index b141e9d7..b4c84c23 100644 --- a/examples/studio/assistant/async_assistant.py +++ b/examples/studio/assistant/async_assistant.py @@ -4,9 +4,7 @@ async def main(): - ai21_client = AsyncAI21Client( - api_key="413NQymWvgp83hNaqbA3EwYgqUjvREgn", api_host="https://api-stage.ai21.com/studio/v1" - ) + ai21_client = AsyncAI21Client() assistant = await ai21_client.beta.assistants.create(name="My Assistant") diff --git a/tests/integration_tests/clients/test_studio.py b/tests/integration_tests/clients/test_studio.py index 5cd0d4dd..4a581b0e 100644 --- a/tests/integration_tests/clients/test_studio.py +++ b/tests/integration_tests/clients/test_studio.py @@ -55,6 +55,8 @@ def test_studio(test_file_name: str): ("chat/async_stream_chat_completions.py",), ("conversational_rag/conversational_rag.py",), ("conversational_rag/async_conversational_rag.py",), + ("assistant/assistant.py",), + ("assistant/async_assistant.py",), ], ids=[ "when_chat__should_return_ok", @@ -62,6 +64,8 @@ def test_studio(test_file_name: str): "when_stream_chat_completions__should_return_ok", "when_conversational_rag__should_return_ok", "when_async_conversational_rag__should_return_ok", + "when_assistant__should_return_ok", + "when_async_assistant__should_return_ok", ], ) async def test_async_studio(test_file_name: str): From 84919b6ea2a95c2269bfc31bb6987fb63b658ba9 Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 3 Dec 2024 14:24:04 +0200 Subject: [PATCH 10/13] docs: :memo: update docs about Assistants --- README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 8658ad0e..71856588 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,7 @@ - [Installation](#Installation) 💿 - [Usage - Chat Completions](#Usage) - [Conversational RAG (Beta)](#Conversational-RAG-Beta) +- [Assistants (Beta)](#Assistants-Beta) - [Older Models Support Usage](#Older-Models-Support-Usage) - [More Models](#More-Models) - [Streaming](#Streaming) @@ -419,7 +420,7 @@ else: pass ``` -For a more detailed example, see the chat [sync](examples/studio/assistant/assistant.py) and [async](examples/studio/assistant/async_assistant.py) examples. +For a more detailed example, see [sync](examples/studio/assistant/assistant.py) and [async](examples/studio/assistant/async_assistant.py) examples. --- From af1e2e303f8f8caf364a867ade1906e53aebfc97 Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 3 Dec 2024 15:16:38 +0200 Subject: [PATCH 11/13] docs: :memo: better example --- examples/studio/assistant/assistant.py | 12 ++++++++++-- examples/studio/assistant/async_assistant.py | 11 +++++++++++ 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/examples/studio/assistant/assistant.py b/examples/studio/assistant/assistant.py index edaebd93..041c62fe 100644 --- a/examples/studio/assistant/assistant.py +++ b/examples/studio/assistant/assistant.py @@ -1,5 +1,9 @@ +import time + from ai21 import AI21Client +TIMEOUT = 20 + def main(): ai21_client = AI21Client() @@ -23,16 +27,20 @@ def main(): assistant_id=assistant.id, ) + start = time.time() + while run.status == "in_progress": run = ai21_client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id) + if time.time() - start > TIMEOUT: + break + time.sleep(1) if run.status == "completed": messages = ai21_client.beta.threads.messages.list(thread_id=thread.id) print("Messages:") print("\n".join(f"{msg.role}: {msg.content['text']}" for msg in messages.results)) else: - # handle error or required action - pass + raise Exception("Run failed") if __name__ == "__main__": diff --git a/examples/studio/assistant/async_assistant.py b/examples/studio/assistant/async_assistant.py index b4c84c23..a70681d8 100644 --- a/examples/studio/assistant/async_assistant.py +++ b/examples/studio/assistant/async_assistant.py @@ -1,8 +1,12 @@ import asyncio +import time from ai21 import AsyncAI21Client +TIMEOUT = 20 + + async def main(): ai21_client = AsyncAI21Client() @@ -25,13 +29,20 @@ async def main(): assistant_id=assistant.id, ) + start = time.time() + while run.status == "in_progress": run = await ai21_client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id) + if time.time() - start > TIMEOUT: + break + time.sleep(1) if run.status == "completed": messages = await ai21_client.beta.threads.messages.list(thread_id=thread.id) print("Messages:") print("\n".join(f"{msg.role}: {msg.content['text']}" for msg in messages.results)) + else: + raise Exception("Run failed") if __name__ == "__main__": From 3c3cfbf5ba5d08899ac143f4ff2e7885b32a3a8c Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 3 Dec 2024 15:32:59 +0200 Subject: [PATCH 12/13] chore: :truck: move files --- ai21/clients/common/{assistant => beta}/__init__.py | 0 .../{studio/resources => common/beta}/assistant/__init__.py | 0 ai21/clients/common/{ => beta}/assistant/assistants.py | 0 ai21/clients/common/{ => beta}/assistant/messages.py | 0 ai21/clients/common/{ => beta}/assistant/runs.py | 0 ai21/clients/common/{ => beta}/assistant/threads.py | 2 +- ai21/clients/studio/resources/beta/assistant/__init__.py | 0 .../studio_assistant.py => beta/assistant/assistant.py} | 2 +- .../studio_thread.py => beta/assistant/thread.py} | 6 +++--- .../assistant/thread_message.py} | 2 +- .../studio_thread_run.py => beta/assistant/thread_run.py} | 2 +- ai21/clients/studio/resources/beta/async_beta.py | 4 ++-- ai21/clients/studio/resources/beta/beta.py | 4 ++-- 13 files changed, 11 insertions(+), 11 deletions(-) rename ai21/clients/common/{assistant => beta}/__init__.py (100%) rename ai21/clients/{studio/resources => common/beta}/assistant/__init__.py (100%) rename ai21/clients/common/{ => beta}/assistant/assistants.py (100%) rename ai21/clients/common/{ => beta}/assistant/messages.py (100%) rename ai21/clients/common/{ => beta}/assistant/runs.py (100%) rename ai21/clients/common/{ => beta}/assistant/threads.py (88%) create mode 100644 ai21/clients/studio/resources/beta/assistant/__init__.py rename ai21/clients/studio/resources/{assistant/studio_assistant.py => beta/assistant/assistant.py} (98%) rename ai21/clients/studio/resources/{assistant/studio_thread.py => beta/assistant/thread.py} (86%) rename ai21/clients/studio/resources/{assistant/studio_thread_message.py => beta/assistant/thread_message.py} (95%) rename ai21/clients/studio/resources/{assistant/studio_thread_run.py => beta/assistant/thread_run.py} (98%) diff --git a/ai21/clients/common/assistant/__init__.py b/ai21/clients/common/beta/__init__.py similarity index 100% rename from ai21/clients/common/assistant/__init__.py rename to ai21/clients/common/beta/__init__.py diff --git a/ai21/clients/studio/resources/assistant/__init__.py b/ai21/clients/common/beta/assistant/__init__.py similarity index 100% rename from ai21/clients/studio/resources/assistant/__init__.py rename to ai21/clients/common/beta/assistant/__init__.py diff --git a/ai21/clients/common/assistant/assistants.py b/ai21/clients/common/beta/assistant/assistants.py similarity index 100% rename from ai21/clients/common/assistant/assistants.py rename to ai21/clients/common/beta/assistant/assistants.py diff --git a/ai21/clients/common/assistant/messages.py b/ai21/clients/common/beta/assistant/messages.py similarity index 100% rename from ai21/clients/common/assistant/messages.py rename to ai21/clients/common/beta/assistant/messages.py diff --git a/ai21/clients/common/assistant/runs.py b/ai21/clients/common/beta/assistant/runs.py similarity index 100% rename from ai21/clients/common/assistant/runs.py rename to ai21/clients/common/beta/assistant/runs.py diff --git a/ai21/clients/common/assistant/threads.py b/ai21/clients/common/beta/assistant/threads.py similarity index 88% rename from ai21/clients/common/assistant/threads.py rename to ai21/clients/common/beta/assistant/threads.py index e7359cc9..75674fc9 100644 --- a/ai21/clients/common/assistant/threads.py +++ b/ai21/clients/common/beta/assistant/threads.py @@ -3,7 +3,7 @@ from abc import ABC, abstractmethod from typing import List -from ai21.clients.common.assistant.messages import Messages +from ai21.clients.common.beta.assistant.messages import Messages from ai21.models.assistant.message import Message from ai21.models.responses.thread_response import ThreadResponse diff --git a/ai21/clients/studio/resources/beta/assistant/__init__.py b/ai21/clients/studio/resources/beta/assistant/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ai21/clients/studio/resources/assistant/studio_assistant.py b/ai21/clients/studio/resources/beta/assistant/assistant.py similarity index 98% rename from ai21/clients/studio/resources/assistant/studio_assistant.py rename to ai21/clients/studio/resources/beta/assistant/assistant.py index 259302bf..7b49a4c6 100644 --- a/ai21/clients/studio/resources/assistant/studio_assistant.py +++ b/ai21/clients/studio/resources/beta/assistant/assistant.py @@ -2,7 +2,7 @@ from typing import List -from ai21.clients.common.assistant.assistants import Assistants +from ai21.clients.common.beta.assistant.assistants import Assistants from ai21.clients.studio.resources.studio_resource import ( AsyncStudioResource, StudioResource, diff --git a/ai21/clients/studio/resources/assistant/studio_thread.py b/ai21/clients/studio/resources/beta/assistant/thread.py similarity index 86% rename from ai21/clients/studio/resources/assistant/studio_thread.py rename to ai21/clients/studio/resources/beta/assistant/thread.py index e2ac565a..b2fdf57b 100644 --- a/ai21/clients/studio/resources/assistant/studio_thread.py +++ b/ai21/clients/studio/resources/beta/assistant/thread.py @@ -2,9 +2,9 @@ from typing import List -from ai21.clients.common.assistant.threads import Threads -from ai21.clients.studio.resources.assistant.studio_thread_message import ThreadMessage, AsyncThreadMessage -from ai21.clients.studio.resources.assistant.studio_thread_run import AsyncThreadRun, ThreadRun +from ai21.clients.common.beta.assistant.threads import Threads +from ai21.clients.studio.resources.beta.assistant.thread_message import ThreadMessage, AsyncThreadMessage +from ai21.clients.studio.resources.beta.assistant.thread_run import AsyncThreadRun, ThreadRun from ai21.clients.studio.resources.studio_resource import StudioResource, AsyncStudioResource from ai21.http_client.async_http_client import AsyncAI21HTTPClient from ai21.http_client.http_client import AI21HTTPClient diff --git a/ai21/clients/studio/resources/assistant/studio_thread_message.py b/ai21/clients/studio/resources/beta/assistant/thread_message.py similarity index 95% rename from ai21/clients/studio/resources/assistant/studio_thread_message.py rename to ai21/clients/studio/resources/beta/assistant/thread_message.py index bc0494dd..f8a8ee10 100644 --- a/ai21/clients/studio/resources/assistant/studio_thread_message.py +++ b/ai21/clients/studio/resources/beta/assistant/thread_message.py @@ -1,6 +1,6 @@ from __future__ import annotations -from ai21.clients.common.assistant.messages import Messages +from ai21.clients.common.beta.assistant.messages import Messages from ai21.clients.studio.resources.studio_resource import StudioResource, AsyncStudioResource from ai21.models.assistant.message import ThreadMessageRole, MessageContentText from ai21.models.responses.message_response import MessageResponse, ListMessageResponse diff --git a/ai21/clients/studio/resources/assistant/studio_thread_run.py b/ai21/clients/studio/resources/beta/assistant/thread_run.py similarity index 98% rename from ai21/clients/studio/resources/assistant/studio_thread_run.py rename to ai21/clients/studio/resources/beta/assistant/thread_run.py index 092ba242..71955363 100644 --- a/ai21/clients/studio/resources/assistant/studio_thread_run.py +++ b/ai21/clients/studio/resources/beta/assistant/thread_run.py @@ -2,7 +2,7 @@ from typing import List -from ai21.clients.common.assistant.runs import Runs +from ai21.clients.common.beta.assistant.runs import Runs from ai21.clients.studio.resources.studio_resource import StudioResource, AsyncStudioResource from ai21.models.assistant.assistant import Optimization from ai21.models.assistant.run import ToolOutput diff --git a/ai21/clients/studio/resources/beta/async_beta.py b/ai21/clients/studio/resources/beta/async_beta.py index 9afbc4e0..5c93b9f6 100644 --- a/ai21/clients/studio/resources/beta/async_beta.py +++ b/ai21/clients/studio/resources/beta/async_beta.py @@ -1,5 +1,5 @@ -from ai21.clients.studio.resources.assistant.studio_assistant import AsyncAssistant -from ai21.clients.studio.resources.assistant.studio_thread import AsyncThread +from ai21.clients.studio.resources.beta.assistant.assistant import AsyncAssistant +from ai21.clients.studio.resources.beta.assistant.thread import AsyncThread from ai21.clients.studio.resources.studio_conversational_rag import AsyncStudioConversationalRag from ai21.clients.studio.resources.studio_resource import AsyncStudioResource from ai21.http_client.async_http_client import AsyncAI21HTTPClient diff --git a/ai21/clients/studio/resources/beta/beta.py b/ai21/clients/studio/resources/beta/beta.py index 18ed1c0c..62b55e80 100644 --- a/ai21/clients/studio/resources/beta/beta.py +++ b/ai21/clients/studio/resources/beta/beta.py @@ -1,5 +1,5 @@ -from ai21.clients.studio.resources.assistant.studio_assistant import Assistant -from ai21.clients.studio.resources.assistant.studio_thread import Thread +from ai21.clients.studio.resources.beta.assistant.assistant import Assistant +from ai21.clients.studio.resources.beta.assistant.thread import Thread from ai21.clients.studio.resources.studio_conversational_rag import StudioConversationalRag from ai21.clients.studio.resources.studio_resource import StudioResource from ai21.http_client.http_client import AI21HTTPClient From 196fc042c7e65f56ad68c86e5bbaa17d2a102e0d Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 3 Dec 2024 17:34:30 +0200 Subject: [PATCH 13/13] docs: :memo: examples fixes --- examples/studio/assistant/assistant.py | 2 +- examples/studio/assistant/async_assistant.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/studio/assistant/assistant.py b/examples/studio/assistant/assistant.py index 041c62fe..ebe3100f 100644 --- a/examples/studio/assistant/assistant.py +++ b/examples/studio/assistant/assistant.py @@ -40,7 +40,7 @@ def main(): print("Messages:") print("\n".join(f"{msg.role}: {msg.content['text']}" for msg in messages.results)) else: - raise Exception("Run failed") + raise Exception(f"Run failed. Status: {run.status}") if __name__ == "__main__": diff --git a/examples/studio/assistant/async_assistant.py b/examples/studio/assistant/async_assistant.py index a70681d8..3f68805f 100644 --- a/examples/studio/assistant/async_assistant.py +++ b/examples/studio/assistant/async_assistant.py @@ -42,7 +42,7 @@ async def main(): print("Messages:") print("\n".join(f"{msg.role}: {msg.content['text']}" for msg in messages.results)) else: - raise Exception("Run failed") + raise Exception(f"Run failed. Status: {run.status}") if __name__ == "__main__":