fix: 🚚 move classes and such

Author: benshuk

ai21/clients/common/assistant/assistants.py

@@ -3,12 +3,8 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List
 
-from ai21.models.assistant.assistant import Optimization, Tool
-from ai21.models.responses.assistant_response import (
-    Assistant,
-    ToolResources,
-    ListAssistant,
-)
+from ai21.models.assistant.assistant import Optimization, Tool, ToolResources
+from ai21.models.responses.assistant_response import Assistant, ListAssistant
 from ai21.types import NotGiven, NOT_GIVEN
 from ai21.utils.typing import remove_not_given
 

ai21/clients/common/assistant/threads.py

@@ -3,19 +3,15 @@
 from abc import ABC, abstractmethod
 from typing import List
 
-from ai21.models.assistant.thread_message import CreateThreadMessagePayload
+from ai21.models.assistant.message import Message
 from ai21.models.responses.thread_response import Thread
 
 
 class Threads(ABC):
     _module_name = "threads"
 
     @abstractmethod
-    def create(
-        self,
-        messages: List[CreateThreadMessagePayload],
-        **kwargs,
-    ) -> Thread:
+    def create(self, messages: List[Message], **kwargs) -> Thread:
         pass
 
     @abstractmethod

ai21/clients/studio/resources/assistant/studio_assistant.py

@@ -7,12 +7,8 @@
     AsyncStudioResource,
     StudioResource,
 )
-from ai21.models.responses.assistant_response import (
-    Assistant,
-    Tool,
-    ToolResources,
-    ListAssistant,
-)
+from ai21.models.assistant.assistant import Tool, ToolResources
+from ai21.models.responses.assistant_response import Assistant, ListAssistant
 from ai21.types import NotGiven, NOT_GIVEN
 
 

ai21/clients/studio/resources/assistant/studio_thread.py

@@ -4,15 +4,12 @@
 
 from ai21.clients.common.assistant.threads import Threads
 from ai21.clients.studio.resources.studio_resource import StudioResource, AsyncStudioResource
-from ai21.models.responses.thread_response import CreateThreadMessagePayload, Thread
+from ai21.models.assistant.message import Message
+from ai21.models.responses.thread_response import Thread
 
 
 class StudioThread(StudioResource, Threads):
-    def create(
-        self,
-        messages: List[CreateThreadMessagePayload],
-        **kwargs,
-    ) -> Thread:
+    def create(self, messages: List[Message], **kwargs) -> Thread:
         body = dict(messages=messages)
 
         return self._post(path=f"/{self._module_name}", body=body, response_cls=Thread)
@@ -22,11 +19,7 @@ def get(self, thread_id: str) -> Thread:
 
 
 class AsyncStudioThread(AsyncStudioResource, Threads):
-    async def create(
-        self,
-        messages: List[CreateThreadMessagePayload],
-        **kwargs,
-    ) -> Thread:
+    async def create(self, messages: List[Message], **kwargs) -> Thread:
         body = dict(messages=messages)
 
         return await self._post(path=f"/{self._module_name}", body=body, response_cls=Thread)

ai21/models/assistant/message.py

@@ -0,0 +1,29 @@
+from datetime import datetime
+from typing import Literal, Optional
+
+from typing_extensions import TypedDict
+
+from ai21.models.ai21_base_model import AI21BaseModel
+
+ThreadMessageRole = Literal["assistant", "user"]
+
+
+class MessageContentText(TypedDict):
+    type: Literal["text"]
+    text: str
+
+
+class Message(TypedDict):
+    role: ThreadMessageRole
+    content: MessageContentText
+
+
+class MessageResponse(AI21BaseModel):
+    id: str
+    created_at: datetime
+    updated_at: datetime
+    object: Literal["message"] = "message"
+    role: ThreadMessageRole
+    content: MessageContentText
+    run_id: Optional[str] = None
+    assistant_id: Optional[str] = None

ai21/models/assistant/thread_message.py

@@ -1,19 +1,7 @@
 from datetime import datetime
-from typing import Optional, List, Literal
+from typing import List, Literal
 
 from ai21.models.ai21_base_model import AI21BaseModel
-from ai21.models.assistant.thread_message import ThreadMessageRole, ThreadMessageContentText
-
-
-class ThreadMessage(AI21BaseModel):
-    id: str
-    created_at: datetime
-    updated_at: datetime
-    object: Literal["message"] = "message"
-    role: ThreadMessageRole
-    content: ThreadMessageContentText
-    run_id: Optional[str] = None
-    assistant_id: Optional[str] = None
 
 
 class Thread(AI21BaseModel):