From 18d2befae672972385c29983214e67410e5caa79 Mon Sep 17 00:00:00 2001 From: Ashpreet Bedi Date: Tue, 9 Jan 2024 11:31:54 +0000 Subject: [PATCH] openai assistant refactor --- phi/assistant/__init__.py | 1 - phi/assistant/file/__init__.py | 1 - phi/assistant/openai/__init__.py | 1 + phi/assistant/{ => openai}/assistant.py | 56 ++++++++++++------------ phi/assistant/{ => openai}/exceptions.py | 0 phi/assistant/openai/file/__init__.py | 1 + phi/assistant/{ => openai}/file/file.py | 2 +- phi/assistant/{ => openai}/file/local.py | 0 phi/assistant/{ => openai}/file/url.py | 2 +- phi/assistant/{ => openai}/message.py | 4 +- phi/assistant/{ => openai}/row.py | 8 ++-- phi/assistant/{ => openai}/run.py | 22 +++++----- phi/assistant/{ => openai}/thread.py | 20 ++++----- phi/assistant/{ => openai}/tool.py | 0 phi/assistant/python.py | 0 15 files changed, 59 insertions(+), 59 deletions(-) delete mode 100644 phi/assistant/__init__.py delete mode 100644 phi/assistant/file/__init__.py create mode 100644 phi/assistant/openai/__init__.py rename phi/assistant/{ => openai}/assistant.py (87%) rename phi/assistant/{ => openai}/exceptions.py (100%) create mode 100644 phi/assistant/openai/file/__init__.py rename phi/assistant/{ => openai}/file/file.py (99%) rename phi/assistant/{ => openai}/file/local.py (100%) rename phi/assistant/{ => openai}/file/url.py (97%) rename phi/assistant/{ => openai}/message.py (98%) rename phi/assistant/{ => openai}/row.py (91%) rename phi/assistant/{ => openai}/run.py (95%) rename phi/assistant/{ => openai}/thread.py (93%) rename phi/assistant/{ => openai}/tool.py (100%) create mode 100644 phi/assistant/python.py diff --git a/phi/assistant/__init__.py b/phi/assistant/__init__.py deleted file mode 100644 index a4af82d530..0000000000 --- a/phi/assistant/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from phi.assistant.assistant import Assistant diff --git a/phi/assistant/file/__init__.py b/phi/assistant/file/__init__.py deleted file mode 100644 index 68f5c3916e..0000000000 --- a/phi/assistant/file/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from phi.assistant.file.file import File diff --git a/phi/assistant/openai/__init__.py b/phi/assistant/openai/__init__.py new file mode 100644 index 0000000000..96b6f9a711 --- /dev/null +++ b/phi/assistant/openai/__init__.py @@ -0,0 +1 @@ +from phi.assistant.openai.assistant import OpenAiAssistant diff --git a/phi/assistant/assistant.py b/phi/assistant/openai/assistant.py similarity index 87% rename from phi/assistant/assistant.py rename to phi/assistant/openai/assistant.py index 9886bac180..a675f168c8 100644 --- a/phi/assistant/assistant.py +++ b/phi/assistant/openai/assistant.py @@ -3,8 +3,8 @@ from pydantic import BaseModel, ConfigDict, field_validator, model_validator -from phi.assistant.file import File -from phi.assistant.exceptions import AssistantIdNotSet +from phi.assistant.openai.file import File +from phi.assistant.openai.exceptions import AssistantIdNotSet from phi.tools import Tool, ToolRegistry from phi.tools.function import Function from phi.utils.log import logger, set_log_level_to_debug @@ -18,13 +18,13 @@ raise -class Assistant(BaseModel): +class OpenAiAssistant(BaseModel): # -*- LLM settings model: str = "gpt-4-1106-preview" openai: Optional[OpenAI] = None - # -*- Assistant settings - # Assistant id which can be referenced in API endpoints. + # -*- OpenAiAssistant settings + # OpenAiAssistant id which can be referenced in API endpoints. id: Optional[str] = None # The object type, populated by the API. Always assistant. object: Optional[str] = None @@ -35,15 +35,15 @@ class Assistant(BaseModel): # The system instructions that the assistant uses. The maximum length is 32768 characters. instructions: Optional[str] = None - # -*- Assistant Tools + # -*- OpenAiAssistant Tools # A list of tools provided to the assistant. There can be a maximum of 128 tools per assistant. # Tools can be of types code_interpreter, retrieval, or function. tools: Optional[List[Union[Tool, ToolRegistry, Callable, Dict]]] = None - # -*- Functions available to the Assistant to call + # -*- Functions available to the OpenAiAssistant to call # Functions extracted from the tools which can be executed locally by the assistant. functions: Optional[Dict[str, Function]] = None - # -*- Assistant Files + # -*- OpenAiAssistant Files # A list of file IDs attached to this assistant. # There can be a maximum of 20 files attached to the assistant. # Files are ordered by their creation date in ascending order. @@ -51,14 +51,14 @@ class Assistant(BaseModel): # Files attached to this assistant. files: Optional[List[File]] = None - # -*- Assistant Storage + # -*- OpenAiAssistant Storage # storage: Optional[AssistantStorage] = None # Create table if it doesn't exist # create_storage: bool = True # AssistantRow from the database: DO NOT SET THIS MANUALLY # database_row: Optional[AssistantRow] = None - # -*- Assistant Knowledge Base + # -*- OpenAiAssistant Knowledge Base # knowledge_base: Optional[KnowledgeBase] = None # Set of 16 key-value pairs that can be attached to an object. @@ -92,18 +92,18 @@ def client(self) -> OpenAI: return self.openai or OpenAI() @model_validator(mode="after") - def extract_functions_from_tools(self) -> "Assistant": + def extract_functions_from_tools(self) -> "OpenAiAssistant": if self.tools is not None: for tool in self.tools: if self.functions is None: self.functions = {} if isinstance(tool, ToolRegistry): self.functions.update(tool.functions) - logger.debug(f"Functions from {tool.name} added to Assistant.") + logger.debug(f"Functions from {tool.name} added to OpenAiAssistant.") elif callable(tool): f = Function.from_callable(tool) self.functions[f.name] = f - logger.debug(f"Function {f.name} added to Assistant") + logger.debug(f"Function {f.name} added to OpenAiAssistant") return self def __enter__(self): @@ -137,7 +137,7 @@ def get_tools_for_api(self) -> Optional[List[Dict[str, Any]]]: tools_for_api.append({"type": "function", "function": _f.to_dict()}) return tools_for_api - def create(self) -> "Assistant": + def create(self) -> "OpenAiAssistant": request_body: Dict[str, Any] = {} if self.name is not None: request_body["name"] = self.name @@ -163,7 +163,7 @@ def create(self) -> "Assistant": **request_body, ) self.load_from_openai(self.openai_assistant) - logger.debug(f"Assistant created: {self.id}") + logger.debug(f"OpenAiAssistant created: {self.id}") return self def get_id(self) -> Optional[str]: @@ -172,7 +172,7 @@ def get_id(self) -> Optional[str]: def get_from_openai(self) -> OpenAIAssistantType: _assistant_id = self.get_id() if _assistant_id is None: - raise AssistantIdNotSet("Assistant.id not set") + raise AssistantIdNotSet("OpenAiAssistant.id not set") self.openai_assistant = self.client.beta.assistants.retrieve( assistant_id=_assistant_id, @@ -180,20 +180,20 @@ def get_from_openai(self) -> OpenAIAssistantType: self.load_from_openai(self.openai_assistant) return self.openai_assistant - def get(self, use_cache: bool = True) -> "Assistant": + def get(self, use_cache: bool = True) -> "OpenAiAssistant": if self.openai_assistant is not None and use_cache: return self self.get_from_openai() return self - def get_or_create(self, use_cache: bool = True) -> "Assistant": + def get_or_create(self, use_cache: bool = True) -> "OpenAiAssistant": try: return self.get(use_cache=use_cache) except AssistantIdNotSet: return self.create() - def update(self) -> "Assistant": + def update(self) -> "OpenAiAssistant": try: assistant_to_update = self.get_from_openai() if assistant_to_update is not None: @@ -227,11 +227,11 @@ def update(self) -> "Assistant": **request_body, ) self.load_from_openai(self.openai_assistant) - logger.debug(f"Assistant updated: {self.id}") + logger.debug(f"OpenAiAssistant updated: {self.id}") return self - raise ValueError("Assistant not available") + raise ValueError("OpenAiAssistant not available") except AssistantIdNotSet: - logger.warning("Assistant not available") + logger.warning("OpenAiAssistant not available") raise def delete(self) -> OpenAIAssistantDeleted: @@ -241,10 +241,10 @@ def delete(self) -> OpenAIAssistantDeleted: deletion_status = self.client.beta.assistants.delete( assistant_id=assistant_to_delete.id, ) - logger.debug(f"Assistant deleted: {deletion_status.id}") + logger.debug(f"OpenAiAssistant deleted: {deletion_status.id}") return deletion_status except AssistantIdNotSet: - logger.warning("Assistant not available") + logger.warning("OpenAiAssistant not available") raise def to_dict(self) -> Dict[str, Any]: @@ -275,18 +275,18 @@ def __str__(self) -> str: return json.dumps(self.to_dict(), indent=4) def __repr__(self) -> str: - return f"" + return f"" # # def run(self, thread: Optional["Thread"]) -> "Thread": - # from phi.assistant.thread import Thread + # from phi.assistant.openai.thread import Thread # # return Thread(assistant=self, thread=thread).run() def print_response(self, message: str, markdown: bool = False) -> None: """Print a response from the assistant""" - from phi.assistant.thread import Thread + from phi.assistant.openai.thread import Thread thread = Thread() thread.print_response(message=message, assistant=self, markdown=markdown) @@ -300,7 +300,7 @@ def cli_app( exit_on: Tuple[str, ...] = ("exit", "bye"), ) -> None: from rich.prompt import Prompt - from phi.assistant.thread import Thread + from phi.assistant.openai.thread import Thread thread = Thread() while True: diff --git a/phi/assistant/exceptions.py b/phi/assistant/openai/exceptions.py similarity index 100% rename from phi/assistant/exceptions.py rename to phi/assistant/openai/exceptions.py diff --git a/phi/assistant/openai/file/__init__.py b/phi/assistant/openai/file/__init__.py new file mode 100644 index 0000000000..976eac5824 --- /dev/null +++ b/phi/assistant/openai/file/__init__.py @@ -0,0 +1 @@ +from phi.assistant.openai.file.file import File diff --git a/phi/assistant/file/file.py b/phi/assistant/openai/file/file.py similarity index 99% rename from phi/assistant/file/file.py rename to phi/assistant/openai/file/file.py index 871f7e0ce1..de2bafe460 100644 --- a/phi/assistant/file/file.py +++ b/phi/assistant/openai/file/file.py @@ -3,7 +3,7 @@ from pydantic import BaseModel, ConfigDict -from phi.assistant.exceptions import FileIdNotSet +from phi.assistant.openai.exceptions import FileIdNotSet from phi.utils.log import logger try: diff --git a/phi/assistant/file/local.py b/phi/assistant/openai/file/local.py similarity index 100% rename from phi/assistant/file/local.py rename to phi/assistant/openai/file/local.py diff --git a/phi/assistant/file/url.py b/phi/assistant/openai/file/url.py similarity index 97% rename from phi/assistant/file/url.py rename to phi/assistant/openai/file/url.py index d2bb949e6c..8e9e422400 100644 --- a/phi/assistant/file/url.py +++ b/phi/assistant/openai/file/url.py @@ -1,7 +1,7 @@ from pathlib import Path from typing import Any, Optional -from phi.assistant.file import File +from phi.assistant.openai.file import File from phi.utils.log import logger diff --git a/phi/assistant/message.py b/phi/assistant/openai/message.py similarity index 98% rename from phi/assistant/message.py rename to phi/assistant/openai/message.py index b7293038c1..b1e75091b1 100644 --- a/phi/assistant/message.py +++ b/phi/assistant/openai/message.py @@ -3,8 +3,8 @@ from pydantic import BaseModel, ConfigDict -from phi.assistant.file import File -from phi.assistant.exceptions import ThreadIdNotSet, MessageIdNotSet +from phi.assistant.openai.file import File +from phi.assistant.openai.exceptions import ThreadIdNotSet, MessageIdNotSet from phi.utils.log import logger try: diff --git a/phi/assistant/row.py b/phi/assistant/openai/row.py similarity index 91% rename from phi/assistant/row.py rename to phi/assistant/openai/row.py index 352f808a84..65ebe08b29 100644 --- a/phi/assistant/row.py +++ b/phi/assistant/openai/row.py @@ -4,9 +4,9 @@ class AssistantRow(BaseModel): - """Interface between Assistant class and the database""" + """Interface between OpenAiAssistant class and the database""" - # Assistant id which can be referenced in API endpoints. + # OpenAiAssistant id which can be referenced in API endpoints. id: str # The object type, which is always assistant. object: str @@ -18,13 +18,13 @@ class AssistantRow(BaseModel): instructions: Optional[str] = None # LLM data (name, model, etc.) llm: Optional[Dict[str, Any]] = None - # Assistant Tools + # OpenAiAssistant Tools tools: Optional[List[Dict[str, Any]]] = None # Files attached to this assistant. files: Optional[List[Dict[str, Any]]] = None # Metadata attached to this assistant. metadata: Optional[Dict[str, Any]] = None - # Assistant Memory + # OpenAiAssistant Memory memory: Optional[Dict[str, Any]] = None # True if this assistant is active is_active: Optional[bool] = None diff --git a/phi/assistant/run.py b/phi/assistant/openai/run.py similarity index 95% rename from phi/assistant/run.py rename to phi/assistant/openai/run.py index b37e778204..9b9b7f2ab9 100644 --- a/phi/assistant/run.py +++ b/phi/assistant/openai/run.py @@ -3,8 +3,8 @@ from pydantic import BaseModel, ConfigDict, model_validator -from phi.assistant.assistant import Assistant -from phi.assistant.exceptions import ThreadIdNotSet, AssistantIdNotSet, RunIdNotSet +from phi.assistant.openai.assistant import OpenAiAssistant +from phi.assistant.openai.exceptions import ThreadIdNotSet, AssistantIdNotSet, RunIdNotSet from phi.tools import Tool, ToolRegistry from phi.tools.function import Function from phi.utils.functions import get_function_call @@ -33,8 +33,8 @@ class Run(BaseModel): # The ID of the thread that was executed on as a part of this run. thread_id: Optional[str] = None - # Assistant used for this run - assistant: Optional[Assistant] = None + # OpenAiAssistant used for this run + assistant: Optional[OpenAiAssistant] = None # The ID of the assistant used for execution of this run. assistant_id: Optional[str] = None @@ -106,11 +106,11 @@ def extract_functions_from_tools(self) -> "Run": self.functions = {} if isinstance(tool, ToolRegistry): self.functions.update(tool.functions) - logger.debug(f"Functions from {tool.name} added to Assistant.") + logger.debug(f"Functions from {tool.name} added to OpenAiAssistant.") elif callable(tool): f = Function.from_callable(tool) self.functions[f.name] = f - logger.debug(f"Function {f.name} added to Assistant") + logger.debug(f"Function {f.name} added to OpenAiAssistant") return self def load_from_openai(self, openai_run: OpenAIRun): @@ -147,7 +147,7 @@ def get_tools_for_api(self) -> Optional[List[Dict[str, Any]]]: return tools_for_api def create( - self, thread_id: Optional[str] = None, assistant: Optional[Assistant] = None, assistant_id: Optional[str] = None + self, thread_id: Optional[str] = None, assistant: Optional[OpenAiAssistant] = None, assistant_id: Optional[str] = None ) -> "Run": _thread_id = thread_id or self.thread_id if _thread_id is None: @@ -157,7 +157,7 @@ def create( if _assistant_id is None: _assistant_id = self.assistant.get_id() if self.assistant is not None else self.assistant_id if _assistant_id is None: - raise AssistantIdNotSet("Assistant.id not set") + raise AssistantIdNotSet("OpenAiAssistant.id not set") request_body: Dict[str, Any] = {} if self.model is not None: @@ -206,7 +206,7 @@ def get_or_create( self, use_cache: bool = True, thread_id: Optional[str] = None, - assistant: Optional[Assistant] = None, + assistant: Optional[OpenAiAssistant] = None, assistant_id: Optional[str] = None, ) -> "Run": try: @@ -264,7 +264,7 @@ def wait( def run( self, thread_id: Optional[str] = None, - assistant: Optional[Assistant] = None, + assistant: Optional[OpenAiAssistant] = None, assistant_id: Optional[str] = None, wait: bool = True, callback: Optional[Callable[[OpenAIRun], None]] = None, @@ -284,7 +284,7 @@ def run( # -*- Check if run requires action if self.status == "requires_action": if self.assistant is None: - logger.warning("Assistant not available to complete required_action") + logger.warning("OpenAiAssistant not available to complete required_action") return self if self.required_action is not None: if self.required_action.type == "submit_tool_outputs": diff --git a/phi/assistant/thread.py b/phi/assistant/openai/thread.py similarity index 93% rename from phi/assistant/thread.py rename to phi/assistant/openai/thread.py index da0a7ac8dd..87ed57281b 100644 --- a/phi/assistant/thread.py +++ b/phi/assistant/openai/thread.py @@ -2,10 +2,10 @@ from pydantic import BaseModel, ConfigDict -from phi.assistant.run import Run -from phi.assistant.message import Message -from phi.assistant.assistant import Assistant -from phi.assistant.exceptions import ThreadIdNotSet +from phi.assistant.openai.run import Run +from phi.assistant.openai.message import Message +from phi.assistant.openai.assistant import OpenAiAssistant +from phi.assistant.openai.exceptions import ThreadIdNotSet from phi.utils.log import logger try: @@ -25,8 +25,8 @@ class Thread(BaseModel): # The object type, populated by the API. Always thread. object: Optional[str] = None - # Assistant used for this thread - assistant: Optional[Assistant] = None + # OpenAiAssistant used for this thread + assistant: Optional[OpenAiAssistant] = None # The ID of the assistant for this thread. assistant_id: Optional[str] = None @@ -162,7 +162,7 @@ def add(self, messages: List[Union[Message, Dict]]) -> None: def run( self, message: Optional[Union[str, Message]] = None, - assistant: Optional[Assistant] = None, + assistant: Optional[OpenAiAssistant] = None, assistant_id: Optional[str] = None, run: Optional[Run] = None, wait: bool = True, @@ -232,7 +232,7 @@ def print_messages(self) -> None: table.add_column("User") table.add_column(m.get_content_with_files()) elif m.role == "assistant": - table.add_row("Assistant", Markdown(m.get_content_with_files())) + table.add_row("OpenAiAssistant", Markdown(m.get_content_with_files())) table.add_section() else: table.add_row(m.role, Markdown(m.get_content_with_files())) @@ -240,7 +240,7 @@ def print_messages(self) -> None: console.print(table) def print_response( - self, message: str, assistant: Assistant, current_message_only: bool = False, markdown: bool = False + self, message: str, assistant: OpenAiAssistant, current_message_only: bool = False, markdown: bool = False ) -> None: from rich.progress import Progress, SpinnerColumn, TextColumn @@ -263,7 +263,7 @@ def print_response( total_messages = len(response_messages) for idx, response_message in enumerate(response_messages[::-1], start=1): response_message.pprint( - title=f"[bold] :robot: Assistant ({idx}/{total_messages}) [/bold]", markdown=markdown + title=f"[bold] :robot: OpenAiAssistant ({idx}/{total_messages}) [/bold]", markdown=markdown ) else: for m in self.messages[::-1]: diff --git a/phi/assistant/tool.py b/phi/assistant/openai/tool.py similarity index 100% rename from phi/assistant/tool.py rename to phi/assistant/openai/tool.py diff --git a/phi/assistant/python.py b/phi/assistant/python.py new file mode 100644 index 0000000000..e69de29bb2