cai/swarm/types.py

55 lines
1.6 KiB
Python

"""
This module contains type definitions for the Swarm library.
"""
from typing import List, Callable, Union, Optional
from openai.types.chat import ChatCompletionMessage # noqa: F401, E501 # pylint: disable=import-error, unused-import
from openai.types.chat.chat_completion_message_tool_call import ( # noqa: F401, E501 # pylint: disable=import-error, unused-import
ChatCompletionMessageToolCall,
Function,
)
# Third-party imports
from pydantic import BaseModel # pylint: disable=import-error
AgentFunction = Callable[[], Union[str, "Agent", dict]]
class Agent(BaseModel): # pylint: disable=too-few-public-methods
"""
Represents an agent in the Swarm.
"""
name: str = "Agent"
# model: str = "gpt-4o"
model: str = "qwen2.5:14b"
instructions: Union[str, Callable[[], str]] = "You are a helpful agent."
functions: List[AgentFunction] = []
tool_choice: str = None
parallel_tool_calls: bool = True
class Response(BaseModel): # pylint: disable=too-few-public-methods
"""
Represents a response from the Swarm.
"""
messages: List = []
agent: Optional[Agent] = None
context_variables: dict = {}
class Result(BaseModel): # pylint: disable=too-few-public-methods
"""
Encapsulates the possible return values for an agent function.
Attributes:
value (str): The result value as a string.
agent (Agent): The agent instance, if applicable.
context_variables (dict): A dictionary of context variables.
"""
value: str = ""
agent: Optional[Agent] = None
context_variables: dict = {}