Spaces:
Runtime error
Runtime error
| import json | |
| import os | |
| import subprocess | |
| import sys | |
| import time | |
| from typing import Any, Callable, Dict, List, Optional | |
| from loguru import logger | |
| from swarms.structs.agent import Agent | |
| def check_openai_package(): | |
| """Check if the OpenAI package is installed, and install it if not.""" | |
| try: | |
| import openai | |
| return openai | |
| except ImportError: | |
| logger.info( | |
| "OpenAI package not found. Attempting to install..." | |
| ) | |
| # Attempt to install the OpenAI package | |
| try: | |
| subprocess.check_call( | |
| [sys.executable, "-m", "pip", "install", "openai"] | |
| ) | |
| logger.info("OpenAI package installed successfully.") | |
| import openai # Re-import the package after installation | |
| return openai | |
| except subprocess.CalledProcessError as e: | |
| logger.error(f"Failed to install OpenAI package: {e}") | |
| raise RuntimeError( | |
| "OpenAI package installation failed." | |
| ) from e | |
| class OpenAIAssistant(Agent): | |
| """ | |
| OpenAI Assistant wrapper for the swarms framework. | |
| Integrates OpenAI's Assistants API with the swarms architecture. | |
| Example: | |
| >>> assistant = OpenAIAssistant( | |
| ... name="Math Tutor", | |
| ... instructions="You are a personal math tutor.", | |
| ... model="gpt-4o", | |
| ... tools=[{"type": "code_interpreter"}] | |
| ... ) | |
| >>> response = assistant.run("Solve 3x + 11 = 14") | |
| """ | |
| def __init__( | |
| self, | |
| name: str, | |
| description: str = "Standard openai assistant wrapper", | |
| instructions: Optional[str] = None, | |
| model: str = "gpt-4o", | |
| tools: Optional[List[Dict[str, Any]]] = None, | |
| file_ids: Optional[List[str]] = None, | |
| metadata: Optional[Dict[str, Any]] = None, | |
| functions: Optional[List[Dict[str, Any]]] = None, | |
| *args, | |
| **kwargs, | |
| ): | |
| """Initialize the OpenAI Assistant. | |
| Args: | |
| name: Name of the assistant | |
| instructions: System instructions for the assistant | |
| model: Model to use (default: gpt-4o) | |
| tools: List of tools to enable (code_interpreter, retrieval) | |
| file_ids: List of file IDs to attach | |
| metadata: Additional metadata | |
| functions: List of custom functions to make available | |
| """ | |
| self.name = name | |
| self.description = description | |
| self.instructions = instructions | |
| self.model = model | |
| self.tools = tools | |
| self.file_ids = file_ids | |
| self.metadata = metadata | |
| self.functions = functions | |
| super().__init__(*args, **kwargs) | |
| # Initialize tools list with any provided functions | |
| self.tools = tools or [] | |
| if functions: | |
| for func in functions: | |
| self.tools.append( | |
| {"type": "function", "function": func} | |
| ) | |
| # Create the OpenAI Assistant | |
| openai = check_openai_package() | |
| self.client = openai.OpenAI( | |
| api_key=os.getenv("OPENAI_API_KEY") | |
| ) | |
| self.assistant = self.client.beta.assistants.create( | |
| name=name, | |
| instructions=instructions, | |
| model=model, | |
| tools=self.tools, | |
| # file_ids=file_ids or [], | |
| metadata=metadata or {}, | |
| ) | |
| # Store available functions | |
| self.available_functions: Dict[str, Callable] = {} | |
| def add_function( | |
| self, | |
| func: Callable, | |
| description: str, | |
| parameters: Dict[str, Any], | |
| ) -> None: | |
| """Add a function that the assistant can call. | |
| Args: | |
| func: The function to make available to the assistant | |
| description: Description of what the function does | |
| parameters: JSON schema describing the function parameters | |
| """ | |
| func_dict = { | |
| "name": func.__name__, | |
| "description": description, | |
| "parameters": parameters, | |
| } | |
| # Add to tools list | |
| self.tools.append({"type": "function", "function": func_dict}) | |
| # Store function reference | |
| self.available_functions[func.__name__] = func | |
| # Update assistant with new tools | |
| self.assistant = self.client.beta.assistants.update( | |
| assistant_id=self.assistant.id, tools=self.tools | |
| ) | |
| def _handle_tool_calls(self, run, thread_id: str) -> None: | |
| """Handle any required tool calls during a run. | |
| This method processes any tool calls required by the assistant during execution. | |
| It extracts function calls, executes them with provided arguments, and submits | |
| the results back to the assistant. | |
| Args: | |
| run: The current run object from the OpenAI API | |
| thread_id: ID of the current conversation thread | |
| Returns: | |
| Updated run object after processing tool calls | |
| Raises: | |
| Exception: If there are errors executing the tool calls | |
| """ | |
| while run.status == "requires_action": | |
| tool_calls = ( | |
| run.required_action.submit_tool_outputs.tool_calls | |
| ) | |
| tool_outputs = [] | |
| for tool_call in tool_calls: | |
| if tool_call.type == "function": | |
| # Get function details | |
| function_name = tool_call.function.name | |
| function_args = json.loads( | |
| tool_call.function.arguments | |
| ) | |
| # Call function if available | |
| if function_name in self.available_functions: | |
| function_response = self.available_functions[ | |
| function_name | |
| ](**function_args) | |
| tool_outputs.append( | |
| { | |
| "tool_call_id": tool_call.id, | |
| "output": str(function_response), | |
| } | |
| ) | |
| # Submit outputs back to the run | |
| run = self.client.beta.threads.runs.submit_tool_outputs( | |
| thread_id=thread_id, | |
| run_id=run.id, | |
| tool_outputs=tool_outputs, | |
| ) | |
| # Wait for processing | |
| run = self._wait_for_run(run) | |
| return run | |
| def _wait_for_run(self, run) -> Any: | |
| """Wait for a run to complete and handle any required actions. | |
| This method polls the OpenAI API to check the status of a run until it completes | |
| or fails. It handles intermediate states like required actions and implements | |
| exponential backoff. | |
| Args: | |
| run: The run object to monitor | |
| Returns: | |
| The completed run object | |
| Raises: | |
| Exception: If the run fails or expires | |
| """ | |
| while True: | |
| run = self.client.beta.threads.runs.retrieve( | |
| thread_id=run.thread_id, run_id=run.id | |
| ) | |
| if run.status == "completed": | |
| break | |
| elif run.status == "requires_action": | |
| run = self._handle_tool_calls(run, run.thread_id) | |
| if run.status == "completed": | |
| break | |
| elif run.status in ["failed", "expired"]: | |
| raise Exception( | |
| f"Run failed with status: {run.status}" | |
| ) | |
| time.sleep(3) # Wait 3 seconds before checking again | |
| return run | |
| def _ensure_thread(self): | |
| """Ensure a thread exists for the conversation. | |
| This method checks if there is an active thread for the current conversation. | |
| If no thread exists, it creates a new one. This maintains conversation context | |
| across multiple interactions. | |
| Side Effects: | |
| Sets self.thread if it doesn't exist | |
| """ | |
| self.thread = self.client.beta.threads.create() | |
| def add_message( | |
| self, content: str, file_ids: Optional[List[str]] = None | |
| ) -> None: | |
| """Add a message to the thread. | |
| This method adds a new user message to the conversation thread. It ensures | |
| a thread exists before adding the message and handles file attachments. | |
| Args: | |
| content: The text content of the message to add | |
| file_ids: Optional list of file IDs to attach to the message. These must be | |
| files that have been previously uploaded to OpenAI. | |
| Side Effects: | |
| Creates a new thread if none exists | |
| Adds the message to the thread in OpenAI's system | |
| """ | |
| self._ensure_thread() | |
| self.client.beta.threads.messages.create( | |
| thread_id=self.thread.id, | |
| role="user", | |
| content=content, | |
| # file_ids=file_ids or [], | |
| ) | |
| def _get_response(self) -> str: | |
| """Get the latest assistant response from the thread.""" | |
| messages = self.client.beta.threads.messages.list( | |
| thread_id=self.thread.id, order="desc", limit=1 | |
| ) | |
| if not messages.data: | |
| return "" | |
| message = messages.data[0] | |
| if message.role == "assistant": | |
| return message.content[0].text.value | |
| return "" | |
| def run(self, task: str, *args, **kwargs) -> str: | |
| """Run a task using the OpenAI Assistant. | |
| Args: | |
| task: The task or prompt to send to the assistant | |
| Returns: | |
| The assistant's response as a string | |
| """ | |
| self._ensure_thread() | |
| # Add the user message | |
| self.add_message(task) | |
| # Create and run the assistant | |
| run = self.client.beta.threads.runs.create( | |
| thread_id=self.thread.id, | |
| assistant_id=self.assistant.id, | |
| instructions=self.instructions, | |
| ) | |
| # Wait for completion | |
| run = self._wait_for_run(run) | |
| # Only get and return the response if run completed successfully | |
| if run.status == "completed": | |
| return self._get_response() | |
| return "" | |
| def call(self, task: str, *args, **kwargs) -> str: | |
| """Alias for run() to maintain compatibility with different agent interfaces.""" | |
| return self.run(task, *args, **kwargs) | |