diff --git a/atomic_agents/agents/base_agent.py b/atomic_agents/agents/base_agent.py index 4e2b425c..c96ecfdc 100644 --- a/atomic_agents/agents/base_agent.py +++ b/atomic_agents/agents/base_agent.py @@ -8,7 +8,7 @@ from atomic_agents.lib.components.system_prompt_generator import SystemPromptContextProviderBase, SystemPromptGenerator -class BaseAgentIO(BaseModel): +class BaseIOSchema(BaseModel): """ Base class for input and output schemas for chat agents. """ @@ -21,7 +21,7 @@ def __rich__(self): return JSON(json_str) -class BaseAgentInputSchema(BaseAgentIO): +class BaseAgentInputSchema(BaseIOSchema): chat_message: str = Field( ..., description="The chat message sent by the user to the assistant.", @@ -36,7 +36,7 @@ class Config: } -class BaseAgentOutputSchema(BaseAgentIO): +class BaseAgentOutputSchema(BaseIOSchema): chat_message: str = Field( ..., description=( @@ -56,7 +56,7 @@ class Config: class BaseAgentConfig(BaseModel): client: instructor.client.Instructor = Field(..., description="Client for interacting with the language model.") - model: str = Field("gpt-3.5-turbo", description="The model to use for generating responses.") + model: str = Field("gpt-4o-mini", description="The model to use for generating responses.") memory: Optional[AgentMemory] = Field(None, description="Memory component for storing chat history.") system_prompt_generator: Optional[SystemPromptGenerator] = Field( None, description="Component for generating system prompts." @@ -76,8 +76,8 @@ class BaseAgent: generating system prompts, and obtaining responses from a language model. Attributes: - input_schema (Type[BaseAgentIO]): Schema for the input data. - output_schema (Type[BaseAgentIO]): Schema for the output data. + input_schema (Type[BaseIOSchema]): Schema for the input data. + output_schema (Type[BaseIOSchema]): Schema for the output data. client: Client for interacting with the language model. model (str): The model to use for generating responses. memory (AgentMemory): Memory component for storing chat history. @@ -133,15 +133,15 @@ def get_response(self, response_model=None) -> Type[BaseModel]: response = self.client.chat.completions.create(model=self.model, messages=messages, response_model=response_model) return response - def run(self, user_input: Optional[Type[BaseAgentIO]] = None) -> Type[BaseAgentIO]: + def run(self, user_input: Optional[Type[BaseIOSchema]] = None) -> Type[BaseIOSchema]: """ Runs the chat agent with the given user input. Args: - user_input (Optional[Type[BaseAgentIO]]): The input from the user. If not provided, skips adding to memory. + user_input (Optional[Type[BaseIOSchema]]): The input from the user. If not provided, skips adding to memory. Returns: - Type[BaseAgentIO]: The response from the chat agent. + Type[BaseIOSchema]: The response from the chat agent. """ if user_input: self.current_user_input = user_input @@ -165,9 +165,9 @@ def get_context_provider(self, provider_name: str) -> Type[SystemPromptContextPr Raises: KeyError: If the context provider is not found. """ - if provider_name not in self.system_prompt_generator.system_prompt_info.context_providers: + if provider_name not in self.system_prompt_generator.context_providers: raise KeyError(f"Context provider '{provider_name}' not found.") - return self.system_prompt_generator.system_prompt_info.context_providers[provider_name] + return self.system_prompt_generator.context_providers[provider_name] def register_context_provider(self, provider_name: str, provider: SystemPromptContextProviderBase): """ @@ -177,7 +177,7 @@ def register_context_provider(self, provider_name: str, provider: SystemPromptCo provider_name (str): The name of the context provider. provider (SystemPromptContextProviderBase): The context provider instance. """ - self.system_prompt_generator.system_prompt_info.context_providers[provider_name] = provider + self.system_prompt_generator.context_providers[provider_name] = provider def unregister_context_provider(self, provider_name: str): """ @@ -186,7 +186,7 @@ def unregister_context_provider(self, provider_name: str): Args: provider_name (str): The name of the context provider to remove. """ - if provider_name in self.system_prompt_generator.system_prompt_info.context_providers: - del self.system_prompt_generator.system_prompt_info.context_providers[provider_name] + if provider_name in self.system_prompt_generator.context_providers: + del self.system_prompt_generator.context_providers[provider_name] else: raise KeyError(f"Context provider '{provider_name}' not found.") diff --git a/atomic_agents/agents/tool_interface_agent.py b/atomic_agents/agents/tool_interface_agent.py index b44907d4..d9e92659 100644 --- a/atomic_agents/agents/tool_interface_agent.py +++ b/atomic_agents/agents/tool_interface_agent.py @@ -2,9 +2,9 @@ from pydantic import Field, create_model -from atomic_agents.agents.base_agent import BaseAgentIO, BaseAgent, BaseAgentConfig -from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator, SystemPromptInfo -from atomic_agents.lib.tools.base import BaseTool +from atomic_agents.agents.base_agent import BaseIOSchema, BaseAgent, BaseAgentConfig +from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator +from atomic_agents.lib.tools.base_tool import BaseTool from atomic_agents.lib.utils.format_tool_message import format_tool_message @@ -13,18 +13,6 @@ class ToolInterfaceAgentConfig(BaseAgentConfig): return_raw_output: bool = False -class ToolInputModel(BaseAgentIO): - tool_input: str = Field(..., description="Tool input. Presented as a single question or instruction") - - class Config: - title = "Default Tool" - description = "Default tool description" - json_schema_extra = { - "title": "Default Tool", - "description": "Default tool description", - } - - class ToolInterfaceAgent(BaseAgent): """ A specialized chat agent designed to interact with a specific tool. @@ -60,7 +48,7 @@ def __init__(self, config: ToolInterfaceAgentConfig): alias=f"tool_input_{self.tool_instance.tool_name}", ), ), - __base__=ToolInputModel, + __base__=BaseIOSchema, ) # Manually set the configuration attributes @@ -86,19 +74,17 @@ def __init__(self, config: ToolInterfaceAgentConfig): output_instructions.append("Return the raw output of the tool.") self.system_prompt_generator = SystemPromptGenerator( - system_prompt_info=SystemPromptInfo( - background=[ - f"This AI agent is designed to interact with the {self.tool_instance.tool_name} tool.", - f"Tool description: {self.tool_instance.tool_description}", - ], - steps=[ - "Get the user input.", - "Convert the input to the proper parameters to call the tool.", - "Call the tool with the parameters.", - "Respond to the user", - ], - output_instructions=output_instructions, - ) + background=[ + f"This AI agent is designed to interact with the {self.tool_instance.tool_name} tool.", + f"Tool description: {self.tool_instance.tool_description}", + ], + steps=[ + "Get the user input.", + "Convert the input to the proper parameters to call the tool.", + "Call the tool with the parameters.", + "Respond to the user", + ], + output_instructions=output_instructions, ) def get_response(self, response_model=None): diff --git a/atomic_agents/lib/components/system_prompt_generator.py b/atomic_agents/lib/components/system_prompt_generator.py index 0f69f636..2f0f6c4a 100644 --- a/atomic_agents/lib/components/system_prompt_generator.py +++ b/atomic_agents/lib/components/system_prompt_generator.py @@ -1,5 +1,4 @@ from abc import ABC, abstractmethod -from dataclasses import dataclass, field from typing import Dict, List, Optional @@ -15,21 +14,20 @@ def __repr__(self) -> str: return self.get_info() -@dataclass -class SystemPromptInfo: - background: List[str] - steps: List[str] = field(default_factory=list) - output_instructions: List[str] = field(default_factory=list) - context_providers: Dict[str, SystemPromptContextProviderBase] = field(default_factory=dict) - - class SystemPromptGenerator: - def __init__(self, system_prompt_info: Optional[SystemPromptInfo] = None): - self.system_prompt_info = system_prompt_info or SystemPromptInfo( - background=["This is a conversation with a helpful and friendly AI assistant."] - ) - - self.system_prompt_info.output_instructions.extend( + def __init__( + self, + background: Optional[List[str]] = None, + steps: Optional[List[str]] = None, + output_instructions: Optional[List[str]] = None, + context_providers: Optional[Dict[str, SystemPromptContextProviderBase]] = None, + ): + self.background = background or ["This is a conversation with a helpful and friendly AI assistant."] + self.steps = steps or [] + self.output_instructions = output_instructions or [] + self.context_providers = context_providers or {} + + self.output_instructions.extend( [ "Always respond using the proper JSON schema.", "Always use the available additional information and context to enhance the response.", @@ -38,9 +36,9 @@ def __init__(self, system_prompt_info: Optional[SystemPromptInfo] = None): def generate_prompt(self) -> str: sections = [ - ("IDENTITY and PURPOSE", self.system_prompt_info.background), - ("INTERNAL ASSISTANT STEPS", self.system_prompt_info.steps), - ("OUTPUT INSTRUCTIONS", self.system_prompt_info.output_instructions), + ("IDENTITY and PURPOSE", self.background), + ("INTERNAL ASSISTANT STEPS", self.steps), + ("OUTPUT INSTRUCTIONS", self.output_instructions), ] prompt_parts = [] @@ -51,9 +49,9 @@ def generate_prompt(self) -> str: prompt_parts.extend(f"- {item}" for item in content) prompt_parts.append("") - if self.system_prompt_info.context_providers: + if self.context_providers: prompt_parts.append("# EXTRA INFORMATION AND CONTEXT") - for provider in self.system_prompt_info.context_providers.values(): + for provider in self.context_providers.values(): info = provider.get_info() if info: prompt_parts.append(f"## {provider.title}") diff --git a/atomic_agents/lib/models/web_document.py b/atomic_agents/lib/models/web_document.py index 38af39fe..f21c1e96 100644 --- a/atomic_agents/lib/models/web_document.py +++ b/atomic_agents/lib/models/web_document.py @@ -1,6 +1,6 @@ from pydantic import BaseModel, Field -from atomic_agents.agents.base_agent import BaseAgentIO +from atomic_agents.agents.base_agent import BaseIOSchema class WebDocumentMetadata(BaseModel): @@ -11,6 +11,6 @@ class WebDocumentMetadata(BaseModel): author: str = Field(default="") -class WebDocument(BaseAgentIO): +class WebDocument(BaseIOSchema): content: str metadata: WebDocumentMetadata = Field(default_factory=lambda: WebDocumentMetadata(url="")) diff --git a/atomic_agents/lib/tools/base.py b/atomic_agents/lib/tools/base_tool.py similarity index 71% rename from atomic_agents/lib/tools/base.py rename to atomic_agents/lib/tools/base_tool.py index 08401cec..b0a23b25 100644 --- a/atomic_agents/lib/tools/base.py +++ b/atomic_agents/lib/tools/base_tool.py @@ -2,7 +2,7 @@ from pydantic import BaseModel -from atomic_agents.agents.base_agent import BaseAgentIO +from atomic_agents.agents.base_agent import BaseIOSchema class BaseToolConfig(BaseModel): @@ -15,15 +15,15 @@ class BaseTool: Base class for all tools in the Atomic Agents framework. Attributes: - input_schema (Type[BaseAgentIO]): The schema for the input data. - output_schema (Type[BaseAgentIO]): The schema for the output data. + input_schema (Type[BaseIOSchema]): The schema for the input data. + output_schema (Type[BaseIOSchema]): The schema for the output data. tool_name (str): The name of the tool, derived from the input schema's title. tool_description (str): The description of the tool, derived from the input schema's description or overridden by the user. """ - input_schema: Type[BaseAgentIO] - output_schema: Type[BaseAgentIO] + input_schema: Type[BaseIOSchema] + output_schema: Type[BaseIOSchema] def __init__(self, config: BaseToolConfig = BaseToolConfig()): """ @@ -35,15 +35,15 @@ def __init__(self, config: BaseToolConfig = BaseToolConfig()): self.tool_name = config.title or self.input_schema.Config.title self.tool_description = config.description or self.input_schema.Config.description - def run(self, params: Type[BaseAgentIO]) -> BaseAgentIO: + def run(self, params: Type[BaseIOSchema]) -> BaseIOSchema: """ Runs the tool with the given parameters. This method should be implemented by subclasses. Args: - params (BaseAgentIO): The input parameters for the tool, adhering to the input schema. + params (BaseIOSchema): The input parameters for the tool, adhering to the input schema. Returns: - BaseAgentIO: The output of the tool, adhering to the output schema. + BaseIOSchema: The output of the tool, adhering to the output schema. Raises: NotImplementedError: If the method is not implemented by a subclass. diff --git a/atomic_agents/lib/tools/calculator_tool.py b/atomic_agents/lib/tools/calculator_tool.py index 2a9a85e2..a87c0721 100644 --- a/atomic_agents/lib/tools/calculator_tool.py +++ b/atomic_agents/lib/tools/calculator_tool.py @@ -2,15 +2,14 @@ from rich.console import Console from sympy import sympify -from atomic_agents.agents.base_agent import BaseAgentIO -from atomic_agents.lib.tools.base import BaseTool, BaseToolConfig +from atomic_agents.agents.base_agent import BaseIOSchema +from atomic_agents.lib.tools.base_tool import BaseTool, BaseToolConfig + ################ # INPUT SCHEMA # ################ - - -class CalculatorToolSchema(BaseAgentIO): +class CalculatorToolInputSchema(BaseIOSchema): expression: str = Field(..., description="Mathematical expression to evaluate. For example, '2 + 2'.") class Config: @@ -27,9 +26,7 @@ class Config: #################### # OUTPUT SCHEMA(S) # #################### - - -class CalculatorToolOutputSchema(BaseAgentIO): +class CalculatorToolOutputSchema(BaseIOSchema): result: str = Field(..., description="Result of the calculation.") @@ -45,11 +42,11 @@ class CalculatorTool(BaseTool): Tool for performing calculations based on the provided mathematical expression. Attributes: - input_schema (CalculatorToolSchema): The schema for the input data. + input_schema (CalculatorToolInputSchema): The schema for the input data. output_schema (CalculatorToolOutputSchema): The schema for the output data. """ - input_schema = CalculatorToolSchema + input_schema = CalculatorToolInputSchema output_schema = CalculatorToolOutputSchema def __init__(self, config: CalculatorToolConfig = CalculatorToolConfig()): @@ -61,12 +58,12 @@ def __init__(self, config: CalculatorToolConfig = CalculatorToolConfig()): """ super().__init__(config) - def run(self, params: CalculatorToolSchema) -> CalculatorToolOutputSchema: + def run(self, params: CalculatorToolInputSchema) -> CalculatorToolOutputSchema: """ Runs the CalculatorTool with the given parameters. Args: - params (CalculatorToolSchema): The input parameters for the tool, adhering to the input schema. + params (CalculatorToolInputSchema): The input parameters for the tool, adhering to the input schema. Returns: CalculatorToolOutputSchema: The output of the tool, adhering to the output schema. @@ -83,4 +80,4 @@ def run(self, params: CalculatorToolSchema) -> CalculatorToolOutputSchema: ################# if __name__ == "__main__": rich_console = Console() - rich_console.print(CalculatorTool().run(CalculatorToolSchema(expression="2 + 2"))) + rich_console.print(CalculatorTool().run(CalculatorToolInputSchema(expression="2 + 2"))) diff --git a/atomic_agents/lib/tools/content_scraping_tool.py b/atomic_agents/lib/tools/content_scraping_tool.py index 43cc5c3b..8aa17cc5 100644 --- a/atomic_agents/lib/tools/content_scraping_tool.py +++ b/atomic_agents/lib/tools/content_scraping_tool.py @@ -8,8 +8,8 @@ from rich.console import Console from rich.markdown import Markdown -from atomic_agents.agents.base_agent import BaseAgentIO -from atomic_agents.lib.tools.base import BaseTool, BaseToolConfig +from atomic_agents.agents.base_agent import BaseIOSchema +from atomic_agents.lib.tools.base_tool import BaseTool, BaseToolConfig from atomic_agents.lib.utils.scraping.pdf_to_markdown import PdfToMarkdownConverter from atomic_agents.lib.utils.scraping.url_to_markdown import UrlToMarkdownConverter @@ -18,7 +18,7 @@ ################ -class ContentScrapingToolSchema(BaseAgentIO): +class ContentScrapingToolInputSchema(BaseIOSchema): url: str = Field(..., description="URL of the web page or PDF to scrape.") class Config: @@ -35,7 +35,7 @@ class Config: #################### -class ContentScrapingToolOutputSchema(BaseAgentIO): +class ContentScrapingToolOutputSchema(BaseIOSchema): content: str metadata: Optional[dict] = None @@ -52,11 +52,11 @@ class ContentScrapingTool(BaseTool): Tool for scraping web pages or PDFs and converting content to markdown. Attributes: - input_schema (ContentScrapingToolSchema): The schema for the input data. + input_schema (ContentScrapingToolInputSchema): The schema for the input data. output_schema (ContentScrapingToolOutputSchema): The schema for the output data. """ - input_schema = ContentScrapingToolSchema + input_schema = ContentScrapingToolInputSchema output_schema = ContentScrapingToolOutputSchema def __init__(self, config: ContentScrapingToolConfig = ContentScrapingToolConfig()): @@ -68,12 +68,12 @@ def __init__(self, config: ContentScrapingToolConfig = ContentScrapingToolConfig """ super().__init__(config) - def run(self, params: ContentScrapingToolSchema) -> ContentScrapingToolOutputSchema: + def run(self, params: ContentScrapingToolInputSchema) -> ContentScrapingToolOutputSchema: """ Runs the ContentScrapingTool with the given parameters. Args: - params (ContentScrapingToolSchema): The input parameters for the tool, adhering to the input schema. + params (ContentScrapingToolInputSchema): The input parameters for the tool, adhering to the input schema. Returns: ContentScrapingToolOutputSchema: The output of the tool, adhering to the output schema. @@ -102,7 +102,7 @@ def run(self, params: ContentScrapingToolSchema) -> ContentScrapingToolOutputSch # TEST WEB PAGE # ################# result = client.chat.completions.create( - model="gpt-3.5-turbo", + model="gpt-4o-mini", response_model=ContentScrapingTool.input_schema, messages=[ { @@ -119,7 +119,7 @@ def run(self, params: ContentScrapingToolSchema) -> ContentScrapingToolOutputSch # TEST PDF URL # ################ result = client.chat.completions.create( - model="gpt-3.5-turbo", + model="gpt-4o-mini", response_model=ContentScrapingTool.input_schema, messages=[ { diff --git a/atomic_agents/lib/tools/search/searx_tool.py b/atomic_agents/lib/tools/search/searxng_tool.py similarity index 73% rename from atomic_agents/lib/tools/search/searx_tool.py rename to atomic_agents/lib/tools/search/searxng_tool.py index e1fa4c66..40cfa02d 100644 --- a/atomic_agents/lib/tools/search/searx_tool.py +++ b/atomic_agents/lib/tools/search/searxng_tool.py @@ -5,22 +5,22 @@ from pydantic import Field from rich.console import Console -from atomic_agents.agents.base_agent import BaseAgentIO -from atomic_agents.lib.tools.base import BaseTool, BaseToolConfig +from atomic_agents.agents.base_agent import BaseIOSchema +from atomic_agents.lib.tools.base_tool import BaseTool, BaseToolConfig ################ # INPUT SCHEMA # ################ -class SearxNGSearchToolSchema(BaseAgentIO): +class SearxNGToolInputSchema(BaseIOSchema): queries: List[str] = Field(..., description="List of search queries.") category: Optional[Literal["general", "news", "social_media"]] = Field( "general", description="Category of the search queries." ) class Config: - title = "SearxNGSearchTool" + title = "SearxNGTool" description = ( "Tool for searching for information, news, references, and other content on the SearxNG. " "Returns a list of search results with a short description or content snippet and URLs for further exploration." @@ -31,65 +31,61 @@ class Config: #################### # OUTPUT SCHEMA(S) # #################### - - -class SearxNGSearchResultSchema(BaseAgentIO): +class SearxNGResultSchema(BaseIOSchema): url: str title: str content: Optional[str] = None -class SearxNGSearchToolOutputSchema(BaseAgentIO): - results: List[SearxNGSearchResultSchema] +class SearxNGToolOutputSchema(BaseIOSchema): + results: List[SearxNGResultSchema] category: Optional[str] = None ############## # TOOL LOGIC # ############## - - -class SearxNGSearchToolConfig(BaseToolConfig): +class SearxNGToolConfig(BaseToolConfig): base_url: str = "" max_results: int = 10 -class SearxNGSearchTool(BaseTool): +class SearxNGTool(BaseTool): """ Tool for performing searches on SearxNG based on the provided queries and category. Attributes: - input_schema (SearxNGSearchToolSchema): The schema for the input data. - output_schema (SearxNGSearchToolOutputSchema): The schema for the output data. + input_schema (SearxNGToolInputSchema): The schema for the input data. + output_schema (SearxNGToolOutputSchema): The schema for the output data. max_results (int): The maximum number of search results to return. base_url (str): The base URL for the SearxNG instance to use. """ - input_schema = SearxNGSearchToolSchema - output_schema = SearxNGSearchToolOutputSchema + input_schema = SearxNGToolInputSchema + output_schema = SearxNGToolOutputSchema - def __init__(self, config: SearxNGSearchToolConfig = SearxNGSearchToolConfig()): + def __init__(self, config: SearxNGToolConfig = SearxNGToolConfig()): """ - Initializes the SearxNGSearchTool. + Initializes the SearxNGTool. Args: - config (SearxNGSearchToolConfig): + config (SearxNGToolConfig): Configuration for the tool, including base URL, max results, and optional title and description overrides. """ super().__init__(config) self.base_url = config.base_url self.max_results = config.max_results - def run(self, params: SearxNGSearchToolSchema, max_results: Optional[int] = None) -> SearxNGSearchToolOutputSchema: + def run(self, params: SearxNGToolInputSchema, max_results: Optional[int] = None) -> SearxNGToolOutputSchema: """ - Runs the SearxNGSearchTool with the given parameters. + Runs the SearxNGTool with the given parameters. Args: - params (SearxNGSearchToolSchema): The input parameters for the tool, adhering to the input schema. + params (SearxNGToolInputSchema): The input parameters for the tool, adhering to the input schema. max_results (Optional[int]): The maximum number of search results to return. Returns: - SearxNGSearchToolOutputSchema: The output of the tool, adhering to the output schema. + SearxNGToolOutputSchema: The output of the tool, adhering to the output schema. Raises: ValueError: If the base URL is not provided. @@ -148,7 +144,7 @@ def run(self, params: SearxNGSearchToolSchema, max_results: Optional[int] = None filtered_results = filtered_results[: max_results or self.max_results] - return SearxNGSearchToolOutputSchema(results=filtered_results, category=params.category) + return SearxNGToolOutputSchema(results=filtered_results, category=params.category) ################# @@ -156,11 +152,9 @@ def run(self, params: SearxNGSearchToolSchema, max_results: Optional[int] = None ################# if __name__ == "__main__": rich_console = Console() - search_tool_instance = SearxNGSearchTool( - config=SearxNGSearchToolConfig(base_url=os.getenv("SEARXNG_BASE_URL"), max_results=5) - ) + search_tool_instance = SearxNGTool(config=SearxNGToolConfig(base_url=os.getenv("SEARXNG_BASE_URL"), max_results=5)) - search_input = SearxNGSearchTool.input_schema( + search_input = SearxNGTool.input_schema( queries=["Python programming", "Machine learning", "Artificial intelligence"], category="news", ) diff --git a/atomic_agents/lib/tools/search/serper_tool.py b/atomic_agents/lib/tools/search/serper_tool.py index 73ef6f3e..43df5557 100644 --- a/atomic_agents/lib/tools/search/serper_tool.py +++ b/atomic_agents/lib/tools/search/serper_tool.py @@ -6,19 +6,19 @@ from pydantic import Field from rich.console import Console -from atomic_agents.agents.base_agent import BaseAgentIO -from atomic_agents.lib.tools.base import BaseTool, BaseToolConfig +from atomic_agents.agents.base_agent import BaseIOSchema +from atomic_agents.lib.tools.base_tool import BaseTool, BaseToolConfig ################ # INPUT SCHEMA # ################ -class SerperSearchToolSchema(BaseAgentIO): +class SerperToolSchema(BaseIOSchema): queries: List[str] = Field(..., description="List of search queries.") class Config: - title = "SerperSearchTool" + title = "SerperTool" description = ( "Tool for searching for information, news, references, and other content using the Serper API. " "Returns a list of search results with a short description or content snippet and URLs for further exploration." @@ -31,15 +31,15 @@ class Config: #################### -class SerperSearchResultSchema(BaseAgentIO): +class SerperResultSchema(BaseIOSchema): url: str title: str content: Optional[str] = None position: Optional[int] = None -class SerperSearchToolOutputSchema(BaseAgentIO): - results: List[SerperSearchResultSchema] +class SerperToolOutputSchema(BaseIOSchema): + results: List[SerperResultSchema] ############## @@ -47,47 +47,47 @@ class SerperSearchToolOutputSchema(BaseAgentIO): ############## -class SerperSearchToolConfig(BaseToolConfig): +class SerperToolConfig(BaseToolConfig): api_key: str = "" max_results: int = 10 -class SerperSearchTool(BaseTool): +class SerperTool(BaseTool): """ Tool for performing searches using the Serper API based on the provided queries. Attributes: - input_schema (SerperSearchToolSchema): The schema for the input data. - output_schema (SerperSearchToolOutputSchema): The schema for the output data. + input_schema (SerperToolSchema): The schema for the input data. + output_schema (SerperToolOutputSchema): The schema for the output data. api_key (str): The API key for the Serper API. max_results (int): The maximum number of search results to return. """ - input_schema = SerperSearchToolSchema - output_schema = SerperSearchToolOutputSchema + input_schema = SerperToolSchema + output_schema = SerperToolOutputSchema - def __init__(self, config: SerperSearchToolConfig = SerperSearchToolConfig()): + def __init__(self, config: SerperToolConfig = SerperToolConfig()): """ - Initializes the SerperSearchTool. + Initializes the SerperTool. Args: - config (SerperSearchToolConfig): + config (SerperToolConfig): Configuration for the tool, including API key, max results, and optional title and description overrides. """ super().__init__(config) self.api_key = config.api_key self.max_results = config.max_results - def run(self, params: SerperSearchToolSchema, max_results: Optional[int] = None) -> SerperSearchToolOutputSchema: + def run(self, params: SerperToolSchema, max_results: Optional[int] = None) -> SerperToolOutputSchema: """ - Runs the SerperSearchTool with the given parameters. + Runs the SerperTool with the given parameters. Args: - params (SerperSearchToolSchema): The input parameters for the tool, adhering to the input schema. + params (SerperToolSchema): The input parameters for the tool, adhering to the input schema. max_results (Optional[int]): The maximum number of search results to return. Returns: - SerperSearchToolOutputSchema: The output of the tool, adhering to the output schema. + SerperToolOutputSchema: The output of the tool, adhering to the output schema. Raises: ValueError: If the API key is not provided. @@ -129,7 +129,7 @@ def run(self, params: SerperSearchToolSchema, max_results: Optional[int] = None) # Convert to output schema format output_results = [ - SerperSearchResultSchema( + SerperResultSchema( url=result["link"], title=result["title"], content=result.get("snippet"), @@ -138,7 +138,7 @@ def run(self, params: SerperSearchToolSchema, max_results: Optional[int] = None) for result in filtered_results ] - return SerperSearchToolOutputSchema(results=output_results) + return SerperToolOutputSchema(results=output_results) ################# @@ -146,9 +146,9 @@ def run(self, params: SerperSearchToolSchema, max_results: Optional[int] = None) ################# if __name__ == "__main__": rich_console = Console() - search_tool_instance = SerperSearchTool(config=SerperSearchToolConfig(api_key=os.getenv("SERPER_API_KEY"), max_results=5)) + search_tool_instance = SerperTool(config=SerperToolConfig(api_key=os.getenv("SERPER_API_KEY"), max_results=5)) - search_input = SerperSearchTool.input_schema(queries=["Python programming", "Machine learning", "Quantum computing"]) + search_input = SerperTool.input_schema(queries=["Python programming", "Machine learning", "Quantum computing"]) output = search_tool_instance.run(search_input) diff --git a/atomic_agents/lib/tools/yelp_restaurant_finder_tool.py b/atomic_agents/lib/tools/yelp_restaurant_finder_tool.py index 1fcdd372..3f9a3bee 100644 --- a/atomic_agents/lib/tools/yelp_restaurant_finder_tool.py +++ b/atomic_agents/lib/tools/yelp_restaurant_finder_tool.py @@ -6,8 +6,8 @@ from pydantic import Field from rich.console import Console -from atomic_agents.agents.base_agent import BaseAgentIO -from atomic_agents.lib.tools.base import BaseTool, BaseToolConfig +from atomic_agents.agents.base_agent import BaseIOSchema +from atomic_agents.lib.tools.base_tool import BaseTool, BaseToolConfig ################ # INPUT SCHEMA # @@ -43,7 +43,7 @@ class PriceRange(Enum): FOUR = "4" -class YelpSearchToolSchema(BaseAgentIO): +class YelpSearchToolInputSchema(BaseIOSchema): location: str = Field(..., description="Location to search for food.") term: Optional[str] = Field(None, description="Search term (e.g., 'pizza', 'sushi').") categories: Optional[List[YelpCategory]] = Field(None, description="Categories to filter by (e.g., 'italian, mexican').") @@ -75,7 +75,7 @@ class Config: #################### -class YelpSearchResultSchema(BaseAgentIO): +class YelpSearchResultSchema(BaseIOSchema): name: str url: str rating: float @@ -85,7 +85,7 @@ class YelpSearchResultSchema(BaseAgentIO): categories: List[str] -class YelpSearchToolOutputSchema(BaseAgentIO): +class YelpSearchToolOutputSchema(BaseIOSchema): results: List[YelpSearchResultSchema] class Config: @@ -112,13 +112,13 @@ class YelpSearchTool(BaseTool): Tool for performing searches using the Yelp API based on the provided queries. Attributes: - input_schema (YelpSearchToolSchema): The schema for the input data. + input_schema (YelpSearchToolInputSchema): The schema for the input data. output_schema (YelpSearchToolOutputSchema): The schema for the output data. api_key (str): The API key for the Yelp API. max_results (int): The maximum number of search results to return. """ - input_schema = YelpSearchToolSchema + input_schema = YelpSearchToolInputSchema output_schema = YelpSearchToolOutputSchema def __init__(self, config: YelpSearchToolConfig = YelpSearchToolConfig()): @@ -133,12 +133,12 @@ def __init__(self, config: YelpSearchToolConfig = YelpSearchToolConfig()): self.api_key = config.api_key self.max_results = config.max_results - def run(self, params: YelpSearchToolSchema) -> YelpSearchToolOutputSchema: + def run(self, params: YelpSearchToolInputSchema) -> YelpSearchToolOutputSchema: """ Runs the YelpSearchTool with the given parameters. Args: - params (YelpSearchToolSchema): The input parameters for the tool, adhering to the input schema. + params (YelpSearchToolInputSchema): The input parameters for the tool, adhering to the input schema. max_results (Optional[int]): The maximum number of search results to return. Returns: diff --git a/atomic_agents/lib/tools/yt_transcript_scraper.py b/atomic_agents/lib/tools/yt_transcript_scraper_tool.py similarity index 89% rename from atomic_agents/lib/tools/yt_transcript_scraper.py rename to atomic_agents/lib/tools/yt_transcript_scraper_tool.py index 8d67ed4e..6410a5f1 100644 --- a/atomic_agents/lib/tools/yt_transcript_scraper.py +++ b/atomic_agents/lib/tools/yt_transcript_scraper_tool.py @@ -6,15 +6,15 @@ from rich.console import Console from youtube_transcript_api import NoTranscriptFound, TranscriptsDisabled, YouTubeTranscriptApi -from atomic_agents.agents.base_agent import BaseAgentIO -from atomic_agents.lib.tools.base import BaseTool, BaseToolConfig +from atomic_agents.agents.base_agent import BaseIOSchema +from atomic_agents.lib.tools.base_tool import BaseTool, BaseToolConfig ################ # INPUT SCHEMA # ################ -class YouTubeTranscriptToolSchema(BaseAgentIO): +class YouTubeTranscriptToolInputSchema(BaseIOSchema): video_url: str = Field(..., description="URL of the YouTube video to fetch the transcript for.") language: Optional[str] = Field(None, description="Language code for the transcript (e.g., 'en' for English).") @@ -32,7 +32,7 @@ class Config: #################### -class YouTubeTranscriptToolOutputSchema(BaseAgentIO): +class YouTubeTranscriptToolOutputSchema(BaseIOSchema): transcript: str duration: float comments: List[str] @@ -61,11 +61,11 @@ class YouTubeTranscriptTool(BaseTool): Tool for fetching the transcript of a YouTube video using the YouTube Transcript API. Attributes: - input_schema (YouTubeTranscriptToolSchema): The schema for the input data. + input_schema (YouTubeTranscriptToolInputSchema): The schema for the input data. output_schema (YouTubeTranscriptToolOutputSchema): The schema for the output data. """ - input_schema = YouTubeTranscriptToolSchema + input_schema = YouTubeTranscriptToolInputSchema output_schema = YouTubeTranscriptToolOutputSchema def __init__(self, config: YouTubeTranscriptToolConfig): @@ -78,12 +78,12 @@ def __init__(self, config: YouTubeTranscriptToolConfig): super().__init__(config) self.api_key = config.api_key - def run(self, params: YouTubeTranscriptToolSchema) -> YouTubeTranscriptToolOutputSchema: + def run(self, params: YouTubeTranscriptToolInputSchema) -> YouTubeTranscriptToolOutputSchema: """ Runs the YouTubeTranscriptTool with the given parameters. Args: - params (YouTubeTranscriptToolSchema): The input parameters for the tool, adhering to the input schema. + params (YouTubeTranscriptToolInputSchema): The input parameters for the tool, adhering to the input schema. Returns: YouTubeTranscriptToolOutputSchema: The output of the tool, adhering to the output schema. diff --git a/docs/api/README.md b/docs/api/README.md index 13094043..e7a95638 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -15,11 +15,11 @@ A full list of `Atomic_agents` project modules. - [Models](atomic_agents/lib/models/index.md#models) - [WebDocument](atomic_agents/lib/models/web_document.md#webdocument) - [Tools](atomic_agents/lib/tools/index.md#tools) - - [Base](atomic_agents/lib/tools/base.md#base) + - [BaseTool](atomic_agents/lib/tools/base_tool.md#basetool) - [CalculatorTool](atomic_agents/lib/tools/calculator_tool.md#calculatortool) - [ContentScrapingTool](atomic_agents/lib/tools/content_scraping_tool.md#contentscrapingtool) - [Yelp Restaurant Finder Tool](atomic_agents/lib/tools/yelp_restaurant_finder_tool.md#yelp-restaurant-finder-tool) - - [Yt Transcript Scraper](atomic_agents/lib/tools/yt_transcript_scraper.md#yt-transcript-scraper) + - [Yt Transcript Scraper Tool](atomic_agents/lib/tools/yt_transcript_scraper_tool.md#yt-transcript-scraper-tool) - [Utils](atomic_agents/lib/utils/index.md#utils) - [Format Tool Message](atomic_agents/lib/utils/format_tool_message.md#format-tool-message) - [Scraping](atomic_agents/lib/utils/scraping/index.md#scraping) diff --git a/docs/api/atomic_agents/agents/base_agent.md b/docs/api/atomic_agents/agents/base_agent.md index 82e97ab1..29ac9c69 100644 --- a/docs/api/atomic_agents/agents/base_agent.md +++ b/docs/api/atomic_agents/agents/base_agent.md @@ -13,9 +13,9 @@ - [BaseAgent().run](#baseagent()run) - [BaseAgent().unregister_context_provider](#baseagent()unregister_context_provider) - [BaseAgentConfig](#baseagentconfig) - - [BaseAgentIO](#baseagentio) - [BaseAgentInputSchema](#baseagentinputschema) - [BaseAgentOutputSchema](#baseagentoutputschema) + - [BaseIOSchema](#baseioschema) ## BaseAgent @@ -28,8 +28,8 @@ generating system prompts, and obtaining responses from a language model. #### Attributes -- `input_schema` *Type[BaseAgentIO]* - Schema for the input data. -- `output_schema` *Type[BaseAgentIO]* - Schema for the output data. +- `input_schema` *Type[BaseIOSchema]* - Schema for the input data. +- `output_schema` *Type[BaseIOSchema]* - Schema for the output data. - `client` - Client for interacting with the language model. - `model` *str* - The model to use for generating responses. - `memory` *AgentMemory* - Memory component for storing chat history. @@ -141,21 +141,21 @@ Runs the chat agent with the given user input. #### Arguments -- `user_input` *Optional[Type[BaseAgentIO]]* - The input from the user. If not provided, skips adding to memory. +- `user_input` *Optional[Type[BaseIOSchema]]* - The input from the user. If not provided, skips adding to memory. #### Returns -- `Type[BaseAgentIO]` - The response from the chat agent. +- `Type[BaseIOSchema]` - The response from the chat agent. #### Signature ```python -def run(self, user_input: Optional[Type[BaseAgentIO]] = None) -> Type[BaseAgentIO]: ... +def run(self, user_input: Optional[Type[BaseIOSchema]] = None) -> Type[BaseIOSchema]: ... ``` #### See also -- [BaseAgentIO](#baseagentio) +- [BaseIOSchema](#baseioschema) ### BaseAgent().unregister_context_provider @@ -187,46 +187,46 @@ class BaseAgentConfig(BaseModel): ... -## BaseAgentIO - -[Show source in base_agent.py:11](../../../../atomic_agents/agents/base_agent.py#L11) +## BaseAgentInputSchema -Base class for input and output schemas for chat agents. +[Show source in base_agent.py:24](../../../../atomic_agents/agents/base_agent.py#L24) #### Signature ```python -class BaseAgentIO(BaseModel): ... +class BaseAgentInputSchema(BaseIOSchema): ... ``` +#### See also +- [BaseIOSchema](#baseioschema) -## BaseAgentInputSchema -[Show source in base_agent.py:24](../../../../atomic_agents/agents/base_agent.py#L24) + +## BaseAgentOutputSchema + +[Show source in base_agent.py:39](../../../../atomic_agents/agents/base_agent.py#L39) #### Signature ```python -class BaseAgentInputSchema(BaseAgentIO): ... +class BaseAgentOutputSchema(BaseIOSchema): ... ``` #### See also -- [BaseAgentIO](#baseagentio) +- [BaseIOSchema](#baseioschema) -## BaseAgentOutputSchema +## BaseIOSchema -[Show source in base_agent.py:39](../../../../atomic_agents/agents/base_agent.py#L39) +[Show source in base_agent.py:11](../../../../atomic_agents/agents/base_agent.py#L11) + +Base class for input and output schemas for chat agents. #### Signature ```python -class BaseAgentOutputSchema(BaseAgentIO): ... -``` - -#### See also - -- [BaseAgentIO](#baseagentio) \ No newline at end of file +class BaseIOSchema(BaseModel): ... +``` \ No newline at end of file diff --git a/docs/api/atomic_agents/agents/tool_interface_agent.md b/docs/api/atomic_agents/agents/tool_interface_agent.md index c2144b47..609d6fcb 100644 --- a/docs/api/atomic_agents/agents/tool_interface_agent.md +++ b/docs/api/atomic_agents/agents/tool_interface_agent.md @@ -5,30 +5,13 @@ > Auto-generated documentation for [atomic_agents.agents.tool_interface_agent](../../../../atomic_agents/agents/tool_interface_agent.py) module. - [ToolInterfaceAgent](#toolinterfaceagent) - - [ToolInputModel](#toolinputmodel) - [ToolInterfaceAgent](#toolinterfaceagent-1) - - [ToolInterfaceAgent()._get_and_handle_response](#toolinterfaceagent()_get_and_handle_response) + - [ToolInterfaceAgent().get_response](#toolinterfaceagent()get_response) - [ToolInterfaceAgentConfig](#toolinterfaceagentconfig) -## ToolInputModel - -[Show source in tool_interface_agent.py:16](../../../../atomic_agents/agents/tool_interface_agent.py#L16) - -#### Signature - -```python -class ToolInputModel(BaseAgentIO): ... -``` - -#### See also - -- [BaseAgentIO](./base_agent.md#baseagentio) - - - ## ToolInterfaceAgent -[Show source in tool_interface_agent.py:28](../../../../atomic_agents/agents/tool_interface_agent.py#L28) +[Show source in tool_interface_agent.py:16](../../../../atomic_agents/agents/tool_interface_agent.py#L16) A specialized chat agent designed to interact with a specific tool. @@ -52,9 +35,9 @@ class ToolInterfaceAgent(BaseAgent): - [BaseAgent](./base_agent.md#baseagent) - [ToolInterfaceAgentConfig](#toolinterfaceagentconfig) -### ToolInterfaceAgent()._get_and_handle_response +### ToolInterfaceAgent().get_response -[Show source in tool_interface_agent.py:104](../../../../atomic_agents/agents/tool_interface_agent.py#L104) +[Show source in tool_interface_agent.py:90](../../../../atomic_agents/agents/tool_interface_agent.py#L90) Handles obtaining and processing the response from the tool. @@ -69,7 +52,7 @@ the raw tool output; otherwise, it processes the output and returns the response #### Signature ```python -def _get_and_handle_response(self): ... +def get_response(self, response_model=None): ... ``` diff --git a/docs/api/atomic_agents/lib/components/system_prompt_generator.md b/docs/api/atomic_agents/lib/components/system_prompt_generator.md index 1eb4ba60..465afc3d 100644 --- a/docs/api/atomic_agents/lib/components/system_prompt_generator.md +++ b/docs/api/atomic_agents/lib/components/system_prompt_generator.md @@ -9,11 +9,10 @@ - [SystemPromptContextProviderBase().get_info](#systempromptcontextproviderbase()get_info) - [SystemPromptGenerator](#systempromptgenerator-1) - [SystemPromptGenerator().generate_prompt](#systempromptgenerator()generate_prompt) - - [SystemPromptInfo](#systempromptinfo) ## SystemPromptContextProviderBase -[Show source in system_prompt_generator.py:6](../../../../../atomic_agents/lib/components/system_prompt_generator.py#L6) +[Show source in system_prompt_generator.py:5](../../../../../atomic_agents/lib/components/system_prompt_generator.py#L5) #### Signature @@ -24,7 +23,7 @@ class SystemPromptContextProviderBase(ABC): ### SystemPromptContextProviderBase().get_info -[Show source in system_prompt_generator.py:10](../../../../../atomic_agents/lib/components/system_prompt_generator.py#L10) +[Show source in system_prompt_generator.py:9](../../../../../atomic_agents/lib/components/system_prompt_generator.py#L9) #### Signature @@ -37,37 +36,31 @@ def get_info(self) -> str: ... ## SystemPromptGenerator -[Show source in system_prompt_generator.py:26](../../../../../atomic_agents/lib/components/system_prompt_generator.py#L26) +[Show source in system_prompt_generator.py:17](../../../../../atomic_agents/lib/components/system_prompt_generator.py#L17) #### Signature ```python class SystemPromptGenerator: - def __init__(self, system_prompt_info: Optional[SystemPromptInfo] = None): ... + def __init__( + self, + background: Optional[List[str]] = None, + steps: Optional[List[str]] = None, + output_instructions: Optional[List[str]] = None, + context_providers: Optional[Dict[str, SystemPromptContextProviderBase]] = None, + ): ... ``` #### See also -- [SystemPromptInfo](#systempromptinfo) +- [SystemPromptContextProviderBase](#systempromptcontextproviderbase) ### SystemPromptGenerator().generate_prompt -[Show source in system_prompt_generator.py:39](../../../../../atomic_agents/lib/components/system_prompt_generator.py#L39) +[Show source in system_prompt_generator.py:37](../../../../../atomic_agents/lib/components/system_prompt_generator.py#L37) #### Signature ```python def generate_prompt(self) -> str: ... -``` - - - -## SystemPromptInfo - -[Show source in system_prompt_generator.py:19](../../../../../atomic_agents/lib/components/system_prompt_generator.py#L19) - -#### Signature - -```python -class SystemPromptInfo: ... ``` \ No newline at end of file diff --git a/docs/api/atomic_agents/lib/models/web_document.md b/docs/api/atomic_agents/lib/models/web_document.md index 677672c9..157149ea 100644 --- a/docs/api/atomic_agents/lib/models/web_document.md +++ b/docs/api/atomic_agents/lib/models/web_document.md @@ -15,12 +15,12 @@ #### Signature ```python -class WebDocument(BaseAgentIO): ... +class WebDocument(BaseIOSchema): ... ``` #### See also -- [BaseAgentIO](../../agents/base_agent.md#baseagentio) +- [BaseIOSchema](../../agents/base_agent.md#baseioschema) diff --git a/docs/api/atomic_agents/lib/tools/base.md b/docs/api/atomic_agents/lib/tools/base_tool.md similarity index 55% rename from docs/api/atomic_agents/lib/tools/base.md rename to docs/api/atomic_agents/lib/tools/base_tool.md index 8f226661..b7e53653 100644 --- a/docs/api/atomic_agents/lib/tools/base.md +++ b/docs/api/atomic_agents/lib/tools/base_tool.md @@ -1,24 +1,24 @@ -# Base +# BaseTool -[Atomic_agents Index](../../../README.md#atomic_agents-index) / [Atomic Agents](../../index.md#atomic-agents) / [Lib](../index.md#lib) / [Tools](./index.md#tools) / Base +[Atomic_agents Index](../../../README.md#atomic_agents-index) / [Atomic Agents](../../index.md#atomic-agents) / [Lib](../index.md#lib) / [Tools](./index.md#tools) / BaseTool -> Auto-generated documentation for [atomic_agents.lib.tools.base](../../../../../atomic_agents/lib/tools/base.py) module. +> Auto-generated documentation for [atomic_agents.lib.tools.base_tool](../../../../../atomic_agents/lib/tools/base_tool.py) module. -- [Base](#base) - - [BaseTool](#basetool) +- [BaseTool](#basetool) + - [BaseTool](#basetool-1) - [BaseTool().run](#basetool()run) - [BaseToolConfig](#basetoolconfig) ## BaseTool -[Show source in base.py:13](../../../../../atomic_agents/lib/tools/base.py#L13) +[Show source in base_tool.py:13](../../../../../atomic_agents/lib/tools/base_tool.py#L13) Base class for all tools in the Atomic Agents framework. #### Attributes -- `input_schema` *Type[BaseAgentIO]* - The schema for the input data. -- `output_schema` *Type[BaseAgentIO]* - The schema for the output data. +- `input_schema` *Type[BaseIOSchema]* - The schema for the input data. +- `output_schema` *Type[BaseIOSchema]* - The schema for the output data. - `tool_name` *str* - The name of the tool, derived from the input schema's title. tool_description (str): The description of the tool, derived from the input schema's description or overridden by the user. @@ -36,17 +36,17 @@ class BaseTool: ### BaseTool().run -[Show source in base.py:38](../../../../../atomic_agents/lib/tools/base.py#L38) +[Show source in base_tool.py:38](../../../../../atomic_agents/lib/tools/base_tool.py#L38) Runs the tool with the given parameters. This method should be implemented by subclasses. #### Arguments -- `params` *BaseAgentIO* - The input parameters for the tool, adhering to the input schema. +- `params` *BaseIOSchema* - The input parameters for the tool, adhering to the input schema. #### Returns -- `BaseAgentIO` - The output of the tool, adhering to the output schema. +- `BaseIOSchema` - The output of the tool, adhering to the output schema. #### Raises @@ -55,18 +55,18 @@ Runs the tool with the given parameters. This method should be implemented by su #### Signature ```python -def run(self, params: Type[BaseAgentIO]) -> BaseAgentIO: ... +def run(self, params: Type[BaseIOSchema]) -> BaseIOSchema: ... ``` #### See also -- [BaseAgentIO](../../agents/base_agent.md#baseagentio) +- [BaseIOSchema](../../agents/base_agent.md#baseioschema) ## BaseToolConfig -[Show source in base.py:8](../../../../../atomic_agents/lib/tools/base.py#L8) +[Show source in base_tool.py:8](../../../../../atomic_agents/lib/tools/base_tool.py#L8) #### Signature diff --git a/docs/api/atomic_agents/lib/tools/calculator_tool.md b/docs/api/atomic_agents/lib/tools/calculator_tool.md index c3403a37..cbd0f6a3 100644 --- a/docs/api/atomic_agents/lib/tools/calculator_tool.md +++ b/docs/api/atomic_agents/lib/tools/calculator_tool.md @@ -8,18 +8,18 @@ - [CalculatorTool](#calculatortool-1) - [CalculatorTool().run](#calculatortool()run) - [CalculatorToolConfig](#calculatortoolconfig) + - [CalculatorToolInputSchema](#calculatortoolinputschema) - [CalculatorToolOutputSchema](#calculatortooloutputschema) - - [CalculatorToolSchema](#calculatortoolschema) ## CalculatorTool -[Show source in calculator_tool.py:43](../../../../../atomic_agents/lib/tools/calculator_tool.py#L43) +[Show source in calculator_tool.py:40](../../../../../atomic_agents/lib/tools/calculator_tool.py#L40) Tool for performing calculations based on the provided mathematical expression. #### Attributes -- `input_schema` *CalculatorToolSchema* - The schema for the input data. +- `input_schema` *CalculatorToolInputSchema* - The schema for the input data. - `output_schema` *CalculatorToolOutputSchema* - The schema for the output data. #### Signature @@ -31,18 +31,18 @@ class CalculatorTool(BaseTool): #### See also -- [BaseTool](./base.md#basetool) +- [BaseTool](./base_tool.md#basetool) - [CalculatorToolConfig](#calculatortoolconfig) ### CalculatorTool().run -[Show source in calculator_tool.py:64](../../../../../atomic_agents/lib/tools/calculator_tool.py#L64) +[Show source in calculator_tool.py:61](../../../../../atomic_agents/lib/tools/calculator_tool.py#L61) Runs the CalculatorTool with the given parameters. #### Arguments -- `params` *CalculatorToolSchema* - The input parameters for the tool, adhering to the input schema. +- `params` *CalculatorToolInputSchema* - The input parameters for the tool, adhering to the input schema. #### Returns @@ -51,19 +51,19 @@ Runs the CalculatorTool with the given parameters. #### Signature ```python -def run(self, params: CalculatorToolSchema) -> CalculatorToolOutputSchema: ... +def run(self, params: CalculatorToolInputSchema) -> CalculatorToolOutputSchema: ... ``` #### See also +- [CalculatorToolInputSchema](#calculatortoolinputschema) - [CalculatorToolOutputSchema](#calculatortooloutputschema) -- [CalculatorToolSchema](#calculatortoolschema) ## CalculatorToolConfig -[Show source in calculator_tool.py:39](../../../../../atomic_agents/lib/tools/calculator_tool.py#L39) +[Show source in calculator_tool.py:36](../../../../../atomic_agents/lib/tools/calculator_tool.py#L36) #### Signature @@ -73,36 +73,36 @@ class CalculatorToolConfig(BaseToolConfig): ... #### See also -- [BaseToolConfig](./base.md#basetoolconfig) +- [BaseToolConfig](./base_tool.md#basetoolconfig) -## CalculatorToolOutputSchema +## CalculatorToolInputSchema -[Show source in calculator_tool.py:32](../../../../../atomic_agents/lib/tools/calculator_tool.py#L32) +[Show source in calculator_tool.py:12](../../../../../atomic_agents/lib/tools/calculator_tool.py#L12) #### Signature ```python -class CalculatorToolOutputSchema(BaseAgentIO): ... +class CalculatorToolInputSchema(BaseIOSchema): ... ``` #### See also -- [BaseAgentIO](../../agents/base_agent.md#baseagentio) +- [BaseIOSchema](../../agents/base_agent.md#baseioschema) -## CalculatorToolSchema +## CalculatorToolOutputSchema -[Show source in calculator_tool.py:13](../../../../../atomic_agents/lib/tools/calculator_tool.py#L13) +[Show source in calculator_tool.py:29](../../../../../atomic_agents/lib/tools/calculator_tool.py#L29) #### Signature ```python -class CalculatorToolSchema(BaseAgentIO): ... +class CalculatorToolOutputSchema(BaseIOSchema): ... ``` #### See also -- [BaseAgentIO](../../agents/base_agent.md#baseagentio) \ No newline at end of file +- [BaseIOSchema](../../agents/base_agent.md#baseioschema) \ No newline at end of file diff --git a/docs/api/atomic_agents/lib/tools/content_scraping_tool.md b/docs/api/atomic_agents/lib/tools/content_scraping_tool.md index 54adb965..f435934c 100644 --- a/docs/api/atomic_agents/lib/tools/content_scraping_tool.md +++ b/docs/api/atomic_agents/lib/tools/content_scraping_tool.md @@ -8,19 +8,19 @@ - `result` - ################ TEST WEB PAGE # - ################: client.chat.completions.create(model='gpt-3.5-turbo', response_model=ContentScrapingTool.input_schema, messages=[{'role': 'user', 'content': 'Scrape the content of https://brainblendai.com'}]) + ################: client.chat.completions.create(model='gpt-4o-mini', response_model=ContentScrapingTool.input_schema, messages=[{'role': 'user', 'content': 'Scrape the content of https://brainblendai.com'}]) - `result` - ############### TEST PDF URL # - ###############: client.chat.completions.create(model='gpt-3.5-turbo', response_model=ContentScrapingTool.input_schema, messages=[{'role': 'user', 'content': 'Scrape the content of https://pdfobject.com/pdf/sample.pdf'}]) + ###############: client.chat.completions.create(model='gpt-4o-mini', response_model=ContentScrapingTool.input_schema, messages=[{'role': 'user', 'content': 'Scrape the content of https://pdfobject.com/pdf/sample.pdf'}]) - [ContentScrapingTool](#contentscrapingtool) - [ContentScrapingTool](#contentscrapingtool-1) - [ContentScrapingTool().run](#contentscrapingtool()run) - [ContentScrapingToolConfig](#contentscrapingtoolconfig) + - [ContentScrapingToolInputSchema](#contentscrapingtoolinputschema) - [ContentScrapingToolOutputSchema](#contentscrapingtooloutputschema) - - [ContentScrapingToolSchema](#contentscrapingtoolschema) ## ContentScrapingTool @@ -30,7 +30,7 @@ Tool for scraping web pages or PDFs and converting content to markdown. #### Attributes -- `input_schema` *ContentScrapingToolSchema* - The schema for the input data. +- `input_schema` *ContentScrapingToolInputSchema* - The schema for the input data. - `output_schema` *ContentScrapingToolOutputSchema* - The schema for the output data. #### Signature @@ -44,7 +44,7 @@ class ContentScrapingTool(BaseTool): #### See also -- [BaseTool](./base.md#basetool) +- [BaseTool](./base_tool.md#basetool) - [ContentScrapingToolConfig](#contentscrapingtoolconfig) ### ContentScrapingTool().run @@ -55,7 +55,7 @@ Runs the ContentScrapingTool with the given parameters. #### Arguments -- `params` *ContentScrapingToolSchema* - The input parameters for the tool, adhering to the input schema. +- `params` *ContentScrapingToolInputSchema* - The input parameters for the tool, adhering to the input schema. #### Returns @@ -64,13 +64,15 @@ Runs the ContentScrapingTool with the given parameters. #### Signature ```python -def run(self, params: ContentScrapingToolSchema) -> ContentScrapingToolOutputSchema: ... +def run( + self, params: ContentScrapingToolInputSchema +) -> ContentScrapingToolOutputSchema: ... ``` #### See also +- [ContentScrapingToolInputSchema](#contentscrapingtoolinputschema) - [ContentScrapingToolOutputSchema](#contentscrapingtooloutputschema) -- [ContentScrapingToolSchema](#contentscrapingtoolschema) @@ -86,36 +88,36 @@ class ContentScrapingToolConfig(BaseToolConfig): ... #### See also -- [BaseToolConfig](./base.md#basetoolconfig) +- [BaseToolConfig](./base_tool.md#basetoolconfig) -## ContentScrapingToolOutputSchema +## ContentScrapingToolInputSchema -[Show source in content_scraping_tool.py:38](../../../../../atomic_agents/lib/tools/content_scraping_tool.py#L38) +[Show source in content_scraping_tool.py:21](../../../../../atomic_agents/lib/tools/content_scraping_tool.py#L21) #### Signature ```python -class ContentScrapingToolOutputSchema(BaseAgentIO): ... +class ContentScrapingToolInputSchema(BaseIOSchema): ... ``` #### See also -- [BaseAgentIO](../../agents/base_agent.md#baseagentio) +- [BaseIOSchema](../../agents/base_agent.md#baseioschema) -## ContentScrapingToolSchema +## ContentScrapingToolOutputSchema -[Show source in content_scraping_tool.py:21](../../../../../atomic_agents/lib/tools/content_scraping_tool.py#L21) +[Show source in content_scraping_tool.py:38](../../../../../atomic_agents/lib/tools/content_scraping_tool.py#L38) #### Signature ```python -class ContentScrapingToolSchema(BaseAgentIO): ... +class ContentScrapingToolOutputSchema(BaseIOSchema): ... ``` #### See also -- [BaseAgentIO](../../agents/base_agent.md#baseagentio) \ No newline at end of file +- [BaseIOSchema](../../agents/base_agent.md#baseioschema) \ No newline at end of file diff --git a/docs/api/atomic_agents/lib/tools/index.md b/docs/api/atomic_agents/lib/tools/index.md index cee3a859..17fc170a 100644 --- a/docs/api/atomic_agents/lib/tools/index.md +++ b/docs/api/atomic_agents/lib/tools/index.md @@ -9,8 +9,8 @@ ## Modules -- [Base](./base.md) +- [BaseTool](./base_tool.md) - [CalculatorTool](./calculator_tool.md) - [ContentScrapingTool](./content_scraping_tool.md) - [Yelp Restaurant Finder Tool](./yelp_restaurant_finder_tool.md) -- [Yt Transcript Scraper](./yt_transcript_scraper.md) \ No newline at end of file +- [Yt Transcript Scraper Tool](./yt_transcript_scraper_tool.md) \ No newline at end of file diff --git a/docs/api/atomic_agents/lib/tools/yelp_restaurant_finder_tool.md b/docs/api/atomic_agents/lib/tools/yelp_restaurant_finder_tool.md index c283f896..ac052d99 100644 --- a/docs/api/atomic_agents/lib/tools/yelp_restaurant_finder_tool.md +++ b/docs/api/atomic_agents/lib/tools/yelp_restaurant_finder_tool.md @@ -11,8 +11,8 @@ - [YelpSearchTool](#yelpsearchtool) - [YelpSearchTool().run](#yelpsearchtool()run) - [YelpSearchToolConfig](#yelpsearchtoolconfig) + - [YelpSearchToolInputSchema](#yelpsearchtoolinputschema) - [YelpSearchToolOutputSchema](#yelpsearchtooloutputschema) - - [YelpSearchToolSchema](#yelpsearchtoolschema) ## PriceRange @@ -45,12 +45,12 @@ class YelpCategory(Enum): ... #### Signature ```python -class YelpSearchResultSchema(BaseAgentIO): ... +class YelpSearchResultSchema(BaseIOSchema): ... ``` #### See also -- [BaseAgentIO](../../agents/base_agent.md#baseagentio) +- [BaseIOSchema](../../agents/base_agent.md#baseioschema) @@ -62,7 +62,7 @@ Tool for performing searches using the Yelp API based on the provided queries. #### Attributes -- `input_schema` *YelpSearchToolSchema* - The schema for the input data. +- `input_schema` *YelpSearchToolInputSchema* - The schema for the input data. - `output_schema` *YelpSearchToolOutputSchema* - The schema for the output data. - `api_key` *str* - The API key for the Yelp API. - `max_results` *int* - The maximum number of search results to return. @@ -76,7 +76,7 @@ class YelpSearchTool(BaseTool): #### See also -- [BaseTool](./base.md#basetool) +- [BaseTool](./base_tool.md#basetool) - [YelpSearchToolConfig](#yelpsearchtoolconfig) ### YelpSearchTool().run @@ -87,7 +87,7 @@ Runs the YelpSearchTool with the given parameters. #### Arguments -- `params` *YelpSearchToolSchema* - The input parameters for the tool, adhering to the input schema. +- `params` *YelpSearchToolInputSchema* - The input parameters for the tool, adhering to the input schema. - `max_results` *Optional[int]* - The maximum number of search results to return. #### Returns @@ -102,13 +102,13 @@ Runs the YelpSearchTool with the given parameters. #### Signature ```python -def run(self, params: YelpSearchToolSchema) -> YelpSearchToolOutputSchema: ... +def run(self, params: YelpSearchToolInputSchema) -> YelpSearchToolOutputSchema: ... ``` #### See also +- [YelpSearchToolInputSchema](#yelpsearchtoolinputschema) - [YelpSearchToolOutputSchema](#yelpsearchtooloutputschema) -- [YelpSearchToolSchema](#yelpsearchtoolschema) @@ -124,36 +124,36 @@ class YelpSearchToolConfig(BaseToolConfig): ... #### See also -- [BaseToolConfig](./base.md#basetoolconfig) +- [BaseToolConfig](./base_tool.md#basetoolconfig) -## YelpSearchToolOutputSchema +## YelpSearchToolInputSchema -[Show source in yelp_restaurant_finder_tool.py:88](../../../../../atomic_agents/lib/tools/yelp_restaurant_finder_tool.py#L88) +[Show source in yelp_restaurant_finder_tool.py:46](../../../../../atomic_agents/lib/tools/yelp_restaurant_finder_tool.py#L46) #### Signature ```python -class YelpSearchToolOutputSchema(BaseAgentIO): ... +class YelpSearchToolInputSchema(BaseIOSchema): ... ``` #### See also -- [BaseAgentIO](../../agents/base_agent.md#baseagentio) +- [BaseIOSchema](../../agents/base_agent.md#baseioschema) -## YelpSearchToolSchema +## YelpSearchToolOutputSchema -[Show source in yelp_restaurant_finder_tool.py:46](../../../../../atomic_agents/lib/tools/yelp_restaurant_finder_tool.py#L46) +[Show source in yelp_restaurant_finder_tool.py:88](../../../../../atomic_agents/lib/tools/yelp_restaurant_finder_tool.py#L88) #### Signature ```python -class YelpSearchToolSchema(BaseAgentIO): ... +class YelpSearchToolOutputSchema(BaseIOSchema): ... ``` #### See also -- [BaseAgentIO](../../agents/base_agent.md#baseagentio) \ No newline at end of file +- [BaseIOSchema](../../agents/base_agent.md#baseioschema) \ No newline at end of file diff --git a/docs/api/atomic_agents/lib/tools/yt_transcript_scraper.md b/docs/api/atomic_agents/lib/tools/yt_transcript_scraper_tool.md similarity index 56% rename from docs/api/atomic_agents/lib/tools/yt_transcript_scraper.md rename to docs/api/atomic_agents/lib/tools/yt_transcript_scraper_tool.md index 6575ca45..e918f70a 100644 --- a/docs/api/atomic_agents/lib/tools/yt_transcript_scraper.md +++ b/docs/api/atomic_agents/lib/tools/yt_transcript_scraper_tool.md @@ -1,27 +1,27 @@ -# Yt Transcript Scraper +# Yt Transcript Scraper Tool -[Atomic_agents Index](../../../README.md#atomic_agents-index) / [Atomic Agents](../../index.md#atomic-agents) / [Lib](../index.md#lib) / [Tools](./index.md#tools) / Yt Transcript Scraper +[Atomic_agents Index](../../../README.md#atomic_agents-index) / [Atomic Agents](../../index.md#atomic-agents) / [Lib](../index.md#lib) / [Tools](./index.md#tools) / Yt Transcript Scraper Tool -> Auto-generated documentation for [atomic_agents.lib.tools.yt_transcript_scraper](../../../../../atomic_agents/lib/tools/yt_transcript_scraper.py) module. +> Auto-generated documentation for [atomic_agents.lib.tools.yt_transcript_scraper_tool](../../../../../atomic_agents/lib/tools/yt_transcript_scraper_tool.py) module. -- [Yt Transcript Scraper](#yt-transcript-scraper) +- [Yt Transcript Scraper Tool](#yt-transcript-scraper-tool) - [YouTubeTranscriptTool](#youtubetranscripttool) - [YouTubeTranscriptTool.extract_video_id](#youtubetranscripttoolextract_video_id) - [YouTubeTranscriptTool().fetch_video_metadata](#youtubetranscripttool()fetch_video_metadata) - [YouTubeTranscriptTool().run](#youtubetranscripttool()run) - [YouTubeTranscriptToolConfig](#youtubetranscripttoolconfig) + - [YouTubeTranscriptToolInputSchema](#youtubetranscripttoolinputschema) - [YouTubeTranscriptToolOutputSchema](#youtubetranscripttooloutputschema) - - [YouTubeTranscriptToolSchema](#youtubetranscripttoolschema) ## YouTubeTranscriptTool -[Show source in yt_transcript_scraper.py:59](../../../../../atomic_agents/lib/tools/yt_transcript_scraper.py#L59) +[Show source in yt_transcript_scraper_tool.py:59](../../../../../atomic_agents/lib/tools/yt_transcript_scraper_tool.py#L59) Tool for fetching the transcript of a YouTube video using the YouTube Transcript API. #### Attributes -- `input_schema` *YouTubeTranscriptToolSchema* - The schema for the input data. +- `input_schema` *YouTubeTranscriptToolInputSchema* - The schema for the input data. - `output_schema` *YouTubeTranscriptToolOutputSchema* - The schema for the output data. #### Signature @@ -33,12 +33,12 @@ class YouTubeTranscriptTool(BaseTool): #### See also -- [BaseTool](./base.md#basetool) +- [BaseTool](./base_tool.md#basetool) - [YouTubeTranscriptToolConfig](#youtubetranscripttoolconfig) ### YouTubeTranscriptTool.extract_video_id -[Show source in yt_transcript_scraper.py:115](../../../../../atomic_agents/lib/tools/yt_transcript_scraper.py#L115) +[Show source in yt_transcript_scraper_tool.py:115](../../../../../atomic_agents/lib/tools/yt_transcript_scraper_tool.py#L115) Extracts the video ID from a YouTube URL. @@ -59,7 +59,7 @@ def extract_video_id(url: str) -> str: ... ### YouTubeTranscriptTool().fetch_video_metadata -[Show source in yt_transcript_scraper.py:128](../../../../../atomic_agents/lib/tools/yt_transcript_scraper.py#L128) +[Show source in yt_transcript_scraper_tool.py:128](../../../../../atomic_agents/lib/tools/yt_transcript_scraper_tool.py#L128) Fetches metadata for a YouTube video. @@ -79,13 +79,13 @@ def fetch_video_metadata(self, video_id: str) -> dict: ... ### YouTubeTranscriptTool().run -[Show source in yt_transcript_scraper.py:81](../../../../../atomic_agents/lib/tools/yt_transcript_scraper.py#L81) +[Show source in yt_transcript_scraper_tool.py:81](../../../../../atomic_agents/lib/tools/yt_transcript_scraper_tool.py#L81) Runs the YouTubeTranscriptTool with the given parameters. #### Arguments -- `params` *YouTubeTranscriptToolSchema* - The input parameters for the tool, adhering to the input schema. +- `params` *YouTubeTranscriptToolInputSchema* - The input parameters for the tool, adhering to the input schema. #### Returns @@ -99,20 +99,20 @@ Runs the YouTubeTranscriptTool with the given parameters. ```python def run( - self, params: YouTubeTranscriptToolSchema + self, params: YouTubeTranscriptToolInputSchema ) -> YouTubeTranscriptToolOutputSchema: ... ``` #### See also +- [YouTubeTranscriptToolInputSchema](#youtubetranscripttoolinputschema) - [YouTubeTranscriptToolOutputSchema](#youtubetranscripttooloutputschema) -- [YouTubeTranscriptToolSchema](#youtubetranscripttoolschema) ## YouTubeTranscriptToolConfig -[Show source in yt_transcript_scraper.py:52](../../../../../atomic_agents/lib/tools/yt_transcript_scraper.py#L52) +[Show source in yt_transcript_scraper_tool.py:52](../../../../../atomic_agents/lib/tools/yt_transcript_scraper_tool.py#L52) #### Signature @@ -122,36 +122,36 @@ class YouTubeTranscriptToolConfig(BaseToolConfig): ... #### See also -- [BaseToolConfig](./base.md#basetoolconfig) +- [BaseToolConfig](./base_tool.md#basetoolconfig) -## YouTubeTranscriptToolOutputSchema +## YouTubeTranscriptToolInputSchema -[Show source in yt_transcript_scraper.py:35](../../../../../atomic_agents/lib/tools/yt_transcript_scraper.py#L35) +[Show source in yt_transcript_scraper_tool.py:17](../../../../../atomic_agents/lib/tools/yt_transcript_scraper_tool.py#L17) #### Signature ```python -class YouTubeTranscriptToolOutputSchema(BaseAgentIO): ... +class YouTubeTranscriptToolInputSchema(BaseIOSchema): ... ``` #### See also -- [BaseAgentIO](../../agents/base_agent.md#baseagentio) +- [BaseIOSchema](../../agents/base_agent.md#baseioschema) -## YouTubeTranscriptToolSchema +## YouTubeTranscriptToolOutputSchema -[Show source in yt_transcript_scraper.py:17](../../../../../atomic_agents/lib/tools/yt_transcript_scraper.py#L17) +[Show source in yt_transcript_scraper_tool.py:35](../../../../../atomic_agents/lib/tools/yt_transcript_scraper_tool.py#L35) #### Signature ```python -class YouTubeTranscriptToolSchema(BaseAgentIO): ... +class YouTubeTranscriptToolOutputSchema(BaseIOSchema): ... ``` #### See also -- [BaseAgentIO](../../agents/base_agent.md#baseagentio) \ No newline at end of file +- [BaseIOSchema](../../agents/base_agent.md#baseioschema) \ No newline at end of file diff --git a/docs/api/atomic_agents/lib/utils/format_tool_message.md b/docs/api/atomic_agents/lib/utils/format_tool_message.md index a34734f2..b3f21cc2 100644 --- a/docs/api/atomic_agents/lib/utils/format_tool_message.md +++ b/docs/api/atomic_agents/lib/utils/format_tool_message.md @@ -9,14 +9,13 @@ ## format_tool_message -[Show source in format_tool_message.py:7](../../../../../atomic_agents/lib/utils/format_tool_message.py#L7) +[Show source in format_tool_message.py:8](../../../../../atomic_agents/lib/utils/format_tool_message.py#L8) Formats a message for a tool call. #### Arguments -- `tool_name` *str* - The type of the tool. -- `tool_calls` *Dict* - A dictionary containing the tool calls. +- `tool_call` *Type[BaseModel]* - The Pydantic model instance representing the tool call. - `tool_id` *str, optional* - The unique identifier for the tool call. If not provided, a random UUID will be generated. #### Returns diff --git a/docs/guides/creating_a_new_tool.md b/docs/guides/creating_a_new_tool.md index 9e6bdeb7..9185371f 100644 --- a/docs/guides/creating_a_new_tool.md +++ b/docs/guides/creating_a_new_tool.md @@ -17,13 +17,13 @@ A tool in the `atomic_agents` framework consists of the following components: #### 1. Define the Input Schema -The input schema is a Pydantic `BaseModel` that specifies the input parameters for the tool. It should inherit from `BaseAgentIO` and include a `Config` class with `title`, `description`, and `json_schema_extra` attributes. Here's an example: +The input schema is a Pydantic `BaseModel` that specifies the input parameters for the tool. It should inherit from `BaseIOSchema` and include a `Config` class with `title`, `description`, and `json_schema_extra` attributes. Here's an example: ```python from pydantic import Field -from atomic_agents.agents.base_agent import BaseAgentIO +from atomic_agents.agents.base_agent import BaseIOSchema -class MyToolInputSchema(BaseAgentIO): +class MyToolInputSchema(BaseIOSchema): parameter: str = Field(..., description="Description of the parameter.") list_param: list[str] = Field(..., description="A list of strings.") @@ -38,10 +38,10 @@ class MyToolInputSchema(BaseAgentIO): #### 2. Define the Output Schema -The output schema is also a Pydantic `BaseModel` that inherits from `BaseAgentIO`. It specifies the structure of the tool's output. Here's an example: +The output schema is also a Pydantic `BaseModel` that inherits from `BaseIOSchema`. It specifies the structure of the tool's output. Here's an example: ```python -class MyToolOutputSchema(BaseAgentIO): +class MyToolOutputSchema(BaseIOSchema): result: str = Field(..., description="Result of the tool's operation.") details: dict = Field(..., description="Additional details about the result.") ``` @@ -107,12 +107,12 @@ Here is a complete example of a new tool called `MyTool`: # my_tool.py from pydantic import Field -from atomic_agents.agents.base_agent import BaseAgentIO +from atomic_agents.agents.base_agent import BaseIOSchema from atomic_agents.lib.tools.base import BaseTool, BaseToolConfig from rich.console import Console # Input Schema -class MyToolInputSchema(BaseAgentIO): +class MyToolInputSchema(BaseIOSchema): parameter: str = Field(..., description="Description of the parameter.") list_param: list[str] = Field(..., description="A list of strings.") @@ -125,7 +125,7 @@ class MyToolInputSchema(BaseAgentIO): } # Output Schema -class MyToolOutputSchema(BaseAgentIO): +class MyToolOutputSchema(BaseIOSchema): result: str = Field(..., description="Result of the tool's operation.") details: dict = Field(..., description="Additional details about the result.") @@ -164,4 +164,4 @@ if __name__ == "__main__": ### Conclusion -By following these steps, you can create a new tool in the `atomic_agents` framework. Define the input and output schemas, implement the tool logic, and provide an example usage to demonstrate how the tool works. Remember to inherit from the appropriate base classes (`BaseAgentIO`, `BaseToolConfig`, and `BaseTool`) and use the `rich` library for console output in your example usage. +By following these steps, you can create a new tool in the `atomic_agents` framework. Define the input and output schemas, implement the tool logic, and provide an example usage to demonstrate how the tool works. Remember to inherit from the appropriate base classes (`BaseIOSchema`, `BaseToolConfig`, and `BaseTool`) and use the `rich` library for console output in your example usage. diff --git a/examples/basic_custom_chatbot.py b/examples/basic_custom_chatbot.py index 10f4c23e..9a0e26fc 100644 --- a/examples/basic_custom_chatbot.py +++ b/examples/basic_custom_chatbot.py @@ -6,12 +6,12 @@ from rich.console import Console from atomic_agents.lib.components.agent_memory import AgentMemory from atomic_agents.agents.base_agent import BaseAgent, BaseAgentConfig -from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator, SystemPromptInfo +from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator import instructor import openai # Define system prompt information including background, steps, and output instructions -system_prompt = SystemPromptInfo( +system_prompt_generator = SystemPromptGenerator( background=[ 'This assistant is a general-purpose AI designed to be helpful and friendly.', ], @@ -25,8 +25,6 @@ 'Always answer in rhyming verse.' ] ) -# Initialize the system prompt generator with the defined system prompt and dynamic info providers -system_prompt_generator = SystemPromptGenerator(system_prompt) # Initialize chat memory to store conversation history memory = AgentMemory() @@ -49,7 +47,7 @@ api_key=api_key or os.getenv('OPENAI_API_KEY') )), system_prompt_generator=system_prompt_generator, - model='gpt-3.5-turbo', + model='gpt-4o-mini', memory=memory, ) ) diff --git a/examples/conversing_agents.py b/examples/conversing_agents.py index f6b76e80..97b59742 100644 --- a/examples/conversing_agents.py +++ b/examples/conversing_agents.py @@ -6,16 +6,16 @@ import openai from rich.console import Console from atomic_agents.agents.base_agent import BaseAgentConfig, BaseAgent, BaseAgentInputSchema, BaseAgentOutputSchema -from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator, SystemPromptInfo +from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator # Initialize the system prompts. For simplicity, we use the same system prompt for both agents. -agent1_system_prompt_info = SystemPromptInfo( +agent1_system_prompt_generator = SystemPromptGenerator( background=['You are an agent that adds 1 to the number said by the previous person in the conversation.'], steps=['If the previous message contains a number, add 1 to it.'], output_instructions=['Respond with only the resulting added number and never anything else. Do not include any additional text.'] ) -agent2_system_prompt_info = SystemPromptInfo( +agent2_system_prompt_generator = SystemPromptGenerator( background=['You are an agent that adds 1 to the number said by the previous person in the conversation.'], steps=['If the previous message contains a number, add 1 to it.'], output_instructions=['Respond with only the resulting added number and never anything else. Do not include any additional text.'] @@ -27,15 +27,15 @@ config=BaseAgentConfig( client=instructor.from_groq(groq.Groq(api_key=os.getenv('GROQ_API_KEY'))), model=os.getenv('GROQ_CHAT_MODEL'), - system_prompt_generator=SystemPromptGenerator(system_prompt_info=agent1_system_prompt_info) + system_prompt_generator=agent1_system_prompt_generator ) ) agent2 = BaseAgent( config=BaseAgentConfig( client=instructor.from_openai(openai.OpenAI()), - model='gpt-3.5-turbo', - system_prompt_generator=SystemPromptGenerator(system_prompt_info=agent2_system_prompt_info) + model='gpt-4o-mini', + system_prompt_generator=agent2_system_prompt_generator ) ) diff --git a/examples/deep_research_multi_agent/agents/info_refiner_agent.py b/examples/deep_research_multi_agent/agents/info_refiner_agent.py index c33acd5a..27b41885 100644 --- a/examples/deep_research_multi_agent/agents/info_refiner_agent.py +++ b/examples/deep_research_multi_agent/agents/info_refiner_agent.py @@ -1,11 +1,11 @@ from pydantic import BaseModel, Field -from atomic_agents.agents.base_agent import BaseAgentIO +from atomic_agents.agents.base_agent import BaseIOSchema import instructor import openai from atomic_agents.agents.base_agent import BaseAgent, BaseAgentConfig -from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator, SystemPromptInfo +from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator -class RefineAnswerInputSchema(BaseAgentIO): +class RefineAnswerInputSchema(BaseIOSchema): question: str = Field(..., description='The question that was asked.') answer: str = Field(..., description='The initial answer to the question.') @@ -16,25 +16,23 @@ class RefineAnswerOutputSchema(BaseModel): refine_answer_agent = BaseAgent( BaseAgentConfig( client=instructor.from_openai(openai.OpenAI()), - model='gpt-3.5-turbo', + model='gpt-4o-mini', system_prompt_generator=SystemPromptGenerator( - SystemPromptInfo( - background=[ - "You are an intelligent answer refinement expert.", - "Your task is to expand and elaborate on an existing answer to a question using additional context from vector DB chunks." - ], - steps=[ - "You will receive a question or instruction, the initial answer, and additional context from vector DB chunks.", - "Correct any inaccuracies and factual errors in the initial answer.", - "Expand and elaborate on the initial answer using the additional context to provide a more comprehensive and detailed response." - ], - output_instructions=[ - "Ensure the refined answer is clear, concise, and well-structured.", - "Ensure the refined answer is directly relevant to the question and incorporates the additional context provided.", - "Add new information and details to make the final answer more elaborate and informative.", - "Do not make up any new information; only use the information present in the context." - ] - ) + background=[ + "You are an intelligent answer refinement expert.", + "Your task is to expand and elaborate on an existing answer to a question using additional context from vector DB chunks." + ], + steps=[ + "You will receive a question or instruction, the initial answer, and additional context from vector DB chunks.", + "Correct any inaccuracies and factual errors in the initial answer.", + "Expand and elaborate on the initial answer using the additional context to provide a more comprehensive and detailed response." + ], + output_instructions=[ + "Ensure the refined answer is clear, concise, and well-structured.", + "Ensure the refined answer is directly relevant to the question and incorporates the additional context provided.", + "Add new information and details to make the final answer more elaborate and informative.", + "Do not make up any new information; only use the information present in the context." + ] ), input_schema=RefineAnswerInputSchema, output_schema=RefineAnswerOutputSchema diff --git a/examples/deep_research_multi_agent/agents/query_agent.py b/examples/deep_research_multi_agent/agents/query_agent.py index ad106924..d24e352e 100644 --- a/examples/deep_research_multi_agent/agents/query_agent.py +++ b/examples/deep_research_multi_agent/agents/query_agent.py @@ -1,10 +1,10 @@ import instructor import openai -from atomic_agents.agents.base_agent import BaseAgentIO, BaseAgent, BaseAgentConfig -from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator, SystemPromptInfo -from atomic_agents.lib.tools.search.searx_tool import SearxNGSearchTool +from atomic_agents.agents.base_agent import BaseIOSchema, BaseAgent, BaseAgentConfig +from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator +from atomic_agents.lib.tools.search.searxng_tool import SearxNGTool -class QueryAgentInputSchema(BaseAgentIO): +class QueryAgentInputSchema(BaseIOSchema): instruction: str = 'A detailed instruction or request to generate deep research queries for.' num_queries: int = 'The number of queries to generate.' @@ -12,25 +12,23 @@ class QueryAgentInputSchema(BaseAgentIO): query_agent = BaseAgent( BaseAgentConfig( client=instructor.from_openai(openai.OpenAI()), - model='gpt-3.5-turbo', + model='gpt-4o-mini', system_prompt_generator=SystemPromptGenerator( - SystemPromptInfo( - background=[ - "You are an intelligent query generation expert.", - "Your task is to generate a specified number of diverse and highly relevant queries based on a given instruction or request.", - "The queries should cover different aspects of the instruction to ensure comprehensive exploration." - ], - steps=[ - "You will receive a detailed instruction or request and the number of queries to generate.", - "Generate the requested number of queries in a JSON format." - ], - output_instructions=[ - "Ensure clarity and conciseness in each query.", - "Ensure each query is unique and as diverse as possible while remaining relevant to the instruction." - ] - ) + background=[ + "You are an intelligent query generation expert.", + "Your task is to generate a specified number of diverse and highly relevant queries based on a given instruction or request.", + "The queries should cover different aspects of the instruction to ensure comprehensive exploration." + ], + steps=[ + "You will receive a detailed instruction or request and the number of queries to generate.", + "Generate the requested number of queries in a JSON format." + ], + output_instructions=[ + "Ensure clarity and conciseness in each query.", + "Ensure each query is unique and as diverse as possible while remaining relevant to the instruction." + ] ), input_schema=QueryAgentInputSchema, - output_schema=SearxNGSearchTool.input_schema + output_schema=SearxNGTool.input_schema ) ) \ No newline at end of file diff --git a/examples/deep_research_multi_agent/agents/question_answering_agent.py b/examples/deep_research_multi_agent/agents/question_answering_agent.py index b7f05889..06a4e178 100644 --- a/examples/deep_research_multi_agent/agents/question_answering_agent.py +++ b/examples/deep_research_multi_agent/agents/question_answering_agent.py @@ -1,35 +1,33 @@ import instructor import openai from pydantic import BaseModel, Field -from atomic_agents.agents.base_agent import BaseAgentIO, BaseAgent, BaseAgentConfig -from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator, SystemPromptInfo +from atomic_agents.agents.base_agent import BaseIOSchema, BaseAgent, BaseAgentConfig +from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator -class AnswerAgentInputSchema(BaseAgentIO): +class AnswerAgentInputSchema(BaseIOSchema): question: str = Field(..., description='A question that needs to be answered based on the provided context.') -class AnswerAgentOutputSchema(BaseAgentIO): +class AnswerAgentOutputSchema(BaseIOSchema): markdown_output: str = Field(..., description='The answer to the question in markdown format.') # Create the answer agent answer_agent = BaseAgent( BaseAgentConfig( client=instructor.from_openai(openai.OpenAI()), - model='gpt-3.5-turbo', + model='gpt-4o-mini', system_prompt_generator=SystemPromptGenerator( - SystemPromptInfo( - background=[ - "You are an intelligent answering expert.", - "Your task is to provide accurate and detailed answers to user questions based on the given context." - ], - steps=[ - "You will receive a question and the context information.", - "Generate a detailed and accurate answer based on the context." - ], - output_instructions=[ - "Ensure clarity and conciseness in each answer.", - "Ensure the answer is directly relevant to the question and context provided." - ] - ) + background=[ + "You are an intelligent answering expert.", + "Your task is to provide accurate and detailed answers to user questions based on the given context." + ], + steps=[ + "You will receive a question and the context information.", + "Generate a detailed and accurate answer based on the context." + ], + output_instructions=[ + "Ensure clarity and conciseness in each answer.", + "Ensure the answer is directly relevant to the question and context provided." + ] ), input_schema=AnswerAgentInputSchema, output_schema=AnswerAgentOutputSchema diff --git a/examples/deep_research_multi_agent/agents/top_urls_selector_agent.py b/examples/deep_research_multi_agent/agents/top_urls_selector_agent.py index ec66863d..6f7dc1e2 100644 --- a/examples/deep_research_multi_agent/agents/top_urls_selector_agent.py +++ b/examples/deep_research_multi_agent/agents/top_urls_selector_agent.py @@ -1,37 +1,35 @@ from typing import List import instructor import openai -from pydantic import BaseModel, Field, HttpUrl -from atomic_agents.agents.base_agent import BaseAgentIO, BaseAgent, BaseAgentConfig -from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator, SystemPromptInfo +from pydantic import Field, HttpUrl +from atomic_agents.agents.base_agent import BaseIOSchema, BaseAgent, BaseAgentConfig +from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator -class TopUrlsSelectorInputSchema(BaseAgentIO): +class TopUrlsSelectorInputSchema(BaseIOSchema): user_input: str = Field(..., description='The user input or question.') num_urls: int = Field(..., description='The number of top URLs to select.') -class TopUrlsSelectorOutputSchema(BaseAgentIO): +class TopUrlsSelectorOutputSchema(BaseIOSchema): top_urls: List[HttpUrl] = Field(..., description='The list of top URLs selected based on the user input.') # Create the top URLs selector agent top_urls_selector_agent = BaseAgent( BaseAgentConfig( client=instructor.from_openai(openai.OpenAI()), - model='gpt-3.5-turbo', + model='gpt-4o-mini', system_prompt_generator=SystemPromptGenerator( - SystemPromptInfo( - background=[ - "You are an intelligent URL selection expert.", - "Your task is to select the best URLs from a list of search results that are most relevant to the user input or question." - ], - steps=[ - "You will receive the user input or question, the list of search results, and the number of best URLs to select.", - "Analyze the search results and select the best URLs that are most relevant to the user input or question.", - ], - output_instructions=[ - "Ensure each selected URL is highly relevant to the user input or question.", - "Ensure the selected URLs are diverse yet as relevant as possible." - ] - ) + background=[ + "You are an intelligent URL selection expert.", + "Your task is to select the best URLs from a list of search results that are most relevant to the user input or question." + ], + steps=[ + "You will receive the user input or question, the list of search results, and the number of best URLs to select.", + "Analyze the search results and select the best URLs that are most relevant to the user input or question.", + ], + output_instructions=[ + "Ensure each selected URL is highly relevant to the user input or question.", + "Ensure the selected URLs are diverse yet as relevant as possible." + ] ), input_schema=TopUrlsSelectorInputSchema, output_schema=TopUrlsSelectorOutputSchema diff --git a/examples/deep_research_multi_agent/main.py b/examples/deep_research_multi_agent/main.py index bdc150d7..5df5d89c 100644 --- a/examples/deep_research_multi_agent/main.py +++ b/examples/deep_research_multi_agent/main.py @@ -5,7 +5,7 @@ from examples.deep_research_multi_agent.agents.info_refiner_agent import refine_answer_agent from examples.deep_research_multi_agent.providers import search_results_provider, vector_db_chunks_provider from examples.deep_research_multi_agent.in_memory_faiss import InMemFaiss -from atomic_agents.lib.tools.search.searx_tool import SearxNGSearchTool, SearxNGSearchToolConfig +from atomic_agents.lib.tools.search.searxng_tool import SearxNGTool, SearxNGToolConfig from rich.console import Console import asyncio @@ -14,7 +14,7 @@ async def main(): console = Console() # Initialize the search tool - search_tool = SearxNGSearchTool(SearxNGSearchToolConfig(base_url=os.getenv('SEARXNG_BASE_URL'), max_results=30)) + search_tool = SearxNGTool(SearxNGToolConfig(base_url=os.getenv('SEARXNG_BASE_URL'), max_results=30)) in_mem_faiss = InMemFaiss(openai_api_key=os.getenv('OPENAI_API_KEY')) user_input = input("Enter your question (or type 'exit' to quit): ") diff --git a/examples/mesop_example.py b/examples/mesop_example.py index 594ad85d..f096d1cc 100644 --- a/examples/mesop_example.py +++ b/examples/mesop_example.py @@ -16,29 +16,29 @@ from dataclasses import field from atomic_agents.agents.tool_interface_agent import ToolInterfaceAgent, ToolInterfaceAgentConfig -from atomic_agents.lib.tools.search.searx_tool import SearxNGSearchTool, SearxNGSearchToolConfig +from atomic_agents.lib.tools.search.searxng_tool import SearxNGTool, SearxNGToolConfig # Initialize the console, client, and tools console = Console() client = instructor.from_openai(openai.OpenAI()) -def initialize_searx_tool(): +def initialize_searxng_tool(): base_url = os.getenv('SEARXNG_BASE_URL') - config = SearxNGSearchToolConfig(base_url=base_url, max_results=10) - return SearxNGSearchTool(config) + config = SearxNGToolConfig(base_url=base_url, max_results=10) + return SearxNGTool(config) -searx_tool = initialize_searx_tool() +searxng_tool = initialize_searxng_tool() -def initialize_agent(client, searx_tool): +def initialize_agent(client, searxng_tool): agent_config = ToolInterfaceAgentConfig( client=client, - model='gpt-3.5-turbo', - tool_instance=searx_tool, + model='gpt-4o-mini', + tool_instance=searxng_tool, return_raw_output=False ) return ToolInterfaceAgent(config=agent_config) -agent = initialize_agent(client, searx_tool) +agent = initialize_agent(client, searxng_tool) # Mesop state class @me.stateclass @@ -48,7 +48,7 @@ class State: # Mesop transform function def transform(input: str, history: list[mel.ChatMessage]): state = me.state(State) - response = agent.run(agent.input_schema(tool_input_SearxNGSearchTool=input)) + response = agent.run(agent.input_schema(tool_input_SearxNGTool=input)) state.chat_history.append(mel.ChatMessage(content=f"You: {input}")) state.chat_history.append(mel.ChatMessage(content=f"Agent: {response.chat_message}")) yield response.chat_message @@ -56,11 +56,11 @@ def transform(input: str, history: list[mel.ChatMessage]): # Mesop page @me.page( - path="/chat", + path="/", title="ToolInterfaceAgent Chat" ) def page(): - mel.chat(transform, title="ToolInterfaceAgent with SearxNGSearchTool", bot_user="Agent") + mel.chat(transform, title="ToolInterfaceAgent with SearxNGTool", bot_user="Agent") # Main function to run the Mesop app if __name__ == "__main__": diff --git a/examples/minimal_chatbot.py b/examples/minimal_chatbot.py index 2ff1b6e9..004b850c 100644 --- a/examples/minimal_chatbot.py +++ b/examples/minimal_chatbot.py @@ -10,7 +10,7 @@ agent = BaseAgent( config=BaseAgentConfig( client=instructor.from_openai(openai.OpenAI()), - model='gpt-3.5-turbo' + model='gpt-4o-mini' ) ) diff --git a/examples/minimal_multi_agent_tool_use.py b/examples/minimal_multi_agent_tool_use.py index e67f55c9..bfdbefad 100644 --- a/examples/minimal_multi_agent_tool_use.py +++ b/examples/minimal_multi_agent_tool_use.py @@ -4,24 +4,24 @@ import openai from pydantic import create_model from rich.console import Console -from atomic_agents.agents.base_agent import BaseAgentIO, BaseAgent, BaseAgentConfig +from atomic_agents.agents.base_agent import BaseIOSchema, BaseAgent, BaseAgentConfig from atomic_agents.agents.tool_interface_agent import ToolInterfaceAgent, ToolInterfaceAgentConfig -from atomic_agents.lib.tools.search.searx_tool import SearxNGSearchTool, SearxNGSearchToolConfig +from atomic_agents.lib.tools.search.searxng_tool import SearxNGTool, SearxNGToolConfig from atomic_agents.lib.tools.calculator_tool import CalculatorTool, CalculatorToolConfig console = Console() client = instructor.from_openai(openai.OpenAI()) -searx_tool = SearxNGSearchTool(SearxNGSearchToolConfig(base_url=os.getenv('SEARXNG_BASE_URL'), max_results=10)) +searxng_tool = SearxNGTool(SearxNGToolConfig(base_url=os.getenv('SEARXNG_BASE_URL'), max_results=10)) calc_tool = CalculatorTool(CalculatorToolConfig()) -search_agent_config = ToolInterfaceAgentConfig(client=client, model='gpt-3.5-turbo', tool_instance=searx_tool, return_raw_output=False) -calculator_agent_config = ToolInterfaceAgentConfig(client=client, model='gpt-3.5-turbo', tool_instance=calc_tool, return_raw_output=False) +search_agent_config = ToolInterfaceAgentConfig(client=client, model='gpt-4o-mini', tool_instance=searxng_tool, return_raw_output=False) +calculator_agent_config = ToolInterfaceAgentConfig(client=client, model='gpt-4o-mini', tool_instance=calc_tool, return_raw_output=False) searx_agent = ToolInterfaceAgent(config=search_agent_config) calc_agent = ToolInterfaceAgent(config=calculator_agent_config) -UnionResponse = create_model('UnionResponse', __base__=BaseAgentIO, response=(Union[searx_agent.input_schema, calc_agent.input_schema], ...)) +UnionResponse = create_model('UnionResponse', __base__=BaseIOSchema, response=(Union[searx_agent.input_schema, calc_agent.input_schema], ...)) -orchestration_agent = BaseAgent(config=BaseAgentConfig(client=client, model='gpt-3.5-turbo', output_schema=UnionResponse)) +orchestration_agent = BaseAgent(config=BaseAgentConfig(client=client, model='gpt-4o-mini', output_schema=UnionResponse)) while True: user_input = input('You: ') diff --git a/examples/notebooks/multi_agent_quickstart.ipynb b/examples/notebooks/multi_agent_quickstart.ipynb index b7117b56..7ec1abf1 100644 --- a/examples/notebooks/multi_agent_quickstart.ipynb +++ b/examples/notebooks/multi_agent_quickstart.ipynb @@ -65,9 +65,9 @@ "import openai\n", "from pydantic import create_model\n", "from rich.console import Console\n", - "from atomic_agents.agents.base_agent import BaseAgentIO, BaseAgent, BaseAgentConfig\n", + "from atomic_agents.agents.base_agent import BaseIOSchema, BaseAgent, BaseAgentConfig\n", "from atomic_agents.agents.tool_interface_agent import ToolInterfaceAgent, ToolInterfaceAgentConfig\n", - "from atomic_agents.lib.tools.search.searx_tool import SearxNGSearchTool, SearxNGSearchToolConfig\n", + "from atomic_agents.lib.tools.search.searxng_tool import SearxNGTool, SearxNGToolConfig\n", "from atomic_agents.lib.tools.calculator_tool import CalculatorTool, CalculatorToolConfig" ] }, @@ -93,7 +93,7 @@ "client = instructor.from_openai(openai.OpenAI())\n", "\n", "# Initialize the SearxNG search tool\n", - "searx_tool = SearxNGSearchTool(SearxNGSearchToolConfig(base_url=os.getenv('SEARXNG_BASE_URL'), max_results=10))\n", + "searxng_tool = SearxNGTool(SearxNGToolConfig(base_url=os.getenv('SEARXNG_BASE_URL'), max_results=10))\n", "\n", "# Initialize the calculator tool\n", "calc_tool = CalculatorTool(CalculatorToolConfig())" @@ -115,10 +115,10 @@ "outputs": [], "source": [ "# Configure the search agent\n", - "search_agent_config = ToolInterfaceAgentConfig(client=client, model='gpt-3.5-turbo', tool_instance=searx_tool, return_raw_output=False)\n", + "search_agent_config = ToolInterfaceAgentConfig(client=client, model='gpt-4o-mini', tool_instance=searxng_tool, return_raw_output=False)\n", "\n", "# Configure the calculator agent\n", - "calculator_agent_config = ToolInterfaceAgentConfig(client=client, model='gpt-3.5-turbo', tool_instance=calc_tool, return_raw_output=False)\n", + "calculator_agent_config = ToolInterfaceAgentConfig(client=client, model='gpt-4o-mini', tool_instance=calc_tool, return_raw_output=False)\n", "\n", "# Create the search agent and calculator agent\n", "searx_agent = ToolInterfaceAgent(config=search_agent_config)\n", @@ -141,7 +141,7 @@ "outputs": [], "source": [ "# Define a union response schema\n", - "UnionResponse = create_model('UnionResponse', __base__=BaseAgentIO, response=(Union[searx_agent.input_schema, calc_agent.input_schema], ...))" + "UnionResponse = create_model('UnionResponse', __base__=BaseIOSchema, response=(Union[searx_agent.input_schema, calc_agent.input_schema], ...))" ] }, { @@ -167,7 +167,7 @@ "outputs": [], "source": [ "# Create the orchestration agent\n", - "orchestration_agent = BaseAgent(config=BaseAgentConfig(client=client, model='gpt-3.5-turbo', output_schema=UnionResponse))" + "orchestration_agent = BaseAgent(config=BaseAgentConfig(client=client, model='gpt-4o-mini', output_schema=UnionResponse))" ] }, { diff --git a/examples/notebooks/quickstart.ipynb b/examples/notebooks/quickstart.ipynb index 39900755..428d7b87 100644 --- a/examples/notebooks/quickstart.ipynb +++ b/examples/notebooks/quickstart.ipynb @@ -67,7 +67,7 @@ "import openai\n", "from atomic_agents.lib.components.agent_memory import AgentMemory\n", "from atomic_agents.agents.base_agent import BaseAgent, BaseAgentConfig\n", - "from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator, SystemPromptInfo" + "from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator" ] }, { @@ -88,7 +88,7 @@ "metadata": {}, "outputs": [], "source": [ - "system_prompt = SystemPromptInfo(\n", + "system_prompt_generator = SystemPromptGenerator(\n", " background=[\n", " 'This assistant is a general-purpose AI designed to be helpful and friendly.',\n", " ],\n", @@ -102,7 +102,6 @@ " 'Always answer in rhyming verse.'\n", " ]\n", ")\n", - "system_prompt_generator = SystemPromptGenerator(system_prompt)\n", "\n", "from rich.markdown import Markdown\n", "Markdown(system_prompt_generator.generate_prompt())" @@ -162,7 +161,7 @@ " config=BaseAgentConfig(\n", " client=client,\n", " system_prompt_generator=system_prompt_generator,\n", - " model='gpt-3.5-turbo',\n", + " model='gpt-4o-mini',\n", " memory=memory,\n", " )\n", ")" diff --git a/examples/notebooks/what_is_atomic_agents.ipynb b/examples/notebooks/what_is_atomic_agents.ipynb index 59a650ac..c6823be9 100644 --- a/examples/notebooks/what_is_atomic_agents.ipynb +++ b/examples/notebooks/what_is_atomic_agents.ipynb @@ -90,7 +90,7 @@ "Next, we will import the necessary components for creating the chatbot. Each component serves a specific purpose:\n", "- `AgentMemory`: Manages the chat history.\n", "- `BaseAgent`: The base class to create a custom chatbot. Can be extended for additional functionality if needed.\n", - "- `SystemPromptGenerator` and `SystemPromptInfo`: To define and generate system prompts." + "- `SystemPromptGenerator`: To define and generate system prompts." ] }, { @@ -100,7 +100,7 @@ "outputs": [], "source": [ "import os\n", - "from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator, SystemPromptInfo" + "from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator" ] }, { @@ -120,7 +120,7 @@ "source": [ "from rich.console import Console\n", "\n", - "system_prompt = SystemPromptInfo(\n", + "system_prompt_generator = SystemPromptGenerator(\n", " background=[\n", " 'This assistant is a general-purpose AI designed to be helpful and friendly.',\n", " ],\n", @@ -134,7 +134,6 @@ " 'Always answer in rhyming verse.'\n", " ]\n", ")\n", - "system_prompt_generator = SystemPromptGenerator(system_prompt)\n", "\n", "Console().print(system_prompt_generator.generate_prompt())" ] @@ -235,7 +234,7 @@ " config=BaseAgentConfig(\n", " client=client,\n", " system_prompt_generator=system_prompt_generator,\n", - " model='gpt-3.5-turbo',\n", + " model='gpt-4o-mini',\n", " memory=memory,\n", " )\n", ")\n", @@ -286,7 +285,7 @@ "\n", "First off, as mentioned before, each tool and each agent has an **input schema** and an **output schema**.\n", "\n", - "These schemas all extend from what is called the `BaseAgentIO` class. This class itself is a Pydantic model with some predefined methods to facilitate storing a structured response in memory and to facilitate printing the response to the console using the `rich` library." + "These schemas all extend from what is called the `BaseIOSchema` class. This class itself is a Pydantic model with some predefined methods to facilitate storing a structured response in memory and to facilitate printing the response to the console using the `rich` library." ] }, { @@ -298,7 +297,7 @@ "from pydantic import BaseModel, Field\n", "from rich.json import JSON\n", "\n", - "class BaseAgentIO(BaseModel):\n", + "class BaseIOSchema(BaseModel):\n", " \"\"\"\n", " Base class for input and output schemas for chat agents.\n", " \"\"\"\n", @@ -309,7 +308,7 @@ " json_str = self.model_dump_json()\n", " return JSON(json_str)\n", " \n", - "class BaseAgentInputSchema(BaseAgentIO):\n", + "class BaseAgentInputSchema(BaseIOSchema):\n", " chat_message: str = Field(\n", " ...,\n", " description='The chat message exchanged between the user and the chat agent. '\n", @@ -324,7 +323,7 @@ " \"description\": description,\n", " }\n", "\n", - "class BaseAgentOutputSchema(BaseAgentIO):\n", + "class BaseAgentOutputSchema(BaseIOSchema):\n", " chat_message: str = Field(\n", " ...,\n", " description='The chat message exchanged between the user and the chat agent. '\n", @@ -345,7 +344,7 @@ "metadata": {}, "source": [ "In the code above, you can see that the `BaseAgentInputSchema` and `BaseAgentOutputSchema` both:\n", - "- Extend from `BaseAgentIO`\n", + "- Extend from `BaseIOSchema`\n", "- Have a `chat_message` field with a `description` property \n", "- Have a Config class with a `title`, `description`, and `json_schema_extra` property.\n", "\n", @@ -433,7 +432,7 @@ " THREE = \"3\"\n", " FOUR = \"4\"\n", "\n", - "class YelpSearchToolSchema(BaseAgentIO):\n", + "class YelpSearchToolInputSchema(BaseIOSchema):\n", " location: str = Field(..., description=\"Location to search for food.\")\n", " term: Optional[str] = Field(None, description=\"Search term (e.g., 'pizza', 'sushi').\")\n", " categories: Optional[List[YelpCategory]] = Field(None, description=\"Categories to filter by (e.g., 'italian, mexican').\")\n", @@ -475,16 +474,16 @@ "source": [ "import json\n", "from pydantic import Field, create_model\n", - "from atomic_agents.agents.base_agent import BaseAgentIO, BaseAgent, BaseAgentConfig\n", + "from atomic_agents.agents.base_agent import BaseIOSchema, BaseAgent, BaseAgentConfig\n", "from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator, SystemPromptInfo\n", - "from atomic_agents.lib.tools.base import BaseTool\n", + "from atomic_agents.lib.tools.base_tool import BaseTool\n", "from atomic_agents.lib.utils.format_tool_message import format_tool_message\n", "\n", "class ToolInterfaceAgentConfig(BaseAgentConfig):\n", " tool_instance: BaseTool\n", " return_raw_output: bool = False\n", "\n", - "class ToolInputModel(BaseAgentIO):\n", + "class ToolInputModel(BaseIOSchema):\n", " tool_input: str = Field(..., description=\"Tool input. Presented as a single question or instruction\")\n", "\n", " class Config:\n", @@ -538,24 +537,23 @@ " self.output_schema = self.tool_instance.output_schema\n", " \n", " self.system_prompt_generator = SystemPromptGenerator(\n", - " system_prompt_info=SystemPromptInfo(\n", - " background=[\n", - " f\"This AI agent is designed to interact with the {self.tool_instance.tool_name} tool.\",\n", - " f\"Tool description: {self.tool_instance.tool_description}\"\n", - " ],\n", - " steps=[\n", - " \"Get the user input.\",\n", - " \"Convert the input to the proper parameters to call the tool.\",\n", - " \"Call the tool with the parameters.\",\n", - " \"Respond to the user\"\n", - " ],\n", - " output_instructions=[\n", - " \"Make sure the tool call will maximize the utility of the tool in the context of the user input.\",\n", - " \"Process the output of the tool into a human readable format and/or use it to respond to the user input.\" if self.return_raw_output else \"Return the raw output of the tool.\"\n", - " ]\n", - " )\n", + " background=[\n", + " f\"This AI agent is designed to interact with the {self.tool_instance.tool_name} tool.\",\n", + " f\"Tool description: {self.tool_instance.tool_description}\"\n", + " ],\n", + " steps=[\n", + " \"Get the user input.\",\n", + " \"Convert the input to the proper parameters to call the tool.\",\n", + " \"Call the tool with the parameters.\",\n", + " \"Respond to the user\"\n", + " ],\n", + " output_instructions=[\n", + " \"Make sure the tool call will maximize the utility of the tool in the context of the user input.\",\n", + " \"Process the output of the tool into a human readable format and/or use it to respond to the user input.\" if self.return_raw_output else \"Return the raw output of the tool.\"\n", + " ]\n", " )\n", " \n", + " \n", " def _get_and_handle_response(self):\n", " \"\"\"\n", " Handles obtaining and processing the response from the tool.\n", @@ -603,24 +601,24 @@ "from rich.console import Console\n", "\n", "from atomic_agents.agents.tool_interface_agent import ToolInterfaceAgent, ToolInterfaceAgentConfig\n", - "from atomic_agents.lib.tools.search.searx_tool import SearxNGSearchTool, SearxNGSearchToolConfig\n", + "from atomic_agents.lib.tools.search.searxng_tool import SearxNGTool, SearxNGToolConfig\n", "\n", - "def initialize_searx_tool():\n", + "def initialize_searxng_tool():\n", " \"\"\"\n", - " Initialize the SearxNGSearchTool with configuration.\n", + " Initialize the SearxNGTool with configuration.\n", " \"\"\"\n", " base_url = os.getenv('SEARXNG_BASE_URL')\n", - " config = SearxNGSearchToolConfig(base_url=base_url, max_results=10)\n", - " return SearxNGSearchTool(config)\n", + " config = SearxNGToolConfig(base_url=base_url, max_results=10)\n", + " return SearxNGTool(config)\n", "\n", - "def initialize_agent(client, searx_tool):\n", + "def initialize_agent(client, searxng_tool):\n", " \"\"\"\n", - " Initialize the ToolInterfaceAgent with the given client and SearxNGSearchTool.\n", + " Initialize the ToolInterfaceAgent with the given client and SearxNGTool.\n", " \"\"\"\n", " agent_config = ToolInterfaceAgentConfig(\n", " client=client,\n", - " model='gpt-3.5-turbo',\n", - " tool_instance=searx_tool,\n", + " model='gpt-4o-mini',\n", + " tool_instance=searxng_tool,\n", " return_raw_output=False\n", " )\n", " return ToolInterfaceAgent(config=agent_config)\n", @@ -632,10 +630,10 @@ " base_url='http://localhost:1234/v1'\n", " )\n", " )\n", - " searx_tool = initialize_searx_tool()\n", - " agent = initialize_agent(client, searx_tool)\n", + " searxng_tool = initialize_searxng_tool()\n", + " agent = initialize_agent(client, searxng_tool)\n", "\n", - " console.print(\"ToolInterfaceAgent with SearxNGSearchTool is ready.\")\n", + " console.print(\"ToolInterfaceAgent with SearxNGTool is ready.\")\n", "\n", " while True:\n", " user_input = input('You: ')\n", diff --git a/examples/notebooks/yelp_agent.ipynb b/examples/notebooks/yelp_agent.ipynb index 5f8ef9f8..1c75c9a3 100644 --- a/examples/notebooks/yelp_agent.ipynb +++ b/examples/notebooks/yelp_agent.ipynb @@ -68,9 +68,9 @@ "from rich.console import Console\n", "\n", "from atomic_agents.lib.components.agent_memory import AgentMemory\n", - "from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator, SystemPromptInfo\n", + "from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator\n", "from atomic_agents.agents.base_agent import BaseAgent, BaseAgentOutputSchema, BaseAgentConfig\n", - "from atomic_agents.lib.tools.yelp_restaurant_finder_tool import YelpSearchTool, YelpSearchToolConfig, YelpSearchToolSchema" + "from atomic_agents.lib.tools.yelp_restaurant_finder_tool import YelpSearchTool, YelpSearchToolConfig, YelpSearchToolInputSchema" ] }, { @@ -88,8 +88,8 @@ "metadata": {}, "outputs": [], "source": [ - "# Define system prompt information including background, steps, and output instructions\n", - "system_prompt = SystemPromptInfo(\n", + "# Initialize the system prompt generator with the defined system prompt and dynamic info providers\n", + "system_prompt_generator = SystemPromptGenerator(\n", " background=[\n", " 'This assistant is a restaurant finder AI designed to help users find the best restaurants based on their preferences by asking clarifying questions.',\n", " ],\n", @@ -124,9 +124,6 @@ "metadata": {}, "outputs": [], "source": [ - "# Initialize the system prompt generator with the defined system prompt and dynamic info providers\n", - "system_prompt_generator = SystemPromptGenerator(system_prompt)\n", - "\n", "# Initialize chat memory to store conversation history\n", "memory = AgentMemory()\n", "# Define initial memory with a greeting message from the assistant\n", @@ -194,7 +191,7 @@ "source": [ "# Define a custom response schema that can handle both chat responses and Yelp search tool responses\n", "class ResponseSchema(BaseModel):\n", - " chosen_schema: Union[BaseAgentOutputSchema, YelpSearchToolSchema] = Field(..., description='The response from the chat agent.')\n", + " chosen_schema: Union[BaseAgentOutputSchema, YelpSearchToolInputSchema] = Field(..., description='The response from the chat agent.')\n", "\n", " class Config:\n", " title = 'OutputSchema'\n", @@ -225,9 +222,9 @@ " config=BaseAgentConfig(\n", " client=client, \n", " system_prompt_generator=system_prompt_generator,\n", - " model='gpt-3.5-turbo',\n", + " model='gpt-4o-mini',\n", " memory=memory,\n", - " output_schema=OutputSchema\n", + " output_schema=ResponseSchema\n", " )\n", ")\n", "\n", @@ -260,7 +257,7 @@ " response = agent.run(agent.input_schema(chat_message=user_input).model_dump())\n", " \n", " # Check the type of the response schema\n", - " if isinstance(response.chosen_schema, YelpSearchToolSchema):\n", + " if isinstance(response.chosen_schema, YelpSearchToolInputSchema):\n", " output = yelp_tool.run(response.chosen_schema)\n", " agent.memory.add_message('system', f'The following information was found: {output.results}\\n\\n Please summarize the results for the user.')\n", " output = agent.run().chosen_schema.chat_message\n", diff --git a/examples/notebooks/youtube_knowledge_extraction.ipynb b/examples/notebooks/youtube_knowledge_extraction.ipynb index 8593625e..55689f07 100644 --- a/examples/notebooks/youtube_knowledge_extraction.ipynb +++ b/examples/notebooks/youtube_knowledge_extraction.ipynb @@ -101,8 +101,8 @@ "from typing import List, Optional\n", "\n", "from atomic_agents.agents.base_agent import BaseAgent, BaseAgentConfig\n", - "from atomic_agents.lib.components.system_prompt_generator import SystemPromptContextProviderBase, SystemPromptGenerator, SystemPromptInfo\n", - "from atomic_agents.lib.tools.yt_transcript_scraper import YouTubeTranscriptTool, YouTubeTranscriptToolConfig, YouTubeTranscriptToolSchema" + "from atomic_agents.lib.components.system_prompt_generator import SystemPromptContextProviderBase, SystemPromptGenerator\n", + "from atomic_agents.lib.tools.yt_transcript_scraper_tool import YouTubeTranscriptTool, YouTubeTranscriptToolConfig, YouTubeTranscriptToolInputSchema" ] }, { @@ -151,8 +151,7 @@ " return f'VIDEO TRANSCRIPT: \"{self.transcript}\"\\n\\nDURATION: {self.duration}\\n\\nMETADATA: {self.metadata}'\n", " \n", "transcript_provider = YtTranscriptProvider(title='YouTube Transcript')\n", - "\n", - "system_prompt_info = SystemPromptInfo(\n", + "system_prompt_generator = SystemPromptGenerator(\n", " background=[\n", " 'This Assistant is an expert at extracting knowledge and other insightful and interesting information from YouTube transcripts.'\n", " ],\n", @@ -213,8 +212,8 @@ "agent = BaseAgent(\n", " config=BaseAgentConfig(\n", " client=client,\n", - " system_prompt_generator=SystemPromptGenerator(system_prompt_info),\n", - " model='gpt-3.5-turbo',\n", + " system_prompt_generator=system_prompt_generator,\n", + " model='gpt-4o-mini',\n", " output_schema=ResponseModel\n", " )\n", ")" @@ -236,7 +235,7 @@ "outputs": [], "source": [ "video_url = input('Enter the YouTube video URL: ')\n", - "scraped_transcript = yt_scraper_tool.run(YouTubeTranscriptToolSchema(video_url=video_url))\n", + "scraped_transcript = yt_scraper_tool.run(YouTubeTranscriptToolInputSchema(video_url=video_url))\n", "transcript_provider.transcript = scraped_transcript.transcript\n", "transcript_provider.duration = scraped_transcript.duration\n", "transcript_provider.metadata = scraped_transcript.metadata\n", diff --git a/examples/product_finder.py b/examples/product_finder.py index a75b430c..afbc69d4 100644 --- a/examples/product_finder.py +++ b/examples/product_finder.py @@ -7,16 +7,16 @@ from rich.console import Console from atomic_agents.lib.components.agent_memory import AgentMemory -from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator, SystemPromptInfo -from atomic_agents.agents.base_agent import BaseAgent, BaseAgentOutputSchema, BaseAgentConfig -from atomic_agents.lib.tools.search.searx_tool import SearxNGSearchTool, SearxNGSearchToolConfig, SearxNGSearchToolSchema +from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator +from atomic_agents.agents.base_agent import BaseAgent, BaseAgentOutputSchema, BaseAgentConfig, BaseIOSchema +from atomic_agents.lib.tools.search.searxng_tool import SearxNGTool, SearxNGToolConfig, SearxNGToolInputSchema # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # Define system prompt information -system_prompt = SystemPromptInfo( +system_prompt_generator = SystemPromptGenerator( background=[ 'This assistant is a product finder AI designed to help users find products based on their preferences by asking clarifying questions.', ], @@ -24,19 +24,17 @@ 'Greet the user and introduce yourself as a product finder assistant.', 'Ask the user questions to gather information about the product they are looking for.', 'Use the chat responses to gather all necessary information from the user.', - 'Once sufficient information is gathered, use the SearxNGSearchTool to search for products.', + 'Once sufficient information is gathered, use the SearxNGTool to search for products.', 'Summarize the search results and provide recommendations to the user.', ], output_instructions=[ + 'Always think in steps before answering using internal reasoning.', 'Provide helpful and relevant information to assist the user.', 'Be friendly and respectful in all interactions.', 'Ensure that the chat responses are used to ask clarifying questions and gather information, and the search tool is used to find products.' ] ) -# Initialize the system prompt generator -system_prompt_generator = SystemPromptGenerator(system_prompt) - # Initialize chat memory memory = AgentMemory() initial_memory = [ @@ -49,12 +47,13 @@ # Initialize the client client = instructor.from_openai(openai.OpenAI()) -# Initialize the SearxNGSearchTool -searx_tool = SearxNGSearchTool(SearxNGSearchToolConfig(base_url=os.getenv('SEARXNG_BASE_URL'), max_results=5)) +# Initialize the SearxNGTool +searxng_tool = SearxNGTool(SearxNGToolConfig(base_url=os.getenv('SEARXNG_BASE_URL'), max_results=5)) # Define a custom response schema -class OutputSchema(BaseModel): - chosen_schema: Union[BaseAgentOutputSchema, SearxNGSearchToolSchema] = Field(..., description='The response from the chat agent.') +class OutputSchema(BaseIOSchema): + internal_reasoning: List[str] = Field(..., description='The internal reasoning behind the response.') + chosen_schema: Union[BaseAgentOutputSchema, SearxNGToolInputSchema] = Field(..., description='The response from the chat agent. Every response must use chosen_schema to indicate the type of response (A chat message, or a search request)') class Config: title = 'OutputSchema' @@ -64,7 +63,7 @@ class Config: agent_config = BaseAgentConfig( client=client, system_prompt_generator=system_prompt_generator, - model='gpt-3.5-turbo', + model='gpt-4o-mini', memory=memory, output_schema=OutputSchema ) @@ -85,8 +84,8 @@ class Config: logger.info(f'Chosen schema: {response.chosen_schema}') - if isinstance(response.chosen_schema, SearxNGSearchToolSchema): - search_results = searx_tool.run(response.chosen_schema) + if isinstance(response.chosen_schema, SearxNGToolInputSchema): + search_results = searxng_tool.run(response.chosen_schema) agent.memory.add_message('assistant', f'INTERNAL THOUGHT: I have found the following products: {search_results.results}\n\n I will now summarize the results for the user.') output = agent.run().chosen_schema.chat_message diff --git a/examples/shared_context.py b/examples/shared_context.py index 54b367e9..7ffb3b15 100644 --- a/examples/shared_context.py +++ b/examples/shared_context.py @@ -21,14 +21,14 @@ def update_context(self, key: str, value: str): agent_one = BaseAgent( config=BaseAgentConfig( client=instructor.from_openai(openai.OpenAI()), - model='gpt-3.5-turbo', + model='gpt-4o-mini', ) ) agent_two = BaseAgent( config=BaseAgentConfig( client=instructor.from_openai(openai.OpenAI()), - model='gpt-3.5-turbo', + model='gpt-4o-mini', ) ) diff --git a/examples/simple_web_search_agent.py b/examples/simple_web_search_agent.py index 722c5200..897fb5f5 100644 --- a/examples/simple_web_search_agent.py +++ b/examples/simple_web_search_agent.py @@ -5,24 +5,24 @@ from rich.console import Console from atomic_agents.agents.tool_interface_agent import ToolInterfaceAgent, ToolInterfaceAgentConfig -from atomic_agents.lib.tools.search.searx_tool import SearxNGSearchTool, SearxNGSearchToolConfig +from atomic_agents.lib.tools.search.searxng_tool import SearxNGTool, SearxNGToolConfig -def initialize_searx_tool(): +def initialize_searxng_tool(): """ - Initialize the SearxNGSearchTool with configuration. + Initialize the SearxNGTool with configuration. """ base_url = os.getenv('SEARXNG_BASE_URL') - config = SearxNGSearchToolConfig(base_url=base_url, max_results=10) - return SearxNGSearchTool(config) + config = SearxNGToolConfig(base_url=base_url, max_results=10) + return SearxNGTool(config) -def initialize_agent(client, searx_tool): +def initialize_agent(client, searxng_tool): """ - Initialize the ToolInterfaceAgent with the given client and SearxNGSearchTool. + Initialize the ToolInterfaceAgent with the given client and SearxNGTool. """ agent_config = ToolInterfaceAgentConfig( client=client, - model='gpt-3.5-turbo', - tool_instance=searx_tool, + model='gpt-4o-mini', + tool_instance=searxng_tool, return_raw_output=False ) return ToolInterfaceAgent(config=agent_config) @@ -30,10 +30,10 @@ def initialize_agent(client, searx_tool): def main(): console = Console() client = instructor.from_openai(openai.OpenAI()) - searx_tool = initialize_searx_tool() - agent = initialize_agent(client, searx_tool) + searxng_tool = initialize_searxng_tool() + agent = initialize_agent(client, searxng_tool) - console.print("ToolInterfaceAgent with SearxNGSearchTool is ready.") + console.print("ToolInterfaceAgent with SearxNGTool is ready.") while True: user_input = input('You: ') @@ -42,7 +42,7 @@ def main(): break # Fix this - response = agent.run(agent.input_schema(tool_input_SearxNGSearchTool=user_input)) + response = agent.run(agent.input_schema(tool_input_SearxNGTool=user_input)) console.print(f'Agent: {response.chat_message}') if __name__ == "__main__": diff --git a/examples/system_prompt_generator_demo.py b/examples/system_prompt_generator_demo.py index f293d333..6a8502ef 100644 --- a/examples/system_prompt_generator_demo.py +++ b/examples/system_prompt_generator_demo.py @@ -1,6 +1,6 @@ from datetime import datetime from unittest.mock import Mock -from atomic_agents.lib.components.system_prompt_generator import SystemPromptContextProviderBase, SystemPromptGenerator, SystemPromptInfo +from atomic_agents.lib.components.system_prompt_generator import SystemPromptContextProviderBase, SystemPromptGenerator class CurrentDateProvider(SystemPromptContextProviderBase): @@ -19,7 +19,7 @@ def get_info(self) -> str: "Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat." ) -system_prompt = SystemPromptInfo( +system_prompt_generator = SystemPromptGenerator( background=[ 'This assistant is a general-purpose AI designed to be helpful and friendly.', ], @@ -38,5 +38,4 @@ def get_info(self) -> str: } ) -system_prompt_generator = SystemPromptGenerator(system_prompt) print(system_prompt_generator.generate_prompt()) diff --git a/examples/yelp_restaurant_finder.py b/examples/yelp_restaurant_finder.py index 4656c1bb..fc83b204 100644 --- a/examples/yelp_restaurant_finder.py +++ b/examples/yelp_restaurant_finder.py @@ -7,16 +7,16 @@ from rich.console import Console from atomic_agents.lib.components.agent_memory import AgentMemory -from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator, SystemPromptInfo +from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator from atomic_agents.agents.base_agent import BaseAgent, BaseAgentOutputSchema, BaseAgentConfig -from atomic_agents.lib.tools.yelp_restaurant_finder_tool import YelpSearchTool, YelpSearchToolConfig, YelpSearchToolSchema +from atomic_agents.lib.tools.yelp_restaurant_finder_tool import YelpSearchTool, YelpSearchToolConfig, YelpSearchToolInputSchema # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # Define system prompt information including background, steps, and output instructions -system_prompt = SystemPromptInfo( +system_prompt_generator = SystemPromptGenerator( background=[ 'This assistant is a restaurant finder AI designed to help users find the best restaurants based on their preferences by asking clarifying questions.', ], @@ -34,9 +34,6 @@ ] ) -# Initialize the system prompt generator with the defined system prompt and dynamic info providers -system_prompt_generator = SystemPromptGenerator(system_prompt) - # Initialize chat memory to store conversation history memory = AgentMemory() # Define initial memory with a greeting message from the assistant @@ -57,7 +54,7 @@ # Define a custom response schema that can handle both chat responses and Yelp search tool responses class OutputSchema(BaseModel): - chosen_schema: Union[BaseAgentOutputSchema, YelpSearchToolSchema] = Field(..., description='The response from the chat agent.') + chosen_schema: Union[BaseAgentOutputSchema, YelpSearchToolInputSchema] = Field(..., description='The response from the chat agent.') class Config: title = 'OutputSchema' @@ -71,7 +68,7 @@ class Config: agent_config = BaseAgentConfig( client=client, system_prompt_generator=system_prompt_generator, - model='gpt-3.5-turbo', + model='gpt-4o-mini', memory=memory, output_schema=OutputSchema ) @@ -94,7 +91,7 @@ class Config: logger.info(f'Chosen schema: {response.chosen_schema}') # Check the type of the response schema - if isinstance(response.chosen_schema, YelpSearchToolSchema): + if isinstance(response.chosen_schema, YelpSearchToolInputSchema): output = yelp_tool.run(response.chosen_schema) # In this example, we will add a simple "internal thought" to the chat memory followed by an empty agent.run() call. diff --git a/examples/youtube_knowledge_extraction.py b/examples/youtube_knowledge_extraction.py index 9ee85efe..4ef2de01 100644 --- a/examples/youtube_knowledge_extraction.py +++ b/examples/youtube_knowledge_extraction.py @@ -7,8 +7,8 @@ from typing import List, Optional from atomic_agents.agents.base_agent import BaseAgent, BaseAgentConfig -from atomic_agents.lib.components.system_prompt_generator import SystemPromptContextProviderBase, SystemPromptGenerator, SystemPromptInfo -from atomic_agents.lib.tools.yt_transcript_scraper import YouTubeTranscriptTool, YouTubeTranscriptToolConfig, YouTubeTranscriptToolSchema +from atomic_agents.lib.components.system_prompt_generator import SystemPromptContextProviderBase, SystemPromptGenerator +from atomic_agents.lib.tools.yt_transcript_scraper_tool import YouTubeTranscriptTool, YouTubeTranscriptToolConfig, YouTubeTranscriptToolInputSchema console = Console() @@ -31,7 +31,7 @@ def get_info(self) -> str: transcript_provider = YtTranscriptProvider(title='YouTube Transcript') # Define the system prompt information -system_prompt_info = SystemPromptInfo( +system_prompt_generator = SystemPromptGenerator( background=[ 'This Assistant is an expert at extracting knowledge and other insightful and interesting information from YouTube transcripts.' ], @@ -61,14 +61,14 @@ class ResponseModel(BaseModel): agent = BaseAgent( config = BaseAgentConfig( client=client, - model='gpt-3.5-turbo', - system_prompt_generator=SystemPromptGenerator(system_prompt_info), + model='gpt-4o-mini', + system_prompt_generator=system_prompt_generator, output_schema=ResponseModel ) ) video_url = input('Enter the YouTube video URL: ') -scraped_transcript = yt_scraper_tool.run(YouTubeTranscriptToolSchema(video_url=video_url)) +scraped_transcript = yt_scraper_tool.run(YouTubeTranscriptToolInputSchema(video_url=video_url)) transcript_provider.transcript = scraped_transcript.transcript transcript_provider.duration = scraped_transcript.duration transcript_provider.metadata = scraped_transcript.metadata diff --git a/tests/agents/test_base_agent.py b/tests/agents/test_base_agent.py index acbce59b..c3ce1283 100644 --- a/tests/agents/test_base_agent.py +++ b/tests/agents/test_base_agent.py @@ -3,7 +3,7 @@ from pydantic import BaseModel import instructor from atomic_agents.agents.base_agent import ( - BaseAgentIO, + BaseIOSchema, BaseAgent, BaseAgentConfig, BaseAgentInputSchema, @@ -12,7 +12,6 @@ AgentMemory, SystemPromptContextProviderBase ) -from atomic_agents.lib.components.system_prompt_generator import SystemPromptInfo @pytest.fixture def mock_instructor(): @@ -32,15 +31,14 @@ def mock_memory(): def mock_system_prompt_generator(): mock = Mock(spec=SystemPromptGenerator) mock.generate_prompt.return_value = "Mocked system prompt" - mock.system_prompt_info = Mock(spec=SystemPromptInfo) - mock.system_prompt_info.context_providers = {} + mock.context_providers = {} return mock @pytest.fixture def agent_config(mock_instructor, mock_memory, mock_system_prompt_generator): return BaseAgentConfig( client=mock_instructor, - model="gpt-3.5-turbo", + model="gpt-4o-mini", memory=mock_memory, system_prompt_generator=mock_system_prompt_generator ) @@ -51,7 +49,7 @@ def agent(agent_config): def test_initialization(agent, mock_instructor, mock_memory, mock_system_prompt_generator): assert agent.client == mock_instructor - assert agent.model == "gpt-3.5-turbo" + assert agent.model == "gpt-4o-mini" assert agent.memory == mock_memory assert agent.system_prompt_generator == mock_system_prompt_generator assert agent.input_schema == BaseAgentInputSchema @@ -75,7 +73,7 @@ def test_get_response(agent, mock_instructor, mock_memory, mock_system_prompt_ge assert response == mock_response mock_instructor.chat.completions.create.assert_called_once_with( - model="gpt-3.5-turbo", + model="gpt-4o-mini", messages=[ {'role': 'system', 'content': 'System prompt'}, {'role': 'user', 'content': 'Hello'} @@ -102,7 +100,7 @@ def test_run(agent): def test_get_context_provider(agent, mock_system_prompt_generator): mock_provider = Mock(spec=SystemPromptContextProviderBase) - mock_system_prompt_generator.system_prompt_info.context_providers = { + mock_system_prompt_generator.context_providers = { 'test_provider': mock_provider } @@ -116,17 +114,17 @@ def test_register_context_provider(agent, mock_system_prompt_generator): mock_provider = Mock(spec=SystemPromptContextProviderBase) agent.register_context_provider('new_provider', mock_provider) - assert 'new_provider' in mock_system_prompt_generator.system_prompt_info.context_providers - assert mock_system_prompt_generator.system_prompt_info.context_providers['new_provider'] == mock_provider + assert 'new_provider' in mock_system_prompt_generator.context_providers + assert mock_system_prompt_generator.context_providers['new_provider'] == mock_provider def test_unregister_context_provider(agent, mock_system_prompt_generator): mock_provider = Mock(spec=SystemPromptContextProviderBase) - mock_system_prompt_generator.system_prompt_info.context_providers = { + mock_system_prompt_generator.context_providers = { 'test_provider': mock_provider } agent.unregister_context_provider('test_provider') - assert 'test_provider' not in mock_system_prompt_generator.system_prompt_info.context_providers + assert 'test_provider' not in mock_system_prompt_generator.context_providers with pytest.raises(KeyError): agent.unregister_context_provider('non_existent_provider') @@ -140,7 +138,7 @@ class CustomOutputSchema(BaseModel): custom_config = BaseAgentConfig( client=mock_instructor, - model="gpt-3.5-turbo", + model="gpt-4o-mini", input_schema=CustomInputSchema, output_schema=CustomOutputSchema ) @@ -152,7 +150,7 @@ class CustomOutputSchema(BaseModel): def test_base_agent_io_str_and_rich(): - class TestIO(BaseAgentIO): + class TestIO(BaseIOSchema): field: str test_io = TestIO(field="test") @@ -184,4 +182,4 @@ def test_run(agent, mock_memory): ]) if __name__ == '__main__': - pytest.main() \ No newline at end of file + pytest.main() diff --git a/tests/agents/test_tool_interface_agent.py b/tests/agents/test_tool_interface_agent.py index 3edbc871..d912ae82 100644 --- a/tests/agents/test_tool_interface_agent.py +++ b/tests/agents/test_tool_interface_agent.py @@ -3,20 +3,19 @@ from pydantic import Field import instructor -from atomic_agents.agents.tool_interface_agent import ToolInterfaceAgent, ToolInterfaceAgentConfig, ToolInputModel -from atomic_agents.lib.tools.base import BaseTool -from atomic_agents.agents.base_agent import BaseAgentIO, BaseAgentOutputSchema -from atomic_agents.lib.components.agent_memory import Message +from atomic_agents.agents.tool_interface_agent import ToolInterfaceAgent, ToolInterfaceAgentConfig +from atomic_agents.lib.tools.base_tool import BaseTool +from atomic_agents.agents.base_agent import BaseIOSchema, BaseAgentOutputSchema class MockTool(BaseTool): - class InputSchema(BaseAgentIO): + class InputSchema(BaseIOSchema): query: str = Field(..., description="The query to process") class Config: title = "MockTool" description = "A mock tool for testing" - class OutputSchema(BaseAgentIO): + class OutputSchema(BaseIOSchema): result: str = Field(..., description="The result of the tool operation") input_schema = InputSchema @@ -39,7 +38,7 @@ def mock_instructor(): def tool_interface_agent(mock_tool, mock_instructor): config = ToolInterfaceAgentConfig( client=mock_instructor, - model="gpt-3.5-turbo", + model="gpt-4o-mini", tool_instance=mock_tool, return_raw_output=False ) @@ -49,7 +48,7 @@ def test_tool_interface_agent_initialization(tool_interface_agent, mock_tool): assert tool_interface_agent.tool_instance == mock_tool assert tool_interface_agent.return_raw_output == False assert tool_interface_agent.input_schema.__name__ == "MockTool" - assert issubclass(tool_interface_agent.input_schema, ToolInputModel) + assert issubclass(tool_interface_agent.input_schema, BaseIOSchema) assert tool_interface_agent.output_schema == BaseAgentOutputSchema def test_tool_interface_agent_input_schema(tool_interface_agent): @@ -80,7 +79,7 @@ def test_tool_interface_agent_get_response_without_raw_output(mock_format_tool_m def test_tool_interface_agent_get_response_with_raw_output(mock_format_tool_message, mock_get_response, mock_tool, mock_instructor): config = ToolInterfaceAgentConfig( client=mock_instructor, - model="gpt-3.5-turbo", + model="gpt-4o-mini", tool_instance=mock_tool, return_raw_output=True ) @@ -123,10 +122,5 @@ def test_tool_interface_agent_memory_updates(mock_format_tool_message, mock_get_ # Check if the mocked result is in the memory assert any("Mocked result" in msg.content for msg in history if isinstance(msg.content, str)) - # Print the memory contents for debugging - print("Memory contents:") - for msg in history: - print(f"Role: {msg.role}, Content: {msg.content}") - # Assert that at least 3 messages were added (TOOL CALL, TOOL RESPONSE, and possibly the final response) - assert len(history) >= 3 \ No newline at end of file + assert len(history) >= 3 diff --git a/tests/lib/components/test_system_prompt_generator.py b/tests/lib/components/test_system_prompt_generator.py index 86109cd2..49fa2827 100644 --- a/tests/lib/components/test_system_prompt_generator.py +++ b/tests/lib/components/test_system_prompt_generator.py @@ -1,212 +1,135 @@ import pytest -from datetime import datetime -from unittest.mock import Mock - -from atomic_agents.lib.components.system_prompt_generator import ( - SystemPromptContextProviderBase, - SystemPromptInfo, - SystemPromptGenerator, -) +from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator, SystemPromptContextProviderBase class MockContextProvider(SystemPromptContextProviderBase): - def get_info(self): - return "Mock info" - -class CurrentDateProvider(SystemPromptContextProviderBase): - def __init__(self, title, format: str = '%Y-%m-%d %H:%M:%S'): + def __init__(self, title: str, info: str): super().__init__(title) - self.format = format - - def get_current_datetime(self): - return datetime.now() + self._info = info def get_info(self) -> str: - current_datetime = self.get_current_datetime() - return f'The current date, in the format "{self.format}", is {current_datetime.strftime(self.format)}' + return self._info + +def test_system_prompt_generator_default_initialization(): + generator = SystemPromptGenerator() + assert generator.background == ["This is a conversation with a helpful and friendly AI assistant."] + assert generator.steps == [] + assert generator.output_instructions == [ + "Always respond using the proper JSON schema.", + "Always use the available additional information and context to enhance the response." + ] + assert generator.context_providers == {} + +def test_system_prompt_generator_custom_initialization(): + background = ["Custom background"] + steps = ["Step 1", "Step 2"] + output_instructions = ["Custom instruction"] + context_providers = { + "provider1": MockContextProvider("Provider 1", "Info 1"), + "provider2": MockContextProvider("Provider 2", "Info 2") + } + + generator = SystemPromptGenerator( + background=background, + steps=steps, + output_instructions=output_instructions, + context_providers=context_providers + ) -class LoremIpsumProvider(SystemPromptContextProviderBase): - def get_info(self) -> str: - return ( - "Lorem ipsum dolor sit amet, consectetur adipiscing elit. " - "Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. " - "Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat." - ) - -@pytest.fixture -def default_generator(): - return SystemPromptGenerator() - -@pytest.fixture -def custom_info(): - current_date_provider = CurrentDateProvider(title="Current date", format="%Y-%m-%d %H:%M:%S") - return SystemPromptInfo( - background=["Custom background"], - steps=["Custom step"], + assert generator.background == background + assert generator.steps == steps + assert generator.output_instructions == [ + "Custom instruction", + "Always respond using the proper JSON schema.", + "Always use the available additional information and context to enhance the response." + ] + assert generator.context_providers == context_providers + +def test_generate_prompt_without_context_providers(): + generator = SystemPromptGenerator( + background=["Background info"], + steps=["Step 1", "Step 2"], + output_instructions=["Custom instruction"] + ) + + expected_prompt = """# IDENTITY and PURPOSE +- Background info + +# INTERNAL ASSISTANT STEPS +- Step 1 +- Step 2 + +# OUTPUT INSTRUCTIONS +- Custom instruction +- Always respond using the proper JSON schema. +- Always use the available additional information and context to enhance the response.""" + + assert generator.generate_prompt() == expected_prompt + +def test_generate_prompt_with_context_providers(): + generator = SystemPromptGenerator( + background=["Background info"], + steps=["Step 1"], output_instructions=["Custom instruction"], context_providers={ - "test": MockContextProvider("Test Provider"), - "date": current_date_provider, - "lorem": LoremIpsumProvider(title="Lorem Ipsum") + "provider1": MockContextProvider("Provider 1", "Info 1"), + "provider2": MockContextProvider("Provider 2", "Info 2") } ) -@pytest.fixture -def custom_generator(custom_info): - return SystemPromptGenerator(custom_info) - -class TestSystemPromptContextProviderBase: - def test_abstract_base_class(self): - with pytest.raises(TypeError): - SystemPromptContextProviderBase("Test") - - def test_repr(self): - class ConcreteProvider(SystemPromptContextProviderBase): - def get_info(self): - return "Test Info" - - provider = ConcreteProvider("Test") - assert repr(provider) == "Test Info" - -class TestSystemPromptInfo: - def test_default_values(self): - info = SystemPromptInfo(background=["Test background"]) - assert info.background == ["Test background"] - assert info.steps == [] - assert info.output_instructions == [] - assert info.context_providers == {} - - def test_custom_values(self): - info = SystemPromptInfo( - background=["Test background"], - steps=["Step 1", "Step 2"], - output_instructions=["Instruction 1"], - context_providers={"test": None} - ) - assert info.background == ["Test background"] - assert info.steps == ["Step 1", "Step 2"] - assert info.output_instructions == ["Instruction 1"] - assert info.context_providers == {"test": None} - -class TestSystemPromptGeneratorBasic: - def test_default_initialization(self, default_generator): - assert default_generator.system_prompt_info.background == ['This is a conversation with a helpful and friendly AI assistant.'] - assert len(default_generator.system_prompt_info.output_instructions) == 2 - - def test_custom_initialization(self, custom_generator, custom_info): - assert custom_generator.system_prompt_info == custom_info - assert len(custom_generator.system_prompt_info.output_instructions) == 3 - - def test_generate_prompt_default(self, default_generator): - prompt = default_generator.generate_prompt() - assert "# IDENTITY and PURPOSE" in prompt - assert "# OUTPUT INSTRUCTIONS" in prompt - assert "# INTERNAL ASSISTANT STEPS" not in prompt - assert "# EXTRA INFORMATION AND CONTEXT" not in prompt - - def test_generate_prompt_custom(self, custom_generator): - prompt = custom_generator.generate_prompt() - - assert "# IDENTITY and PURPOSE" in prompt - assert "# INTERNAL ASSISTANT STEPS" in prompt - assert "# OUTPUT INSTRUCTIONS" in prompt - assert "# EXTRA INFORMATION AND CONTEXT" in prompt - assert "Custom background" in prompt - assert "Custom step" in prompt - assert "Custom instruction" in prompt - assert "Mock info" in prompt - - def test_empty_system_prompt_info(self): - empty_info = SystemPromptInfo(background=[]) - empty_generator = SystemPromptGenerator(empty_info) - prompt = empty_generator.generate_prompt() - expected_prompt = "# OUTPUT INSTRUCTIONS\n- Always respond using the proper JSON schema.\n- Always use the available additional information and context to enhance the response." - assert prompt == expected_prompt - - def test_partial_system_prompt_info(self): - partial_info = SystemPromptInfo( - background=["Test background"], - steps=["Test step"], - ) - partial_generator = SystemPromptGenerator(partial_info) - prompt = partial_generator.generate_prompt() - assert "# IDENTITY and PURPOSE" in prompt - assert "Test background" in prompt - assert "# INTERNAL ASSISTANT STEPS" in prompt - assert "Test step" in prompt - assert "# OUTPUT INSTRUCTIONS" in prompt - assert "# EXTRA INFORMATION AND CONTEXT" not in prompt - -class TestSystemPromptGeneratorContextProviders: - def test_generate_prompt_with_context_providers(self, custom_generator): - prompt = custom_generator.generate_prompt() - - assert "# IDENTITY and PURPOSE" in prompt - assert "# INTERNAL ASSISTANT STEPS" in prompt - assert "# OUTPUT INSTRUCTIONS" in prompt - assert "# EXTRA INFORMATION AND CONTEXT" in prompt - assert "Custom background" in prompt - assert "Custom step" in prompt - assert "Custom instruction" in prompt - assert "Mock info" in prompt - assert "## Current date" in prompt - assert "The current date, in the format" in prompt - assert "## Lorem Ipsum" in prompt - assert "Lorem ipsum dolor sit amet" in prompt - - def test_current_date_provider(self, custom_generator, monkeypatch): - mock_datetime = datetime(2023, 5, 1, 12, 0, 0) - monkeypatch.setattr(custom_generator.system_prompt_info.context_providers['date'], 'get_current_datetime', lambda: mock_datetime) - prompt = custom_generator.generate_prompt() - assert 'The current date, in the format "%Y-%m-%d %H:%M:%S", is 2023-05-01 12:00:00' in prompt - - def test_custom_output_instructions(self): - custom_instructions = ["Custom instruction 1", "Custom instruction 2"] - custom_info = SystemPromptInfo( - background=["Test background"], - output_instructions=custom_instructions - ) - custom_generator = SystemPromptGenerator(custom_info) - prompt = custom_generator.generate_prompt() - for instruction in custom_instructions: - assert instruction in prompt - assert "Always respond using the proper JSON schema." in prompt - assert "Always use the available additional information and context to enhance the response." in prompt - - def test_empty_context_provider(self): - class EmptyProvider(SystemPromptContextProviderBase): - def get_info(self): - return "" - - info = SystemPromptInfo( - background=["Test"], - context_providers={"empty": EmptyProvider("Empty")} - ) - generator = SystemPromptGenerator(info) - prompt = generator.generate_prompt() - assert "## Empty" not in prompt - - def test_multiple_context_providers(self): - class Provider1(SystemPromptContextProviderBase): - def get_info(self): - return "Info 1" - - class Provider2(SystemPromptContextProviderBase): - def get_info(self): - return "Info 2" - - info = SystemPromptInfo( - background=["Test"], - context_providers={ - "provider1": Provider1("Provider 1"), - "provider2": Provider2("Provider 2") - } - ) - generator = SystemPromptGenerator(info) - prompt = generator.generate_prompt() - assert "## Provider 1" in prompt - assert "Info 1" in prompt - assert "## Provider 2" in prompt - assert "Info 2" in prompt - -if __name__ == '__main__': - pytest.main() \ No newline at end of file + expected_prompt = """# IDENTITY and PURPOSE +- Background info + +# INTERNAL ASSISTANT STEPS +- Step 1 + +# OUTPUT INSTRUCTIONS +- Custom instruction +- Always respond using the proper JSON schema. +- Always use the available additional information and context to enhance the response. + +# EXTRA INFORMATION AND CONTEXT +## Provider 1 +Info 1 + +## Provider 2 +Info 2""" + + assert generator.generate_prompt() == expected_prompt + +def test_generate_prompt_with_empty_sections(): + generator = SystemPromptGenerator( + background=[], + steps=[], + output_instructions=[] + ) + + expected_prompt = """# IDENTITY and PURPOSE +- This is a conversation with a helpful and friendly AI assistant. + +# OUTPUT INSTRUCTIONS +- Always respond using the proper JSON schema. +- Always use the available additional information and context to enhance the response.""" + + assert generator.generate_prompt() == expected_prompt + +def test_context_provider_repr(): + provider = MockContextProvider("Test Provider", "Test Info") + assert repr(provider) == "Test Info" + +def test_generate_prompt_with_empty_context_provider(): + empty_provider = MockContextProvider("Empty Provider", "") + generator = SystemPromptGenerator( + background=["Background"], + context_providers={"empty": empty_provider} + ) + + expected_prompt = """# IDENTITY and PURPOSE +- Background + +# OUTPUT INSTRUCTIONS +- Always respond using the proper JSON schema. +- Always use the available additional information and context to enhance the response. + +# EXTRA INFORMATION AND CONTEXT""" + + assert generator.generate_prompt() == expected_prompt \ No newline at end of file diff --git a/tests/lib/models/test_web_document.py b/tests/lib/models/test_web_document.py index 5efd2409..d451fb4a 100644 --- a/tests/lib/models/test_web_document.py +++ b/tests/lib/models/test_web_document.py @@ -1,7 +1,7 @@ import pytest from pydantic import ValidationError from atomic_agents.lib.models.web_document import WebDocumentMetadata, WebDocument -from atomic_agents.agents.base_agent import BaseAgentIO +from atomic_agents.agents.base_agent import BaseIOSchema def test_web_document_metadata_creation(): metadata = WebDocumentMetadata(url="https://example.com") @@ -49,7 +49,7 @@ def test_web_document_content_required(): def test_web_document_inheritance(): doc = WebDocument(content="This is the content") assert isinstance(doc, WebDocument) - assert isinstance(doc, BaseAgentIO) + assert isinstance(doc, BaseIOSchema) def test_web_document_metadata_serialization(): metadata = WebDocumentMetadata(url="https://example.com", title="Example") diff --git a/tests/lib/tools/test_base.py b/tests/lib/tools/test_base_tool.py similarity index 91% rename from tests/lib/tools/test_base.py rename to tests/lib/tools/test_base_tool.py index 93ce33b5..fed8103e 100644 --- a/tests/lib/tools/test_base.py +++ b/tests/lib/tools/test_base_tool.py @@ -1,15 +1,15 @@ import pytest from pydantic import BaseModel -from atomic_agents.lib.tools.base import BaseToolConfig, BaseTool -from atomic_agents.agents.base_agent import BaseAgentIO +from atomic_agents.lib.tools.base_tool import BaseToolConfig, BaseTool +from atomic_agents.agents.base_agent import BaseIOSchema # Mock classes for testing -class MockInputSchema(BaseAgentIO): +class MockInputSchema(BaseIOSchema): class Config: title = "Mock Input" description = "Mock input description" -class MockOutputSchema(BaseAgentIO): +class MockOutputSchema(BaseIOSchema): result: str class MockTool(BaseTool): diff --git a/tests/lib/tools/test_calculator_tool.py b/tests/lib/tools/test_calculator_tool.py index 71405509..9290dc01 100644 --- a/tests/lib/tools/test_calculator_tool.py +++ b/tests/lib/tools/test_calculator_tool.py @@ -1,6 +1,6 @@ import pytest -from atomic_agents.lib.tools.calculator_tool import CalculatorTool, CalculatorToolSchema, CalculatorToolOutputSchema, CalculatorToolConfig -from atomic_agents.lib.tools.base import BaseTool, BaseToolConfig +from atomic_agents.lib.tools.calculator_tool import CalculatorTool, CalculatorToolInputSchema, CalculatorToolOutputSchema, CalculatorToolConfig +from atomic_agents.lib.tools.base_tool import BaseTool, BaseToolConfig def test_calculator_tool_initialization(): tool = CalculatorTool() @@ -16,7 +16,7 @@ def test_calculator_tool_with_custom_config(): def test_calculator_tool_input_schema(): tool = CalculatorTool() - assert tool.input_schema == CalculatorToolSchema + assert tool.input_schema == CalculatorToolInputSchema def test_calculator_tool_output_schema(): tool = CalculatorTool() @@ -32,33 +32,33 @@ def test_calculator_tool_output_schema(): ]) def test_calculator_tool_run_valid_expressions(expression, expected_result): tool = CalculatorTool() - result = tool.run(CalculatorToolSchema(expression=expression)) + result = tool.run(CalculatorToolInputSchema(expression=expression)) assert isinstance(result, CalculatorToolOutputSchema) assert result.result == expected_result def test_calculator_tool_run_complex_expression(): tool = CalculatorTool() expression = "sqrt(3**2 + 4**2)" - result = tool.run(CalculatorToolSchema(expression=expression)) + result = tool.run(CalculatorToolInputSchema(expression=expression)) assert isinstance(result, CalculatorToolOutputSchema) assert result.result == "5.00000000000000" def test_calculator_tool_run_invalid_expression(): tool = CalculatorTool() invalid_expression = "invalid + expression" - result = tool.run(CalculatorToolSchema(expression=invalid_expression)) + result = tool.run(CalculatorToolInputSchema(expression=invalid_expression)) assert isinstance(result, CalculatorToolOutputSchema) assert set(result.result.split(' + ')) == set(invalid_expression.split(' + ')) def test_calculator_tool_run_division_by_zero(): tool = CalculatorTool() - result = tool.run(CalculatorToolSchema(expression="1/0")) + result = tool.run(CalculatorToolInputSchema(expression="1/0")) assert isinstance(result, CalculatorToolOutputSchema) assert result.result == "zoo" # SymPy returns "zoo" for complex infinity def test_calculator_tool_run_with_variables(): tool = CalculatorTool() - result = tool.run(CalculatorToolSchema(expression="x + y")) + result = tool.run(CalculatorToolInputSchema(expression="x + y")) assert isinstance(result, CalculatorToolOutputSchema) assert result.result == "x + y" # The tool doesn't evaluate expressions with variables diff --git a/tests/lib/tools/test_content_scraping_tool.py b/tests/lib/tools/test_content_scraping_tool.py index 36711bfe..bef15c0f 100644 --- a/tests/lib/tools/test_content_scraping_tool.py +++ b/tests/lib/tools/test_content_scraping_tool.py @@ -2,11 +2,11 @@ from unittest.mock import patch, MagicMock from atomic_agents.lib.tools.content_scraping_tool import ( ContentScrapingTool, - ContentScrapingToolSchema, + ContentScrapingToolInputSchema, ContentScrapingToolOutputSchema, ContentScrapingToolConfig, ) -from atomic_agents.lib.tools.base import BaseTool, BaseToolConfig +from atomic_agents.lib.tools.base_tool import BaseTool, BaseToolConfig from atomic_agents.lib.utils.scraping.pdf_to_markdown import PdfToMarkdownConverter from atomic_agents.lib.utils.scraping.url_to_markdown import UrlToMarkdownConverter @@ -24,7 +24,7 @@ def test_content_scraping_tool_with_custom_config(): def test_content_scraping_tool_input_schema(): tool = ContentScrapingTool() - assert tool.input_schema == ContentScrapingToolSchema + assert tool.input_schema == ContentScrapingToolInputSchema def test_content_scraping_tool_output_schema(): tool = ContentScrapingTool() @@ -46,7 +46,7 @@ def test_content_scraping_tool_run(mock_head, content_type, expected_converter): with patch.object(expected_converter, "convert", return_value=mock_document): tool = ContentScrapingTool() - result = tool.run(ContentScrapingToolSchema(url="https://example.com")) + result = tool.run(ContentScrapingToolInputSchema(url="https://example.com")) assert isinstance(result, ContentScrapingToolOutputSchema) assert result.content == "Mocked content" @@ -64,14 +64,14 @@ def test_content_scraping_tool_run_unknown_content_type(mock_head): with patch.object(UrlToMarkdownConverter, "convert", return_value=mock_document): tool = ContentScrapingTool() - result = tool.run(ContentScrapingToolSchema(url="https://example.com")) + result = tool.run(ContentScrapingToolInputSchema(url="https://example.com")) assert isinstance(result, ContentScrapingToolOutputSchema) assert result.content == "Mocked content" assert result.metadata == {"key": "value"} def test_content_scraping_tool_schema_fields(): - schema = ContentScrapingToolSchema(url="https://example.com") + schema = ContentScrapingToolInputSchema(url="https://example.com") assert schema.url == "https://example.com" def test_content_scraping_tool_output_schema_fields(): diff --git a/tests/lib/tools/test_yelp_restaurant_finder_tool.py b/tests/lib/tools/test_yelp_restaurant_finder_tool.py index 03117c75..0c8811df 100644 --- a/tests/lib/tools/test_yelp_restaurant_finder_tool.py +++ b/tests/lib/tools/test_yelp_restaurant_finder_tool.py @@ -5,13 +5,13 @@ from atomic_agents.lib.tools.yelp_restaurant_finder_tool import ( YelpSearchTool, YelpSearchToolConfig, - YelpSearchToolSchema, + YelpSearchToolInputSchema, YelpSearchToolOutputSchema, YelpCategory, PriceRange ) -from atomic_agents.lib.tools.base import BaseTool -from atomic_agents.agents.base_agent import BaseAgentIO +from atomic_agents.lib.tools.base_tool import BaseTool +from atomic_agents.agents.base_agent import BaseIOSchema # Sample data SAMPLE_YELP_RESPONSE = { @@ -39,14 +39,14 @@ def test_yelp_search_tool_initialization(yelp_search_tool): assert yelp_search_tool.max_results == 10 def test_yelp_search_tool_input_schema(): - assert issubclass(YelpSearchToolSchema, BaseAgentIO) - assert "location" in YelpSearchToolSchema.model_fields - assert "term" in YelpSearchToolSchema.model_fields - assert "categories" in YelpSearchToolSchema.model_fields - assert "price" in YelpSearchToolSchema.model_fields + assert issubclass(YelpSearchToolInputSchema, BaseIOSchema) + assert "location" in YelpSearchToolInputSchema.model_fields + assert "term" in YelpSearchToolInputSchema.model_fields + assert "categories" in YelpSearchToolInputSchema.model_fields + assert "price" in YelpSearchToolInputSchema.model_fields def test_yelp_search_tool_output_schema(): - assert issubclass(YelpSearchToolOutputSchema, BaseAgentIO) + assert issubclass(YelpSearchToolOutputSchema, BaseIOSchema) assert "results" in YelpSearchToolOutputSchema.model_fields @patch('atomic_agents.lib.tools.yelp_restaurant_finder_tool.requests.get') @@ -56,7 +56,7 @@ def test_run_success(mock_get, yelp_search_tool): mock_response.json.return_value = SAMPLE_YELP_RESPONSE mock_get.return_value = mock_response - input_data = YelpSearchToolSchema( + input_data = YelpSearchToolInputSchema( location="Sample City", term="pizza", categories=[YelpCategory.ITALIAN], @@ -88,7 +88,7 @@ def test_run_api_error(mock_get, yelp_search_tool): mock_response.reason = "Bad Request" mock_get.return_value = mock_response - input_data = YelpSearchToolSchema(location="Sample City") + input_data = YelpSearchToolInputSchema(location="Sample City") with pytest.raises(Exception, match="Failed to fetch search results: 400 Bad Request"): yelp_search_tool.run(input_data) @@ -100,7 +100,7 @@ def test_run_with_all_parameters(mock_get, yelp_search_tool): mock_response.json.return_value = SAMPLE_YELP_RESPONSE mock_get.return_value = mock_response - input_data = YelpSearchToolSchema( + input_data = YelpSearchToolInputSchema( location="Sample City", term="pizza", categories=[YelpCategory.ITALIAN, YelpCategory.PIZZA], @@ -126,7 +126,7 @@ def test_run_no_results(mock_get, yelp_search_tool): mock_response.json.return_value = {"businesses": []} mock_get.return_value = mock_response - input_data = YelpSearchToolSchema(location="Sample City") + input_data = YelpSearchToolInputSchema(location="Sample City") result = yelp_search_tool.run(input_data) assert isinstance(result, YelpSearchToolOutputSchema) @@ -135,7 +135,7 @@ def test_run_no_results(mock_get, yelp_search_tool): def test_run_missing_api_key(): config = YelpSearchToolConfig(api_key="", max_results=10) tool = YelpSearchTool(config) - input_data = YelpSearchToolSchema(location="Sample City") + input_data = YelpSearchToolInputSchema(location="Sample City") with pytest.raises(ValueError, match="API key is required to use the Yelp API."): tool.run(input_data) \ No newline at end of file diff --git a/tests/lib/tools/test_yt_transcript_scraper.py b/tests/lib/tools/test_yt_transcript_scraper_tool.py similarity index 71% rename from tests/lib/tools/test_yt_transcript_scraper.py rename to tests/lib/tools/test_yt_transcript_scraper_tool.py index 7c3c6ab5..321a0dfe 100644 --- a/tests/lib/tools/test_yt_transcript_scraper.py +++ b/tests/lib/tools/test_yt_transcript_scraper_tool.py @@ -2,15 +2,15 @@ from unittest.mock import patch, Mock from pydantic import ValidationError from youtube_transcript_api import YouTubeTranscriptApi, NoTranscriptFound, TranscriptsDisabled -from atomic_agents.lib.tools.yt_transcript_scraper import YouTubeTranscriptToolSchema -from atomic_agents.lib.tools.yt_transcript_scraper import ( +from atomic_agents.lib.tools.yt_transcript_scraper_tool import YouTubeTranscriptToolInputSchema +from atomic_agents.lib.tools.yt_transcript_scraper_tool import ( YouTubeTranscriptTool, YouTubeTranscriptToolConfig, - YouTubeTranscriptToolSchema, + YouTubeTranscriptToolInputSchema, YouTubeTranscriptToolOutputSchema, ) -from atomic_agents.lib.tools.base import BaseTool -from atomic_agents.agents.base_agent import BaseAgentIO +from atomic_agents.lib.tools.base_tool import BaseTool +from atomic_agents.agents.base_agent import BaseIOSchema from youtube_transcript_api import YouTubeTranscriptApi @@ -37,12 +37,12 @@ def test_youtube_transcript_tool_initialization(youtube_transcript_tool): assert youtube_transcript_tool.api_key == "dummy_api_key" def test_youtube_transcript_tool_input_schema(): - assert issubclass(YouTubeTranscriptToolSchema, BaseAgentIO) - assert "video_url" in YouTubeTranscriptToolSchema.model_fields - assert "language" in YouTubeTranscriptToolSchema.model_fields + assert issubclass(YouTubeTranscriptToolInputSchema, BaseIOSchema) + assert "video_url" in YouTubeTranscriptToolInputSchema.model_fields + assert "language" in YouTubeTranscriptToolInputSchema.model_fields def test_youtube_transcript_tool_output_schema(): - assert issubclass(YouTubeTranscriptToolOutputSchema, BaseAgentIO) + assert issubclass(YouTubeTranscriptToolOutputSchema, BaseIOSchema) assert "transcript" in YouTubeTranscriptToolOutputSchema.model_fields assert "duration" in YouTubeTranscriptToolOutputSchema.model_fields assert "comments" in YouTubeTranscriptToolOutputSchema.model_fields @@ -53,15 +53,15 @@ def test_extract_video_id(): video_id = YouTubeTranscriptTool.extract_video_id(url) assert video_id == "dQw4w9WgXcQ" -@patch('atomic_agents.lib.tools.yt_transcript_scraper.YouTubeTranscriptTool.extract_video_id') -@patch('atomic_agents.lib.tools.yt_transcript_scraper.YouTubeTranscriptTool.fetch_video_metadata') -@patch('atomic_agents.lib.tools.yt_transcript_scraper.YouTubeTranscriptApi.get_transcript') +@patch('atomic_agents.lib.tools.yt_transcript_scraper_tool.YouTubeTranscriptTool.extract_video_id') +@patch('atomic_agents.lib.tools.yt_transcript_scraper_tool.YouTubeTranscriptTool.fetch_video_metadata') +@patch('atomic_agents.lib.tools.yt_transcript_scraper_tool.YouTubeTranscriptApi.get_transcript') def test_run(mock_get_transcript, mock_fetch_metadata, mock_extract_video_id, youtube_transcript_tool): mock_extract_video_id.return_value = "dQw4w9WgXcQ" mock_fetch_metadata.return_value = SAMPLE_VIDEO_INFO mock_get_transcript.return_value = SAMPLE_TRANSCRIPT - input_data = YouTubeTranscriptToolSchema(video_url="https://www.youtube.com/watch?v=dQw4w9WgXcQ") + input_data = YouTubeTranscriptToolInputSchema(video_url="https://www.youtube.com/watch?v=dQw4w9WgXcQ") result = youtube_transcript_tool.run(input_data) assert isinstance(result, YouTubeTranscriptToolOutputSchema) @@ -70,15 +70,15 @@ def test_run(mock_get_transcript, mock_fetch_metadata, mock_extract_video_id, yo assert result.comments == [] assert result.metadata == SAMPLE_VIDEO_INFO -@patch('atomic_agents.lib.tools.yt_transcript_scraper.YouTubeTranscriptTool.extract_video_id') -@patch('atomic_agents.lib.tools.yt_transcript_scraper.YouTubeTranscriptTool.fetch_video_metadata') -@patch('atomic_agents.lib.tools.yt_transcript_scraper.YouTubeTranscriptApi.get_transcript') +@patch('atomic_agents.lib.tools.yt_transcript_scraper_tool.YouTubeTranscriptTool.extract_video_id') +@patch('atomic_agents.lib.tools.yt_transcript_scraper_tool.YouTubeTranscriptTool.fetch_video_metadata') +@patch('atomic_agents.lib.tools.yt_transcript_scraper_tool.YouTubeTranscriptApi.get_transcript') def test_run_with_language(mock_get_transcript, mock_fetch_metadata, mock_extract_video_id, youtube_transcript_tool): mock_extract_video_id.return_value = "dQw4w9WgXcQ" mock_fetch_metadata.return_value = SAMPLE_VIDEO_INFO mock_get_transcript.return_value = SAMPLE_TRANSCRIPT - input_data = YouTubeTranscriptToolSchema( + input_data = YouTubeTranscriptToolInputSchema( video_url="https://www.youtube.com/watch?v=dQw4w9WgXcQ", language="en" ) @@ -86,15 +86,15 @@ def test_run_with_language(mock_get_transcript, mock_fetch_metadata, mock_extrac mock_get_transcript.assert_called_once_with("dQw4w9WgXcQ", languages=["en"]) -@patch('atomic_agents.lib.tools.yt_transcript_scraper.YouTubeTranscriptTool.extract_video_id') -@patch('atomic_agents.lib.tools.yt_transcript_scraper.YouTubeTranscriptTool.fetch_video_metadata') -@patch('atomic_agents.lib.tools.yt_transcript_scraper.YouTubeTranscriptApi.get_transcript') +@patch('atomic_agents.lib.tools.yt_transcript_scraper_tool.YouTubeTranscriptTool.extract_video_id') +@patch('atomic_agents.lib.tools.yt_transcript_scraper_tool.YouTubeTranscriptTool.fetch_video_metadata') +@patch('atomic_agents.lib.tools.yt_transcript_scraper_tool.YouTubeTranscriptApi.get_transcript') def test_run_transcript_not_found(mock_get_transcript, mock_fetch_metadata, mock_extract_video_id, youtube_transcript_tool): mock_extract_video_id.return_value = "dQw4w9WgXcQ" mock_fetch_metadata.return_value = SAMPLE_VIDEO_INFO mock_get_transcript.side_effect = Exception("No transcript found") - input_data = YouTubeTranscriptToolSchema(video_url="https://www.youtube.com/watch?v=dQw4w9WgXcQ") + input_data = YouTubeTranscriptToolInputSchema(video_url="https://www.youtube.com/watch?v=dQw4w9WgXcQ") with pytest.raises(Exception, match="No transcript found"): youtube_transcript_tool.run(input_data) @@ -107,32 +107,32 @@ def test_youtube_transcript_tool_config_validation(): assert config.api_key == "dummy_api_key" @patch.object(YouTubeTranscriptApi, 'get_transcript') -@patch('atomic_agents.lib.tools.yt_transcript_scraper.YouTubeTranscriptTool.fetch_video_metadata') -@patch('atomic_agents.lib.tools.yt_transcript_scraper.YouTubeTranscriptTool.extract_video_id') +@patch('atomic_agents.lib.tools.yt_transcript_scraper_tool.YouTubeTranscriptTool.fetch_video_metadata') +@patch('atomic_agents.lib.tools.yt_transcript_scraper_tool.YouTubeTranscriptTool.extract_video_id') def test_run_no_transcript_found(mock_extract_video_id, mock_fetch_metadata, mock_get_transcript, youtube_transcript_tool): mock_extract_video_id.return_value = "dQw4w9WgXcQ" mock_fetch_metadata.return_value = SAMPLE_VIDEO_INFO mock_get_transcript.side_effect = NoTranscriptFound("dQw4w9WgXcQ", "en", "Transcript not found") - input_data = YouTubeTranscriptToolSchema(video_url="https://www.youtube.com/watch?v=dQw4w9WgXcQ") + input_data = YouTubeTranscriptToolInputSchema(video_url="https://www.youtube.com/watch?v=dQw4w9WgXcQ") with pytest.raises(Exception, match="Failed to fetch transcript for video 'dQw4w9WgXcQ': "): youtube_transcript_tool.run(input_data) @patch.object(YouTubeTranscriptApi, 'get_transcript') -@patch('atomic_agents.lib.tools.yt_transcript_scraper.YouTubeTranscriptTool.fetch_video_metadata') -@patch('atomic_agents.lib.tools.yt_transcript_scraper.YouTubeTranscriptTool.extract_video_id') +@patch('atomic_agents.lib.tools.yt_transcript_scraper_tool.YouTubeTranscriptTool.fetch_video_metadata') +@patch('atomic_agents.lib.tools.yt_transcript_scraper_tool.YouTubeTranscriptTool.extract_video_id') def test_run_transcripts_disabled(mock_extract_video_id, mock_fetch_metadata, mock_get_transcript, youtube_transcript_tool): mock_extract_video_id.return_value = "dQw4w9WgXcQ" mock_fetch_metadata.return_value = SAMPLE_VIDEO_INFO mock_get_transcript.side_effect = TranscriptsDisabled("Transcripts are disabled for this video") - input_data = YouTubeTranscriptToolSchema(video_url="https://www.youtube.com/watch?v=dQw4w9WgXcQ") + input_data = YouTubeTranscriptToolInputSchema(video_url="https://www.youtube.com/watch?v=dQw4w9WgXcQ") with pytest.raises(Exception, match="Failed to fetch transcript for video 'dQw4w9WgXcQ': "): youtube_transcript_tool.run(input_data) -@patch('atomic_agents.lib.tools.yt_transcript_scraper.build') +@patch('atomic_agents.lib.tools.yt_transcript_scraper_tool.build') def test_fetch_video_metadata_success(mock_build, youtube_transcript_tool): mock_youtube = Mock() mock_build.return_value = mock_youtube @@ -162,7 +162,7 @@ def test_fetch_video_metadata_success(mock_build, youtube_transcript_tool): mock_build.assert_called_once_with("youtube", "v3", developerKey=youtube_transcript_tool.api_key) mock_youtube.videos().list.assert_called_once_with(part="snippet", id=video_id) -@patch('atomic_agents.lib.tools.yt_transcript_scraper.build') +@patch('atomic_agents.lib.tools.yt_transcript_scraper_tool.build') def test_fetch_video_metadata_no_items(mock_build, youtube_transcript_tool): mock_youtube = Mock() mock_build.return_value = mock_youtube