================================================================================ ATOMIC AGENTS SOURCE CODE ================================================================================ This file contains the complete source code for the Atomic Agents framework. Generated for use with Large Language Models and AI assistants. Project Repository: https://github.com/eigenwise/atomic-agents ### File: atomic-agents/atomic_agents/__init__.py ```python """ Atomic Agents - A modular framework for building AI agents. """ # Core exports - base classes only from .agents.atomic_agent import AtomicAgent, AgentConfig, BasicChatInputSchema, BasicChatOutputSchema from .base import BaseIOSchema, BaseTool, BaseToolConfig, VideoURL # Version info - read from pyproject.toml via package metadata from importlib.metadata import version as _version __version__ = _version("atomic-agents") __all__ = [ "AtomicAgent", "AgentConfig", "BasicChatInputSchema", "BasicChatOutputSchema", "BaseIOSchema", "BaseTool", "BaseToolConfig", "VideoURL", ] ``` ### File: atomic-agents/atomic_agents/agents/__init__.py ```python """Agent implementations and configurations.""" from .atomic_agent import ( AtomicAgent, AgentConfig, BasicChatInputSchema, BasicChatOutputSchema, ) __all__ = [ "AtomicAgent", "AgentConfig", "BasicChatInputSchema", "BasicChatOutputSchema", ] ``` ### File: atomic-agents/atomic_agents/agents/atomic_agent.py ```python import instructor from instructor import Mode from instructor.processing.multimodal import Image, Audio, PDF from pydantic import BaseModel, Field from typing import Optional, Type, Generator, AsyncGenerator, get_args, get_origin, Dict, List, Callable, Union, Any import logging from atomic_agents.context.chat_history import ChatHistory from atomic_agents.context.base_chat_history import BaseChatHistory from atomic_agents.context.system_prompt_generator import ( BaseDynamicContextProvider, SystemPromptGenerator, BaseSystemPromptGenerator, ) from atomic_agents.base.base_io_schema import BaseIOSchema from atomic_agents.utils.token_counter import get_token_counter, TokenCountResult import json from instructor.dsl.partial import PartialBase from jiter import from_json # Instructor modes whose prepared request carries the output schema as a # tool/function definition instead of text inside the messages. Token accounting # must agree with the actual request, so these modes are counted with a tools # definition while all other modes append the schema to the system message. # Keep aligned with Instructor's per-mode request preparation; the parametrized # test in tests/agents/test_agent_mode_consistency.py prepares a request per # mode and asserts this set matches what is actually sent. # Mode.COHERE_TOOLS is deliberately absent: despite its name, its prepared # request embeds the schema in an instruction message and sends no tool. _TOOL_MODES = frozenset( { Mode.FUNCTIONS, Mode.PARALLEL_TOOLS, Mode.TOOLS, Mode.TOOLS_STRICT, Mode.RESPONSES_TOOLS, Mode.RESPONSES_TOOLS_WITH_INBUILT_TOOLS, Mode.ANTHROPIC_TOOLS, Mode.ANTHROPIC_REASONING_TOOLS, Mode.ANTHROPIC_PARALLEL_TOOLS, Mode.MISTRAL_TOOLS, Mode.VERTEXAI_TOOLS, Mode.VERTEXAI_PARALLEL_TOOLS, Mode.GEMINI_TOOLS, Mode.GENAI_TOOLS, Mode.XAI_TOOLS, Mode.CEREBRAS_TOOLS, Mode.FIREWORKS_TOOLS, Mode.WRITER_TOOLS, Mode.BEDROCK_TOOLS, } ) # Events emitted by AtomicAgent itself, not accepted by Instructor's hook registry. _LOCAL_HOOK_EVENTS = frozenset({"token:counted"}) def _is_instructor_hook(event: str) -> bool: return event not in _LOCAL_HOOK_EVENTS # Instructor multimodal objects that can be lowered to a chat-format content part. _InstructorMedia = Union[Image, Audio, PDF] def model_from_chunks_patched(cls, json_chunks, **kwargs): kwargs.setdefault("strict", None) potential_object = "" partial_model = cls.get_partial_model() for chunk in json_chunks: potential_object += chunk obj = from_json((potential_object or "{}").encode(), partial_mode="trailing-strings") obj = partial_model.model_validate(obj, **kwargs) yield obj async def model_from_chunks_async_patched(cls, json_chunks, **kwargs): kwargs.setdefault("strict", None) potential_object = "" partial_model = cls.get_partial_model() async for chunk in json_chunks: potential_object += chunk obj = from_json((potential_object or "{}").encode(), partial_mode="trailing-strings") obj = partial_model.model_validate(obj, **kwargs) yield obj PartialBase.model_from_chunks = classmethod(model_from_chunks_patched) PartialBase.model_from_chunks_async = classmethod(model_from_chunks_async_patched) class BasicChatInputSchema(BaseIOSchema): """This schema represents the input from the user to the AI agent.""" chat_message: str = Field( ..., description="The chat message sent by the user to the assistant.", ) class BasicChatOutputSchema(BaseIOSchema): """This schema represents the response generated by the chat agent.""" chat_message: str = Field( ..., description=( "The chat message exchanged between the user and the chat agent. " "This contains the markdown-enabled response generated by the chat agent." ), ) class AgentConfig(BaseModel): client: instructor.Instructor = Field(..., description="Client for interacting with the language model.") name: Optional[str] = Field( default=None, description="Stable name for this agent in logs and traces, for example the role it plays. Defaults to the class name.", ) model: str = Field(default="gpt-5.6-terra", description="The model to use for generating responses.") history: Optional[BaseChatHistory] = Field(default=None, description="History component for storing chat history.") system_prompt_generator: Optional[BaseSystemPromptGenerator] = Field( default=None, description=( "Component for generating system prompts. " "Defaults to SystemPromptGenerator if no subclass of BaseSystemPromptGenerator is passed." ), ) system_role: Optional[str] = Field( default="system", description="The role of the system in the conversation. None means no system prompt." ) assistant_role: str = Field( default="assistant", description="The role of the assistant in the conversation. Use 'model' for Gemini, 'assistant' for OpenAI/Anthropic.", ) tool_result_role: Optional[str] = Field( default=None, description=( "The role to use for mid-conversation tool results and context injections. " "Defaults to 'user' when assistant_role is 'model' (Gemini), otherwise 'system'. " "Set explicitly to override auto-detection." ), ) model_config = {"arbitrary_types_allowed": True} mode: Optional[Mode] = Field( default=None, description=( "The Instructor mode used for structured outputs (TOOLS, JSON, etc.). " "Drives token accounting only: modes that transmit the schema as a tool " "definition count it as a tool, others append it to the system message. " "When None, the client's mode is used; an explicit mode must agree with " "the client's mode on how the schema is transmitted, or a warning is logged." ), ) model_api_parameters: Optional[dict] = Field(default=None, description="Additional parameters passed to the API provider.") max_context_tokens: Optional[int] = Field( default=None, description=( "Maximum tokens for the full context (system prompt + history + tools). " "When exceeded, oldest conversation turns are automatically trimmed to stay within limit. " "Uses LiteLLM's provider-agnostic token counter — works with any supported model." ), ) class AtomicAgent[InputSchema: BaseIOSchema, OutputSchema: BaseIOSchema]: """ Base class for chat agents with full Instructor hook system integration. This class provides the core functionality for handling chat interactions, including managing history, generating system prompts, and obtaining responses from a language model. It includes comprehensive hook system support for monitoring and error handling. Type Parameters: InputSchema: Schema for the user input, must be a subclass of BaseIOSchema. OutputSchema: Schema for the agent's output, must be a subclass of BaseIOSchema. Attributes: client: Client for interacting with the language model. model (str): The model to use for generating responses. history (BaseChatHistory): History component for storing chat history. system_prompt_generator (BaseSystemPromptGenerator): Component for generating system prompts. system_role (Optional[str]): The role of the system in the conversation. None means no system prompt. assistant_role (str): The role of the assistant in the conversation. Use 'model' for Gemini, 'assistant' for OpenAI/Anthropic. initial_history (ChatHistory): Initial state of the history. current_user_input (Optional[InputSchema]): The current user input being processed. model_api_parameters (dict): Additional parameters passed to the API provider. - Use this for parameters like 'temperature', 'max_tokens', etc. max_context_tokens (Optional[int]): Maximum tokens for the full context. When exceeded, oldest conversation turns are automatically trimmed. Uses LiteLLM's token counter. Hook System: The AtomicAgent integrates with Instructor's hook system to provide comprehensive monitoring and error handling capabilities. Supported events include: - 'parse:error': Triggered when Pydantic validation fails - 'completion:kwargs': Triggered before completion request - 'completion:response': Triggered after completion response - 'completion:error': Triggered on completion errors - 'completion:last_attempt': Triggered on final retry attempt Hook Methods: - register_hook(event, handler): Register a hook handler for an event - unregister_hook(event, handler): Remove a hook handler - clear_hooks(event=None): Clear hooks for specific event or all events - enable_hooks()/disable_hooks(): Control hook processing - hooks_enabled: Property to check if hooks are enabled Example: ```python # Basic usage agent = AtomicAgent[InputSchema, OutputSchema](config) # Register parse error hook for intelligent retry handling def handle_parse_error(error): print(f"Validation failed: {error}") # Implement custom retry logic, logging, etc. agent.register_hook("parse:error", handle_parse_error) # Now parse:error hooks will fire on validation failures response = agent.run(user_input) ``` """ def __init__(self, config: AgentConfig): """ Initializes the AtomicAgent. Args: config (AgentConfig): Configuration for the chat agent. """ self.name = config.name self.client = config.client self.model = config.model self.history = config.history if config.history is not None else ChatHistory() self.system_prompt_generator = config.system_prompt_generator or SystemPromptGenerator() self.system_role = config.system_role self.assistant_role = config.assistant_role self.tool_result_role = ( config.tool_result_role if config.tool_result_role is not None else self._default_tool_result_role(config.assistant_role) ) self.initial_history = self.history.copy() self.current_user_input = None self.mode = self._resolve_mode(config) self.model_api_parameters = config.model_api_parameters or {} self.max_context_tokens = config.max_context_tokens # Hook management attributes self._hook_handlers: Dict[str, List[Callable]] = {} self._hooks_enabled: bool = True @staticmethod def _default_tool_result_role(assistant_role: str) -> str: """ Picks the role for tool results when AgentConfig does not override it. Gemini drops mid-conversation "system" messages, so Gemini backends (identified by assistant_role="model") receive tool results as user messages; every other backend uses "system". """ return "user" if assistant_role == "model" else "system" def reset_history(self): """ Resets the history to its initial state. """ self.history = self.initial_history.copy() @staticmethod def _client_mode(client: instructor.Instructor) -> Optional[Mode]: """ The mode of an Instructor client, when it exposes one. `AgentConfig` only accepts Instructor clients, but a client double or a custom client can still lack a comparable `mode` attribute; None is returned for them. """ mode = getattr(client, "mode", None) return mode if isinstance(mode, Mode) else None @staticmethod def _modes_disagree(config_mode: Mode, client_mode: Optional[Mode]) -> bool: """ Whether the two modes transmit the output schema differently. Only the transmission matters for token accounting: two modes that both send the schema as a tool definition are accounted for identically even when they target different providers. """ if client_mode is None: return False return (config_mode in _TOOL_MODES) != (client_mode in _TOOL_MODES) @staticmethod def _warn_mode_mismatch(config_mode: Mode, client_mode: Mode) -> None: """Warns that token accounting will silently disagree with the API calls.""" logging.getLogger(__name__).warning( "AgentConfig.mode (%s) and the Instructor client's mode (%s) disagree on how the " "output schema is transmitted, so token accounting will silently disagree with " "the actual API calls. Create the Instructor client with %s as well, or omit " "AgentConfig.mode to follow the client automatically.", config_mode.name, client_mode.name, config_mode.name, ) @staticmethod def _resolve_mode(config: AgentConfig) -> Mode: """ Resolves the effective Instructor mode used for token accounting. The API call format is decided by the mode the Instructor client was created with, so `AgentConfig.mode` follows the client by default. An explicitly configured mode that disagrees with the client on how the schema is transmitted warns, because token accounting then silently disagrees with the actual API calls. Args: config (AgentConfig): Configuration for the chat agent. Returns: Mode: The mode to use for token accounting. """ client_mode = AtomicAgent._client_mode(config.client) if config.mode is None: return client_mode if client_mode is not None else Mode.TOOLS if AtomicAgent._modes_disagree(config.mode, client_mode): AtomicAgent._warn_mode_mismatch(config.mode, client_mode) return config.mode def add_tool_result(self, content: BaseIOSchema) -> None: """ Adds a tool result or context injection to the chat history using the backend-appropriate role. This method should be used instead of ``history.add_message("system", ...)`` when injecting tool execution results, resource contents, or other mid-conversation context into the agent's history. It automatically uses the correct role for the configured backend (e.g. ``"user"`` for Gemini, ``"system"`` for OpenAI/Anthropic). Args: content (BaseIOSchema): The tool result or context to inject. """ self.history.add_message(self.tool_result_role, content) def _resolve_schema_types(self) -> tuple[Type[BaseIOSchema], Type[BaseIOSchema]]: """Resolve AtomicAgent's schemas through the instance's specialized bases.""" def resolve(base: Any, bindings: Dict[Any, Any]) -> Optional[tuple[Any, Any]]: origin = get_origin(base) or base arguments = tuple(bindings.get(argument, argument) for argument in get_args(base)) if origin is AtomicAgent: return (arguments[0], arguments[1]) if len(arguments) == 2 else None bindings = {**bindings, **dict(zip(getattr(origin, "__parameters__", ()), arguments))} # Read only bases declared on this class; inherited __orig_bases__ # can skip a specialization introduced by an intermediate subclass. for parent in origin.__dict__.get("__orig_bases__", origin.__bases__): schemas = resolve(parent, bindings) if schemas is not None: return schemas return None return resolve(getattr(self, "__orig_class__", type(self)), {}) or (BasicChatInputSchema, BasicChatOutputSchema) @property def input_schema(self) -> Type[BaseIOSchema]: """Return the concrete input schema, or the default for an untyped agent.""" return self._resolve_schema_types()[0] @property def output_schema(self) -> Type[BaseIOSchema]: """Return the concrete output schema, or the default for an untyped agent.""" return self._resolve_schema_types()[1] def _build_system_messages(self) -> List[Dict]: """ Builds the system message(s) based on the configured system role. Returns: List[Dict]: A list containing the system message, or an empty list if system_role is None. """ if self.system_role is None: return [] return [ { "role": self.system_role, "content": self.system_prompt_generator.generate_prompt(), } ] def _trim_context(self) -> None: """ Trim oldest conversation turns to stay within max_context_tokens limit. Called before building the messages list. Uses the full context token count (system prompt + history + tools) via get_context_token_count(). Removes oldest turns one at a time until the context fits within the limit, preserving the current turn so the request never loses its latest input. Turn-preserving: always removes complete turns, never individual messages. Raises: ValueError: If a single turn itself exceeds max_context_tokens. """ if self.max_context_tokens is None: return result = self.get_context_token_count() total_tokens = result.total if total_tokens <= self.max_context_tokens: return logger = logging.getLogger(__name__) # Collect unique turn_ids in order (oldest first) turn_ids_ordered = [] seen = set() for msg in self.history.history: if msg.turn_id and msg.turn_id not in seen: turn_ids_ordered.append(msg.turn_id) seen.add(msg.turn_id) # Remove oldest turns until within limit for turn_id in turn_ids_ordered: if total_tokens <= self.max_context_tokens: break if turn_id == self.history.current_turn_id: continue self.history.delete_turn_id(turn_id) new_result = self.get_context_token_count() removed = total_tokens - new_result.total total_tokens = new_result.total logger.warning( "Context exceeded max_context_tokens (%d). " "Trimmed turn '%s' (%d tokens). New total: %d.", self.max_context_tokens, turn_id, removed, total_tokens, ) if total_tokens > self.max_context_tokens: raise ValueError( f"max_context_tokens ({self.max_context_tokens}) is smaller than the " f"minimum required for a single turn ({total_tokens} tokens). " "Increase max_context_tokens or reduce system prompt size." ) def _prepare_messages(self): self.messages = self._build_system_messages() history = self.history.get_history() # Remap "system" role messages in history when the backend doesn't support # mid-conversation system messages (e.g. Gemini). The initial system prompt # is built separately via _build_system_messages() and is unaffected. if self.tool_result_role != "system": logger = logging.getLogger(__name__) for msg in history: if msg["role"] == "system": logger.debug( "Remapping mid-conversation 'system' message to '%s' " "(backend does not support mid-conversation system messages).", self.tool_result_role, ) msg["role"] = self.tool_result_role self.messages += history def _get_completion_kwargs(self) -> Dict[str, Any]: """ Build kwargs for Instructor completion calls. Instructor defaults `strict=True`, which forces enum fields to receive enum instances instead of allowing Pydantic's normal coercion from strings. We default to `strict=None` here so the output schema's own Pydantic behavior applies unless callers explicitly override it via `model_api_parameters`. """ completion_kwargs = dict(self.model_api_parameters) completion_kwargs.setdefault("strict", None) return completion_kwargs def _build_tools_definition(self) -> Optional[List[Dict[str, Any]]]: """ Build the tools definition that Instructor sends for TOOLS mode. This uses Instructor's actual schema generation to create the exact tools parameter that would be sent to the LLM for TOOLS mode. For JSON modes, returns None as the schema is embedded in messages. Returns: Optional[List[Dict[str, Any]]]: Tools definition for TOOLS mode, or None for JSON modes. """ from instructor.processing.schema import generate_openai_schema # Only return tools for modes that transmit the schema as a tool definition if self.mode in _TOOL_MODES: return [ { "type": "function", "function": generate_openai_schema(self.output_schema), } ] return None def _build_schema_for_json_mode(self) -> str: """ Build the schema context for JSON modes (appended to system message). This matches exactly how Instructor formats the schema for JSON/MD_JSON modes. Returns: str: JSON schema string formatted as Instructor does. """ from textwrap import dedent schema = self.output_schema.model_json_schema() return dedent( f""" As a genius expert, your task is to understand the content and provide the parsed objects in json that match the following json_schema: {json.dumps(schema, indent=2, ensure_ascii=False)} Make sure to return an instance of the JSON, not the schema itself """ ).strip() def _serialize_history_for_token_count(self) -> List[Dict[str, Any]]: """ Serialize conversation history for token counting, handling multimodal content. This method converts instructor multimodal objects (Image, Audio, PDF) to the OpenAI format that LiteLLM's token counter expects. Text content is also converted to the proper multimodal text format when mixed with media. Content-part dicts (e.g. video) become text placeholders because LiteLLM cannot count them. Returns: List[Dict[str, Any]]: History messages in LiteLLM-compatible format. """ history = self.history.get_history() serialized = [] for message in history: content = message.get("content") if isinstance(content, list): # Multimodal content - convert to the chat format the counter accepts serialized.append( { "role": message["role"], "content": [self._token_count_content_part(item) for item in content], } ) else: # Simple text content - keep as is serialized.append(message) return serialized @staticmethod def _token_count_content_part(item: object) -> Dict[str, Any]: """ Converts one history content part to the chat format the token counter accepts. Args: item: A text part, an instructor multimodal object, or a content-part dict. Anything else is counted as its string form. Returns: Dict[str, Any]: The content part in LiteLLM-compatible format. """ if isinstance(item, str): # Text content - wrap in OpenAI text format return {"type": "text", "text": item} if isinstance(item, (Image, Audio, PDF)): return AtomicAgent._serialize_media_for_token_count(item) if isinstance(item, dict): # get_history() emits pre-lowered content-part dicts (e.g. video). # LiteLLM's token counter only accepts text and image_url parts, # so estimate them with a placeholder. return {"type": "text", "text": f"[{item.get('type', 'unknown')} content]"} # Unknown type - convert to string return {"type": "text", "text": str(item)} @staticmethod def _serialize_media_for_token_count(item: _InstructorMedia) -> Dict[str, Any]: """ Serializes an instructor media object to the chat format the token counter accepts. LiteLLM's counter only accepts chat-format content parts, so media is normalized to that format regardless of the agent's mode (Responses-format parts such as input_image raise). Serialization failures degrade to a text placeholder so counting still works. """ try: return item.to_openai(Mode.JSON) except Exception as e: # Log the error and use placeholder for token estimation media_type = type(item).__name__ logging.getLogger(__name__).warning( f"Failed to serialize {media_type} for token counting: {e}. Using placeholder for estimation." ) return {"type": "text", "text": f"[{media_type.lower()} content]"} def get_context_token_count(self) -> TokenCountResult: """ Get the accurate token count for the current context. This method computes the token count by serializing the context exactly as Instructor does, including: - System prompt - Conversation history (with multimodal content serialized properly) - Tools/schema overhead (using Instructor's actual schema generation) For TOOLS mode: Uses the actual tools parameter that Instructor sends. For JSON modes: Appends the schema to the system message as Instructor does. Works with any model supported by LiteLLM including OpenAI, Anthropic, Google, and 100+ other providers. Returns: TokenCountResult: A named tuple containing: - total: Total tokens in the complete message request; when tools are provided without messages, this is schema-only overhead and excludes request framing - system_prompt: System message tokens, including request framing when a system message is present - history: Incremental tokens added by conversation history, including request framing when no system message is present - tools: Incremental tokens added by tool definitions (TOOLS mode only) - model: The model used for counting - max_tokens: Maximum context window (if known) - utilization: Percentage of context used (if max_tokens known) The breakdown is additive: ``system_prompt + history + tools == total``. When tools are provided without system or history messages, ``total`` and ``tools`` report schema-only overhead and exclude request framing. Example: ```python agent = AtomicAgent[InputSchema, OutputSchema](config) # Get accurate token count at any time result = agent.get_context_token_count() print(f"Total: {result.total} tokens") print(f"System: {result.system_prompt} tokens") print(f"History: {result.history} tokens") print(f"Tools: {result.tools} tokens") if result.utilization: print(f"Context usage: {result.utilization:.1%}") ``` Note: The 'token:counted' hook event is dispatched, allowing for monitoring and logging of token usage. """ counter = get_token_counter() # Build system messages system_messages = self._build_system_messages() # Handle schema serialization based on mode tools = self._build_tools_definition() if tools is None: # JSON mode - append schema to system message like Instructor does schema_context = self._build_schema_for_json_mode() if system_messages: system_messages = [ { "role": system_messages[0]["role"], "content": system_messages[0]["content"] + "\n\n" + schema_context, } ] else: system_messages = [{"role": "system", "content": schema_context}] result = counter.count_context( model=self.model, system_messages=system_messages, history_messages=self._serialize_history_for_token_count(), tools=tools, ) # Dispatch hook for monitoring self._dispatch_hook("token:counted", result) return result def run(self, user_input: Optional[InputSchema] = None) -> OutputSchema: """ Runs the chat agent with the given user input synchronously. Args: user_input (Optional[InputSchema]): The input from the user. If not provided, skips adding to history. Returns: OutputSchema: The response from the chat agent. """ assert not isinstance( self.client, instructor.AsyncInstructor ), "The run method is not supported for async clients. Use run_async instead." if user_input: self.history.initialize_turn() self.current_user_input = user_input self.history.add_message("user", user_input) self._trim_context() self._prepare_messages() response = self.client.chat.completions.create( messages=self.messages, model=self.model, response_model=self.output_schema, **self._get_completion_kwargs(), ) self.history.add_message(self.assistant_role, response) self._prepare_messages() return response def run_stream(self, user_input: Optional[InputSchema] = None) -> Generator[OutputSchema, None, OutputSchema]: """ Runs the chat agent with the given user input, supporting streaming output. Args: user_input (Optional[InputSchema]): The input from the user. If not provided, skips adding to history. Yields: OutputSchema: Partial responses from the chat agent. Returns: OutputSchema: The final response from the chat agent. """ assert not isinstance( self.client, instructor.AsyncInstructor ), "The run_stream method is not supported for async clients. Use run_async instead." if user_input: self.history.initialize_turn() self.current_user_input = user_input self.history.add_message("user", user_input) self._trim_context() self._prepare_messages() response_stream = self.client.chat.completions.create_partial( model=self.model, messages=self.messages, response_model=self.output_schema, **self._get_completion_kwargs(), stream=True, ) last_response = None for partial_response in response_stream: last_response = partial_response yield partial_response if last_response: # model_dump() keys the data by field name, so accept field names # for output schemas whose fields have an alias. full_response_content = self.output_schema.model_validate(last_response.model_dump(), by_name=True) self.history.add_message(self.assistant_role, full_response_content) self._prepare_messages() return full_response_content async def run_async(self, user_input: Optional[InputSchema] = None) -> OutputSchema: """ Runs the chat agent asynchronously with the given user input. Args: user_input (Optional[InputSchema]): The input from the user. If not provided, skips adding to history. Returns: OutputSchema: The response from the chat agent. Raises: NotAsyncIterableError: If used as an async generator (in an async for loop). Use run_async_stream() method instead for streaming responses. """ assert isinstance(self.client, instructor.AsyncInstructor), "The run_async method is for async clients." if user_input: self.history.initialize_turn() self.current_user_input = user_input self.history.add_message("user", user_input) self._trim_context() self._prepare_messages() response = await self.client.chat.completions.create( model=self.model, messages=self.messages, response_model=self.output_schema, **self._get_completion_kwargs() ) self.history.add_message(self.assistant_role, response) self._prepare_messages() return response async def run_async_stream(self, user_input: Optional[InputSchema] = None) -> AsyncGenerator[OutputSchema, None]: """ Runs the chat agent asynchronously with the given user input, supporting streaming output. Args: user_input (Optional[InputSchema]): The input from the user. If not provided, skips adding to history. Yields: OutputSchema: Partial responses from the chat agent. """ assert isinstance(self.client, instructor.AsyncInstructor), "The run_async method is for async clients." if user_input: self.history.initialize_turn() self.current_user_input = user_input self.history.add_message("user", user_input) self._trim_context() self._prepare_messages() response_stream = self.client.chat.completions.create_partial( model=self.model, messages=self.messages, response_model=self.output_schema, **self._get_completion_kwargs(), stream=True, ) last_response = None async for partial_response in response_stream: last_response = partial_response yield partial_response if last_response: # See run_stream(): accept field names for aliased output schemas. full_response_content = self.output_schema.model_validate(last_response.model_dump(), by_name=True) self.history.add_message(self.assistant_role, full_response_content) self._prepare_messages() def get_context_provider(self, provider_name: str) -> Type[BaseDynamicContextProvider]: """ Retrieves a context provider by name. Args: provider_name (str): The name of the context provider. Returns: BaseDynamicContextProvider: The context provider if found. Raises: KeyError: If the context provider is not found. """ 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.context_providers[provider_name] def register_context_provider(self, provider_name: str, provider: BaseDynamicContextProvider): """ Registers a new context provider. Args: provider_name (str): The name of the context provider. provider (BaseDynamicContextProvider): The context provider instance. """ self.system_prompt_generator.context_providers[provider_name] = provider def unregister_context_provider(self, provider_name: str): """ Unregisters an existing context provider. Args: provider_name (str): The name of the context provider to remove. """ 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.") # Hook Management Methods def register_hook(self, event: str, handler: Callable) -> None: """ Registers a hook handler for a specific event. Local events use a snapshot of the registered handlers for each dispatch. Args: event (str): The event name (e.g., 'parse:error', 'completion:kwargs', etc.) handler (Callable): The callback function to handle the event """ if event not in self._hook_handlers: self._hook_handlers[event] = [] self._hook_handlers[event].append(handler) # Register with instructor client if it supports hooks if _is_instructor_hook(event) and self._hooks_enabled and hasattr(self.client, "on"): self.client.on(event, handler) def unregister_hook(self, event: str, handler: Callable) -> None: """ Unregisters a hook handler for a specific event. Args: event (str): The event name handler (Callable): The callback function to remove """ if event in self._hook_handlers and handler in self._hook_handlers[event]: self._hook_handlers[event].remove(handler) # Remove from instructor client if it supports hooks if _is_instructor_hook(event) and self._hooks_enabled and hasattr(self.client, "off"): self.client.off(event, handler) def clear_hooks(self, event: Optional[str] = None) -> None: """ Clears hook handlers for a specific event or all events. Args: event (Optional[str]): The event name to clear, or None to clear all """ events = [event] if event else list(self._hook_handlers) for registered_event in events: for handler in list(self._hook_handlers.get(registered_event, [])): self.unregister_hook(registered_event, handler) if not event: self._hook_handlers.clear() def _dispatch_hook(self, event: str, *args, **kwargs) -> None: """ Internal method to dispatch hook events with error isolation. Args: event (str): The event name *args: Arguments to pass to handlers **kwargs: Keyword arguments to pass to handlers """ if not self._hooks_enabled or event not in self._hook_handlers: return for handler in tuple(self._hook_handlers[event]): try: handler(*args, **kwargs) except Exception as e: # Log error but don't interrupt main flow logger = logging.getLogger(__name__) logger.warning(f"Hook handler for '{event}' raised exception: {e}") def enable_hooks(self) -> None: """Enable hook processing.""" if self._hooks_enabled: return self._hooks_enabled = True if hasattr(self.client, "on"): for event, handlers in self._hook_handlers.items(): if not _is_instructor_hook(event): continue for handler in handlers: self.client.on(event, handler) def disable_hooks(self) -> None: """Disable hook processing.""" if not self._hooks_enabled: return self._hooks_enabled = False if hasattr(self.client, "off"): for event, handlers in self._hook_handlers.items(): if not _is_instructor_hook(event): continue for handler in handlers: self.client.off(event, handler) @property def hooks_enabled(self) -> bool: """Check if hooks are enabled.""" return self._hooks_enabled if __name__ == "__main__": from rich.console import Console from rich.panel import Panel from rich.table import Table from rich.syntax import Syntax from rich import box from openai import OpenAI, AsyncOpenAI import instructor import asyncio from rich.live import Live def _create_schema_table(title: str, schema: Type[BaseModel]) -> Table: """Create a table displaying schema information. Args: title (str): Title of the table schema (Type[BaseModel]): Schema to display Returns: Table: Rich table containing schema information """ schema_table = Table(title=title, box=box.ROUNDED) schema_table.add_column("Field", style="cyan") schema_table.add_column("Type", style="magenta") schema_table.add_column("Description", style="green") for field_name, field in schema.model_fields.items(): schema_table.add_row(field_name, str(field.annotation), field.description or "") return schema_table def _create_config_table(agent: AtomicAgent) -> Table: """Create a table displaying agent configuration. Args: agent (AtomicAgent): Agent instance Returns: Table: Rich table containing configuration information """ info_table = Table(title="Agent Configuration", box=box.ROUNDED) info_table.add_column("Property", style="cyan") info_table.add_column("Value", style="yellow") info_table.add_row("Model", agent.model) info_table.add_row("History", str(type(agent.history).__name__)) info_table.add_row("System Prompt Generator", str(type(agent.system_prompt_generator).__name__)) return info_table def display_agent_info(agent: AtomicAgent): """Display information about the agent's configuration and schemas.""" console = Console() console.print( Panel.fit( "[bold blue]Agent Information[/bold blue]", border_style="blue", padding=(1, 1), ) ) # Display input schema input_schema_table = _create_schema_table("Input Schema", agent.input_schema) console.print(input_schema_table) # Display output schema output_schema_table = _create_schema_table("Output Schema", agent.output_schema) console.print(output_schema_table) # Display configuration info_table = _create_config_table(agent) console.print(info_table) # Display system prompt system_prompt = agent.system_prompt_generator.generate_prompt() console.print( Panel( Syntax(system_prompt, "markdown", theme="monokai", line_numbers=True), title="Sample System Prompt", border_style="green", expand=False, ) ) async def chat_loop(streaming: bool = False): """Interactive chat loop with the AI agent. Args: streaming (bool): Whether to use streaming mode for responses """ if streaming: client = instructor.from_openai(AsyncOpenAI()) config = AgentConfig(client=client, model="gpt-5.6-terra") agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) else: client = instructor.from_openai(OpenAI()) config = AgentConfig(client=client, model="gpt-5.6-terra") agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) # Display agent information before starting the chat display_agent_info(agent) console = Console() console.print( Panel.fit( "[bold blue]Interactive Chat Mode[/bold blue]\n" f"[cyan]Streaming: {streaming}[/cyan]\n" "Type 'exit' to quit", border_style="blue", padding=(1, 1), ) ) while True: user_message = console.input("\n[bold green]You:[/bold green] ") if user_message.lower() == "exit": console.print("[yellow]Goodbye![/yellow]") break user_input = agent.input_schema(chat_message=user_message) console.print("[bold blue]Assistant:[/bold blue]") if streaming: with Live(console=console, refresh_per_second=4) as live: # Use run_async_stream instead of run_async for streaming responses async for partial_response in agent.run_async_stream(user_input): response_json = partial_response.model_dump() json_str = json.dumps(response_json, indent=2) live.update(json_str) else: response = agent.run(user_input) response_json = response.model_dump() json_str = json.dumps(response_json, indent=2) console.print(json_str) console = Console() console.print("\n[bold]Starting chat loop...[/bold]") asyncio.run(chat_loop(streaming=True)) ``` ### File: atomic-agents/atomic_agents/base/__init__.py ```python """Base classes for Atomic Agents.""" from .base_io_schema import BaseIOSchema from .base_tool import BaseTool, BaseToolConfig from .base_resource import BaseResource, BaseResourceConfig from .base_prompt import BasePrompt, BasePromptConfig from .multimodal import VideoURL __all__ = [ "BaseIOSchema", "BaseTool", "BaseToolConfig", "BaseResource", "BaseResourceConfig", "BasePrompt", "BasePromptConfig", "VideoURL", ] ``` ### File: atomic-agents/atomic_agents/base/base_io_schema.py ```python import inspect from pydantic import BaseModel from rich.json import JSON class BaseIOSchema(BaseModel): """Base schema for input/output in the Atomic Agents framework.""" def __str__(self): return self.model_dump_json() def __rich__(self): json_str = self.model_dump_json() return JSON(json_str) @classmethod def __pydantic_init_subclass__(cls, **kwargs): super().__pydantic_init_subclass__(**kwargs) cls._validate_description() @classmethod def _validate_description(cls): description = cls.__doc__ if not description or not description.strip(): if not cls._is_generated_by_instructor(): raise ValueError(f"{cls.__name__} must have a non-empty docstring to serve as its description") @classmethod def _is_generated_by_instructor(cls) -> bool: return cls.__module__.startswith("instructor.") or hasattr(cls, "from_streaming_response") @classmethod def model_json_schema(cls, *args, **kwargs): schema = super().model_json_schema(*args, **kwargs) if "description" not in schema and cls.__doc__: schema["description"] = inspect.cleandoc(cls.__doc__) if "title" not in schema: schema["title"] = cls.__name__ return schema ``` ### File: atomic-agents/atomic_agents/base/base_prompt.py ```python from typing import Optional, Type, get_args, get_origin from abc import ABC, abstractmethod from pydantic import BaseModel from atomic_agents.base.base_io_schema import BaseIOSchema class BasePromptConfig(BaseModel): """ Configuration for a prompt. Attributes: title (Optional[str]): Overrides the default title of the prompt. description (Optional[str]): Overrides the default description of the prompt. """ title: Optional[str] = None description: Optional[str] = None class BasePrompt[InputSchema: BaseIOSchema, OutputSchema: BaseIOSchema](ABC): """ Base class for prompts within the Atomic Agents framework. Prompts enable agents to perform specific tasks by providing a standardized interface for input and output. Each prompt is defined with specific input and output schemas that enforce type safety and provide documentation. Type Parameters: InputSchema: Schema defining the input data, must be a subclass of BaseIOSchema. OutputSchema: Schema defining the output data, must be a subclass of BaseIOSchema. Attributes: config (BasePromptConfig): Configuration for the prompt, including optional title and description overrides. input_schema (Type[InputSchema]): Schema class defining the input data (derived from generic type parameter). output_schema (Type[OutputSchema]): Schema class defining the output data (derived from generic type parameter). prompt_name (str): The name of the prompt, derived from the input schema's title or overridden by the config. prompt_description (str): Description of the prompt, derived from the input schema's description or overridden by the config. """ def __init__(self, config: Optional[BasePromptConfig] = None): """ Initializes the BasePrompt with an optional configuration override. Args: config (BasePromptConfig, optional): Configuration for the prompt, including optional title and description overrides. If omitted, a new configuration is created for this instance. """ self.config = config if config is not None else BasePromptConfig() def __init_subclass__(cls, **kwargs): """ Hook called when a class is subclassed. Captures generic type parameters during class creation and stores them as class attributes to work around the unreliable __orig_class__ attribute in modern Python generic syntax. """ super().__init_subclass__(**kwargs) if hasattr(cls, "__orig_bases__"): for base in cls.__orig_bases__: if get_origin(base) is BasePrompt: args = get_args(base) if len(args) == 2: cls._input_schema_cls = args[0] cls._output_schema_cls = args[1] break @property def input_schema(self) -> Type[InputSchema]: """ Returns the input schema class for the prompt. Returns: Type[InputSchema]: The input schema class. """ # Inheritance pattern: MyPrompt(BasePrompt[Schema1, Schema2]) if hasattr(self.__class__, "_input_schema_cls"): return self.__class__._input_schema_cls # Dynamic instantiation: MockPrompt[Schema1, Schema2]() if hasattr(self, "__orig_class__"): TI, _ = get_args(self.__orig_class__) return TI # No type info available: MockPrompt() return BaseIOSchema @property def output_schema(self) -> Type[OutputSchema]: """ Returns the output schema class for the prompt. Returns: Type[OutputSchema]: The output schema class. """ # Inheritance pattern: MyPrompt(BasePrompt[Schema1, Schema2]) if hasattr(self.__class__, "_output_schema_cls"): return self.__class__._output_schema_cls # Dynamic instantiation: MockPrompt[Schema1, Schema2]() if hasattr(self, "__orig_class__"): _, TO = get_args(self.__orig_class__) return TO # No type info available: MockPrompt() return BaseIOSchema @property def prompt_name(self) -> str: """ Returns the name of the prompt. Returns: str: The name of the prompt. """ return self.config.title or self.input_schema.model_json_schema()["title"] @property def prompt_description(self) -> str: """ Returns the description of the prompt. Returns: str: The description of the prompt. """ return self.config.description or self.input_schema.model_json_schema()["description"] @abstractmethod def generate(self, params: InputSchema) -> OutputSchema: """ Executes the prompt with the provided parameters. Args: params (InputSchema): Input parameters adhering to the input schema. Returns: OutputSchema: Output resulting from executing the prompt, adhering to the output schema. Raises: NotImplementedError: If the method is not implemented by a subclass. """ pass ``` ### File: atomic-agents/atomic_agents/base/base_resource.py ```python from typing import Optional, Type, get_args, get_origin from abc import ABC, abstractmethod from pydantic import BaseModel from atomic_agents.base.base_io_schema import BaseIOSchema class BaseResourceConfig(BaseModel): """ Configuration for a resource. Attributes: title (Optional[str]): Overrides the default title of the resource. description (Optional[str]): Overrides the default description of the resource. """ title: Optional[str] = None description: Optional[str] = None class BaseResource[InputSchema: BaseIOSchema, OutputSchema: BaseIOSchema](ABC): """ Base class for resources within the Atomic Agents framework. Resources enable agents to perform specific tasks by providing a standardized interface for input and output. Each resource is defined with specific input and output schemas that enforce type safety and provide documentation. Type Parameters: InputSchema: Schema defining the input data, must be a subclass of BaseIOSchema. OutputSchema: Schema defining the output data, must be a subclass of BaseIOSchema. Attributes: config (BaseResourceConfig): Configuration for the resource, including optional title and description overrides. input_schema (Type[InputSchema]): Schema class defining the input data (derived from generic type parameter). output_schema (Type[OutputSchema]): Schema class defining the output data (derived from generic type parameter). resource_name (str): The name of the resource, derived from the input schema's title or overridden by the config. resource_description (str): Description of the resource, derived from the input schema's description or overridden by the config. """ def __init__(self, config: Optional[BaseResourceConfig] = None): """ Initializes the BaseResource with an optional configuration override. Args: config (BaseResourceConfig, optional): Configuration for the resource, including optional title and description overrides. If omitted, a new configuration is created for this instance. """ self.config = config if config is not None else BaseResourceConfig() def __init_subclass__(cls, **kwargs): """ Hook called when a class is subclassed. Captures generic type parameters during class creation and stores them as class attributes to work around the unreliable __orig_class__ attribute in modern Python generic syntax. """ super().__init_subclass__(**kwargs) if hasattr(cls, "__orig_bases__"): for base in cls.__orig_bases__: if get_origin(base) is BaseResource: args = get_args(base) if len(args) == 2: cls._input_schema_cls = args[0] cls._output_schema_cls = args[1] break @property def input_schema(self) -> Type[InputSchema]: """ Returns the input schema class for the resource. Returns: Type[InputSchema]: The input schema class. """ # Inheritance pattern: MyResource(BaseResource[Schema1, Schema2]) if hasattr(self.__class__, "_input_schema_cls"): return self.__class__._input_schema_cls # Dynamic instantiation: MockResource[Schema1, Schema2]() if hasattr(self, "__orig_class__"): TI, _ = get_args(self.__orig_class__) return TI # No type info available: MockResource() return BaseIOSchema @property def output_schema(self) -> Type[OutputSchema]: """ Returns the output schema class for the resource. Returns: Type[OutputSchema]: The output schema class. """ # Inheritance pattern: MyResource(BaseResource[Schema1, Schema2]) if hasattr(self.__class__, "_output_schema_cls"): return self.__class__._output_schema_cls # Dynamic instantiation: MockResource[Schema1, Schema2]() if hasattr(self, "__orig_class__"): _, TO = get_args(self.__orig_class__) return TO # No type info available: MockResource() return BaseIOSchema @property def resource_name(self) -> str: """ Returns the name of the resource. Returns: str: The name of the resource. """ return self.config.title or self.input_schema.model_json_schema()["title"] @property def resource_description(self) -> str: """ Returns the description of the resource. Returns: str: The description of the resource. """ return self.config.description or self.input_schema.model_json_schema()["description"] @abstractmethod def read(self, params: InputSchema) -> OutputSchema: """ Executes the resource with the provided parameters. Args: params (InputSchema): Input parameters adhering to the input schema. Returns: OutputSchema: Output resulting from executing the resource, adhering to the output schema. Raises: NotImplementedError: If the method is not implemented by a subclass. """ pass ``` ### File: atomic-agents/atomic_agents/base/base_tool.py ```python from typing import Optional, Type, get_args, get_origin from abc import ABC, abstractmethod from pydantic import BaseModel from atomic_agents.base.base_io_schema import BaseIOSchema class BaseToolConfig(BaseModel): """ Configuration for a tool. Attributes: title (Optional[str]): Overrides the default title of the tool. description (Optional[str]): Overrides the default description of the tool. """ title: Optional[str] = None description: Optional[str] = None class BaseTool[InputSchema: BaseIOSchema, OutputSchema: BaseIOSchema](ABC): """ Base class for tools within the Atomic Agents framework. Tools enable agents to perform specific tasks by providing a standardized interface for input and output. Each tool is defined with specific input and output schemas that enforce type safety and provide documentation. Type Parameters: InputSchema: Schema defining the input data, must be a subclass of BaseIOSchema. OutputSchema: Schema defining the output data, must be a subclass of BaseIOSchema. Attributes: config (BaseToolConfig): Configuration for the tool, including optional title and description overrides. input_schema (Type[InputSchema]): Schema class defining the input data (derived from generic type parameter). output_schema (Type[OutputSchema]): Schema class defining the output data (derived from generic type parameter). tool_name (str): The name of the tool, derived from the input schema's title or overridden by the config. tool_description (str): Description of the tool, derived from the input schema's description or overridden by the config. """ def __init__(self, config: Optional[BaseToolConfig] = None): """ Initializes the BaseTool with an optional configuration override. Args: config (BaseToolConfig, optional): Configuration for the tool, including optional title and description overrides. """ self.config = config if config is not None else BaseToolConfig() def __init_subclass__(cls, **kwargs): """ Hook called when a class is subclassed. Captures generic type parameters during class creation and stores them as class attributes to work around the unreliable __orig_class__ attribute in modern Python generic syntax. """ super().__init_subclass__(**kwargs) if hasattr(cls, "__orig_bases__"): for base in cls.__orig_bases__: if get_origin(base) is BaseTool: args = get_args(base) if len(args) == 2: cls._input_schema_cls = args[0] cls._output_schema_cls = args[1] break @property def input_schema(self) -> Type[InputSchema]: """ Returns the input schema class for the tool. Returns: Type[InputSchema]: The input schema class. """ # Inheritance pattern: MyTool(BaseTool[Schema1, Schema2]) if hasattr(self.__class__, "_input_schema_cls"): return self.__class__._input_schema_cls # Dynamic instantiation: MockTool[Schema1, Schema2]() if hasattr(self, "__orig_class__"): TI, _ = get_args(self.__orig_class__) return TI # No type info available: MockTool() return BaseIOSchema @property def output_schema(self) -> Type[OutputSchema]: """ Returns the output schema class for the tool. Returns: Type[OutputSchema]: The output schema class. """ # Inheritance pattern: MyTool(BaseTool[Schema1, Schema2]) if hasattr(self.__class__, "_output_schema_cls"): return self.__class__._output_schema_cls # Dynamic instantiation: MockTool[Schema1, Schema2]() if hasattr(self, "__orig_class__"): _, TO = get_args(self.__orig_class__) return TO # No type info available: MockTool() return BaseIOSchema @property def tool_name(self) -> str: """ Returns the name of the tool. Returns: str: The name of the tool. """ return self.config.title or self.input_schema.model_json_schema()["title"] @property def tool_description(self) -> str: """ Returns the description of the tool. Returns: str: The description of the tool. """ return self.config.description or self.input_schema.model_json_schema()["description"] @abstractmethod def run(self, params: InputSchema) -> OutputSchema: """ Executes the tool with the provided parameters. Args: params (InputSchema): Input parameters adhering to the input schema. Returns: OutputSchema: Output resulting from executing the tool, adhering to the output schema. Raises: NotImplementedError: If the method is not implemented by a subclass. """ pass ``` ### File: atomic-agents/atomic_agents/base/multimodal.py ```python """Multimodal content types that Instructor does not provide.""" from typing import Any, Dict, Optional from pydantic import BaseModel, Field class VideoURL(BaseModel): """ Video reference sent to the LLM as an OpenAI-compatible ``video_url`` content part. Instructor ships Image, Audio, and PDF types but no video type (see https://github.com/567-labs/instructor/discussions/2520), so this class fills the gap for providers that accept ``video_url`` content parts, such as MiniMax and Qwen-VL. Attributes: url (str): HTTP(S) or data: URL of the video. fps (Optional[float]): Frame sampling rate, for providers that accept it. detail (Optional[str]): Detail level, for providers that accept it. """ url: str = Field(..., description="HTTP(S) or data: URL of the video.") fps: Optional[float] = Field(default=None, description="Frame sampling rate, for providers that accept it.") detail: Optional[str] = Field(default=None, description="Detail level, for providers that accept it.") def to_openai(self) -> Dict[str, Any]: """ Build the OpenAI-compatible content part for this video. Returns: Dict[str, Any]: A ``{"type": "video_url", "video_url": {...}}`` content part, omitting optional parameters that were not set. """ return {"type": "video_url", "video_url": self.model_dump(exclude_none=True)} ``` ### File: atomic-agents/atomic_agents/connectors/__init__.py ```python # Only expose the subpackages; no direct re‑exports. from . import mcp # ensure pkg_resources-style discovery __all__ = ["mcp"] ``` ### File: atomic-agents/atomic_agents/connectors/mcp/__init__.py ```python from .mcp_factory import ( MCPFactory, MCPToolOutputSchema, fetch_mcp_tools, fetch_mcp_tools_async, fetch_mcp_resources, fetch_mcp_resources_async, fetch_mcp_prompts, fetch_mcp_prompts_async, create_mcp_orchestrator_schema, fetch_mcp_attributes_with_schema, ) from .schema_transformer import SchemaTransformer from .mcp_definition_service import ( MCPTransportType, MCPToolDefinition, MCPResourceDefinition, MCPPromptDefinition, MCPDefinitionService, ) __all__ = [ "MCPFactory", "MCPToolOutputSchema", "fetch_mcp_tools", "fetch_mcp_tools_async", "fetch_mcp_resources", "fetch_mcp_resources_async", "fetch_mcp_prompts", "fetch_mcp_prompts_async", "create_mcp_orchestrator_schema", "fetch_mcp_attributes_with_schema", "SchemaTransformer", "MCPTransportType", "MCPToolDefinition", "MCPResourceDefinition", "MCPPromptDefinition", "MCPDefinitionService", ] ``` ### File: atomic-agents/atomic_agents/connectors/mcp/mcp_definition_service.py ```python """Module for fetching tool definitions from MCP endpoints.""" import logging import re import shlex from contextlib import AsyncExitStack from typing import List, NamedTuple, Optional, Dict, Any, Callable, Awaitable, TypeVar, Union from enum import Enum from mcp import ClientSession, StdioServerParameters from mcp.client.sse import sse_client from mcp.client.stdio import stdio_client from mcp.client.streamable_http import streamablehttp_client from mcp.shared.exceptions import McpError import mcp.types as types from urllib.parse import urlsplit, urlunsplit logger = logging.getLogger(__name__) PageType = TypeVar( "PageType", types.ListToolsResult, types.ListResourcesResult, types.ListResourceTemplatesResult, types.ListPromptsResult, ) ItemType = TypeVar("ItemType", types.Tool, types.Resource, types.ResourceTemplate, types.Prompt) async def _fetch_all_pages( first_page: PageType, list_next: Callable[[str], Awaitable[PageType]], get_items: Callable[[PageType], List[ItemType]], ) -> List[ItemType]: """Collect an MCP list response, following its opaque continuation cursors.""" response = first_page items = list(get_items(response) or []) cursor = getattr(response, "nextCursor", None) seen_cursors: set[str] = set() while isinstance(cursor, str): if cursor in seen_cursors: raise RuntimeError("MCP pagination returned a repeated cursor") seen_cursors.add(cursor) response = await list_next(cursor) items.extend(get_items(response) or []) cursor = getattr(response, "nextCursor", None) return items def _transport_endpoint(endpoint: Optional[str], suffix: str) -> str: """Resolve a base URL while preserving an explicit transport endpoint.""" if endpoint is None: raise ValueError("Endpoint is required") parts = urlsplit(endpoint) if parts.path.rstrip("/").endswith(suffix.rstrip("/")): return endpoint path = parts.path.rstrip("/") + suffix return urlunsplit((parts.scheme, parts.netloc, path, parts.query, parts.fragment)) class MCPTransportType(Enum): """Enum for MCP transport types.""" SSE = "sse" HTTP_STREAM = "http_stream" STDIO = "stdio" class MCPAttributeType: """MCP attribute types.""" TOOL = "tool" RESOURCE = "resource" PROMPT = "prompt" class MCPToolDefinition(NamedTuple): """Definition of an MCP tool.""" name: str description: Optional[str] input_schema: Dict[str, Any] output_schema: Optional[Dict[str, Any]] = None class MCPResourceDefinition(NamedTuple): """Definition of an MCP resource.""" name: str description: Optional[str] uri: str input_schema: Dict[str, Any] mime_type: Optional[str] = None class MCPPromptDefinition(NamedTuple): """Definition of an MCP prompt/template.""" name: str description: Optional[str] input_schema: Dict[str, Any] # required: List[str] # A list of required argument names class MCPDefinitionService: """Service for fetching tool definitions from MCP endpoints.""" def __init__( self, endpoint: Optional[str] = None, transport_type: MCPTransportType = MCPTransportType.HTTP_STREAM, working_directory: Optional[str] = None, ): """ Initialize the service. Args: endpoint: MCP server base URL or explicit /mcp or /sse endpoint (for SSE/HTTP stream), or command string (for STDIO). Query parameters are preserved. transport_type: Type of transport to use (SSE, HTTP_STREAM, or STDIO) working_directory: Optional working directory to use when running STDIO commands """ self.endpoint = endpoint self.transport_type = transport_type self.working_directory = working_directory async def fetch_tool_definitions(self) -> List[MCPToolDefinition]: """ Fetch tool definitions from the configured endpoint. Returns: List of tool definitions Raises: ConnectionError: If connection to the MCP server fails ValueError: If the STDIO command string is empty RuntimeError: For other unexpected errors """ if not self.endpoint: raise ValueError("Endpoint is required") definitions = [] stack = AsyncExitStack() try: if self.transport_type == MCPTransportType.STDIO: # STDIO transport command_parts = shlex.split(self.endpoint) if not command_parts: raise ValueError("STDIO command string cannot be empty.") command = command_parts[0] args = command_parts[1:] logger.info(f"Attempting STDIO connection with command='{command}', args={args}") server_params = StdioServerParameters(command=command, args=args, env=None, cwd=self.working_directory) stdio_transport = await stack.enter_async_context(stdio_client(server_params)) read_stream, write_stream = stdio_transport elif self.transport_type == MCPTransportType.HTTP_STREAM: # HTTP Stream transport - use trailing slash to avoid redirect # See: https://github.com/modelcontextprotocol/python-sdk/issues/732 transport_endpoint = _transport_endpoint(self.endpoint, "/mcp/") logger.info(f"Attempting HTTP Stream connection to {transport_endpoint}") transport = await stack.enter_async_context(streamablehttp_client(transport_endpoint)) read_stream, write_stream, _ = transport elif self.transport_type == MCPTransportType.SSE: # SSE transport (deprecated) transport_endpoint = _transport_endpoint(self.endpoint, "/sse") logger.info(f"Attempting SSE connection to {transport_endpoint}") transport = await stack.enter_async_context(sse_client(transport_endpoint)) read_stream, write_stream = transport else: available_types = [t.value for t in MCPTransportType] raise ValueError(f"Unknown transport type: {self.transport_type}. Available types: {available_types}") session = await stack.enter_async_context(ClientSession(read_stream, write_stream)) definitions = await self.fetch_tool_definitions_from_session(session) except ConnectionError as e: logger.error(f"Error fetching MCP tool definitions from {self.endpoint}: {e}", exc_info=True) raise except Exception as e: logger.error(f"Unexpected error fetching MCP tool definitions from {self.endpoint}: {e}", exc_info=True) raise RuntimeError(f"Unexpected error during tool definition fetching: {e}") from e finally: await stack.aclose() return definitions @staticmethod async def fetch_tool_definitions_from_session(session: ClientSession) -> List[MCPToolDefinition]: """ Fetch tool definitions from an existing session. Args: session: MCP client session Returns: List of tool definitions Raises: Exception: If listing tools fails """ definitions: List[MCPToolDefinition] = [] try: # `initialize` is idempotent – calling it twice is safe and # ensures the session is ready. await session.initialize() mcp_tools = await _fetch_all_pages( await session.list_tools(), lambda cursor: session.list_tools(cursor=cursor), lambda response: response.tools ) for mcp_tool in mcp_tools: # Capture outputSchema if the MCP server provides one output_schema = getattr(mcp_tool, "outputSchema", None) definitions.append( MCPToolDefinition( name=mcp_tool.name, description=mcp_tool.description, input_schema=mcp_tool.inputSchema or {"type": "object", "properties": {}}, output_schema=output_schema, ) ) if not definitions: logger.warning("No tool definitions found on MCP server") except Exception as e: logger.error("Failed to list tools via MCP session: %s", e, exc_info=True) raise return definitions async def fetch_resource_definitions(self) -> List[MCPResourceDefinition]: """ Fetch resource definitions from the configured endpoint. Returns: List of resource definitions """ if not self.endpoint: raise ValueError("Endpoint is required") resources: List[MCPResourceDefinition] = [] stack = AsyncExitStack() try: if self.transport_type == MCPTransportType.STDIO: command_parts = shlex.split(self.endpoint) if not command_parts: raise ValueError("STDIO command string cannot be empty.") command = command_parts[0] args = command_parts[1:] server_params = StdioServerParameters(command=command, args=args, env=None, cwd=self.working_directory) stdio_transport = await stack.enter_async_context(stdio_client(server_params)) read_stream, write_stream = stdio_transport elif self.transport_type == MCPTransportType.HTTP_STREAM: transport_endpoint = _transport_endpoint(self.endpoint, "/mcp/") transport = await stack.enter_async_context(streamablehttp_client(transport_endpoint)) read_stream, write_stream, _ = transport elif self.transport_type == MCPTransportType.SSE: transport_endpoint = _transport_endpoint(self.endpoint, "/sse") transport = await stack.enter_async_context(sse_client(transport_endpoint)) read_stream, write_stream = transport else: available_types = [t.value for t in MCPTransportType] raise ValueError(f"Unknown transport type: {self.transport_type}. Available types: {available_types}") session = await stack.enter_async_context(ClientSession(read_stream, write_stream)) resources = await self.fetch_resource_definitions_from_session(session) except ConnectionError as e: logger.error(f"Error fetching MCP resources from {self.endpoint}: {e}", exc_info=True) raise except Exception as e: logger.error(f"Unexpected error fetching MCP resources from {self.endpoint}: {e}", exc_info=True) raise RuntimeError(f"Unexpected error during resource fetching: {e}") from e finally: await stack.aclose() return resources @staticmethod async def fetch_resource_definitions_from_session(session: ClientSession) -> List[MCPResourceDefinition]: """ Fetch resource definitions from an existing session. Args: session: MCP client session Returns: List of resource definitions """ resources: List[MCPResourceDefinition] = [] try: await session.initialize() resources_iterable: List[Union[types.Resource, types.ResourceTemplate]] = list( await _fetch_all_pages( await session.list_resources(), lambda cursor: session.list_resources(cursor=cursor), lambda response: response.resources, ) ) try: first_page = await session.list_resource_templates() except McpError as error: if error.error.code != types.METHOD_NOT_FOUND: raise else: res_templates = await _fetch_all_pages( first_page, lambda cursor: session.list_resource_templates(cursor=cursor), lambda response: response.resourceTemplates, ) resources_iterable.extend(res_templates) for mcp_resource in resources_iterable: # Support both attribute-style objects and dict-like responses if hasattr(mcp_resource, "name"): name = mcp_resource.name description = mcp_resource.description if isinstance(mcp_resource, types.ResourceTemplate): uri = str(mcp_resource.uriTemplate) parse_placeholders = True else: uri = str(mcp_resource.uri) parse_placeholders = not isinstance(mcp_resource, types.Resource) elif isinstance(mcp_resource, dict): # assume mapping name = mcp_resource["name"] description = mcp_resource.get("description") uri = str(mcp_resource.get("uri", "")) parse_placeholders = True else: raise ValueError(f"Unexpected resource format: {mcp_resource}") # Preserve encoded URI identity; only explicit template braces define parameters. placeholders = re.findall(r"\{([^}]+)\}", uri) if parse_placeholders else [] properties: Dict[str, Any] = {} for param_name in placeholders: properties[param_name] = {"type": "string", "description": f"URI parameter {param_name}"} resources.append( MCPResourceDefinition( name=name, description=description, uri=uri, mime_type=getattr(mcp_resource, "mimeType", None), input_schema={"type": "object", "properties": properties, "required": list(placeholders)}, ) ) if not resources: logger.warning("No resources found on MCP server") except Exception as e: logger.error("Failed to list resources via MCP session: %s", e, exc_info=True) raise return resources async def fetch_prompt_definitions(self) -> List[MCPPromptDefinition]: """ Fetch prompt/template definitions from the configured endpoint. Returns: List of prompt definitions """ if not self.endpoint: raise ValueError("Endpoint is required") prompts: List[MCPPromptDefinition] = [] stack = AsyncExitStack() try: if self.transport_type == MCPTransportType.STDIO: command_parts = shlex.split(self.endpoint) if not command_parts: raise ValueError("STDIO command string cannot be empty.") command = command_parts[0] args = command_parts[1:] server_params = StdioServerParameters(command=command, args=args, env=None, cwd=self.working_directory) stdio_transport = await stack.enter_async_context(stdio_client(server_params)) read_stream, write_stream = stdio_transport elif self.transport_type == MCPTransportType.HTTP_STREAM: transport_endpoint = _transport_endpoint(self.endpoint, "/mcp/") transport = await stack.enter_async_context(streamablehttp_client(transport_endpoint)) read_stream, write_stream, _ = transport elif self.transport_type == MCPTransportType.SSE: transport_endpoint = _transport_endpoint(self.endpoint, "/sse") transport = await stack.enter_async_context(sse_client(transport_endpoint)) read_stream, write_stream = transport else: available_types = [t.value for t in MCPTransportType] raise ValueError(f"Unknown transport type: {self.transport_type}. Available types: {available_types}") session = await stack.enter_async_context(ClientSession(read_stream, write_stream)) prompts = await self.fetch_prompt_definitions_from_session(session) except ConnectionError as e: logger.error(f"Error fetching MCP prompts from {self.endpoint}: {e}", exc_info=True) raise except Exception as e: logger.error(f"Unexpected error fetching MCP prompts from {self.endpoint}: {e}", exc_info=True) raise RuntimeError(f"Unexpected error during prompt fetching: {e}") from e finally: await stack.aclose() return prompts @staticmethod async def fetch_prompt_definitions_from_session(session: ClientSession) -> List[MCPPromptDefinition]: """ Fetch prompt/template definitions from an existing session. Args: session: MCP client session Returns: List of prompt definitions """ prompts: List[MCPPromptDefinition] = [] try: await session.initialize() mcp_prompts = await _fetch_all_pages( await session.list_prompts(), lambda cursor: session.list_prompts(cursor=cursor), lambda response: response.prompts, ) for mcp_prompt in mcp_prompts: arguments: List[types.PromptArgument] = mcp_prompt.arguments or [] prompts.append( MCPPromptDefinition( name=mcp_prompt.name, description=mcp_prompt.description, input_schema={ "type": "object", "properties": {arg.name: {"type": "string", "description": arg.description} for arg in arguments}, "required": [arg.name for arg in arguments if arg.required], }, ) ) if not prompts: logger.warning("No prompts found on MCP server") except Exception as e: logger.error("Failed to list prompts via MCP session: %s", e, exc_info=True) raise return prompts ``` ### File: atomic-agents/atomic_agents/connectors/mcp/mcp_factory.py ```python import asyncio import json import logging from typing import Any, Dict, List, Type, Optional, Union, Tuple, cast, Protocol, Callable from contextlib import AsyncExitStack import shlex import types from pydantic import create_model, Field, BaseModel, JsonValue, ValidationError from mcp import ClientSession, StdioServerParameters from mcp.client.sse import sse_client from mcp.client.stdio import stdio_client from mcp.client.streamable_http import streamablehttp_client import mcp.types from atomic_agents.base.base_io_schema import BaseIOSchema from atomic_agents.base import BaseTool, BaseResource, BasePrompt from atomic_agents.connectors.mcp.schema_transformer import SchemaTransformer from atomic_agents.connectors.mcp.mcp_definition_service import ( _transport_endpoint, MCPAttributeType, MCPDefinitionService, MCPToolDefinition, MCPTransportType, MCPResourceDefinition, MCPPromptDefinition, ) logger = logging.getLogger(__name__) class MCPToolOutputSchema(BaseIOSchema): """Generic output schema for dynamically generated MCP tools. Used as a fallback when the MCP server does not provide an outputSchema definition. Tools with MCP-provided outputSchema will have typed output schemas instead. When a result has no content blocks, its structuredContent is returned directly. """ result: Any = Field(..., description="The result returned by the MCP tool.") class MCPResourceOutputSchema(BaseIOSchema): """Generic output schema for dynamically generated MCP resources.""" content: Any = Field(..., description="The content of the MCP resource.") mime_type: Optional[str] = Field(None, description="The MIME type of the resource.") class MCPPromptOutputSchema(BaseIOSchema): """Generic output schema for dynamically generated MCP prompts.""" content: str = Field(..., description="The content of the MCP prompt.") class _ModelDumpPayload(Protocol): model_dump: Callable[[], dict[str, JsonValue]] class _TextToolContent(Protocol): text: JsonValue | bytes | bytearray class _DataToolContent(Protocol): data: JsonValue _ToolContentItem = BaseModel | JsonValue | _TextToolContent | _DataToolContent _ToolResult = BaseModel | JsonValue | tuple[object, ...] | list[object] _ToolContent = BaseModel | JsonValue | tuple[_ToolContentItem, ...] | list[_ToolContentItem] class _GeneratedMCPTool(Protocol): mcp_tool_name: str mcp_endpoint: str | None transport_type: MCPTransportType output_schema: type[BaseIOSchema] def _validated_tool_envelope(tool_result: _ToolResult) -> mcp.types.CallToolResult | None: if isinstance(tool_result, mcp.types.CallToolResult): return tool_result if isinstance(tool_result, dict): try: return mcp.types.CallToolResult.model_validate(tool_result, strict=True) except ValidationError: return None return None def _tool_error_diagnostic(envelope: mcp.types.CallToolResult) -> str: text_parts = [block.text.strip() for block in envelope.content if isinstance(block, mcp.types.TextContent)] return "\n".join(filter(None, text_parts)) or "MCP tool reported an error without text content" def _raise_tool_execution_error(tool_result: _ToolResult) -> None: envelope = _validated_tool_envelope(tool_result) if envelope is not None and envelope.isError: raise RuntimeError(_tool_error_diagnostic(envelope)) def _validate_structured_tool_output( structured_data: JsonValue | _ModelDumpPayload, output_schema: type[BaseIOSchema], tool_name: str ) -> BaseIOSchema: if isinstance(structured_data, dict): return output_schema(**structured_data) elif hasattr(structured_data, "model_dump"): model_dump: Callable[[], dict[str, JsonValue]] = getattr(structured_data, "model_dump") return output_schema(**model_dump()) logger.error( f"Unexpected structuredContent type for tool '{tool_name}': " f"got {type(structured_data).__name__}, expected dict or BaseModel. " f"Content: {structured_data!r}" ) raise TypeError( f"MCP tool '{tool_name}' returned structuredContent with unexpected type " f"{type(structured_data).__name__}. Expected dict or BaseModel." ) def _log_invalid_tool_json(text: str | bytes | bytearray, tool_name: str, error: json.JSONDecodeError) -> None: logger.debug( f"Tool '{tool_name}' content is not valid JSON: {error}. " f"Content preview: {text[:200]!r}..." if len(text) > 200 else f"Content: {text!r}" ) def _validate_json_text_tool_output( text: JsonValue | bytes | bytearray, output_schema: type[BaseIOSchema], tool_name: str ) -> BaseIOSchema | None: try: # Invalid adapter text must keep json.loads' native TypeError and the existing data fallback. parsed: JsonValue = json.loads(text) # pyright: ignore[reportArgumentType] if isinstance(parsed, dict): return output_schema(**parsed) logger.debug( f"Tool '{tool_name}' content parsed as JSON but was " f"{type(parsed).__name__}, not dict. Trying other extraction methods." ) except json.JSONDecodeError as error: if isinstance(text, (str, bytes, bytearray)): _log_invalid_tool_json(text, tool_name, error) except TypeError as error: logger.warning(f"Tool '{tool_name}' content.text has unexpected type {type(text).__name__}: {error}") return None def _validate_tool_content_item( content_item: _ToolContentItem, output_schema: type[BaseIOSchema], tool_name: str ) -> BaseIOSchema | None: if hasattr(content_item, "text"): text: JsonValue | bytes | bytearray = getattr(content_item, "text") parsed_output = _validate_json_text_tool_output(text, output_schema, tool_name) if parsed_output is not None: return parsed_output if hasattr(content_item, "data"): data: JsonValue = getattr(content_item, "data") if isinstance(data, dict): return output_schema(**data) return None def _validate_tool_content_output( content: _ToolContent, output_schema: type[BaseIOSchema], tool_name: str ) -> BaseIOSchema | None: if content and isinstance(content, (list, tuple)) and len(content) > 0: return _validate_tool_content_item(content[0], output_schema, tool_name) return None def _validate_model_tool_output( tool_result: BaseModel, output_schema: type[BaseIOSchema], tool_name: str ) -> BaseIOSchema | None: structured_data: JsonValue | _ModelDumpPayload = getattr(tool_result, "structuredContent", None) if structured_data is not None: return _validate_structured_tool_output(structured_data, output_schema, tool_name) elif hasattr(tool_result, "content"): content: _ToolContent = getattr(tool_result, "content") return _validate_tool_content_output(content, output_schema, tool_name) return None def _validate_mapping_tool_output(tool_result: dict[str, JsonValue], output_schema: type[BaseIOSchema]) -> BaseIOSchema: if "structuredContent" in tool_result: # Malformed mapping payloads must keep the constructor's native keyword-unpacking TypeError. return output_schema(**tool_result["structuredContent"]) # pyright: ignore[reportCallIssue] elif "content" in tool_result: content = tool_result["content"] if isinstance(content, dict): return output_schema(**content) return output_schema(**tool_result) def _validate_typed_tool_output(tool_result: _ToolResult, output_schema: type[BaseIOSchema], tool_name: str) -> BaseIOSchema: if isinstance(tool_result, BaseModel): parsed_output = _validate_model_tool_output(tool_result, output_schema, tool_name) if parsed_output is not None: return parsed_output elif isinstance(tool_result, dict): return _validate_mapping_tool_output(tool_result, output_schema) logger.error( f"Could not parse structured output for tool '{tool_name}'. " f"Expected typed output but got: type={type(tool_result).__name__}, " f"value={tool_result!r}" ) raise ValueError( f"MCP tool '{tool_name}' has outputSchema but returned unparseable result. " f"Received type: {type(tool_result).__name__}. Check MCP server implementation." ) def _render_generic_tool_output(tool_result: _ToolResult, output_schema: type[BaseIOSchema]) -> BaseIOSchema: envelope = _validated_tool_envelope(tool_result) if envelope is not None and not envelope.content and envelope.structuredContent is not None: return output_schema(**{"result": envelope.structuredContent}) content: _ToolContent | _ToolResult if isinstance(tool_result, BaseModel) and hasattr(tool_result, "content"): content = getattr(tool_result, "content") elif isinstance(tool_result, dict) and "content" in tool_result: content = tool_result["content"] else: content = tool_result return output_schema(**{"result": content}) class MCPFactory: """Factory for creating MCP tool classes.""" def __init__( self, mcp_endpoint: Optional[str] = None, transport_type: MCPTransportType = MCPTransportType.HTTP_STREAM, client_session: Optional[ClientSession] = None, event_loop: Optional[asyncio.AbstractEventLoop] = None, working_directory: Optional[str] = None, ): """ Initialize the factory. Args: mcp_endpoint: MCP server base URL or explicit /mcp or /sse endpoint (for SSE/HTTP stream), or the full command to run the server (for STDIO). Query parameters are preserved. transport_type: Type of transport to use (SSE, HTTP_STREAM, or STDIO) client_session: Optional pre-initialized ClientSession for reuse event_loop: Optional event loop for running asynchronous operations working_directory: Optional working directory to use when running STDIO commands """ self.mcp_endpoint = mcp_endpoint self.transport_type = transport_type self.client_session = client_session self.event_loop = event_loop self.schema_transformer = SchemaTransformer() self.working_directory = working_directory # Validate configuration if client_session is not None and event_loop is None: raise ValueError("When `client_session` is provided an `event_loop` must also be supplied.") if not mcp_endpoint and client_session is None: raise ValueError("`mcp_endpoint` must be provided when no `client_session` is supplied.") def create_tools(self) -> List[Type[BaseTool]]: """ Create tool classes from the configured endpoint or session. Returns: List of dynamically generated BaseTool subclasses """ tool_definitions = self._fetch_tool_definitions() if not tool_definitions: return [] return self._create_tool_classes(tool_definitions) def _fetch_tool_definitions(self) -> List[MCPToolDefinition]: """ Fetch tool definitions using the appropriate method. Returns: List of tool definitions """ if self.client_session is not None: # Use existing session async def _gather_defs(): return await MCPDefinitionService.fetch_tool_definitions_from_session(self.client_session) # pragma: no cover return cast(asyncio.AbstractEventLoop, self.event_loop).run_until_complete(_gather_defs()) # pragma: no cover else: # Create new connection service = MCPDefinitionService( self.mcp_endpoint, self.transport_type, self.working_directory, ) return asyncio.run(service.fetch_tool_definitions()) def _create_tool_classes(self, tool_definitions: List[MCPToolDefinition]) -> List[Type[BaseTool]]: """ Create tool classes from definitions. Args: tool_definitions: List of tool definitions Returns: List of dynamically generated BaseTool subclasses """ generated_tools = [] for definition in tool_definitions: try: tool_name = definition.name tool_description = definition.description or f"Dynamically generated tool for MCP tool: {tool_name}" input_schema_dict = definition.input_schema # Create input schema InputSchema = self.schema_transformer.create_model_from_schema( input_schema_dict, f"{tool_name}InputSchema", tool_name, f"Input schema for {tool_name}", attribute_type=MCPAttributeType.TOOL, ) # Create output schema - use MCP-provided schema if available, otherwise fallback to generic. # When a typed output schema is used, _has_typed_output_schema is set on the tool class # to enable structured content extraction at runtime (see result processing below). output_schema_dict: Optional[Dict[str, Any]] = definition.output_schema has_typed_output_schema = False if output_schema_dict: # Use the schema transformer to create a proper typed output schema OutputSchema = self.schema_transformer.create_model_from_schema( output_schema_dict, f"{tool_name}OutputSchema", tool_name, f"Output schema for {tool_name}", attribute_type=MCPAttributeType.TOOL, is_output_schema=True, ) has_typed_output_schema = True else: # Fallback to generic output schema OutputSchema = type( f"{tool_name}OutputSchema", (MCPToolOutputSchema,), {"__doc__": f"Output schema for {tool_name}"} ) # Async implementation async def run_tool_async(self: _GeneratedMCPTool, params: BaseIOSchema) -> BaseIOSchema: bound_tool_name = self.mcp_tool_name bound_mcp_endpoint = self.mcp_endpoint # May be None when using external session bound_transport_type = self.transport_type persistent_session: Optional[ClientSession] = getattr(self, "_client_session", None) bound_working_directory = getattr(self, "working_directory", None) # Get arguments, excluding tool_name arguments = params.model_dump(exclude={"tool_name"}, exclude_none=True) async def _connect_and_call(): stack = AsyncExitStack() try: if bound_transport_type == MCPTransportType.STDIO: # Split the command string into the command and its arguments command_parts = shlex.split(bound_mcp_endpoint) if not command_parts: raise ValueError("STDIO command string cannot be empty.") command = command_parts[0] args = command_parts[1:] logger.debug(f"Executing tool '{bound_tool_name}' via STDIO: command='{command}', args={args}") server_params = StdioServerParameters( command=command, args=args, env=None, cwd=bound_working_directory ) stdio_transport = await stack.enter_async_context(stdio_client(server_params)) read_stream, write_stream = stdio_transport elif bound_transport_type == MCPTransportType.HTTP_STREAM: # HTTP Stream transport - use trailing slash to avoid redirect # See: https://github.com/modelcontextprotocol/python-sdk/issues/732 http_endpoint = _transport_endpoint(bound_mcp_endpoint, "/mcp/") logger.debug(f"Executing tool '{bound_tool_name}' via HTTP Stream: endpoint={http_endpoint}") http_transport = await stack.enter_async_context(streamablehttp_client(http_endpoint)) read_stream, write_stream, _ = http_transport elif bound_transport_type == MCPTransportType.SSE: # SSE transport (deprecated) sse_endpoint = _transport_endpoint(bound_mcp_endpoint, "/sse") logger.debug(f"Executing tool '{bound_tool_name}' via SSE: endpoint={sse_endpoint}") sse_transport = await stack.enter_async_context(sse_client(sse_endpoint)) read_stream, write_stream = sse_transport else: available_types = [t.value for t in MCPTransportType] raise ValueError( f"Unknown transport type: {bound_transport_type}. Available transport types: {available_types}" ) session = await stack.enter_async_context(ClientSession(read_stream, write_stream)) await session.initialize() # Ensure arguments is a dict, even if empty call_args = arguments if isinstance(arguments, dict) else {} tool_result = await session.call_tool(name=bound_tool_name, arguments=call_args) return tool_result finally: await stack.aclose() async def _call_with_persistent_session(): # Ensure arguments is a dict, even if empty call_args = arguments if isinstance(arguments, dict) else {} return await persistent_session.call_tool(name=bound_tool_name, arguments=call_args) try: if persistent_session is not None: # Use the always‑on session/loop supplied at construction time. tool_result = await _call_with_persistent_session() else: # Legacy behaviour – open a fresh connection per invocation. tool_result = await _connect_and_call() _raise_tool_execution_error(tool_result) output_schema = self.output_schema if getattr(self, "_has_typed_output_schema", False): return _validate_typed_tool_output(tool_result, output_schema, bound_tool_name) return _render_generic_tool_output(tool_result, output_schema) except Exception as error: logger.error(f"Error executing MCP tool '{bound_tool_name}': {error}", exc_info=True) raise RuntimeError(f"Failed to execute MCP tool '{bound_tool_name}': {error}") from error # Create sync wrapper def run_tool_sync(self, params: InputSchema) -> OutputSchema: # type: ignore persistent_session: Optional[ClientSession] = getattr(self, "_client_session", None) loop: Optional[asyncio.AbstractEventLoop] = getattr(self, "_event_loop", None) if persistent_session is not None: # Use the always‑on session/loop supplied at construction time. try: return cast(asyncio.AbstractEventLoop, loop).run_until_complete(self.arun(params)) except AttributeError as e: raise RuntimeError(f"Failed to execute MCP tool '{self.mcp_tool_name}': {e}") from e else: # Legacy behaviour – run in new event loop. return asyncio.run(self.arun(params)) # Create the tool class using types.new_class() instead of type() attrs = { "arun": run_tool_async, "run": run_tool_sync, "__doc__": tool_description, "mcp_tool_name": tool_name, "mcp_endpoint": self.mcp_endpoint, "transport_type": self.transport_type, "_client_session": self.client_session, "_event_loop": self.event_loop, "working_directory": self.working_directory, "_has_typed_output_schema": has_typed_output_schema, } # Create the class using new_class() for proper generic type support tool_class = types.new_class( tool_name, (BaseTool[InputSchema, OutputSchema],), {}, lambda ns: ns.update(attrs) ) # Add the input_schema and output_schema class attributes explicitly # since they might not be properly inherited with types.new_class setattr(tool_class, "input_schema", InputSchema) setattr(tool_class, "output_schema", OutputSchema) generated_tools.append(tool_class) except Exception as e: logger.error(f"Error generating class for tool '{definition.name}': {e}", exc_info=True) continue return generated_tools def create_orchestrator_schema( self, tools: Optional[List[Type[BaseTool]]] = None, resources: Optional[List[Type[BaseResource]]] = None, prompts: Optional[List[Type[BasePrompt]]] = None, ) -> Optional[Type[BaseIOSchema]]: """ Create an orchestrator schema for the given tools. Args: tools: List of tool classes resources: List of resource classes prompts: List of prompt classes Returns: Orchestrator schema or None if no tools provided """ if tools is None and resources is None and prompts is None: logger.warning("No tools/resources/prompts provided to create orchestrator schema") return None if tools is None: tools = [] if resources is None: resources = [] if prompts is None: prompts = [] tool_schemas = [ToolClass.input_schema for ToolClass in tools] resource_schemas = [ResourceClass.input_schema for ResourceClass in resources] prompt_schemas = [PromptClass.input_schema for PromptClass in prompts] # Build runtime Union types for each attribute group when present field_defs = {} if tool_schemas: ToolUnion = Union[tuple(tool_schemas)] field_defs["tool_parameters"] = ( ToolUnion, Field( ..., description="The parameters for the selected tool, matching its specific schema (which includes the 'tool_name').", ), ) if resource_schemas: ResourceUnion = Union[tuple(resource_schemas)] field_defs["resource_parameters"] = ( ResourceUnion, Field( ..., description="The parameters for the selected resource, matching its specific schema (which includes the 'resource_name').", ), ) if prompt_schemas: PromptUnion = Union[tuple(prompt_schemas)] field_defs["prompt_parameters"] = ( PromptUnion, Field( ..., description="The parameters for the selected prompt, matching its specific schema (which includes the 'prompt_name').", ), ) if not field_defs: logger.warning("No schemas available to create orchestrator union") return None # Dynamically create the output schema with the appropriate fields orchestrator_schema = create_model( "MCPOrchestratorOutputSchema", __doc__="Output schema for the MCP Orchestrator Agent. Contains the parameters for the selected tool/resource/prompt.", __base__=BaseIOSchema, **field_defs, ) return orchestrator_schema def create_resources(self) -> List[Type[BaseResource]]: """ Create resource classes from the configured endpoint or session. Returns: List of dynamically generated resource classes """ resource_definitions = self._fetch_resource_definitions() if not resource_definitions: return [] return self._create_resource_classes(resource_definitions) def _fetch_resource_definitions(self) -> List[MCPResourceDefinition]: """ Fetch resource definitions using the appropriate method. Returns: List of resource definitions """ if self.client_session is not None: # Use existing session async def _gather_defs(): return await MCPDefinitionService.fetch_resource_definitions_from_session( self.client_session ) # pragma: no cover return cast(asyncio.AbstractEventLoop, self.event_loop).run_until_complete(_gather_defs()) # pragma: no cover else: # Create new connection service = MCPDefinitionService( self.mcp_endpoint, self.transport_type, self.working_directory, ) return asyncio.run(service.fetch_resource_definitions()) def _create_resource_classes(self, resource_definitions: List[MCPResourceDefinition]) -> List[Type[BaseResource]]: """ Create resource classes from definitions. Args: resource_definitions: List of resource definitions Returns: List of dynamically generated resource classes """ generated_resources = [] for definition in resource_definitions: try: resource_name = definition.name resource_description = ( definition.description or f"Dynamically generated resource for MCP resource: {resource_name}" ) uri = definition.uri mime_type = definition.mime_type InputSchema = self.schema_transformer.create_model_from_schema( definition.input_schema, f"{resource_name}InputSchema", resource_name, f"Input schema for {resource_name}", attribute_type=MCPAttributeType.RESOURCE, ) # Create output schema OutputSchema = type( f"{resource_name}OutputSchema", (MCPResourceOutputSchema,), {"__doc__": f"Output schema for {resource_name}"}, ) # Async implementation async def read_resource_async(self, params: InputSchema) -> OutputSchema: # type: ignore bound_uri = self.uri bound_mcp_endpoint = self.mcp_endpoint # May be None when using external session bound_transport_type = self.transport_type persistent_session: Optional[ClientSession] = getattr(self, "_client_session", None) bound_working_directory = getattr(self, "working_directory", None) arguments = params.model_dump(exclude={"resource_name"}, exclude_none=True) async def _connect_and_read(): stack = AsyncExitStack() try: if bound_transport_type == MCPTransportType.STDIO: # Split the command string into the command and its arguments command_parts = shlex.split(bound_mcp_endpoint) if not command_parts: raise ValueError("STDIO command string cannot be empty.") command = command_parts[0] args = command_parts[1:] logger.debug( f"Reading resource '{self.mcp_resource_name}' via STDIO: command='{command}', args={args}" ) server_params = StdioServerParameters( command=command, args=args, env=None, cwd=bound_working_directory ) stdio_transport = await stack.enter_async_context(stdio_client(server_params)) read_stream, write_stream = stdio_transport elif bound_transport_type == MCPTransportType.HTTP_STREAM: # HTTP Stream transport - use trailing slash to avoid redirect # See: https://github.com/modelcontextprotocol/python-sdk/issues/732 http_endpoint = _transport_endpoint(bound_mcp_endpoint, "/mcp/") logger.debug( f"Reading resource '{self.mcp_resource_name}' via HTTP Stream: endpoint={http_endpoint}" ) http_transport = await stack.enter_async_context(streamablehttp_client(http_endpoint)) read_stream, write_stream, _ = http_transport elif bound_transport_type == MCPTransportType.SSE: # SSE transport (deprecated) sse_endpoint = _transport_endpoint(bound_mcp_endpoint, "/sse") logger.debug(f"Reading resource '{self.mcp_resource_name}' via SSE: endpoint={sse_endpoint}") sse_transport = await stack.enter_async_context(sse_client(sse_endpoint)) read_stream, write_stream = sse_transport else: available_types = [t.value for t in MCPTransportType] raise ValueError( f"Unknown transport type: {bound_transport_type}. Available transport types: {available_types}" ) session = await stack.enter_async_context(ClientSession(read_stream, write_stream)) await session.initialize() # Substitute URI placeholders with provided parameters when available. call_args = arguments if isinstance(arguments, dict) else {} # If params contain keys, format the URI template. try: concrete_uri = bound_uri.format(**call_args) if call_args else bound_uri except Exception: concrete_uri = bound_uri resource_result: mcp.types.ReadResourceResult = await session.read_resource(uri=concrete_uri) return resource_result finally: await stack.aclose() async def _read_with_persistent_session(): call_args = arguments if isinstance(arguments, dict) else {} try: concrete_uri_p = bound_uri.format(**call_args) if call_args else bound_uri except Exception: concrete_uri_p = bound_uri return await persistent_session.read_resource(uri=concrete_uri_p) try: if persistent_session is not None: # Use the always‑on session/loop supplied at construction time. resource_result = await _read_with_persistent_session() else: # Legacy behaviour – open a fresh connection per invocation. resource_result = await _connect_and_read() # Process the result if isinstance(resource_result, BaseModel) and hasattr(resource_result, "contents"): actual_content = resource_result.contents # MCP stores mimeType in each content item, not on the result itself if actual_content and len(actual_content) > 0: # Get mimeType from the first content item first_content = actual_content[0] actual_mime = getattr(first_content, "mimeType", self._mime_type) else: actual_mime = self._mime_type elif isinstance(resource_result, dict) and "contents" in resource_result: actual_content = resource_result["contents"] actual_mime = resource_result.get("mime_type", self._mime_type) else: actual_content = resource_result actual_mime = self._mime_type return self.output_schema(content=actual_content, mime_type=actual_mime) except Exception as e: logger.error(f"Error reading MCP resource '{self.mcp_resource_name}': {e}", exc_info=True) raise RuntimeError(f"Failed to read MCP resource '{self.mcp_resource_name}': {e}") from e # Create sync wrapper def read_resource_sync(self, params: InputSchema) -> OutputSchema: # type: ignore persistent_session: Optional[ClientSession] = getattr(self, "_client_session", None) loop: Optional[asyncio.AbstractEventLoop] = getattr(self, "_event_loop", None) if persistent_session is not None: # Use the always‑on session/loop supplied at construction time. try: return cast(asyncio.AbstractEventLoop, loop).run_until_complete(self.aread(params)) except AttributeError as e: raise RuntimeError(f"Failed to read MCP resource '{self.mcp_resource_name}': {e}") from e else: # Legacy behaviour – run in new event loop. return asyncio.run(self.aread(params)) # Create the resource class using types.new_class() instead of type() attrs = { "aread": read_resource_async, "read": read_resource_sync, "__doc__": resource_description, "mcp_resource_name": resource_name, "mcp_endpoint": self.mcp_endpoint, "transport_type": self.transport_type, "_client_session": self.client_session, "_event_loop": self.event_loop, "working_directory": self.working_directory, "uri": uri, "_mime_type": mime_type, } # Create the class using new_class() for proper generic type support resource_class = types.new_class( resource_name, (BaseResource[InputSchema, OutputSchema],), {}, lambda ns: ns.update(attrs) ) # Add the input_schema and output_schema class attributes explicitly setattr(resource_class, "input_schema", InputSchema) setattr(resource_class, "output_schema", OutputSchema) generated_resources.append(resource_class) except Exception as e: logger.error(f"Error generating class for resource '{definition.name}': {e}", exc_info=True) continue return generated_resources def create_prompts(self) -> List[Type[BasePrompt]]: """ Create prompt classes from the configured endpoint or session. Returns: List of dynamically generated prompt classes """ prompt_definitions = self._fetch_prompt_definitions() if not prompt_definitions: return [] return self._create_prompt_classes(prompt_definitions) def _fetch_prompt_definitions(self) -> List[MCPPromptDefinition]: """ Fetch prompt definitions using the appropriate method. Returns: List of prompt definitions """ if self.client_session is not None: # Use existing session async def _gather_defs(): return await MCPDefinitionService.fetch_prompt_definitions_from_session( self.client_session ) # pragma: no cover return cast(asyncio.AbstractEventLoop, self.event_loop).run_until_complete(_gather_defs()) # pragma: no cover else: # Create new connection service = MCPDefinitionService( self.mcp_endpoint, self.transport_type, self.working_directory, ) return asyncio.run(service.fetch_prompt_definitions()) def _create_prompt_classes(self, prompt_definitions: List[MCPPromptDefinition]) -> List[Type[BasePrompt]]: """ Create prompt classes from definitions. Args: prompt_definitions: List of prompt definitions Returns: List of dynamically generated prompt classes """ generated_prompts = [] for definition in prompt_definitions: try: prompt_name = definition.name prompt_description = definition.description or f"Dynamically generated prompt for MCP prompt: {prompt_name}" InputSchema = self.schema_transformer.create_model_from_schema( definition.input_schema, f"{prompt_name}InputSchema", prompt_name, f"Input schema for {prompt_name}", attribute_type=MCPAttributeType.PROMPT, ) # Create output schema OutputSchema = type( f"{prompt_name}OutputSchema", (MCPPromptOutputSchema,), {"__doc__": f"Output schema for {prompt_name}"} ) # Async implementation async def generate_prompt_async(self, params: InputSchema) -> OutputSchema: # type: ignore bound_prompt_name = self.mcp_prompt_name bound_mcp_endpoint = self.mcp_endpoint # May be None when using external session bound_transport_type = self.transport_type persistent_session: Optional[ClientSession] = getattr(self, "_client_session", None) bound_working_directory = getattr(self, "working_directory", None) # Get arguments arguments = params.model_dump(exclude={"prompt_name"}, exclude_none=True) async def _connect_and_generate(): stack = AsyncExitStack() try: if bound_transport_type == MCPTransportType.STDIO: # Split the command string into the command and its arguments command_parts = shlex.split(bound_mcp_endpoint) if not command_parts: raise ValueError("STDIO command string cannot be empty.") command = command_parts[0] args = command_parts[1:] logger.debug( f"Getting prompt '{bound_prompt_name}' via STDIO: command='{command}', args={args}" ) server_params = StdioServerParameters( command=command, args=args, env=None, cwd=bound_working_directory ) stdio_transport = await stack.enter_async_context(stdio_client(server_params)) read_stream, write_stream = stdio_transport elif bound_transport_type == MCPTransportType.HTTP_STREAM: # HTTP Stream transport - use trailing slash to avoid redirect # See: https://github.com/modelcontextprotocol/python-sdk/issues/732 http_endpoint = _transport_endpoint(bound_mcp_endpoint, "/mcp/") logger.debug(f"Getting prompt '{bound_prompt_name}' via HTTP Stream: endpoint={http_endpoint}") http_transport = await stack.enter_async_context(streamablehttp_client(http_endpoint)) read_stream, write_stream, _ = http_transport elif bound_transport_type == MCPTransportType.SSE: # SSE transport (deprecated) sse_endpoint = _transport_endpoint(bound_mcp_endpoint, "/sse") logger.debug(f"Getting prompt '{bound_prompt_name}' via SSE: endpoint={sse_endpoint}") sse_transport = await stack.enter_async_context(sse_client(sse_endpoint)) read_stream, write_stream = sse_transport else: available_types = [t.value for t in MCPTransportType] raise ValueError( f"Unknown transport type: {bound_transport_type}. Available transport types: {available_types}" ) session = await stack.enter_async_context(ClientSession(read_stream, write_stream)) await session.initialize() # Ensure arguments is a dict, even if empty call_args = arguments if isinstance(arguments, dict) else {} prompt_result = await session.get_prompt(name=bound_prompt_name, arguments=call_args) return prompt_result finally: await stack.aclose() async def _get_with_persistent_session(): # Ensure arguments is a dict, even if empty call_args = arguments if isinstance(arguments, dict) else {} return await persistent_session.get_prompt(name=bound_prompt_name, arguments=call_args) try: if persistent_session is not None: # Use the always‑on session/loop supplied at construction time. prompt_result = await _get_with_persistent_session() else: # Legacy behaviour – open a fresh connection per invocation. prompt_result = await _connect_and_generate() # Process the result messages = None if isinstance(prompt_result, BaseModel) and hasattr(prompt_result, "messages"): messages = prompt_result.messages elif isinstance(prompt_result, dict) and "messages" in prompt_result: messages = prompt_result["messages"] else: raise Exception("Prompt response has no messages.") texts = [] for message in messages: if isinstance(message, BaseModel) and hasattr(message, "content"): content = message.content # type: ignore elif isinstance(message, dict) and "content" in message: content = message["content"] else: content = message if isinstance(content, str): texts.append(content) elif isinstance(content, dict): texts.append(content.get("text", "")) elif getattr(content, "text", None): texts.append(content.text) # type: ignore else: texts.append(str(content)) final_content = "\n\n".join(texts) return self.output_schema(content=final_content) except Exception as e: logger.error(f"Error getting MCP prompt '{bound_prompt_name}': {e}", exc_info=True) raise RuntimeError(f"Failed to get MCP prompt '{bound_prompt_name}': {e}") from e # Create sync wrapper def generate_prompt_sync(self, params: InputSchema) -> OutputSchema: # type: ignore persistent_session: Optional[ClientSession] = getattr(self, "_client_session", None) loop: Optional[asyncio.AbstractEventLoop] = getattr(self, "_event_loop", None) if persistent_session is not None: # Use the always‑on session/loop supplied at construction time. try: return cast(asyncio.AbstractEventLoop, loop).run_until_complete(self.agenerate(params)) except AttributeError as e: raise RuntimeError(f"Failed to get MCP prompt '{self.mcp_prompt_name}': {e}") from e else: # Legacy behaviour – run in new event loop. return asyncio.run(self.agenerate(params)) # Create the prompt class using types.new_class() instead of type() attrs = { "agenerate": generate_prompt_async, "generate": generate_prompt_sync, "__doc__": prompt_description, "mcp_prompt_name": prompt_name, "mcp_endpoint": self.mcp_endpoint, "transport_type": self.transport_type, "_client_session": self.client_session, "_event_loop": self.event_loop, "working_directory": self.working_directory, } # Create the class using new_class() for proper generic type support prompt_class = types.new_class( prompt_name, (BasePrompt[InputSchema, OutputSchema],), {}, lambda ns: ns.update(attrs) ) # Add the input_schema and output_schema class attributes explicitly setattr(prompt_class, "input_schema", InputSchema) setattr(prompt_class, "output_schema", OutputSchema) generated_prompts.append(prompt_class) except Exception as e: logger.error(f"Error generating class for prompt '{definition.name}': {e}", exc_info=True) continue return generated_prompts # Public API functions def fetch_mcp_tools( mcp_endpoint: Optional[str] = None, transport_type: MCPTransportType = MCPTransportType.HTTP_STREAM, *, client_session: Optional[ClientSession] = None, event_loop: Optional[asyncio.AbstractEventLoop] = None, working_directory: Optional[str] = None, ) -> List[Type[BaseTool]]: """ Connects to an MCP server via SSE, HTTP Stream or STDIO, discovers tool definitions, and dynamically generates synchronous Atomic Agents compatible BaseTool subclasses for each tool. Each generated tool will establish its own connection when its `run` method is called. Args: mcp_endpoint: URL of the MCP server or command for STDIO. transport_type: Type of transport to use (SSE, HTTP_STREAM, or STDIO). client_session: Optional pre-initialized ClientSession for reuse. event_loop: Optional event loop for running asynchronous operations. working_directory: Optional working directory for STDIO. """ factory = MCPFactory(mcp_endpoint, transport_type, client_session, event_loop, working_directory) return factory.create_tools() async def fetch_mcp_tools_async( mcp_endpoint: Optional[str] = None, transport_type: MCPTransportType = MCPTransportType.STDIO, *, client_session: Optional[ClientSession] = None, working_directory: Optional[str] = None, ) -> List[Type[BaseTool]]: """ Asynchronously connects to an MCP server and dynamically generates BaseTool subclasses for each tool. Must be called within an existing asyncio event loop context. Args: mcp_endpoint: URL of the MCP server (for HTTP/SSE) or command for STDIO. transport_type: Type of transport to use (SSE, HTTP_STREAM, or STDIO). client_session: Optional pre-initialized ClientSession for reuse. working_directory: Optional working directory for STDIO transport. """ if client_session is not None: tool_defs = await MCPDefinitionService.fetch_tool_definitions_from_session(client_session) factory = MCPFactory(mcp_endpoint, transport_type, client_session, asyncio.get_running_loop(), working_directory) else: service = MCPDefinitionService(mcp_endpoint, transport_type, working_directory) tool_defs = await service.fetch_tool_definitions() factory = MCPFactory(mcp_endpoint, transport_type, None, None, working_directory) return factory._create_tool_classes(tool_defs) def create_mcp_orchestrator_schema( tools: Optional[List[Type[BaseTool]]] = None, resources: Optional[List[Type[BaseResource]]] = None, prompts: Optional[List[Type[BasePrompt]]] = None, ) -> Optional[Type[BaseIOSchema]]: """ Creates a schema for the MCP Orchestrator's output using the Union of all tool input schemas. Args: tools: List of dynamically generated MCP tool classes Returns: A Pydantic model class to be used as the output schema for an orchestrator agent """ # Bypass constructor validation since orchestrator schema does not require endpoint or session factory = object.__new__(MCPFactory) return MCPFactory.create_orchestrator_schema(factory, tools, resources, prompts) def fetch_mcp_attributes_with_schema( mcp_endpoint: Optional[str] = None, transport_type: MCPTransportType = MCPTransportType.HTTP_STREAM, *, client_session: Optional[ClientSession] = None, event_loop: Optional[asyncio.AbstractEventLoop] = None, working_directory: Optional[str] = None, ) -> Tuple[List[Type[BaseTool]], List[Type[BaseResource]], List[Type[BasePrompt]], Optional[Type[BaseIOSchema]]]: """ Fetches MCP tools and creates an orchestrator schema for them. Returns both as a tuple. Args: mcp_endpoint: URL of the MCP server or command for STDIO. transport_type: Type of transport to use (SSE, HTTP_STREAM, or STDIO). client_session: Optional pre-initialized ClientSession for reuse. event_loop: Optional event loop for running asynchronous operations. working_directory: Optional working directory for STDIO. Returns: A tuple containing: - List of dynamically generated tool classes - List of dynamically generated resource classes - List of dynamically generated prompt classes - Orchestrator output schema with Union of tool input schemas, or None if no tools found. """ factory = MCPFactory(mcp_endpoint, transport_type, client_session, event_loop, working_directory) tools = factory.create_tools() resources = factory.create_resources() prompts = factory.create_prompts() if not tools and not resources and not prompts: return [], [], [], None orchestrator_schema = factory.create_orchestrator_schema(tools, resources, prompts) return tools, resources, prompts, orchestrator_schema # Resource / Prompt convenience API def fetch_mcp_resources( mcp_endpoint: Optional[str] = None, transport_type: MCPTransportType = MCPTransportType.HTTP_STREAM, *, client_session: Optional[ClientSession] = None, event_loop: Optional[asyncio.AbstractEventLoop] = None, working_directory: Optional[str] = None, ) -> List[Type[BaseResource]]: """ Fetch resource classes from an MCP server (sync). """ factory = MCPFactory(mcp_endpoint, transport_type, client_session, event_loop, working_directory) return factory.create_resources() async def fetch_mcp_resources_async( mcp_endpoint: Optional[str] = None, transport_type: MCPTransportType = MCPTransportType.HTTP_STREAM, *, client_session: Optional[ClientSession] = None, working_directory: Optional[str] = None, ) -> List[Type[BaseResource]]: """ Async version of fetch_mcp_resources. Call from within an event loop. """ if client_session is not None: resource_defs = await MCPDefinitionService.fetch_resource_definitions_from_session(client_session) factory = MCPFactory(mcp_endpoint, transport_type, client_session, asyncio.get_running_loop(), working_directory) else: service = MCPDefinitionService(mcp_endpoint, transport_type, working_directory) resource_defs = await service.fetch_resource_definitions() factory = MCPFactory(mcp_endpoint, transport_type, None, None, working_directory) return factory._create_resource_classes(resource_defs) def fetch_mcp_prompts( mcp_endpoint: Optional[str] = None, transport_type: MCPTransportType = MCPTransportType.HTTP_STREAM, *, client_session: Optional[ClientSession] = None, event_loop: Optional[asyncio.AbstractEventLoop] = None, working_directory: Optional[str] = None, ) -> List[Type[BasePrompt]]: """ Fetch prompt classes from an MCP server (sync). """ factory = MCPFactory(mcp_endpoint, transport_type, client_session, event_loop, working_directory) return factory.create_prompts() async def fetch_mcp_prompts_async( mcp_endpoint: Optional[str] = None, transport_type: MCPTransportType = MCPTransportType.HTTP_STREAM, *, client_session: Optional[ClientSession] = None, working_directory: Optional[str] = None, ) -> List[Type[BasePrompt]]: """ Async version of fetch_mcp_prompts. Call from within an event loop. """ if client_session is not None: prompt_defs = await MCPDefinitionService.fetch_prompt_definitions_from_session(client_session) factory = MCPFactory(mcp_endpoint, transport_type, client_session, asyncio.get_running_loop(), working_directory) else: service = MCPDefinitionService(mcp_endpoint, transport_type, working_directory) prompt_defs = await service.fetch_prompt_definitions() factory = MCPFactory(mcp_endpoint, transport_type, None, None, working_directory) return factory._create_prompt_classes(prompt_defs) ``` ### File: atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py ```python """Module for transforming JSON schemas to Pydantic models.""" import logging from typing import Any, Dict, List, Optional, Type, Tuple, Literal, Union, Annotated, cast from atomic_agents.connectors.mcp.mcp_definition_service import MCPAttributeType from jsonschema import Draft202012Validator, validators from pydantic import BeforeValidator, Field, TypeAdapter, WithJsonSchema, create_model from pydantic_core import SchemaError from referencing import Registry, Resource from referencing.exceptions import Unresolvable from referencing.jsonschema import DRAFT202012 from atomic_agents.base.base_io_schema import BaseIOSchema logger = logging.getLogger(__name__) def _reject_false_schema(value: Any) -> Any: raise ValueError("Value is forbidden by a false JSON Schema") _FalseSchemaValue = Annotated[Any, BeforeValidator(_reject_false_schema), WithJsonSchema({"not": {}})] def _validate_ref(validator: Any, reference: str, value: Any, schema: Any) -> Any: try: yield from Draft202012Validator.VALIDATORS["$ref"](validator, reference, value, schema) except Unresolvable: # Match the transformer's Any fallback only for the unresolved reference. return _ScalarChoiceValidator = validators.extend(Draft202012Validator, {"$ref": _validate_ref}) # JSON type mapping JSON_TYPE_MAP = { "string": str, "number": float, "integer": int, "boolean": bool, "array": list, "object": dict, "null": type(None), } _NUMERIC_CONSTRAINTS = { "minimum": "ge", "maximum": "le", "exclusiveMinimum": "gt", "exclusiveMaximum": "lt", "multipleOf": "multiple_of", } _STRING_CONSTRAINTS = {"minLength": "min_length", "maxLength": "max_length"} _ARRAY_CONSTRAINTS = {"minItems": "min_length", "maxItems": "max_length"} def _constrained_value_type(python_type: Type, json_type: str, prop_schema: Dict[str, Any]) -> Type: """Attach supported constraints without changing the regex engine or accepting boolean bounds.""" constraints = {} if json_type in {"integer", "number"}: for source, target in _NUMERIC_CONSTRAINTS.items(): value = prop_schema.get(source) # Draft-04 boolean exclusive bounds are not numeric bounds. if isinstance(value, (int, float)) and not isinstance(value, bool): constraints[target] = value else: keywords = _STRING_CONSTRAINTS if json_type == "string" else _ARRAY_CONSTRAINTS if json_type == "array" else {} constraints.update({target: prop_schema[source] for source, target in keywords.items() if source in prop_schema}) if json_type == "string" and isinstance(prop_schema.get("pattern"), str): pattern = prop_schema["pattern"] try: TypeAdapter(Annotated[str, Field(pattern=pattern)]) except SchemaError: logger.warning("Unsupported JSON Schema pattern %r; skipping local pattern validation", pattern) else: constraints["pattern"] = pattern if constraints: # Recursive callers retain the value type but discard returned Field metadata. return Annotated[python_type, Field(**constraints)] return python_type class SchemaTransformer: """Class for transforming JSON schemas to Pydantic models.""" @staticmethod def _matches_scalar_type(value: Any, json_types: List[str]) -> bool: """Match JSON scalar types without treating Python bools as numbers.""" if isinstance(value, bool): return "boolean" in json_types if isinstance(value, (int, float)): return "number" in json_types or ("integer" in json_types and (isinstance(value, int) or value.is_integer())) return (value is None and "null" in json_types) or (isinstance(value, str) and "string" in json_types) @staticmethod def _scalar_choice_type(prop_schema: Dict[str, Any], choices: List[Any], root_schema: Dict[str, Any]) -> Type: """Keep only choices satisfying all advertised constraints, including compositions.""" declared_type = prop_schema.get("type") json_types = declared_type if isinstance(declared_type, list) else [declared_type] # The finite set can be intersected once, without coercing values through union types. # Bind local references to the original root, and never retrieve remote schemas. validator = _ScalarChoiceValidator(root_schema, registry=Registry()).evolve(schema=prop_schema) choices = [value for value in choices if validator.is_valid(value)] if not choices: raise ValueError("Scalar choices do not match the declared JSON Schema constraints") def validate_choice(value: Any) -> Any: if declared_type is not None and not SchemaTransformer._matches_scalar_type(value, json_types): raise ValueError("Value does not match the declared JSON Schema type") for choice in choices: if isinstance(value, bool) == isinstance(choice, bool) and value == choice: # Give Literal the matching choice so 1.0 cannot resolve to True. return choice raise ValueError("Value does not match the JSON Schema choices") choice_type = Annotated[Literal[tuple(choices)], BeforeValidator(validate_choice)] if declared_type is not None: # Keep type metadata inside the choice branch when Optional wraps it. choice_type = Annotated[choice_type, Field(json_schema_extra={"type": declared_type})] return choice_type @staticmethod def _resolve_ref(ref_path: str, root_schema: Dict[str, Any], model_cache: Dict[str, Type]) -> Type: """Resolve a $ref to its Pydantic-compatible value type.""" if ref_path in model_cache: return model_cache[ref_path] # Resolve the complete local pointer, including escaped tokens and array # indices. An empty registry never retrieves remote schemas. resource = Resource.from_contents(root_schema, default_specification=DRAFT202012) try: ref_schema = Registry().resolver_with_root(resource).lookup(ref_path).contents except Unresolvable: logger.warning(f"Could not resolve $ref: {ref_path}") return Any # Avoid infinite recursion by adding a placeholder before conversion. model_cache[ref_path] = Any if isinstance(ref_schema, dict) and ( ref_schema.get("properties") or (ref_schema.get("type") == "object" and ref_schema.get("additionalProperties") is False) ): ref_name = ref_path.split("/")[-1].replace("~1", "/").replace("~0", "~") model_name = ref_schema.get("title", ref_name) model = SchemaTransformer._create_nested_model(ref_schema, model_name, root_schema, model_cache) else: model, _ = SchemaTransformer.json_to_pydantic_field(ref_schema, True, root_schema, model_cache) model_cache[ref_path] = model return model @staticmethod def _create_nested_model( schema: Dict[str, Any], model_name: str, root_schema: Dict[str, Any], model_cache: Dict[str, Type] ) -> Type: """Create a nested Pydantic model from a schema.""" fields = {} required_fields = set(schema.get("required", [])) properties = schema.get("properties", {}) for prop_name, prop_schema in properties.items(): is_required = prop_name in required_fields fields[prop_name] = SchemaTransformer.json_to_pydantic_field(prop_schema, is_required, root_schema, model_cache) return create_model(model_name, **fields) @staticmethod def json_to_pydantic_field( prop_schema: Union[Dict[str, Any], bool], required: bool, root_schema: Optional[Dict[str, Any]] = None, model_cache: Optional[Dict[str, Type]] = None, ) -> Tuple[Type, Field]: """ Convert a JSON schema property to a Pydantic field. Args: prop_schema: JSON schema for the property required: Whether the field is required root_schema: Full root schema for resolving $refs model_cache: Cache for resolved models Returns: Tuple of (type, Field) """ if isinstance(prop_schema, bool): value_type = Any if prop_schema else _FalseSchemaValue return value_type, Field(default=... if required else None) if root_schema is None: root_schema = {} if model_cache is None: model_cache = {} description = prop_schema.get("description") default = prop_schema.get("default") python_type: Any = Any # Handle $ref if "$ref" in prop_schema: python_type = SchemaTransformer._resolve_ref(prop_schema["$ref"], root_schema, model_cache) # Scalar value constraints must also appear in the model-facing schema. elif "const" in prop_schema and isinstance(prop_schema["const"], (str, int, float, bool, type(None))): python_type = SchemaTransformer._scalar_choice_type(prop_schema, [prop_schema["const"]], root_schema) elif prop_schema.get("enum") and all( isinstance(value, (str, int, float, bool, type(None))) for value in prop_schema["enum"] ): python_type = SchemaTransformer._scalar_choice_type(prop_schema, prop_schema["enum"], root_schema) # Handle oneOf/anyOf (unions) elif "oneOf" in prop_schema or "anyOf" in prop_schema: union_schemas = prop_schema.get("oneOf", prop_schema.get("anyOf", [])) if union_schemas: union_types = [] for union_schema in union_schemas: member_type, _ = SchemaTransformer.json_to_pydantic_field(union_schema, True, root_schema, model_cache) union_types.append(member_type) if len(union_types) == 1: python_type = union_types[0] else: python_type = Union[tuple(union_types)] # A JSON Schema type array is a union of its listed types. elif isinstance(prop_schema.get("type"), list): union_types = [ SchemaTransformer.json_to_pydantic_field({**prop_schema, "type": json_type}, True, root_schema, model_cache)[0] for json_type in prop_schema["type"] ] python_type = Union[tuple(union_types)] # Handle regular types else: json_type = prop_schema.get("type") if json_type in JSON_TYPE_MAP: python_type = JSON_TYPE_MAP[json_type] if json_type == "array": items_schema = prop_schema.get("items", {}) item_type, _ = SchemaTransformer.json_to_pydantic_field(items_schema, True, root_schema, model_cache) python_type = List[item_type] elif json_type == "object": python_type = Dict[str, Any] value_schema = prop_schema.get("additionalProperties", {}) if isinstance(value_schema, dict): value_type, _ = SchemaTransformer.json_to_pydantic_field(value_schema, True, root_schema, model_cache) python_type = Dict[str, value_type] python_type = _constrained_value_type(python_type, json_type, prop_schema) field_kwargs = {"description": description} if required: field_kwargs["default"] = ... elif default is not None: field_kwargs["default"] = default else: if python_type is not _FalseSchemaValue: python_type = Optional[python_type] field_kwargs["default"] = None return (python_type, Field(**field_kwargs)) @staticmethod def create_model_from_schema( schema: Dict[str, Any], model_name: str, tool_name_literal: str, docstring: Optional[str] = None, attribute_type: str = MCPAttributeType.TOOL, is_output_schema: bool = False, ) -> Type[BaseIOSchema]: """ Dynamically create a Pydantic model from a JSON schema. Args: schema: JSON schema model_name: Name for the model tool_name_literal: Tool name to use for the Literal type docstring: Optional docstring for the model attribute_type: Type of MCP attribute (tool, resource, prompt) is_output_schema: If True, skip adding the tool_name/resource_name/prompt_name literal field. Output schemas represent tool responses and don't need an identifier field since the tool has already been selected and executed. Input schemas need the identifier for discriminated unions when selecting among multiple tools in an orchestrator. Returns: Pydantic model class """ fields = {} required_fields = set(schema.get("required", [])) properties = schema.get("properties") model_cache: Dict[str, Type] = {} if properties: for prop_name, prop_schema in properties.items(): is_required = prop_name in required_fields fields[prop_name] = SchemaTransformer.json_to_pydantic_field(prop_schema, is_required, schema, model_cache) elif schema.get("type") == "object" and not properties: pass elif schema: logger.warning( f"Schema for {model_name} is not a typical object with properties. Fields might be empty beyond tool_name." ) # Only add the attribute identifier field for input schemas if not is_output_schema: tool_name_type = cast(Type[str], Literal[tool_name_literal]) fields[f"{attribute_type}_name"] = ( tool_name_type, Field(..., description=f"Required identifier for the {tool_name_literal} {attribute_type}."), ) # Create the model model = create_model( model_name, __base__=BaseIOSchema, __doc__=docstring or f"Dynamically generated Pydantic model for {model_name}", __config__={"title": tool_name_literal}, **fields, ) return model ``` ### File: atomic-agents/atomic_agents/context/__init__.py ```python from .base_chat_history import BaseChatHistory from .chat_history import Message, ChatHistory from .system_prompt_generator import ( BaseDynamicContextProvider, SystemPromptGenerator, BaseSystemPromptGenerator, ) __all__ = [ "Message", "BaseChatHistory", "ChatHistory", "SystemPromptGenerator", "BaseDynamicContextProvider", "BaseSystemPromptGenerator", ] ``` ### File: atomic-agents/atomic_agents/context/base_chat_history.py ```python from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Dict, List, Optional from atomic_agents.base.base_io_schema import BaseIOSchema if TYPE_CHECKING: from atomic_agents.context.chat_history import Message class BaseChatHistory(ABC): """ Declares the public contract of a chat history. This is the full method surface of the built-in ``ChatHistory``, of which ``AtomicAgent`` itself uses a subset (``add_message``, ``get_history``, ``initialize_turn``, ``delete_turn_id``, ``copy``, plus direct reads of the ``history`` / ``current_turn_id`` attributes). The remaining methods (``dump`` / ``load`` / ``get_message_count`` / ``get_current_turn_id``) round out the standard history surface that tooling and callers rely on. This is an interface-only abstract base class: it defines no state and no behavior, only the methods a chat history implementation must provide. ``ChatHistory`` is the framework's built-in, in-memory implementation. Implement this interface directly (or subclass ``ChatHistory``) to plug in a custom persistent or backend-specific memory store. Every implementation must maintain the following two attributes, since callers (including ``AtomicAgent._trim_context``) read them directly: Attributes: history (List[Message]): The in-memory working list of messages that make up the chat history. ``AtomicAgent._trim_context`` reads this list directly (via ``msg.turn_id``) to decide which turns to trim when a context-length limit is exceeded. current_turn_id (Optional[str]): The ID of the current turn, or ``None`` if no turn has been started yet. Set by ``initialize_turn()``. """ history: List["Message"] current_turn_id: Optional[str] @abstractmethod def initialize_turn(self) -> None: """ Starts a new turn, generating and storing a new turn ID. Implementations must set ``self.current_turn_id`` to a new, unique value each time this is called. """ raise NotImplementedError @abstractmethod def add_message(self, role: str, content: BaseIOSchema) -> None: """ Adds a message to the chat history, tagging it with the current turn ID. If no turn has been started yet (``current_turn_id`` is ``None``), implementations should start one (e.g. by calling ``initialize_turn()``) before recording the message. Args: role (str): The role of the message sender (e.g. 'user', 'system', 'assistant'). content (BaseIOSchema): The content of the message. """ raise NotImplementedError @abstractmethod def get_history(self) -> List[Dict]: """ Retrieves the chat history in the wire format expected by the LLM client. Returns: List[Dict]: The list of messages as dictionaries with 'role' and 'content' keys, suitable for passing to a chat completion call. """ raise NotImplementedError @abstractmethod def get_current_turn_id(self) -> Optional[str]: """ Returns the current turn ID. Returns: Optional[str]: The current turn ID, or ``None`` if no turn has been started. """ raise NotImplementedError @abstractmethod def delete_turn_id(self, turn_id: str): """ Deletes all messages belonging to the given turn ID. Args: turn_id (str): The turn ID whose messages should be removed. Raises: ValueError: If the specified turn ID is not found in the history. """ raise NotImplementedError @abstractmethod def get_message_count(self) -> int: """ Returns the number of messages currently in the chat history. Returns: int: The number of messages. """ raise NotImplementedError @abstractmethod def dump(self) -> str: """ Serializes the entire chat history to a JSON string. The returned string must be a valid argument to ``load()``, such that ``instance.load(instance.dump())`` round-trips the full history state. Returns: str: A JSON string representation of the chat history. """ raise NotImplementedError @abstractmethod def load(self, serialized_data: str) -> None: """ Deserializes a JSON string produced by ``dump()`` and loads it into this instance, replacing any existing history state. Args: serialized_data (str): A JSON string representation of the chat history, as produced by ``dump()``. Raises: ValueError: If the serialized data is invalid or cannot be deserialized. """ raise NotImplementedError @abstractmethod def copy(self) -> "BaseChatHistory": """ Creates a copy of the chat history, independent of the original. ``AtomicAgent`` calls ``copy()`` to snapshot ``initial_history`` at construction and to restore it in ``reset_history()``. The built-in ``ChatHistory.copy()`` returns a plain ``ChatHistory``; a subclass that carries extra state (a database handle, a ``session_id``, a store reference) MUST override ``copy()`` to return its own type, otherwise ``reset_history()`` silently replaces the backend with a plain in-memory history and later writes stop reaching the store. Returns: BaseChatHistory: A copy of the chat history. """ raise NotImplementedError ``` ### File: atomic-agents/atomic_agents/context/chat_history.py ```python import json import uuid from copy import deepcopy from enum import Enum from pathlib import Path from typing import Dict, List, Optional, Tuple, Type, Union from instructor.processing.multimodal import PDF, Image, Audio from pydantic import BaseModel, Field from atomic_agents.base.base_io_schema import BaseIOSchema from atomic_agents.base.multimodal import VideoURL from atomic_agents.context.base_chat_history import BaseChatHistory MULTIMODAL_TYPES = (Image, Audio, PDF, VideoURL) class Message(BaseModel): """ Represents a message in the chat history. Attributes: role (str): The role of the message sender (e.g., 'user', 'system', 'tool'). content (BaseIOSchema): The content of the message. turn_id (Optional[str]): Unique identifier for the turn this message belongs to. """ role: str content: BaseIOSchema turn_id: Optional[str] = None class ChatHistory(BaseChatHistory): """ Manages the chat history for an AI agent. Attributes: history (List[Message]): A list of messages representing the chat history. max_messages (Optional[int]): Maximum number of messages to keep in history. current_turn_id (Optional[str]): The ID of the current turn. """ def __init__(self, max_messages: Optional[int] = None): """ Initializes the ChatHistory with an empty history and optional constraints. Args: max_messages (Optional[int]): Maximum number of messages to keep in history. When exceeded, oldest messages are removed first. """ self.history: List[Message] = [] self.max_messages = self._validate_max_messages(max_messages) self.current_turn_id: Optional[str] = None @staticmethod def _validate_max_messages(value: Optional[int]) -> Optional[int]: if value is not None and (not isinstance(value, int) or isinstance(value, bool) or value < 0): raise ValueError("max_messages must be a non-negative integer or None") return value def initialize_turn(self) -> None: """ Initializes a new turn by generating a random turn ID. """ self.current_turn_id = str(uuid.uuid4()) def add_message( self, role: str, content: BaseIOSchema, ) -> None: """ Adds a message to the chat history and manages overflow. Args: role (str): The role of the message sender. content (BaseIOSchema): The content of the message. """ if self.current_turn_id is None: self.initialize_turn() message = Message( role=role, content=content, turn_id=self.current_turn_id, ) self.history.append(message) self._manage_overflow() def _manage_overflow(self) -> None: """ Manages the chat history overflow based on max_messages constraint. """ if self.max_messages is not None: while len(self.history) > self.max_messages: self.history.pop(0) def get_history(self) -> List[Dict]: """ Retrieves the chat history, handling both regular and multimodal content. Returns: List[Dict]: The list of messages in the chat history as dictionaries. Each dictionary has 'role' and 'content' keys, where 'content' contains either a single JSON string or a mixed array of JSON and multimodal objects. Note: This method supports multimodal content at any nesting depth by recursively extracting multimodal objects and using Pydantic's model_dump_json(exclude=...) for proper serialization of remaining fields. """ history = [] for message in self.history: input_content = message.content multimodal_objects, exclude_spec = self._extract_multimodal_info(input_content) if multimodal_objects: processed_content = [] # Message.content is a BaseIOSchema, so extraction never excludes the whole object. assert not isinstance(exclude_spec, bool) content_json = input_content.model_dump_json(exclude=exclude_spec) if content_json and content_json != "{}": processed_content.append(content_json) # Instructor forwards dict content parts to the provider unchanged but # raises on object types it does not know, so VideoURL is converted here. processed_content.extend( multimodal_object.to_openai() if isinstance(multimodal_object, VideoURL) else multimodal_object for multimodal_object in multimodal_objects ) history.append({"role": message.role, "content": processed_content}) else: content_json = input_content.model_dump_json() history.append({"role": message.role, "content": content_json}) return history @staticmethod def _extract_multimodal_info(obj) -> Tuple[List[Union[Image, Audio, PDF, VideoURL]], Union[Dict, bool, None]]: """ Recursively extract multimodal objects and build a Pydantic-compatible exclude spec. Walks the object tree to find all multimodal types (Instructor's Image, Audio, and PDF, plus VideoURL) at any nesting depth, collecting them into a flat list and building an exclude specification that can be passed to model_dump_json(exclude=...). Args: obj: The object to inspect (BaseIOSchema, list, dict, or primitive). Returns: tuple: (multimodal_objects, exclude_spec) where: - multimodal_objects: flat list of all multimodal objects found - exclude_spec: Pydantic exclude dict, True (exclude entirely), or None """ if isinstance(obj, MULTIMODAL_TYPES): return [obj], True if hasattr(obj, "__class__") and hasattr(obj.__class__, "model_fields"): all_objects = [] exclude = {} for field_name, field_info in obj.__class__.model_fields.items(): if field_info.exclude: continue if hasattr(obj, field_name): field_value = getattr(obj, field_name) objects, sub_exclude = ChatHistory._extract_multimodal_info(field_value) if objects: all_objects.extend(objects) exclude[field_name] = sub_exclude return all_objects, (exclude if exclude else None) if isinstance(obj, (list, tuple)): all_objects = [] exclude = {} for i, item in enumerate(obj): objects, sub_exclude = ChatHistory._extract_multimodal_info(item) if objects: all_objects.extend(objects) exclude[i] = sub_exclude if not all_objects: return [], None # If every item in the list is fully multimodal, exclude the entire field if len(exclude) == len(obj) and all(v is True for v in exclude.values()): return all_objects, True return all_objects, exclude if isinstance(obj, dict): all_objects = [] exclude = {} for k, v in obj.items(): objects, sub_exclude = ChatHistory._extract_multimodal_info(v) if objects: all_objects.extend(objects) exclude[k] = sub_exclude if not all_objects: return [], None # If every value in the dict is fully multimodal, exclude the entire field if len(exclude) == len(obj) and all(v is True for v in exclude.values()): return all_objects, True return all_objects, exclude return [], None def copy(self) -> "ChatHistory": """ Creates a copy of the chat history. Returns: ChatHistory: A copy of the chat history. """ new_history = ChatHistory(max_messages=self.max_messages) new_history.history = deepcopy(self.history) new_history.current_turn_id = self.current_turn_id return new_history def get_current_turn_id(self) -> Optional[str]: """ Returns the current turn ID. Returns: Optional[str]: The current turn ID, or None if not set. """ return self.current_turn_id def delete_turn_id(self, turn_id: str): """ Delete messages from the history by its turn ID. Args: turn_id (str): The turn ID of the message to delete. Returns: str: A success message with the deleted turn ID. Raises: ValueError: If the specified turn ID is not found in the history. """ initial_length = len(self.history) self.history = [msg for msg in self.history if msg.turn_id != turn_id] if len(self.history) == initial_length: raise ValueError(f"Turn ID {turn_id} not found in history.") # Update current_turn_id if necessary if not len(self.history): self.current_turn_id = None elif turn_id == self.current_turn_id: # Always update to the last message's turn_id self.current_turn_id = self.history[-1].turn_id def get_message_count(self) -> int: """ Returns the number of messages in the chat history. Returns: int: The number of messages. """ return len(self.history) def dump(self) -> str: """ Serializes the entire ChatHistory instance to a JSON string. Returns: str: A JSON string representation of the ChatHistory. """ serialized_history = [] for message in self.history: content_class = message.content.__class__ serialized_message = { "role": message.role, "content": { "class_name": f"{content_class.__module__}.{content_class.__name__}", "data": message.content.model_dump_json(), }, "turn_id": message.turn_id, } serialized_history.append(serialized_message) history_data = { "history": serialized_history, "max_messages": self.max_messages, "current_turn_id": self.current_turn_id, } return json.dumps(history_data) def load(self, serialized_data: str) -> None: """ Deserializes a JSON string and loads it into the ChatHistory instance. If loading fails, the existing history and metadata remain unchanged. Args: serialized_data (str): A JSON string representation of the ChatHistory. Raises: ValueError: If the serialized data is invalid or cannot be deserialized. """ try: history_data = json.loads(serialized_data) loaded_history: List[Message] = [] try: max_messages = self._validate_max_messages(history_data["max_messages"]) except ValueError as error: raise ValueError(f"Invalid serialized data: {error}") from error current_turn_id = history_data["current_turn_id"] if current_turn_id is not None and not isinstance(current_turn_id, str): raise ValueError("Invalid serialized data: current_turn_id must be a string or null") messages = history_data["history"] if not isinstance(messages, list): raise ValueError("Invalid serialized data: history must be a list") for message_data in messages: content_info = message_data["content"] content_class = self._get_class_from_string(content_info["class_name"]) content_instance = content_class.model_validate_json(content_info["data"]) # Process any Image objects to convert string paths back to Path objects self._process_multimodal_paths(content_instance) message = Message(role=message_data["role"], content=content_instance, turn_id=message_data["turn_id"]) loaded_history.append(message) # Replace the state only after every message has been deserialized successfully. self.history = loaded_history self.max_messages = max_messages self.current_turn_id = current_turn_id self._manage_overflow() except (json.JSONDecodeError, KeyError, AttributeError, TypeError) as e: raise ValueError(f"Invalid serialized data: {e}") @staticmethod def _get_class_from_string(class_string: str) -> Type[BaseIOSchema]: """ Retrieves a class object from its string representation. Args: class_string (str): The fully qualified class name. Returns: Type[BaseIOSchema]: The class object. Raises: AttributeError: If the class cannot be found. """ module_name, class_name = class_string.rsplit(".", 1) module = __import__(module_name, fromlist=[class_name]) content_class = getattr(module, class_name) if not isinstance(content_class, type) or not issubclass(content_class, BaseIOSchema): raise ValueError(f"Invalid serialized data: {class_string} is not a BaseIOSchema subclass") return content_class def _process_multimodal_paths(self, obj): """ Process multimodal objects to convert string paths to Path objects. Note: this is necessary only for PDF and Image instructor types. The from_path behavior is slightly different for Audio as it keeps the source as a string. Args: obj: The object to process. """ if isinstance(obj, (Image, PDF)) and isinstance(obj.source, str): # Check if the string looks like a file path (not a URL or base64 data) if not obj.source.startswith(("http://", "https://", "data:")): obj.source = Path(obj.source) elif isinstance(obj, (list, tuple)): # Process each item in the sequence for item in obj: self._process_multimodal_paths(item) elif isinstance(obj, dict): # Process each value in the dictionary for value in obj.values(): self._process_multimodal_paths(value) elif hasattr(obj, "__class__") and hasattr(obj.__class__, "model_fields"): # Process each field of the Pydantic model for field_name in obj.__class__.model_fields: if hasattr(obj, field_name): self._process_multimodal_paths(getattr(obj, field_name)) elif hasattr(obj, "__dict__") and not isinstance(obj, Enum): # Process each attribute of the object for attr_name, attr_value in obj.__dict__.items(): if attr_name != "__pydantic_fields_set__": # Skip pydantic internal fields self._process_multimodal_paths(attr_value) if __name__ == "__main__": import instructor from typing import List as TypeList, Dict as TypeDict import os # Define complex test schemas class NestedSchema(BaseIOSchema): """A nested schema for testing""" nested_field: str = Field(..., description="A nested field") nested_int: int = Field(..., description="A nested integer") class ComplexInputSchema(BaseIOSchema): """Complex Input Schema""" text_field: str = Field(..., description="A text field") number_field: float = Field(..., description="A number field") list_field: TypeList[str] = Field(..., description="A list of strings") nested_field: NestedSchema = Field(..., description="A nested schema") class ComplexOutputSchema(BaseIOSchema): """Complex Output Schema""" response_text: str = Field(..., description="A response text") calculated_value: int = Field(..., description="A calculated value") data_dict: TypeDict[str, NestedSchema] = Field(..., description="A dictionary of nested schemas") # Add a new multimodal schema for testing class MultimodalSchema(BaseIOSchema): """Schema for testing multimodal content""" instruction_text: str = Field(..., description="The instruction text") images: List[instructor.Image] = Field(..., description="The images to analyze") # Create and populate the original history with complex data original_history = ChatHistory(max_messages=10) # Add a complex input message original_history.add_message( "user", ComplexInputSchema( text_field="Hello, this is a complex input", number_field=3.14159, list_field=["item1", "item2", "item3"], nested_field=NestedSchema(nested_field="Nested input", nested_int=42), ), ) # Add a complex output message original_history.add_message( "assistant", ComplexOutputSchema( response_text="This is a complex response", calculated_value=100, data_dict={ "key1": NestedSchema(nested_field="Nested output 1", nested_int=10), "key2": NestedSchema(nested_field="Nested output 2", nested_int=20), }, ), ) # Test multimodal functionality if test image exists test_image_path = os.path.join("test_images", "test.jpg") if os.path.exists(test_image_path): # Add a multimodal message original_history.add_message( "user", MultimodalSchema( instruction_text="Please analyze this image", images=[instructor.Image.from_path(test_image_path)] ), ) # Continue with existing tests... dumped_data = original_history.dump() print("Dumped data:") print(dumped_data) # Create a new history and load the dumped data loaded_history = ChatHistory() loaded_history.load(dumped_data) # Print detailed information about the loaded history print("\nLoaded history details:") for i, message in enumerate(loaded_history.history): print(f"\nMessage {i + 1}:") print(f"Role: {message.role}") print(f"Turn ID: {message.turn_id}") print(f"Content type: {type(message.content).__name__}") print("Content:") for field, value in message.content.model_dump().items(): print(f" {field}: {value}") # Final verification print("\nFinal verification:") print(f"Max messages: {loaded_history.max_messages}") print(f"Current turn ID: {loaded_history.get_current_turn_id()}") print("Last message content:") last_message = loaded_history.history[-1] print(last_message.content.model_dump()) ``` ### File: atomic-agents/atomic_agents/context/system_prompt_generator.py ```python from abc import ABC, abstractmethod from typing import Dict, List, Optional class BaseDynamicContextProvider(ABC): def __init__(self, title: str): self.title = title @abstractmethod def get_info(self) -> str: pass def __repr__(self) -> str: return self.get_info() class BaseSystemPromptGenerator(ABC): def __init__(self, context_providers: Optional[Dict[str, BaseDynamicContextProvider]] = None): self.context_providers = context_providers or {} @abstractmethod def generate_prompt(self) -> str: pass def __repr__(self) -> str: return f"{self.__class__.__name__} (providers={list(self.context_providers)})" class SystemPromptGenerator(BaseSystemPromptGenerator): def __init__( self, background: Optional[List[str]] = None, steps: Optional[List[str]] = None, output_instructions: Optional[List[str]] = None, context_providers: Optional[Dict[str, BaseDynamicContextProvider]] = None, ): super().__init__(context_providers=context_providers) self.background = background or ["This is a conversation with a helpful and friendly AI assistant."] self.steps = steps or [] self.output_instructions = list(output_instructions) if output_instructions else [] self.output_instructions.extend( [ "Always respond using the proper JSON schema.", "Always use the available additional information and context to enhance the response.", ] ) def generate_prompt(self) -> str: sections = [ ("IDENTITY and PURPOSE", self.background), ("INTERNAL ASSISTANT STEPS", self.steps), ("OUTPUT INSTRUCTIONS", self.output_instructions), ] prompt_parts = [] for title, content in sections: if content: prompt_parts.append(f"# {title}") prompt_parts.extend(f"- {item}" for item in content) prompt_parts.append("") if self.context_providers: prompt_parts.append("# EXTRA INFORMATION AND CONTEXT") for provider in self.context_providers.values(): info = provider.get_info() if info: prompt_parts.append(f"## {provider.title}") prompt_parts.append(info) prompt_parts.append("") return "\n".join(prompt_parts).strip() ``` ### File: atomic-agents/atomic_agents/observability/__init__.py ```python """Opt-in tracing for any OpenTelemetry backend using GenAI attributes. Install the observability extra for tracing. The optional phoenix extra adds Arize Phoenix integration and OpenInference attributes. """ import importlib.util import logging from importlib.metadata import PackageNotFoundError, version from typing import TYPE_CHECKING if TYPE_CHECKING: from atomic_agents.observability.tracing import PayloadVisibility, TracingConfig, using_conversation from opentelemetry.trace import TracerProvider logger = logging.getLogger(__name__) _INSTALL_HINT = ( "Atomic Agents tracing needs opentelemetry-api. " 'Install it with: pip install "atomic-agents[observability]" ' '(add "atomic-agents[phoenix]" for Arize Phoenix / OpenInference attributes)' ) _OPENINFERENCE_DISTRIBUTION = "openinference-instrumentation" _MINIMUM_OPENINFERENCE_VERSION = (0, 1, 38) def _opentelemetry_is_installed() -> bool: try: return importlib.util.find_spec("opentelemetry.trace") is not None except ModuleNotFoundError: return False def _openinference_is_too_old() -> bool: try: installed = version(_OPENINFERENCE_DISTRIBUTION) except PackageNotFoundError: return True release = tuple(int(part) for part in installed.split(".")[:3] if part.isdigit()) return release < _MINIMUM_OPENINFERENCE_VERSION def _load_tracing_module(): from atomic_agents.observability import tracing return tracing def __getattr__(name: str) -> object: if name in ("TracingConfig", "using_conversation"): return getattr(_load_tracing_module(), name) raise AttributeError(name) def _warn_unknown_options(options: dict[str, object]) -> None: for key in options: if key != "skip_dep_check": logger.warning("Unknown Atomic Agents instrumentation option %s ignored", key) class AtomicAgentsInstrumentor: """ Installs AGENT, TOOL and LLM-attempt spans around ``AtomicAgent`` and ``BaseTool`` calls. The public surface mirrors OpenTelemetry's ``BaseInstrumentor`` (``instrument`` / ``uninstrument``) without depending on it, because Arize Phoenix only guarantees ``opentelemetry-api`` and ``openinference-instrumentation`` in an auto-instrumented environment. """ def instrument( self, *, tracer_provider: "TracerProvider | None" = None, config: "TracingConfig | PayloadVisibility | None" = None, **options: object, ) -> None: """ Patch the framework classes. Safe to call more than once; a second call with different arguments replaces the tracer provider/config and logs a warning. Args: tracer_provider: An OpenTelemetry ``TracerProvider``. Defaults to the global provider. config: Payload visibility and truncation settings. OpenInference TraceConfig is also accepted. **options: Launcher options accepted for OpenTelemetry auto-instrumentation compatibility. Raises: ModuleNotFoundError: When the ``observability`` extra is not installed, except during OpenTelemetry auto-discovery, which skips unavailable optional instrumentation. """ _warn_unknown_options(options) if not _opentelemetry_is_installed(): if options.get("skip_dep_check") is True: logger.debug(_INSTALL_HINT) return raise ModuleNotFoundError(_INSTALL_HINT) _load_tracing_module().install(tracer_provider=tracer_provider, config=config) def uninstrument(self) -> None: """Restore the original framework methods. No-op when nothing is installed.""" if not self.is_instrumented_by_opentelemetry: return _load_tracing_module().uninstall() @property def is_instrumented_by_opentelemetry(self) -> bool: if not _opentelemetry_is_installed(): return False return _load_tracing_module().is_installed() __all__ = ["AtomicAgentsInstrumentor", "TracingConfig", "using_conversation"] ``` ### File: atomic-agents/atomic_agents/observability/agent_spans.py ```python """Agent call and stream spans.""" from __future__ import annotations import asyncio import functools from contextlib import aclosing, contextmanager from typing import AsyncGenerator, Awaitable, Callable, Coroutine, Generator, Generic, Iterator, TypeVar from opentelemetry import context as context_api from opentelemetry import trace as trace_api from opentelemetry.trace import Span, Status, StatusCode from atomic_agents.agents.atomic_agent import AtomicAgent from . import installation, llm_attempt_spans from .tracing_config import TracingConfig from .span_attributes import ATTR_STREAM_CLOSED_EARLY, _agent_display_name, _agent_start_attributes, _suppressed from .span_payloads import _record_failure, _set_input_safely, _set_output_safely OutputSchema = TypeVar("OutputSchema") CallResult = TypeVar("CallResult") AGENT_SYNC_METHOD = "run" AGENT_ASYNC_METHOD = "run_async" AGENT_SYNC_STREAM_METHOD = "run_stream" AGENT_ASYNC_STREAM_METHOD = "run_async_stream" def _user_input_from_call(args: tuple[object, ...], kwargs: dict[str, object]) -> object: return args[0] if args else kwargs.get("user_input") def _finish_agent_span(span: Span, output: object, config: TracingConfig) -> None: _set_output_safely(span, output, config=config) span.set_status(Status(StatusCode.OK)) def _mark_closed_early(span: Span, output: object, config: TracingConfig) -> None: span.set_attribute(ATTR_STREAM_CLOSED_EARLY, True) _set_output_safely(span, output, config=config) @contextmanager def _agent_span( agent: AtomicAgent, args: tuple[object, ...], kwargs: dict[str, object] ) -> Iterator[llm_attempt_spans._AgentRun]: installed = installation._current_installation() with installed.tracer.start_as_current_span( f"invoke_agent {_agent_display_name(agent)}", attributes=_agent_start_attributes(agent), record_exception=False, set_status_on_exception=False, ) as span: _set_input_safely(span, _user_input_from_call(args, kwargs)) run = llm_attempt_spans.new_run(span, agent) try: with llm_attempt_spans.active_run(run): yield run except (asyncio.CancelledError, GeneratorExit): span.set_attribute("atomic_agents.cancelled", True) raise except BaseException as error: _record_failure(span, error, run.config) raise finally: llm_attempt_spans.finish_run(run) def _trace_agent_call(wrapped: Callable[..., OutputSchema]) -> Callable[..., OutputSchema]: @functools.wraps(wrapped) def traced(agent: AtomicAgent, *args: object, **kwargs: object) -> OutputSchema: if not installation.is_installed() or _suppressed(): return wrapped(agent, *args, **kwargs) with _agent_span(agent, args, kwargs) as run: response = wrapped(agent, *args, **kwargs) _finish_agent_span(run.span, response, run.config) return response return traced def _trace_agent_async_call( wrapped: Callable[..., Awaitable[OutputSchema]], ) -> Callable[..., Coroutine[object, object, OutputSchema]]: @functools.wraps(wrapped) async def traced(agent: AtomicAgent, *args: object, **kwargs: object) -> OutputSchema: if not installation.is_installed() or _suppressed(): return await wrapped(agent, *args, **kwargs) with _agent_span(agent, args, kwargs) as run: response = await wrapped(agent, *args, **kwargs) _finish_agent_span(run.span, response, run.config) return response return traced def _start_agent_span(agent: AtomicAgent, args: tuple[object, ...], kwargs: dict[str, object]) -> Span: if _suppressed(): return trace_api.INVALID_SPAN span = installation._current_installation().tracer.start_span( f"invoke_agent {_agent_display_name(agent)}", attributes=_agent_start_attributes(agent), ) _set_input_safely(span, _user_input_from_call(args, kwargs)) return span def _call_within_span(run: llm_attempt_spans._AgentRun, function: Callable[..., CallResult], *args: object) -> CallResult: span = run.span token = context_api.attach(trace_api.set_span_in_context(span)) try: with llm_attempt_spans.active_run(run): return function(*args) finally: context_api.detach(token) async def _await_within_span(run: llm_attempt_spans._AgentRun, awaitable: Awaitable[CallResult]) -> CallResult: span = run.span token = context_api.attach(trace_api.set_span_in_context(span)) try: with llm_attempt_spans.active_run(run): return await awaitable finally: context_api.detach(token) class _SyncStreamCursor(Generic[OutputSchema]): def __init__(self, run: llm_attempt_spans._AgentRun, inner: Generator[OutputSchema, None, OutputSchema | None]): self.run = run self.span = run.span self.inner = inner self.last: OutputSchema self.output: OutputSchema | None = None self.final: OutputSchema | None = None def advance(self) -> bool: try: self.last = _call_within_span(self.run, next, self.inner) self.output = self.last except StopIteration as stop: self.final = stop.value return False return True def close_early(self) -> None: _call_within_span(self.run, self.inner.close) _mark_closed_early(self.span, self.output, self.run.config) self.span.set_status(Status(StatusCode.OK)) class _AsyncStreamCursor(Generic[OutputSchema]): """Async twin of :class:`_SyncStreamCursor`; async generators carry no return value, so the last partial is the output.""" def __init__(self, run: llm_attempt_spans._AgentRun, inner: AsyncGenerator[OutputSchema, None]): self.run = run self.span = run.span self.inner = inner self.last: OutputSchema self.output: OutputSchema | None = None async def advance(self) -> bool: try: self.last = await _await_within_span(self.run, self.inner.__anext__()) self.output = self.last except StopAsyncIteration: return False return True async def close_early(self) -> None: await _await_within_span(self.run, self.inner.aclose()) _mark_closed_early(self.span, self.output, self.run.config) self.span.set_status(Status(StatusCode.OK)) def mark_cancelled(self) -> None: _mark_closed_early(self.span, self.output, self.run.config) def _trace_agent_stream( wrapped: Callable[..., Generator[OutputSchema, None, OutputSchema | None]], ) -> Callable[..., Generator[OutputSchema, None, OutputSchema | None]]: @functools.wraps(wrapped) def traced(agent: AtomicAgent, *args: object, **kwargs: object) -> Generator[OutputSchema, None, OutputSchema | None]: if not installation.is_installed(): return (yield from wrapped(agent, *args, **kwargs)) span = _start_agent_span(agent, args, kwargs) run = llm_attempt_spans.new_run(span, agent) cursor = _SyncStreamCursor(run, wrapped(agent, *args, **kwargs)) try: while cursor.advance(): yield cursor.last _finish_agent_span(span, cursor.final, run.config) return cursor.final except GeneratorExit: cursor.close_early() raise except asyncio.CancelledError: _mark_closed_early(span, cursor.output, run.config) raise except BaseException as error: _record_failure(span, error, run.config) raise finally: llm_attempt_spans.finish_run(run) span.end() return traced def _trace_agent_async_stream( wrapped: Callable[..., AsyncGenerator[OutputSchema, None]], ) -> Callable[..., AsyncGenerator[OutputSchema, None]]: @functools.wraps(wrapped) async def traced(agent: AtomicAgent, *args: object, **kwargs: object) -> AsyncGenerator[OutputSchema, None]: if not installation.is_installed(): async with aclosing(wrapped(agent, *args, **kwargs)) as stream: async for output in stream: yield output return span = _start_agent_span(agent, args, kwargs) run = llm_attempt_spans.new_run(span, agent) cursor = _AsyncStreamCursor(run, wrapped(agent, *args, **kwargs)) try: while await cursor.advance(): yield cursor.last _finish_agent_span(span, cursor.output, run.config) except GeneratorExit: await cursor.close_early() raise except asyncio.CancelledError: cursor.mark_cancelled() raise except BaseException as error: _record_failure(span, error, run.config) raise finally: llm_attempt_spans.finish_run(run) span.end() return traced _AGENT_WRAPPERS = { AGENT_SYNC_METHOD: _trace_agent_call, AGENT_ASYNC_METHOD: _trace_agent_async_call, AGENT_SYNC_STREAM_METHOD: _trace_agent_stream, AGENT_ASYNC_STREAM_METHOD: _trace_agent_async_stream, } ``` ### File: atomic-agents/atomic_agents/observability/installation.py ```python """Tracing installation and method restoration.""" import logging from dataclasses import replace from typing import TYPE_CHECKING, Callable, Optional, Type from opentelemetry import trace as trace_api from opentelemetry.trace import Tracer from opentelemetry.util.types import AttributeValue from atomic_agents import __version__ from atomic_agents.agents.atomic_agent import AtomicAgent from atomic_agents.base.base_tool import BaseTool from . import agent_spans, llm_attempt_spans, span_attributes, tool_spans from .tracing_config import PayloadVisibility, TracingConfig if TYPE_CHECKING: from opentelemetry.trace import TracerProvider logger = logging.getLogger("atomic_agents.observability.tracing") TRACED_MARKER = "__atomic_agents_traced__" INHERITED_MARKER = "__atomic_agents_inherited__" TRACER_NAME = "atomic_agents" def _current_installation() -> "_Installation": assert _installation is not None return _installation class _Installation: def __init__( self, tracer: Tracer, tracer_provider: Optional["TracerProvider"], config: TracingConfig, original_init_subclass: classmethod, context_attributes: Callable[[], dict[str, AttributeValue]] | None, ): self.context_attributes = context_attributes self.tracer = tracer self.tracer_provider = tracer_provider self.config = config self.original_init_subclass = original_init_subclass _installation: Optional[_Installation] = None def is_installed() -> bool: return _installation is not None def install( *, tracer_provider: Optional["TracerProvider"] = None, config: TracingConfig | PayloadVisibility | None = None ) -> None: global _installation if _installation is not None: tracer_provider = _installation.tracer_provider if tracer_provider is None else tracer_provider _replace_installation(_installation, tracer_provider, config) return config = TracingConfig.coerce(config) context_attributes = span_attributes._import_context_attributes() if config.openinference_attributes is None: config = replace(config, openinference_attributes=context_attributes is not None) tracer = trace_api.get_tracer(TRACER_NAME, __version__, tracer_provider) original_init_subclass = BaseTool.__dict__["__init_subclass__"] _installation = _Installation(tracer, tracer_provider, config, original_init_subclass, context_attributes) _instrument_agent_class() tool_spans._instrument_existing_tool_classes() tool_spans._hook_future_tool_classes(original_init_subclass) def uninstall() -> None: global _installation if _installation is None: return setattr(BaseTool, "__init_subclass__", _installation.original_init_subclass) for tool_class in tool_spans._all_subclasses(BaseTool): tool_spans._restore_tool_class(tool_class) for method_name in agent_spans._AGENT_WRAPPERS: _restore_method(AtomicAgent, method_name) llm_attempt_spans.unregister_clients() _installation = None def _replace_installation( installed: _Installation, tracer_provider: Optional["TracerProvider"], config: TracingConfig | PayloadVisibility | None ) -> None: config = installed.config if config is None else TracingConfig.coerce(config) if config.openinference_attributes is None: config = replace(config, openinference_attributes=installed.config.openinference_attributes) if tracer_provider is installed.tracer_provider and config == installed.config: return installed.tracer = trace_api.get_tracer(TRACER_NAME, __version__, tracer_provider) installed.tracer_provider = tracer_provider installed.config = config installed.context_attributes = span_attributes._import_context_attributes() logger.warning( "AtomicAgentsInstrumentor tracer provider/config replaced; " "in-flight spans finish on the provider they started with." ) def _instrument_agent_class() -> None: for method_name, wrap in agent_spans._AGENT_WRAPPERS.items(): original = AtomicAgent.__dict__[method_name] if not getattr(original, TRACED_MARKER, False): setattr(AtomicAgent, method_name, _mark_traced(wrap(original), original, False)) def _mark_traced(wrapper: Callable, original: Callable, inherited: bool) -> Callable: setattr(wrapper, TRACED_MARKER, True) setattr(wrapper, INHERITED_MARKER, inherited) setattr(wrapper, "__wrapped__", original) return wrapper def _restore_method(owner: Type, method_name: str) -> None: current = owner.__dict__.get(method_name) if current is None or not getattr(current, TRACED_MARKER, False): return if getattr(current, INHERITED_MARKER, False): delattr(owner, method_name) else: setattr(owner, method_name, current.__wrapped__) ``` ### File: atomic-agents/atomic_agents/observability/llm_attempt_spans.py ```python """Instructor attempt spans and cumulative agent usage.""" from __future__ import annotations import contextvars import functools import logging import threading import weakref from contextlib import contextmanager from dataclasses import dataclass, field from typing import Callable, Iterator, ParamSpec, Protocol from opentelemetry import context as context_api, trace as trace_api from opentelemetry.trace import Span, SpanKind, Status, StatusCode, Tracer from opentelemetry.util.types import AttributeValue from atomic_agents.agents.atomic_agent import AtomicAgent from . import installation from .message_payloads import write_messages from .span_attributes import _best_effort_attributes, _optional_attributes, _provider_attributes, _without_none from .span_payloads import _bounded_text, _failure_message, _record_failure, _write_payload from .tracing_config import TracingConfig logger = logging.getLogger("atomic_agents.observability.tracing") _PARAMETERS = ParamSpec("_PARAMETERS") _RESPONSE_MODEL = "gen_ai.response.model" _ATTEMPT = "atomic_agents.llm.attempt" _VALIDATION_ERROR = "atomic_agents.llm.last_validation_error" class _PromptTokenDetails(Protocol): cached_tokens: int | None class _CompletionTokenDetails(Protocol): reasoning_tokens: int | None class _InputOutputUsage(Protocol): input_tokens: int output_tokens: int input_tokens_details: _PromptTokenDetails | None output_tokens_details: _CompletionTokenDetails | None class _PromptCompletionUsage(Protocol): prompt_tokens: int completion_tokens: int prompt_tokens_details: _PromptTokenDetails | None completion_tokens_details: _CompletionTokenDetails | None _Usage = _InputOutputUsage | _PromptCompletionUsage class _Response(Protocol): model: str usage: _Usage class _HookClient(Protocol): def on(self, event: str, handler: Callable[..., None]) -> object: """Append a handler without replacing application hooks.""" def off(self, event: str, handler: Callable[..., None]) -> object: """Remove only this handler, preserving application hooks.""" @dataclass class _AgentRun: span: Span tracer: Tracer config: TracingConfig provider: dict[str, AttributeValue] usage_totals: dict[str, int] = field(default_factory=dict) attempt_count: int = 0 suppressed_attempt: bool = False response_received: bool = False open_attempt: Span | None = None context_token: contextvars.Token | None = None _active_run: contextvars.ContextVar[_AgentRun | None] = contextvars.ContextVar("atomic_agents_llm_run", default=None) _clients: weakref.WeakSet[_HookClient] = weakref.WeakSet() _registration_lock = threading.Lock() def _never_raise(handler: Callable[_PARAMETERS, None]) -> Callable[_PARAMETERS, None]: @functools.wraps(handler) def guarded(*args: _PARAMETERS.args, **kwargs: _PARAMETERS.kwargs) -> None: try: handler(*args, **kwargs) except Exception: logger.warning("Instructor tracing hook failed: %s", handler.__name__, exc_info=True) return guarded @_never_raise def register_client(client: _HookClient) -> None: if not hasattr(client, "on"): return # Shared clients must acquire all handlers once before another run can use them. with _registration_lock: if client not in _clients: _register_handlers(client) _clients.add(client) def _register_handlers(client: _HookClient) -> None: registered: list[tuple[str, Callable[..., None]]] = [] try: for event, handler in _HANDLERS.items(): client.on(event, handler) registered.append((event, handler)) except Exception: for event, handler in registered: client.off(event, handler) raise @_never_raise def unregister_clients() -> None: with _registration_lock: for client in list(_clients): _unregister_client(client) _clients.clear() @_never_raise def _unregister_client(client: _HookClient) -> None: for event, handler in _HANDLERS.items(): client.off(event, handler) def new_run(span: Span, agent: AtomicAgent) -> _AgentRun: installed = installation._current_installation() register_client(getattr(agent, "client")) provider = _best_effort_attributes( {"gen_ai.provider.name": lambda: _provider_attributes(agent).get("gen_ai.provider.name")} ) return _AgentRun(span, installed.tracer, installed.config, provider) @contextmanager def active_run(run: _AgentRun) -> Iterator[None]: token = _active_run.set(None if run.span is trace_api.INVALID_SPAN else run) _attach_attempt(run) try: yield finally: _detach_attempt(run) _active_run.reset(token) def _attach_attempt(run: _AgentRun) -> None: if run.open_attempt is not None: run.context_token = context_api.attach(trace_api.set_span_in_context(run.open_attempt)) def _detach_attempt(run: _AgentRun) -> None: if run.context_token is not None: context_api.detach(run.context_token) run.context_token = None @_never_raise def finish_attempt(run: _AgentRun, status: Status | None = None) -> None: if run.open_attempt is None: return if run.response_received: status = Status(StatusCode.OK) run.response_received = False if status is not None: run.open_attempt.set_status(status) run.open_attempt.end() run.open_attempt = None _detach_attempt(run) @_never_raise def _completion_kwargs(*args: object, **kwargs: object) -> None: run = _active_run.get() if run is None: return finish_attempt(run) # Suppressed kwargs must also swallow later response events outside that context. run.suppressed_attempt = bool(context_api.get_value("suppress_instrumentation")) if run.suppressed_attempt: return run.attempt_count += 1 if not run.config.llm_spans: return model = _request_model(kwargs) attributes: dict[str, AttributeValue | None] = { "gen_ai.operation.name": "chat", "gen_ai.request.model": model, _ATTEMPT: run.attempt_count, } attributes.update(run.provider) attributes.update(_optional_attributes("LLM", {})) run.open_attempt = run.tracer.start_span( f"chat {model}" if model else "chat", kind=SpanKind.CLIENT, attributes=_without_none(attributes) ) _attach_attempt(run) write_messages(run.open_attempt, kwargs.get("messages"), "input", run.config) def _request_model(kwargs: dict[str, object]) -> str | None: model = kwargs.get("model") return str(model) if model is not None else None def _usage_attributes(usage: _Usage | None) -> dict[str, int]: readers = { "input_tokens": getattr(usage, "input_tokens", getattr(usage, "prompt_tokens", None)), "output_tokens": getattr(usage, "output_tokens", getattr(usage, "completion_tokens", None)), "cache_read.input_tokens": getattr( getattr(usage, "input_tokens_details", getattr(usage, "prompt_tokens_details", None)), "cached_tokens", None ), "reasoning.output_tokens": getattr( getattr(usage, "output_tokens_details", getattr(usage, "completion_tokens_details", None)), "reasoning_tokens", None, ), } return {"gen_ai.usage." + key: value for key, value in readers.items() if isinstance(value, int)} @_never_raise def _completion_response(response: _Response, **kwargs: object) -> None: run = _active_run.get() if run is None or run.suppressed_attempt: return attributes = _without_none({_RESPONSE_MODEL: getattr(response, "model", None)}) run.span.set_attributes(attributes) usage = _usage_attributes(getattr(response, "usage", None)) for key, value in usage.items(): run.usage_totals[key] = run.usage_totals.get(key, 0) + value _complete_response_attempt(run, usage, attributes) if run.open_attempt is not None: write_messages(run.open_attempt, response, "output", run.config) def _complete_response_attempt(run: _AgentRun, usage: dict[str, int], attributes: dict[str, AttributeValue]) -> None: if run.open_attempt is None: return run.open_attempt.set_attributes(attributes) run.open_attempt.set_attributes(usage) run.response_received = True @_never_raise def _completion_error(error: BaseException, **kwargs: object) -> None: run = _active_run.get() if run is None or run.suppressed_attempt or run.open_attempt is None: return run.response_received = False _record_failure(run.open_attempt, error, run.config) finish_attempt(run) @_never_raise def _parse_error(error: BaseException, **kwargs: object) -> None: run = _active_run.get() if run is None or run.suppressed_attempt: return text = "__REDACTED__" if run.config.hide_outputs else _failure_message(error, run.config) _write_payload(run.span, _VALIDATION_ERROR, None, text, run.config.hide_outputs, run.config) event = {_ATTEMPT: run.attempt_count, _VALIDATION_ERROR: _bounded_text(text, run.config)} run.span.add_event("validation_error", event) _mark_attempt_validation_error(run, error) def _mark_attempt_validation_error(run: _AgentRun, error: BaseException) -> None: if run.open_attempt is None: return run.response_received = False message = _bounded_text(_failure_message(error, run.config), run.config) run.open_attempt.add_event("validation_error", {_ATTEMPT: run.attempt_count, _VALIDATION_ERROR: message}) run.open_attempt.set_attribute("error.type", type(error).__qualname__) run.open_attempt.set_status(Status(StatusCode.ERROR, message)) @_never_raise def finish_run(run: _AgentRun) -> None: finish_attempt(run) run.span.set_attributes(run.usage_totals) run.span.set_attribute("atomic_agents.llm.attempts", run.attempt_count) _HANDLERS: dict[str, Callable[..., None]] = { "completion:kwargs": _completion_kwargs, "completion:response": _completion_response, "completion:error": _completion_error, "parse:error": _parse_error, } ``` ### File: atomic-agents/atomic_agents/observability/message_payloads.py ```python """Bounded, opt-in GenAI chat message attributes.""" import json from opentelemetry import context as context_api from opentelemetry.trace import Span from .tracing_config import TracingConfig def _field(value: object, name: str) -> object: if isinstance(value, dict): return value.get(name) return getattr(value, name, None) def _text_part(content: object) -> dict[str, object] | None: if isinstance(content, str) and content: return {"type": "text", "content": content} return None def _text_block(block: object) -> dict[str, object] | None: if _field(block, "type") != "text": return None return _text_part(_field(block, "text")) def _text_parts(content: object) -> list[dict[str, object]]: if isinstance(content, list): return [part for block in content if (part := _text_block(block)) is not None] part = _text_part(content) return [part] if part is not None else [] def _call_id(part: dict[str, object], value: object, field: str) -> dict[str, object]: identifier = _field(value, field) if isinstance(identifier, str): part["id"] = identifier return part def _reject_non_finite_constant(constant: str) -> object: raise ValueError(f"non-finite JSON constant: {constant}") def _tool_call_arguments(arguments: str) -> object: # Truncated or non-finite arguments are what Instructor rejected, so their raw text stays visible. try: return json.loads(arguments, parse_constant=_reject_non_finite_constant) except ValueError: return arguments def _tool_call(call: object) -> dict[str, object] | None: if _field(call, "type") != "function": return None function = _field(call, "function") name = _field(function, "name") if not isinstance(name, str) or not name: return None part: dict[str, object] = {"type": "tool_call", "name": name} arguments = _field(function, "arguments") if isinstance(arguments, str): part["arguments"] = _tool_call_arguments(arguments) return _call_id(part, call, "id") def _tool_calls(message: object) -> list[dict[str, object]]: calls = _field(message, "tool_calls") if not isinstance(calls, list): return [] return [part for call in calls if (part := _tool_call(call)) is not None] def _message_parts(message: object, role: str) -> list[dict[str, object]]: content = _field(message, "content") if role == "tool": if not isinstance(content, str): return [] return [_call_id({"type": "tool_call_response", "response": content}, message, "tool_call_id")] return _text_parts(content) + _tool_calls(message) def _message(message: object) -> dict[str, object] | None: role = _field(message, "role") if not isinstance(role, str): return None parts = _message_parts(message, role) return {"role": role, "parts": parts} if parts else None def _messages(values: object) -> list[dict[str, object]]: if not isinstance(values, list): return [] return [message for value in values if (message := _message(value)) is not None] def _output_messages(response: object) -> list[dict[str, object]]: choices = _field(response, "choices") if not isinstance(choices, list): return [] return [message for choice in choices if (message := _message(_field(choice, "message"))) is not None] def _messages_hidden(direction: str, config: TracingConfig) -> bool: if direction == "output": return config.hide_outputs # Request messages replay model output: earlier assistant turns, and on an Instructor # re-ask the rejected answer plus the validation error quoting it. return config.hide_inputs or config.hide_outputs def write_messages(span: Span, value: object, direction: str, config: TracingConfig) -> None: if not config.capture_messages or not config.llm_spans or context_api.get_value("suppress_instrumentation"): return if not _messages_hidden(direction, config): _write_visible_messages(span, value, direction, config) def _write_visible_messages(span: Span, value: object, direction: str, config: TracingConfig) -> None: try: messages = _messages(value) if direction == "input" else _output_messages(value) if not messages: return serialized = json.dumps(messages, ensure_ascii=False, separators=(",", ":"), allow_nan=False) key = f"gen_ai.{direction}.messages" if len(serialized) <= config.max_attribute_length: span.set_attribute(key, serialized) else: span.set_attribute(key + ".original_length", len(serialized)) except Exception: return ``` ### File: atomic-agents/atomic_agents/observability/span_attributes.py ```python """Span attributes and conversation context.""" import contextvars import hashlib import json import logging from collections.abc import Mapping from contextlib import contextmanager from typing import Callable, Dict, Iterator, List from opentelemetry import context as context_api from opentelemetry.util.types import AttributeValue from atomic_agents.agents.atomic_agent import AtomicAgent from . import installation logger = logging.getLogger("atomic_agents.observability.tracing") ATTR_MODEL = "atomic_agents.model" ATTR_MODE = "atomic_agents.mode" ATTR_INPUT_SCHEMA = "atomic_agents.input_schema" ATTR_OUTPUT_SCHEMA = "atomic_agents.output_schema" ATTR_HISTORY_MESSAGE_COUNT = "atomic_agents.history.message_count" ATTR_CONTEXT_PROVIDERS = "atomic_agents.context_providers" ATTR_SYSTEM_PROMPT_SHA256 = "atomic_agents.system_prompt.sha256" ATTR_STREAM_CLOSED_EARLY = "atomic_agents.stream.closed_early" ATTR_OPERATION = "gen_ai.operation.name" ATTR_AGENT_NAME = "gen_ai.agent.name" ATTR_REQUEST_MODEL = "gen_ai.request.model" ATTR_PROVIDER_NAME = "gen_ai.provider.name" ATTR_CONVERSATION_ID = "gen_ai.conversation.id" ATTR_TOOL_NAME = "gen_ai.tool.name" ATTR_TOOL_DESCRIPTION = "gen_ai.tool.description" ATTR_TOOL_TYPE = "gen_ai.tool.type" ATTR_TOOL_ARGUMENTS = "gen_ai.tool.call.arguments" ATTR_TOOL_RESULT = "gen_ai.tool.call.result" ATTR_SPAN_KIND = "openinference.span.kind" ATTR_INPUT_VALUE = "input.value" ATTR_INPUT_MIME_TYPE = "input.mime_type" ATTR_OUTPUT_VALUE = "output.value" ATTR_OUTPUT_MIME_TYPE = "output.mime_type" ATTR_OI_AGENT_NAME = "agent.name" ATTR_OI_TOOL_NAME = "tool.name" ATTR_OI_TOOL_DESCRIPTION = "tool.description" ATTR_TOOL_PARAMETERS = "tool.parameters" ATTR_SESSION_ID = "session.id" ATTR_INPUT = "atomic_agents.input" ATTR_OUTPUT = "atomic_agents.output" ATTR_AGENT_CLASS = "atomic_agents.agent.class" ATTR_ERROR_TYPE = "error.type" ATTR_TOOL_CLASS = "atomic_agents.tool.class" PROVIDER_NAMES = { "azure_openai": "azure.ai.openai", "vertexai": "gcp.vertex_ai", "gemini": "gcp.gemini", "genai": "gcp.gemini", "bedrock": "aws.bedrock", "mistral": "mistral_ai", "xai": "x_ai", } def _import_context_attributes() -> Callable[[], dict[str, AttributeValue]] | None: try: from openinference.instrumentation import get_attributes_from_context except ImportError: return None from atomic_agents.observability import _openinference_is_too_old if _openinference_is_too_old(): logger.warning("OpenInference context attributes disabled: 0.1.38 or newer is required.") return None return lambda: dict(get_attributes_from_context()) def _context_attributes() -> dict[str, AttributeValue]: reader = installation._current_installation().context_attributes return reader() if reader is not None else {} def _suppressed() -> bool: return bool(context_api.get_value("suppress_instrumentation")) _conversation_id: contextvars.ContextVar[str | None] = contextvars.ContextVar("atomic_agents_conversation_id", default=None) @contextmanager def using_conversation(conversation_id: str) -> Iterator[None]: token = _conversation_id.set(conversation_id) try: yield finally: _conversation_id.reset(token) def _agent_display_name(agent: AtomicAgent) -> str: attributes = _best_effort_attributes({ATTR_AGENT_NAME: lambda: getattr(agent, "name", None)}) return str(attributes.get(ATTR_AGENT_NAME) or type(agent).__name__) def _agent_class_name(agent: AtomicAgent) -> str: return f"{type(agent).__name__}[{agent.input_schema.__name__}, {agent.output_schema.__name__}]" def _best_effort_attributes( readers: dict[str, Callable[[], AttributeValue | None]], ) -> dict[str, AttributeValue]: attributes: dict[str, AttributeValue] = {} for key, read in readers.items(): try: value = read() if value is not None: attributes[key] = value except Exception as error: logger.debug("Could not collect tracing attribute %s: %s", key, error) return attributes def _without_none(attributes: Mapping[str, AttributeValue | None]) -> dict[str, AttributeValue]: return {key: value for key, value in attributes.items() if value is not None} def _agent_start_attributes(agent: AtomicAgent) -> dict[str, AttributeValue]: attributes: dict[str, AttributeValue] = { ATTR_OPERATION: "invoke_agent", ATTR_AGENT_NAME: _agent_display_name(agent), ATTR_AGENT_CLASS: _agent_class_name(agent), } attributes.update( _best_effort_attributes( { ATTR_REQUEST_MODEL: lambda: agent.model, ATTR_MODEL: lambda: agent.model, ATTR_MODE: lambda: getattr(agent.mode, "value", None), ATTR_INPUT_SCHEMA: lambda: agent.input_schema.__name__, ATTR_OUTPUT_SCHEMA: lambda: agent.output_schema.__name__, ATTR_HISTORY_MESSAGE_COUNT: lambda: len(agent.history.history), ATTR_CONTEXT_PROVIDERS: lambda: _context_provider_names(agent), ATTR_SYSTEM_PROMPT_SHA256: lambda: _system_prompt_version(agent).get(ATTR_SYSTEM_PROMPT_SHA256), ATTR_PROVIDER_NAME: lambda: _provider_attributes(agent).get(ATTR_PROVIDER_NAME), } ) ) attributes.update(_optional_attributes("AGENT", {ATTR_OI_AGENT_NAME: _agent_display_name(agent)})) return _without_none(attributes) def _context_provider_names(agent: AtomicAgent) -> List[str]: providers = getattr(agent.system_prompt_generator, "context_providers", None) or {} return sorted(str(name) for name in providers) def _system_prompt_version(agent: AtomicAgent) -> Dict[str, str]: """Hashes only the static prompt sections, so the value identifies the prompt version across turns even though context providers change the rendered prompt on every call.""" generator = agent.system_prompt_generator sections = {name: getattr(generator, name, None) for name in ("background", "steps", "output_instructions")} if agent.system_role is None or all(value is None for value in sections.values()): return {} digest = hashlib.sha256(json.dumps(sections, sort_keys=True, default=str).encode("utf-8")).hexdigest() return {ATTR_SYSTEM_PROMPT_SHA256: digest} def _provider_attributes(agent: AtomicAgent) -> dict[str, str]: provider = getattr(getattr(agent.client, "provider", None), "value", None) return {ATTR_PROVIDER_NAME: PROVIDER_NAMES.get(provider, provider)} if isinstance(provider, str) else {} def _optional_attributes(kind: str, attributes: dict[str, AttributeValue]) -> dict[str, AttributeValue]: context = _context_attributes() conversation = _conversation_attributes(context) if not installation._current_installation().config.openinference_attributes: return conversation return _without_none({ATTR_SPAN_KIND: kind, **attributes, **context, **conversation}) def _conversation_attributes(context: dict[str, AttributeValue]) -> dict[str, AttributeValue]: conversation_id: AttributeValue | None = _conversation_id.get() if conversation_id is None: conversation_id = context.get(ATTR_SESSION_ID) if conversation_id is None: return {} attributes = {ATTR_CONVERSATION_ID: conversation_id} if installation._current_installation().config.openinference_attributes: attributes[ATTR_SESSION_ID] = conversation_id return attributes ``` ### File: atomic-agents/atomic_agents/observability/span_payloads.py ```python """Payload serialization, redaction, and bounded failures.""" import json import logging import traceback from collections.abc import Iterator from instructor.core.exceptions import InstructorError from opentelemetry.trace import Span, Status, StatusCode from pydantic import BaseModel, ValidationError from . import installation from .tracing_config import TracingConfig from .span_attributes import ( ATTR_INPUT_VALUE, ATTR_INPUT_MIME_TYPE, ATTR_OUTPUT_VALUE, ATTR_OUTPUT_MIME_TYPE, ATTR_INPUT, ATTR_OUTPUT, ATTR_ERROR_TYPE, ) logger = logging.getLogger("atomic_agents.observability.tracing") TRUNCATION_MARKER = "...[truncated]" def _serialize_payload(value: object) -> str | None: try: return value.model_dump_json() if isinstance(value, BaseModel) else json.dumps(value, default=str) except Exception as error: logger.warning("Could not serialize tracing payload: %s", error) return None def _record_payload( span: Span, key_pairs: tuple[tuple[str, str | None], ...], value: object, hide: bool, config: TracingConfig ) -> None: if value is None: return serialized = "__REDACTED__" if hide else _serialize_payload(value) if serialized is None: return for key, mime_key in key_pairs: _write_payload(span, key, mime_key, serialized, hide, config) def _write_payload(span: Span, key: str, mime_key: str | None, serialized: str, hide: bool, config: TracingConfig) -> None: if hide: span.set_attribute(key, "__REDACTED__") return if len(serialized) > config.max_attribute_length: span.set_attribute(key + ".original_length", len(serialized)) serialized = _bounded_text(serialized, config) span.set_attribute(key, serialized) if mime_key is not None: span.set_attribute(mime_key, "application/json") def _payload_keys(primary: str, output: bool, config: TracingConfig) -> tuple[tuple[str, str | None], ...]: keys: tuple[tuple[str, str | None], ...] = ((primary, None),) if config.openinference_attributes: keys += ((ATTR_OUTPUT_VALUE, ATTR_OUTPUT_MIME_TYPE),) if output else ((ATTR_INPUT_VALUE, ATTR_INPUT_MIME_TYPE),) return keys def _set_input_safely(span: Span, value: object, primary: str = ATTR_INPUT) -> None: config = installation._current_installation().config _record_payload(span, _payload_keys(primary, False, config), value, config.hide_inputs, config) def _set_output_safely(span: Span, value: object, primary: str = ATTR_OUTPUT, *, config: TracingConfig) -> None: _record_payload(span, _payload_keys(primary, True, config), value, config.hide_outputs, config) def _bounded_text(value: str, config: TracingConfig) -> str: limit = config.max_attribute_length if len(value) > limit: marker = TRUNCATION_MARKER[:limit] return value[: limit - len(marker)] + marker return value def _exception_chain(error: BaseException) -> Iterator[BaseException]: seen: set[int] = set() while id(error) not in seen: seen.add(id(error)) yield error following = error.__cause__ or error.__context__ if following is None: return error = following def _instructor_attempt_count(error: InstructorError) -> int: if error.failed_attempts: return len(error.failed_attempts) return getattr(error, "n_attempts", 0) def _failure_count(error: BaseException) -> str | None: if isinstance(error, InstructorError): return f"{_instructor_attempt_count(error)} attempts" if isinstance(error, ValidationError): return f"{error.error_count()} errors" return None def _failure_message(error: BaseException, config: TracingConfig) -> str: if not config.hide_outputs: return str(error) count = next((count for cause in _exception_chain(error) if (count := _failure_count(cause)) is not None), "1 errors") return f"{type(error).__qualname__}: {count} (__REDACTED__)" def _failure_stacktrace(error: BaseException, config: TracingConfig) -> str: lines = traceback.format_exception(error, chain=not config.hide_outputs) if config.hide_outputs: lines = _redacted_traceback(error, config) return _bounded_text("".join(lines), config) def _redacted_traceback(error: BaseException, config: TracingConfig) -> list[str]: frames = [ traceback.FrameSummary(frame.filename, frame.lineno, frame.name, line="") for frame in traceback.extract_tb(error.__traceback__) ] return ["Traceback (most recent call last):\n", *traceback.format_list(frames), _failure_message(error, config) + "\n"] def _record_failure(span: Span, error: BaseException, config: TracingConfig) -> None: message = _bounded_text(_failure_message(error, config), config) span.set_attribute(ATTR_ERROR_TYPE, type(error).__qualname__) span.record_exception( error, attributes={"exception.message": message, "exception.stacktrace": _failure_stacktrace(error, config)} ) span.set_status(Status(StatusCode.ERROR, message)) ``` ### File: atomic-agents/atomic_agents/observability/tool_spans.py ```python """Tool spans and tool class instrumentation.""" import asyncio import contextvars import functools import inspect import json import logging from contextlib import contextmanager from types import FunctionType from typing import Awaitable, Callable, Coroutine, Iterator, List, Optional, Type, TypeVar from opentelemetry.trace import Span, Status, StatusCode from opentelemetry.util.types import AttributeValue from atomic_agents.base.base_tool import BaseTool from . import installation from .tracing_config import TracingConfig from .span_attributes import _best_effort_attributes, _optional_attributes, _suppressed, _without_none from .span_payloads import _bounded_text, _record_failure, _set_input_safely, _set_output_safely from .span_attributes import ( ATTR_OPERATION, ATTR_TOOL_NAME, ATTR_TOOL_DESCRIPTION, ATTR_TOOL_TYPE, ATTR_TOOL_ARGUMENTS, ATTR_TOOL_RESULT, ATTR_OI_TOOL_NAME, ATTR_OI_TOOL_DESCRIPTION, ATTR_TOOL_PARAMETERS, ATTR_TOOL_CLASS, ) logger = logging.getLogger("atomic_agents.observability.tracing") OutputSchema = TypeVar("OutputSchema") TOOL_METHOD_NAMES = ("run", "run_async", "arun") _active_tool: contextvars.ContextVar[Optional[BaseTool]] = contextvars.ContextVar("atomic_agents_active_tool", default=None) def _all_subclasses(base: Type) -> List[Type]: found: List[Type] = [] pending = list(base.__subclasses__()) while pending: current = pending.pop(0) if current not in found: found.append(current) pending.extend(current.__subclasses__()) return found def _instrument_existing_tool_classes() -> None: for tool_class in _all_subclasses(BaseTool): _instrument_tool_class(tool_class) def _hook_future_tool_classes(original_init_subclass: classmethod) -> None: def traced_init_subclass(cls, **kwargs): original_init_subclass.__func__(cls, **kwargs) _instrument_tool_class(cls) setattr(BaseTool, "__init_subclass__", classmethod(traced_init_subclass)) def _instrument_tool_class(tool_class: Type) -> None: for method_name in TOOL_METHOD_NAMES: method = inspect.getattr_static(tool_class, method_name, None) if _is_traceable_function(method): inherited = method_name not in tool_class.__dict__ setattr(tool_class, method_name, installation._mark_traced(_trace_tool_method(method), method, inherited)) def _is_traceable_function(method: object) -> bool: if not isinstance(method, FunctionType) or inspect.isasyncgenfunction(method): return False if getattr(method, "__isabstractmethod__", False): return False return not getattr(method, installation.TRACED_MARKER, False) def _restore_tool_class(tool_class: Type) -> None: for method_name in TOOL_METHOD_NAMES: installation._restore_method(tool_class, method_name) def _tool_parameter_attributes(tool: BaseTool) -> dict[str, AttributeValue]: parameters = _best_effort_attributes({ATTR_TOOL_PARAMETERS: lambda: json.dumps(tool.input_schema.model_json_schema())}) value = parameters.get(ATTR_TOOL_PARAMETERS) if isinstance(value, str): config = installation._current_installation().config parameters[ATTR_TOOL_PARAMETERS] = _bounded_text(value, config) if len(value) > config.max_attribute_length: parameters[ATTR_TOOL_PARAMETERS + ".original_length"] = len(value) return parameters def _tool_display_name(tool: BaseTool) -> str: try: return tool.tool_name except Exception as error: logger.debug("Could not collect tracing tool name: %s", error) return type(tool).__name__ def _tool_start_attributes(tool: BaseTool) -> dict[str, AttributeValue]: attributes: dict[str, AttributeValue] = { ATTR_OPERATION: "execute_tool", ATTR_TOOL_NAME: _tool_display_name(tool), ATTR_TOOL_CLASS: type(tool).__name__, ATTR_TOOL_TYPE: "function", **_optional_attributes("TOOL", {}), } attributes.update(_best_effort_attributes({ATTR_TOOL_DESCRIPTION: lambda: tool.tool_description})) if installation._current_installation().config.openinference_attributes: attributes.update( _best_effort_attributes( { ATTR_OI_TOOL_NAME: lambda: attributes[ATTR_TOOL_NAME], ATTR_OI_TOOL_DESCRIPTION: lambda: tool.tool_description, } ) ) attributes.update(_tool_parameter_attributes(tool)) return _without_none(attributes) @contextmanager def _tool_span(tool: BaseTool, args: tuple[object, ...], kwargs: dict[str, object]) -> Iterator[tuple[Span, TracingConfig]]: installed = installation._current_installation() config = installed.config attributes = _tool_start_attributes(tool) with installed.tracer.start_as_current_span( f"execute_tool {attributes[ATTR_TOOL_NAME]}", attributes=attributes, record_exception=False, set_status_on_exception=False, ) as span: _set_input_safely(span, _tool_input_from_call(args, kwargs), ATTR_TOOL_ARGUMENTS) try: yield span, config except (asyncio.CancelledError, GeneratorExit): span.set_attribute("atomic_agents.cancelled", True) raise except BaseException as error: _record_failure(span, error, config) raise def _tool_input_from_call(args: tuple[object, ...], kwargs: dict[str, object]) -> object: return args[0] if args else kwargs.get("params") def _trace_tool_method( wrapped: Callable[..., OutputSchema], ) -> Callable[..., OutputSchema] | Callable[..., Coroutine[object, object, object]]: if inspect.iscoroutinefunction(wrapped): return _trace_tool_async_method(wrapped) @functools.wraps(wrapped) def traced(tool: BaseTool, *args: object, **kwargs: object) -> OutputSchema: if not installation.is_installed() or _active_tool.get() is tool or _suppressed(): return wrapped(tool, *args, **kwargs) with _tool_span(tool, args, kwargs) as (span, config): result = _run_as_active_tool(tool, wrapped, args, kwargs) _set_output_safely(span, result, ATTR_TOOL_RESULT, config=config) span.set_status(Status(StatusCode.OK)) return result return traced def _trace_tool_async_method( wrapped: Callable[..., Awaitable[OutputSchema]], ) -> Callable[..., Coroutine[object, object, OutputSchema]]: @functools.wraps(wrapped) async def traced(tool: BaseTool, *args: object, **kwargs: object) -> OutputSchema: if not installation.is_installed() or _active_tool.get() is tool or _suppressed(): return await wrapped(tool, *args, **kwargs) with _tool_span(tool, args, kwargs) as (span, config): result = await _await_as_active_tool(tool, wrapped, args, kwargs) _set_output_safely(span, result, ATTR_TOOL_RESULT, config=config) span.set_status(Status(StatusCode.OK)) return result return traced def _run_as_active_tool( tool: BaseTool, wrapped: Callable[..., OutputSchema], args: tuple[object, ...], kwargs: dict[str, object] ) -> OutputSchema: token = _active_tool.set(tool) try: return wrapped(tool, *args, **kwargs) finally: _active_tool.reset(token) async def _await_as_active_tool( tool: BaseTool, wrapped: Callable[..., Awaitable[OutputSchema]], args: tuple[object, ...], kwargs: dict[str, object] ) -> OutputSchema: token = _active_tool.set(tool) try: return await wrapped(tool, *args, **kwargs) finally: _active_tool.reset(token) ``` ### File: atomic-agents/atomic_agents/observability/tracing.py ```python """Compatibility exports for Atomic Agents tracing.""" from . import installation, span_attributes, tracing_config TRACER_NAME = installation.TRACER_NAME install = installation.install is_installed = installation.is_installed uninstall = installation.uninstall PayloadVisibility = tracing_config.PayloadVisibility TracingConfig = tracing_config.TracingConfig using_conversation = span_attributes.using_conversation ``` ### File: atomic-agents/atomic_agents/observability/tracing_config.py ```python """Tracing configuration and payload visibility.""" import logging import os from dataclasses import dataclass, replace from typing import Protocol logger = logging.getLogger("atomic_agents.observability.tracing") DEFAULT_MAX_ATTRIBUTE_LENGTH = 4096 class PayloadVisibility(Protocol): hide_inputs: bool hide_outputs: bool @dataclass(frozen=True) class TracingConfig: hide_inputs: bool = False hide_outputs: bool = False max_attribute_length: int = DEFAULT_MAX_ATTRIBUTE_LENGTH openinference_attributes: bool | None = None llm_spans: bool = True capture_messages: bool = False @classmethod def from_environment(cls) -> "TracingConfig": return cls( hide_inputs=_visibility_environment("INPUTS").lower() in ("true", "1"), hide_outputs=_visibility_environment("OUTPUTS").lower() in ("true", "1"), max_attribute_length=_attribute_length_environment(), llm_spans=os.environ.get("ATOMIC_AGENTS_TRACE_LLM_SPANS", "true").lower() not in ("false", "0"), capture_messages=os.environ.get("ATOMIC_AGENTS_TRACE_CAPTURE_MESSAGES", "").lower() in ("true", "1"), ) @classmethod def coerce(cls, config: "TracingConfig | PayloadVisibility | None") -> "TracingConfig": if config is None: return cls.from_environment() if isinstance(config, cls): return config return replace(cls.from_environment(), hide_inputs=config.hide_inputs, hide_outputs=config.hide_outputs) def _visibility_environment(direction: str) -> str: return os.environ.get(f"ATOMIC_AGENTS_TRACE_HIDE_{direction}", os.environ.get(f"OPENINFERENCE_HIDE_{direction}", "")) def _attribute_length_environment() -> int: value = os.environ.get("ATOMIC_AGENTS_TRACE_MAX_ATTRIBUTE_LENGTH", str(DEFAULT_MAX_ATTRIBUTE_LENGTH)) try: length = int(value) if length >= 1: return length except ValueError: pass logger.warning("Invalid ATOMIC_AGENTS_TRACE_MAX_ATTRIBUTE_LENGTH value %r; using %d", value, DEFAULT_MAX_ATTRIBUTE_LENGTH) return DEFAULT_MAX_ATTRIBUTE_LENGTH ``` ### File: atomic-agents/atomic_agents/utils/__init__.py ```python """Utility functions.""" from .format_tool_message import format_tool_message from .token_counter import TokenCounter, TokenCountResult, TokenCountError, get_token_counter __all__ = [ "format_tool_message", "TokenCounter", "TokenCountResult", "TokenCountError", "get_token_counter", ] ``` ### File: atomic-agents/atomic_agents/utils/format_tool_message.py ```python import json import uuid from pydantic import BaseModel from typing import Dict, Optional, Type def format_tool_message(tool_call: Type[BaseModel], tool_id: Optional[str] = None) -> Dict: """ Formats a message for a tool call, using Pydantic's JSON representations for arguments. Args: 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: Dict: A formatted message dictionary for the tool call. """ if tool_id is None: tool_id = str(uuid.uuid4()) # Use the same schema title that Instructor advertises as the function name. tool_name = tool_call.__class__.model_json_schema().get("title", tool_call.__class__.__name__) return { "id": tool_id, "type": "function", "function": { "name": tool_name, "arguments": json.dumps(tool_call.model_dump(mode="json"), separators=(", ", ": ")), }, } ``` ### File: atomic-agents/atomic_agents/utils/token_counter.py ```python """Token counting utilities for provider-agnostic context measurement.""" import logging from typing import Any, Dict, List, NamedTuple, Optional logger = logging.getLogger(__name__) class TokenCountError(Exception): """Exception raised when token counting fails.""" pass class TokenCountResult(NamedTuple): """ Result of a token count operation. Attributes: total: Total number of tokens in the complete message request. When tools are provided without messages, this is schema-only overhead and excludes request framing. system_prompt: Tokens in the system messages, including request framing when a system message is present (0 otherwise). history: Incremental tokens added by the history messages. Includes request framing when there are no system messages. tools: Incremental tokens added by tools to the complete message request. When there are no messages, this is the schema-only token overhead. model: The model used for tokenization. max_tokens: Maximum context window for the model (None if unknown). utilization: Percentage of context window used (None if max_tokens unknown). """ total: int system_prompt: int history: int tools: int model: str max_tokens: Optional[int] = None utilization: Optional[float] = None # Module-level singleton for efficiency _token_counter_instance: Optional["TokenCounter"] = None def get_token_counter() -> "TokenCounter": """Get the singleton TokenCounter instance.""" global _token_counter_instance if _token_counter_instance is None: _token_counter_instance = TokenCounter() return _token_counter_instance class TokenCounter: """ Utility class for counting tokens using LiteLLM's provider-agnostic tokenizer. This class provides methods for counting tokens in messages, text, tools, and retrieving model context limits. It uses LiteLLM's token_counter which automatically selects the appropriate tokenizer based on the model. Works with any model supported by LiteLLM including: - OpenAI (gpt-4, gpt-3.5-turbo, etc.) - Anthropic (claude-3-opus, claude-3-sonnet, etc.) - Google (gemini-pro, gemini-1.5-pro, etc.) - And 100+ other providers Example: ```python counter = TokenCounter() # Count tokens in messages messages = [{"role": "user", "content": "Hello, world!"}] count = counter.count_messages("gpt-4", messages) # Count tokens with tools (for TOOLS mode) tools = [{"type": "function", "function": {...}}] count = counter.count_messages("gpt-4", messages, tools=tools) # Get max tokens for a model max_tokens = counter.get_max_tokens("gpt-4") ``` """ def count_messages( self, model: str, messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None, ) -> int: """ Count the number of tokens in a list of messages and optional tools. Args: model: The model identifier (e.g., "gpt-4", "anthropic/claude-3-opus"). messages: List of message dictionaries with 'role' and 'content' keys. tools: Optional list of tool definitions (for TOOLS mode). Returns: The number of tokens in the messages (and tools if provided). Raises: TokenCountError: If token counting fails. """ if not model: raise ValueError("model is required for token counting") try: from litellm import token_counter if tools: return token_counter(model=model, messages=messages, tools=tools) return token_counter(model=model, messages=messages) except ImportError as e: raise ImportError("litellm is required for token counting. " "Install it with: pip install litellm") from e except Exception as e: raise TokenCountError(f"Failed to count tokens for model '{model}': {e}") from e def count_text(self, model: str, text: str) -> int: """ Count the number of tokens in a text string. Args: model: The model identifier. text: The text to tokenize. Returns: The number of tokens in the text. Raises: TokenCountError: If token counting fails. """ messages = [{"role": "user", "content": text}] return self.count_messages(model, messages) def get_max_tokens(self, model: str) -> Optional[int]: """ Get the maximum context window size for a model. Args: model: The model identifier. Returns: The maximum number of tokens, or None if unknown. Raises: TypeError: If model is None or not a string. ImportError: If litellm is not installed. """ if not isinstance(model, str): raise TypeError(f"model must be a string, got {type(model).__name__}") try: from litellm import get_model_info except ImportError as e: raise ImportError("litellm is required for token counting. " "Install it with: pip install litellm") from e try: info = get_model_info(model) # Use max_input_tokens (context window) not max_tokens (output limit) max_input = info.get("max_input_tokens") return max_input if max_input is not None else info.get("max_tokens") except Exception as e: logger.warning(f"Could not determine max tokens for model '{model}': {e}") return None def count_context( self, model: str, system_messages: List[Dict[str, Any]], history_messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None, ) -> TokenCountResult: """ Count tokens with breakdown by system prompt, history, and tools. The breakdown is additive: ``system_prompt + history + tools == total``. Request framing is attributed to the system prompt when present, or to history otherwise. Tool tokens are the increment over the full message request without tools. When tools are provided without any messages, ``total`` and ``tools`` report schema-only overhead and exclude request framing. Args: model: The model identifier. system_messages: System prompt messages (may be empty). history_messages: Conversation history messages. tools: Optional list of tool definitions (for TOOLS mode). Returns: TokenCountResult with breakdown and utilization metrics. Raises: TokenCountError: If token counting fails. """ system_tokens = self.count_messages(model, system_messages) if system_messages else 0 all_messages = system_messages + history_messages # Tokenizers include request-level framing, so independently counting the # system prompt and history would count that overhead twice. Count the # complete message list and attribute the incremental tokens to history. if history_messages: messages_tokens = self.count_messages(model, all_messages) history_tokens = messages_tokens - system_tokens else: messages_tokens = system_tokens history_tokens = 0 tools_tokens = 0 total_tokens = messages_tokens if tools: if all_messages: total_tokens = self.count_messages(model, all_messages, tools=tools) tools_tokens = total_tokens - messages_tokens else: # Preserve schema-only counting by subtracting the same placeholder # request without tools. Keep the placeholder out of the total. placeholder = [{"role": "user", "content": ""}] with_tools_tokens = self.count_messages(model, placeholder, tools=tools) without_tools_tokens = self.count_messages(model, placeholder) tools_tokens = with_tools_tokens - without_tools_tokens total_tokens = tools_tokens max_tokens = self.get_max_tokens(model) # Prevent division by zero utilization = (total_tokens / max_tokens) if max_tokens and max_tokens > 0 else None return TokenCountResult( total=total_tokens, system_prompt=system_tokens, history=history_tokens, tools=tools_tokens, model=model, max_tokens=max_tokens, utilization=utilization, ) ``` ### File: atomic-agents/tests/agents/test_agent_config_name.py ```python import instructor from openai import OpenAI from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema def test_agent_name_defaults_to_none() -> None: with OpenAI(api_key="synthetic-test-key") as provider_client: client = instructor.from_openai(provider_client) configuration = AgentConfig(client=client) agent = AtomicAgent[BaseIOSchema, BaseIOSchema](configuration) assert configuration.name is None assert agent.name is None def test_explicit_agent_name_is_stored() -> None: with OpenAI(api_key="synthetic-test-key") as provider_client: client = instructor.from_openai(provider_client) configuration = AgentConfig(client=client, name="orchestrator_planning") agent = AtomicAgent[BaseIOSchema, BaseIOSchema](configuration) assert configuration.name == "orchestrator_planning" assert agent.name == "orchestrator_planning" ``` ### File: atomic-agents/tests/agents/test_agent_hooks.py ```python import instructor import pytest from unittest.mock import Mock from atomic_agents import AgentConfig, AtomicAgent @pytest.mark.parametrize("event", [None, "completion:kwargs"]) def test_clear_hooks_preserves_foreign_handlers(event): client = instructor.Instructor(client=Mock(), create=Mock()) agent = AtomicAgent(AgentConfig(client=client)) foreign = Mock() owned = Mock() other_owned = Mock() client.on("completion:kwargs", foreign) agent.register_hook("completion:kwargs", owned) agent.register_hook("completion:response", other_owned) agent.clear_hooks(event) client.hooks.emit_completion_arguments(model="test-model") client.hooks.emit_completion_response(Mock()) foreign.assert_called_once_with(model="test-model") owned.assert_not_called() assert other_owned.call_count == (0 if event is None else 1) assert agent._hook_handlers.get("completion:kwargs", []) == [] def test_instructor_hooks_stop_and_resume_without_duplicates(): client = instructor.Instructor(client=Mock(), create=Mock()) agent = AtomicAgent(AgentConfig(client=client)) foreign = Mock() owned = Mock() client.on("completion:kwargs", foreign) agent.register_hook("completion:kwargs", owned) client.hooks.emit_completion_arguments(model="before") owned.assert_called_once_with(model="before") agent.disable_hooks() agent.disable_hooks() client.hooks.emit_completion_arguments(model="disabled") assert owned.call_count == 1 assert foreign.call_count == 2 agent.enable_hooks() agent.enable_hooks() client.hooks.emit_completion_arguments(model="resumed") assert owned.call_count == 2 owned.assert_called_with(model="resumed") assert foreign.call_count == 3 def test_instructor_hooks_registered_while_disabled_wait_for_enable(): client = instructor.Instructor(client=Mock(), create=Mock()) agent = AtomicAgent(AgentConfig(client=client)) handler = Mock() agent.disable_hooks() agent.register_hook("completion:kwargs", handler) client.hooks.emit_completion_arguments(model="disabled") handler.assert_not_called() agent.enable_hooks() client.hooks.emit_completion_arguments(model="resumed") handler.assert_called_once_with(model="resumed") @pytest.mark.parametrize("remove", ["unregister", "clear_event", "clear_all"]) def test_removing_disabled_instructor_hooks_does_not_restore_them(remove): client = instructor.Instructor(client=Mock(), create=Mock()) agent = AtomicAgent(AgentConfig(client=client)) removed = Mock() retained = Mock() agent.register_hook("completion:kwargs", removed) agent.register_hook("completion:response", retained) agent.disable_hooks() if remove == "unregister": agent.unregister_hook("completion:kwargs", removed) elif remove == "clear_event": agent.clear_hooks("completion:kwargs") else: agent.clear_hooks() agent.enable_hooks() client.hooks.emit_completion_arguments(model="resumed") client.hooks.emit_completion_response("response") removed.assert_not_called() assert retained.call_count == (0 if remove == "clear_all" else 1) @pytest.fixture def token_hook_agent(monkeypatch): """Use real Instructor hook validation without making provider/tokenizer calls.""" from atomic_agents.utils.token_counter import TokenCountResult client = instructor.Instructor(client=Mock(), create=Mock()) agent = AtomicAgent(AgentConfig(client=client)) result = TokenCountResult(total=100, system_prompt=30, history=70, tools=0, model="gpt-5-mini") counter = Mock() counter.count_context.return_value = result monkeypatch.setattr("atomic_agents.agents.atomic_agent.get_token_counter", lambda: counter) return agent, result @pytest.mark.parametrize("initially_disabled", [False, True]) def test_local_token_hook_lifecycle_with_real_instructor(token_hook_agent, initially_disabled): agent, result = token_hook_agent local = Mock() native = Mock() foreign = Mock() agent.client.on("completion:kwargs", foreign) if initially_disabled: agent.disable_hooks() agent.register_hook("token:counted", local) agent.register_hook("completion:kwargs", native) if initially_disabled: agent.get_context_token_count() local.assert_not_called() agent.enable_hooks() assert agent.get_context_token_count() == result local.assert_called_once_with(result) agent.client.hooks.emit_completion_arguments(model="before") native.assert_called_once_with(model="before") agent.disable_hooks() agent.get_context_token_count() agent.client.hooks.emit_completion_arguments(model="disabled") assert local.call_count == native.call_count == 1 assert foreign.call_count == 2 agent.enable_hooks() agent.get_context_token_count() agent.client.hooks.emit_completion_arguments(model="resumed") assert local.call_count == native.call_count == 2 assert foreign.call_count == 3 @pytest.mark.parametrize("disabled", [False, True]) @pytest.mark.parametrize("remove", ["unregister", "clear_event", "clear_all"]) def test_remove_local_token_hook_with_real_instructor(token_hook_agent, disabled, remove): agent, _ = token_hook_agent local = Mock() native = Mock() foreign = Mock() agent.client.on("completion:kwargs", foreign) agent.register_hook("token:counted", local) agent.register_hook("completion:kwargs", native) if disabled: agent.disable_hooks() if remove == "unregister": agent.unregister_hook("token:counted", local) elif remove == "clear_event": agent.clear_hooks("token:counted") else: agent.clear_hooks() if disabled: agent.enable_hooks() agent.get_context_token_count() agent.client.hooks.emit_completion_arguments(model="after") local.assert_not_called() foreign.assert_called_once_with(model="after") assert native.call_count == (0 if remove == "clear_all" else 1) def test_duplicate_local_token_hooks_with_real_instructor(token_hook_agent): agent, result = token_hook_agent local = Mock() agent.register_hook("token:counted", local) agent.register_hook("token:counted", local) agent.get_context_token_count() assert local.call_count == 2 local.assert_called_with(result) agent.unregister_hook("token:counted", local) agent.get_context_token_count() assert local.call_count == 3 agent.clear_hooks("token:counted") agent.get_context_token_count() assert local.call_count == 3 @pytest.mark.parametrize( "mutation,next_handlers", [ ("unregister", ["retained"]), ("unregister_raising", ["retained"]), ("register", ["mutating", "retained", "late"]), ("clear_event", []), ("clear_all", []), ], ) def test_local_token_hook_mutations_apply_to_next_dispatch(token_hook_agent, mutation, next_handlers): agent, result = token_hook_agent calls = [] added = False def late(value): calls.append(("late", value)) def mutating(value): nonlocal added calls.append(("mutating", value)) if mutation.startswith("unregister"): agent.unregister_hook("token:counted", mutating) if mutation == "unregister_raising": raise ValueError("one-shot callback failed") elif mutation == "register": if not added: added = True agent.register_hook("token:counted", late) else: agent.clear_hooks("token:counted" if mutation == "clear_event" else None) def retained(value): calls.append(("retained", value)) agent.register_hook("token:counted", mutating) agent.register_hook("token:counted", retained) assert agent.get_context_token_count() == result assert calls == [("mutating", result), ("retained", result)] calls.clear() assert agent.get_context_token_count() == result assert calls == [(name, result) for name in next_handlers] def test_unknown_hook_is_still_validated_by_instructor(token_hook_agent): agent, _ = token_hook_agent with pytest.raises(ValueError, match="Invalid hook name"): agent.register_hook("token:countedd", Mock()) ``` ### File: atomic-agents/tests/agents/test_agent_mode_consistency.py ```python """Tests for AgentConfig mode resolution and client/config mode consistency (issue #282). The mode on the Instructor client decides the API call format, while ``AgentConfig.mode`` only drives token accounting. These tests cover deriving the mode from the client, warning when an explicit mode disagrees with the client on how the schema is transmitted, and keeping the accounting aligned with the requests Instructor actually prepares. """ import logging from unittest.mock import Mock, patch import instructor import pytest from instructor.processing.response import handle_response_model from openai import OpenAI from pydantic import Field from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema, BasicChatInputSchema, BasicChatOutputSchema from atomic_agents.agents.atomic_agent import _TOOL_MODES from atomic_agents.context import ChatHistory DUMMY_KEY = "not-a-real-key" LOGGER_NAME = "atomic_agents.agents.atomic_agent" COUNTING_MODEL = "gpt-4o-mini" # Request parameters Instructor uses to transmit the output schema as a tool or # function definition; xAI carries it under the private _xai_tool kwarg instead. TOOL_REQUEST_PARAMS = ("tools", "functions", "toolConfig", "_xai_tool") class InputWithImage(BaseIOSchema): """Input with image.""" chat_message: str = Field(..., description="Text input") image: instructor.Image = Field(..., description="Image to analyze") def _client(mode): """Build a real Instructor client with the given mode (no request is ever made).""" return instructor.from_openai(OpenAI(api_key=DUMMY_KEY, base_url="https://example.invalid/v1"), mode=mode) def _mock_client(mode): """A client double that exposes `mode` like a real Instructor client does.""" client = Mock(spec=instructor.Instructor) client.mode = mode return client def _agent(config): return AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) def _history_with_image(): history = ChatHistory() history.add_message( "user", InputWithImage(chat_message="Look", image=instructor.Image.from_url("https://example.com/x.png")) ) return history def _prepared_request(mode, messages=None): """Return the kwargs Instructor would send for `mode`, prepared without any provider call.""" if messages is None: messages = [{"role": "user", "content": "hello"}] return handle_response_model( response_model=BasicChatOutputSchema, mode=mode, messages=[dict(message) for message in messages], model="test-model", )[1] def _mode_id(mode): """Readable test id for each Instructor mode.""" return mode.name def _agent_warnings(caplog): """Assertion oracle: the warning records the agent module emitted during a test.""" return [record for record in caplog.records if record.name == LOGGER_NAME and record.levelno == logging.WARNING] def _assert_warns_once(caplog, *fragments): """Assertion oracle: exactly one agent warning was logged, mentioning every fragment.""" warnings = _agent_warnings(caplog) assert len(warnings) == 1 message = warnings[0].getMessage() for fragment in fragments: assert fragment in message class TestModeDerivation: """AgentConfig.mode defaults to the client's mode.""" def test_none_derives_client_mode(self): agent = _agent(AgentConfig(client=_client(instructor.Mode.JSON), model=COUNTING_MODEL)) assert agent.mode == instructor.Mode.JSON def test_none_derives_tools_mode(self): agent = _agent(AgentConfig(client=_client(instructor.Mode.TOOLS), model=COUNTING_MODEL)) assert agent.mode == instructor.Mode.TOOLS def test_none_derives_native_provider_mode(self): agent = _agent(AgentConfig(client=_mock_client(instructor.Mode.ANTHROPIC_TOOLS), model=COUNTING_MODEL)) assert agent.mode == instructor.Mode.ANTHROPIC_TOOLS def test_derived_mode_drives_token_accounting(self): json_agent = _agent(AgentConfig(client=_client(instructor.Mode.JSON), model=COUNTING_MODEL)) tools_agent = _agent(AgentConfig(client=_client(instructor.Mode.TOOLS), model=COUNTING_MODEL)) assert json_agent._build_tools_definition() is None assert tools_agent._build_tools_definition() is not None def test_derived_native_mode_counts_schema_as_tools(self): agent = _agent(AgentConfig(client=_mock_client(instructor.Mode.ANTHROPIC_TOOLS), model=COUNTING_MODEL)) assert agent._build_tools_definition() is not None assert agent.get_context_token_count().tools > 0 class TestClientWithoutMode: """Clients that expose no Mode attribute fall back to TOOLS without warnings.""" def test_mock_client_falls_back_to_tools_mode(self, caplog): mock_client = Mock(spec=instructor.Instructor) # spec exposes no instance attribute `mode` agent = _agent(AgentConfig(client=mock_client, model=COUNTING_MODEL)) assert agent.mode == instructor.Mode.TOOLS assert not [r for r in caplog.records if r.name == LOGGER_NAME and r.levelno >= logging.WARNING] def test_explicit_mode_is_kept_for_mock_client(self, caplog): mock_client = Mock(spec=instructor.Instructor) agent = _agent(AgentConfig(client=mock_client, model=COUNTING_MODEL, mode=instructor.Mode.JSON)) assert agent.mode == instructor.Mode.JSON assert not [r for r in caplog.records if r.name == LOGGER_NAME and r.levelno >= logging.WARNING] class TestModeConsistencyWarning: """Explicit modes are checked against how the client transmits the schema.""" def test_matching_mode_does_not_warn(self, caplog): with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): agent = _agent(AgentConfig(client=_client(instructor.Mode.JSON), model=COUNTING_MODEL, mode=instructor.Mode.JSON)) assert agent.mode == instructor.Mode.JSON assert not [r for r in caplog.records if r.name == LOGGER_NAME] def test_mismatching_accounting_warns(self, caplog): with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): agent = _agent(AgentConfig(client=_client(instructor.Mode.JSON), model=COUNTING_MODEL, mode=instructor.Mode.TOOLS)) assert agent.mode == instructor.Mode.TOOLS _assert_warns_once(caplog, "TOOLS", "JSON") def test_same_accounting_different_mode_does_not_warn(self, caplog): # TOOLS_STRICT on the config and TOOLS on the client both send a tools definition. with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): agent = _agent( AgentConfig(client=_client(instructor.Mode.TOOLS), model=COUNTING_MODEL, mode=instructor.Mode.TOOLS_STRICT) ) assert agent.mode == instructor.Mode.TOOLS_STRICT assert not [r for r in caplog.records if r.name == LOGGER_NAME] def test_native_provider_mode_matches_tools_accounting(self, caplog): # ANTHROPIC_TOOLS on the client and TOOLS on the config are counted the same way. with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): agent = _agent( AgentConfig( client=_mock_client(instructor.Mode.ANTHROPIC_TOOLS), model=COUNTING_MODEL, mode=instructor.Mode.TOOLS ) ) assert agent.mode == instructor.Mode.TOOLS assert not [r for r in caplog.records if r.name == LOGGER_NAME] def test_native_provider_mode_against_json_accounting_warns(self, caplog): with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): _agent( AgentConfig( client=_mock_client(instructor.Mode.ANTHROPIC_TOOLS), model=COUNTING_MODEL, mode=instructor.Mode.JSON ) ) _assert_warns_once(caplog, "AgentConfig.mode (JSON)", "client's mode (ANTHROPIC_TOOLS)") def test_warning_mentions_both_directions(self, caplog): with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): _agent(AgentConfig(client=_client(instructor.Mode.TOOLS), model=COUNTING_MODEL, mode=instructor.Mode.JSON)) _assert_warns_once(caplog, "AgentConfig.mode (JSON)", "client's mode (TOOLS)") class TestMultimodalTokenCounting: """Media serialization stays in the chat format the token counter accepts.""" def test_image_history_counts_in_responses_mode(self): # Regression check: Responses-format input_image parts used to reach the # chat token counter and make get_context_token_count() raise. agent = _agent( AgentConfig(client=_client(instructor.Mode.RESPONSES_TOOLS), model=COUNTING_MODEL, history=_history_with_image()) ) result = agent.get_context_token_count() assert result.history > 0 assert result.total > 0 def test_media_parts_stay_chat_format_in_responses_mode(self): agent = _agent( AgentConfig(client=_client(instructor.Mode.RESPONSES_TOOLS), model=COUNTING_MODEL, history=_history_with_image()) ) serialized = agent._serialize_history_for_token_count() assert [part["type"] for part in serialized[0]["content"]] == ["text", "image_url"] def test_image_history_counts_in_tools_mode(self): agent = _agent(AgentConfig(client=_client(instructor.Mode.TOOLS), model=COUNTING_MODEL, history=_history_with_image())) result = agent.get_context_token_count() assert result.history > 0 assert result.tools > 0 def test_media_serialization_failure_falls_back_to_placeholder(self, caplog): agent = _agent(AgentConfig(client=_client(instructor.Mode.TOOLS), model=COUNTING_MODEL)) history = [{"role": "user", "content": [instructor.Image.from_url("https://example.com/x.png")]}] with ( caplog.at_level(logging.WARNING, logger=LOGGER_NAME), patch.object(instructor.Image, "to_openai", side_effect=RuntimeError("serialization failed")), patch.object(agent.history, "get_history", return_value=history), ): serialized = agent._serialize_history_for_token_count() assert serialized == [{"role": "user", "content": [{"type": "text", "text": "[image content]"}]}] assert any("Failed to serialize" in record.getMessage() for record in caplog.records) def test_unknown_content_part_becomes_text(self): agent = _agent(AgentConfig(client=_client(instructor.Mode.TOOLS), model=COUNTING_MODEL)) class UnknownPart: def __str__(self): return "unknown-part-7" history = [{"role": "user", "content": [UnknownPart()]}] with patch.object(agent.history, "get_history", return_value=history): serialized = agent._serialize_history_for_token_count() assert serialized == [{"role": "user", "content": [{"type": "text", "text": "unknown-part-7"}]}] class TestAccountingMatchesPreparedRequests: """The schema must be counted the same way Instructor transmits it.""" @pytest.mark.parametrize("mode", list(instructor.Mode), ids=_mode_id) def test_tools_accounting_matches_prepared_request(self, mode): try: kwargs = _prepared_request(mode) except Exception as e: # Preparing some modes needs their provider SDK (not installed here) or an # Iterable response model (parallel tool calls); those cannot be checked # offline, and the accounting set covers them by construction. pytest.skip(f"{mode.name} cannot be prepared without its provider setup: {type(e).__name__}: {e}") sends_tool_definition = any(param in kwargs for param in TOOL_REQUEST_PARAMS) assert sends_tool_definition == (mode in _TOOL_MODES) def test_anthropic_tools_request_carries_tool_without_schema_prose(self): kwargs = _prepared_request( instructor.Mode.ANTHROPIC_TOOLS, messages=[{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "hello"}], ) assert kwargs["tools"] assert kwargs["tool_choice"] assert not any("json_schema" in str(block.get("text", "")) for block in kwargs["system"]) assert not any("json_schema" in str(message.get("content", "")) for message in kwargs["messages"]) def test_xai_tools_request_carries_schema_as_tool(self): # xAI has no OpenAI-style tools parameter: the schema travels as a tool # definition under Instructor's private _xai_tool kwarg, which the # accounting set must recognize. kwargs = _prepared_request(instructor.Mode.XAI_TOOLS) assert kwargs["_xai_tool"]["name"] == BasicChatOutputSchema.__name__ assert "chat_message" in kwargs["_xai_tool"]["parameters"]["properties"] ``` ### File: atomic-agents/tests/agents/test_atomic_agent.py ```python import json import pytest from unittest.mock import AsyncMock, Mock, call, patch from enum import Enum from pydantic import BaseModel, Field from pydantic import ValidationError import instructor from atomic_agents import ( BaseIOSchema, AtomicAgent, AgentConfig, BasicChatInputSchema, BasicChatOutputSchema, ) from atomic_agents.context import ( ChatHistory, SystemPromptGenerator, BaseDynamicContextProvider, BaseSystemPromptGenerator, ) from atomic_agents.utils.token_counter import TokenCountResult from instructor.dsl.partial import PartialBase @pytest.fixture def mock_instructor(): mock = Mock(spec=instructor.Instructor) # Set up the nested mock structure mock.chat = Mock() mock.chat.completions = Mock() mock.chat.completions.create = Mock(return_value=BasicChatOutputSchema(chat_message="Test output")) # Make create_partial return an iterable mock_response = BasicChatOutputSchema(chat_message="Test output") mock_iter = Mock() mock_iter.__iter__ = Mock(return_value=iter([mock_response])) mock.chat.completions.create_partial.return_value = mock_iter return mock @pytest.fixture def mock_instructor_async(): mock = Mock(spec=instructor.AsyncInstructor) # Configure chat.completions structure mock.chat = Mock() mock.chat.completions = Mock() # Make create method awaitable by using an async function async def mock_create(*args, **kwargs): return BasicChatOutputSchema(chat_message="Test output") mock.chat.completions.create = mock_create # Mock the create_partial method to return an async generator async def mock_create_partial(*args, **kwargs): yield BasicChatOutputSchema(chat_message="Test output") mock.chat.completions.create_partial = mock_create_partial return mock @pytest.fixture def mock_history(): mock = Mock(spec=ChatHistory) mock.get_history.return_value = [] mock.add_message = Mock() mock.copy = Mock(return_value=Mock(spec=ChatHistory)) mock.initialize_turn = Mock() return mock @pytest.fixture def mock_system_prompt_generator(): mock = Mock(spec=SystemPromptGenerator) mock.generate_prompt.return_value = "Mocked system prompt" mock.context_providers = {} return mock @pytest.fixture def agent_config(mock_instructor, mock_history, mock_system_prompt_generator): return AgentConfig( client=mock_instructor, model="gpt-5-mini", history=mock_history, system_prompt_generator=mock_system_prompt_generator, ) @pytest.fixture def agent(agent_config): return AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](agent_config) @pytest.fixture def agent_config_async(mock_instructor_async, mock_history, mock_system_prompt_generator): return AgentConfig( client=mock_instructor_async, model="gpt-5-mini", history=mock_history, system_prompt_generator=mock_system_prompt_generator, ) @pytest.fixture def agent_async(agent_config_async): return AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](agent_config_async) @pytest.fixture def mock_custom_system_prompt_generator(): class MockSystemPromptGenerator(BaseSystemPromptGenerator): def __init__(self, context_providers=None, system_prompt=""): super().__init__(context_providers=context_providers) self.system_prompt = system_prompt def generate_prompt(self) -> str: return self.system_prompt return MockSystemPromptGenerator(system_prompt="Custom Prompt") @pytest.fixture def mock_context_provider(): class MockContextProvider(BaseDynamicContextProvider): def __init__(self, title: str, info: str): super().__init__(title) self._info = info def get_info(self) -> str: return self._info return MockContextProvider(title="Mock Provider", info="Test") def test_initialization(agent, mock_instructor, mock_history, mock_system_prompt_generator): assert agent.client == mock_instructor assert agent.model == "gpt-5-mini" assert agent.history == mock_history assert agent.system_prompt_generator == mock_system_prompt_generator assert "max_tokens" not in agent.model_api_parameters def test_initial_history_with_local_schema_can_be_reset(mock_instructor): class LocalInput(BaseIOSchema): """Input schema created inside an application factory.""" chat_message: str history = ChatHistory() history.add_message("user", LocalInput(chat_message="Initial context")) agent = AtomicAgent[LocalInput, BasicChatOutputSchema](AgentConfig(client=mock_instructor, history=history)) result = agent.run(LocalInput(chat_message="New question")) assert result.chat_message == "Test output" assert agent.history.get_message_count() == 3 agent.reset_history() assert agent.history.get_message_count() == 1 assert isinstance(agent.history.history[0].content, LocalInput) assert agent.history.history[0].content.chat_message == "Initial context" agent.history.history[0].content.chat_message = "Changed" agent.reset_history() assert agent.history.history[0].content.chat_message == "Initial context" # model_api_parameters should have priority over other settings def test_initialization_temperature_priority(mock_instructor, mock_history, mock_system_prompt_generator): config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=mock_history, system_prompt_generator=mock_system_prompt_generator, model_api_parameters={"temperature": 1.0}, ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) assert agent.model_api_parameters["temperature"] == 1.0 def test_initialization_without_temperature(mock_instructor, mock_history, mock_system_prompt_generator): config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=mock_history, system_prompt_generator=mock_system_prompt_generator, model_api_parameters={"temperature": 0.5}, ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) assert agent.model_api_parameters["temperature"] == 0.5 def test_initialization_without_max_tokens(mock_instructor, mock_history, mock_system_prompt_generator): config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=mock_history, system_prompt_generator=mock_system_prompt_generator, model_api_parameters={"max_tokens": 1024}, ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) assert agent.model_api_parameters["max_tokens"] == 1024 def test_run_uses_pydantic_default_strictness_for_enum_output(mock_history, mock_system_prompt_generator): class Topic(Enum): FOOD = "food" OTHER = "other" class EnumOutputSchema(BaseIOSchema): """Output schema with an enum field for validation tests.""" topic: Topic = Field(...) enum_instructor = Mock(spec=instructor.Instructor) enum_instructor.chat = Mock() enum_instructor.chat.completions = Mock() def mock_create(*args, **kwargs): return kwargs["response_model"].model_validate({"topic": "food"}, strict=kwargs["strict"]) enum_instructor.chat.completions.create.side_effect = mock_create config = AgentConfig( client=enum_instructor, model="gpt-5-mini", history=mock_history, system_prompt_generator=mock_system_prompt_generator, ) agent = AtomicAgent[BasicChatInputSchema, EnumOutputSchema](config) result = agent.run() assert result.topic is Topic.FOOD assert enum_instructor.chat.completions.create.call_args.kwargs["strict"] is None def test_run_respects_explicit_strict_override_for_enum_output(mock_history, mock_system_prompt_generator): class Topic(Enum): FOOD = "food" OTHER = "other" class EnumOutputSchema(BaseIOSchema): """Output schema with an enum field for validation tests.""" topic: Topic = Field(...) enum_instructor = Mock(spec=instructor.Instructor) enum_instructor.chat = Mock() enum_instructor.chat.completions = Mock() def mock_create(*args, **kwargs): return kwargs["response_model"].model_validate({"topic": "food"}, strict=kwargs["strict"]) enum_instructor.chat.completions.create.side_effect = mock_create config = AgentConfig( client=enum_instructor, model="gpt-5-mini", history=mock_history, system_prompt_generator=mock_system_prompt_generator, model_api_parameters={"strict": True}, ) agent = AtomicAgent[BasicChatInputSchema, EnumOutputSchema](config) with pytest.raises(ValidationError): agent.run() assert enum_instructor.chat.completions.create.call_args.kwargs["strict"] is True def test_initialization_system_role_equals_developer(mock_instructor, mock_history, mock_system_prompt_generator): config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=mock_history, system_prompt_generator=mock_system_prompt_generator, system_role="developer", model_api_parameters={}, # No temperature specified ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) _ = agent._prepare_messages() assert isinstance(agent.messages, list) and agent.messages[0]["role"] == "developer" def test_initialization_system_role_equals_None(mock_instructor, mock_history, mock_system_prompt_generator): config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=mock_history, system_prompt_generator=mock_system_prompt_generator, system_role=None, model_api_parameters={}, # No temperature specified ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) _ = agent._prepare_messages() assert isinstance(agent.messages, list) and len(agent.messages) == 0 def test_reset_history(agent, mock_history): initial_history = agent.initial_history agent.reset_history() assert agent.history != initial_history mock_history.copy.assert_called_once() def test_get_context_provider(agent, mock_system_prompt_generator): mock_provider = Mock(spec=BaseDynamicContextProvider) mock_system_prompt_generator.context_providers = {"test_provider": mock_provider} result = agent.get_context_provider("test_provider") assert result == mock_provider with pytest.raises(KeyError): agent.get_context_provider("non_existent_provider") def test_register_context_provider(agent, mock_system_prompt_generator): mock_provider = Mock(spec=BaseDynamicContextProvider) agent.register_context_provider("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=BaseDynamicContextProvider) 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.context_providers with pytest.raises(KeyError): agent.unregister_context_provider("non_existent_provider") def test_no_type_parameters(mock_instructor): custom_config = AgentConfig( client=mock_instructor, model="gpt-5-mini", ) custom_agent = AtomicAgent(custom_config) assert custom_agent.input_schema == BasicChatInputSchema assert custom_agent.output_schema == BasicChatOutputSchema def test_custom_input_output_schemas(mock_instructor): class CustomInputSchema(BaseModel): custom_field: str class CustomOutputSchema(BaseModel): result: str custom_config = AgentConfig( client=mock_instructor, model="gpt-5-mini", ) custom_agent = AtomicAgent[CustomInputSchema, CustomOutputSchema](custom_config) assert custom_agent.input_schema == CustomInputSchema assert custom_agent.output_schema == CustomOutputSchema def test_subclass_with_custom_constructor(mock_instructor): """Test that generic types are preserved in subclasses with custom constructors.""" class CustomInputSchema(BaseModel): custom_field: str class CustomOutputSchema(BaseModel): result: str class MyAgent(AtomicAgent[CustomInputSchema, CustomOutputSchema]): def __init__(self, extra_param: str): self.extra_param = extra_param config = AgentConfig( client=mock_instructor, model="gpt-5-mini", ) super().__init__(config) agent = MyAgent("test_value") # These would fail without the __init_subclass__ fix assert agent.input_schema == CustomInputSchema assert agent.output_schema == CustomOutputSchema assert agent.extra_param == "test_value" def test_base_agent_io_str_and_rich(): class TestIO(BaseIOSchema): """TestIO docstring""" field: str test_io = TestIO(field="test") assert str(test_io) == '{"field":"test"}' assert test_io.__rich__() is not None # Just check if it returns something, as we can't easily compare Rich objects def test_base_io_schema_empty_docstring(): with pytest.raises(ValueError, match="must have a non-empty docstring"): class EmptyDocStringSchema(BaseIOSchema): """""" pass def test_base_io_schema_model_json_schema_no_description(): class TestSchema(BaseIOSchema): """Test schema docstring.""" field: str # Mock the superclass model_json_schema to return a schema without a description with patch("pydantic.BaseModel.model_json_schema", return_value={}): schema = TestSchema.model_json_schema() assert "description" in schema assert schema["description"] == "Test schema docstring." def test_run(agent, mock_history): # Use the agent fixture that's already configured correctly mock_input = BasicChatInputSchema(chat_message="Test input") result = agent.run(mock_input) # Assertions assert result.chat_message == "Test output" assert agent.current_user_input == mock_input mock_history.add_message.assert_has_calls([call("user", mock_input), call("assistant", result)]) def test_messages_sync_after_run(mock_instructor, mock_system_prompt_generator): """ Test that agent.messages includes the assistant response after run() completes. Regression test for GitHub issue #194: https://github.com/eigenwise/atomic-agents/issues/194 The issue was that agent.messages only contained the system prompt and user message after run(), while agent.history.get_history() correctly included the assistant response. """ # Use real ChatHistory instead of mock to verify actual message synchronization real_history = ChatHistory() config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=real_history, system_prompt_generator=mock_system_prompt_generator, ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) mock_input = BasicChatInputSchema(chat_message="Test input") result = agent.run(mock_input) # Verify the result is returned correctly assert result.chat_message == "Test output" # Verify agent.messages is in sync with history.get_history() history_messages = agent.history.get_history() # agent.messages should contain: system prompt + history (user + assistant) assert len(agent.messages) == 3, f"Expected 3 messages (system + user + assistant), got {len(agent.messages)}" # First message should be the system prompt assert agent.messages[0]["role"] == "system" # Second message should be the user input assert agent.messages[1]["role"] == "user" # Third message should be the assistant response (the key fix for issue #194) assert agent.messages[2]["role"] == "assistant" # Verify consistency: agent.messages[-2:] should match history.get_history() assert len(history_messages) == 2, f"Expected 2 history messages, got {len(history_messages)}" assert agent.messages[1:] == history_messages def test_run_stream(mock_instructor, mock_history): # Create a AgentConfig with system_role set to None config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=mock_history, system_prompt_generator=None, # No system prompt generator ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) mock_input = BasicChatInputSchema(chat_message="Test input") mock_output = BasicChatOutputSchema(chat_message="Test output") for result in agent.run_stream(mock_input): pass assert result == mock_output assert agent.current_user_input == mock_input mock_history.add_message.assert_has_calls([call("user", mock_input), call("assistant", mock_output)]) class AliasedOutputSchema(BaseIOSchema): """Output schema whose field is set through an alias.""" chat_message: str = Field(..., alias="chatMessage", description="The reply text") def test_run_stream_with_aliased_output_schema(mock_instructor, mock_history): streamed = AliasedOutputSchema(chatMessage="Test output") mock_instructor.chat.completions.create_partial.return_value = iter([streamed]) config = AgentConfig(client=mock_instructor, model="gpt-5-mini", history=mock_history) agent = AtomicAgent[BasicChatInputSchema, AliasedOutputSchema](config) stream = agent.run_stream(BasicChatInputSchema(chat_message="Test input")) with pytest.raises(StopIteration) as stop: while True: next(stream) assert stop.value.value == streamed mock_history.add_message.assert_called_with("assistant", streamed) @pytest.mark.asyncio async def test_run_async(agent_async, mock_history): # Create a mock input mock_input = BasicChatInputSchema(chat_message="Test input") mock_output = BasicChatOutputSchema(chat_message="Test output") # Get response from run_async method response = await agent_async.run_async(mock_input) # Assertions assert response == mock_output assert agent_async.current_user_input == mock_input mock_history.add_message.assert_has_calls([call("user", mock_input), call("assistant", mock_output)]) @pytest.mark.asyncio async def test_run_async_stream(agent_async, mock_history): # Create a mock input mock_input = BasicChatInputSchema(chat_message="Test input") mock_output = BasicChatOutputSchema(chat_message="Test output") responses = [] # Get response from run_async_stream method async for response in agent_async.run_async_stream(mock_input): responses.append(response) # Assertions assert len(responses) == 1 assert responses[0] == mock_output assert agent_async.current_user_input == mock_input # Verify that both user input and assistant response were added to history mock_history.add_message.assert_any_call("user", mock_input) # Create the expected full response content to check full_response_content = agent_async.output_schema(**responses[0].model_dump()) mock_history.add_message.assert_any_call("assistant", full_response_content) @pytest.mark.asyncio async def test_run_async_stream_with_aliased_output_schema(mock_instructor_async, mock_history): streamed = AliasedOutputSchema(chatMessage="Test output") async def mock_create_partial(*args, **kwargs): yield streamed mock_instructor_async.chat.completions.create_partial = mock_create_partial config = AgentConfig(client=mock_instructor_async, model="gpt-5-mini", history=mock_history) agent = AtomicAgent[BasicChatInputSchema, AliasedOutputSchema](config) responses = [response async for response in agent.run_async_stream(BasicChatInputSchema(chat_message="Test input"))] assert responses == [streamed] mock_history.add_message.assert_called_with("assistant", streamed) def test_model_from_chunks_patched(): class TestPartialModel(PartialBase): @classmethod def get_partial_model(cls): class PartialModel(BaseModel): field: str return PartialModel chunks = ['{"field": "hel', 'lo"}'] expected_values = ["hel", "hello"] generator = TestPartialModel.model_from_chunks(chunks) results = [result.field for result in generator] assert results == expected_values @pytest.mark.asyncio async def test_model_from_chunks_async_patched(): class TestPartialModel(PartialBase): @classmethod def get_partial_model(cls): class PartialModel(BaseModel): field: str return PartialModel async def async_gen(): yield '{"field": "hel' yield 'lo"}' expected_values = ["hel", "hello"] generator = TestPartialModel.model_from_chunks_async(async_gen()) results = [] async for result in generator: results.append(result.field) assert results == expected_values # Hook System Tests def test_hook_initialization(agent): """Test that hook system is properly initialized.""" # Verify hook attributes exist and are properly initialized assert hasattr(agent, "_hook_handlers") assert hasattr(agent, "_hooks_enabled") assert isinstance(agent._hook_handlers, dict) assert agent._hooks_enabled is True assert len(agent._hook_handlers) == 0 def test_hook_registration(agent): """Test hook registration and unregistration functionality.""" # Test registration handler_called = [] def test_handler(error): handler_called.append(error) agent.register_hook("parse:error", test_handler) # Verify internal storage assert "parse:error" in agent._hook_handlers assert test_handler in agent._hook_handlers["parse:error"] # Test unregistration agent.unregister_hook("parse:error", test_handler) assert test_handler not in agent._hook_handlers["parse:error"] def test_hook_registration_with_instructor_client(mock_instructor): """Test that hooks are registered with instructor client when available.""" # Add hook methods to mock instructor mock_instructor.on = Mock() mock_instructor.off = Mock() mock_instructor.clear = Mock() config = AgentConfig(client=mock_instructor, model="gpt-5-mini") agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) def test_handler(error): pass # Test registration delegates to instructor client agent.register_hook("parse:error", test_handler) mock_instructor.on.assert_called_once_with("parse:error", test_handler) # Test unregistration delegates to instructor client agent.unregister_hook("parse:error", test_handler) mock_instructor.off.assert_called_once_with("parse:error", test_handler) def test_multiple_hook_handlers(agent): """Test multiple handlers for the same event.""" handler1_calls = [] handler2_calls = [] def handler1(error): handler1_calls.append(error) def handler2(error): handler2_calls.append(error) # Register multiple handlers agent.register_hook("parse:error", handler1) agent.register_hook("parse:error", handler2) # Verify both are registered assert len(agent._hook_handlers["parse:error"]) == 2 assert handler1 in agent._hook_handlers["parse:error"] assert handler2 in agent._hook_handlers["parse:error"] # Test dispatch to both handlers test_error = Exception("test error") agent._dispatch_hook("parse:error", test_error) assert len(handler1_calls) == 1 assert len(handler2_calls) == 1 assert handler1_calls[0] is test_error assert handler2_calls[0] is test_error def test_hook_clear_specific_event(agent): """Test clearing hooks for a specific event.""" def handler1(): pass def handler2(): pass # Register handlers for different events agent.register_hook("parse:error", handler1) agent.register_hook("completion:error", handler2) # Clear specific event agent.clear_hooks("parse:error") # Verify only parse:error was cleared assert len(agent._hook_handlers["parse:error"]) == 0 assert handler2 in agent._hook_handlers["completion:error"] def test_hook_clear_all_events(agent): """Test clearing all hooks.""" def handler1(): pass def handler2(): pass # Register handlers for different events agent.register_hook("parse:error", handler1) agent.register_hook("completion:error", handler2) # Clear all hooks agent.clear_hooks() # Verify all hooks are cleared assert len(agent._hook_handlers) == 0 def test_hook_enable_disable(agent): """Test hook enable/disable functionality.""" # Test initial state assert agent.hooks_enabled is True # Test disable agent.disable_hooks() assert agent.hooks_enabled is False assert agent._hooks_enabled is False # Test enable agent.enable_hooks() assert agent.hooks_enabled is True assert agent._hooks_enabled is True def test_hook_dispatch_when_disabled(agent): """Test that hooks don't execute when disabled.""" handler_called = [] def test_handler(error): handler_called.append(error) agent.register_hook("parse:error", test_handler) # Disable hooks agent.disable_hooks() # Dispatch should not call handler agent._dispatch_hook("parse:error", Exception("test")) assert len(handler_called) == 0 # Re-enable and test agent.enable_hooks() agent._dispatch_hook("parse:error", Exception("test")) assert len(handler_called) == 1 def test_hook_error_isolation(agent): """Test that hook handler errors don't interrupt main flow.""" good_handler_called = [] def bad_handler(error): raise RuntimeError("Handler error") def good_handler(error): good_handler_called.append(error) # Register both handlers agent.register_hook("test:event", bad_handler) agent.register_hook("test:event", good_handler) # Dispatch should not raise exception with patch("logging.getLogger") as mock_logger: mock_log = Mock() mock_logger.return_value = mock_log agent._dispatch_hook("test:event", Exception("test")) # Verify error was logged mock_log.warning.assert_called_once() # Verify good handler still executed assert len(good_handler_called) == 1 def test_hook_dispatch_nonexistent_event(agent): """Test dispatching to nonexistent event.""" # Should not raise exception agent._dispatch_hook("nonexistent:event", Exception("test")) def test_hook_unregister_nonexistent_handler(agent): """Test unregistering handler that doesn't exist.""" def test_handler(): pass # Should not raise exception agent.unregister_hook("parse:error", test_handler) def test_agent_initialization_includes_hooks(mock_instructor, mock_history, mock_system_prompt_generator): """Test that agent initialization properly sets up hook system.""" config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=mock_history, system_prompt_generator=mock_system_prompt_generator, ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) # Verify hook system is initialized assert hasattr(agent, "_hook_handlers") assert hasattr(agent, "_hooks_enabled") assert agent._hooks_enabled is True assert isinstance(agent._hook_handlers, dict) assert len(agent._hook_handlers) == 0 # Verify hook management methods exist assert hasattr(agent, "register_hook") assert hasattr(agent, "unregister_hook") assert hasattr(agent, "clear_hooks") assert hasattr(agent, "enable_hooks") assert hasattr(agent, "disable_hooks") assert hasattr(agent, "hooks_enabled") assert hasattr(agent, "_dispatch_hook") def test_backward_compatibility_no_breaking_changes(mock_instructor, mock_history, mock_system_prompt_generator): """Test that hook system addition doesn't break existing functionality.""" # Ensure mock_history.get_history() returns an empty list mock_history.get_history.return_value = [] # Ensure the copy method returns a properly configured mock copied_mock = Mock(spec=ChatHistory) copied_mock.get_history.return_value = [] mock_history.copy.return_value = copied_mock config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=mock_history, system_prompt_generator=mock_system_prompt_generator, ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) # Test that all existing attributes still exist and work assert agent.client == mock_instructor assert agent.model == "gpt-5-mini" assert agent.history == mock_history assert agent.system_prompt_generator == mock_system_prompt_generator # Test that existing methods still work # Note: reset_history() changes the history object, so we skip it to focus on core functionality # Properties should work assert agent.input_schema == BasicChatInputSchema assert agent.output_schema == BasicChatOutputSchema # Run method should work (with hooks enabled by default) user_input = BasicChatInputSchema(chat_message="test") response = agent.run(user_input) # Verify the response is valid assert response is not None # Verify the call was made correctly mock_instructor.chat.completions.create.assert_called() # Test context provider methods still work from atomic_agents.context import BaseDynamicContextProvider class TestProvider(BaseDynamicContextProvider): def get_info(self): return "test" provider = TestProvider(title="Test") agent.register_context_provider("test", provider) retrieved = agent.get_context_provider("test") assert retrieved == provider agent.unregister_context_provider("test") # Should raise KeyError for non-existent provider with pytest.raises(KeyError): agent.get_context_provider("test") # Token Counting Tests @patch("atomic_agents.agents.atomic_agent.get_token_counter") def test_get_context_token_count_basic(mock_get_token_counter, mock_instructor, mock_history, mock_system_prompt_generator): """Test basic token counting functionality.""" mock_history.get_history.return_value = [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}, ] mock_system_prompt_generator.generate_prompt.return_value = "You are a helpful assistant." mock_counter_instance = Mock() mock_get_token_counter.return_value = mock_counter_instance mock_counter_instance.count_context.return_value = TokenCountResult( total=100, system_prompt=30, history=70, tools=0, model="gpt-5-mini", max_tokens=8192, utilization=0.0122 ) config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=mock_history, system_prompt_generator=mock_system_prompt_generator, ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) result = agent.get_context_token_count() assert result.total == 100 assert result.system_prompt == 30 assert result.history == 70 assert result.tools == 0 assert result.model == "gpt-5-mini" assert result.max_tokens == 8192 assert result.utilization is not None # Should have utilization since max_tokens is known mock_counter_instance.count_context.assert_called_once() @patch("atomic_agents.agents.atomic_agent.get_token_counter") def test_get_context_token_count_includes_schema_overhead( mock_get_token_counter, mock_instructor, mock_history, mock_system_prompt_generator ): """Test that token counting includes the output schema overhead for JSON mode.""" from instructor import Mode mock_history.get_history.return_value = [] mock_system_prompt_generator.generate_prompt.return_value = "System prompt" mock_counter_instance = Mock() mock_get_token_counter.return_value = mock_counter_instance mock_counter_instance.count_context.return_value = TokenCountResult( total=50, system_prompt=50, history=0, tools=0, model="gpt-5-mini" ) config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=mock_history, system_prompt_generator=mock_system_prompt_generator, mode=Mode.JSON, # JSON mode appends schema to system message ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) agent.get_context_token_count() # Verify count_context was called with system message that includes schema call_args = mock_counter_instance.count_context.call_args system_messages = call_args.kwargs["system_messages"] assert len(system_messages) == 1 # System message should contain both the prompt AND the schema (for JSON mode) assert "System prompt" in system_messages[0]["content"] assert "chat_message" in system_messages[0]["content"] # Schema field from BasicChatOutputSchema @patch("atomic_agents.agents.atomic_agent.get_token_counter") def test_get_context_token_count_no_system_prompt(mock_get_token_counter, mock_instructor, mock_history): """Test token counting when system_role is None (schema still included for JSON mode).""" from instructor import Mode mock_history.get_history.return_value = [{"role": "user", "content": "Hello"}] mock_counter_instance = Mock() mock_get_token_counter.return_value = mock_counter_instance mock_counter_instance.count_context.return_value = TokenCountResult( total=25, system_prompt=20, history=5, tools=0, model="gpt-5-mini" ) config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=mock_history, system_role=None, # No system prompt mode=Mode.JSON, # JSON mode to test schema in system message ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) agent.get_context_token_count() # Even without system prompt, schema should be included (for JSON mode) call_args = mock_counter_instance.count_context.call_args system_messages = call_args.kwargs["system_messages"] assert len(system_messages) == 1 # Schema is added as system message assert "chat_message" in system_messages[0]["content"] # Schema content @patch("atomic_agents.agents.atomic_agent.get_token_counter") def test_get_context_token_count_uses_configured_model( mock_get_token_counter, mock_instructor, mock_history, mock_system_prompt_generator ): """Test that token counting uses the agent's configured model.""" mock_history.get_history.return_value = [] mock_counter_instance = Mock() mock_get_token_counter.return_value = mock_counter_instance mock_counter_instance.count_context.return_value = TokenCountResult( total=15, system_prompt=15, history=0, tools=0, model="claude-3-opus-20240229" ) config = AgentConfig( client=mock_instructor, model="claude-3-opus-20240229", # Different model history=mock_history, system_prompt_generator=mock_system_prompt_generator, ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) agent.get_context_token_count() call_args = mock_counter_instance.count_context.call_args assert call_args.kwargs["model"] == "claude-3-opus-20240229" @patch("atomic_agents.agents.atomic_agent.get_token_counter") def test_get_context_token_count_returns_token_count_result(mock_get_token_counter, agent): """Test that get_context_token_count returns a TokenCountResult.""" mock_counter_instance = Mock() mock_get_token_counter.return_value = mock_counter_instance mock_counter_instance.count_context.return_value = TokenCountResult( total=50, system_prompt=20, history=30, tools=0, model="gpt-5-mini" ) result = agent.get_context_token_count() assert isinstance(result, TokenCountResult) assert hasattr(result, "total") assert hasattr(result, "system_prompt") assert hasattr(result, "history") assert hasattr(result, "tools") assert hasattr(result, "model") assert hasattr(result, "max_tokens") assert hasattr(result, "utilization") @patch("atomic_agents.agents.atomic_agent.get_token_counter") def test_get_context_token_count_hook_dispatch(mock_get_token_counter, agent): """Test that token:counted hook is dispatched.""" mock_counter_instance = Mock() mock_get_token_counter.return_value = mock_counter_instance expected_result = TokenCountResult(total=100, system_prompt=30, history=70, tools=0, model="gpt-5-mini") mock_counter_instance.count_context.return_value = expected_result hook_called = [] def hook_handler(result): hook_called.append(result) agent.register_hook("token:counted", hook_handler) agent.get_context_token_count() assert len(hook_called) == 1 assert hook_called[0] == expected_result assert hook_called[0].total == 100 @patch("atomic_agents.agents.atomic_agent.get_token_counter") def test_get_context_token_count_hook_not_called_when_disabled(mock_get_token_counter, agent): """Test that token:counted hook is not called when hooks are disabled.""" mock_counter_instance = Mock() mock_get_token_counter.return_value = mock_counter_instance mock_counter_instance.count_context.return_value = TokenCountResult( total=100, system_prompt=30, history=70, tools=0, model="gpt-5-mini" ) hook_called = [] def hook_handler(result): hook_called.append(result) agent.register_hook("token:counted", hook_handler) agent.disable_hooks() agent.get_context_token_count() assert len(hook_called) == 0 @patch("atomic_agents.agents.atomic_agent.get_token_counter") def test_get_context_token_count_multimodal_content(mock_get_token_counter, mock_instructor): """Test that multimodal content is properly serialized for token counting.""" from instructor.processing.multimodal import Image mock_counter_instance = Mock() mock_get_token_counter.return_value = mock_counter_instance mock_counter_instance.count_context.return_value = TokenCountResult( total=150, system_prompt=50, history=100, tools=0, model="gpt-4-vision-preview" ) # Create a multimodal input schema class MultimodalInputSchema(BaseIOSchema): """Input with image.""" text: str = Field(..., description="Text input") image: instructor.Image = Field(..., description="Image to analyze") # Create agent with multimodal schema config = AgentConfig( client=mock_instructor, model="gpt-4-vision-preview", ) agent = AtomicAgent[MultimodalInputSchema, BasicChatOutputSchema](config) # Add multimodal message to history test_image = Image(source="https://example.com/test.png", media_type="image/png") multimodal_input = MultimodalInputSchema(text="Describe this image", image=test_image) agent.history.add_message("user", multimodal_input) # Get token count agent.get_context_token_count() # Verify count_context was called assert mock_counter_instance.count_context.called # Get the history_messages passed to count_context call_args = mock_counter_instance.count_context.call_args history_messages = call_args.kwargs["history_messages"] # Verify multimodal content was serialized properly assert len(history_messages) == 1 assert history_messages[0]["role"] == "user" # Content should be a list with text and image content = history_messages[0]["content"] assert isinstance(content, list) # Should have text and image entries content_types = [item.get("type") for item in content] assert "text" in content_types assert "image_url" in content_types # Verify image was converted to OpenAI format image_entry = next(item for item in content if item.get("type") == "image_url") assert "image_url" in image_entry assert image_entry["image_url"]["url"] == "https://example.com/test.png" @patch("atomic_agents.agents.atomic_agent.get_token_counter") def test_get_context_token_count_image_conversion_failure(mock_get_token_counter, mock_instructor, caplog): from logging import WARNING from instructor.processing.multimodal import Image class ImageInputSchema(BaseIOSchema): """Input with text and an image.""" text: str image: Image image = Image(source="https://example.com/test.png", media_type="image/png") history = ChatHistory() history.add_message("system", BasicChatInputSchema(chat_message="Earlier context")) history.add_message("user", ImageInputSchema(text="Describe this image", image=image)) history.add_message("assistant", BasicChatOutputSchema(chat_message="Later reply")) agent = AtomicAgent[ImageInputSchema, BasicChatOutputSchema](AgentConfig(client=mock_instructor, history=history)) expected_result = TokenCountResult(total=150, system_prompt=50, history=100, tools=0, model=agent.model) mock_get_token_counter.return_value.count_context.return_value = expected_result with patch.object( Image, "to_openai", autospec=True, side_effect=RuntimeError("controlled conversion failure") ) as convert_image: result = agent.get_context_token_count() convert_image.assert_called_once_with(image, instructor.Mode.JSON) assert result == expected_result assert caplog.record_tuples == [ ( "atomic_agents.agents.atomic_agent", WARNING, "Failed to serialize Image for token counting: controlled conversion failure. Using placeholder for estimation.", ) ] assert mock_get_token_counter.return_value.count_context.call_args.kwargs["history_messages"] == [ {"role": "system", "content": '{"chat_message":"Earlier context"}'}, { "role": "user", "content": [ {"type": "text", "text": '{"text":"Describe this image"}'}, {"type": "text", "text": "[image content]"}, ], }, {"role": "assistant", "content": '{"chat_message":"Later reply"}'}, ] @patch("atomic_agents.agents.atomic_agent.get_token_counter") def test_get_context_token_count_custom_history_stringifies_unusual_content(mock_get_token_counter, mock_instructor): from copy import deepcopy from decimal import Decimal class CustomHistory(ChatHistory): def get_history(self) -> list[dict[str, str | list[str | Decimal]]]: return [ {"role": "tool", "content": ["Before value", Decimal("12.50"), "After value"]}, {"role": "assistant", "content": "Remaining message"}, ] def copy(self) -> "CustomHistory": return deepcopy(self) history = CustomHistory() agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](AgentConfig(client=mock_instructor, history=history)) expected_result = TokenCountResult(total=150, system_prompt=50, history=100, tools=0, model=agent.model) mock_get_token_counter.return_value.count_context.return_value = expected_result result = agent.get_context_token_count() assert result == expected_result assert mock_get_token_counter.return_value.count_context.call_args.kwargs["history_messages"] == [ { "role": "tool", "content": [ {"type": "text", "text": "Before value"}, {"type": "text", "text": "12.50"}, {"type": "text", "text": "After value"}, ], }, {"role": "assistant", "content": "Remaining message"}, ] def test_serialize_history_for_token_count_video_placeholder(agent, mock_history): video_part = {"type": "video_url", "video_url": {"url": "https://example.com/clip.mp4"}} mock_history.get_history.return_value = [{"role": "user", "content": ['{"prompt":"Summarize the video"}', video_part]}] serialized = agent._serialize_history_for_token_count() assert serialized[0]["content"] == [ {"type": "text", "text": '{"prompt":"Summarize the video"}'}, {"type": "text", "text": "[video_url content]"}, ] # --- Tests for tool_result_role and Gemini system message remapping (issue #221) --- def test_tool_result_role_defaults_to_system_for_openai(mock_instructor, mock_system_prompt_generator): config = AgentConfig( client=mock_instructor, model="gpt-5-mini", system_prompt_generator=mock_system_prompt_generator, assistant_role="assistant", ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) assert agent.tool_result_role == "system" def test_tool_result_role_defaults_to_user_for_gemini(mock_instructor, mock_system_prompt_generator): config = AgentConfig( client=mock_instructor, model="gemini-2.0-flash", system_prompt_generator=mock_system_prompt_generator, assistant_role="model", ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) assert agent.tool_result_role == "user" def test_tool_result_role_explicit_override(mock_instructor, mock_system_prompt_generator): config = AgentConfig( client=mock_instructor, model="gemini-2.0-flash", system_prompt_generator=mock_system_prompt_generator, assistant_role="model", tool_result_role="system", ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) assert agent.tool_result_role == "system" def test_add_tool_result_uses_correct_role(mock_instructor, mock_system_prompt_generator): history = ChatHistory() config = AgentConfig( client=mock_instructor, model="gemini-2.0-flash", system_prompt_generator=mock_system_prompt_generator, assistant_role="model", history=history, ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) content = BasicChatInputSchema(chat_message="Tool result data") agent.add_tool_result(content) messages = history.get_history() assert len(messages) == 1 assert messages[0]["role"] == "user" def test_add_tool_result_uses_system_for_openai(mock_instructor, mock_system_prompt_generator): history = ChatHistory() config = AgentConfig( client=mock_instructor, model="gpt-5-mini", system_prompt_generator=mock_system_prompt_generator, history=history, ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) content = BasicChatInputSchema(chat_message="Tool result data") agent.add_tool_result(content) messages = history.get_history() assert len(messages) == 1 assert messages[0]["role"] == "system" def test_prepare_messages_remaps_system_to_user_for_gemini(mock_instructor, mock_system_prompt_generator): history = ChatHistory() config = AgentConfig( client=mock_instructor, model="gemini-2.0-flash", system_prompt_generator=mock_system_prompt_generator, assistant_role="model", history=history, ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) # Simulate legacy pattern: add_message("system", ...) directly on history history.add_message("system", BasicChatInputSchema(chat_message="Tool output")) agent._prepare_messages() # The system prompt message (first) should keep its role, # but mid-conversation "system" messages in history should be remapped to "user" history_messages = [m for m in agent.messages if m.get("content") != mock_system_prompt_generator.generate_prompt()] assert len(history_messages) == 1 assert history_messages[0]["role"] == "user" def test_prepare_messages_keeps_system_for_openai(mock_instructor, mock_system_prompt_generator): history = ChatHistory() config = AgentConfig( client=mock_instructor, model="gpt-5-mini", system_prompt_generator=mock_system_prompt_generator, history=history, ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) history.add_message("system", BasicChatInputSchema(chat_message="Tool output")) agent._prepare_messages() history_messages = [m for m in agent.messages if m.get("content") != mock_system_prompt_generator.generate_prompt()] assert len(history_messages) == 1 assert history_messages[0]["role"] == "system" # --- max_context_tokens and _trim_context tests --- @pytest.mark.asyncio @pytest.mark.parametrize("method", ["run", "run_stream", "run_async", "run_async_stream"]) @pytest.mark.parametrize("input_size", [15, 40]) @patch("atomic_agents.agents.atomic_agent.get_token_counter") async def test_context_budget_includes_new_input( mock_get_token_counter, method, input_size, mock_instructor, mock_instructor_async, mock_system_prompt_generator ): """All completion APIs budget the actual request and preserve the latest input.""" history = ChatHistory() history.add_message("user", BasicChatInputSchema(chat_message="o" * 20)) client = mock_instructor_async if "async" in method else mock_instructor response = BasicChatOutputSchema(chat_message="Response") if "async" in method: client.chat.completions.create = AsyncMock(return_value=response) async def partials(**kwargs): yield response client.chat.completions.create_partial = Mock(side_effect=partials) def count_context(model, system_messages, history_messages, tools): history_tokens = sum(len(json.loads(message["content"])["chat_message"]) for message in history_messages) return TokenCountResult(total=10 + history_tokens, system_prompt=10, history=history_tokens, tools=0, model=model) mock_get_token_counter.return_value.count_context.side_effect = count_context agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema]( AgentConfig( client=client, model="gpt-5-mini", history=history, system_prompt_generator=mock_system_prompt_generator, max_context_tokens=30, ) ) user_input = BasicChatInputSchema(chat_message="n" * input_size) async def invoke(): if method == "run": agent.run(user_input) elif method == "run_stream": list(agent.run_stream(user_input)) elif method == "run_async": await agent.run_async(user_input) else: async for _ in agent.run_async_stream(user_input): pass if input_size == 40: with pytest.raises(ValueError, match="max_context_tokens"): await invoke() client.chat.completions.create.assert_not_called() client.chat.completions.create_partial.assert_not_called() assert history.history[-1].content == user_input else: await invoke() completion = client.chat.completions.create_partial if "stream" in method else client.chat.completions.create request_history = completion.call_args.kwargs["messages"][1:] assert request_history == [{"role": "user", "content": user_input.model_dump_json()}] assert count_context(agent.model, [], request_history, None).total <= agent.max_context_tokens def test_trim_context_no_op_when_max_context_tokens_unset(mock_instructor, mock_system_prompt_generator): """When max_context_tokens is None, _trim_context returns without any action.""" history = ChatHistory() history.add_message("user", BasicChatInputSchema(chat_message="Hello")) config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=history, system_prompt_generator=mock_system_prompt_generator, # max_context_tokens not set — default None ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) # Should not raise, should not modify history agent._trim_context() assert history.get_message_count() == 1 @patch("atomic_agents.agents.atomic_agent.get_token_counter") def test_trim_context_no_op_when_within_limit(mock_get_token_counter, mock_instructor, mock_system_prompt_generator): """When context is within max_context_tokens, history is not modified.""" history = ChatHistory() history.add_message("user", BasicChatInputSchema(chat_message="Hello")) mock_counter_instance = Mock() mock_get_token_counter.return_value = mock_counter_instance mock_counter_instance.count_context.return_value = TokenCountResult( total=100, system_prompt=30, history=70, tools=0, model="gpt-5-mini", max_tokens=8192, utilization=0.01 ) config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=history, system_prompt_generator=mock_system_prompt_generator, max_context_tokens=200, # well above current 100 ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) agent._trim_context() # History untouched assert history.get_message_count() == 1 # delete_turn_id never called mock_counter_instance.count_context.assert_called_once() @patch("atomic_agents.agents.atomic_agent.get_token_counter") def test_trim_context_trims_oldest_turn_when_over_limit(mock_get_token_counter, mock_instructor, mock_system_prompt_generator): """When context exceeds max_context_tokens, oldest turn is removed.""" history = ChatHistory() # Turn 1 (user message) history.initialize_turn() history.add_message("user", BasicChatInputSchema(chat_message="First request")) # Turn 2 (user message) history.initialize_turn() history.add_message("user", BasicChatInputSchema(chat_message="Second request")) turn1_id = history.history[0].turn_id mock_counter_instance = Mock() mock_get_token_counter.return_value = mock_counter_instance # First call: over limit; Second call: within limit after one turn removed mock_counter_instance.count_context.side_effect = [ TokenCountResult( total=500, system_prompt=100, history=400, tools=0, model="gpt-5-mini", max_tokens=8192, utilization=0.06 ), TokenCountResult( total=150, system_prompt=100, history=50, tools=0, model="gpt-5-mini", max_tokens=8192, utilization=0.018 ), ] config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=history, system_prompt_generator=mock_system_prompt_generator, max_context_tokens=200, ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) agent._trim_context() # Turn 1 should be gone; Turn 2 remains assert history.get_message_count() == 1 remaining = history.history[0] assert remaining.content.chat_message == "Second request" assert remaining.turn_id != turn1_id @patch("atomic_agents.agents.atomic_agent.get_token_counter") def test_trim_context_raises_when_single_turn_exceeds_limit( mock_get_token_counter, mock_instructor, mock_system_prompt_generator ): """When even one turn exceeds max_context_tokens, ValueError is raised.""" history = ChatHistory() history.initialize_turn() history.add_message("user", BasicChatInputSchema(chat_message="Very long message that exceeds the limit")) mock_counter_instance = Mock() mock_get_token_counter.return_value = mock_counter_instance # Even after trimming all turns, still over limit mock_counter_instance.count_context.return_value = TokenCountResult( total=500, system_prompt=400, history=100, tools=0, model="gpt-5-mini", max_tokens=8192, utilization=0.06 ) config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=history, system_prompt_generator=mock_system_prompt_generator, max_context_tokens=100, # artificially low ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) with pytest.raises(ValueError, match="max_context_tokens.*smaller than the minimum"): agent._trim_context() @patch("atomic_agents.agents.atomic_agent.get_token_counter") def test_trim_context_preserves_new_user_message_in_run(mock_get_token_counter, mock_instructor, mock_system_prompt_generator): """Old turns can be trimmed while the new user input is preserved.""" history = ChatHistory() # Existing turn in history history.initialize_turn() history.add_message("user", BasicChatInputSchema(chat_message="Old message")) mock_counter_instance = Mock() mock_get_token_counter.return_value = mock_counter_instance mock_counter_instance.count_context.side_effect = [ TokenCountResult( total=500, system_prompt=100, history=400, tools=0, model="gpt-5-mini", max_tokens=8192, utilization=0.06 ), TokenCountResult( total=50, system_prompt=100, history=0, tools=0, model="gpt-5-mini", max_tokens=8192, utilization=0.006 ), ] config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=history, system_prompt_generator=mock_system_prompt_generator, max_context_tokens=200, ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) # Mock chat completion mock_instructor.chat.completions.create.return_value = BasicChatOutputSchema(chat_message="Response") agent.run(BasicChatInputSchema(chat_message="New message")) # The old turn was trimmed and the new message was preserved. assert history.get_message_count() == 2 # assistant response + new user message assert history.history[0].content.chat_message == "New message" # new message is first # --- Test BaseSystemPromptGenerator integration --- def test_custom_system_prompt_generator_reaches_agent( mock_instructor, mock_custom_system_prompt_generator, mock_context_provider ): agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema]( AgentConfig( client=mock_instructor, model="gpt-5-mini", system_prompt_generator=mock_custom_system_prompt_generator, ) ) agent.register_context_provider("test_provider", mock_context_provider) assert agent.system_prompt_generator == mock_custom_system_prompt_generator assert agent._build_system_messages() == [{"content": "Custom Prompt", "role": "system"}] assert agent.system_prompt_generator.context_providers == {"test_provider": mock_context_provider} ``` ### File: atomic-agents/tests/agents/test_config_defaults_typing.py ```python from typing import TypedDict import instructor from openai import OpenAI from atomic_agents import AgentConfig class ModelSettings(TypedDict): model: str def test_optional_agent_settings_have_defaults() -> None: settings: ModelSettings = {"model": "test-model"} with OpenAI(api_key="synthetic-test-key") as provider_client: client = instructor.from_openai(provider_client) direct_configuration = AgentConfig(client=client, model="test-model") unpacked_configuration = AgentConfig(client=client, **settings) assert (direct_configuration.model_api_parameters, direct_configuration.max_context_tokens) == (None, None) assert (unpacked_configuration.model_api_parameters, unpacked_configuration.max_context_tokens) == (None, None) def test_default_model_pins_string_gpt_5_6_terra() -> None: with OpenAI(api_key="synthetic-test-key") as provider_client: client = instructor.from_openai(provider_client) configuration = AgentConfig(client=client) assert configuration.model == "gpt-5.6-terra" ``` ### File: atomic-agents/tests/agents/test_generic_agent_schemas.py ```python """Concrete schema bindings survive reusable generic agent subclasses.""" import json import httpx import instructor import openai import pytest from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema, BasicChatInputSchema, BasicChatOutputSchema class StageInput(BaseIOSchema): """Input to a composed stage.""" text: str class FirstOutput(BaseIOSchema): """Output from the first stage.""" answer: str class SecondOutput(BaseIOSchema): """Output from the second stage.""" answer: str class ConcreteAgent(AtomicAgent[StageInput, FirstOutput]): def __init__(self, config): super().__init__(config) class ConcreteChild(ConcreteAgent): pass class OutputTemplate[T: BaseIOSchema](AtomicAgent[StageInput, T]): def __init__(self, config): super().__init__(config) class TwoSchemaTemplate[I: BaseIOSchema, O: BaseIOSchema](AtomicAgent[I, O]): pass class ReorderedTemplate[O: BaseIOSchema, I: BaseIOSchema](AtomicAgent[I, O]): pass class ForwardTemplate[T: BaseIOSchema](OutputTemplate[T]): pass class FinalAgent(ForwardTemplate[FirstOutput]): pass class UntypedAgent(AtomicAgent): pass @pytest.fixture def offline_client(): requests = [] def respond(request): body = json.loads(request.content) function = body["tools"][0]["function"] assert function["name"] in {"FirstOutput", "SecondOutput"} assert set(function["parameters"]["properties"]) == {"answer"} requests.append(body) return httpx.Response( 200, json={ "id": "schema-binding-test", "object": "chat.completion", "created": 0, "model": "offline-fixture", "choices": [ { "index": 0, "finish_reason": "tool_calls", "message": { "role": "assistant", "content": None, "tool_calls": [ { "id": "call_test", "type": "function", "function": { "name": function["name"], "arguments": json.dumps({"answer": "literal fixture result"}), }, } ], }, } ], }, ) with openai.OpenAI( api_key="test-placeholder", base_url="http://localhost/v1", max_retries=0, http_client=httpx.Client(transport=httpx.MockTransport(respond)), ) as provider: yield instructor.from_openai(provider), requests @pytest.mark.parametrize( "factory, expected_output", [ pytest.param(AtomicAgent[StageInput, FirstOutput], FirstOutput, id="direct"), pytest.param(ConcreteAgent, FirstOutput, id="concrete-custom-constructor"), pytest.param(ConcreteChild, FirstOutput, id="concrete-descendant"), pytest.param(OutputTemplate[FirstOutput], FirstOutput, id="single-output-parameter"), pytest.param(OutputTemplate[SecondOutput], SecondOutput, id="other-output-specialization"), pytest.param(TwoSchemaTemplate[StageInput, FirstOutput], FirstOutput, id="two-parameters"), pytest.param(ReorderedTemplate[FirstOutput, StageInput], FirstOutput, id="reordered-parameters"), pytest.param(ForwardTemplate[FirstOutput], FirstOutput, id="forwarded-parameter"), pytest.param(FinalAgent, FirstOutput, id="concrete-template-descendant"), ], ) def test_specialized_agents_use_concrete_schemas(offline_client, factory, expected_output): client, requests = offline_client agent = factory(AgentConfig(client=client, model="offline-fixture")) assert agent.input_schema is StageInput assert agent.output_schema is expected_output result = agent.run(StageInput(text="ordinary stage input")) assert isinstance(result, expected_output) assert result.answer == "literal fixture result" assert len(requests) == 1 assert agent.history.history[-1].content is result def test_specializations_keep_their_own_bindings(offline_client): client, _ = offline_client config = AgentConfig(client=client) first = OutputTemplate[FirstOutput](config) second = OutputTemplate[SecondOutput](config) assert first.output_schema is FirstOutput assert second.output_schema is SecondOutput assert first.output_schema is FirstOutput assert FinalAgent(config).output_schema is FirstOutput assert second.output_schema is SecondOutput @pytest.mark.parametrize("factory", [AtomicAgent, UntypedAgent]) def test_untyped_agents_keep_default_schemas(offline_client, factory): client, _ = offline_client agent = factory(AgentConfig(client=client)) assert agent.input_schema is BasicChatInputSchema assert agent.output_schema is BasicChatOutputSchema ``` ### File: atomic-agents/tests/agents/test_instructor_transport.py ```python """Drive a real Instructor client through a mocked HTTP transport. The rest of the suite mocks the Instructor client itself, so it never sees how Instructor wraps the response model or which keyword arguments it forwards into streaming. Both broke on Instructor 1.17 without any mocked test noticing. """ import json import httpx import instructor import openai import pytest from atomic_agents import AgentConfig, AtomicAgent, BasicChatInputSchema, BasicChatOutputSchema ASSISTANT_ARGUMENTS = json.dumps({"chat_message": "hi there"}) def tool_call_chunk(arguments_piece: str) -> str: delta = { "tool_calls": [ { "index": 0, "id": "call-1", "type": "function", "function": {"name": "BasicChatOutputSchema", "arguments": arguments_piece}, } ] } chunk = {"id": "chunk", "object": "chat.completion.chunk", "created": 0, "model": "m"} chunk["choices"] = [{"index": 0, "delta": delta, "finish_reason": None}] return "data: " + json.dumps(chunk) + "\n\n" def respond_with_tool_call(request: httpx.Request) -> httpx.Response: body = json.loads(request.content) if body.get("stream"): frames = tool_call_chunk(ASSISTANT_ARGUMENTS[:10]) + tool_call_chunk(ASSISTANT_ARGUMENTS[10:]) + "data: [DONE]\n\n" return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=frames.encode()) message = { "role": "assistant", "content": None, "tool_calls": [ { "id": "call-1", "type": "function", "function": {"name": "BasicChatOutputSchema", "arguments": ASSISTANT_ARGUMENTS}, } ], } completion = {"id": "completion", "object": "chat.completion", "created": 0, "model": "m"} completion["choices"] = [{"index": 0, "message": message, "finish_reason": "tool_calls"}] completion["usage"] = {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} return httpx.Response(200, json=completion) def agent_over_mock_transport(**model_api_parameters) -> AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema]: raw_client = openai.OpenAI(api_key="test", http_client=httpx.Client(transport=httpx.MockTransport(respond_with_tool_call))) config = AgentConfig(client=instructor.from_openai(raw_client), model="m", model_api_parameters=model_api_parameters) return AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config=config) def async_agent_over_mock_transport(**model_api_parameters: bool) -> AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema]: raw_client = openai.AsyncOpenAI( api_key="test", http_client=httpx.AsyncClient(transport=httpx.MockTransport(respond_with_tool_call)) ) config = AgentConfig(client=instructor.from_openai(raw_client), model="m", model_api_parameters=model_api_parameters) return AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config=config) def test_run_parses_the_tool_call_through_instructor(): response = agent_over_mock_transport().run(BasicChatInputSchema(chat_message="hey")) assert response.chat_message == "hi there" @pytest.mark.asyncio async def test_run_async_parses_the_tool_call_through_instructor(): response = await async_agent_over_mock_transport().run_async(BasicChatInputSchema(chat_message="hey")) assert response.chat_message == "hi there" def test_run_stream_yields_partials_then_the_complete_response(): partials = list(agent_over_mock_transport().run_stream(BasicChatInputSchema(chat_message="hey"))) assert partials[-1].chat_message == "hi there" def test_run_stream_honors_a_caller_strict_setting(): partials = list(agent_over_mock_transport(strict=True).run_stream(BasicChatInputSchema(chat_message="hey"))) assert partials[-1].chat_message == "hi there" @pytest.mark.asyncio async def test_run_async_stream_yields_the_complete_response() -> None: agent = async_agent_over_mock_transport() partials = [response async for response in agent.run_async_stream(BasicChatInputSchema(chat_message="hey"))] assert partials[-1].chat_message == "hi there" @pytest.mark.asyncio async def test_run_async_stream_honors_a_caller_strict_setting() -> None: agent = async_agent_over_mock_transport(strict=True) partials = [response async for response in agent.run_async_stream(BasicChatInputSchema(chat_message="hey"))] assert partials[-1].chat_message == "hi there" ``` ### File: atomic-agents/tests/agents/test_minimax_integration.py ```python """Integration tests for MiniMax provider with Atomic Agents. These tests require a valid MINIMAX_API_KEY environment variable. Run with: pytest -m integration tests/agents/test_minimax_integration.py """ import os import pytest import instructor from pydantic import Field from atomic_agents import ( AtomicAgent, AgentConfig, BasicChatInputSchema, BasicChatOutputSchema, BaseIOSchema, ) pytestmark = pytest.mark.skipif( not os.getenv("MINIMAX_API_KEY"), reason="MINIMAX_API_KEY not set", ) def _make_minimax_agent(model="MiniMax-M3", **agent_kwargs): """Helper to create a MiniMax-backed agent.""" from openai import OpenAI raw = OpenAI( base_url="https://api.minimax.io/v1", api_key=os.environ["MINIMAX_API_KEY"], ) client = instructor.from_openai(raw, mode=instructor.Mode.JSON) config = AgentConfig(client=client, model=model, mode=instructor.Mode.JSON, **agent_kwargs) return AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) @pytest.mark.integration class TestMiniMaxLiveChat: """Live integration tests against MiniMax API.""" def test_basic_chat(self): """Test a basic chat completion with MiniMax.""" agent = _make_minimax_agent() response = agent.run(BasicChatInputSchema(chat_message="Say hello in one word.")) assert response.chat_message assert len(response.chat_message) > 0 def test_multi_turn_conversation(self): """Test multi-turn conversation with MiniMax.""" agent = _make_minimax_agent() r1 = agent.run(BasicChatInputSchema(chat_message="Remember the number 42.")) assert r1.chat_message # first response should be non-empty r2 = agent.run(BasicChatInputSchema(chat_message="What number did I ask you to remember?")) assert r2.chat_message # The model should recall "42" from the conversation history assert "42" in r2.chat_message def test_custom_output_schema(self): """Test structured output with a custom schema via MiniMax.""" from openai import OpenAI class AnalysisOutput(BaseIOSchema): """Analysis output schema.""" sentiment: str = Field(..., description="One of: positive, negative, neutral") confidence: float = Field(..., description="Confidence score between 0 and 1") raw = OpenAI( base_url="https://api.minimax.io/v1", api_key=os.environ["MINIMAX_API_KEY"], ) client = instructor.from_openai(raw, mode=instructor.Mode.JSON) config = AgentConfig(client=client, model="MiniMax-M3", mode=instructor.Mode.JSON) agent = AtomicAgent[BasicChatInputSchema, AnalysisOutput](config) response = agent.run(BasicChatInputSchema(chat_message="I love this product, it's amazing!")) assert response.sentiment in ("positive", "negative", "neutral") assert 0 <= response.confidence <= 1 ``` ### File: atomic-agents/tests/agents/test_minimax_provider.py ```python """Unit tests for MiniMax provider integration with Atomic Agents.""" import os import pytest from unittest.mock import Mock, patch import instructor from atomic_agents import ( AtomicAgent, AgentConfig, BasicChatInputSchema, BasicChatOutputSchema, ) from atomic_agents.context import SystemPromptGenerator def _create_minimax_client(api_key="test-key"): """Create a MiniMax client via OpenAI-compatible interface.""" from openai import OpenAI return instructor.from_openai(OpenAI(base_url="https://api.minimax.io/v1", api_key=api_key)) class TestMiniMaxClientSetup: """Tests for MiniMax client initialization.""" def test_minimax_client_creation(self): """Test that MiniMax client can be created with correct base_url.""" from openai import OpenAI raw_client = OpenAI(base_url="https://api.minimax.io/v1", api_key="test-key") assert raw_client.base_url == "https://api.minimax.io/v1/" def test_minimax_instructor_wrapping(self): """Test that MiniMax client can be wrapped with instructor.""" client = _create_minimax_client() assert isinstance(client, instructor.Instructor) def test_minimax_agent_config(self): """Test that AgentConfig accepts MiniMax client and model.""" client = _create_minimax_client() config = AgentConfig( client=client, model="MiniMax-M3", ) assert config.model == "MiniMax-M3" assert config.assistant_role == "assistant" def test_minimax_agent_initialization(self): """Test that AtomicAgent can be initialized with MiniMax config.""" client = _create_minimax_client() config = AgentConfig( client=client, model="MiniMax-M3", ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) assert agent.model == "MiniMax-M3" assert agent.assistant_role == "assistant" def test_minimax_m27_legacy_model(self): """Test that the legacy M2.7 model variant still works.""" client = _create_minimax_client() config = AgentConfig( client=client, model="MiniMax-M2.7", ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) assert agent.model == "MiniMax-M2.7" class TestMiniMaxAgentBehavior: """Tests for agent behavior with MiniMax provider.""" @pytest.fixture def mock_minimax_instructor(self): mock = Mock(spec=instructor.Instructor) mock.chat = Mock() mock.chat.completions = Mock() mock.chat.completions.create = Mock(return_value=BasicChatOutputSchema(chat_message="MiniMax response")) mock_response = BasicChatOutputSchema(chat_message="MiniMax response") mock_iter = Mock() mock_iter.__iter__ = Mock(return_value=iter([mock_response])) mock.chat.completions.create_partial.return_value = mock_iter return mock @pytest.fixture def minimax_agent(self, mock_minimax_instructor): config = AgentConfig( client=mock_minimax_instructor, model="MiniMax-M3", ) return AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) def test_run_with_minimax(self, minimax_agent, mock_minimax_instructor): """Test that agent.run works with MiniMax mock client.""" user_input = BasicChatInputSchema(chat_message="Hello from MiniMax test") response = minimax_agent.run(user_input) assert response.chat_message == "MiniMax response" mock_minimax_instructor.chat.completions.create.assert_called_once() def test_run_passes_correct_model(self, minimax_agent, mock_minimax_instructor): """Test that the correct model name is passed to the API.""" user_input = BasicChatInputSchema(chat_message="Test") minimax_agent.run(user_input) call_kwargs = mock_minimax_instructor.chat.completions.create.call_args assert call_kwargs.kwargs["model"] == "MiniMax-M3" def test_run_stream_with_minimax(self, minimax_agent): """Test that streaming works with MiniMax mock client.""" user_input = BasicChatInputSchema(chat_message="Stream test") responses = list(minimax_agent.run_stream(user_input)) assert len(responses) == 1 assert responses[0].chat_message == "MiniMax response" def test_history_tracking_with_minimax(self, minimax_agent): """Test that chat history is properly tracked.""" user_input = BasicChatInputSchema(chat_message="First message") minimax_agent.run(user_input) history = minimax_agent.history.get_history() assert len(history) == 2 # user + assistant def test_system_prompt_with_minimax(self, mock_minimax_instructor): """Test that system prompt works correctly with MiniMax.""" spg = SystemPromptGenerator( background=["You are a helpful MiniMax-powered assistant."], ) config = AgentConfig( client=mock_minimax_instructor, model="MiniMax-M3", system_prompt_generator=spg, ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) prompt = agent.system_prompt_generator.generate_prompt() assert "MiniMax" in prompt def test_model_api_parameters_with_minimax(self, mock_minimax_instructor): """Test that custom API parameters are passed through.""" config = AgentConfig( client=mock_minimax_instructor, model="MiniMax-M3", model_api_parameters={"temperature": 0.7, "max_tokens": 1024}, ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) assert agent.model_api_parameters["temperature"] == 0.7 assert agent.model_api_parameters["max_tokens"] == 1024 def test_minimax_reset_history(self, minimax_agent): """Test that history reset works with MiniMax agent.""" user_input = BasicChatInputSchema(chat_message="Test") minimax_agent.run(user_input) minimax_agent.reset_history() history = minimax_agent.history.get_history() assert len(history) == 0 class TestMiniMaxProviderSetup: """Tests for the provider setup function from the quickstart example.""" def test_setup_client_minimax_by_number(self): """Test setup_client with provider number '7'.""" import sys sys.path.insert( 0, os.path.join( os.path.dirname(__file__), "..", "..", "..", "atomic-examples", "quickstart", "quickstart", ), ) # We can't import the example directly (it has top-level console input), # but we can verify the pattern works from openai import OpenAI api_key = "test-minimax-key" raw_client = OpenAI(base_url="https://api.minimax.io/v1", api_key=api_key) client = instructor.from_openai(raw_client) assert isinstance(client, instructor.Instructor) def test_minimax_env_var_detection(self): """Test that MINIMAX_API_KEY env var can be used.""" with patch.dict(os.environ, {"MINIMAX_API_KEY": "test-env-key"}): api_key = os.getenv("MINIMAX_API_KEY") assert api_key == "test-env-key" ``` ### File: atomic-agents/tests/agents/test_minimax_transport.py ```python """Exercise MiniMax defaults and explicit models through the real SDK transport.""" import json from pathlib import Path import runpy import httpx import instructor import openai import pytest from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema, BasicChatInputSchema, BasicChatOutputSchema, VideoURL @pytest.fixture def transport_requests(): requests = [] def respond(request): payload = json.loads(request.content) requests.append((request, payload)) return httpx.Response( 200, json={ "id": "offline-completion", "object": "chat.completion", "created": 0, "model": payload["model"], "choices": [ { "index": 0, "message": {"role": "assistant", "content": '{"chat_message":"Offline response"}'}, "finish_reason": "stop", } ], }, ) with httpx.Client(transport=httpx.MockTransport(respond)) as client: yield client, requests @pytest.mark.parametrize("provider", ["7", "minimax"]) def test_quickstart_default_reaches_http_transport(monkeypatch, transport_requests, provider): """Run the actual interactive example, including two complete agent turns.""" import dotenv from rich.console import Console http_client, requests = transport_requests sdk_init = openai.OpenAI.__init__ def offline_init(self, **kwargs): sdk_init(self, **kwargs, http_client=http_client) monkeypatch.setattr(openai.OpenAI, "__init__", offline_init) monkeypatch.setattr(dotenv, "load_dotenv", lambda: None) monkeypatch.setenv("MINIMAX_API_KEY", "offline-test-key") inputs = iter([provider, "Remember 42", "What number?", "/exit"]) monkeypatch.setattr(Console, "input", lambda *args, **kwargs: next(inputs)) example = ( Path(__file__).resolve().parents[3] / "atomic-examples/quickstart/quickstart/4_basic_chatbot_different_providers.py" ) namespace = runpy.run_path(str(example)) assert len(requests) == 2 for request, payload in requests: assert str(request.url) == "https://api.minimax.io/v1/chat/completions" assert payload["model"] == "MiniMax-M3.1-Flash" assert payload["max_tokens"] == 2048 assert payload["response_format"] == {"type": "json_object"} assert "service_tier" not in payload assert "thinking" not in payload assert "reasoning_effort" not in payload assert "reasoning_split" not in payload second_messages = requests[1][1]["messages"] assert any("Remember 42" in str(message["content"]) for message in second_messages) assert any("Offline response" in str(message["content"]) for message in second_messages) assert namespace["agent"].model == "MiniMax-M3.1-Flash" # The existing non-MiniMax provider choice retains its model and parameters. monkeypatch.setenv("OPENAI_API_KEY", "offline-test-key") _, model, parameters, role = namespace["setup_client"]("openai") assert model == "gpt-5.6-terra" assert parameters == {"reasoning_effort": "low", "max_tokens": 2048} assert role == "assistant" @pytest.mark.parametrize("base_url", ["https://api.minimax.io/v1", "https://api.minimaxi.com/v1"]) @pytest.mark.parametrize("model", ["MiniMax-M3.1-Flash", "MiniMax-M3", "MiniMax-M2.7", "custom-saved-model"]) def test_explicit_model_and_parameters_reach_http_transport(transport_requests, base_url, model): """Configured IDs and parameters remain caller-owned on either configured endpoint.""" http_client, requests = transport_requests raw = openai.OpenAI(base_url=base_url, api_key="offline-test-key", http_client=http_client) client = instructor.from_openai(raw, mode=instructor.Mode.JSON) parameters = {"max_tokens": 2048} if model == "MiniMax-M3.1-Flash": parameters["reasoning_effort"] = "high" elif model == "MiniMax-M3": parameters["extra_body"] = {"thinking": {"type": "disabled"}} agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema]( AgentConfig(client=client, model=model, mode=instructor.Mode.JSON, model_api_parameters=parameters) ) assert agent.run(BasicChatInputSchema(chat_message="Hello")).chat_message == "Offline response" request, payload = requests[0] assert str(request.url) == base_url + "/chat/completions" assert payload["model"] == model assert payload["max_tokens"] == 2048 assert "service_tier" not in payload if model == "MiniMax-M3.1-Flash": assert payload["reasoning_effort"] == "high" elif model == "MiniMax-M3": assert payload["thinking"] == {"type": "disabled"} def test_video_history_preserved_at_http_boundary(transport_requests): """The default update does not change existing typed video serialization.""" class VideoInput(BaseIOSchema): """A question about a video.""" question: str video: VideoURL http_client, requests = transport_requests raw = openai.OpenAI(base_url="https://api.minimax.io/v1", api_key="offline-test-key", http_client=http_client) client = instructor.from_openai(raw, mode=instructor.Mode.JSON) agent = AtomicAgent[VideoInput, BasicChatOutputSchema]( AgentConfig(client=client, model="MiniMax-M3.1-Flash", mode=instructor.Mode.JSON) ) video = VideoURL(url="https://example.test/video.mp4", fps=1) agent.run(VideoInput(question="Summarize this video", video=video)) user_message = next(message for message in requests[0][1]["messages"] if message["role"] == "user") assert video.to_openai() in user_message["content"] ``` ### File: atomic-agents/tests/base/test_base_prompt.py ```python from pydantic import ConfigDict from atomic_agents.base import BaseIOSchema, BasePrompt, BasePromptConfig class GreetingInput(BaseIOSchema): """Name of the person who will receive the greeting.""" model_config = ConfigDict(title="Greeting request") name: str class GreetingOutput(BaseIOSchema): """Greeting addressed to the requested person.""" greeting: str class GreetingPrompt(BasePrompt[GreetingInput, GreetingOutput]): def generate(self, params: GreetingInput) -> GreetingOutput: return GreetingOutput(greeting=f"Hello, {params.name}!") class InheritedGreetingPrompt(GreetingPrompt): pass class EmptyPrompt[InputSchema: BaseIOSchema, OutputSchema: BaseIOSchema](BasePrompt): def generate(self, params: InputSchema) -> OutputSchema: return self.output_schema() def test_base_prompt_inherited_schema_identity(): prompt = GreetingPrompt() assert prompt.input_schema is GreetingInput assert prompt.output_schema is GreetingOutput def test_base_prompt_descendant_preserves_schema_identity(): prompt = InheritedGreetingPrompt() assert prompt.input_schema is GreetingInput assert prompt.output_schema is GreetingOutput def test_base_prompt_explicit_type_parameters_resolve_schema_identity(): prompt = EmptyPrompt[GreetingInput, GreetingOutput]() assert prompt.input_schema is GreetingInput assert prompt.output_schema is GreetingOutput def test_base_prompt_without_type_parameters_defaults_to_base_schema(): prompt = EmptyPrompt() assert prompt.input_schema is BaseIOSchema assert prompt.output_schema is BaseIOSchema def test_base_prompt_without_type_parameters_uses_base_schema_metadata(): prompt = EmptyPrompt() assert prompt.prompt_name == "BaseIOSchema" assert prompt.prompt_description == "Base schema for input/output in the Atomic Agents framework." def test_base_prompt_metadata_comes_from_input_schema(): prompt = GreetingPrompt() assert prompt.prompt_name == "Greeting request" assert prompt.prompt_description == "Name of the person who will receive the greeting." def test_base_prompt_explicit_type_parameters_use_input_schema_metadata(): prompt = EmptyPrompt[GreetingInput, GreetingOutput]() assert prompt.prompt_name == "Greeting request" assert prompt.prompt_description == "Name of the person who will receive the greeting." def test_base_prompt_config_overrides_both_metadata_fields(): config = BasePromptConfig(title="Welcome", description="Welcome a new member.") prompt = GreetingPrompt(config=config) assert prompt.config is config assert prompt.prompt_name == "Welcome" assert prompt.prompt_description == "Welcome a new member." def test_base_prompt_title_override_preserves_schema_description(): prompt = GreetingPrompt(config=BasePromptConfig(title="Welcome")) assert prompt.prompt_name == "Welcome" assert prompt.prompt_description == "Name of the person who will receive the greeting." def test_base_prompt_description_override_preserves_schema_title(): prompt = GreetingPrompt(config=BasePromptConfig(description="Welcome a new member.")) assert prompt.prompt_name == "Greeting request" assert prompt.prompt_description == "Welcome a new member." def test_base_prompt_empty_metadata_overrides_use_schema_defaults(): prompt = GreetingPrompt(config=BasePromptConfig(title="", description="")) assert prompt.prompt_name == "Greeting request" assert prompt.prompt_description == "Name of the person who will receive the greeting." def test_greeting_prompt_generate_returns_greeting_schema_and_value(): prompt = GreetingPrompt() result = prompt.generate(GreetingInput(name="Ada")) assert type(result) is GreetingOutput assert result.greeting == "Hello, Ada!" def test_untyped_prompt_generate_returns_empty_base_schema(): prompt = EmptyPrompt() result = prompt.generate(BaseIOSchema()) assert type(result) is BaseIOSchema assert result.model_dump() == {} ``` ### File: atomic-agents/tests/base/test_base_resource.py ```python import pytest from pydantic import ConfigDict from atomic_agents.base import BaseIOSchema, BaseResource, BaseResourceConfig class ReadInputSchema(BaseIOSchema): """Value requested from a resource.""" model_config = ConfigDict(title="Read value") value: str class ReadOutputSchema(BaseIOSchema): """Value returned by a resource.""" value: str class InheritedEchoResource(BaseResource[ReadInputSchema, ReadOutputSchema]): def read(self, input_data: ReadInputSchema) -> ReadOutputSchema: return ReadOutputSchema(value=input_data.value) class ParameterizedEchoResource[InputSchema: BaseIOSchema, OutputSchema: BaseIOSchema](BaseResource): def read(self, input_data: InputSchema) -> OutputSchema: return self.output_schema.model_validate(input_data.model_dump()) def test_inherited_resource_schema_identity() -> None: resource = InheritedEchoResource() assert resource.input_schema is ReadInputSchema assert resource.output_schema is ReadOutputSchema def test_parameterized_resource_schema_identity() -> None: resource = ParameterizedEchoResource[ReadInputSchema, ReadOutputSchema]() assert resource.input_schema is ReadInputSchema assert resource.output_schema is ReadOutputSchema def test_untyped_resource_schema_defaults() -> None: resource = ParameterizedEchoResource() assert resource.input_schema is BaseIOSchema assert resource.output_schema is BaseIOSchema def test_untyped_resource_metadata_defaults() -> None: resource = ParameterizedEchoResource() assert resource.resource_name == "BaseIOSchema" assert resource.resource_description == "Base schema for input/output in the Atomic Agents framework." def test_resource_metadata_comes_from_input_schema() -> None: resource = InheritedEchoResource() assert resource.resource_name == "Read value" assert resource.resource_description == "Value requested from a resource." def test_resource_title_override_preserves_schema_description() -> None: resource = InheritedEchoResource(config=BaseResourceConfig(title="Echo resource")) assert resource.resource_name == "Echo resource" assert resource.resource_description == "Value requested from a resource." def test_resource_description_override_preserves_schema_title() -> None: resource = InheritedEchoResource(config=BaseResourceConfig(description="Configured description")) assert resource.resource_name == "Read value" assert resource.resource_description == "Configured description" def test_resource_metadata_can_override_both_schema_values() -> None: resource = InheritedEchoResource(config=BaseResourceConfig(title="Echo resource", description="Configured description")) assert resource.resource_name == "Echo resource" assert resource.resource_description == "Configured description" def test_empty_resource_metadata_falls_back_to_input_schema() -> None: resource = InheritedEchoResource(config=BaseResourceConfig(title="", description="")) assert resource.resource_name == "Read value" assert resource.resource_description == "Value requested from a resource." @pytest.mark.parametrize("resource", [InheritedEchoResource(), ParameterizedEchoResource[ReadInputSchema, ReadOutputSchema]()]) def test_concrete_resource_read_returns_typed_value(resource: BaseResource[ReadInputSchema, ReadOutputSchema]) -> None: result = resource.read(ReadInputSchema(value="requested value")) assert type(result) is ReadOutputSchema assert result.value == "requested value" def test_untyped_resource_read_returns_base_schema() -> None: result = ParameterizedEchoResource().read(BaseIOSchema()) assert type(result) is BaseIOSchema assert result.model_dump() == {} ``` ### File: atomic-agents/tests/base/test_base_tool.py ```python from pydantic import BaseModel from atomic_agents import BaseToolConfig, BaseTool, BaseIOSchema # Mock classes for testing class MockInputSchema(BaseIOSchema): """Mock input schema for testing""" query: str class MockOutputSchema(BaseIOSchema): """Mock output schema for testing""" result: str class MockTool[InputSchema: BaseIOSchema, OutputSchema: BaseIOSchema](BaseTool): def run(self, params: InputSchema) -> OutputSchema: if self.output_schema == MockOutputSchema: return MockOutputSchema(result="Mock result") elif self.output_schema == BaseIOSchema: return BaseIOSchema() else: raise ValueError("Unsupported output schema") def test_base_tool_config_creation(): config = BaseToolConfig() assert config.title is None assert config.description is None def test_base_tool_config_with_values(): config = BaseToolConfig(title="Test Tool", description="Test description") assert config.title == "Test Tool" assert config.description == "Test description" def test_base_tool_initialization_without_type_parameters(): tool = MockTool() assert tool.tool_name == "BaseIOSchema" assert tool.tool_description == "Base schema for input/output in the Atomic Agents framework." assert tool.output_schema == BaseIOSchema def test_base_tool_initialization(): tool = MockTool[MockInputSchema, MockOutputSchema]() assert tool.tool_name == "MockInputSchema" assert tool.tool_description == "Mock input schema for testing" def test_base_tool_with_config(): config = BaseToolConfig(title="Custom Title", description="Custom description") tool = MockTool[MockInputSchema, MockOutputSchema](config=config) assert tool.tool_name == "Custom Title" assert tool.tool_description == "Custom description" def test_base_tool_with_custom_title(): config = BaseToolConfig(title="Custom Tool Name") tool = MockTool[MockInputSchema, MockOutputSchema](config=config) assert tool.tool_name == "Custom Tool Name" assert tool.tool_description == "Mock input schema for testing" def test_mock_tool_run(): tool = MockTool[MockInputSchema, MockOutputSchema]() result = tool.run(MockInputSchema(query="mock query")) assert isinstance(result, MockOutputSchema) assert result.result == "Mock result" def test_base_tool_input_schema(): tool = MockTool[MockInputSchema, MockOutputSchema]() assert tool.input_schema == MockInputSchema def test_base_tool_output_schema(): tool = MockTool[MockInputSchema, MockOutputSchema]() assert tool.output_schema == MockOutputSchema def test_base_tool_inheritance(): tool = MockTool[MockInputSchema, MockOutputSchema]() assert isinstance(tool, BaseTool) def test_base_tool_config_is_pydantic_model(): assert issubclass(BaseToolConfig, BaseModel) def test_base_tool_config_optional_fields(): config = BaseToolConfig() assert hasattr(config, "title") assert hasattr(config, "description") # Test for GitHub issue #161 fix: proper schema resolution def test_base_tool_schema_resolution(): """Test that input_schema and output_schema return correct types (not BaseIOSchema)""" class CustomInput(BaseIOSchema): """Custom input schema for testing""" name: str class CustomOutput(BaseIOSchema): """Custom output schema for testing""" result: str class TestTool(BaseTool[CustomInput, CustomOutput]): def run(self, params: CustomInput) -> CustomOutput: return CustomOutput(result=f"processed_{params.name}") tool = TestTool() # These should return the specific types, not BaseIOSchema assert tool.input_schema == CustomInput assert tool.output_schema == CustomOutput assert tool.input_schema != BaseIOSchema assert tool.output_schema != BaseIOSchema ``` ### File: atomic-agents/tests/base/test_default_config_isolation.py ```python import pytest from atomic_agents.base import BaseIOSchema, BasePrompt, BaseResource, BaseTool class ValueSchema(BaseIOSchema): """A value passed through a component.""" value: str class EchoTool(BaseTool[ValueSchema, ValueSchema]): def run(self, params: ValueSchema) -> ValueSchema: return params class EchoResource(BaseResource[ValueSchema, ValueSchema]): def read(self, params: ValueSchema) -> ValueSchema: return params class EchoPrompt(BasePrompt[ValueSchema, ValueSchema]): def generate(self, params: ValueSchema) -> ValueSchema: return params @pytest.mark.parametrize( "component_class, name_attribute, description_attribute", [ (EchoTool, "tool_name", "tool_description"), (EchoResource, "resource_name", "resource_description"), (EchoPrompt, "prompt_name", "prompt_description"), ], ) def test_default_config_changes_do_not_affect_other_instances( component_class: type[EchoTool] | type[EchoResource] | type[EchoPrompt], name_attribute: str, description_attribute: str, ) -> None: first = component_class() second = component_class() original_title, original_description = first.config.title, first.config.description try: first.config.title = "Customized instance" first.config.description = "Only this instance should change." assert getattr(first, name_attribute) == "Customized instance" assert getattr(first, description_attribute) == "Only this instance should change." for other in (second, component_class()): assert getattr(other, name_attribute) == "ValueSchema" assert getattr(other, description_attribute) == ValueSchema.__doc__ finally: first.config.title, first.config.description = original_title, original_description @pytest.mark.parametrize("component_class", [EchoTool, EchoResource, EchoPrompt]) def test_explicit_config_is_retained(component_class: type[EchoTool] | type[EchoResource] | type[EchoPrompt]) -> None: config = component_class().config component = component_class(config=config) assert component.config is config ``` ### File: atomic-agents/tests/base/test_multimodal.py ```python from atomic_agents import VideoURL def test_to_openai_with_url_only(): content_part = VideoURL(url="https://example.com/clip.mp4").to_openai() assert content_part == {"type": "video_url", "video_url": {"url": "https://example.com/clip.mp4"}} def test_to_openai_with_sampling_parameters(): content_part = VideoURL(url="https://example.com/clip.mp4", fps=1.0, detail="low").to_openai() assert content_part == { "type": "video_url", "video_url": {"url": "https://example.com/clip.mp4", "fps": 1.0, "detail": "low"}, } ``` ### File: atomic-agents/tests/connectors/mcp/test_boolean_schemas.py ```python import pytest from jsonschema import Draft202012Validator from pydantic import ValidationError from atomic_agents.connectors.mcp import SchemaTransformer def make_model(value_schema, required=True, definitions=None): schema = { "type": "object", "properties": {"value": value_schema}, "required": ["value"] if required else [], "$defs": definitions or {}, } return SchemaTransformer.create_model_from_schema(schema, "BooleanInput", "boolean"), schema @pytest.mark.parametrize("required", [True, False]) @pytest.mark.parametrize("value", [None, False, 1, "text", [1, None], {"x": True}]) def test_true_property_accepts_every_json_value(required, value): model, original = make_model(True, required) payload = {"tool_name": "boolean", "value": value} assert Draft202012Validator(original).is_valid(payload) assert Draft202012Validator(model.model_json_schema()).is_valid(payload) assert model.model_validate(payload).value == value @pytest.mark.parametrize("required", [True, False]) @pytest.mark.parametrize("value", [None, False, 1, "text", [], {}]) def test_false_property_rejects_every_present_json_value(required, value): model, original = make_model(False, required) payload = {"tool_name": "boolean", "value": value} assert not Draft202012Validator(original).is_valid(payload) assert not Draft202012Validator(model.model_json_schema()).is_valid(payload) with pytest.raises(ValidationError): model.model_validate(payload) @pytest.mark.parametrize("value_schema", [True, False]) def test_optional_boolean_property_can_be_omitted(value_schema): model, _ = make_model(value_schema, required=False) result = model(tool_name="boolean") assert result.model_dump(exclude_unset=True) == {"tool_name": "boolean"} assert Draft202012Validator(model.model_json_schema()).is_valid({"tool_name": "boolean"}) @pytest.mark.parametrize("value_schema", [True, False]) def test_required_boolean_property_cannot_be_omitted(value_schema): model, _ = make_model(value_schema) with pytest.raises(ValidationError): model(tool_name="boolean") @pytest.mark.parametrize("keyword", ["anyOf", "oneOf"]) def test_false_union_member_keeps_the_other_branch(keyword): model, _ = make_model({keyword: [False, {"type": "string"}]}) assert model(tool_name="boolean", value="text").value == "text" with pytest.raises(ValidationError): model(tool_name="boolean", value={}) assert not Draft202012Validator(model.model_json_schema()).is_valid({"tool_name": "boolean", "value": {}}) def test_true_union_member_accepts_other_values(): model, _ = make_model({"anyOf": [True, {"type": "string"}]}) assert model(tool_name="boolean", value={"x": 1}).value == {"x": 1} @pytest.mark.parametrize("items, valid, invalid", [(True, [1, None, "x"], None), (False, [], [None])]) def test_boolean_array_items(items, valid, invalid): model, _ = make_model({"type": "array", "items": items}) assert model(tool_name="boolean", value=valid).value == valid if invalid is not None: with pytest.raises(ValidationError): model(tool_name="boolean", value=invalid) assert not Draft202012Validator(model.model_json_schema()).is_valid({"tool_name": "boolean", "value": invalid}) @pytest.mark.parametrize("definition", [True, False]) @pytest.mark.parametrize("required", [True, False]) def test_references_to_boolean_definitions(definition, required): model, _ = make_model({"$ref": "#/$defs/Value"}, required=required, definitions={"Value": definition}) if not required: assert model(tool_name="boolean").model_dump(exclude_unset=True) == {"tool_name": "boolean"} if definition: assert model(tool_name="boolean", value={"x": 1}).value == {"x": 1} else: with pytest.raises(ValidationError): model(tool_name="boolean", value={"x": 1}) with pytest.raises(ValidationError): model(tool_name="boolean", value=None) ``` ### File: atomic-agents/tests/connectors/mcp/test_generated_class_binding.py ```python """Generated MCP classes must retain their own output type and metadata.""" import asyncio from unittest.mock import AsyncMock, MagicMock import pytest from mcp import ClientSession from mcp.types import GetPromptResult, PromptMessage, ReadResourceResult, TextContent, TextResourceContents from atomic_agents.connectors.mcp import MCPFactory, MCPPromptDefinition, MCPResourceDefinition, MCPToolDefinition @pytest.fixture(params=[False, True], ids=["owned", "borrowed"]) def connection(request, monkeypatch): session = AsyncMock(spec=ClientSession) session.__aenter__.return_value = session transport = AsyncMock() transport.__aenter__.return_value = (None, None, None) monkeypatch.setattr("atomic_agents.connectors.mcp.mcp_factory.ClientSession", MagicMock(return_value=session)) monkeypatch.setattr("atomic_agents.connectors.mcp.mcp_factory.streamablehttp_client", MagicMock(return_value=transport)) with asyncio.Runner() as runner: factory = MCPFactory( "http://localhost", client_session=session if request.param else None, event_loop=runner.get_loop() ) yield factory, session, runner def resource_classes(factory): schema = {"type": "object", "properties": {}} return factory._create_resource_classes( [ MCPResourceDefinition("first", None, "file:///first", schema, "text/plain"), MCPResourceDefinition("last", None, "file:///last", schema, "application/json"), ] ) def prompt_classes(factory): schema = {"type": "object", "properties": {}} return factory._create_prompt_classes( [MCPPromptDefinition("first", None, schema), MCPPromptDefinition("last", None, schema)] ) @pytest.mark.parametrize("synchronous", [False, True], ids=["async", "sync"]) def test_resource_output_class_matches_each_resource(connection, synchronous): factory, session, runner = connection classes = resource_classes(factory) for resource_class in classes: resource = resource_class() session.read_resource.return_value = ReadResourceResult( contents=[TextResourceContents(uri=resource.uri, text=resource.mcp_resource_name, mimeType="text/custom")] ) params = resource.input_schema(resource_name=resource.mcp_resource_name) result = resource.read(params) if synchronous else runner.run(resource.aread(params)) assert type(result) is resource.output_schema assert result.content[0].text == resource.mcp_resource_name assert result.mime_type == "text/custom" session.read_resource.assert_awaited_with(uri=resource.uri) @pytest.mark.parametrize("payload", [ReadResourceResult(contents=[]), {"contents": []}, "raw content"]) def test_resource_mime_fallback_matches_each_resource(connection, payload): factory, session, runner = connection session.read_resource.return_value = payload for resource_class, mime_type in zip(resource_classes(factory), ["text/plain", "application/json"]): resource = resource_class() result = runner.run(resource.aread(resource.input_schema(resource_name=resource.mcp_resource_name))) assert result.mime_type == mime_type @pytest.mark.parametrize("synchronous", [False, True], ids=["async", "sync"]) def test_prompt_output_class_matches_each_prompt(connection, synchronous): factory, session, runner = connection classes = prompt_classes(factory) for prompt_class in classes: prompt = prompt_class() session.get_prompt.return_value = GetPromptResult( messages=[PromptMessage(role="user", content=TextContent(type="text", text=prompt.mcp_prompt_name))] ) params = prompt.input_schema(prompt_name=prompt.mcp_prompt_name) result = prompt.generate(params) if synchronous else runner.run(prompt.agenerate(params)) assert type(result) is prompt.output_schema assert result.content == prompt.mcp_prompt_name session.get_prompt.assert_awaited_with(name=prompt.mcp_prompt_name, arguments={}) @pytest.mark.parametrize( "names,index", [(("first", "last"), 0), (("first", "last"), 1), (("only",), 0)], ids=["first", "last", "single"] ) @pytest.mark.parametrize( "kind,method,action", [("tool", "run", "execute"), ("resource", "read", "read"), ("prompt", "generate", "get")] ) def test_sync_wrapper_error_names_each_generated_class(kind, method, action, names, index): """A missing borrowed-session loop must not report the last definition's name.""" schema = {"type": "object", "properties": {}} definitions = { "tool": [MCPToolDefinition(name, None, schema) for name in names], "resource": [MCPResourceDefinition(name, None, f"file:///{name}", schema) for name in names], "prompt": [MCPPromptDefinition(name, None, schema) for name in names], } session = AsyncMock(spec=ClientSession) with asyncio.Runner() as runner: factory = MCPFactory(client_session=session, event_loop=runner.get_loop()) classes = getattr(factory, f"_create_{kind}_classes")(definitions[kind]) assert len(classes) == len(names) component_class, name = classes[index], names[index] component = component_class() component._event_loop = None params = component.input_schema(**{f"{kind}_name": name}) with pytest.raises(RuntimeError) as caught: getattr(component, method)(params) assert str(caught.value) == ( f"Failed to {action} MCP {kind} '{name}': " "'NoneType' object has no attribute 'run_until_complete'" ) assert isinstance(caught.value.__cause__, AttributeError) session.call_tool.assert_not_awaited() session.read_resource.assert_not_awaited() session.get_prompt.assert_not_awaited() ``` ### File: atomic-agents/tests/connectors/mcp/test_generated_class_session.py ```python """Verify multiple generated MCP classes across real JSON-round-tripped sessions. In-memory streams exercise discovery and the client/server protocol, not network or stdio transports. Each session is borrowed by every generated component. """ import asyncio from collections.abc import AsyncIterator from contextlib import AsyncExitStack, asynccontextmanager from datetime import timedelta import json import anyio from mcp import ClientSession from mcp import types as mt from mcp.server import Server from mcp.server.lowlevel.helper_types import ReadResourceContents from mcp.shared.message import SessionMessage import pytest from atomic_agents.connectors.mcp.mcp_factory import fetch_mcp_resources_async, fetch_mcp_prompts_async @asynccontextmanager async def connected(server: Server, frames: list[dict]) -> AsyncIterator[ClientSession]: """Relay serialized messages while owning both sessions and their streams.""" async with AsyncExitStack() as stack: channels = [anyio.create_memory_object_stream(1) for _ in range(4)] for pair in channels: for stream in pair: await stack.enter_async_context(stream) client_send, client_out = channels[0] server_send, server_out = channels[1] to_server, server_read = channels[2] to_client, client_read = channels[3] async def relay(source, target) -> None: async for message in source: encoded = message.message.model_dump_json(by_alias=True, exclude_none=True) frames.append(json.loads(encoded)) await target.send(SessionMessage(mt.JSONRPCMessage.model_validate_json(encoded))) async with anyio.create_task_group() as group: group.start_soon(relay, client_out, to_server) group.start_soon(relay, server_out, to_client) group.start_soon(server.run, server_read, server_send, server.create_initialization_options()) try: async with ClientSession(client_read, client_send, read_timeout_seconds=timedelta(seconds=3)) as session: await session.initialize() yield session finally: group.cancel_scope.cancel() @pytest.mark.parametrize("empty", [True, False], ids=["empty-fallback", "explicit-content-mime"]) def test_discovered_resources_keep_output_binding_and_mime(empty: bool) -> None: async def scenario() -> None: server = Server("resource-binding-test") frames = [] names = ["first", "last"] uris = {name: f"memory://binding/{name}" for name in names} fallback = dict(zip(names, ["text/plain", "application/json"])) calls = [] @server.list_resources() async def resources() -> list[mt.Resource]: return [mt.Resource(name=name, uri=uris[name], mimeType=fallback[name]) for name in names] @server.list_resource_templates() async def templates() -> list[mt.ResourceTemplate]: return [] @server.read_resource() async def read_resource(uri) -> list[ReadResourceContents]: calls.append(str(uri)) name = str(uri).rsplit("/", 1)[1] return [] if empty else [ReadResourceContents(content=f"{name}: 世界", mime_type=f"application/x-{name}")] async with connected(server, frames) as session: classes = await fetch_mcp_resources_async(client_session=session) assert [cls.mcp_resource_name for cls in classes] == names instances = {cls.mcp_resource_name: cls() for cls in classes} assert instances["first"].output_schema is not instances["last"].output_schema order = ["last", "first", "last", "first"] for name in order: resource = instances[name] result = await resource.aread(resource.input_schema(resource_name=name)) assert result.mime_type == (fallback[name] if empty else f"application/x-{name}") assert type(result) is resource.output_schema if empty: assert result.content == [] else: assert len(result.content) == 1 assert str(result.content[0].uri) == uris[name] assert result.content[0].text == f"{name}: 世界" await session.send_ping() assert calls == [uris[name] for name in order] requests = [frame["params"]["uri"] for frame in frames if frame.get("method") == "resources/read"] assert requests == calls assert any(frame.get("method") == "resources/list" for frame in frames) assert sum(frame.get("method") == "ping" for frame in frames) == len(order) asyncio.run(scenario()) def test_discovered_prompts_keep_output_binding_and_arguments() -> None: async def scenario() -> None: server = Server("prompt-binding-test") frames = [] names = ["first", "last"] calls = [] @server.list_prompts() async def prompts() -> list[mt.Prompt]: return [mt.Prompt(name=name, arguments=[mt.PromptArgument(name="text", required=True)]) for name in names] @server.get_prompt() async def get_prompt(name: str, arguments: dict[str, str] | None) -> mt.GetPromptResult: assert arguments is not None calls.append((name, arguments)) return mt.GetPromptResult( messages=[ mt.PromptMessage(role="user", content=mt.TextContent(type="text", text=f"{name}: {arguments['text']}")), mt.PromptMessage(role="assistant", content=mt.TextContent(type="text", text=f"second: {name}")), ] ) async with connected(server, frames) as session: classes = await fetch_mcp_prompts_async(client_session=session) assert [cls.mcp_prompt_name for cls in classes] == names instances = {cls.mcp_prompt_name: cls() for cls in classes} assert instances["first"].output_schema is not instances["last"].output_schema order = ["last", "first", "last", "first"] for index, name in enumerate(order): prompt = instances[name] text = f"世界 {index}" result = await prompt.agenerate(prompt.input_schema(prompt_name=name, text=text)) assert type(result) is prompt.output_schema assert result.content == f"{name}: {text}\n\nsecond: {name}" await session.send_ping() expected = [(name, {"text": f"世界 {index}"}) for index, name in enumerate(order)] assert calls == expected requests = [frame["params"] for frame in frames if frame.get("method") == "prompts/get"] assert requests == [{"name": name, "arguments": arguments} for name, arguments in expected] assert any(frame.get("method") == "prompts/list" for frame in frames) assert sum(frame.get("method") == "ping" for frame in frames) == len(order) asyncio.run(scenario()) ``` ### File: atomic-agents/tests/connectors/mcp/test_local_reference_paths.py ```python import pytest from pydantic import ValidationError from atomic_agents.connectors.mcp import SchemaTransformer def _model(schema): return SchemaTransformer.create_model_from_schema(schema, "PointerOutput", "pointer", is_output_schema=True) def test_definition_containers_do_not_share_cache_entries(): schema = { "type": "object", "properties": {"count": {"$ref": "#/$defs/Value"}, "name": {"$ref": "#/definitions/Value"}}, "required": ["count", "name"], "$defs": {"Value": {"type": "integer"}}, "definitions": {"Value": {"type": "string"}}, } model = _model(schema) assert model(count=7, name="seven").model_dump() == {"count": 7, "name": "seven"} with pytest.raises(ValidationError): model(count="invalid", name="seven") with pytest.raises(ValidationError): model(count=7, name={}) @pytest.mark.parametrize( "path, definition_key", [ ("#/$defs/Container/properties/count", "Container"), ("#/$defs/ns~1Container/properties/count", "ns/Container"), ("#/$defs/ns~01Container/properties/count", "ns~1Container"), ("#/$defs/Container%20Name/properties/count", "Container Name"), ], ) def test_full_pointer_resolves_nested_schema_instead_of_same_named_definition(path, definition_key): schema = { "type": "object", "properties": {"value": {"$ref": path}, "values": {"type": "array", "items": {"$ref": path}}}, "required": ["value", "values"], "$defs": { definition_key: {"type": "object", "properties": {"count": {"type": "integer"}}}, "count": {"type": "boolean"}, }, } model = _model(schema) assert model(value=7, values=[8]).model_dump() == {"value": 7, "values": [8]} assert model.model_json_schema()["properties"]["value"]["type"] == "integer" with pytest.raises(ValidationError): model(value="invalid", values=[]) with pytest.raises(ValidationError): model(value=7, values=[{}]) def test_pointer_can_select_schema_from_an_array(): schema = { "type": "object", "properties": {"value": {"$ref": "#/$defs/Choices/anyOf/1"}}, "required": ["value"], "$defs": {"Choices": {"anyOf": [{"type": "string"}, {"type": "integer"}]}}, } model = _model(schema) assert model(value=7).value == 7 with pytest.raises(ValidationError): model(value={}) @pytest.mark.parametrize( "reference", ["https://example.invalid/schema#/$defs/Value", "#/$defs/Missing/Value", "#/definitions/Value"] ) def test_unresolved_reference_does_not_bind_an_unrelated_local_definition(reference, caplog): schema = { "type": "object", "properties": {"value": {"$ref": reference}}, "required": ["value"], "$defs": {"Value": {"type": "integer"}}, } model = _model(schema) assert model(value={"unresolved": True}).value == {"unresolved": True} assert "Could not resolve $ref" in caplog.text def test_recursive_definition_retains_existing_nonrecursive_fields(): schema = { "type": "object", "properties": {"node": {"$ref": "#/$defs/Node"}}, "required": ["node"], "$defs": { "Node": { "type": "object", "properties": {"value": {"type": "integer"}, "next": {"$ref": "#/$defs/Node"}}, "required": ["value"], } }, } model = _model(schema) assert model(node={"value": 7}).node.value == 7 with pytest.raises(ValidationError): model(node={"value": "invalid"}) ``` ### File: atomic-agents/tests/connectors/mcp/test_mapping_schemas.py ```python import asyncio from unittest.mock import AsyncMock import pytest from mcp import ClientSession from mcp.types import CallToolResult from pydantic import ValidationError from atomic_agents.connectors.mcp import MCPFactory, MCPToolDefinition, SchemaTransformer @pytest.mark.parametrize("definitions_key", ["$defs", "definitions"]) @pytest.mark.parametrize("inline", [False, True], ids=["referenced", "inline"]) @pytest.mark.parametrize("empty_properties", [False, True]) def test_typed_mapping_preserves_values_and_validates_items(definitions_key, inline, empty_properties): mapping = {"type": "object", "additionalProperties": {"type": "integer"}} if empty_properties: mapping["properties"] = {} schema = { "type": "object", "properties": {"counts": mapping if inline else {"$ref": f"#/{definitions_key}/Counts"}}, "required": ["counts"], definitions_key: {"Counts": mapping}, } model = SchemaTransformer.create_model_from_schema(schema, "CountsOutput", "count", is_output_schema=True) counts = {"apples": 3, "oranges": 7} assert model(counts=counts).model_dump() == {"counts": counts} assert model(counts={}).model_dump() == {"counts": {}} assert model.model_json_schema()["properties"]["counts"]["additionalProperties"] == {"type": "integer"} with pytest.raises(ValidationError): model(counts={"apples": "invalid"}) @pytest.mark.parametrize("additional_properties", [None, True, {}]) def test_untyped_referenced_mapping_preserves_arbitrary_values(additional_properties): mapping = {"type": "object"} if additional_properties is not None: mapping["additionalProperties"] = additional_properties schema = { "type": "object", "properties": {"data": {"$ref": "#/$defs/Data"}}, "required": ["data"], "$defs": {"Data": mapping}, } model = SchemaTransformer.create_model_from_schema(schema, "DataOutput", "data", is_output_schema=True) data = {"name": "sample", "count": 3, "nested": [False, None, {"value": 7}]} assert model(data=data).model_dump() == {"data": data} @pytest.mark.asyncio async def test_generated_tool_preserves_mapping_arguments_and_structured_results(): schema = { "type": "object", "properties": {"groups": {"$ref": "#/$defs/Groups"}}, "required": ["groups"], "$defs": { "Groups": {"type": "object", "additionalProperties": {"type": "array", "items": {"$ref": "#/$defs/Item"}}}, "Item": {"type": "object", "properties": {"count": {"type": "integer"}}, "required": ["count"]}, }, } session = AsyncMock(spec=ClientSession) session.call_tool.return_value = CallToolResult(content=[], structuredContent={"groups": {"result": [{"count": 7}]}}) factory = MCPFactory(client_session=session, event_loop=asyncio.get_running_loop()) tool = factory._create_tool_classes([MCPToolDefinition("group", "Group items", schema, schema)])[0]() groups = {"input": [{"count": 3}]} params = tool.input_schema(tool_name="group", groups=groups) result = await tool.arun(params) session.call_tool.assert_awaited_once_with(name="group", arguments={"groups": groups}) assert result.model_dump() == {"groups": {"result": [{"count": 7}]}} with pytest.raises(ValidationError): tool.input_schema(tool_name="group", groups={"input": [{"count": "invalid"}]}) with pytest.raises(ValidationError): tool.output_schema(groups={"result": [{}]}) ``` ### File: atomic-agents/tests/connectors/mcp/test_mcp_definition_service.py ```python from contextlib import AsyncExitStack from unittest.mock import DEFAULT, AsyncMock, MagicMock, patch import pytest import mcp.types as types from mcp.shared.exceptions import McpError from atomic_agents.connectors.mcp import ( MCPDefinitionService, MCPToolDefinition, MCPResourceDefinition, MCPPromptDefinition, MCPTransportType, ) class MockAsyncContextManager: def __init__(self, return_value=None): self.return_value = return_value self.enter_called = False self.exit_called = False async def __aenter__(self): self.enter_called = True return self.return_value async def __aexit__(self, exc_type, exc_val, exc_tb): self.exit_called = True return False @pytest.fixture def mock_client_session(): mock_session = AsyncMock() # Setup mock responses mock_tool = MagicMock() mock_tool.name = "TestTool" mock_tool.description = "Test tool description" mock_tool.inputSchema = { "type": "object", "properties": {"param1": {"type": "string", "description": "A string parameter"}}, "required": ["param1"], } mock_response = MagicMock() mock_response.tools = [mock_tool] mock_session.list_tools.return_value = mock_response # Setup tool result mock_tool_result = MagicMock() mock_tool_result.content = "Tool result" mock_session.call_tool.return_value = mock_tool_result # Same for resources and prompts mock_resource = MagicMock() mock_resource.name = "TestResource" mock_resource.description = "A test resource" mock_resource.input_schema = {"type": "object", "properties": {"id": {"type": "string"}}} mock_response.resources = [mock_resource] mock_response.uri = "resource://TestResource/{id}" mock_session.list_resources.return_value = mock_response mock_session.list_resource_templates.return_value = types.ListResourceTemplatesResult(resourceTemplates=[]) mock_prompt = MagicMock() mock_prompt.name = "welcome" mock_prompt.description = "Welcome prompt" arguments = [{"name": "id", "description": "The user's ID", "required": True}] mock_prompt.input_schema = { "type": "object", "properties": {arg["name"]: {"type": "string", "description": arg["description"]} for arg in arguments}, "required": [arg["name"] for arg in arguments if arg["required"]], } # ensure list_prompts returns the same response object mock_response.prompts = [mock_prompt] mock_session.list_prompts.return_value = mock_response return mock_session class TestToolDefinitionService: @pytest.mark.asyncio @patch("atomic_agents.connectors.mcp.mcp_definition_service.sse_client") @patch("atomic_agents.connectors.mcp.mcp_definition_service.ClientSession") async def test_fetch_via_sse(self, mock_client_session_cls, mock_sse_client, mock_client_session): # Setup mock_transport = MockAsyncContextManager(return_value=(AsyncMock(), AsyncMock())) mock_sse_client.return_value = mock_transport mock_session = MockAsyncContextManager(return_value=mock_client_session) mock_client_session_cls.return_value = mock_session # Create service service = MCPDefinitionService("http://test-endpoint", transport_type=MCPTransportType.SSE) # Mock the fetch_tool_definitions_from_session to return directly original_method = service.fetch_tool_definitions_from_session service.fetch_tool_definitions_from_session = AsyncMock( return_value=[ MCPToolDefinition( name="MockTool", description="Mock tool for testing", input_schema={"type": "object", "properties": {"param": {"type": "string"}}}, ) ] ) # Execute result = await service.fetch_tool_definitions() # Verify assert len(result) == 1 assert isinstance(result[0], MCPToolDefinition) assert result[0].name == "MockTool" assert result[0].description == "Mock tool for testing" # Restore the original method service.fetch_tool_definitions_from_session = original_method # Same for resources and prompts original_method_resources = service.fetch_resource_definitions_from_session service.fetch_resource_definitions_from_session = AsyncMock( return_value=[ MCPResourceDefinition( name="MockResource", description="Mock resource for testing", uri="resource://MockResource", input_schema={"type": "object", "properties": {}, "required": []}, ) ] ) resource_result = await service.fetch_resource_definitions() assert len(resource_result) == 1 assert isinstance(resource_result[0], MCPResourceDefinition) assert resource_result[0].name == "MockResource" assert resource_result[0].description == "Mock resource for testing" service.fetch_resource_definitions_from_session = original_method_resources original_method_prompts = service.fetch_prompt_definitions_from_session service.fetch_prompt_definitions_from_session = AsyncMock( return_value=[ MCPPromptDefinition( name="welcome", description="Welcome prompt", input_schema={"type": "object", "properties": {}, "required": []}, ) ] ) prompt_result = await service.fetch_prompt_definitions() assert len(prompt_result) == 1 assert isinstance(prompt_result[0], MCPPromptDefinition) assert prompt_result[0].name == "welcome" assert prompt_result[0].description == "Welcome prompt" service.fetch_prompt_definitions_from_session = original_method_prompts @pytest.mark.asyncio @patch("atomic_agents.connectors.mcp.mcp_definition_service.streamablehttp_client") @patch("atomic_agents.connectors.mcp.mcp_definition_service.ClientSession") async def test_fetch_via_http_stream(self, mock_client_session_cls, mock_http_client, mock_client_session): # Setup mock_transport = MockAsyncContextManager(return_value=(AsyncMock(), AsyncMock(), AsyncMock())) mock_http_client.return_value = mock_transport mock_session = MockAsyncContextManager(return_value=mock_client_session) mock_client_session_cls.return_value = mock_session # Create service with HTTP_STREAM transport service = MCPDefinitionService("http://test-endpoint", transport_type=MCPTransportType.HTTP_STREAM) # Mock the fetch_tool_definitions_from_session to return directly original_method = service.fetch_tool_definitions_from_session service.fetch_tool_definitions_from_session = AsyncMock( return_value=[ MCPToolDefinition( name="MockTool", description="Mock tool for testing", input_schema={"type": "object", "properties": {"param": {"type": "string"}}}, ) ] ) # Execute result = await service.fetch_tool_definitions() # Verify assert len(result) == 1 assert isinstance(result[0], MCPToolDefinition) assert result[0].name == "MockTool" assert result[0].description == "Mock tool for testing" # Verify HTTP client was called with correct endpoint (should have /mcp/ suffix) mock_http_client.assert_called_once_with("http://test-endpoint/mcp/") # Restore the original method service.fetch_tool_definitions_from_session = original_method # Same for resources and prompts original_method_resources = service.fetch_resource_definitions_from_session service.fetch_resource_definitions_from_session = AsyncMock( return_value=[ MCPResourceDefinition( name="MockResource", description="Mock resource for testing", uri="resource://MockResource", input_schema={"type": "object", "properties": {}, "required": []}, ) ] ) resource_result = await service.fetch_resource_definitions() assert len(resource_result) == 1 assert isinstance(resource_result[0], MCPResourceDefinition) assert resource_result[0].name == "MockResource" assert resource_result[0].description == "Mock resource for testing" service.fetch_resource_definitions_from_session = original_method_resources original_method_prompts = service.fetch_prompt_definitions_from_session service.fetch_prompt_definitions_from_session = AsyncMock( return_value=[ MCPPromptDefinition( name="welcome", description="Welcome prompt", input_schema={"type": "object", "properties": {}, "required": []}, ) ] ) prompt_result = await service.fetch_prompt_definitions() assert len(prompt_result) == 1 assert isinstance(prompt_result[0], MCPPromptDefinition) assert prompt_result[0].name == "welcome" assert prompt_result[0].description == "Welcome prompt" service.fetch_prompt_definitions_from_session = original_method_prompts @pytest.mark.asyncio async def test_fetch_via_stdio(self): # Create service service = MCPDefinitionService("command arg1 arg2", MCPTransportType.STDIO) # Mock the fetch_tool_definitions_from_session method service.fetch_tool_definitions_from_session = AsyncMock( return_value=[ MCPToolDefinition( name="MockTool", description="Mock tool for testing", input_schema={"type": "object", "properties": {"param": {"type": "string"}}}, ) ] ) service.fetch_resource_definitions_from_session = AsyncMock( return_value=[ MCPResourceDefinition( name="MockResource", description="Mock resource for testing", uri="resource://MockResource", input_schema={"type": "object", "properties": {"id": {"type": "string"}}}, ) ] ) service.fetch_prompt_definitions_from_session = AsyncMock( return_value=[ MCPPromptDefinition( name="welcome", description="Welcome prompt", # arguments=[{"name": "id", "description": "The user's ID", "required": True}], input_schema={"type": "object", "properties": {"id": {"type": "string"}}}, ) ] ) # Patch the stdio_client to avoid actual subprocess execution with patch("atomic_agents.connectors.mcp.mcp_definition_service.stdio_client") as mock_stdio: mock_transport = MockAsyncContextManager(return_value=(AsyncMock(), AsyncMock())) mock_stdio.return_value = mock_transport with patch("atomic_agents.connectors.mcp.mcp_definition_service.ClientSession") as mock_session_cls: mock_session = MockAsyncContextManager(return_value=AsyncMock()) mock_session_cls.return_value = mock_session # Execute result = await service.fetch_tool_definitions() # Verify assert len(result) == 1 assert result[0].name == "MockTool" # Same for resources and prompts resource_result = await service.fetch_resource_definitions() assert len(resource_result) == 1 assert resource_result[0].name == "MockResource" prompt_result = await service.fetch_prompt_definitions() assert len(prompt_result) == 1 assert prompt_result[0].name == "welcome" @pytest.mark.asyncio async def test_stdio_empty_command(self): # Create service with empty command service = MCPDefinitionService("", MCPTransportType.STDIO) # Test that ValueError is raised for empty command with pytest.raises(ValueError, match="Endpoint is required"): await service.fetch_tool_definitions() with pytest.raises(ValueError, match="Endpoint is required"): await service.fetch_resource_definitions() with pytest.raises(ValueError, match="Endpoint is required"): await service.fetch_prompt_definitions() @pytest.mark.asyncio @pytest.mark.parametrize( ("fetch_method", "error_action"), [ ("fetch_tool_definitions", "tool definition fetching"), ("fetch_resource_definitions", "resource fetching"), ("fetch_prompt_definitions", "prompt fetching"), ], ) async def test_unsupported_transport_rejects_before_connecting_and_closes_stack( self, fetch_method: str, error_action: str, mock_client_session: AsyncMock, monkeypatch: pytest.MonkeyPatch ): service = MCPDefinitionService("http://test-endpoint") monkeypatch.setattr(service, "transport_type", "unsupported") stack = AsyncExitStack() with ( patch.multiple( "atomic_agents.connectors.mcp.mcp_definition_service", stdio_client=DEFAULT, streamablehttp_client=DEFAULT, sse_client=DEFAULT, ClientSession=DEFAULT, AsyncExitStack=DEFAULT, ) as connection_factories, patch.object(stack, "aclose", wraps=stack.aclose) as close_stack, patch.object(stack, "enter_async_context", wraps=stack.enter_async_context) as enter_context, ): connection_factories["AsyncExitStack"].return_value = stack connection_factories["ClientSession"].return_value = MockAsyncContextManager(mock_client_session) with pytest.raises(RuntimeError) as raised_error: await getattr(service, fetch_method)() expected_cause = "Unknown transport type: unsupported. Available types: ['sse', 'http_stream', 'stdio']" assert isinstance(raised_error.value.__cause__, ValueError) assert str(raised_error.value.__cause__) == expected_cause assert str(raised_error.value) == f"Unexpected error during {error_action}: {expected_cause}" connection_factories["stdio_client"].assert_not_called() connection_factories["streamablehttp_client"].assert_not_called() connection_factories["sse_client"].assert_not_called() connection_factories["ClientSession"].assert_not_called() enter_context.assert_not_awaited() close_stack.assert_awaited_once_with() assert mock_client_session.mock_calls == [] @pytest.mark.asyncio async def test_fetch_tool_definitions_from_session(self, mock_client_session): # Execute using the static method result = await MCPDefinitionService.fetch_tool_definitions_from_session(mock_client_session) # Verify assert len(result) == 1 assert isinstance(result[0], MCPToolDefinition) assert result[0].name == "TestTool" # Verify session initialization mock_client_session.initialize.assert_called_once() mock_client_session.list_tools.assert_called_once() @pytest.mark.asyncio async def test_fetch_resource_definitions_from_session(self, mock_client_session): result = await MCPDefinitionService.fetch_resource_definitions_from_session(mock_client_session) assert len(result) == 1 assert isinstance(result[0], MCPResourceDefinition) assert result[0].name == "TestResource" mock_client_session.initialize.assert_called() mock_client_session.list_resources.assert_called_once() @pytest.mark.asyncio async def test_mapping_adapter_preserves_uri_and_requires_only_explicit_parameters(self, mock_client_session: AsyncMock): adapter_response = MagicMock( resources=[{"name": "user_document", "uri": "documents://user%20files/%7Buser_id%7D/{document_id}"}] ) mock_client_session.list_resources.return_value = adapter_response definitions = await MCPDefinitionService.fetch_resource_definitions_from_session(mock_client_session) assert definitions == [ MCPResourceDefinition( name="user_document", description=None, uri="documents://user%20files/%7Buser_id%7D/{document_id}", input_schema={ "type": "object", "properties": { "document_id": {"type": "string", "description": "URI parameter document_id"}, }, "required": ["document_id"], }, ) ] mock_client_session.initialize.assert_awaited_once_with() mock_client_session.list_resources.assert_awaited_once_with() mock_client_session.list_resource_templates.assert_awaited_once_with() @pytest.mark.asyncio async def test_mapping_adapter_rejects_unrecognized_resource_without_creating_definition( self, mock_client_session: AsyncMock ): adapter_response = MagicMock(resources=[42]) mock_client_session.list_resources.return_value = adapter_response with patch("atomic_agents.connectors.mcp.mcp_definition_service.MCPResourceDefinition") as resource_definition: with pytest.raises(ValueError, match="^Unexpected resource format: 42$"): await MCPDefinitionService.fetch_resource_definitions_from_session(mock_client_session) resource_definition.assert_not_called() mock_client_session.initialize.assert_awaited_once_with() mock_client_session.list_resources.assert_awaited_once_with() mock_client_session.list_resource_templates.assert_awaited_once_with() @pytest.mark.asyncio async def test_fetch_resource_definitions_from_session_includes_templates_with_static_resources(self): mock_session = AsyncMock() mock_session.list_resources.return_value = types.ListResourcesResult( resources=[types.Resource(name="current_config", uri="config://current")] ) mock_session.list_resource_templates.return_value = types.ListResourceTemplatesResult( resourceTemplates=[types.ResourceTemplate(name="user_profile", uriTemplate="users://{user_id}")] ) result = await MCPDefinitionService.fetch_resource_definitions_from_session(mock_session) assert [definition.name for definition in result] == ["current_config", "user_profile"] assert result[1].uri == "users://{user_id}" assert result[1].input_schema == { "type": "object", "properties": {"user_id": {"type": "string", "description": "URI parameter user_id"}}, "required": ["user_id"], } mock_session.list_resources.assert_awaited_once() mock_session.list_resource_templates.assert_awaited_once() @pytest.mark.asyncio async def test_fetch_resource_definitions_from_session_allows_unsupported_templates(self): mock_session = AsyncMock() mock_session.list_resources.return_value = types.ListResourcesResult( resources=[types.Resource(name="current_config", uri="config://current")] ) mock_session.list_resource_templates.side_effect = McpError( types.ErrorData(code=types.METHOD_NOT_FOUND, message="Method not found") ) result = await MCPDefinitionService.fetch_resource_definitions_from_session(mock_session) assert [definition.name for definition in result] == ["current_config"] mock_session.list_resource_templates.assert_awaited_once() @pytest.mark.asyncio @pytest.mark.parametrize("code", [types.METHOD_NOT_FOUND, types.INTERNAL_ERROR]) async def test_resource_template_continuation_error_is_propagated(self, code): mock_session = AsyncMock() mock_session.list_resources.return_value = types.ListResourcesResult(resources=[]) error = McpError(types.ErrorData(code=code, message="Continuation failed")) mock_session.list_resource_templates.side_effect = [ types.ListResourceTemplatesResult( resourceTemplates=[types.ResourceTemplate(name="user", uriTemplate="users://{id}")], nextCursor="next-page", ), error, ] with pytest.raises(McpError) as exc_info: await MCPDefinitionService.fetch_resource_definitions_from_session(mock_session) assert exc_info.value is error assert mock_session.list_resource_templates.await_count == 2 assert mock_session.list_resource_templates.await_args.kwargs == {"cursor": "next-page"} @pytest.mark.asyncio async def test_fetch_resource_definitions_from_session_propagates_template_errors(self): mock_session = AsyncMock() mock_session.list_resources.return_value = types.ListResourcesResult( resources=[types.Resource(name="current_config", uri="config://current")] ) template_error = McpError(types.ErrorData(code=types.INTERNAL_ERROR, message="Template listing failed")) mock_session.list_resource_templates.side_effect = template_error with pytest.raises(McpError) as exc_info: await MCPDefinitionService.fetch_resource_definitions_from_session(mock_session) assert exc_info.value is template_error @pytest.mark.asyncio async def test_fetch_prompt_definitions_from_session(self, mock_client_session): result = await MCPDefinitionService.fetch_prompt_definitions_from_session(mock_client_session) assert len(result) == 1 assert isinstance(result[0], MCPPromptDefinition) assert result[0].name == "welcome" mock_client_session.initialize.assert_called() mock_client_session.list_prompts.assert_called_once() @pytest.mark.asyncio async def test_session_exception(self): mock_session = AsyncMock() mock_session.initialize.side_effect = Exception("Session error") with pytest.raises(Exception, match="Session error"): await MCPDefinitionService.fetch_tool_definitions_from_session(mock_session) with pytest.raises(Exception, match="Session error"): await MCPDefinitionService.fetch_resource_definitions_from_session(mock_session) with pytest.raises(Exception, match="Session error"): await MCPDefinitionService.fetch_prompt_definitions_from_session(mock_session) @pytest.mark.asyncio async def test_null_input_schema(self, mock_client_session): # Create a tool with null inputSchema mock_tool = MagicMock() mock_tool.name = "NullSchemaTool" mock_tool.description = "Tool with null schema" mock_tool.inputSchema = None mock_response = MagicMock() mock_response.tools = [mock_tool] mock_client_session.list_tools.return_value = mock_response # Execute result = await MCPDefinitionService.fetch_tool_definitions_from_session(mock_client_session) # Verify default empty schema is created assert len(result) == 1 assert result[0].name == "NullSchemaTool" # input_schema is {"type": "object", "properties": {}, "required": []} assert result[0].input_schema.get("type") == "object" assert result[0].input_schema.get("properties") == {} # Same for resources and prompts mock_resource = MagicMock() mock_resource.name = "NullSchemaResource" mock_resource.description = "Resource with null schema" mock_resource.uri = "resource://NullSchemaResource" mock_resource.input_schema = None mock_response.resources = [mock_resource] # ensure the session will return this response for list_resources mock_client_session.list_resources.return_value = mock_response resource_result = await MCPDefinitionService.fetch_resource_definitions_from_session(mock_client_session) assert len(resource_result) == 1 assert resource_result[0].name == "NullSchemaResource" assert resource_result[0].input_schema.get("type") == "object" assert resource_result[0].input_schema.get("properties") == {} assert resource_result[0].uri == "resource://NullSchemaResource" # prompts mock_prompt = MagicMock() mock_prompt.name = "NullSchemaPrompt" mock_prompt.description = "Prompt with null schema" mock_prompt.arguments = None mock_prompt.input_schema = None mock_response.prompts = [mock_prompt] mock_client_session.list_prompts.return_value = mock_response prompt_result = await MCPDefinitionService.fetch_prompt_definitions_from_session(mock_client_session) assert len(prompt_result) == 1 assert prompt_result[0].name == "NullSchemaPrompt" assert prompt_result[0].description == "Prompt with null schema" assert prompt_result[0].input_schema.get("type") == "object" assert prompt_result[0].input_schema.get("properties") == {} @pytest.mark.asyncio async def test_stdio_command_parts_empty(self): svc = MCPDefinitionService(" ", MCPTransportType.STDIO) with pytest.raises( RuntimeError, match="Unexpected error during tool definition fetching: STDIO command string cannot be empty" ): await svc.fetch_tool_definitions() with pytest.raises( RuntimeError, match="Unexpected error during resource fetching: STDIO command string cannot be empty" ): await svc.fetch_resource_definitions() with pytest.raises( RuntimeError, match="Unexpected error during prompt fetching: STDIO command string cannot be empty" ): await svc.fetch_prompt_definitions() @pytest.mark.asyncio async def test_sse_connection_error(self): with patch("atomic_agents.connectors.mcp.mcp_definition_service.sse_client", side_effect=ConnectionError): svc = MCPDefinitionService("http://host", transport_type=MCPTransportType.SSE) with pytest.raises(ConnectionError): await svc.fetch_tool_definitions() with pytest.raises(ConnectionError): await svc.fetch_resource_definitions() with pytest.raises(ConnectionError): await svc.fetch_prompt_definitions() @pytest.mark.asyncio async def test_http_stream_connection_error(self): with patch("atomic_agents.connectors.mcp.mcp_definition_service.streamablehttp_client", side_effect=ConnectionError): svc = MCPDefinitionService("http://host", transport_type=MCPTransportType.HTTP_STREAM) with pytest.raises(ConnectionError): await svc.fetch_tool_definitions() with pytest.raises(ConnectionError): await svc.fetch_resource_definitions() with pytest.raises(ConnectionError): await svc.fetch_prompt_definitions() @pytest.mark.asyncio async def test_generic_error_wrapped(self): with patch("atomic_agents.connectors.mcp.mcp_definition_service.sse_client", side_effect=OSError("BOOM")): svc = MCPDefinitionService("http://host", transport_type=MCPTransportType.SSE) with pytest.raises(RuntimeError): await svc.fetch_tool_definitions() with pytest.raises(RuntimeError): await svc.fetch_resource_definitions() with pytest.raises(RuntimeError): await svc.fetch_prompt_definitions() # Helper class for no-tools test class _NoToolsResponse: """Response object that simulates an empty tools list""" tools = [] class _NoResourcesResponse: """Response object that simulates an empty resources list""" resources = [] class _NoPromptsResponse: """Response object that simulates an empty prompts list""" prompts = [] @pytest.mark.asyncio async def test_fetch_tool_definitions_from_session_no_tools(caplog): """Test handling of empty tools list from session""" sess = AsyncMock() sess.initialize = AsyncMock() sess.list_tools = AsyncMock(return_value=_NoToolsResponse()) result = await MCPDefinitionService.fetch_tool_definitions_from_session(sess) assert result == [] assert "No tool definitions found on MCP server" in caplog.text @pytest.mark.asyncio async def test_fetch_resources_from_session_no_resources(caplog): """Test handling of empty resources list from session""" sess = AsyncMock() sess.initialize = AsyncMock() sess.list_resources = AsyncMock(return_value=_NoResourcesResponse()) sess.list_resource_templates = AsyncMock(return_value=types.ListResourceTemplatesResult(resourceTemplates=[])) result = await MCPDefinitionService.fetch_resource_definitions_from_session(sess) assert result == [] assert "No resources found on MCP server" in caplog.text @pytest.mark.asyncio async def test_fetch_prompts_from_session_no_prompts(caplog): """Test handling of empty prompts list from session""" sess = AsyncMock() sess.initialize = AsyncMock() sess.list_prompts = AsyncMock(return_value=_NoPromptsResponse()) result = await MCPDefinitionService.fetch_prompt_definitions_from_session(sess) assert result == [] assert "No prompts found on MCP server" in caplog.text @pytest.mark.asyncio async def test_fetch_resources_from_session(caplog): """Test fetching resources via session""" sess = AsyncMock() sess.initialize = AsyncMock() # Mock resource object as SimpleNamespace-like dict with a URI template mock_resource = MagicMock() mock_resource.name = "TestResource" mock_resource.description = "A test resource" mock_resource.uri = "resource://TestResource/{id}" mock_response = MagicMock() mock_response.resources = [mock_resource] sess.list_resources = AsyncMock(return_value=mock_response) sess.list_resource_templates = AsyncMock(return_value=types.ListResourceTemplatesResult(resourceTemplates=[])) result = await MCPDefinitionService.fetch_resource_definitions_from_session(sess) assert len(result) == 1 rd = result[0] assert rd.name == "TestResource" assert rd.description == "A test resource" assert rd.input_schema["properties"]["id"]["type"] == "string" @pytest.mark.asyncio async def test_fetch_prompts_from_session(caplog): """Test fetching prompts via session""" sess = AsyncMock() sess.initialize = AsyncMock() # Some MCP clients may return prompt objects or dicts; provide arguments as objects mock_prompt = MagicMock() mock_prompt.name = "welcome" mock_prompt.description = "Welcome prompt" arg = MagicMock() arg.name = "name" arg.description = "The user's name" arg.required = True mock_prompt.arguments = [arg] mock_response = MagicMock() mock_response.prompts = [mock_prompt] sess.list_prompts = AsyncMock(return_value=mock_response) result = await MCPDefinitionService.fetch_prompt_definitions_from_session(sess) assert len(result) == 1 pd = result[0] assert pd.name == "welcome" # validate input_schema was constructed from arguments assert pd.input_schema["properties"]["name"]["description"] == "The user's name" @pytest.mark.asyncio async def test_fetch_tool_definitions_with_output_schema(): """Test that outputSchema is captured from MCP tools when available""" sess = AsyncMock() sess.initialize = AsyncMock() # Create a mock tool with outputSchema mock_tool = MagicMock() mock_tool.name = "StructuredTool" mock_tool.description = "A tool with structured output" mock_tool.inputSchema = { "type": "object", "properties": {"query": {"type": "string", "description": "Search query"}}, "required": ["query"], } mock_tool.outputSchema = { "type": "object", "properties": { "results": {"type": "array", "items": {"type": "string"}, "description": "Search results"}, "count": {"type": "integer", "description": "Number of results"}, }, "required": ["results", "count"], } mock_response = MagicMock() mock_response.tools = [mock_tool] sess.list_tools = AsyncMock(return_value=mock_response) result = await MCPDefinitionService.fetch_tool_definitions_from_session(sess) assert len(result) == 1 td = result[0] assert td.name == "StructuredTool" assert td.output_schema is not None assert td.output_schema["properties"]["results"]["type"] == "array" assert td.output_schema["properties"]["count"]["type"] == "integer" @pytest.mark.asyncio async def test_fetch_tool_definitions_without_output_schema(): """Test that output_schema is None when MCP tool doesn't provide outputSchema""" sess = AsyncMock() sess.initialize = AsyncMock() # Create a mock tool without outputSchema mock_tool = MagicMock() mock_tool.name = "SimpleTool" mock_tool.description = "A simple tool without structured output" mock_tool.inputSchema = {"type": "object", "properties": {}} # Simulate tool without outputSchema attribute del mock_tool.outputSchema mock_response = MagicMock() mock_response.tools = [mock_tool] sess.list_tools = AsyncMock(return_value=mock_response) result = await MCPDefinitionService.fetch_tool_definitions_from_session(sess) assert len(result) == 1 td = result[0] assert td.name == "SimpleTool" assert td.output_schema is None ``` ### File: atomic-agents/tests/connectors/mcp/test_mcp_factory.py ```python import pytest from pydantic import BaseModel, ValidationError from mcp import ClientSession, StdioServerParameters from mcp.types import ( CallToolResult, GetPromptResult, ImageContent, ListResourcesResult, ListResourceTemplatesResult, ListToolsResult, PromptMessage, ReadResourceResult, ResourceTemplate, TextContent, Tool, ) from atomic_agents.base import BasePrompt, BaseResource from unittest.mock import AsyncMock, MagicMock, call import asyncio from atomic_agents.connectors.mcp import ( fetch_mcp_tools, fetch_mcp_resources, fetch_mcp_prompts, create_mcp_orchestrator_schema, fetch_mcp_attributes_with_schema, fetch_mcp_tools_async, fetch_mcp_resources_async, fetch_mcp_prompts_async, MCPFactory, ) from atomic_agents.connectors.mcp import ( MCPToolDefinition, MCPResourceDefinition, MCPPromptDefinition, MCPDefinitionService, MCPTransportType, ) @pytest.fixture(params=[False, True], ids=["generic", "typed"]) def error_test_tool(request, monkeypatch): output_schema = None if request.param: output_schema = {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]} definition = MCPToolDefinition( name="SearchTool", description="Search", input_schema={"type": "object", "properties": {}}, output_schema=output_schema ) factory = MCPFactory("http://example.com") tool = factory._create_tool_classes([definition])[0]() session = AsyncMock(spec=ClientSession) session.__aenter__.return_value = session transport = AsyncMock() transport.__aenter__.return_value = (None, None, None) monkeypatch.setattr("atomic_agents.connectors.mcp.mcp_factory.ClientSession", MagicMock(return_value=session)) monkeypatch.setattr("atomic_agents.connectors.mcp.mcp_factory.streamablehttp_client", MagicMock(return_value=transport)) return tool, session, transport, request.param @pytest.fixture(params=[False, True], ids=["owned", "borrowed"]) def error_test_connection(request, error_test_tool): tool, session, transport, typed = error_test_tool loop = asyncio.new_event_loop() if request.param: tool._client_session = session tool._event_loop = loop yield tool, session, transport, typed, request.param loop.close() @pytest.mark.parametrize("synchronous", [False, True], ids=["async", "sync"]) @pytest.mark.parametrize("mapping", [False, True], ids=["sdk", "mapping"]) def test_execution_error_precedes_output_parsing(error_test_connection, synchronous, mapping, caplog): tool, session, transport, typed, borrowed = error_test_connection envelope = CallToolResult( isError=True, content=[ TextContent(type="text", text="Search unavailable"), ImageContent(type="image", data="PRIVATE_BINARY", mimeType="image/png", text=None), ], structuredContent={"name": "PRIVATE_STRUCTURED"}, _meta={"secret": "PRIVATE_META"}, ) session.call_tool.return_value = envelope.model_dump() if mapping else envelope params = tool.input_schema(tool_name="SearchTool") with pytest.raises(RuntimeError, match="Failed to execute MCP tool 'SearchTool': Search unavailable") as raised: if synchronous: tool.run(params) else: asyncio.run(tool.arun(params)) assert isinstance(raised.value.__cause__, RuntimeError) assert str(raised.value.__cause__) == "Search unavailable" assert "PRIVATE_" not in str(raised.value) + caplog.text session.call_tool.assert_awaited_once_with(name="SearchTool", arguments={}) owned_session_calls = int(not borrowed) assert session.initialize.await_count == owned_session_calls assert session.__aexit__.await_count == owned_session_calls assert transport.__aexit__.await_count == owned_session_calls @pytest.mark.parametrize("mapping", [False, True], ids=["sdk", "mapping"]) @pytest.mark.parametrize( "content", [[], [ImageContent(type="image", data="PRIVATE_BINARY", mimeType="image/png")], [TextContent(type="text", text=" \n ")]], ) def test_execution_error_without_text_has_safe_fallback(error_test_connection, mapping, content, caplog): tool, session, transport, typed, borrowed = error_test_connection envelope = CallToolResult(isError=True, content=content, structuredContent={"name": "PRIVATE_STRUCTURED"}) session.call_tool.return_value = envelope.model_dump() if mapping else envelope with pytest.raises(RuntimeError, match="MCP tool reported an error without text content") as raised: asyncio.run(tool.arun(tool.input_schema(tool_name="SearchTool"))) assert "PRIVATE_" not in str(raised.value) + caplog.text @pytest.mark.parametrize("mapping", [False, True], ids=["sdk", "mapping"]) @pytest.mark.parametrize("synchronous", [False, True], ids=["async", "sync"]) def test_successful_envelope_keeps_output_precedence(error_test_connection, mapping, synchronous): tool, session, transport, typed, borrowed = error_test_connection envelope = CallToolResult( isError=False, content=[TextContent(type="text", text='{"name":"text"}')], structuredContent={"name": "structured"} ) session.call_tool.return_value = envelope.model_dump() if mapping else envelope params = tool.input_schema(tool_name="SearchTool") output = tool.run(params) if synchronous else asyncio.run(tool.arun(params)) if typed: assert output.name == "structured" else: expected_content = envelope.model_dump()["content"] if mapping else envelope.content assert output.result == expected_content session.call_tool.assert_awaited_once_with(name="SearchTool", arguments={}) @pytest.mark.parametrize("error_test_tool", [False], indirect=True, ids=["generic"]) @pytest.mark.parametrize("structured", [None, {}, {"name": "structured", "count": 0, "items": ["a", "b"]}]) @pytest.mark.parametrize("mapping", [False, True], ids=["sdk", "mapping"]) @pytest.mark.parametrize("synchronous", [False, True], ids=["async", "sync"]) def test_generic_tool_keeps_structured_only_result(error_test_connection, structured, mapping, synchronous): tool, session, transport, typed, borrowed = error_test_connection envelope = CallToolResult(content=[], structuredContent=structured) session.call_tool.return_value = envelope.model_dump() if mapping else envelope params = tool.input_schema(tool_name="SearchTool") output = tool.run(params) if synchronous else asyncio.run(tool.arun(params)) assert output.result == ([] if structured is None else structured) session.call_tool.assert_awaited_once_with(name="SearchTool", arguments={}) @pytest.mark.parametrize("content", [None, "article"]) @pytest.mark.parametrize("synchronous", [False, True], ids=["async", "sync"]) def test_raw_error_named_dictionary_stays_application_data(error_test_connection, content, synchronous): tool, session, transport, typed, borrowed = error_test_connection payload = {"isError": True, "name": "test_value"} if content is not None: payload["content"] = content session.call_tool.return_value = payload params = tool.input_schema(tool_name="SearchTool") output = tool.run(params) if synchronous else asyncio.run(tool.arun(params)) if typed: assert output.name == "test_value" else: assert output.result == (payload if content is None else content) session.call_tool.assert_awaited_once_with(name="SearchTool", arguments={}) class DummySession: pass def test_fetch_mcp_tools_no_endpoint_raises(): with pytest.raises(ValueError): fetch_mcp_tools() def test_fetch_mcp_resources_no_endpoint_raises(): with pytest.raises(ValueError): fetch_mcp_resources() def test_fetch_mcp_prompts_no_endpoint_raises(): with pytest.raises(ValueError): fetch_mcp_prompts() def test_fetch_mcp_tools_event_loop_without_client_session_raises(): with pytest.raises(ValueError): fetch_mcp_tools(None, MCPTransportType.HTTP_STREAM, client_session=DummySession(), event_loop=None) def test_fetch_mcp_resources_event_loop_without_client_session_raises(): with pytest.raises(ValueError): fetch_mcp_resources(None, MCPTransportType.HTTP_STREAM, client_session=DummySession(), event_loop=None) def test_fetch_mcp_prompts_event_loop_without_client_session_raises(): with pytest.raises(ValueError): fetch_mcp_prompts(None, MCPTransportType.HTTP_STREAM, client_session=DummySession(), event_loop=None) def test_fetch_mcp_tools_empty_definitions(monkeypatch): monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: []) tools = fetch_mcp_tools("http://example.com", MCPTransportType.HTTP_STREAM) assert tools == [] def test_fetch_mcp_resources_empty_definitions(monkeypatch): monkeypatch.setattr(MCPFactory, "_fetch_resource_definitions", lambda self: []) resources = fetch_mcp_resources("http://example.com", MCPTransportType.HTTP_STREAM) assert resources == [] def test_fetch_mcp_prompts_empty_definitions(monkeypatch): monkeypatch.setattr(MCPFactory, "_fetch_prompt_definitions", lambda self: []) prompts = fetch_mcp_prompts("http://example.com", MCPTransportType.HTTP_STREAM) assert prompts == [] def test_fetch_mcp_tools_with_definitions_http(monkeypatch): input_schema = {"type": "object", "properties": {}, "required": []} definitions = [MCPToolDefinition(name="ToolX", description="Dummy tool", input_schema=input_schema)] monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: definitions) tools = fetch_mcp_tools("http://example.com", MCPTransportType.HTTP_STREAM) assert len(tools) == 1 tool_cls = tools[0] # verify class attributes assert tool_cls.mcp_endpoint == "http://example.com" assert tool_cls.transport_type == MCPTransportType.HTTP_STREAM # input_schema has only tool_name field Model = tool_cls.input_schema assert "tool_name" in Model.model_fields # output_schema has result field (generic schema) OutModel = tool_cls.output_schema assert "result" in OutModel.model_fields # verify _has_typed_output_schema is False for generic schema assert tool_cls._has_typed_output_schema is False @pytest.mark.parametrize("nullable_schema_location", ["input", "output"]) def test_fetch_mcp_tools_preserves_nullable_schema(monkeypatch, nullable_schema_location): """A tool must not disappear when its input or output contains a type array.""" nullable_schema = { "type": "object", "properties": {"query": {"type": ["string", "null"]}}, "required": ["query"], } definition = MCPToolDefinition( name="SearchTool", description="Search tool", input_schema=nullable_schema if nullable_schema_location == "input" else {"type": "object"}, output_schema=nullable_schema if nullable_schema_location == "output" else None, ) monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: [definition]) tools = fetch_mcp_tools("http://example.com", MCPTransportType.HTTP_STREAM) assert len(tools) == 1 tool_cls = tools[0] if nullable_schema_location == "input": assert tool_cls.input_schema(tool_name="SearchTool", query=None).query is None else: assert tool_cls._has_typed_output_schema is True assert tool_cls.output_schema(query=None).query is None def test_fetch_mcp_tools_with_typed_output_schema(monkeypatch): """Test that tools with outputSchema get typed output models""" input_schema = {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]} output_schema = { "type": "object", "properties": { "results": {"type": "array", "items": {"type": "string"}, "description": "Search results"}, "count": {"type": "integer", "description": "Number of results"}, }, "required": ["results", "count"], } definitions = [ MCPToolDefinition( name="SearchTool", description="A tool with typed output", input_schema=input_schema, output_schema=output_schema ) ] monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: definitions) tools = fetch_mcp_tools("http://example.com", MCPTransportType.HTTP_STREAM) assert len(tools) == 1 tool_cls = tools[0] # verify class attributes assert tool_cls.mcp_endpoint == "http://example.com" assert tool_cls._has_typed_output_schema is True # input_schema has tool_name and query fields Model = tool_cls.input_schema assert "tool_name" in Model.model_fields assert "query" in Model.model_fields # output_schema has typed fields instead of generic 'result' OutModel = tool_cls.output_schema assert "results" in OutModel.model_fields assert "count" in OutModel.model_fields # Should NOT have the generic 'result' field assert "result" not in OutModel.model_fields # Should NOT have the tool_name field (output schemas don't need it) assert "tool_name" not in OutModel.model_fields def test_fetch_mcp_tools_mixed_output_schemas(monkeypatch): """Test that tools with and without outputSchema are handled correctly together""" input_schema = {"type": "object", "properties": {}, "required": []} output_schema = { "type": "object", "properties": {"data": {"type": "string"}}, "required": ["data"], } definitions = [ MCPToolDefinition(name="GenericTool", description="No output schema", input_schema=input_schema), MCPToolDefinition( name="TypedTool", description="With output schema", input_schema=input_schema, output_schema=output_schema ), ] monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: definitions) tools = fetch_mcp_tools("http://example.com", MCPTransportType.HTTP_STREAM) assert len(tools) == 2 # First tool should have generic output generic_tool = tools[0] assert generic_tool._has_typed_output_schema is False assert "result" in generic_tool.output_schema.model_fields # Second tool should have typed output typed_tool = tools[1] assert typed_tool._has_typed_output_schema is True assert "data" in typed_tool.output_schema.model_fields assert "result" not in typed_tool.output_schema.model_fields # ============================================================================= # Tests for typed output schema result processing # ============================================================================= class MockStructuredContentResult(BaseModel): """Mock MCP result with structuredContent attribute (MCP spec primary path)""" structuredContent: dict class MockContentItem(BaseModel): """Mock content item with text attribute""" text: str class MockContentItemWithData(BaseModel): """Mock content item with data attribute""" data: dict class MockContentResult(BaseModel): """Mock MCP result with content array""" content: list @pytest.mark.asyncio async def test_typed_output_schema_with_structured_content_dict(monkeypatch): """Test result processing when tool returns BaseModel with structuredContent as dict""" input_schema = {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]} output_schema = { "type": "object", "properties": { "results": {"type": "array", "items": {"type": "string"}}, "count": {"type": "integer"}, }, "required": ["results", "count"], } definitions = [ MCPToolDefinition(name="SearchTool", description="Search tool", input_schema=input_schema, output_schema=output_schema) ] monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: definitions) tools = fetch_mcp_tools("http://example.com", MCPTransportType.HTTP_STREAM) tool_cls = tools[0] tool_instance = tool_cls() mock_result = MockStructuredContentResult(structuredContent={"results": ["a", "b"], "count": 2}) import atomic_agents.connectors.mcp.mcp_factory as factory_module class MockClientSession: def __init__(self, *args): pass async def __aenter__(self): return self async def __aexit__(self, *args): pass async def initialize(self): pass async def call_tool(self, name, arguments): return mock_result class MockHttpClient: async def __aenter__(self): return (None, None, None) async def __aexit__(self, *args): pass monkeypatch.setattr(factory_module, "ClientSession", MockClientSession) monkeypatch.setattr(factory_module, "streamablehttp_client", lambda *args, **kwargs: MockHttpClient()) InputSchema = tool_cls.input_schema params = InputSchema(tool_name="SearchTool", query="test") result = await tool_instance.arun(params) assert result.results == ["a", "b"] assert result.count == 2 @pytest.mark.asyncio async def test_typed_output_schema_with_json_text_content(monkeypatch): """Test result processing when tool returns content[0].text as JSON string""" input_schema = {"type": "object", "properties": {}, "required": []} output_schema = { "type": "object", "properties": {"data": {"type": "string"}}, "required": ["data"], } definitions = [ MCPToolDefinition(name="JsonTool", description="JSON tool", input_schema=input_schema, output_schema=output_schema) ] monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: definitions) tools = fetch_mcp_tools("http://example.com", MCPTransportType.HTTP_STREAM) tool_cls = tools[0] tool_instance = tool_cls() mock_content_item = MockContentItem(text='{"data": "hello"}') mock_result = MockContentResult(content=[mock_content_item]) import atomic_agents.connectors.mcp.mcp_factory as factory_module class MockClientSession: def __init__(self, *args): pass async def __aenter__(self): return self async def __aexit__(self, *args): pass async def initialize(self): pass async def call_tool(self, name, arguments): return mock_result class MockHttpClient: async def __aenter__(self): return (None, None, None) async def __aexit__(self, *args): pass monkeypatch.setattr(factory_module, "ClientSession", MockClientSession) monkeypatch.setattr(factory_module, "streamablehttp_client", lambda *args, **kwargs: MockHttpClient()) InputSchema = tool_cls.input_schema params = InputSchema(tool_name="JsonTool") result = await tool_instance.arun(params) assert result.data == "hello" @pytest.mark.asyncio async def test_typed_output_schema_with_content_data_dict(monkeypatch): """Test result processing when content item has .data attribute as dict""" input_schema = {"type": "object", "properties": {}, "required": []} output_schema = { "type": "object", "properties": {"value": {"type": "integer"}}, "required": ["value"], } definitions = [ MCPToolDefinition(name="DataTool", description="Data tool", input_schema=input_schema, output_schema=output_schema) ] monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: definitions) tools = fetch_mcp_tools("http://example.com", MCPTransportType.HTTP_STREAM) tool_cls = tools[0] tool_instance = tool_cls() mock_content_item = MockContentItemWithData(data={"value": 42}) mock_result = MockContentResult(content=[mock_content_item]) import atomic_agents.connectors.mcp.mcp_factory as factory_module class MockClientSession: def __init__(self, *args): pass async def __aenter__(self): return self async def __aexit__(self, *args): pass async def initialize(self): pass async def call_tool(self, name, arguments): return mock_result class MockHttpClient: async def __aenter__(self): return (None, None, None) async def __aexit__(self, *args): pass monkeypatch.setattr(factory_module, "ClientSession", MockClientSession) monkeypatch.setattr(factory_module, "streamablehttp_client", lambda *args, **kwargs: MockHttpClient()) InputSchema = tool_cls.input_schema params = InputSchema(tool_name="DataTool") result = await tool_instance.arun(params) assert result.value == 42 @pytest.mark.asyncio async def test_typed_output_schema_with_raw_dict(monkeypatch): """Test fallback when tool_result is plain dict""" input_schema = {"type": "object", "properties": {}, "required": []} output_schema = { "type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"], } definitions = [ MCPToolDefinition(name="DictTool", description="Dict tool", input_schema=input_schema, output_schema=output_schema) ] monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: definitions) tools = fetch_mcp_tools("http://example.com", MCPTransportType.HTTP_STREAM) tool_cls = tools[0] tool_instance = tool_cls() mock_result = {"name": "test_value"} import atomic_agents.connectors.mcp.mcp_factory as factory_module class MockClientSession: def __init__(self, *args): pass async def __aenter__(self): return self async def __aexit__(self, *args): pass async def initialize(self): pass async def call_tool(self, name, arguments): return mock_result class MockHttpClient: async def __aenter__(self): return (None, None, None) async def __aexit__(self, *args): pass monkeypatch.setattr(factory_module, "ClientSession", MockClientSession) monkeypatch.setattr(factory_module, "streamablehttp_client", lambda *args, **kwargs: MockHttpClient()) InputSchema = tool_cls.input_schema params = InputSchema(tool_name="DictTool") result = await tool_instance.arun(params) assert result.name == "test_value" @pytest.mark.asyncio async def test_typed_output_schema_raises_on_unparseable_result(monkeypatch): """Test that ValueError is raised when typed schema tool returns unparseable result""" input_schema = {"type": "object", "properties": {}, "required": []} output_schema = { "type": "object", "properties": {"data": {"type": "string"}}, "required": ["data"], } definitions = [ MCPToolDefinition( name="FailingTool", description="Failing tool", input_schema=input_schema, output_schema=output_schema ) ] monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: definitions) tools = fetch_mcp_tools("http://example.com", MCPTransportType.HTTP_STREAM) tool_cls = tools[0] tool_instance = tool_cls() # Return a string which can't be parsed as structured content mock_result = "just a string, not structured" import atomic_agents.connectors.mcp.mcp_factory as factory_module class MockClientSession: def __init__(self, *args): pass async def __aenter__(self): return self async def __aexit__(self, *args): pass async def initialize(self): pass async def call_tool(self, name, arguments): return mock_result class MockHttpClient: async def __aenter__(self): return (None, None, None) async def __aexit__(self, *args): pass monkeypatch.setattr(factory_module, "ClientSession", MockClientSession) monkeypatch.setattr(factory_module, "streamablehttp_client", lambda *args, **kwargs: MockHttpClient()) InputSchema = tool_cls.input_schema params = InputSchema(tool_name="FailingTool") with pytest.raises(RuntimeError) as exc_info: await tool_instance.arun(params) # The ValueError gets wrapped in RuntimeError by the outer exception handler assert "unparseable result" in str(exc_info.value) or "FailingTool" in str(exc_info.value) @pytest.mark.asyncio async def test_typed_output_schema_handles_empty_content_array(monkeypatch): """Test graceful handling when content array is empty""" input_schema = {"type": "object", "properties": {}, "required": []} output_schema = { "type": "object", "properties": {"data": {"type": "string"}}, "required": ["data"], } definitions = [ MCPToolDefinition( name="EmptyContentTool", description="Empty content tool", input_schema=input_schema, output_schema=output_schema ) ] monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: definitions) tools = fetch_mcp_tools("http://example.com", MCPTransportType.HTTP_STREAM) tool_cls = tools[0] tool_instance = tool_cls() # Empty content array - should fall through and raise error mock_result = MockContentResult(content=[]) import atomic_agents.connectors.mcp.mcp_factory as factory_module class MockClientSession: def __init__(self, *args): pass async def __aenter__(self): return self async def __aexit__(self, *args): pass async def initialize(self): pass async def call_tool(self, name, arguments): return mock_result class MockHttpClient: async def __aenter__(self): return (None, None, None) async def __aexit__(self, *args): pass monkeypatch.setattr(factory_module, "ClientSession", MockClientSession) monkeypatch.setattr(factory_module, "streamablehttp_client", lambda *args, **kwargs: MockHttpClient()) InputSchema = tool_cls.input_schema params = InputSchema(tool_name="EmptyContentTool") # Should raise error since we can't extract structured content with pytest.raises(RuntimeError) as exc_info: await tool_instance.arun(params) assert "EmptyContentTool" in str(exc_info.value) def test_fetch_mcp_resources_with_definitions_stdio(monkeypatch): input_schema = {"type": "object", "properties": {}, "required": []} uri = "resource://example-resource" definitions = [MCPResourceDefinition(name="ResY", description="Dummy resource", uri=uri, input_schema=input_schema)] monkeypatch.setattr(MCPFactory, "_fetch_resource_definitions", lambda self: definitions) resources = fetch_mcp_resources("run me", MCPTransportType.STDIO, working_directory="/tmp") assert len(resources) == 1 res_cls = resources[0] # verify class attributes assert res_cls.mcp_endpoint == "run me" assert res_cls.transport_type == MCPTransportType.STDIO assert res_cls.working_directory == "/tmp" # input_schema has only resource_name field Model = res_cls.input_schema assert "resource_name" in Model.model_fields # output_schema has content field for resources OutModel = res_cls.output_schema assert "content" in OutModel.model_fields def test_fetch_mcp_prompts_with_definitions_http(monkeypatch): input_schema = {"type": "object", "properties": {}, "required": []} definitions = [MCPPromptDefinition(name="PromptZ", description="Dummy prompt", input_schema=input_schema)] monkeypatch.setattr(MCPFactory, "_fetch_prompt_definitions", lambda self: definitions) prompts = fetch_mcp_prompts("http://example.com", MCPTransportType.HTTP_STREAM) assert len(prompts) == 1 prompt_cls = prompts[0] # verify class attributes assert prompt_cls.mcp_endpoint == "http://example.com" assert prompt_cls.transport_type == MCPTransportType.HTTP_STREAM # input_schema has only prompt_name field Model = prompt_cls.input_schema assert "prompt_name" in Model.model_fields # output_schema has content field for prompts OutModel = prompt_cls.output_schema assert "content" in OutModel.model_fields def test_create_mcp_orchestrator_schema_empty(): schema = create_mcp_orchestrator_schema([], [], []) assert schema is None def test_create_mcp_orchestrator_schema_with_tools(): class FakeInput(BaseModel): tool_name: str param: int class FakeTool: input_schema = FakeInput mcp_tool_name = "FakeTool" schema = create_mcp_orchestrator_schema(tools=[FakeTool], resources=[], prompts=[]) assert schema is not None assert "tool_parameters" in schema.model_fields inst = schema(tool_parameters=FakeInput(tool_name="FakeTool", param=1)) assert inst.tool_parameters.param == 1 def test_create_mcp_orchestrator_schema_with_resources(): class FakeInput(BaseModel): resource_name: str param: int class FakeResource: input_schema = FakeInput mcp_resource_name = "FakeResource" schema = create_mcp_orchestrator_schema(resources=[FakeResource]) assert schema is not None assert "resource_parameters" in schema.model_fields inst = schema(resource_parameters=FakeInput(resource_name="FakeResource", param=2)) assert inst.resource_parameters.param == 2 def test_create_mcp_orchestrator_schema_with_prompts(): class FakeInput(BaseModel): prompt_name: str param: int class FakePrompt: input_schema = FakeInput mcp_prompt_name = "FakePrompt" schema = create_mcp_orchestrator_schema(prompts=[FakePrompt]) assert schema is not None assert "prompt_parameters" in schema.model_fields inst = schema(prompt_parameters=FakeInput(prompt_name="FakePrompt", param=3)) assert inst.prompt_parameters.param == 3 def test_fetch_mcp_attributes_with_schema_no_endpoint_raises(): with pytest.raises(ValueError): fetch_mcp_attributes_with_schema() def test_fetch_mcp_attributes_with_schema_empty(monkeypatch): monkeypatch.setattr(MCPFactory, "create_tools", lambda self: []) monkeypatch.setattr(MCPFactory, "create_resources", lambda self: []) monkeypatch.setattr(MCPFactory, "create_prompts", lambda self: []) tools, resources, prompts, schema = fetch_mcp_attributes_with_schema("endpoint", MCPTransportType.HTTP_STREAM) assert tools == [] assert resources == [] assert prompts == [] assert schema is None def test_fetch_mcp_attributes_with_schema_nonempty(monkeypatch): dummy_tools = ["a", "b"] dummy_resources = ["c", "d"] dummy_prompts = ["e", "f"] dummy_schema = object() monkeypatch.setattr(MCPFactory, "create_tools", lambda self: dummy_tools) monkeypatch.setattr(MCPFactory, "create_resources", lambda self: dummy_resources) monkeypatch.setattr(MCPFactory, "create_prompts", lambda self: dummy_prompts) monkeypatch.setattr(MCPFactory, "create_orchestrator_schema", lambda self, tools, resources, prompts: dummy_schema) tools, resources, prompts, schema = fetch_mcp_attributes_with_schema("endpoint", MCPTransportType.STDIO) assert tools == dummy_tools assert resources == dummy_resources assert prompts == dummy_prompts assert schema is dummy_schema def test_fetch_mcp_tools_with_stdio_and_working_directory(monkeypatch): input_schema = {"type": "object", "properties": {}, "required": []} tool_definitions = [MCPToolDefinition(name="ToolZ", description=None, input_schema=input_schema)] monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: tool_definitions) tools = fetch_mcp_tools("run me", MCPTransportType.STDIO, working_directory="/tmp") assert len(tools) == 1 tool_cls = tools[0] assert tool_cls.transport_type == MCPTransportType.STDIO assert tool_cls.mcp_endpoint == "run me" assert tool_cls.working_directory == "/tmp" def test_fetch_mcp_resources_with_stdio_and_working_directory(monkeypatch): input_schema = {"type": "object", "properties": {}, "required": []} resource_definitions = [ MCPResourceDefinition(name="ResZ", description=None, uri="resource://ResZ", input_schema=input_schema) ] monkeypatch.setattr(MCPFactory, "_fetch_resource_definitions", lambda self: resource_definitions) resources = fetch_mcp_resources("run me", MCPTransportType.STDIO, working_directory="/tmp") assert len(resources) == 1 res_cls = resources[0] assert res_cls.transport_type == MCPTransportType.STDIO assert res_cls.mcp_endpoint == "run me" assert res_cls.working_directory == "/tmp" def test_fetch_mcp_prompts_with_stdio_and_working_directory(monkeypatch): input_schema = {"type": "object", "properties": {}, "required": []} prompt_definitions = [MCPPromptDefinition(name="PromptZ", description=None, input_schema=input_schema)] monkeypatch.setattr(MCPFactory, "_fetch_prompt_definitions", lambda self: prompt_definitions) prompts = fetch_mcp_prompts("run me", MCPTransportType.STDIO, working_directory="/tmp") assert len(prompts) == 1 prompt_cls = prompts[0] assert prompt_cls.transport_type == MCPTransportType.STDIO assert prompt_cls.mcp_endpoint == "run me" assert prompt_cls.working_directory == "/tmp" @pytest.mark.parametrize("transport_type", [MCPTransportType.HTTP_STREAM, MCPTransportType.STDIO]) def test_run_tool(monkeypatch, transport_type): # Setup dummy transports and session import atomic_agents.connectors.mcp.mcp_factory as mtf class DummyTransportCM: def __init__(self, ret): self.ret = ret async def __aenter__(self): return self.ret async def __aexit__(self, exc_type, exc, tb): pass def dummy_sse_client(endpoint): return DummyTransportCM((None, None)) def dummy_stdio_client(params): return DummyTransportCM((None, None)) class DummySessionCM: def __init__(self, rs=None, ws=None): pass async def initialize(self): pass async def call_tool(self, name, arguments): return {"content": f"{name}-{arguments}-ok"} async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): pass monkeypatch.setattr(mtf, "sse_client", dummy_sse_client) monkeypatch.setattr(mtf, "stdio_client", dummy_stdio_client) monkeypatch.setattr(mtf, "ClientSession", DummySessionCM) # Prepare definitions input_schema = {"type": "object", "properties": {}, "required": []} tool_definitions = [MCPToolDefinition(name="ToolA", description="desc", input_schema=input_schema)] monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: tool_definitions) # Run fetch and execute tool endpoint = "cmd run" if transport_type == MCPTransportType.STDIO else "http://e" tools = fetch_mcp_tools( endpoint, transport_type, working_directory="wd" if transport_type == MCPTransportType.STDIO else None ) tool_cls = tools[0] inst = tool_cls() result = inst.run(tool_cls.input_schema(tool_name="ToolA")) assert result.result == "ToolA-{}-ok" @pytest.mark.parametrize("transport_type", [MCPTransportType.HTTP_STREAM, MCPTransportType.STDIO]) def test_read_resource(monkeypatch, transport_type): # Setup dummy transports and session import atomic_agents.connectors.mcp.mcp_factory as mtf class DummyTransportCM: def __init__(self, ret): self.ret = ret async def __aenter__(self): return self.ret async def __aexit__(self, exc_type, exc, tb): pass def dummy_sse_client(endpoint): return DummyTransportCM((None, None)) def dummy_stdio_client(params): return DummyTransportCM((None, None)) class DummySessionCM: def __init__(self, rs=None, ws=None): pass async def initialize(self): pass async def read_resource(self, *args, **kwargs): return {"content": "resource-ResA-ok"} async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): pass monkeypatch.setattr(mtf, "sse_client", dummy_sse_client) monkeypatch.setattr(mtf, "stdio_client", dummy_stdio_client) monkeypatch.setattr(mtf, "ClientSession", DummySessionCM) # Prepare definitions input_schema = {"type": "object", "properties": {}, "required": []} resource_definitions = [ MCPResourceDefinition(name="ResA", description="desc", uri="resource://ResA", input_schema=input_schema) ] monkeypatch.setattr(MCPFactory, "_fetch_resource_definitions", lambda self: resource_definitions) endpoint = "cmd run" if transport_type == MCPTransportType.STDIO else "http://e" # Read data from resource resources = fetch_mcp_resources( endpoint, transport_type, working_directory="wd" if transport_type == MCPTransportType.STDIO else None ) resource_cls = resources[0] inst = resource_cls() result = inst.read(resource_cls.input_schema(resource_name="ResA")) assert result.content["content"] == "resource-ResA-ok" @pytest.mark.parametrize("transport_type", [MCPTransportType.HTTP_STREAM, MCPTransportType.STDIO]) def test_generate_prompt(monkeypatch, transport_type): # Setup dummy transports and session import atomic_agents.connectors.mcp.mcp_factory as mtf class DummyTransportCM: def __init__(self, ret): self.ret = ret async def __aenter__(self): return self.ret async def __aexit__(self, exc_type, exc, tb): pass def dummy_sse_client(endpoint): return DummyTransportCM((None, None)) def dummy_stdio_client(params): return DummyTransportCM((None, None)) class DummySessionCM: def __init__(self, rs=None, ws=None): pass async def initialize(self): pass async def get_prompt(self, *, name, arguments): class Msg(BaseModel): content: str return {"messages": [Msg(content=f"prompt-{name}-{arguments}-ok")]} async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): pass monkeypatch.setattr(mtf, "sse_client", dummy_sse_client) monkeypatch.setattr(mtf, "stdio_client", dummy_stdio_client) monkeypatch.setattr(mtf, "ClientSession", DummySessionCM) # Prepare definitions input_schema = {"type": "object", "properties": {}, "required": []} prompt_definitions = [MCPPromptDefinition(name="PromptA", description="desc", input_schema=input_schema)] monkeypatch.setattr(MCPFactory, "_fetch_prompt_definitions", lambda self: prompt_definitions) endpoint = "cmd run" if transport_type == MCPTransportType.STDIO else "http://e" # Generate prompt prompts = fetch_mcp_prompts( endpoint, transport_type, working_directory="wd" if transport_type == MCPTransportType.STDIO else None ) prompt_cls = prompts[0] inst = prompt_cls() result = inst.generate(prompt_cls.input_schema(prompt_name="PromptA")) assert result.content == "prompt-PromptA-{}-ok" def test_run_tool_with_persistent_session(monkeypatch): import atomic_agents.connectors.mcp.mcp_factory as mtf # Setup persistent client class DummySessionPersistent: async def call_tool(self, name, arguments): return {"content": "persist-ok"} client = DummySessionPersistent() # Stub definition fetch for persistent definitions = [ MCPToolDefinition(name="ToolB", description=None, input_schema={"type": "object", "properties": {}, "required": []}) ] async def fake_fetch_defs(session): return definitions monkeypatch.setattr(mtf.MCPDefinitionService, "fetch_tool_definitions_from_session", staticmethod(fake_fetch_defs)) # Create and pass an event loop loop = asyncio.new_event_loop() try: tools = fetch_mcp_tools(None, MCPTransportType.HTTP_STREAM, client_session=client, event_loop=loop) tool_cls = tools[0] inst = tool_cls() result = inst.run(tool_cls.input_schema(tool_name="ToolB")) assert result.result == "persist-ok" finally: loop.close() def test_read_resource_with_persistent_session(monkeypatch): import atomic_agents.connectors.mcp.mcp_factory as mtf # Setup persistent client that matches factory expectations class DummySessionPersistent: async def read_resource(self, *, uri): return {"content": "persist-resource-ok"} client = DummySessionPersistent() # Stub definition fetch for persistent definitions = [ MCPResourceDefinition( name="ResB", description=None, uri="resource://ResB", input_schema={"type": "object", "properties": {}, "required": []}, ) ] async def fake_fetch_defs(session): return definitions monkeypatch.setattr(mtf.MCPDefinitionService, "fetch_resource_definitions_from_session", staticmethod(fake_fetch_defs)) # Create and pass an event loop loop = asyncio.new_event_loop() try: resources = fetch_mcp_resources(None, MCPTransportType.HTTP_STREAM, client_session=client, event_loop=loop) res_cls = resources[0] inst = res_cls() result = inst.read(res_cls.input_schema(resource_name="ResB")) assert result.content["content"] == "persist-resource-ok" finally: loop.close() def test_generate_prompt_with_persistent_session(monkeypatch): import atomic_agents.connectors.mcp.mcp_factory as mtf # Setup persistent client class DummySessionPersistent: async def get_prompt(self, *, name, arguments): class Msg(BaseModel): content: str return {"messages": [Msg(content="persist-prompt-ok")]} client = DummySessionPersistent() # Stub definition fetch for persistent definitions = [ MCPPromptDefinition( name="PromptB", description=None, input_schema={"type": "object", "properties": {}, "required": []} ) ] async def fake_fetch_defs(session): return definitions monkeypatch.setattr(mtf.MCPDefinitionService, "fetch_prompt_definitions_from_session", staticmethod(fake_fetch_defs)) # Create and pass an event loop loop = asyncio.new_event_loop() try: prompts = fetch_mcp_prompts(None, MCPTransportType.HTTP_STREAM, client_session=client, event_loop=loop) prompt_cls = prompts[0] inst = prompt_cls() result = inst.generate(prompt_cls.input_schema(prompt_name="PromptB")) assert result.content == "persist-prompt-ok" finally: loop.close() def test_fetch_tool_definitions_via_service(monkeypatch): from atomic_agents.connectors.mcp.mcp_factory import MCPFactory from atomic_agents.connectors.mcp.mcp_definition_service import MCPToolDefinition defs = [MCPToolDefinition(name="X", description="d", input_schema={"type": "object", "properties": {}, "required": []})] def fake_fetch(self): return defs monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", fake_fetch) factory_http = MCPFactory("http://e", MCPTransportType.HTTP_STREAM) assert factory_http._fetch_tool_definitions() == defs factory_stdio = MCPFactory("http://e", MCPTransportType.STDIO, working_directory="/tmp") assert factory_stdio._fetch_tool_definitions() == defs def test_fetch_resource_definitions_via_service(monkeypatch): from atomic_agents.connectors.mcp.mcp_factory import MCPFactory from atomic_agents.connectors.mcp.mcp_definition_service import MCPResourceDefinition defs = [ MCPResourceDefinition( name="Y", description="d", uri="resource://Y", input_schema={"type": "object", "properties": {}, "required": []} ) ] def fake_fetch(self): return defs monkeypatch.setattr(MCPFactory, "_fetch_resource_definitions", fake_fetch) factory_http = MCPFactory("http://e", MCPTransportType.HTTP_STREAM) assert factory_http._fetch_resource_definitions() == defs factory_stdio = MCPFactory("http://e", MCPTransportType.STDIO, working_directory="/tmp") assert factory_stdio._fetch_resource_definitions() == defs def test_fetch_prompt_definitions_via_service(monkeypatch): from atomic_agents.connectors.mcp.mcp_factory import MCPFactory from atomic_agents.connectors.mcp.mcp_definition_service import MCPPromptDefinition defs = [MCPPromptDefinition(name="Z", description="d", input_schema={"type": "object", "properties": {}, "required": []})] def fake_fetch(self): return defs monkeypatch.setattr(MCPFactory, "_fetch_prompt_definitions", fake_fetch) factory_http = MCPFactory("http://e", MCPTransportType.HTTP_STREAM) assert factory_http._fetch_prompt_definitions() == defs factory_stdio = MCPFactory("http://e", MCPTransportType.STDIO, working_directory="/tmp") assert factory_stdio._fetch_prompt_definitions() == defs def test_fetch_tool_definitions_propagates_error(monkeypatch): from atomic_agents.connectors.mcp.mcp_factory import MCPFactory def fake_fetch(self): raise RuntimeError("nope") monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", fake_fetch) factory = MCPFactory("http://e", MCPTransportType.HTTP_STREAM) with pytest.raises(RuntimeError): factory._fetch_tool_definitions() def test_fetch_resource_definitions_propagates_error(monkeypatch): from atomic_agents.connectors.mcp.mcp_factory import MCPFactory def fake_fetch(self): raise RuntimeError("nope") monkeypatch.setattr(MCPFactory, "_fetch_resource_definitions", fake_fetch) factory = MCPFactory("http://e", MCPTransportType.HTTP_STREAM) with pytest.raises(RuntimeError): factory._fetch_resource_definitions() def test_fetch_prompt_definitions_propagates_error(monkeypatch): from atomic_agents.connectors.mcp.mcp_factory import MCPFactory def fake_fetch(self): raise RuntimeError("nope") monkeypatch.setattr(MCPFactory, "_fetch_prompt_definitions", fake_fetch) factory = MCPFactory("http://e", MCPTransportType.HTTP_STREAM) with pytest.raises(RuntimeError): factory._fetch_prompt_definitions() def test_run_tool_handles_special_result_types(monkeypatch): import atomic_agents.connectors.mcp.mcp_factory as mtf class DummyTransportCM: def __init__(self, ret): self.ret = ret async def __aenter__(self): return self.ret async def __aexit__(self, exc_type, exc, tb): pass def dummy_sse_client(endpoint): return DummyTransportCM((None, None)) def dummy_stdio_client(params): return DummyTransportCM((None, None)) class DynamicSession: def __init__(self, *args, **kwargs): pass async def initialize(self): pass async def call_tool(self, name, arguments): class R(BaseModel): content: str return R(content="hello") async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): pass monkeypatch.setattr(mtf, "sse_client", dummy_sse_client) monkeypatch.setattr(mtf, "stdio_client", dummy_stdio_client) monkeypatch.setattr(mtf, "ClientSession", DynamicSession) definitions = [ MCPToolDefinition(name="T", description=None, input_schema={"type": "object", "properties": {}, "required": []}) ] monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: definitions) tool_cls = fetch_mcp_tools("e", MCPTransportType.HTTP_STREAM)[0] result = tool_cls().run(tool_cls.input_schema(tool_name="T")) assert result.result == "hello" # plain result class PlainSession(DynamicSession): async def call_tool(self, name, arguments): return 123 monkeypatch.setattr(mtf, "ClientSession", PlainSession) result2 = fetch_mcp_tools("e", MCPTransportType.HTTP_STREAM)[0]().run(tool_cls.input_schema(tool_name="T")) assert result2.result == 123 def test_run_resource_handles_special_result_types(monkeypatch): import atomic_agents.connectors.mcp.mcp_factory as mtf class DummyTransportCM: def __init__(self, ret): self.ret = ret async def __aenter__(self): return self.ret async def __aexit__(self, exc_type, exc, tb): pass def dummy_sse_client(endpoint): return DummyTransportCM((None, None)) def dummy_stdio_client(params): return DummyTransportCM((None, None)) class DynamicSession: def __init__(self, *args, **kwargs): pass async def initialize(self): pass async def read_resource(self, *, uri): class R(BaseModel): contents: str return R(contents="res-hello") async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): pass monkeypatch.setattr(mtf, "sse_client", dummy_sse_client) monkeypatch.setattr(mtf, "stdio_client", dummy_stdio_client) monkeypatch.setattr(mtf, "ClientSession", DynamicSession) definitions = [ MCPResourceDefinition( name="R", description=None, uri="resource://R", input_schema={"type": "object", "properties": {}, "required": []} ) ] monkeypatch.setattr(MCPFactory, "_fetch_resource_definitions", lambda self: definitions) resource_cls = fetch_mcp_resources("e", MCPTransportType.HTTP_STREAM)[0] result = resource_cls().read(resource_cls.input_schema(resource_name="R")) # resource output schema uses 'content' as the field name; the inner value # may itself be a BaseModel with attribute 'contents' (legacy) or 'content'. def _unwrap_output(out): val = getattr(out, "content", out) if isinstance(val, BaseModel): if hasattr(val, "content"): return val.content if hasattr(val, "contents"): return val.contents return val assert _unwrap_output(result) == "res-hello" # plain result class PlainSession(DynamicSession): async def read_resource(self, *, uri): return 456 monkeypatch.setattr(mtf, "ClientSession", PlainSession) result2 = fetch_mcp_resources("e", MCPTransportType.HTTP_STREAM)[0]().read(resource_cls.input_schema(resource_name="R")) assert _unwrap_output(result2) == 456 def test_run_prompt_handles_special_result_types(monkeypatch): import atomic_agents.connectors.mcp.mcp_factory as mtf class DummyTransportCM: def __init__(self, ret): self.ret = ret async def __aenter__(self): return self.ret async def __aexit__(self, exc_type, exc, tb): pass def dummy_sse_client(endpoint): return DummyTransportCM((None, None)) def dummy_stdio_client(params): return DummyTransportCM((None, None)) class DynamicSession: def __init__(self, *args, **kwargs): pass async def initialize(self): pass async def get_prompt(self, *, name, arguments): class Msg(BaseModel): content: str return {"messages": [Msg(content="prompt-hello")]} async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): pass monkeypatch.setattr(mtf, "sse_client", dummy_sse_client) monkeypatch.setattr(mtf, "stdio_client", dummy_stdio_client) monkeypatch.setattr(mtf, "ClientSession", DynamicSession) definitions = [ MCPPromptDefinition(name="P", description=None, input_schema={"type": "object", "properties": {}, "required": []}) ] monkeypatch.setattr(MCPFactory, "_fetch_prompt_definitions", lambda self: definitions) prompt_cls = fetch_mcp_prompts("e", MCPTransportType.HTTP_STREAM)[0] result = prompt_cls().generate(prompt_cls.input_schema(prompt_name="P")) assert result.content == "prompt-hello" # plain result class PlainSession(DynamicSession): async def get_prompt(self, *, name, arguments): return {"messages": ["plain-hello"]} monkeypatch.setattr(mtf, "ClientSession", PlainSession) result2 = fetch_mcp_prompts("e", MCPTransportType.HTTP_STREAM)[0]().generate(prompt_cls.input_schema(prompt_name="P")) assert result2.content == "plain-hello" def test_run_invalid_stdio_command_raises(monkeypatch): import atomic_agents.connectors.mcp.mcp_factory as mtf class DummyTransportCM: def __init__(self, ret): self.ret = ret async def __aenter__(self): return self.ret async def __aexit__(self, exc_type, exc, tb): pass def dummy_sse_client(endpoint): return DummyTransportCM((None, None)) def dummy_stdio_client(params): return DummyTransportCM((None, None)) monkeypatch.setattr(mtf, "sse_client", dummy_sse_client) monkeypatch.setattr(mtf, "stdio_client", dummy_stdio_client) monkeypatch.setattr( MCPFactory, "_fetch_tool_definitions", lambda self: [ MCPToolDefinition(name="Bad", description=None, input_schema={"type": "object", "properties": {}, "required": []}) ], ) monkeypatch.setattr( MCPFactory, "_fetch_resource_definitions", lambda self: [ MCPResourceDefinition( name="Y", description="d", uri="resource://Y", input_schema={"type": "object", "properties": {}, "required": []}, ) ], ) monkeypatch.setattr( MCPFactory, "_fetch_prompt_definitions", lambda self: [ MCPPromptDefinition(name="Z", description="d", input_schema={"type": "object", "properties": {}, "required": []}) ], ) # Use a blank-space endpoint to bypass init validation but trigger empty command in STDIO tool_cls = fetch_mcp_tools(" ", MCPTransportType.STDIO, working_directory="/wd")[0] with pytest.raises(RuntimeError) as exc: tool_cls().run(tool_cls.input_schema(tool_name="Bad")) assert "STDIO command string cannot be empty" in str(exc.value) resource_cls = fetch_mcp_resources(" ", MCPTransportType.STDIO, working_directory="/wd")[0] with pytest.raises(RuntimeError) as exc: resource_cls().read(resource_cls.input_schema(resource_name="Y")) assert "STDIO command string cannot be empty" in str(exc.value) prompt_cls = fetch_mcp_prompts(" ", MCPTransportType.STDIO, working_directory="/wd")[0] with pytest.raises(RuntimeError) as exc: prompt_cls().generate(prompt_cls.input_schema(prompt_name="Z")) assert "STDIO command string cannot be empty" in str(exc.value) def test_create_tool_classes_skips_invalid(monkeypatch): factory = MCPFactory("endpoint", MCPTransportType.HTTP_STREAM) defs = [ MCPToolDefinition(name="Bad", description=None, input_schema={"type": "object", "properties": {}, "required": []}), MCPToolDefinition(name="Good", description=None, input_schema={"type": "object", "properties": {}, "required": []}), ] class FakeST: def create_model_from_schema(self, schema, model_name, tname, doc, attribute_type="tool"): if tname == "Bad": raise ValueError("fail") return BaseModel factory.schema_transformer = FakeST() tools = factory._create_tool_classes(defs) assert len(tools) == 1 assert tools[0].mcp_tool_name == "Good" def test_create_resource_classes_skips_invalid(monkeypatch): factory = MCPFactory("endpoint", MCPTransportType.HTTP_STREAM) defs = [ MCPResourceDefinition( name="Bad", description=None, uri="resource://Bad", input_schema={"type": "object", "properties": {}, "required": []}, ), MCPResourceDefinition( name="Good", description=None, uri="resource://Good", input_schema={"type": "object", "properties": {}, "required": []}, ), ] class FakeST: def create_model_from_schema(self, schema, model_name, tname, doc, attribute_type="resource"): if tname == "Bad": raise ValueError("fail") return BaseModel factory.schema_transformer = FakeST() resources = factory._create_resource_classes(defs) assert len(resources) == 1 assert resources[0].mcp_resource_name == "Good" def test_create_prompt_classes_skips_invalid(monkeypatch): factory = MCPFactory("endpoint", MCPTransportType.HTTP_STREAM) defs = [ MCPPromptDefinition(name="Bad", description=None, input_schema={"type": "object", "properties": {}, "required": []}), MCPPromptDefinition(name="Good", description=None, input_schema={"type": "object", "properties": {}, "required": []}), ] class FakeST: def create_model_from_schema(self, schema, model_name, tname, doc, attribute_type="prompt"): if tname == "Bad": raise ValueError("fail") return BaseModel factory.schema_transformer = FakeST() prompts = factory._create_prompt_classes(defs) assert len(prompts) == 1 assert prompts[0].mcp_prompt_name == "Good" @pytest.fixture def mocked_mcp_connection(monkeypatch: pytest.MonkeyPatch) -> tuple[AsyncMock, dict[str, MagicMock]]: import atomic_agents.connectors.mcp.mcp_factory as factory_module import atomic_agents.connectors.mcp.mcp_definition_service as definition_module session = AsyncMock(spec=ClientSession) session.__aenter__.return_value = session session.__aexit__.return_value = None session.list_tools.return_value = ListToolsResult( tools=[ Tool( name="EchoTool", inputSchema={ "type": "object", "properties": {"query": {"type": "string"}, "optional": {"type": "string"}}, "required": ["query"], }, ) ] ) session.call_tool.return_value = CallToolResult(content=[TextContent(type="text", text="echo result")]) monkeypatch.setattr(factory_module, "ClientSession", MagicMock(return_value=session)) monkeypatch.setattr(definition_module, "ClientSession", MagicMock(return_value=session)) clients = {} for client_name, streams in ( ("stdio_client", (None, None)), ("sse_client", (None, None)), ("streamablehttp_client", (None, None, None)), ): transport = AsyncMock() transport.__aenter__.return_value = streams transport.__aexit__.return_value = None client = MagicMock(return_value=transport) clients[client_name] = client monkeypatch.setattr(factory_module, client_name, client) monkeypatch.setattr(definition_module, client_name, client) return session, clients @pytest.mark.parametrize( "transport_type,client_name,endpoint,expected_connection", [ pytest.param( MCPTransportType.STDIO, "stdio_client", "mock-mcp --local", StdioServerParameters(command="mock-mcp", args=["--local"], env=None, cwd="workdir"), id="stdio", ), pytest.param(MCPTransportType.SSE, "sse_client", "http://example.test", "http://example.test/sse", id="sse"), pytest.param( MCPTransportType.HTTP_STREAM, "streamablehttp_client", "http://example.test", "http://example.test/mcp/", id="http-stream", ), ], ) def test_generated_tool_discovers_and_calls_sdk_tool_through_transport( mocked_mcp_connection: tuple[AsyncMock, dict[str, MagicMock]], transport_type: MCPTransportType, client_name: str, endpoint: str, expected_connection: str | StdioServerParameters, ) -> None: session, clients = mocked_mcp_connection tool_class = fetch_mcp_tools(endpoint, transport_type, working_directory="workdir")[0] result = tool_class().run(tool_class.input_schema(tool_name="EchoTool", query="hello", optional=None)) assert result.result == [TextContent(type="text", text="echo result")] session.list_tools.assert_awaited_once_with() session.call_tool.assert_awaited_once_with(name="EchoTool", arguments={"query": "hello"}) assert session.initialize.await_count == 2 assert clients[client_name].call_args_list == [call(expected_connection), call(expected_connection)] session.__aexit__.assert_has_awaits([call(session, None, None, None), call(session, None, None, None)]) transport = clients[client_name].return_value transport.__aexit__.assert_has_awaits([call(transport, None, None, None), call(transport, None, None, None)]) @pytest.mark.asyncio @pytest.mark.parametrize( "tool_result,expected_count", [ pytest.param(CallToolResult(content=[TextContent(type="text", text='{"count": 7}')]), 7, id="json-text-fallback"), pytest.param( CallToolResult(content=[TextContent(type="text", text='{"count": 99}')], structuredContent={"count": 7}), 7, id="structured-precedence", ), pytest.param( CallToolResult(content=[TextContent(type="text", text='{"count": 99}')], structuredContent={"count": 0}), 0, id="structured-zero", ), ], ) async def test_generated_tool_uses_sdk_typed_output_on_persistent_session( mocked_mcp_connection: tuple[AsyncMock, dict[str, MagicMock]], tool_result: CallToolResult, expected_count: int ) -> None: session, clients = mocked_mcp_connection session.list_tools.return_value.tools[0].outputSchema = { "type": "object", "properties": {"count": {"type": "integer"}}, "required": ["count"], } session.call_tool.return_value = tool_result tool_class = (await fetch_mcp_tools_async(client_session=session))[0] result = await tool_class().arun(tool_class.input_schema(tool_name="EchoTool", query="hello")) assert result.model_dump() == {"count": expected_count} session.call_tool.assert_awaited_once_with(name="EchoTool", arguments={"query": "hello"}) session.initialize.assert_awaited_once_with() for client in clients.values(): client.assert_not_called() @pytest.mark.asyncio @pytest.mark.parametrize( "tool_result,cause_type,message", [ pytest.param( CallToolResult(content=[TextContent(type="text", text="invalid")]), ValueError, "unparseable result", id="malformed-json", ), pytest.param( CallToolResult(content=[TextContent(type="text", text="[]")]), ValueError, "unparseable result", id="json-array" ), pytest.param( CallToolResult(content=[TextContent(type="text", text="0")]), ValueError, "unparseable result", id="json-zero" ), pytest.param( CallToolResult(content=[TextContent(type="text", text="false")]), ValueError, "unparseable result", id="json-false" ), pytest.param( CallToolResult(content=[TextContent(type="text", text="null")]), ValueError, "unparseable result", id="json-null" ), pytest.param( CallToolResult(content=[TextContent(type="text", text='{"count": "invalid"}')]), ValidationError, "validation error", id="invalid-count", ), pytest.param( CallToolResult(content=[TextContent(type="text", text='{"count": 7}')], structuredContent={}), ValidationError, "validation error", id="empty-structured-precedence", ), ], ) async def test_generated_typed_tool_rejects_invalid_sdk_output( mocked_mcp_connection: tuple[AsyncMock, dict[str, MagicMock]], tool_result: CallToolResult, cause_type: type[Exception], message: str, ) -> None: session, _ = mocked_mcp_connection session.list_tools.return_value.tools[0].outputSchema = { "type": "object", "properties": {"count": {"type": "integer"}}, "required": ["count"], } session.call_tool.return_value = tool_result tool_class = (await fetch_mcp_tools_async(client_session=session))[0] with pytest.raises(RuntimeError, match=message) as failure: await tool_class().arun(tool_class.input_schema(tool_name="EchoTool", query="hello")) assert isinstance(failure.value.__cause__, cause_type) session.call_tool.assert_awaited_once_with(name="EchoTool", arguments={"query": "hello"}) def test__fetch_tool_definitions_service_branch(monkeypatch): """Covers lines 112-113: MCPDefinitionService branch in _fetch_tool_definitions.""" factory = MCPFactory("dummy_endpoint", MCPTransportType.HTTP_STREAM) # Patch fetch_tool_definitions to avoid real async work async def dummy_fetch_tool_definitions(self): return [ MCPToolDefinition(name="COV", description="cov", input_schema={"type": "object", "properties": {}, "required": []}) ] monkeypatch.setattr(MCPDefinitionService, "fetch_tool_definitions", dummy_fetch_tool_definitions) result = factory._fetch_tool_definitions() assert result[0].name == "COV" def test_fetch_resource_definitions_service_branch(monkeypatch): """Covers lines of MCPDefinitionService branch in _fetch_resource_definitions.""" factory = MCPFactory("dummy_endpoint", MCPTransportType.HTTP_STREAM) # Patch fetch_resource_definitions to avoid real async work async def dummy_fetch_resource_definitions(self): return [ MCPResourceDefinition( name="COVR", description="covr", uri="resource://COVR", input_schema={"type": "object", "properties": {}, "required": []}, ) ] monkeypatch.setattr(MCPDefinitionService, "fetch_resource_definitions", dummy_fetch_resource_definitions) result = factory._fetch_resource_definitions() assert result[0].name == "COVR" def test_fetch_prompt_definitions_service_branch(monkeypatch): """Covers lines of MCPDefinitionService branch in _fetch_prompt_definitions.""" factory = MCPFactory("dummy_endpoint", MCPTransportType.HTTP_STREAM) # Patch fetch_prompt_definitions to avoid real async work async def dummy_fetch_prompt_definitions(self): return [ MCPPromptDefinition( name="COVP", description="covp", input_schema={"type": "object", "properties": {}, "required": []} ) ] monkeypatch.setattr(MCPDefinitionService, "fetch_prompt_definitions", dummy_fetch_prompt_definitions) result = factory._fetch_prompt_definitions() assert result[0].name == "COVP" def test_fetch_mcp_tools_rejects_persistent_session_without_event_loop( mocked_mcp_connection: tuple[AsyncMock, dict[str, MagicMock]], ) -> None: session, clients = mocked_mcp_connection with pytest.raises(ValueError, match="When `client_session` is provided an `event_loop` must also be supplied"): fetch_mcp_tools(client_session=session) session.initialize.assert_not_awaited() session.list_tools.assert_not_awaited() session.call_tool.assert_not_awaited() for client in clients.values(): client.assert_not_called() def test_run_tool_with_persistent_session_no_event_loop(monkeypatch): """Covers AttributeError when no event loop is provided for persistent session.""" import atomic_agents.connectors.mcp.mcp_factory as mtf # Setup persistent client class DummySessionPersistent: async def call_tool(self, name, arguments): return {"content": "should not get here"} client = DummySessionPersistent() definitions = [ MCPToolDefinition(name="ToolCOV", description=None, input_schema={"type": "object", "properties": {}, "required": []}) ] async def fake_fetch_defs(session): return definitions monkeypatch.setattr(mtf.MCPDefinitionService, "fetch_tool_definitions_from_session", staticmethod(fake_fetch_defs)) # Create tool with persistent session and a valid event loop loop = asyncio.new_event_loop() try: tools = fetch_mcp_tools(None, MCPTransportType.HTTP_STREAM, client_session=client, event_loop=loop) tool_cls = tools[0] inst = tool_cls() # Remove the event loop to simulate the error path inst._event_loop = None with pytest.raises(RuntimeError) as exc: inst.run(tool_cls.input_schema(tool_name="ToolCOV")) # The error originates as AttributeError but is wrapped in RuntimeError assert "'NoneType' object has no attribute 'run_until_complete'" in str(exc.value) finally: loop.close() def test_run_resource_with_persistent_session_no_event_loop(monkeypatch): """Covers AttributeError when no event loop is provided for persistent session.""" import atomic_agents.connectors.mcp.mcp_factory as mtf # Setup persistent client class DummySessionPersistent: async def read_resource(self, *, uri): return {"content": "should not get here"} client = DummySessionPersistent() definitions = [ MCPResourceDefinition( name="ResCOV", description=None, uri="resource://ResCOV", input_schema={"type": "object", "properties": {}, "required": []}, ) ] async def fake_fetch_defs(session): return definitions monkeypatch.setattr(mtf.MCPDefinitionService, "fetch_resource_definitions_from_session", staticmethod(fake_fetch_defs)) # Create resource with persistent session and a valid event loop loop = asyncio.new_event_loop() try: resources = fetch_mcp_resources(None, MCPTransportType.HTTP_STREAM, client_session=client, event_loop=loop) res_cls = resources[0] inst = res_cls() # Remove the event loop to simulate the error inst._event_loop = None with pytest.raises(RuntimeError) as exc: inst.read(res_cls.input_schema(resource_name="ResCOV")) # The error originates as AttributeError but is wrapped in RuntimeError assert "'NoneType' object has no attribute 'run_until_complete'" in str(exc.value) finally: loop.close() def test_run_prompt_with_persistent_session_no_event_loop(monkeypatch): """Covers AttributeError when no event loop is provided for persistent session.""" import atomic_agents.connectors.mcp.mcp_factory as mtf # Setup persistent client class DummySessionPersistent: async def get_prompt(self, *, name, arguments): return {"content": "should not get here"} client = DummySessionPersistent() definitions = [ MCPPromptDefinition( name="PromptCOV", description=None, input_schema={"type": "object", "properties": {}, "required": []} ) ] async def fake_fetch_defs(session): return definitions monkeypatch.setattr(mtf.MCPDefinitionService, "fetch_prompt_definitions_from_session", staticmethod(fake_fetch_defs)) # Create prompt with persistent session and a valid event loop loop = asyncio.new_event_loop() try: prompts = fetch_mcp_prompts(None, MCPTransportType.HTTP_STREAM, client_session=client, event_loop=loop) prompt_cls = prompts[0] inst = prompt_cls() # Remove the event loop to simulate the error inst._event_loop = None with pytest.raises(RuntimeError) as exc: inst.generate(prompt_cls.input_schema(prompt_name="PromptCOV")) # The error originates as AttributeError but is wrapped in RuntimeError assert "'NoneType' object has no attribute 'run_until_complete'" in str(exc.value) finally: loop.close() def test_http_stream_connection_error_handling(monkeypatch): """Test HTTP stream connection error handling in MCPToolFactory.""" from atomic_agents.connectors.mcp.mcp_definition_service import MCPDefinitionService # Mock MCPDefinitionService.fetch_tool_definitions to raise ConnectionError for HTTP_STREAM original_fetch_tools = MCPDefinitionService.fetch_tool_definitions async def mock_fetch_tool_definitions(self): if self.transport_type == MCPTransportType.HTTP_STREAM: raise ConnectionError("HTTP stream connection failed") return await original_fetch_tools(self) monkeypatch.setattr(MCPDefinitionService, "fetch_tool_definitions", mock_fetch_tool_definitions) factory = MCPFactory("http://test-endpoint", MCPTransportType.HTTP_STREAM) with pytest.raises(ConnectionError, match="HTTP stream connection failed"): factory._fetch_tool_definitions() original_fetch_resources = MCPDefinitionService.fetch_resource_definitions async def mock_fetch_resource_definitions(self): if self.transport_type == MCPTransportType.HTTP_STREAM: raise ConnectionError("HTTP stream connection failed") return await original_fetch_resources(self) monkeypatch.setattr(MCPDefinitionService, "fetch_resource_definitions", mock_fetch_resource_definitions) with pytest.raises(ConnectionError, match="HTTP stream connection failed"): factory._fetch_resource_definitions() original_fetch_prompts = MCPDefinitionService.fetch_prompt_definitions async def mock_fetch_prompt_definitions(self): if self.transport_type == MCPTransportType.HTTP_STREAM: raise ConnectionError("HTTP stream connection failed") return await original_fetch_prompts(self) monkeypatch.setattr(MCPDefinitionService, "fetch_prompt_definitions", mock_fetch_prompt_definitions) with pytest.raises(ConnectionError, match="HTTP stream connection failed"): factory._fetch_prompt_definitions() def test_http_stream_endpoint_formatting(): """Test that HTTP stream endpoints are properly formatted with /mcp/ suffix.""" factory = MCPFactory("http://test-endpoint", MCPTransportType.HTTP_STREAM) # Verify the factory was created with correct transport type assert factory.transport_type == MCPTransportType.HTTP_STREAM # Tests for fetch_mcp_tools_async function @pytest.mark.asyncio async def test_fetch_mcp_tools_async_with_client_session(monkeypatch): """Test fetch_mcp_tools_async with pre-initialized client session.""" import atomic_agents.connectors.mcp.mcp_factory as mtf # Setup persistent client class DummySessionPersistent: async def call_tool(self, name, arguments): return {"content": "async-session-ok"} client = DummySessionPersistent() definitions = [ MCPToolDefinition( name="AsyncTool", description="Test async tool", input_schema={"type": "object", "properties": {}, "required": []} ) ] async def fake_fetch_defs(session): return definitions monkeypatch.setattr(mtf.MCPDefinitionService, "fetch_tool_definitions_from_session", staticmethod(fake_fetch_defs)) # Call fetch_mcp_tools_async with client session tools = await fetch_mcp_tools_async(None, MCPTransportType.HTTP_STREAM, client_session=client) assert len(tools) == 1 tool_cls = tools[0] # Verify the tool was created correctly assert hasattr(tool_cls, "mcp_tool_name") @pytest.mark.asyncio async def test_fetch_mcp_resources_async_with_client_session(monkeypatch): """Test fetch_mcp_resources_async with pre-initialized client session.""" import atomic_agents.connectors.mcp.mcp_factory as mtf # Setup persistent client class DummySessionPersistent: async def read_resource(self, name, uri): return {"content": "async-resource-ok"} client = DummySessionPersistent() definitions = [ MCPResourceDefinition( name="AsyncRes", description="Test async resource", uri="resource://AsyncRes", input_schema={"type": "object", "properties": {}, "required": []}, ) ] async def fake_fetch_defs(session): return definitions monkeypatch.setattr(mtf.MCPDefinitionService, "fetch_resource_definitions_from_session", staticmethod(fake_fetch_defs)) # Call fetch_mcp_resources_async with client session resources = await fetch_mcp_resources_async(None, MCPTransportType.HTTP_STREAM, client_session=client) assert len(resources) == 1 res_cls = resources[0] # Verify the resource was created correctly assert hasattr(res_cls, "mcp_resource_name") @pytest.mark.asyncio async def test_fetch_mcp_prompts_async_with_client_session(monkeypatch): """Test fetch_mcp_prompts_async with pre-initialized client session.""" import atomic_agents.connectors.mcp.mcp_factory as mtf # Setup persistent client class DummySessionPersistent: async def generate_prompt(self, name, arguments): return {"content": "async-prompt-ok"} client = DummySessionPersistent() definitions = [ MCPPromptDefinition( name="AsyncPrompt", description="Test async prompt", input_schema={"type": "object", "properties": {}, "required": []}, ) ] async def fake_fetch_defs(session): return definitions monkeypatch.setattr(mtf.MCPDefinitionService, "fetch_prompt_definitions_from_session", staticmethod(fake_fetch_defs)) # Call fetch_mcp_prompts_async with client session prompts = await fetch_mcp_prompts_async(None, MCPTransportType.HTTP_STREAM, client_session=client) assert len(prompts) == 1 prompt_cls = prompts[0] # Verify the prompt was created correctly assert hasattr(prompt_cls, "mcp_prompt_name") @pytest.mark.asyncio async def test_fetch_mcp_tools_async_without_client_session(monkeypatch): """Test fetch_mcp_tools_async without pre-initialized client session.""" definitions = [ MCPToolDefinition( name="AsyncTool2", description="Test async tool 2", input_schema={"type": "object", "properties": {}, "required": []}, ) ] async def fake_fetch_defs(self): return definitions monkeypatch.setattr(MCPDefinitionService, "fetch_tool_definitions", fake_fetch_defs) # Call fetch_mcp_tools_async without client session tools = await fetch_mcp_tools_async("http://test-endpoint", MCPTransportType.HTTP_STREAM) assert len(tools) == 1 tool_cls = tools[0] # Verify the tool was created correctly assert hasattr(tool_cls, "mcp_tool_name") @pytest.mark.asyncio async def test_fetch_mcp_resources_async_without_client_session(monkeypatch): """Test fetch_mcp_resources_async without pre-initialized client session.""" definitions = [ MCPResourceDefinition( name="AsyncRes2", description="Test async resource 2", uri="resource://AsyncRes2", input_schema={"type": "object", "properties": {}, "required": []}, ) ] async def fake_fetch_defs(self): return definitions monkeypatch.setattr(MCPDefinitionService, "fetch_resource_definitions", fake_fetch_defs) # Call fetch_mcp_resources_async without client session resources = await fetch_mcp_resources_async("http://test-endpoint", MCPTransportType.HTTP_STREAM) assert len(resources) == 1 res_cls = resources[0] # Verify the resource was created correctly assert hasattr(res_cls, "mcp_resource_name") @pytest.mark.asyncio async def test_fetch_mcp_prompts_async_without_client_session(monkeypatch): """Test fetch_mcp_prompts_async without pre-initialized client session.""" definitions = [ MCPPromptDefinition( name="AsyncPrompt2", description="Test async prompt 2", input_schema={"type": "object", "properties": {}, "required": []}, ) ] async def fake_fetch_defs(self): return definitions monkeypatch.setattr(MCPDefinitionService, "fetch_prompt_definitions", fake_fetch_defs) # Call fetch_mcp_prompts_async without client session prompts = await fetch_mcp_prompts_async("http://test-endpoint", MCPTransportType.HTTP_STREAM) assert len(prompts) == 1 prompt_cls = prompts[0] # Verify the prompt was created correctly assert hasattr(prompt_cls, "mcp_prompt_name") @pytest.mark.asyncio async def test_fetch_mcp_tools_async_stdio_transport(monkeypatch): """Test fetch_mcp_tools_async with STDIO transport.""" definitions = [ MCPToolDefinition( name="StdioAsyncTool", description="Test stdio async tool", input_schema={"type": "object", "properties": {}, "required": []}, ) ] async def fake_fetch_defs(self): return definitions monkeypatch.setattr(MCPDefinitionService, "fetch_tool_definitions", fake_fetch_defs) # Call fetch_mcp_tools_async with STDIO transport tools = await fetch_mcp_tools_async("test-command", MCPTransportType.STDIO, working_directory="/tmp") assert len(tools) == 1 tool_cls = tools[0] # Verify the tool was created correctly assert hasattr(tool_cls, "mcp_tool_name") @pytest.mark.asyncio async def test_fetch_mcp_resources_async_stdio_transport(monkeypatch): """Test fetch_mcp_resources_async with STDIO transport.""" definitions = [ MCPResourceDefinition( name="StdioAsyncRes", description="Test stdio async resource", uri="resource://StdioAsyncRes", input_schema={"type": "object", "properties": {}, "required": []}, ) ] async def fake_fetch_defs(self): return definitions monkeypatch.setattr(MCPDefinitionService, "fetch_resource_definitions", fake_fetch_defs) # Call fetch_mcp_resources_async with STDIO transport resources = await fetch_mcp_resources_async("test-command", MCPTransportType.STDIO, working_directory="/tmp") assert len(resources) == 1 res_cls = resources[0] # Verify the resource was created correctly assert hasattr(res_cls, "mcp_resource_name") @pytest.mark.asyncio async def test_fetch_mcp_prompts_async_stdio_transport(monkeypatch): """Test fetch_mcp_prompts_async with STDIO transport.""" definitions = [ MCPPromptDefinition( name="StdioAsyncPrompt", description="Test stdio async prompt", input_schema={"type": "object", "properties": {}, "required": []}, ) ] async def fake_fetch_defs(self): return definitions monkeypatch.setattr(MCPDefinitionService, "fetch_prompt_definitions", fake_fetch_defs) # Call fetch_mcp_prompts_async with STDIO transport prompts = await fetch_mcp_prompts_async("test-command", MCPTransportType.STDIO, working_directory="/tmp") assert len(prompts) == 1 prompt_cls = prompts[0] # Verify the prompt was created correctly assert hasattr(prompt_cls, "mcp_prompt_name") @pytest.mark.asyncio async def test_fetch_mcp_tools_async_empty_definitions(monkeypatch): """Test fetch_mcp_tools_async returns empty list when no definitions found.""" async def fake_fetch_defs(self): return [] monkeypatch.setattr(MCPDefinitionService, "fetch_tool_definitions", fake_fetch_defs) # Call fetch_mcp_tools_async tools = await fetch_mcp_tools_async("http://test-endpoint", MCPTransportType.HTTP_STREAM) assert tools == [] @pytest.mark.asyncio async def test_fetch_mcp_resources_async_empty_definitions(monkeypatch): """Test fetch_mcp_resources_async returns empty list when no definitions found.""" async def fake_fetch_defs(self): return [] monkeypatch.setattr(MCPDefinitionService, "fetch_resource_definitions", fake_fetch_defs) # Call fetch_mcp_resources_async resources = await fetch_mcp_resources_async("http://test-endpoint", MCPTransportType.HTTP_STREAM) assert resources == [] @pytest.mark.asyncio async def test_fetch_mcp_prompts_async_empty_definitions(monkeypatch): """Test fetch_mcp_prompts_async returns empty list when no definitions found.""" async def fake_fetch_defs(self): return [] monkeypatch.setattr(MCPDefinitionService, "fetch_prompt_definitions", fake_fetch_defs) # Call fetch_mcp_prompts_async prompts = await fetch_mcp_prompts_async("http://test-endpoint", MCPTransportType.HTTP_STREAM) assert prompts == [] @pytest.mark.asyncio async def test_fetch_mcp_tools_async_connection_error(monkeypatch): """Test fetch_mcp_tools_async propagates connection errors.""" async def fake_fetch_defs_error(self): raise ConnectionError("Failed to connect to MCP server") monkeypatch.setattr(MCPDefinitionService, "fetch_tool_definitions", fake_fetch_defs_error) # Call fetch_mcp_tools_async and expect ConnectionError with pytest.raises(ConnectionError, match="Failed to connect to MCP server"): await fetch_mcp_tools_async("http://test-endpoint", MCPTransportType.HTTP_STREAM) @pytest.mark.asyncio async def test_fetch_mcp_resources_async_connection_error(monkeypatch): """Test fetch_mcp_resources_async propagates connection errors.""" async def fake_fetch_defs_error(self): raise ConnectionError("Failed to connect to MCP server") monkeypatch.setattr(MCPDefinitionService, "fetch_resource_definitions", fake_fetch_defs_error) # Call fetch_mcp_resources_async and expect ConnectionError with pytest.raises(ConnectionError, match="Failed to connect to MCP server"): await fetch_mcp_resources_async("http://test-endpoint", MCPTransportType.HTTP_STREAM) @pytest.mark.asyncio async def test_fetch_mcp_prompts_async_connection_error(monkeypatch): """Test fetch_mcp_prompts_async propagates connection errors.""" async def fake_fetch_defs_error(self): raise ConnectionError("Failed to connect to MCP server") monkeypatch.setattr(MCPDefinitionService, "fetch_prompt_definitions", fake_fetch_defs_error) # Call fetch_mcp_prompts_async and expect ConnectionError with pytest.raises(ConnectionError, match="Failed to connect to MCP server"): await fetch_mcp_prompts_async("http://test-endpoint", MCPTransportType.HTTP_STREAM) @pytest.mark.asyncio async def test_fetch_mcp_tools_async_runtime_error(monkeypatch): """Test fetch_mcp_tools_async propagates runtime errors.""" async def fake_fetch_defs_error(self): raise RuntimeError("Unexpected error during fetching") monkeypatch.setattr(MCPDefinitionService, "fetch_tool_definitions", fake_fetch_defs_error) # Call fetch_mcp_tools_async and expect RuntimeError with pytest.raises(RuntimeError, match="Unexpected error during fetching"): await fetch_mcp_tools_async("http://test-endpoint", MCPTransportType.HTTP_STREAM) @pytest.mark.asyncio async def test_fetch_mcp_resources_async_runtime_error(monkeypatch): """Test fetch_mcp_resources_async propagates runtime errors.""" async def fake_fetch_defs_error(self): raise RuntimeError("Unexpected error during fetching") monkeypatch.setattr(MCPDefinitionService, "fetch_resource_definitions", fake_fetch_defs_error) # Call fetch_mcp_resources_async and expect RuntimeError with pytest.raises(RuntimeError, match="Unexpected error during fetching"): await fetch_mcp_resources_async("http://test-endpoint", MCPTransportType.HTTP_STREAM) @pytest.mark.asyncio async def test_fetch_mcp_prompts_async_runtime_error(monkeypatch): """Test fetch_mcp_prompts_async propagates runtime errors.""" async def fake_fetch_defs_error(self): raise RuntimeError("Unexpected error during fetching") monkeypatch.setattr(MCPDefinitionService, "fetch_prompt_definitions", fake_fetch_defs_error) # Call fetch_mcp_prompts_async and expect RuntimeError with pytest.raises(RuntimeError, match="Unexpected error during fetching"): await fetch_mcp_prompts_async("http://test-endpoint", MCPTransportType.HTTP_STREAM) @pytest.mark.asyncio async def test_fetch_mcp_tools_async_with_working_directory(monkeypatch): """Test fetch_mcp_tools_async with working directory parameter.""" definitions = [ MCPToolDefinition( name="WorkingDirTool", description="Test tool with working dir", input_schema={"type": "object", "properties": {}, "required": []}, ) ] async def fake_fetch_defs(self): return definitions monkeypatch.setattr(MCPDefinitionService, "fetch_tool_definitions", fake_fetch_defs) # Call fetch_mcp_tools_async with working directory tools = await fetch_mcp_tools_async("test-command", MCPTransportType.STDIO, working_directory="/custom/working/dir") assert len(tools) == 1 tool_cls = tools[0] # Verify the tool was created correctly assert hasattr(tool_cls, "mcp_tool_name") @pytest.mark.asyncio async def test_fetch_mcp_resources_async_with_working_directory(monkeypatch): """Test fetch_mcp_resources_async with working directory parameter.""" definitions = [ MCPResourceDefinition( name="WorkingDirRes", description="Test resource with working dir", uri="resource://WorkingDirRes", input_schema={"type": "object", "properties": {}, "required": []}, ) ] async def fake_fetch_defs(self): return definitions monkeypatch.setattr(MCPDefinitionService, "fetch_resource_definitions", fake_fetch_defs) # Call fetch_mcp_resources_async with working directory resources = await fetch_mcp_resources_async( "test-command", MCPTransportType.STDIO, working_directory="/custom/working/dir" ) assert len(resources) == 1 res_cls = resources[0] # Verify the resource was created correctly assert hasattr(res_cls, "mcp_resource_name") @pytest.mark.asyncio async def test_fetch_mcp_prompts_async_with_working_directory(monkeypatch): """Test fetch_mcp_prompts_async with working directory parameter.""" definitions = [ MCPPromptDefinition( name="WorkingDirPrompt", description="Test prompt with working dir", input_schema={"type": "object", "properties": {}, "required": []}, ) ] async def fake_fetch_defs(self): return definitions monkeypatch.setattr(MCPDefinitionService, "fetch_prompt_definitions", fake_fetch_defs) # Call fetch_mcp_prompts_async with working directory prompts = await fetch_mcp_prompts_async("test-command", MCPTransportType.STDIO, working_directory="/custom/working/dir") assert len(prompts) == 1 prompt_cls = prompts[0] # Verify the prompt was created correctly assert hasattr(prompt_cls, "mcp_prompt_name") @pytest.mark.asyncio async def test_fetch_mcp_tools_async_session_error_propagation(monkeypatch): """Test fetch_mcp_tools_async with client session error propagation.""" import atomic_agents.connectors.mcp.mcp_factory as mtf class DummySessionPersistent: async def call_tool(self, name, arguments): return {"content": "session-ok"} client = DummySessionPersistent() async def fake_fetch_defs_error(session): raise ValueError("Session fetch error") monkeypatch.setattr(mtf.MCPDefinitionService, "fetch_tool_definitions_from_session", staticmethod(fake_fetch_defs_error)) # Call fetch_mcp_tools_async with client session and expect error with pytest.raises(ValueError, match="Session fetch error"): await fetch_mcp_tools_async(None, MCPTransportType.HTTP_STREAM, client_session=client) @pytest.mark.asyncio async def test_fetch_mcp_resources_async_session_error_propagation(monkeypatch): """Test fetch_mcp_resources_async with client session error propagation.""" import atomic_agents.connectors.mcp.mcp_factory as mtf class DummySessionPersistent: async def read_resource(self, name, uri): return {"content": "session-ok"} client = DummySessionPersistent() async def fake_fetch_defs_error(session): raise ValueError("Session fetch error") monkeypatch.setattr( mtf.MCPDefinitionService, "fetch_resource_definitions_from_session", staticmethod(fake_fetch_defs_error) ) # Call fetch_mcp_resources_async with client session and expect error with pytest.raises(ValueError, match="Session fetch error"): await fetch_mcp_resources_async(None, MCPTransportType.HTTP_STREAM, client_session=client) @pytest.mark.asyncio async def test_fetch_mcp_prompts_async_session_error_propagation(monkeypatch): """Test fetch_mcp_prompts_async with client session error propagation.""" import atomic_agents.connectors.mcp.mcp_factory as mtf class DummySessionPersistent: async def generate_prompt(self, name, arguments): return {"content": "session-ok"} client = DummySessionPersistent() async def fake_fetch_defs_error(session): raise ValueError("Session fetch error") monkeypatch.setattr(mtf.MCPDefinitionService, "fetch_prompt_definitions_from_session", staticmethod(fake_fetch_defs_error)) # Call fetch_mcp_prompts_async with client session and expect error with pytest.raises(ValueError, match="Session fetch error"): await fetch_mcp_prompts_async(None, MCPTransportType.HTTP_STREAM, client_session=client) @pytest.mark.asyncio @pytest.mark.parametrize("transport_type", [MCPTransportType.HTTP_STREAM, MCPTransportType.STDIO, MCPTransportType.SSE]) async def test_fetch_mcp_tools_async_all_transport_types(monkeypatch, transport_type): """Test fetch_mcp_tools_async with all supported transport types.""" definitions = [ MCPToolDefinition( name=f"Tool_{transport_type.value}", description=f"Test tool for {transport_type.value}", input_schema={"type": "object", "properties": {}, "required": []}, ) ] async def fake_fetch_defs(self): return definitions monkeypatch.setattr(MCPDefinitionService, "fetch_tool_definitions", fake_fetch_defs) # Determine endpoint based on transport type endpoint = "test-command" if transport_type == MCPTransportType.STDIO else "http://test-endpoint" working_dir = "/tmp" if transport_type == MCPTransportType.STDIO else None # Call fetch_mcp_tools_async with different transport types tools = await fetch_mcp_tools_async(endpoint, transport_type, working_directory=working_dir) assert len(tools) == 1 tool_cls = tools[0] # Verify the tool was created correctly assert hasattr(tool_cls, "mcp_tool_name") @pytest.mark.asyncio @pytest.mark.parametrize("transport_type", [MCPTransportType.HTTP_STREAM, MCPTransportType.STDIO, MCPTransportType.SSE]) async def test_fetch_mcp_resources_async_all_transport_types(monkeypatch, transport_type): """Test fetch_mcp_resources_async with all supported transport types.""" definitions = [ MCPResourceDefinition( name=f"Res_{transport_type.value}", description=f"Test resource for {transport_type.value}", uri=f"resource://Res_{transport_type.value}", input_schema={"type": "object", "properties": {}, "required": []}, ) ] async def fake_fetch_defs(self): return definitions monkeypatch.setattr(MCPDefinitionService, "fetch_resource_definitions", fake_fetch_defs) # Determine endpoint based on transport type endpoint = "test-command" if transport_type == MCPTransportType.STDIO else "http://test-endpoint" working_dir = "/tmp" if transport_type == MCPTransportType.STDIO else None # Call fetch_mcp_resources_async with different transport types resources = await fetch_mcp_resources_async(endpoint, transport_type, working_directory=working_dir) assert len(resources) == 1 res_cls = resources[0] # Verify the resource was created correctly assert hasattr(res_cls, "mcp_resource_name") @pytest.mark.asyncio @pytest.mark.parametrize("transport_type", [MCPTransportType.HTTP_STREAM, MCPTransportType.STDIO, MCPTransportType.SSE]) async def test_fetch_mcp_prompts_async_all_transport_types(monkeypatch, transport_type): """Test fetch_mcp_prompts_async with all supported transport types.""" definitions = [ MCPPromptDefinition( name=f"Prompt_{transport_type.value}", description=f"Test prompt for {transport_type.value}", input_schema={"type": "object", "properties": {}, "required": []}, ) ] async def fake_fetch_defs(self): return definitions monkeypatch.setattr(MCPDefinitionService, "fetch_prompt_definitions", fake_fetch_defs) # Determine endpoint based on transport type endpoint = "test-command" if transport_type == MCPTransportType.STDIO else "http://test-endpoint" working_dir = "/tmp" if transport_type == MCPTransportType.STDIO else None # Call fetch_mcp_prompts_async with different transport types prompts = await fetch_mcp_prompts_async(endpoint, transport_type, working_directory=working_dir) assert len(prompts) == 1 prompt_cls = prompts[0] # Verify the prompt was created correctly assert hasattr(prompt_cls, "mcp_prompt_name") @pytest.mark.asyncio async def test_fetch_mcp_tools_async_multiple_tools(monkeypatch): """Test fetch_mcp_tools_async with multiple tool definitions.""" definitions = [ MCPToolDefinition( name="Tool1", description="First tool", input_schema={"type": "object", "properties": {}, "required": []} ), MCPToolDefinition( name="Tool2", description="Second tool", input_schema={"type": "object", "properties": {"param": {"type": "string"}}, "required": ["param"]}, ), MCPToolDefinition( name="Tool3", description="Third tool", input_schema={ "type": "object", "properties": {"x": {"type": "number"}, "y": {"type": "number"}}, "required": ["x", "y"], }, ), ] async def fake_fetch_defs(self): return definitions monkeypatch.setattr(MCPDefinitionService, "fetch_tool_definitions", fake_fetch_defs) # Call fetch_mcp_tools_async tools = await fetch_mcp_tools_async("http://test-endpoint", MCPTransportType.HTTP_STREAM) assert len(tools) == 3 tool_names = [getattr(tool_cls, "mcp_tool_name", None) for tool_cls in tools] assert "Tool1" in tool_names assert "Tool2" in tool_names assert "Tool3" in tool_names @pytest.mark.asyncio async def test_fetch_mcp_resources_async_multiple_resources(monkeypatch): """Test fetch_mcp_resources_async with multiple resource definitions.""" definitions = [ MCPResourceDefinition( name="Res1", description="First resource", uri="resource://Res1", input_schema={"type": "object", "properties": {}, "required": []}, ), MCPResourceDefinition( name="Res2", description="Second resource", uri="resource://Res2", input_schema={"type": "object", "properties": {"param": {"type": "string"}}, "required": ["param"]}, ), MCPResourceDefinition( name="Res3", description="Third resource", uri="resource://Res3", input_schema={ "type": "object", "properties": {"x": {"type": "number"}, "y": {"type": "number"}}, "required": ["x", "y"], }, ), ] async def fake_fetch_defs(self): return definitions monkeypatch.setattr(MCPDefinitionService, "fetch_resource_definitions", fake_fetch_defs) # Call fetch_mcp_resources_async resources = await fetch_mcp_resources_async("http://test-endpoint", MCPTransportType.HTTP_STREAM) assert len(resources) == 3 res_names = [getattr(res_cls, "mcp_resource_name", None) for res_cls in resources] assert "Res1" in res_names assert "Res2" in res_names assert "Res3" in res_names @pytest.mark.asyncio async def test_fetch_mcp_prompts_async_multiple_prompts(monkeypatch): """Test fetch_mcp_prompts_async with multiple prompt definitions.""" definitions = [ MCPPromptDefinition( name="Prompt1", description="First prompt", input_schema={"type": "object", "properties": {}, "required": []} ), MCPPromptDefinition( name="Prompt2", description="Second prompt", input_schema={"type": "object", "properties": {"param": {"type": "string"}}, "required": ["param"]}, ), MCPPromptDefinition( name="Prompt3", description="Third prompt", input_schema={ "type": "object", "properties": {"x": {"type": "number"}, "y": {"type": "number"}}, "required": ["x", "y"], }, ), ] async def fake_fetch_defs(self): return definitions monkeypatch.setattr(MCPDefinitionService, "fetch_prompt_definitions", fake_fetch_defs) # Call fetch_mcp_prompts_async prompts = await fetch_mcp_prompts_async("http://test-endpoint", MCPTransportType.HTTP_STREAM) assert len(prompts) == 3 prompt_names = [getattr(prompt_cls, "mcp_prompt_name", None) for prompt_cls in prompts] assert "Prompt1" in prompt_names assert "Prompt2" in prompt_names assert "Prompt3" in prompt_names # Tests for arun functionality def test_arun_attribute_exists_on_generated_tools(monkeypatch): """Test that dynamically generated tools have the arun attribute.""" input_schema = {"type": "object", "properties": {}, "required": []} definitions = [MCPToolDefinition(name="TestTool", description="test", input_schema=input_schema)] monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: definitions) # Create tool tools = fetch_mcp_tools("http://test", MCPTransportType.HTTP_STREAM) tool_cls = tools[0] # Verify the class has arun as an attribute assert hasattr(tool_cls, "arun") # Verify instance has arun inst = tool_cls() assert hasattr(inst, "arun") assert callable(getattr(inst, "arun")) def test_arun_attribute_exists_on_generated_resources(monkeypatch): """Test that dynamically generated resources have the arun attribute.""" input_schema = {"type": "object", "properties": {}, "required": []} definitions = [ MCPResourceDefinition(name="TestRes", description="test", uri="resource://TestRes", input_schema=input_schema) ] monkeypatch.setattr(MCPFactory, "_fetch_resource_definitions", lambda self: definitions) # Create resource resources = fetch_mcp_resources("http://test", MCPTransportType.HTTP_STREAM) res_cls = resources[0] # Verify the class has aread as an attribute assert hasattr(res_cls, "aread") # Verify instance has aread inst = res_cls() assert hasattr(inst, "aread") assert callable(getattr(inst, "aread")) def test_arun_attribute_exists_on_generated_prompts(monkeypatch): """Test that dynamically generated prompts have the arun attribute.""" input_schema = {"type": "object", "properties": {}, "required": []} definitions = [MCPPromptDefinition(name="TestPrompt", description="test", input_schema=input_schema)] monkeypatch.setattr(MCPFactory, "_fetch_prompt_definitions", lambda self: definitions) # Create prompt prompts = fetch_mcp_prompts("http://test", MCPTransportType.HTTP_STREAM) prompt_cls = prompts[0] # Verify the class has aread as an attribute assert hasattr(prompt_cls, "agenerate") # Verify instance has aread inst = prompt_cls() assert hasattr(inst, "agenerate") assert callable(getattr(inst, "agenerate")) @pytest.mark.asyncio async def test_arun_tool_async_execution(monkeypatch): """Test that arun method executes tool asynchronously.""" import atomic_agents.connectors.mcp.mcp_factory as mtf class DummyTransportCM: def __init__(self, ret): self.ret = ret async def __aenter__(self): return self.ret async def __aexit__(self, exc_type, exc, tb): pass def dummy_http_client(endpoint): return DummyTransportCM((None, None, None)) class DummySessionCM: def __init__(self, rs=None, ws=None, *args): pass async def initialize(self): pass async def call_tool(self, name, arguments): return {"content": f"async-{name}-{arguments}-ok"} async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): pass monkeypatch.setattr(mtf, "streamablehttp_client", dummy_http_client) monkeypatch.setattr(mtf, "ClientSession", DummySessionCM) # Prepare definitions input_schema = {"type": "object", "properties": {}, "required": []} definitions = [MCPToolDefinition(name="AsyncTool", description="async test", input_schema=input_schema)] monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: definitions) # Create tool and test arun tools = fetch_mcp_tools("http://test", MCPTransportType.HTTP_STREAM) tool_cls = tools[0] inst = tool_cls() # Test arun execution arun_method = getattr(inst, "arun") # type: ignore params = tool_cls.input_schema(tool_name="AsyncTool") # type: ignore result = await arun_method(params) assert result.result == "async-AsyncTool-{}-ok" @pytest.mark.asyncio async def test_aread_resource_async_execution(monkeypatch): """Test that aread method executes resource asynchronously.""" import atomic_agents.connectors.mcp.mcp_factory as mtf class DummyTransportCM: def __init__(self, ret): self.ret = ret async def __aenter__(self): return self.ret async def __aexit__(self, exc_type, exc, tb): pass def dummy_http_client(endpoint): return DummyTransportCM((None, None, None)) class DummySessionCM: def __init__(self, rs=None, ws=None, *args): pass async def initialize(self): pass async def read_resource(self, uri): # If uri is resource://AsyncRes/{id}, name is AsyncRes name = uri.split("/")[2].split("-")[0] return {"content": f"async-{name}-ok"} async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): pass monkeypatch.setattr(mtf, "streamablehttp_client", dummy_http_client) monkeypatch.setattr(mtf, "ClientSession", DummySessionCM) # Prepare definitions input_schema = {"type": "object", "properties": {}, "required": []} definitions = [ MCPResourceDefinition(name="AsyncRes", description="async test", uri="resource://AsyncRes", input_schema=input_schema) ] monkeypatch.setattr(MCPFactory, "_fetch_resource_definitions", lambda self: definitions) # Create resource and test aread resources = fetch_mcp_resources("http://test", MCPTransportType.HTTP_STREAM) res_cls = resources[0] inst = res_cls() # Test aread execution aread_method = getattr(inst, "aread") # type: ignore params = res_cls.input_schema(resource_name="AsyncRes") # type: ignore result = await aread_method(params) assert result.content["content"] == "async-AsyncRes-ok" @pytest.mark.asyncio async def test_agenerate_prompt_async_execution(monkeypatch): """Test that agenerate method executes prompt asynchronously.""" import atomic_agents.connectors.mcp.mcp_factory as mtf class DummyTransportCM: def __init__(self, ret): self.ret = ret async def __aenter__(self): return self.ret async def __aexit__(self, exc_type, exc, tb): pass def dummy_http_client(endpoint): return DummyTransportCM((None, None, None)) class DummySessionCM: def __init__(self, rs=None, ws=None, *args): pass async def initialize(self): pass async def get_prompt(self, *, name, arguments): class Msg(BaseModel): content: str return {"messages": [Msg(content=f"async-{name}-{arguments}-ok")]} async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): pass monkeypatch.setattr(mtf, "streamablehttp_client", dummy_http_client) monkeypatch.setattr(mtf, "ClientSession", DummySessionCM) # Prepare definitions input_schema = {"type": "object", "properties": {}, "required": []} definitions = [MCPPromptDefinition(name="AsyncPrompt", description="async test", input_schema=input_schema)] monkeypatch.setattr(MCPFactory, "_fetch_prompt_definitions", lambda self: definitions) # Create prompt and test agenerate prompts = fetch_mcp_prompts("http://test", MCPTransportType.HTTP_STREAM) prompt_cls = prompts[0] inst = prompt_cls() # Test agenerate execution agenerate_method = getattr(inst, "agenerate") # type: ignore params = prompt_cls.input_schema(prompt_name="AsyncPrompt") # type: ignore result = await agenerate_method(params) assert result.content == "async-AsyncPrompt-{}-ok" @pytest.mark.asyncio async def test_arun_error_handling(monkeypatch): """Test that arun properly handles and wraps errors.""" import atomic_agents.connectors.mcp.mcp_factory as mtf class DummyTransportCM: def __init__(self, ret): self.ret = ret async def __aenter__(self): return self.ret async def __aexit__(self, exc_type, exc, tb): pass def dummy_http_client(endpoint): return DummyTransportCM((None, None, None)) class ErrorSessionCM: def __init__(self, rs=None, ws=None, *args): pass async def initialize(self): pass async def call_tool(self, name, arguments): raise RuntimeError("Tool execution failed") async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): pass monkeypatch.setattr(mtf, "streamablehttp_client", dummy_http_client) monkeypatch.setattr(mtf, "ClientSession", ErrorSessionCM) # Prepare definitions input_schema = {"type": "object", "properties": {}, "required": []} definitions = [MCPToolDefinition(name="ErrorTool", description="error test", input_schema=input_schema)] monkeypatch.setattr(MCPFactory, "_fetch_tool_definitions", lambda self: definitions) # Create tool and test arun error handling tools = fetch_mcp_tools("http://test", MCPTransportType.HTTP_STREAM) tool_cls = tools[0] inst = tool_cls() # Test that arun properly wraps errors arun_method = getattr(inst, "arun") # type: ignore params = tool_cls.input_schema(tool_name="ErrorTool") # type: ignore with pytest.raises(RuntimeError) as exc_info: await arun_method(params) assert "Failed to execute MCP tool 'ErrorTool'" in str(exc_info.value) @pytest.mark.asyncio async def test_resource_aread_error_handling(monkeypatch): """Test that aread properly handles and wraps errors.""" import atomic_agents.connectors.mcp.mcp_factory as mtf class DummyTransportCM: def __init__(self, ret): self.ret = ret async def __aenter__(self): return self.ret async def __aexit__(self, exc_type, exc, tb): pass def dummy_http_client(endpoint): return DummyTransportCM((None, None, None)) class ErrorSessionCM: def __init__(self, rs=None, ws=None, *args): pass async def initialize(self): pass async def read_resource(self, uri): raise RuntimeError("Resource read failed") async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): pass monkeypatch.setattr(mtf, "streamablehttp_client", dummy_http_client) monkeypatch.setattr(mtf, "ClientSession", ErrorSessionCM) # Prepare definitions input_schema = {"type": "object", "properties": {}, "required": []} definitions = [ MCPResourceDefinition(name="ErrorRes", description="error test", uri="resource://ErrorRes", input_schema=input_schema) ] monkeypatch.setattr(MCPFactory, "_fetch_resource_definitions", lambda self: definitions) # Create resource and test aread error handling resources = fetch_mcp_resources("http://test", MCPTransportType.HTTP_STREAM) res_cls = resources[0] inst = res_cls() # Test that aread properly wraps errors aread_method = getattr(inst, "aread") # type: ignore params = res_cls.input_schema(resource_name="ErrorRes") # type: ignore with pytest.raises(RuntimeError) as exc_info: await aread_method(params) assert "Failed to read MCP resource 'ErrorRes'" in str(exc_info.value) @pytest.mark.asyncio async def test_prompt_agenerate_error_handling(monkeypatch): """Test that agenerate properly handles and wraps errors.""" import atomic_agents.connectors.mcp.mcp_factory as mtf class DummyTransportCM: def __init__(self, ret): self.ret = ret async def __aenter__(self): return self.ret async def __aexit__(self, exc_type, exc, tb): pass def dummy_http_client(endpoint): return DummyTransportCM((None, None, None)) class ErrorSessionCM: def __init__(self, rs=None, ws=None, *args): pass async def initialize(self): pass async def get_prompt(self, *, name, arguments): raise RuntimeError("Prompt generation failed") async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): pass monkeypatch.setattr(mtf, "streamablehttp_client", dummy_http_client) monkeypatch.setattr(mtf, "ClientSession", ErrorSessionCM) # Prepare definitions input_schema = {"type": "object", "properties": {}, "required": []} definitions = [MCPPromptDefinition(name="ErrorPrompt", description="error test", input_schema=input_schema)] monkeypatch.setattr(MCPFactory, "_fetch_prompt_definitions", lambda self: definitions) # Create prompt and test agenerate error handling prompts = fetch_mcp_prompts("http://test", MCPTransportType.HTTP_STREAM) prompt_cls = prompts[0] inst = prompt_cls() # Test that agenerate properly wraps errors agenerate_method = getattr(inst, "agenerate") # type: ignore params = prompt_cls.input_schema(prompt_name="ErrorPrompt") # type: ignore with pytest.raises(RuntimeError) as exc_info: await agenerate_method(params) assert "Failed to get MCP prompt 'ErrorPrompt'" in str(exc_info.value) class StructuredModelResult(BaseModel): structuredContent: BaseModel | dict[str, str] | str | None class NamedToolPayload(BaseModel): name: str class TextAndDataContent(BaseModel): text: str | int | None data: dict[str, str] @pytest.fixture def named_output_tool_session(mocked_mcp_connection: tuple[AsyncMock, dict[str, MagicMock]]) -> AsyncMock: session, _ = mocked_mcp_connection session.list_tools.return_value.tools[0].outputSchema = { "type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"], } return session @pytest.mark.asyncio @pytest.mark.parametrize( "tool_result,expected_name", [ (CallToolResult(content=[TextContent(type="text", text='{"name":"text"}')], structuredContent={"name": "sdk"}), "sdk"), (CallToolResult(content=[TextContent(type="text", text='{"name":"text"}')]), "text"), (StructuredModelResult(structuredContent=NamedToolPayload(name="model")), "model"), ( MockContentResult(content=[MockContentItem(text='{"name":"first"}'), MockContentItem(text='{"name":"second"}')]), "first", ), (MockContentResult(content=(MockContentItemWithData(data={"name": "tuple"}),)), "tuple"), (MockContentResult(content=[TextAndDataContent(text="[]", data={"name": "data"})]), "data"), (MockContentResult(content=[TextAndDataContent(text="invalid" * 40, data={"name": "long"})]), "long"), (MockContentResult(content=[TextAndDataContent(text="invalid", data={"name": "short"})]), "short"), (MockContentResult(content=[TextAndDataContent(text=None, data={"name": "none-text"})]), "none-text"), ({"structuredContent": {"name": "structured"}, "content": {"name": "content"}}, "structured"), ({"content": {"name": "content"}, "name": "raw"}, "content"), ({"isError": True, "name": "test_value", "content": "article"}, "test_value"), ], ) async def test_typed_tool_conversion_preserves_precedence( named_output_tool_session: AsyncMock, tool_result: BaseModel | dict[str, object], expected_name: str ) -> None: named_output_tool_session.call_tool.return_value = tool_result tool_class = (await fetch_mcp_tools_async(client_session=named_output_tool_session))[0] result = await tool_class().arun(tool_class.input_schema(tool_name="EchoTool", query="hello")) assert result.model_dump() == {"name": expected_name} @pytest.mark.asyncio @pytest.mark.parametrize( "tool_result,cause_type,message", [ (StructuredModelResult(structuredContent="invalid"), TypeError, "structuredContent with unexpected type str"), (StructuredModelResult(structuredContent=""), TypeError, "structuredContent with unexpected type str"), (MockContentResult(content=[]), ValueError, "unparseable result"), (MockContentResult(content=[MockContentItem(text="invalid")]), ValueError, "unparseable result"), ( MockContentResult(content=[TextAndDataContent(text='{"wrong":"field"}', data={"name": "ignored"})]), ValueError, "validation error", ), ({"structuredContent": None, "name": "ignored"}, TypeError, "mapping"), ("raw string", ValueError, "unparseable result"), ], ) async def test_typed_tool_conversion_preserves_wrapped_failures( named_output_tool_session: AsyncMock, tool_result: BaseModel | dict[str, object] | str, cause_type: type[Exception], message: str, ) -> None: named_output_tool_session.call_tool.return_value = tool_result tool_class = (await fetch_mcp_tools_async(client_session=named_output_tool_session))[0] with pytest.raises(RuntimeError, match="Failed to execute MCP tool 'EchoTool':") as failure: await tool_class().arun(tool_class.input_schema(tool_name="EchoTool", query="hello")) assert isinstance(failure.value.__cause__, cause_type) assert message in str(failure.value) @pytest.mark.asyncio @pytest.mark.parametrize( "tool_result,expected_content", [ ( CallToolResult(content=[TextContent(type="text", text="first"), TextContent(type="text", text="second")]), [TextContent(type="text", text="first"), TextContent(type="text", text="second")], ), ({"isError": True, "name": "test_value", "content": "article"}, "article"), ({"name": "raw"}, {"name": "raw"}), (123, 123), ], ) async def test_generic_tool_conversion_preserves_content( mocked_mcp_connection: tuple[AsyncMock, dict[str, MagicMock]], tool_result: BaseModel | dict[str, object] | int, expected_content: list[TextContent] | dict[str, str] | str | int, ) -> None: session, _ = mocked_mcp_connection session.call_tool.return_value = tool_result tool_class = (await fetch_mcp_tools_async(client_session=session))[0] result = await tool_class().arun(tool_class.input_schema(tool_name="EchoTool", query="hello")) assert result.result == expected_content @pytest.fixture def generated_resource() -> BaseResource: definition = MCPResourceDefinition( name="ReadItem", description="Read an item", uri="resource://items/{item}", input_schema={"type": "object", "properties": {"item": {"type": "string"}}, "required": ["item"]}, mime_type="text/plain", ) return MCPFactory("http://example.test", MCPTransportType.SSE)._create_resource_classes([definition])[0]() @pytest.fixture def generated_prompt() -> BasePrompt: definition = MCPPromptDefinition( name="DescribeItem", description="Describe an item", input_schema={"type": "object", "properties": {"item": {"type": "string"}}, "required": ["item"]}, ) return MCPFactory("http://example.test", MCPTransportType.SSE)._create_prompt_classes([definition])[0]() @pytest.mark.parametrize( "resource_result", [ReadResourceResult(contents=[]), {"contents": [], "mime_type": "text/plain"}], ids=["sdk-empty", "mapping-empty"], ) def test_generated_resource_sse_closes_contexts( mocked_mcp_connection: tuple[AsyncMock, dict[str, MagicMock]], generated_resource: BaseResource, resource_result: ReadResourceResult | dict[str, object], ) -> None: session, clients = mocked_mcp_connection session.read_resource.return_value = resource_result result = generated_resource.read(generated_resource.input_schema(resource_name="ReadItem", item="entry")) assert result.model_dump() == {"content": [], "mime_type": "text/plain"} session.read_resource.assert_awaited_once_with(uri="resource://items/entry") clients["sse_client"].assert_called_once_with("http://example.test/sse") session.initialize.assert_awaited_once_with() session.__aexit__.assert_awaited_once_with(session, None, None, None) transport = clients["sse_client"].return_value transport.__aexit__.assert_awaited_once_with(transport, None, None, None) @pytest.mark.parametrize( "mime_fields,expected_mime", [({"mime_type": "application/json"}, "application/json"), ({}, "text/plain")] ) def test_generated_resource_mapping_adapter_preserves_contents_and_mime( mocked_mcp_connection: tuple[AsyncMock, dict[str, MagicMock]], generated_resource: BaseResource, mime_fields: dict[str, str], expected_mime: str, ) -> None: session, _ = mocked_mcp_connection session.read_resource.return_value = {"contents": [{"text": "item"}], **mime_fields} result = generated_resource.read(generated_resource.input_schema(resource_name="ReadItem", item="entry")) assert result.content == [{"text": "item"}] assert result.mime_type == expected_mime session.read_resource.assert_awaited_once_with(uri="resource://items/entry") @pytest.mark.parametrize( "content", [TextContent(type="text", text="second"), ImageContent(type="image", data="aW1hZ2U=", mimeType="image/png")], ids=["ordered-text", "nontext-fallback"], ) def test_generated_prompt_sse_preserves_sdk_message_order_and_closes_contexts( mocked_mcp_connection: tuple[AsyncMock, dict[str, MagicMock]], generated_prompt: BasePrompt, content: TextContent | ImageContent, ) -> None: session, clients = mocked_mcp_connection session.get_prompt.return_value = GetPromptResult( messages=[ PromptMessage(role="user", content=TextContent(type="text", text="first")), PromptMessage(role="assistant", content=content), ] ) result = generated_prompt.generate(generated_prompt.input_schema(prompt_name="DescribeItem", item="entry")) assert result.content == f"first\n\n{getattr(content, 'text', str(content))}" session.get_prompt.assert_awaited_once_with(name="DescribeItem", arguments={"item": "entry"}) clients["sse_client"].assert_called_once_with("http://example.test/sse") session.initialize.assert_awaited_once_with() session.__aexit__.assert_awaited_once_with(session, None, None, None) transport = clients["sse_client"].return_value transport.__aexit__.assert_awaited_once_with(transport, None, None, None) @pytest.mark.parametrize( "messages,expected_content", [ ([{"content": "first"}, {"content": {"text": "second"}}], "first\n\nsecond"), ([{"content": {}}, {"text": "direct"}, 17], "\n\ndirect\n\n17"), ], ids=["string-and-text-mapping", "empty-and-direct-fallbacks"], ) def test_generated_prompt_mapping_adapter_renders_messages( mocked_mcp_connection: tuple[AsyncMock, dict[str, MagicMock]], generated_prompt: BasePrompt, messages: list[dict[str, object] | int], expected_content: str, ) -> None: session, _ = mocked_mcp_connection session.get_prompt.return_value = {"messages": messages} result = generated_prompt.generate(generated_prompt.input_schema(prompt_name="DescribeItem", item="entry")) assert result.content == expected_content session.get_prompt.assert_awaited_once_with(name="DescribeItem", arguments={"item": "entry"}) @pytest.mark.parametrize("prompt_result", [{}, "missing messages"], ids=["mapping", "plain"]) def test_generated_prompt_rejects_missing_messages_and_closes_contexts( mocked_mcp_connection: tuple[AsyncMock, dict[str, MagicMock]], generated_prompt: BasePrompt, prompt_result: dict[str, object] | str, ) -> None: session, clients = mocked_mcp_connection session.get_prompt.return_value = prompt_result with pytest.raises( RuntimeError, match="Failed to get MCP prompt 'DescribeItem': Prompt response has no messages." ) as failure: generated_prompt.generate(generated_prompt.input_schema(prompt_name="DescribeItem", item="entry")) assert type(failure.value.__cause__) is Exception assert str(failure.value.__cause__) == "Prompt response has no messages." session.get_prompt.assert_awaited_once_with(name="DescribeItem", arguments={"item": "entry"}) session.__aexit__.assert_awaited_once_with(session, None, None, None) transport = clients["sse_client"].return_value transport.__aexit__.assert_awaited_once_with(transport, None, None, None) @pytest.mark.parametrize( "attribute_fixture,method_name,discriminator", [ ("generated_resource", "read", "resource_name"), ("generated_prompt", "generate", "prompt_name"), ], ) def test_generated_attributes_reject_unknown_transport_without_connection( request: pytest.FixtureRequest, mocked_mcp_connection: tuple[AsyncMock, dict[str, MagicMock]], attribute_fixture: str, method_name: str, discriminator: str, ) -> None: session, clients = mocked_mcp_connection attribute: BaseResource | BasePrompt = request.getfixturevalue(attribute_fixture) attribute.transport_type = "unsupported" name = getattr(attribute, "mcp_" + discriminator) parameters = attribute.input_schema(**{discriminator: name, "item": "entry"}) with pytest.raises(RuntimeError, match="Unknown transport type: unsupported") as failure: getattr(attribute, method_name)(parameters) assert isinstance(failure.value.__cause__, ValueError) assert "Available transport types: ['sse', 'http_stream', 'stdio']" in str(failure.value.__cause__) session.__aenter__.assert_not_awaited() for client in clients.values(): client.assert_not_called() def test_generated_tool_rejects_unknown_transport_without_connection( mocked_mcp_connection: tuple[AsyncMock, dict[str, MagicMock]], ) -> None: session, clients = mocked_mcp_connection definition = MCPToolDefinition(name="EchoTool", description="Echo", input_schema={"type": "object", "properties": {}}) tool = MCPFactory("http://example.test")._create_tool_classes([definition])[0]() tool.transport_type = "unsupported" with pytest.raises(RuntimeError, match="Unknown transport type: unsupported") as failure: tool.run(tool.input_schema(tool_name="EchoTool")) assert isinstance(failure.value.__cause__, ValueError) assert "Available transport types: ['sse', 'http_stream', 'stdio']" in str(failure.value.__cause__) session.call_tool.assert_not_awaited() for client in clients.values(): client.assert_not_called() def test_orchestrator_without_attribute_groups_warns(caplog: pytest.LogCaptureFixture) -> None: assert MCPFactory("http://example.test").create_orchestrator_schema() is None assert "No tools/resources/prompts provided to create orchestrator schema" in caplog.text @pytest.mark.asyncio @pytest.mark.parametrize("borrowed", [False, True], ids=["fresh", "persistent"]) async def test_generated_resource_preserves_dotted_sdk_template_when_python_formatting_fails( mocked_mcp_connection: tuple[AsyncMock, dict[str, MagicMock]], borrowed: bool ) -> None: session, clients = mocked_mcp_connection session.list_resources.return_value = ListResourcesResult(resources=[]) session.list_resource_templates.return_value = ListResourceTemplatesResult( resourceTemplates=[ResourceTemplate(name="DottedItem", uriTemplate="resource://items/{item.name}")] ) definitions = await MCPDefinitionService.fetch_resource_definitions_from_session(session) session.initialize.reset_mock() factory = MCPFactory("http://example.test", MCPTransportType.SSE) resource = factory._create_resource_classes(definitions)[0]() if borrowed: resource._client_session = session session.read_resource.return_value = ReadResourceResult(contents=[]) result = await resource.aread(resource.input_schema(resource_name="DottedItem", **{"item.name": "entry"})) assert result.content == [] session.read_resource.assert_awaited_once_with(uri="resource://items/{item.name}") assert (session.initialize.await_count, clients["sse_client"].return_value.__aexit__.await_count) == ( int(not borrowed), int(not borrowed), ) @pytest.mark.asyncio async def test_typed_tool_rejects_model_adapter_without_output_fields(named_output_tool_session: AsyncMock) -> None: named_output_tool_session.call_tool.return_value = NamedToolPayload(name="adapter-without-content") tool_class = (await fetch_mcp_tools_async(client_session=named_output_tool_session))[0] with pytest.raises(RuntimeError, match="unparseable result") as failure: await tool_class().arun(tool_class.input_schema(tool_name="EchoTool", query="hello")) assert isinstance(failure.value.__cause__, ValueError) named_output_tool_session.call_tool.assert_awaited_once_with(name="EchoTool", arguments={"query": "hello"}) ``` ### File: atomic-agents/tests/connectors/mcp/test_non_object_refs.py ```python import pytest from pydantic import ValidationError from atomic_agents.connectors.mcp import SchemaTransformer @pytest.mark.parametrize("definitions_key", ["$defs", "definitions"]) @pytest.mark.parametrize( "definition, valid, invalid", [ ({"type": "string"}, "query", {}), ({"type": "integer"}, 42, {}), ({"type": "array", "items": {"type": "integer"}}, [1, 2], ["invalid"]), ({"anyOf": [{"type": "string"}, {"type": "null"}]}, None, {}), ], ) def test_non_object_definitions_preserve_their_value_type(definitions_key, definition, valid, invalid): reference = {"$ref": f"#/{definitions_key}/Value"} schema = { "type": "object", "properties": {"value": reference, "values": {"type": "array", "items": reference}}, "required": ["value", "values"], definitions_key: {"Value": definition}, } model = SchemaTransformer.create_model_from_schema(schema, "ValuesInput", "values") instance = model(tool_name="values", value=valid, values=[valid]) assert instance.value == valid assert instance.values == [valid] with pytest.raises(ValidationError): model(tool_name="values", value=invalid, values=[]) with pytest.raises(ValidationError): model(tool_name="values", value=valid, values=[invalid]) def test_definition_aliases_resolve_to_the_underlying_scalar(): schema = { "type": "object", "properties": {"value": {"$ref": "#/$defs/Alias"}}, "required": ["value"], "$defs": {"Alias": {"$ref": "#/$defs/Value"}, "Value": {"type": "string"}}, } model = SchemaTransformer.create_model_from_schema(schema, "AliasOutput", "alias", is_output_schema=True) assert model(value="result").model_dump() == {"value": "result"} with pytest.raises(ValidationError): model(value={}) ``` ### File: atomic-agents/tests/connectors/mcp/test_paginated_discovery.py ```python from unittest.mock import AsyncMock, call import pytest from mcp import ClientSession from mcp import types from atomic_agents.connectors.mcp import MCPDefinitionService @pytest.mark.asyncio @pytest.mark.parametrize("empty_first_page", [False, True]) async def test_discover_all_tool_pages(empty_first_page): first = types.Tool(name="first", inputSchema={"type": "object"}) last = types.Tool(name="last", inputSchema={"type": "object"}, outputSchema={"type": "string"}) session = AsyncMock(spec=ClientSession) session.list_tools.side_effect = [ types.ListToolsResult(tools=[] if empty_first_page else [first], nextCursor="tools-next"), types.ListToolsResult(tools=[last]), ] definitions = await MCPDefinitionService.fetch_tool_definitions_from_session(session) assert [definition.name for definition in definitions] == ([] if empty_first_page else ["first"]) + ["last"] assert definitions[-1].output_schema == {"type": "string"} assert session.list_tools.await_args_list == [call(), call(cursor="tools-next")] session.initialize.assert_awaited_once() @pytest.mark.asyncio @pytest.mark.parametrize("empty_first_page", [False, True]) async def test_discover_all_prompt_pages(empty_first_page): session = AsyncMock(spec=ClientSession) session.list_prompts.side_effect = [ types.ListPromptsResult(prompts=[] if empty_first_page else [types.Prompt(name="first")], nextCursor="prompts-next"), types.ListPromptsResult( prompts=[types.Prompt(name="last", arguments=[types.PromptArgument(name="topic", required=True)])] ), ] definitions = await MCPDefinitionService.fetch_prompt_definitions_from_session(session) assert [definition.name for definition in definitions] == ([] if empty_first_page else ["first"]) + ["last"] assert definitions[-1].input_schema["required"] == ["topic"] assert session.list_prompts.await_args_list == [call(), call(cursor="prompts-next")] @pytest.mark.asyncio @pytest.mark.parametrize("empty_first_page", [False, True]) async def test_discover_all_resource_and_template_pages(empty_first_page): session = AsyncMock(spec=ClientSession) session.list_resources.side_effect = [ types.ListResourcesResult( resources=[] if empty_first_page else [types.Resource(name="first", uri="file:///first")], nextCursor="resources-next", ), types.ListResourcesResult(resources=[types.Resource(name="last", uri="file:///last", mimeType="text/plain")]), ] session.list_resource_templates.side_effect = [ types.ListResourceTemplatesResult( resourceTemplates=( [] if empty_first_page else [types.ResourceTemplate(name="template-first", uriTemplate="file:///{id}")] ), nextCursor="templates-next", ), types.ListResourceTemplatesResult( resourceTemplates=[types.ResourceTemplate(name="template-last", uriTemplate="data:///{key}")] ), ] definitions = await MCPDefinitionService.fetch_resource_definitions_from_session(session) expected = ([] if empty_first_page else ["first"]) + ["last"] expected += ([] if empty_first_page else ["template-first"]) + ["template-last"] assert [definition.name for definition in definitions] == expected assert next(definition for definition in definitions if definition.name == "last").mime_type == "text/plain" assert definitions[-1].input_schema["required"] == ["key"] assert session.list_resources.await_args_list == [call(), call(cursor="resources-next")] assert session.list_resource_templates.await_args_list == [call(), call(cursor="templates-next")] ``` ### File: atomic-agents/tests/connectors/mcp/test_pagination_cycles.py ```python from unittest.mock import AsyncMock, call import pytest from mcp import ClientSession, types from atomic_agents.connectors.mcp import MCPDefinitionService CATALOGS = [ ( "list_tools", types.ListToolsResult, "tools", MCPDefinitionService.fetch_tool_definitions_from_session, ), ( "list_prompts", types.ListPromptsResult, "prompts", MCPDefinitionService.fetch_prompt_definitions_from_session, ), ( "list_resources", types.ListResourcesResult, "resources", MCPDefinitionService.fetch_resource_definitions_from_session, ), ( "list_resource_templates", types.ListResourceTemplatesResult, "resourceTemplates", MCPDefinitionService.fetch_resource_definitions_from_session, ), ] @pytest.fixture def session(): client = AsyncMock(spec=ClientSession) client.list_resources.return_value = types.ListResourcesResult(resources=[]) client.list_resource_templates.return_value = types.ListResourceTemplatesResult(resourceTemplates=[]) return client @pytest.mark.asyncio @pytest.mark.parametrize("method,result_type,field,discover", CATALOGS) @pytest.mark.parametrize("cursors", [("same", "same"), ("a", "b", "a"), ("", "")]) async def test_discovery_rejects_cyclic_pagination(session, method, result_type, field, discover, cursors): list_pages = getattr(session, method) list_pages.side_effect = [result_type(**{field: [], "nextCursor": cursor}) for cursor in cursors] + [ AssertionError("The repeated cursor must not trigger another request") ] with pytest.raises(RuntimeError, match="MCP pagination returned a repeated cursor"): await discover(session) assert list_pages.await_args_list == [call()] + [call(cursor=cursor) for cursor in cursors[:-1]] @pytest.mark.asyncio @pytest.mark.parametrize("method,result_type,field,discover", CATALOGS) async def test_empty_string_cursor_is_followed_once(session, method, result_type, field, discover): list_pages = getattr(session, method) list_pages.side_effect = [result_type(**{field: [], "nextCursor": ""}), result_type(**{field: []})] assert await discover(session) == [] assert list_pages.await_args_list == [call(), call(cursor="")] @pytest.mark.asyncio async def test_resource_and_template_cursors_are_independent(session): session.list_resources.side_effect = [ types.ListResourcesResult(resources=[], nextCursor="next"), types.ListResourcesResult(resources=[types.Resource(name="file", uri="file:///data")]), ] session.list_resource_templates.side_effect = [ types.ListResourceTemplatesResult(resourceTemplates=[], nextCursor="next"), types.ListResourceTemplatesResult( resourceTemplates=[types.ResourceTemplate(name="template", uriTemplate="file:///{id}")] ), ] definitions = await MCPDefinitionService.fetch_resource_definitions_from_session(session) assert [definition.name for definition in definitions] == ["file", "template"] assert session.list_resources.await_args_list == [call(), call(cursor="next")] assert session.list_resource_templates.await_args_list == [call(), call(cursor="next")] ``` ### File: atomic-agents/tests/connectors/mcp/test_resource_uri_identity.py ```python """Preserve MCP resource identifiers across catalog discovery and reads.""" from types import SimpleNamespace from unittest.mock import AsyncMock import pytest from mcp import types from atomic_agents.connectors.mcp.mcp_definition_service import MCPDefinitionService from atomic_agents.connectors.mcp.mcp_factory import fetch_mcp_resources_async ENCODED_URIS = [ "resource://bucket/key%2Fvalue", "resource://bucket/a%23b", "resource://bucket/a%3Fb", "resource://bucket/%7Bid%7D", "resource://bucket/%252F", "resource://bucket/caf%C3%A9", "resource://bucket/plain", ] def resource_session(uri: str, adapter: str = "sdk") -> AsyncMock: session = AsyncMock() if adapter == "sdk": response = types.ListResourcesResult(resources=[types.Resource(name="document", uri=uri)]) else: resource = {"name": "document", "description": None, "uri": uri} if adapter == "object": resource = SimpleNamespace(**resource) response = SimpleNamespace(resources=[resource], nextCursor=None) session.list_resources.return_value = response session.list_resource_templates.return_value = types.ListResourceTemplatesResult(resourceTemplates=[]) session.read_resource.return_value = types.ReadResourceResult(contents=[]) return session @pytest.mark.asyncio @pytest.mark.parametrize("uri", ENCODED_URIS) @pytest.mark.parametrize("adapter", ["sdk", "dict", "object"]) async def test_discovery_preserves_literal_resource_uri(uri: str, adapter: str) -> None: session = resource_session(uri, adapter) definition = (await MCPDefinitionService.fetch_resource_definitions_from_session(session))[0] assert definition.uri == uri assert definition.input_schema == {"type": "object", "properties": {}, "required": []} @pytest.mark.asyncio @pytest.mark.parametrize("uri", ENCODED_URIS) @pytest.mark.parametrize("adapter", ["sdk", "dict", "object"]) async def test_generated_resource_reads_the_advertised_identifier(uri: str, adapter: str) -> None: session = resource_session(uri, adapter) resource_class = (await fetch_mcp_resources_async(client_session=session))[0] resource = resource_class() await resource.aread(resource.input_schema(resource_name="document")) session.read_resource.assert_awaited_once_with(uri=uri) request = types.ReadResourceRequest( method="resources/read", params=types.ReadResourceRequestParams(uri=session.read_resource.call_args.kwargs["uri"]) ) assert request.model_dump(mode="json")["params"]["uri"] == uri @pytest.mark.asyncio @pytest.mark.parametrize("literal", ["key%2Fvalue", "a%23b", "%7Bliteral%7D", "%252F", "plain"]) async def test_template_expands_only_the_declared_parameter(literal: str) -> None: uri_template = f"resource://bucket/{literal}/{{item}}" session = resource_session("resource://bucket/static") session.list_resources.return_value = types.ListResourcesResult(resources=[]) session.list_resource_templates.return_value = types.ListResourceTemplatesResult( resourceTemplates=[types.ResourceTemplate(name="document", uriTemplate=uri_template, mimeType="text/plain")] ) definition = (await MCPDefinitionService.fetch_resource_definitions_from_session(session))[0] assert definition.uri == uri_template assert definition.mime_type == "text/plain" assert definition.input_schema["required"] == ["item"] assert set(definition.input_schema["properties"]) == {"item"} resource_class = (await fetch_mcp_resources_async(client_session=session))[0] resource = resource_class() await resource.aread(resource.input_schema(resource_name="document", item="entry")) session.read_resource.assert_awaited_once_with(uri=f"resource://bucket/{literal}/entry") @pytest.mark.asyncio async def test_plain_sdk_resource_with_braces_does_not_gain_parameters() -> None: session = resource_session("resource://bucket/{literal}") resource = session.list_resources.return_value.resources[0] definition = (await MCPDefinitionService.fetch_resource_definitions_from_session(session))[0] assert definition.uri == str(resource.uri) assert definition.input_schema == {"type": "object", "properties": {}, "required": []} @pytest.mark.asyncio @pytest.mark.parametrize("adapter", ["dict", "object"]) async def test_adapter_keeps_explicit_template_placeholders(adapter: str) -> None: session = resource_session("resource://bucket/%7Bliteral%7D/{item}", adapter) definition = (await MCPDefinitionService.fetch_resource_definitions_from_session(session))[0] assert definition.uri == "resource://bucket/%7Bliteral%7D/{item}" assert definition.input_schema["required"] == ["item"] ``` ### File: atomic-agents/tests/connectors/mcp/test_schema_constraints.py ```python import pytest from pydantic import ValidationError from atomic_agents.connectors.mcp import SchemaTransformer CONSTRAINT_CASES = [ ({"type": "integer", "minimum": 0}, 0, -1), ({"type": "integer", "maximum": 10}, 10, 11), ({"type": "number", "exclusiveMinimum": 0}, 0.5, 0), ({"type": "number", "exclusiveMaximum": 10}, 9.5, 10), ({"type": "number", "multipleOf": 0.25}, 1.25, 1.1), ({"type": "string", "minLength": 2}, "é字", "é"), ({"type": "string", "maxLength": 2}, "é字", "é字x"), ({"type": "string", "pattern": "^[a-z]+$"}, "hello", "hello123"), ({"type": "array", "items": {"type": "integer"}, "minItems": 2}, [1, 2], [1]), ({"type": "array", "items": {"type": "integer"}, "maxItems": 2}, [1, 2], [1, 2, 3]), ] def build_model(prop_schema, definitions=None): schema = {"type": "object", "properties": {"value": prop_schema}, "required": ["value"]} if definitions: schema["$defs"] = definitions return SchemaTransformer.create_model_from_schema(schema, "ConstrainedInput", "constrained") @pytest.mark.parametrize("prop_schema,valid,invalid", CONSTRAINT_CASES) def test_property_constraints_validate_values(prop_schema, valid, invalid): model = build_model(prop_schema) assert model(tool_name="constrained", value=valid).value == valid with pytest.raises(ValidationError): model(tool_name="constrained", value=invalid) @pytest.mark.parametrize("prop_schema,valid,invalid", CONSTRAINT_CASES) def test_property_constraints_are_advertised_to_models(prop_schema, valid, invalid): model = build_model(prop_schema) advertised = model.model_json_schema()["properties"]["value"] for key, value in prop_schema.items(): if key not in {"type", "items"}: assert advertised[key] == value @pytest.mark.parametrize( "prop_schema,valid,invalid", [ ({"type": "array", "items": {"type": "integer", "minimum": 2}}, [2, 3], [1, 2]), ({"type": "object", "additionalProperties": {"type": "integer", "minimum": 2}}, {"n": 2}, {"n": 1}), ({"$ref": "#/$defs/Bounded"}, 2, 1), ({"anyOf": [{"type": "integer", "minimum": 2}, {"type": "null"}]}, 2, 1), ({"type": ["integer", "null"], "minimum": 2}, 2, 1), ], ) def test_constraints_survive_recursive_conversion(prop_schema, valid, invalid): model = build_model(prop_schema, {"Bounded": {"type": "integer", "minimum": 2}}) assert model(tool_name="constrained", value=valid).value == valid with pytest.raises(ValidationError): model(tool_name="constrained", value=invalid) def test_numeric_constraints_apply_after_existing_type_coercion(): model = build_model({"type": "integer", "minimum": 5}) assert model(tool_name="constrained", value="5").value == 5 with pytest.raises(ValidationError): model(tool_name="constrained", value="4") def test_multiple_of_preserves_floating_point_tolerance(): model = build_model({"type": "number", "multipleOf": 0.1}) assert model(tool_name="constrained", value=0.3).value == 0.3 def test_optional_constrained_field_can_be_omitted(): model = SchemaTransformer.create_model_from_schema( {"type": "object", "properties": {"value": {"type": "integer", "minimum": 2}}}, "OptionalConstrainedInput", "constrained", ) assert model(tool_name="constrained").value is None with pytest.raises(ValidationError): model(tool_name="constrained", value=1) @pytest.mark.parametrize("pattern", ["(?<=prefix:)[a-z]+$", "["]) def test_unsupported_pattern_preserves_tool_and_other_constraints(pattern, caplog): model = build_model({"type": "string", "pattern": pattern, "minLength": 2}) assert model(tool_name="constrained", value="other:hello").value == "other:hello" assert "Unsupported JSON Schema pattern" in caplog.text assert "pattern" not in model.model_json_schema()["properties"]["value"] with pytest.raises(ValidationError): model(tool_name="constrained", value="x") def test_unicode_property_pattern_builds_and_validates(): model = build_model({"type": "string", "pattern": r"^\p{L}+$"}) assert model(tool_name="constrained", value="é字").value == "é字" with pytest.raises(ValidationError): model(tool_name="constrained", value="é字1") assert model.model_json_schema()["properties"]["value"]["pattern"] == r"^\p{L}+$" @pytest.mark.parametrize("keyword", ["minimum", "maximum", "exclusiveMinimum", "exclusiveMaximum", "multipleOf"]) @pytest.mark.parametrize("invalid_bound", [True, False, "5"]) def test_nonnumeric_bounds_do_not_become_numeric_constraints(keyword, invalid_bound): model = build_model({"type": "number", keyword: invalid_bound}) for value in [-2, 0, 2.5]: assert model(tool_name="constrained", value=value).value == value assert keyword not in model.model_json_schema()["properties"]["value"] @pytest.mark.parametrize("exclusive", [True, False]) def test_draft04_boolean_exclusivity_does_not_override_numeric_bounds(exclusive): model = build_model( {"type": "number", "minimum": 5, "maximum": 10, "exclusiveMinimum": exclusive, "exclusiveMaximum": exclusive} ) assert model(tool_name="constrained", value=5).value == 5 assert model(tool_name="constrained", value=10).value == 10 with pytest.raises(ValidationError): model(tool_name="constrained", value=4) with pytest.raises(ValidationError): model(tool_name="constrained", value=11) def test_nested_quantifier_pattern_rejects_long_input(): model = build_model({"type": "string", "pattern": "^(a+)+$"}) with pytest.raises(ValidationError): model(tool_name="constrained", value="a" * 5000 + "!") ``` ### File: atomic-agents/tests/connectors/mcp/test_schema_transformer.py ```python import logging import pytest from jsonschema import Draft202012Validator from typing import Any, Dict, List, Optional, Union from pydantic import ValidationError from atomic_agents import BaseIOSchema from atomic_agents.connectors.mcp import SchemaTransformer class TestSchemaTransformer: def test_string_type_required(self): prop_schema = {"type": "string", "description": "A string field"} result = SchemaTransformer.json_to_pydantic_field(prop_schema, True) assert result[0] == str assert result[1].description == "A string field" assert result[1].is_required() is True def test_number_type_optional(self): prop_schema = {"type": "number", "description": "A number field"} result = SchemaTransformer.json_to_pydantic_field(prop_schema, False) assert result[0] == Optional[float] assert result[1].description == "A number field" assert result[1].default is None def test_integer_type_with_default(self): prop_schema = {"type": "integer", "description": "An integer field", "default": 42} result = SchemaTransformer.json_to_pydantic_field(prop_schema, False) assert result[0] == int assert result[1].description == "An integer field" assert result[1].default == 42 def test_boolean_type(self): prop_schema = {"type": "boolean", "description": "A boolean field"} result = SchemaTransformer.json_to_pydantic_field(prop_schema, True) assert result[0] == bool assert result[1].description == "A boolean field" assert result[1].is_required() is True def test_array_type_with_string_items(self): prop_schema = {"type": "array", "description": "An array of strings", "items": {"type": "string"}} result = SchemaTransformer.json_to_pydantic_field(prop_schema, True) assert result[0] == List[str] assert result[1].description == "An array of strings" assert result[1].is_required() is True def test_array_type_with_untyped_items(self): prop_schema = {"type": "array", "description": "An array of unknown types", "items": {}} result = SchemaTransformer.json_to_pydantic_field(prop_schema, True) assert result[0] == List[Any] assert result[1].description == "An array of unknown types" assert result[1].is_required() is True def test_object_type(self): prop_schema = {"type": "object", "description": "An object field"} result = SchemaTransformer.json_to_pydantic_field(prop_schema, True) assert result[0] == Dict[str, Any] assert result[1].description == "An object field" assert result[1].is_required() is True def test_unknown_type(self): prop_schema = {"type": "unknown", "description": "An unknown field"} result = SchemaTransformer.json_to_pydantic_field(prop_schema, True) assert result[0] == Any assert result[1].description == "An unknown field" assert result[1].is_required() is True def test_no_type(self): prop_schema = {"description": "A field without type"} result = SchemaTransformer.json_to_pydantic_field(prop_schema, True) assert result[0] == Any assert result[1].description == "A field without type" assert result[1].is_required() is True class TestCreateModelFromSchema: @pytest.mark.parametrize("required", [True, False]) @pytest.mark.parametrize( "prop_schema", [ {"enum": [1, "bad"], "anyOf": [{"type": "integer"}]}, {"enum": [1, 2.5, "bad"], "oneOf": [{"type": "integer"}, {"type": "number"}]}, {"enum": [1, 2, 3, "bad"], "allOf": [{"type": "number"}, {"minimum": 2}]}, {"const": 1, "anyOf": [{"type": "integer"}]}, {"const": 2.5, "oneOf": [{"type": "integer"}, {"type": "number"}]}, {"const": 3, "allOf": [{"type": "number"}, {"minimum": 2}]}, {"enum": [True, 1, "bad"], "anyOf": [{"type": "boolean"}, {"type": "integer"}]}, {"enum": [1, 2], "const": 2}, {"enum": [1, "bad"], "anyOf": [{"$ref": "#/$defs/Integer"}]}, { "enum": [0, 1, 2, 3], "anyOf": [{"type": "integer"}, {"type": "string"}], "oneOf": [{"const": 1}, {"const": 2}], "allOf": [{"minimum": 2}], }, {"enum": [1, "bad"], "anyOf": [False, {"type": "integer"}]}, {"enum": ["brief", None, 1], "anyOf": [{"type": "string"}, {"type": "null"}]}, {"const": None, "anyOf": [{"type": "string"}, {"type": "null"}]}, ], ) def test_scalar_choices_intersect_compositions_in_validation_and_emitted_schema(self, prop_schema, required): schema = { "type": "object", "properties": {"value": prop_schema}, "required": ["value"] if required else [], "$defs": {"Integer": {"type": "integer"}}, } model = SchemaTransformer.create_model_from_schema(schema, "ComposedChoiceInput", "choose") advertised = Draft202012Validator(schema) emitted = Draft202012Validator(model.model_json_schema()) for value in [None, True, False, 0, 1, 1.0, 2, 2.5, 3, "brief", "bad", "unknown", [], {}]: payload = {"tool_name": "choose", "value": value} # Optional fields retain the transformer's existing local None allowance. allowed = advertised.is_valid(payload) or (not required and value is None) assert emitted.is_valid(payload) == allowed if allowed: assert model.model_validate(payload).value == value else: with pytest.raises(ValidationError): model.model_validate(payload) @pytest.mark.parametrize( "prop_schema", [ {"const": "bad", "anyOf": [{"type": "integer"}]}, {"const": 1, "oneOf": [{"type": "integer"}, {"type": "number"}]}, {"enum": [1], "allOf": [{"minimum": 2}]}, ], ) def test_composition_with_no_admissible_scalar_choice_rejects_definition(self, prop_schema): schema = {"type": "object", "properties": {"value": prop_schema}, "required": ["value"]} with pytest.raises(ValueError, match="choices"): SchemaTransformer.create_model_from_schema(schema, "ImpossibleChoiceInput", "choose") @pytest.mark.parametrize("ref_first", [True, False]) def test_scalar_choice_composition_unresolved_reference_keeps_known_sibling_constraint(self, ref_first): branch = {"$ref": "https://example.invalid/schema.json", "type": "integer"} if not ref_first: branch = dict(reversed(branch.items())) schema = { "type": "object", "properties": {"value": {"enum": [1, "bad"], "anyOf": [branch]}}, "required": ["value"], } model = SchemaTransformer.create_model_from_schema(schema, "UnresolvedChoiceInput", "choose") assert model(tool_name="choose", value=1).value == 1 with pytest.raises(ValidationError): model(tool_name="choose", value="bad") assert not Draft202012Validator(model.model_json_schema()).is_valid({"tool_name": "choose", "value": "bad"}) def test_scalar_choice_composition_unresolved_reference_alone_preserves_fallback(self): schema = { "type": "object", "properties": {"value": {"enum": [1, "brief"], "anyOf": [{"$ref": "#/$defs/Missing"}]}}, "required": ["value"], } model = SchemaTransformer.create_model_from_schema(schema, "UnresolvedChoiceInput", "choose") assert model(tool_name="choose", value=1).value == 1 assert model(tool_name="choose", value="brief").value == "brief" with pytest.raises(ValidationError): model(tool_name="choose", value="other") @pytest.mark.parametrize( "prop_schema, invalid", [ ({"type": ["boolean", "integer"], "enum": [True]}, 1), ({"type": ["boolean", "integer"], "enum": [1]}, True), ({"type": ["boolean", "number"], "const": False}, 0.0), ({"type": ["boolean", "number"], "const": 0}, False), ({"enum": [True]}, 1.0), ({"enum": [0]}, False), ({"const": True}, 1), ({"const": 1}, True), ], ) def test_scalar_choices_preserve_json_identity(self, prop_schema, invalid): schema = {"type": "object", "properties": {"value": prop_schema}, "required": ["value"]} model = SchemaTransformer.create_model_from_schema(schema, "ChoiceInput", "choose") payload = {"tool_name": "choose", "value": invalid} assert not Draft202012Validator(schema).is_valid(payload) with pytest.raises(ValidationError): model.model_validate(payload) assert not Draft202012Validator(model.model_json_schema()).is_valid(payload) @pytest.mark.parametrize("choices", [[True, 1], [1, True], [False, 0], [0, False]]) def test_mixed_boolean_and_numeric_choices_keep_the_matching_choice(self, choices): schema = {"type": "object", "properties": {"value": {"enum": choices}}, "required": ["value"]} model = SchemaTransformer.create_model_from_schema(schema, "MixedChoiceInput", "choose") numeric = next(value for value in choices if not isinstance(value, bool)) assert type(model(tool_name="choose", value=float(numeric)).value) is int boolean = next(value for value in choices if isinstance(value, bool)) assert model(tool_name="choose", value=boolean).value is boolean @pytest.mark.parametrize("prop_schema", [{"type": "number", "enum": [1]}, {"const": 1.0}, {"enum": [0, 1]}]) def test_scalar_choices_retain_json_numeric_equality(self, prop_schema): schema = {"type": "object", "properties": {"value": prop_schema}, "required": ["value"]} model = SchemaTransformer.create_model_from_schema(schema, "NumericChoiceInput", "choose") for value in (1, 1.0): payload = {"tool_name": "choose", "value": value} assert Draft202012Validator(schema).is_valid(payload) assert model.model_validate(payload).value == value assert Draft202012Validator(model.model_json_schema()).is_valid(payload) @pytest.mark.parametrize( "prop_schema", [ {"type": "string", "enum": ["brief", "detailed"]}, {"type": "string", "const": "brief"}, {"type": ["string", "null"], "enum": ["brief", None]}, {"type": "number", "enum": [1, 2]}, ], ) def test_optional_scalar_choice_schema_matches_none_validation(self, prop_schema): schema = {"type": "object", "properties": {"value": prop_schema}} model = SchemaTransformer.create_model_from_schema(schema, "OptionalChoiceInput", "choose") payload = {"tool_name": "choose", "value": None} assert model.model_validate(payload).value is None emitted = model.model_json_schema() assert Draft202012Validator(emitted).is_valid(payload) assert Draft202012Validator(emitted).is_valid({"tool_name": "choose"}) assert not Draft202012Validator(emitted).is_valid({"tool_name": "choose", "value": "unknown"}) @pytest.mark.parametrize( "prop_schema, valid, invalid, emitted_type, choices", [ ({"type": "string", "enum": ["brief", None]}, "brief", None, "string", ["brief"]), ({"type": "integer", "enum": [1, "two"]}, 1, "two", "integer", [1]), ({"type": "integer", "enum": [1, True, 2.5]}, 1, True, "integer", [1]), ({"type": "number", "enum": [1, 2.5, False]}, 2.5, False, "number", [1, 2.5]), ({"type": "boolean", "enum": [False, 0, "false"]}, False, 0, "boolean", [False]), ({"type": "integer", "enum": [1.0, 2.5]}, 1, 2.5, "integer", [1.0]), ({"type": "boolean", "const": False}, False, 0, "boolean", [False]), ], ) def test_typed_scalar_choices(self, prop_schema, valid, invalid, emitted_type, choices): """Choice constraints intersect the declared JSON type, including bool/number boundaries.""" schema = {"type": "object", "properties": {"mode": prop_schema}, "required": ["mode"]} model = SchemaTransformer.create_model_from_schema(schema, "TypedModeInput", "configure") assert model(tool_name="configure", mode=valid).mode == valid with pytest.raises(ValueError): model(tool_name="configure", mode=invalid) emitted = model.model_json_schema()["properties"]["mode"] assert emitted["type"] == emitted_type assert emitted.get("enum", [emitted.get("const")]) == choices @pytest.mark.parametrize("prop_schema", [{"type": "string", "const": None}, {"type": "integer", "enum": ["two"]}]) def test_incompatible_scalar_choices_reject_definition(self, prop_schema): """An impossible type/choice combination never becomes an unrestricted input field.""" schema = {"type": "object", "properties": {"mode": prop_schema}, "required": ["mode"]} with pytest.raises(ValueError, match="choices.*declared"): SchemaTransformer.create_model_from_schema(schema, "InvalidModeInput", "configure") def test_nullable_type_array_with_mixed_enum(self): schema = { "type": "object", "properties": {"mode": {"type": ["string", "null"], "enum": ["brief", None, 1]}}, "required": ["mode"], } model = SchemaTransformer.create_model_from_schema(schema, "NullableModeInput", "configure") assert model(tool_name="configure", mode="brief").mode == "brief" assert model(tool_name="configure", mode=None).mode is None with pytest.raises(ValueError): model(tool_name="configure", mode=1) assert model.model_json_schema()["properties"]["mode"]["enum"] == ["brief", None] @pytest.mark.parametrize( "prop_schema, valid, invalid", [ ({"type": "string", "enum": ["brief", "detailed"]}, "brief", "unknown"), ({"type": "integer", "enum": [1, 2]}, 1, 3), ({"enum": ["brief", None]}, None, "unknown"), ({"type": "string", "const": "brief"}, "brief", "detailed"), ({"type": "integer", "const": 0}, 0, 1), ({"type": "boolean", "const": False}, False, True), ({"const": None}, None, "unknown"), ], ) def test_scalar_value_constraints(self, prop_schema, valid, invalid): """MCP parameter choices constrain both validation and the model-facing schema.""" schema = {"type": "object", "properties": {"mode": prop_schema}, "required": ["mode"]} model = SchemaTransformer.create_model_from_schema(schema, "ModeInput", "configure") assert model(tool_name="configure", mode=valid).mode == valid with pytest.raises(ValueError): model(tool_name="configure", mode=invalid) emitted = model.model_json_schema()["properties"]["mode"] choices = prop_schema.get("enum", [prop_schema.get("const")]) if len(choices) == 1: assert emitted["const"] == choices[0] else: assert emitted["enum"] == choices def test_array_item_enum(self): """Enum constraints survive recursive conversion of tool parameter arrays.""" schema = { "type": "object", "properties": {"modes": {"type": "array", "items": {"type": "string", "enum": ["brief", "detailed"]}}}, "required": ["modes"], } model = SchemaTransformer.create_model_from_schema(schema, "ModesInput", "configure") assert model(tool_name="configure", modes=["brief", "detailed"]).modes == ["brief", "detailed"] with pytest.raises(ValueError): model(tool_name="configure", modes=["unknown"]) assert model.model_json_schema()["properties"]["modes"]["items"]["enum"] == ["brief", "detailed"] @pytest.mark.parametrize("definitions_key", ["$defs", "definitions"]) def test_referenced_enum(self, definitions_key): """Enum definitions emitted by MCP servers remain scalar parameter choices.""" schema = { "type": "object", "properties": {"mode": {"$ref": f"#/{definitions_key}/Mode"}}, "required": ["mode"], definitions_key: {"Mode": {"type": "string", "enum": ["brief", "detailed"]}}, } model = SchemaTransformer.create_model_from_schema(schema, "ModeInput", "configure") assert model(tool_name="configure", mode="brief").mode == "brief" with pytest.raises(ValueError): model(tool_name="configure", mode="unknown") assert model.model_json_schema()["properties"]["mode"]["enum"] == ["brief", "detailed"] def test_nullable_enum_with_default(self): schema = { "type": "object", "properties": { "mode": {"anyOf": [{"type": "string", "enum": ["brief", "detailed"]}, {"type": "null"}], "default": "brief"} }, } model = SchemaTransformer.create_model_from_schema(schema, "ModeInput", "configure") assert model(tool_name="configure").mode == "brief" assert model(tool_name="configure", mode=None).mode is None with pytest.raises(ValueError): model(tool_name="configure", mode="unknown") @pytest.mark.parametrize("invalid_field", ["first", "second"]) def test_repeated_reference_reuses_model_and_validates_both_fields(self, invalid_field: str) -> None: reference = {"$ref": "#/$defs/Item"} schema = { "type": "object", "properties": {"first": reference, "second": reference}, "required": ["first", "second"], "$defs": {"Item": {"type": "object", "properties": {"count": {"type": "integer"}}, "required": ["count"]}}, } model = SchemaTransformer.create_model_from_schema(schema, "ItemPairInput", "load_pair") instance = model.model_validate({"tool_name": "load_pair", "first": {"count": 42}, "second": {"count": 7}}) assert model.model_fields["first"].annotation is model.model_fields["second"].annotation assert instance.model_dump() == {"tool_name": "load_pair", "first": {"count": 42}, "second": {"count": 7}} values: dict[str, str | dict[str, int | str]] = { "tool_name": "load_pair", "first": {"count": 42}, "second": {"count": 7}, } values[invalid_field] = {"count": "invalid"} with pytest.raises(ValidationError) as error: model.model_validate(values) assert error.value.errors()[0]["loc"] == (invalid_field, "count") def test_unresolved_reference_warns_and_preserves_value(self, caplog: pytest.LogCaptureFixture) -> None: schema = { "type": "object", "properties": {"payload": {"$ref": "#/$defs/MissingPayload"}}, "required": ["payload"], } with caplog.at_level(logging.WARNING, logger="atomic_agents.connectors.mcp.schema_transformer"): model = SchemaTransformer.create_model_from_schema(schema, "PayloadInput", "load_payload") assert ( "atomic_agents.connectors.mcp.schema_transformer", logging.WARNING, "Could not resolve $ref: #/$defs/MissingPayload", ) in caplog.record_tuples payload = {"count": "untyped", "items": [7, None]} assert model.model_validate({"tool_name": "load_payload", "payload": payload}).model_dump()["payload"] == payload @pytest.mark.parametrize("union_keyword", ["oneOf", "anyOf"]) def test_single_member_union_validates_integer(self, union_keyword: str) -> None: schema = { "type": "object", "properties": {"count": {union_keyword: [{"type": "integer"}]}}, "required": ["count"], } model = SchemaTransformer.create_model_from_schema(schema, "CountInput", "count_items") assert model.model_validate({"tool_name": "count_items", "count": 42}).model_dump()["count"] == 42 with pytest.raises(ValidationError): model.model_validate({"tool_name": "count_items", "count": "invalid"}) @pytest.mark.parametrize("definitions_key", ["$defs", "definitions"]) @pytest.mark.parametrize( "name, reference_token", [("Item", "Item"), ("ns/Item", "ns~1Item"), ("ns~Item", "ns~0Item"), ("ns~1Item", "ns~01Item")], ) def test_ref_names_with_json_pointer_escapes(self, definitions_key, name, reference_token): """Escaped definition names retain nested types and validation for MCP tools.""" reference = {"$ref": f"#/{definitions_key}/{reference_token}"} schema = { "type": "object", "properties": {"item": reference, "items": {"type": "array", "items": reference}}, "required": ["item", "items"], definitions_key: {name: {"type": "object", "properties": {"count": {"type": "integer"}}, "required": ["count"]}}, } model = SchemaTransformer.create_model_from_schema(schema, "ItemsInput", "load_items") instance = model(tool_name="load_items", item={"count": 42}, items=[{"count": 7}]) assert instance.item.count == 42 assert instance.items[0].count == 7 with pytest.raises(ValueError): model(tool_name="load_items", item={"count": "invalid"}, items=[]) with pytest.raises(ValueError): model(tool_name="load_items", item={"count": 42}, items=[{}]) @pytest.mark.parametrize( "prop_schema, valid_value", [ ({"type": "null"}, None), ({"type": ["string", "null"]}, "query"), ({"type": ["null", "string"]}, "query"), ({"anyOf": [{"type": "string"}, {"type": "null"}]}, "query"), ({"oneOf": [{"type": "string"}, {"type": "null"}]}, "query"), ], ) def test_required_nullable_field(self, prop_schema, valid_value): """A nullable MCP parameter still rejects other types and must be present.""" schema = {"type": "object", "properties": {"query": prop_schema}, "required": ["query"]} model = SchemaTransformer.create_model_from_schema(schema, "SearchInput", "search") assert model(tool_name="search", query=valid_value).query == valid_value assert model(tool_name="search", query=None).query is None with pytest.raises(ValueError): model(tool_name="search", query={"unexpected": "object"}) with pytest.raises(ValueError): model(tool_name="search") @pytest.mark.parametrize("types", [["string"], ["string", "integer"]]) def test_type_array_without_null(self, types): """Type arrays preserve the allowed types without implicitly allowing null.""" schema = {"type": "object", "properties": {"value": {"type": types}}, "required": ["value"]} model = SchemaTransformer.create_model_from_schema(schema, "ValueOutput", "value", is_output_schema=True) assert model(value="text").value == "text" if "integer" in types: assert model(value=42).value == 42 with pytest.raises(ValueError): model(value=None) with pytest.raises(ValueError): model(value={}) @pytest.mark.parametrize("default", [None, "fallback"]) def test_optional_nullable_field_default(self, default): """Converting a type array retains the property's default and description.""" schema = { "type": "object", "properties": {"query": {"type": ["string", "null"], "default": default, "description": "Search query"}}, } model = SchemaTransformer.create_model_from_schema(schema, "SearchInput", "search") assert model(tool_name="search").query == default assert model.model_fields["query"].description == "Search query" assert model(tool_name="search", query=None).query is None @pytest.mark.parametrize( "items, valid, invalid", [ ({"type": ["string", "null"]}, ["text", None], [{}]), ({"type": "null"}, [None], ["text"]), ({"type": "array", "items": {"type": ["string", "null"]}}, [["text", None]], [[{}]]), ], ) def test_nullable_array_items(self, items, valid, invalid): """Nullable array types retain their item types at each array depth.""" schema = { "type": "object", "properties": {"values": {"type": ["array", "null"], "items": items}}, "required": ["values"], } model = SchemaTransformer.create_model_from_schema(schema, "ValuesOutput", "values", is_output_schema=True) assert model(values=valid).values == valid assert model(values=None).values is None with pytest.raises(ValueError): model(values=invalid) def test_basic_model_creation(self): schema = { "type": "object", "properties": { "name": {"type": "string", "description": "A name"}, "age": {"type": "integer", "description": "An age"}, }, "required": ["name"], } model = SchemaTransformer.create_model_from_schema(schema, "TestModel", "test_tool") # Check the model structure assert issubclass(model, BaseIOSchema) assert model.__name__ == "TestModel" assert "tool_name" in model.model_fields assert "name" in model.model_fields assert "age" in model.model_fields # Test required vs optional fields assert model.model_fields["name"].is_required() is True assert model.model_fields["age"].is_required() is False # Test type annotations assert model.model_fields["name"].annotation == str assert model.model_fields["age"].annotation == Optional[int] # Test docstring assert model.__doc__ == "Dynamically generated Pydantic model for TestModel" def test_model_with_custom_docstring(self): schema = {"type": "object", "properties": {}} model = SchemaTransformer.create_model_from_schema(schema, "TestModel", "test_tool", docstring="Custom docstring") assert model.__doc__ == "Custom docstring" def test_empty_object_schema(self): schema = {"type": "object"} model = SchemaTransformer.create_model_from_schema(schema, "EmptyModel", "empty_tool") assert issubclass(model, BaseIOSchema) assert model.__name__ == "EmptyModel" assert "tool_name" in model.model_fields assert len(model.model_fields) == 1 # Only the tool_name field def test_non_object_schema(self, caplog): schema = {"type": "string"} model = SchemaTransformer.create_model_from_schema(schema, "StringModel", "string_tool") assert issubclass(model, BaseIOSchema) assert model.__name__ == "StringModel" assert "tool_name" in model.model_fields assert len(model.model_fields) == 1 # Only the tool_name field assert "Schema for StringModel is not a typical object with properties" in caplog.text def test_tool_name_field(self): schema = {"type": "object", "properties": {}} model = SchemaTransformer.create_model_from_schema(schema, "ToolModel", "specific_tool") # Test that tool_name is a Literal type with the correct value assert "tool_name" in model.model_fields tool_instance = model(tool_name="specific_tool") assert tool_instance.tool_name == "specific_tool" # Test that an invalid tool_name raises an error with pytest.raises(ValueError): model(tool_name="wrong_tool") def test_union_type_oneof(self): """Test oneOf creates Union types.""" prop_schema = {"oneOf": [{"type": "string"}, {"type": "integer"}], "description": "A union field"} result = SchemaTransformer.json_to_pydantic_field(prop_schema, True) # Should create Union[str, int] assert result[0] == Union[str, int] assert result[1].description == "A union field" def test_union_type_anyof(self): """Test anyOf creates Union types.""" prop_schema = {"anyOf": [{"type": "boolean"}, {"type": "number"}], "description": "Another union field"} result = SchemaTransformer.json_to_pydantic_field(prop_schema, True) # Should create Union[bool, float] assert result[0] == Union[bool, float] def test_array_with_ref_items(self): """Test arrays with $ref items are resolved.""" root_schema = { "$defs": {"MyObject": {"type": "object", "properties": {"name": {"type": "string"}}, "title": "MyObject"}} } prop_schema = {"type": "array", "items": {"$ref": "#/$defs/MyObject"}, "description": "Array of MyObject"} result = SchemaTransformer.json_to_pydantic_field(prop_schema, True, root_schema) # Should be List[MyObject] not List[Any] assert hasattr(result[0], "__origin__") and result[0].__origin__ is list # The inner type should be the created model, not Any inner_type = result[0].__args__[0] assert inner_type != Any assert hasattr(inner_type, "model_fields") def test_array_with_union_items(self): """Test arrays with oneOf items.""" prop_schema = { "type": "array", "items": {"oneOf": [{"type": "string"}, {"type": "integer"}]}, "description": "Array of union items", } result = SchemaTransformer.json_to_pydantic_field(prop_schema, True) # Should be List[Union[str, int]] assert hasattr(result[0], "__origin__") and result[0].__origin__ is list inner_type = result[0].__args__[0] assert inner_type == Union[str, int] def test_model_with_complex_types(self): """Test create_model_from_schema with complex types.""" schema = { "type": "object", "properties": { "expr": {"oneOf": [{"$ref": "#/$defs/ANode"}, {"$ref": "#/$defs/BNode"}], "description": "Expression node"}, "objects": {"type": "array", "items": {"$ref": "#/$defs/MyObject"}, "description": "List of objects"}, }, "required": ["expr", "objects"], "$defs": { "ANode": {"type": "object", "properties": {"a_value": {"type": "string"}}, "title": "ANode"}, "BNode": {"type": "object", "properties": {"b_value": {"type": "integer"}}, "title": "BNode"}, "MyObject": {"type": "object", "properties": {"name": {"type": "string"}}, "title": "MyObject"}, }, } model = SchemaTransformer.create_model_from_schema(schema, "ComplexModel", "complex_tool") # Check that expr is a Union, not Any expr_field = model.model_fields["expr"] assert expr_field.annotation != Any # Should be Union[ANode, BNode] assert hasattr(expr_field.annotation, "__origin__") and expr_field.annotation.__origin__ is Union # Check that objects is List[MyObject], not List[Any] objects_field = model.model_fields["objects"] assert objects_field.annotation != List[Any] assert hasattr(objects_field.annotation, "__origin__") and objects_field.annotation.__origin__ is list inner_type = objects_field.annotation.__args__[0] assert inner_type != Any def test_output_schema_no_tool_name_field(self): """Test that output schemas don't include tool_name field when is_output_schema=True.""" schema = { "type": "object", "properties": { "results": {"type": "array", "items": {"type": "string"}, "description": "Search results"}, "count": {"type": "integer", "description": "Number of results"}, }, "required": ["results", "count"], } model = SchemaTransformer.create_model_from_schema(schema, "OutputModel", "my_tool", is_output_schema=True) # Output schema should NOT have tool_name field assert "tool_name" not in model.model_fields # But should have the defined fields assert "results" in model.model_fields assert "count" in model.model_fields assert len(model.model_fields) == 2 # Only results and count, no tool_name # Should be instantiable without tool_name instance = model(results=["a", "b"], count=2) assert instance.results == ["a", "b"] assert instance.count == 2 def test_input_schema_has_tool_name_field(self): """Test that input schemas include tool_name field when is_output_schema=False (default).""" schema = { "type": "object", "properties": { "query": {"type": "string", "description": "Search query"}, }, "required": ["query"], } model = SchemaTransformer.create_model_from_schema(schema, "InputModel", "my_tool", is_output_schema=False) # Input schema SHOULD have tool_name field assert "tool_name" in model.model_fields assert "query" in model.model_fields assert len(model.model_fields) == 2 # query and tool_name # Should require tool_name for instantiation instance = model(tool_name="my_tool", query="test") assert instance.tool_name == "my_tool" assert instance.query == "test" def test_output_schema_with_resource_attribute_type(self): """Test that output schemas work with different attribute types.""" from atomic_agents.connectors.mcp.mcp_definition_service import MCPAttributeType schema = { "type": "object", "properties": { "data": {"type": "string", "description": "Some data"}, }, "required": ["data"], } # Output schema for resource - should not have resource_name model = SchemaTransformer.create_model_from_schema( schema, "ResourceOutput", "my_resource", attribute_type=MCPAttributeType.RESOURCE, is_output_schema=True ) assert "resource_name" not in model.model_fields assert "data" in model.model_fields ``` ### File: atomic-agents/tests/connectors/mcp/test_transport_endpoints.py ```python """MCP URL handling must agree between discovery and invocation.""" from unittest.mock import AsyncMock, MagicMock import pytest from mcp import ClientSession from atomic_agents.connectors.mcp import ( MCPDefinitionService, MCPFactory, MCPPromptDefinition, MCPResourceDefinition, MCPToolDefinition, MCPTransportType, ) @pytest.fixture(params=[MCPTransportType.HTTP_STREAM, MCPTransportType.SSE]) def endpoint_case(request): transport_type = request.param suffix = "/mcp/" if transport_type == MCPTransportType.HTTP_STREAM else "/sse" return transport_type, suffix @pytest.mark.asyncio @pytest.mark.parametrize("kind", ["tool", "resource", "prompt"]) @pytest.mark.parametrize("discovery", [True, False], ids=["discovery", "invocation"]) @pytest.mark.parametrize("url_kind", ["root", "mount", "complete", "complete_query", "query"]) async def test_transport_uses_valid_endpoint(monkeypatch, endpoint_case, kind, discovery, url_kind): transport_type, suffix = endpoint_case if url_kind == "root": endpoint, expected = "https://example.com", f"https://example.com{suffix}" elif url_kind == "mount": endpoint, expected = "https://example.com/api/", f"https://example.com/api{suffix}" elif url_kind == "complete": endpoint = expected = f"https://example.com/api{suffix.rstrip('/')}" elif url_kind == "complete_query": endpoint = expected = f"https://example.com/api{suffix.rstrip('/')}/?token=a%2Fb#section" else: endpoint = "https://example.com?token=a%2Fb#section" expected = f"https://example.com{suffix}?token=a%2Fb#section" module = "atomic_agents.connectors.mcp." + ("mcp_definition_service" if discovery else "mcp_factory") session = AsyncMock(spec=ClientSession) session.__aenter__.return_value = session session.call_tool.return_value = {"ok": True} session.read_resource.return_value = {"contents": []} session.get_prompt.return_value = {"messages": []} transport = AsyncMock() transport.__aenter__.return_value = (None, None, None) if transport_type == MCPTransportType.HTTP_STREAM else (None, None) client = MagicMock(return_value=transport) client_name = "streamablehttp_client" if transport_type == MCPTransportType.HTTP_STREAM else "sse_client" monkeypatch.setattr(f"{module}.{client_name}", client) monkeypatch.setattr(f"{module}.ClientSession", MagicMock(return_value=session)) if discovery: service = MCPDefinitionService(endpoint, transport_type) monkeypatch.setattr(service, f"fetch_{kind}_definitions_from_session", AsyncMock(return_value=[])) assert await getattr(service, f"fetch_{kind}_definitions")() == [] else: factory = MCPFactory(endpoint, transport_type) schema = {"type": "object", "properties": {}} if kind == "tool": definition = MCPToolDefinition("sample", "Example tool", schema) instance = factory._create_tool_classes([definition])[0]() await instance.arun(instance.input_schema(tool_name="sample")) elif kind == "resource": definition = MCPResourceDefinition("sample", "Example resource", "example://item", schema) instance = factory._create_resource_classes([definition])[0]() await instance.aread(instance.input_schema(resource_name="sample")) else: definition = MCPPromptDefinition("sample", "Example prompt", schema) instance = factory._create_prompt_classes([definition])[0]() await instance.agenerate(instance.input_schema(prompt_name="sample")) client.assert_called_once_with(expected) session.__aexit__.assert_awaited_once() transport.__aexit__.assert_awaited_once() ``` ### File: atomic-agents/tests/context/test_base_chat_history.py ```python from unittest.mock import Mock import instructor import pytest from pydantic import Field from atomic_agents import ( AgentConfig, AtomicAgent, BaseIOSchema, BasicChatInputSchema, BasicChatOutputSchema, ) from atomic_agents.context import BaseChatHistory, ChatHistory, Message, SystemPromptGenerator class InputSchema(BaseIOSchema): """Test Input Schema""" test_field: str = Field(..., description="A test field") def test_base_chat_history_cannot_be_instantiated_directly(): with pytest.raises(TypeError): BaseChatHistory() def test_chat_history_is_subclass_and_instance_of_base(): assert issubclass(ChatHistory, BaseChatHistory) assert isinstance(ChatHistory(), BaseChatHistory) def test_subclass_missing_abstractmethod_cannot_be_instantiated(): class IncompleteHistory(BaseChatHistory): """A subclass that forgets to implement `copy`.""" def initialize_turn(self) -> None: pass def add_message(self, role, content): pass def get_history(self): return [] def get_current_turn_id(self): return None def delete_turn_id(self, turn_id): pass def get_message_count(self): return 0 def dump(self): return "{}" def load(self, serialized_data): pass # `copy` intentionally omitted with pytest.raises(TypeError): IncompleteHistory() class RecordingChatHistory(ChatHistory): """A custom persistent-memory-style backend used to exercise the extension seam.""" def __init__(self, max_messages=None): super().__init__(max_messages=max_messages) self.recorded_calls = 0 def add_message(self, role, content): self.recorded_calls += 1 super().add_message(role, content) @pytest.fixture def mock_instructor(): mock = Mock(spec=instructor.Instructor) mock.chat = Mock() mock.chat.completions = Mock() mock.chat.completions.create = Mock(return_value=BasicChatOutputSchema(chat_message="Test output")) return mock def test_custom_chat_history_subclass_used_by_agent(mock_instructor): custom_history = RecordingChatHistory() config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=custom_history, system_prompt_generator=SystemPromptGenerator(), ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) assert isinstance(agent.history, BaseChatHistory) result = agent.run(BasicChatInputSchema(chat_message="Hello")) assert result.chat_message == "Test output" # One call for the user message, one for the assistant response. assert custom_history.recorded_calls == 2 assert custom_history.get_message_count() == 2 class MinimalHistory(BaseChatHistory): """A from-scratch BaseChatHistory implementation that does NOT subclass ChatHistory. Exercises the "own everything" path the docs describe, to prove a backend built purely against the ABC drives AtomicAgent end to end. """ def __init__(self, max_messages=None): self.history = [] self.current_turn_id = None self.max_messages = max_messages def initialize_turn(self): self.current_turn_id = "turn" def add_message(self, role, content): if self.current_turn_id is None: self.initialize_turn() self.history.append(Message(role=role, content=content, turn_id=self.current_turn_id)) def get_history(self): return [{"role": m.role, "content": m.content.model_dump_json()} for m in self.history] def get_current_turn_id(self): return self.current_turn_id def delete_turn_id(self, turn_id): self.history = [m for m in self.history if m.turn_id != turn_id] def get_message_count(self): return len(self.history) def dump(self): return "{}" def load(self, serialized_data): pass def copy(self): new = MinimalHistory(max_messages=self.max_messages) new.history = list(self.history) new.current_turn_id = self.current_turn_id return new def test_from_scratch_base_history_drives_agent(mock_instructor): history = MinimalHistory() config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=history, system_prompt_generator=SystemPromptGenerator(), ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) assert isinstance(agent.history, MinimalHistory) result = agent.run(BasicChatInputSchema(chat_message="Hello")) assert result.chat_message == "Test output" assert agent.history.get_message_count() == 2 class StatefulHistory(ChatHistory): """A ChatHistory subclass carrying extra state, with copy() overridden to preserve it.""" def __init__(self, marker, max_messages=None): super().__init__(max_messages=max_messages) self.marker = marker def copy(self): new = StatefulHistory(self.marker, max_messages=self.max_messages) new.load(self.dump()) new.current_turn_id = self.current_turn_id return new def test_reset_history_preserves_backend_when_copy_overridden(mock_instructor): history = StatefulHistory(marker="db-handle") config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=history, system_prompt_generator=SystemPromptGenerator(), ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) # The initial_history snapshot keeps the subclass type and its state. assert isinstance(agent.initial_history, StatefulHistory) assert agent.initial_history.marker == "db-handle" agent.run(BasicChatInputSchema(chat_message="Hi")) agent.reset_history() # After reset, the backend is still the custom type, not a plain ChatHistory. assert isinstance(agent.history, StatefulHistory) assert agent.history.marker == "db-handle" def test_subclass_without_copy_override_downgrades_to_chat_history(): """Pins the known sharp edge: an un-overridden copy() returns a plain ChatHistory. This is exactly why subclasses carrying extra state must override copy(); documented in the Memory guide and warned about in BaseChatHistory.copy's docstring. """ class NaiveHistory(ChatHistory): pass copied = NaiveHistory().copy() assert type(copied) is ChatHistory def test_falsy_custom_history_is_not_discarded(mock_instructor): """AgentConfig must keep a provided history even if it is falsy (e.g. defines __len__).""" class FalsyWhenEmpty(ChatHistory): def __len__(self): return self.get_message_count() history = FalsyWhenEmpty() assert not history # empty -> falsy via __len__ config = AgentConfig( client=mock_instructor, model="gpt-5-mini", history=history, system_prompt_generator=SystemPromptGenerator(), ) agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) assert agent.history is history ``` ### File: atomic-agents/tests/context/test_chat_history.py ```python from enum import Enum, IntEnum from unittest.mock import Mock, call import pytest import json from typing import List, Dict, Union from pathlib import Path from pydantic import Field, ValidationError from atomic_agents.context import ChatHistory, Message from atomic_agents import BaseIOSchema, VideoURL import instructor class TupleAttachments(BaseIOSchema): """A message with typed attachments held in tuples.""" attachments: tuple[instructor.Image, instructor.processing.multimodal.PDF] nested: dict[str, tuple[tuple[instructor.Image, ...], ...]] class InputSchema(BaseIOSchema): """Test Input Schema""" test_field: str = Field(..., description="A test field") class MockOutputSchema(BaseIOSchema): """Test Output Schema""" test_field: str = Field(..., description="A test field") class MockNestedSchema(BaseIOSchema): """Test Nested Schema""" nested_field: str = Field(..., description="A nested field") nested_int: int = Field(..., description="A nested integer") class MockComplexInputSchema(BaseIOSchema): """Test Complex Input Schema""" text_field: str = Field(..., description="A text field") number_field: float = Field(..., description="A number field") list_field: List[str] = Field(..., description="A list of strings") nested_field: MockNestedSchema = Field(..., description="A nested schema") class MockComplexOutputSchema(BaseIOSchema): """Test Complex Output Schema""" response_text: str = Field(..., description="A response text") calculated_value: int = Field(..., description="A calculated value") data_dict: Dict[str, MockNestedSchema] = Field(..., description="A dictionary of nested schemas") class MockMultimodalSchema(BaseIOSchema): """Test schema for multimodal content""" instruction_text: str = Field(..., description="The instruction text") images: List[instructor.Image] = Field(..., description="The images to analyze") pdfs: List[instructor.processing.multimodal.PDF] = Field(..., description="The PDFs to analyze") audio: instructor.processing.multimodal.Audio = Field(..., description="The audio to analyze") class MockImageSchema(BaseIOSchema): """Test schema for serialized image paths.""" image: instructor.Image = Field(..., description="An image with a local path") class ColorEnum(str, Enum): BLUE = "blue" RED = "red" class MockEnumSchema(BaseIOSchema): """Test Input Schema with Enum.""" color: ColorEnum = Field(..., description="Some color.") @pytest.fixture def history(): return ChatHistory(max_messages=5) def test_initialization(history): assert history.history == [] assert history.max_messages == 5 assert history.current_turn_id is None def test_initialize_turn(history): history.initialize_turn() assert history.current_turn_id is not None def test_add_message(history): history.add_message("user", InputSchema(test_field="Hello")) assert len(history.history) == 1 assert history.history[0].role == "user" assert isinstance(history.history[0].content, InputSchema) assert history.history[0].turn_id is not None def test_manage_overflow(history): for i in range(7): history.add_message("user", InputSchema(test_field=f"Message {i}")) assert len(history.history) == 5 assert history.history[0].content.test_field == "Message 2" def test_get_history(history): """ Ensure non-ASCII characters are serialized without Unicode escaping, because it can cause issue with some OpenAI models like GPT-4.1. Reference ticket: https://github.com/eigenwise/atomic-agents/issues/138. """ history.add_message("user", InputSchema(test_field="Hello")) history.add_message("assistant", MockOutputSchema(test_field="Hi there")) history = history.get_history() assert len(history) == 2 assert history[0]["role"] == "user" assert json.loads(history[0]["content"]) == {"test_field": "Hello"} assert json.loads(history[1]["content"]) == {"test_field": "Hi there"} def test_get_history_allow_unicode(history): history.add_message("user", InputSchema(test_field="àéèï")) history.add_message("assistant", MockOutputSchema(test_field="â")) history = history.get_history() assert len(history) == 2 assert history[0]["role"] == "user" assert history[0]["content"] == '{"test_field":"àéèï"}' assert history[1]["content"] == '{"test_field":"â"}' assert json.loads(history[0]["content"]) == {"test_field": "àéèï"} assert json.loads(history[1]["content"]) == {"test_field": "â"} @pytest.mark.parametrize( "media", [ instructor.Image(source="private.jpg", media_type="image/jpeg"), instructor.processing.multimodal.Audio(source="private.mp3", media_type="audio/mp3"), instructor.processing.multimodal.PDF(source="private.pdf", media_type="application/pdf"), VideoURL(url="https://example.com/private.mp4"), ], ids=["image", "audio", "pdf", "video"], ) def test_get_history_omits_excluded_multimodal_fields(history, media): class PrivateMediaSchema(BaseIOSchema): """Message carrying media that must stay local.""" instruction: str private_media: Union[ instructor.Image, instructor.processing.multimodal.Audio, instructor.processing.multimodal.PDF, VideoURL ] = Field(exclude=True) content = PrivateMediaSchema(instruction="public", private_media=media) history.add_message("user", content) assert history.get_history() == [{"role": "user", "content": content.model_dump_json()}] def test_get_history_preserves_visible_media_beside_nested_excluded_media(history): class Document(BaseIOSchema): """Document with a public image and a local source.""" caption: str image: instructor.Image original: instructor.Image = Field(exclude=True) class DocumentsSchema(BaseIOSchema): """Documents provided to the model.""" documents: List[Document] visible = instructor.Image(source="visible.jpg", media_type="image/jpeg") private = instructor.Image(source="private.jpg", media_type="image/jpeg") content = DocumentsSchema(documents=[Document(caption="public", image=visible, original=private)]) history.add_message("user", content) parts = history.get_history()[0]["content"] assert json.loads(parts[0]) == {"documents": [{"caption": "public"}]} assert parts[1:] == [visible] @pytest.mark.parametrize("container", [list, tuple, lambda values: {"media": values[0]}], ids=["list", "tuple", "dict"]) def test_get_history_skips_excluded_containers(history, container): class PrivateContainerSchema(BaseIOSchema): """Message with an excluded collection of source media.""" image: instructor.Image originals: object = Field(exclude=True) visible = instructor.Image(source="visible.jpg", media_type="image/jpeg") private = instructor.Image(source="private.jpg", media_type="image/jpeg") history.add_message("user", PrivateContainerSchema(image=visible, originals=container([private]))) assert history.get_history() == [{"role": "user", "content": [visible]}] def test_copy(history): history.add_message("user", InputSchema(test_field="Hello")) copied_history = history.copy() assert copied_history.max_messages == history.max_messages assert copied_history.current_turn_id == history.current_turn_id assert len(copied_history.history) == len(history.history) assert copied_history.history[0].role == history.history[0].role assert copied_history.history[0].content.test_field == history.history[0].content.test_field def test_copy_with_local_schema_is_independent(history): class LocalSchema(BaseIOSchema): """Message schema defined by an application factory.""" items: List[Dict[str, List[str]]] history.add_message("user", LocalSchema(items=[{"values": ["original"]}])) copied_history = history.copy() assert copied_history.max_messages == history.max_messages assert copied_history.current_turn_id == history.current_turn_id assert copied_history.get_history() == history.get_history() assert isinstance(copied_history.history[0].content, LocalSchema) copied_history.history[0].content.items[0]["values"].append("changed") copied_history.history[0].role = "assistant" copied_history.add_message("user", LocalSchema(items=[])) assert history.history[0].content.items == [{"values": ["original"]}] assert history.history[0].role == "user" assert history.get_message_count() == 1 def test_get_current_turn_id(history): assert history.get_current_turn_id() is None history.initialize_turn() assert history.get_current_turn_id() is not None def test_get_message_count(history): assert history.get_message_count() == 0 history.add_message("user", InputSchema(test_field="Hello")) assert history.get_message_count() == 1 def test_dump_and_load_comprehensive(history): """Comprehensive test for dump/load functionality with complex nested data""" # Test complex nested schemas history.add_message( "user", MockComplexInputSchema( text_field="Complex input", number_field=2.718, list_field=["a", "b", "c"], nested_field=MockNestedSchema(nested_field="Nested input", nested_int=99), ), ) history.add_message( "assistant", MockComplexOutputSchema( response_text="Complex output", calculated_value=200, data_dict={ "key1": MockNestedSchema(nested_field="Nested output 1", nested_int=10), "key2": MockNestedSchema(nested_field="Nested output 2", nested_int=20), }, ), ) # Test get_history format with nested models history_output = history.get_history() assert len(history_output) == 2 assert history_output[0]["role"] == "user" assert history_output[1]["role"] == "assistant" expected_input_content = ( '{"text_field":"Complex input","number_field":2.718,"list_field":["a","b","c"],' '"nested_field":{"nested_field":"Nested input","nested_int":99}}' ) expected_output_content = ( '{"response_text":"Complex output","calculated_value":200,' '"data_dict":{"key1":{"nested_field":"Nested output 1","nested_int":10},' '"key2":{"nested_field":"Nested output 2","nested_int":20}}}' ) assert history_output[0]["content"] == expected_input_content assert history_output[1]["content"] == expected_output_content # Test dump and load dumped_data = history.dump() new_history = ChatHistory() new_history.load(dumped_data) # Verify all properties are preserved assert new_history.max_messages == history.max_messages assert new_history.current_turn_id == history.current_turn_id assert len(new_history.history) == len(history.history) assert isinstance(new_history.history[0].content, MockComplexInputSchema) assert isinstance(new_history.history[1].content, MockComplexOutputSchema) # Verify detailed content assert new_history.history[0].content.text_field == "Complex input" assert new_history.history[0].content.nested_field.nested_int == 99 assert new_history.history[1].content.response_text == "Complex output" assert new_history.history[1].content.data_dict["key1"].nested_field == "Nested output 1" # Test adding new messages to loaded history still works new_history.add_message("user", InputSchema(test_field="New message")) assert len(new_history.history) == 3 assert new_history.history[2].content.test_field == "New message" def test_dump_and_load_multimodal_data(history): import os base_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) test_image = instructor.Image.from_path(path=os.path.join(base_path, "files/image_sample.jpg")) test_pdf = instructor.processing.multimodal.PDF.from_path(path=os.path.join(base_path, "files/pdf_sample.pdf")) test_audio = instructor.processing.multimodal.Audio.from_path(path=os.path.join(base_path, "files/audio_sample.mp3")) # multimodal message history.add_message( role="user", content=MockMultimodalSchema( instruction_text="Analyze this image", images=[test_image], pdfs=[test_pdf], audio=test_audio ), ) dumped_data = history.dump() new_history = ChatHistory() new_history.load(dumped_data) assert new_history.max_messages == history.max_messages assert new_history.current_turn_id == history.current_turn_id assert len(new_history.history) == len(history.history) assert isinstance(new_history.history[0].content, MockMultimodalSchema) assert new_history.history[0].content.instruction_text == history.history[0].content.instruction_text assert new_history.history[0].content.images == history.history[0].content.images assert new_history.history[0].content.pdfs == history.history[0].content.pdfs assert new_history.history[0].content.audio == history.history[0].content.audio def test_dump_and_load_with_enum(history): """Test that get_history works with Enum.""" history.add_message( "user", MockEnumSchema( color=ColorEnum.RED, ), ) dumped_data = history.dump() new_history = ChatHistory() new_history.load(dumped_data) assert new_history.max_messages == history.max_messages assert new_history.current_turn_id == history.current_turn_id assert len(new_history.history) == len(history.history) def test_load_invalid_data(history): with pytest.raises(ValueError): history.load("invalid json") MISSING = object() @pytest.fixture def saved_history(history: ChatHistory) -> tuple[ChatHistory, List[Message], Message, str]: history.add_message("user", InputSchema(test_field="Keep this message")) return history, history.history, history.history[0], history.dump() @pytest.fixture def replacement_history() -> ChatHistory: replacement = ChatHistory(max_messages=2) replacement.current_turn_id = "replacement-turn" replacement.add_message("user", InputSchema(test_field="Replacement user message")) replacement.add_message("assistant", MockOutputSchema(test_field="Replacement assistant message")) return replacement def _assert_history_unchanged( history: ChatHistory, original_history: List[Message], original_message: Message, original_dump: str ) -> None: original_data = json.loads(original_dump) assert history.history is original_history assert history.history[0] is original_message assert history.dump() == original_dump assert history.max_messages == original_data["max_messages"] assert history.current_turn_id == original_data["current_turn_id"] def test_load_invalid_json_preserves_existing_history(saved_history) -> None: history = saved_history[0] with pytest.raises(ValueError, match="^Invalid serialized data: ") as error: history.load("invalid json") assert type(error.value) is ValueError _assert_history_unchanged(*saved_history) @pytest.mark.parametrize( "path, value, error_type, error_match", [ pytest.param(("max_messages",), MISSING, ValueError, "'max_messages'", id="missing_max_messages"), pytest.param(("current_turn_id",), MISSING, ValueError, "'current_turn_id'", id="missing_current_turn_id"), pytest.param(("history",), MISSING, ValueError, "'history'", id="missing_history"), pytest.param(("max_messages",), -1, ValueError, "^Invalid serialized data: max_messages", id="negative_capacity"), pytest.param(("max_messages",), True, ValueError, "^Invalid serialized data: max_messages", id="boolean_capacity"), pytest.param(("max_messages",), "4", ValueError, "^Invalid serialized data: max_messages", id="string_capacity"), pytest.param(("max_messages",), 2.5, ValueError, "^Invalid serialized data: max_messages", id="fractional_capacity"), pytest.param(("current_turn_id",), {}, ValueError, "current_turn_id", id="invalid_current_turn"), pytest.param(("history",), {}, ValueError, "history must be a list", id="empty_mapping"), pytest.param(("history", 0, "content", "class_name"), "builtins.str", ValueError, "BaseIOSchema", id="non_schema"), pytest.param(("history",), None, ValueError, "^Invalid serialized data: ", id="invalid_history"), pytest.param(("history", 0, "content"), MISSING, ValueError, "'content'", id="missing_content"), pytest.param( ("history", 0, "content", "class_name"), f"{InputSchema.__module__}.NonexistentSchema", ValueError, "NonexistentSchema", id="unknown_class", ), pytest.param( ("history", 0, "content", "data"), "invalid json", ValidationError, "Invalid JSON", id="invalid_content_json" ), pytest.param(("history", 0, "content", "data"), "{}", ValidationError, "test_field", id="invalid_content"), pytest.param(("history", 1, "content", "data"), "{}", ValidationError, "test_field", id="later_invalid_content"), pytest.param(("history", 0, "role"), None, ValidationError, "role", id="invalid_role"), pytest.param(("history", 0, "turn_id"), [], ValidationError, "turn_id", id="invalid_turn_id"), pytest.param(("history", 1, "role"), None, ValidationError, "role", id="later_invalid_role"), pytest.param(("history", 1, "turn_id"), [], ValidationError, "turn_id", id="later_invalid_turn_id"), pytest.param( ("history", 1, "content", "class_name"), "nonexistent_chat_history_module.InputSchema", ModuleNotFoundError, "nonexistent_chat_history_module", id="later_missing_module", ), ], ) def test_load_failure_preserves_existing_history( saved_history, replacement_history, path, value, error_type, error_match ) -> None: history = saved_history[0] replacement_data = json.loads(replacement_history.dump()) target = replacement_data for field in path[:-1]: target = target[field] if value is MISSING: del target[path[-1]] else: target[path[-1]] = value with pytest.raises(error_type, match=error_match) as error: history.load(json.dumps(replacement_data)) assert type(error.value) is error_type _assert_history_unchanged(*saved_history) @pytest.mark.parametrize( "normalization_error_type, error_type, error_match", [ pytest.param(TypeError, ValueError, "^Invalid serialized data: late normalization failure$", id="wrapped_type_error"), pytest.param(RuntimeError, RuntimeError, "^late normalization failure$", id="unwrapped_runtime_error"), ], ) def test_load_late_normalization_failure_preserves_existing_history( saved_history, monkeypatch, normalization_error_type, error_type, error_match ) -> None: history = saved_history[0] replacement = ChatHistory(max_messages=2) replacement.current_turn_id = "replacement-turn" for source in ("test/first.jpg", "test/late.jpg"): replacement.add_message("user", MockImageSchema(image=instructor.Image(source=source, media_type="image/jpeg"))) normalize_path = Mock(side_effect=[Path("test/first.jpg"), normalization_error_type("late normalization failure")]) monkeypatch.setattr("atomic_agents.context.chat_history.Path", normalize_path) with pytest.raises(error_type, match=error_match) as error: history.load(replacement.dump()) assert type(error.value) is error_type assert normalize_path.call_args_list == [call("test/first.jpg"), call("test/late.jpg")] _assert_history_unchanged(*saved_history) def test_load_replaces_existing_history(history: ChatHistory, replacement_history: ChatHistory) -> None: history.add_message("user", InputSchema(test_field="Old message")) original_history = history.history history.load(replacement_history.dump()) assert history.history is not original_history assert history.dump() == replacement_history.dump() assert isinstance(history.history[0].content, InputSchema) assert isinstance(history.history[1].content, MockOutputSchema) assert history.max_messages == 2 assert history.current_turn_id == "replacement-turn" @pytest.mark.parametrize("max_messages, current_turn_id", [(None, None), (0, "empty-turn")]) def test_load_empty_history_replaces_existing_history(history, max_messages, current_turn_id): history.add_message("user", InputSchema(test_field="Old message")) replacement = ChatHistory(max_messages=max_messages) replacement.current_turn_id = current_turn_id history.load(replacement.dump()) assert history.history == [] assert history.dump() == replacement.dump() assert history.max_messages == max_messages assert history.current_turn_id == current_turn_id def test_get_class_from_string(): class_string = "tests.context.test_chat_history.InputSchema" cls = ChatHistory._get_class_from_string(class_string) assert cls.__name__ == InputSchema.__name__ assert cls.__module__.endswith("test_chat_history") assert issubclass(cls, BaseIOSchema) def test_get_class_from_string_invalid(): with pytest.raises((ImportError, AttributeError)): ChatHistory._get_class_from_string("invalid.module.Class") def test_message_model(): message = Message(role="user", content=InputSchema(test_field="Test"), turn_id="123") assert message.role == "user" assert isinstance(message.content, InputSchema) assert message.turn_id == "123" def test_history_with_no_max_messages(): unlimited_history = ChatHistory() for i in range(100): unlimited_history.add_message("user", InputSchema(test_field=f"Message {i}")) assert len(unlimited_history.history) == 100 @pytest.mark.parametrize("capacity", [-1, True, "4", 2.5]) def test_history_rejects_invalid_capacity(capacity): with pytest.raises(ValueError, match="max_messages"): ChatHistory(max_messages=capacity) def test_history_accepts_int_subclass_capacity(): class Capacity(IntEnum): ONE = 1 history = ChatHistory(max_messages=Capacity.ONE) history.add_message("user", InputSchema(test_field="old")) history.add_message("user", InputSchema(test_field="new")) assert [message.content.test_field for message in history.history] == ["new"] restored = ChatHistory() restored.load(history.dump()) assert restored.max_messages == Capacity.ONE assert [message.content.test_field for message in restored.history] == ["new"] @pytest.mark.parametrize( "capacity, expected", [(None, ["first", "second", "third"]), (0, []), (1, ["third"]), (2, ["second", "third"])] ) def test_load_trims_history_after_capacity_is_lowered(capacity, expected): original = ChatHistory() for text in ["first", "second", "third"]: original.add_message("user", InputSchema(test_field=text)) original.max_messages = capacity restored = ChatHistory() restored.load(original.dump()) assert restored.max_messages == capacity assert [message.content.test_field for message in restored.history] == expected def test_history_with_zero_max_messages(): zero_max_history = ChatHistory(max_messages=0) for i in range(10): zero_max_history.add_message("user", InputSchema(test_field=f"Message {i}")) assert len(zero_max_history.history) == 0 def test_history_turn_consistency(): history = ChatHistory() history.initialize_turn() turn_id = history.get_current_turn_id() history.add_message("user", InputSchema(test_field="Hello")) history.add_message("assistant", MockOutputSchema(test_field="Hi")) assert history.history[0].turn_id == turn_id assert history.history[1].turn_id == turn_id history.initialize_turn() new_turn_id = history.get_current_turn_id() assert new_turn_id != turn_id history.add_message("user", InputSchema(test_field="Next turn")) assert history.history[2].turn_id == new_turn_id def test_chat_history_delete_turn_id(history): mock_input = InputSchema(test_field="Test input") mock_output = InputSchema(test_field="Test output") history = ChatHistory() initial_turn_id = "123-456" history.current_turn_id = initial_turn_id # Add a message with a specific turn ID history.add_message( "user", mock_input, ) history.history[-1].turn_id = initial_turn_id # Add another message with a different turn ID other_turn_id = "789-012" history.add_message( "assistant", mock_output, ) history.history[-1].turn_id = other_turn_id # Act & Assert: Delete the message with initial_turn_id and verify history.delete_turn_id(initial_turn_id) # The remaining message in history should have the other_turn_id assert len(history.history) == 1 assert history.history[0].turn_id == other_turn_id # If we delete the last message, current_turn_id should become None history.delete_turn_id(other_turn_id) assert history.current_turn_id is None assert len(history.history) == 0 # Assert: Trying to delete a non-existing turn ID should raise a ValueError with pytest.raises(ValueError, match="Turn ID non-existent-id not found in history."): history.delete_turn_id("non-existent-id") def test_get_history_with_multimodal_content(history): """Test that get_history correctly handles multimodal content""" # Create mock multimodal objects mock_image = instructor.Image(source="test_url", media_type="image/jpeg", detail="low") mock_pdf = instructor.processing.multimodal.PDF(source="test_pdf_url", media_type="application/pdf", detail="low") mock_audio = instructor.processing.multimodal.Audio(source="test_audio_url", media_type="audio/mp3", detail="low") # Add a multimodal message history.add_message( "user", MockMultimodalSchema(instruction_text="Analyze this image", images=[mock_image], pdfs=[mock_pdf], audio=mock_audio), ) # Get history and verify format history = history.get_history() assert len(history) == 1 assert history[0]["role"] == "user" assert isinstance(history[0]["content"], list) assert json.loads(history[0]["content"][0]) == {"instruction_text": "Analyze this image"} assert history[0]["content"][1] == mock_image def test_get_history_with_multiple_images_multimodal_content(history): """Test that get_history correctly handles multimodal content""" class MockMultimodalSchemaArbitraryKeys(BaseIOSchema): """Test schema for multimodal content""" instruction_text: str = Field(..., description="The instruction text") some_key_for_images: List[instructor.Image] = Field(..., description="The images to analyze") some_other_key_with_image: instructor.Image = Field(..., description="The images to analyze") # Create a mock image mock_image = instructor.Image(source="test_url", media_type="image/jpeg", detail="low") mock_image_2 = instructor.Image(source="test_url_2", media_type="image/jpeg", detail="low") mock_image_3 = instructor.Image(source="test_url_3", media_type="image/jpeg", detail="low") # Add a multimodal message history.add_message( "user", MockMultimodalSchemaArbitraryKeys( instruction_text="Analyze this image", some_other_key_with_image=mock_image, some_key_for_images=[mock_image_2, mock_image_3], ), ) # Get history and verify format history = history.get_history() assert len(history) == 1 assert history[0]["role"] == "user" assert isinstance(history[0]["content"], list) assert json.loads(history[0]["content"][0]) == {"instruction_text": "Analyze this image"} assert mock_image in history[0]["content"] assert mock_image_2 in history[0]["content"] assert mock_image_3 in history[0]["content"] def test_get_history_with_mixed_content(history): """Test that get_history correctly handles mixed multimodal and non-multimodal items in lists""" # Create a schema with a list that can contain both multimodal and non-multimodal items class MixedContentSchema(BaseIOSchema): """Schema for testing mixed multimodal and non-multimodal content""" instruction_text: str = Field(..., description="The instruction text") mixed_items: List[Union[str, instructor.Image]] = Field(..., description="Mix of strings and images") mock_image = instructor.Image(source="test_url", media_type="image/jpeg", detail="low") # Add a message with mixed content history.add_message( "user", MixedContentSchema(instruction_text="Analyze this", mixed_items=["text_item1", mock_image, "text_item2"]), ) # Get history and verify format result = history.get_history() assert len(result) == 1 assert result[0]["role"] == "user" assert isinstance(result[0]["content"], list) # Should have JSON for non-multimodal items and the image separately json_content = json.loads(result[0]["content"][0]) assert json_content["instruction_text"] == "Analyze this" assert json_content["mixed_items"] == ["text_item1", "text_item2"] assert result[0]["content"][1] == mock_image def test_process_multimodal_paths_comprehensive(): """Comprehensive test for _process_multimodal_paths and load functionality""" history = ChatHistory() # Test 1: Direct Image/PDF objects with file paths vs URLs image_file = instructor.Image(source="test/image.jpg", media_type="image/jpeg") image_url = instructor.Image(source="https://example.com/image.jpg", media_type="image/jpeg") image_data = instructor.Image(source="data:image/jpeg;base64,xyz", media_type="image/jpeg") pdf_file = instructor.processing.multimodal.PDF(source="test/doc.pdf", media_type="application/pdf") history._process_multimodal_paths(image_file) history._process_multimodal_paths(image_url) history._process_multimodal_paths(image_data) history._process_multimodal_paths(pdf_file) assert isinstance(image_file.source, Path) and image_file.source == Path("test/image.jpg") assert isinstance(image_url.source, str) and image_url.source == "https://example.com/image.jpg" assert isinstance(image_data.source, str) and image_data.source == "data:image/jpeg;base64,xyz" assert isinstance(pdf_file.source, Path) and pdf_file.source == Path("test/doc.pdf") # Test 2: Lists with mixed content test_list = [ "regular_string", instructor.Image(source="test/list_image.jpg", media_type="image/jpeg"), instructor.Image(source="https://example.com/url_image.jpg", media_type="image/jpeg"), ] history._process_multimodal_paths(test_list) assert isinstance(test_list[1].source, Path) and test_list[1].source == Path("test/list_image.jpg") assert isinstance(test_list[2].source, str) and test_list[2].source == "https://example.com/url_image.jpg" # Test 3: Dictionaries test_dict = {"image": instructor.Image(source="test/dict_image.jpg", media_type="image/jpeg"), "regular": "text_content"} history._process_multimodal_paths(test_dict) assert isinstance(test_dict["image"].source, Path) and test_dict["image"].source == Path("test/dict_image.jpg") # Test 4: Pydantic model class TestModel(BaseIOSchema): """Test model for multimodal path processing""" image_field: instructor.Image = Field(..., description="Image field") text_field: str = Field(..., description="Text field") model_instance = TestModel( image_field=instructor.Image(source="test/model_image.jpg", media_type="image/jpeg"), text_field="test text" ) history._process_multimodal_paths(model_instance) assert isinstance(model_instance.image_field.source, Path) assert model_instance.image_field.source == Path("test/model_image.jpg") # Test 5: Object with __dict__ class SimpleObject: def __init__(self): self.image = instructor.Image(source="test/obj_image.jpg", media_type="image/jpeg") self.__pydantic_fields_set__ = {"should_be_skipped"} obj = SimpleObject() history._process_multimodal_paths(obj) assert isinstance(obj.image.source, Path) and obj.image.source == Path("test/obj_image.jpg") # Test 6: Enum (should not process __dict__) from enum import Enum class TestEnum(Enum): VALUE1 = "value1" history._process_multimodal_paths(TestEnum.VALUE1) # Should not raise errors assert TestEnum.VALUE1.value == "value1" # Test 7: Load functionality with multimodal file paths original_history = ChatHistory() original_history.add_message( "user", MockMultimodalSchema( instruction_text="Process this file", images=[instructor.Image(source="test/sample.jpg", media_type="image/jpeg")], pdfs=[instructor.processing.multimodal.PDF(source="test/doc.pdf", media_type="application/pdf")], audio=instructor.processing.multimodal.Audio(source="test/audio.mp3", media_type="audio/mp3"), ), ) # Dump and reload dumped = original_history.dump() loaded_history = ChatHistory() loaded_history.load(dumped) # Verify that the loaded images and PDFs have Path objects for file-like sources loaded_message = loaded_history.history[0] loaded_content = loaded_message.content assert isinstance(loaded_content.images[0].source, Path) assert loaded_content.images[0].source == Path("test/sample.jpg") assert isinstance(loaded_content.pdfs[0].source, Path) assert loaded_content.pdfs[0].source == Path("test/doc.pdf") def test_tuple_attachment_paths_survive_history_round_trip(): original = ChatHistory() original.add_message( "user", TupleAttachments( attachments=( instructor.Image(source=Path("image.png"), media_type="image/png"), instructor.processing.multimodal.PDF(source=Path("document.pdf"), media_type="application/pdf"), ), nested={"images": ((instructor.Image(source=Path("nested.png"), media_type="image/png"),),)}, ), ) restored = ChatHistory() restored.load(original.dump()) content = restored.history[0].content assert content.attachments[0].source == Path("image.png") assert content.attachments[1].source == Path("document.pdf") assert content.nested["images"][0][0].source == Path("nested.png") assert isinstance(content.attachments, tuple) parts = restored.get_history()[0]["content"] assert len(parts) == 3 assert all(isinstance(part.source, Path) for part in parts) def test_tuple_url_and_data_sources_survive_history_round_trip(): original = ChatHistory() original.add_message( "user", TupleAttachments( attachments=( instructor.Image(source="https://example.com/image.png", media_type="image/png"), instructor.processing.multimodal.PDF( source="data:application/pdf;base64,JVBERi0=", media_type="application/pdf" ), ), nested={}, ), ) restored = ChatHistory() restored.load(original.dump()) attachments = restored.history[0].content.attachments assert attachments[0].source == "https://example.com/image.png" assert attachments[1].source == "data:application/pdf;base64,JVBERi0=" def test_get_history_nested_pydantic_with_toplevel_multimodal(history): """Issue #208: nested Pydantic model + top-level multimodal causes json.dumps TypeError""" class ContextInfo(BaseIOSchema): """Nested context info.""" label: str = Field(..., description="Label") value: str = Field(..., description="Value") class AgentInput(BaseIOSchema): """Input with multimodal and nested schema.""" instruction: str = Field(..., description="Instruction text") images: List[instructor.Image] = Field(..., description="Images") context: ContextInfo = Field(..., description="Nested context") mock_image = instructor.Image(source="test_url", media_type="image/jpeg", detail="low") context = ContextInfo(label="example", value="nested") content = AgentInput(instruction="Do something", images=[mock_image], context=context) history.add_message("user", content) result = history.get_history() assert len(result) == 1 assert result[0]["role"] == "user" assert isinstance(result[0]["content"], list) json_part = json.loads(result[0]["content"][0]) assert json_part["instruction"] == "Do something" assert json_part["context"] == {"label": "example", "value": "nested"} assert "images" not in json_part assert result[0]["content"][1] == mock_image def test_get_history_deeply_nested_multimodal_only(history): """Issue #141: multimodal inside nested schema with no top-level multimodal""" class Document(BaseIOSchema): """PDF document with owner.""" pdf: instructor.processing.multimodal.PDF = Field(..., description="The PDF data") owner: str = Field(..., description="The PDF owner") class InputSchema(BaseIOSchema): """A list of documents to analyze.""" documents: List[Document] = Field(..., description="List of documents") instruction: str = Field(..., description="What to do") mock_pdf = instructor.processing.multimodal.PDF(source="test_pdf_url", media_type="application/pdf", detail="low") content = InputSchema( documents=[Document(pdf=mock_pdf, owner="Alice")], instruction="Analyze these", ) history.add_message("user", content) result = history.get_history() assert len(result) == 1 assert isinstance(result[0]["content"], list) json_part = json.loads(result[0]["content"][0]) assert json_part["instruction"] == "Analyze these" assert json_part["documents"] == [{"owner": "Alice"}] assert result[0]["content"][1] == mock_pdf def test_get_history_mixed_nested_and_toplevel_multimodal(history): """Both nested and top-level multimodal content""" class Attachment(BaseIOSchema): """An attachment with an image.""" image: instructor.Image = Field(..., description="Attached image") caption: str = Field(..., description="Caption") class MessageInput(BaseIOSchema): """Message with both nested and top-level multimodal.""" text: str = Field(..., description="Message text") inline_image: instructor.Image = Field(..., description="Inline image") attachment: Attachment = Field(..., description="An attachment") img1 = instructor.Image(source="inline_url", media_type="image/jpeg", detail="low") img2 = instructor.Image(source="attached_url", media_type="image/png", detail="low") content = MessageInput( text="Check this out", inline_image=img1, attachment=Attachment(image=img2, caption="See here"), ) history.add_message("user", content) result = history.get_history() assert len(result) == 1 assert isinstance(result[0]["content"], list) assert len(result[0]["content"]) == 3 # JSON + 2 images json_part = json.loads(result[0]["content"][0]) assert json_part["text"] == "Check this out" assert json_part["attachment"] == {"caption": "See here"} assert "inline_image" not in json_part assert img1 in result[0]["content"] assert img2 in result[0]["content"] def test_get_history_list_of_nested_schemas_with_multimodal(history): """Multiple nested schemas each containing multimodal objects""" class Document(BaseIOSchema): """A document with PDF.""" pdf: instructor.processing.multimodal.PDF = Field(..., description="The PDF") title: str = Field(..., description="Document title") class BatchInput(BaseIOSchema): """Batch of documents.""" documents: List[Document] = Field(..., description="Documents to process") pdf1 = instructor.processing.multimodal.PDF(source="doc1.pdf", media_type="application/pdf", detail="low") pdf2 = instructor.processing.multimodal.PDF(source="doc2.pdf", media_type="application/pdf", detail="low") content = BatchInput( documents=[ Document(pdf=pdf1, title="First"), Document(pdf=pdf2, title="Second"), ] ) history.add_message("user", content) result = history.get_history() assert len(result) == 1 assert isinstance(result[0]["content"], list) assert len(result[0]["content"]) == 3 # JSON + 2 PDFs json_part = json.loads(result[0]["content"][0]) assert json_part["documents"] == [{"title": "First"}, {"title": "Second"}] assert pdf1 in result[0]["content"] assert pdf2 in result[0]["content"] def test_get_history_only_multimodal_fields(history): """Schema where ALL fields are multimodal - JSON should be omitted""" class ImagesOnly(BaseIOSchema): """Only images.""" images: List[instructor.Image] = Field(..., description="Images") img1 = instructor.Image(source="url1", media_type="image/jpeg", detail="low") img2 = instructor.Image(source="url2", media_type="image/jpeg", detail="low") content = ImagesOnly(images=[img1, img2]) history.add_message("user", content) result = history.get_history() assert len(result) == 1 assert isinstance(result[0]["content"], list) # No JSON string should be present since all fields are multimodal assert len(result[0]["content"]) == 2 assert all(not isinstance(item, str) for item in result[0]["content"]) assert img1 in result[0]["content"] assert img2 in result[0]["content"] def test_get_history_no_multimodal_unchanged(history): """Non-multimodal schemas should work exactly as before""" class SimpleInput(BaseIOSchema): """Simple input.""" text: str = Field(..., description="Text") count: int = Field(..., description="Count") content = SimpleInput(text="hello", count=42) history.add_message("user", content) result = history.get_history() assert len(result) == 1 assert result[0]["role"] == "user" assert isinstance(result[0]["content"], str) assert json.loads(result[0]["content"]) == {"text": "hello", "count": 42} # --------------------------------------------------------------------------- # Direct unit tests for _extract_multimodal_info # --------------------------------------------------------------------------- def test_extract_multimodal_info_plain_schema(): """No multimodal content returns empty list and None exclude spec""" class Plain(BaseIOSchema): """Plain schema.""" text: str = Field(..., description="Text") objs, spec = ChatHistory._extract_multimodal_info(Plain(text="hello")) assert objs == [] assert spec is None def test_extract_multimodal_info_toplevel_image(): """Top-level Image returns the object and True exclude spec""" img = instructor.Image(source="url", media_type="image/jpeg", detail="low") objs, spec = ChatHistory._extract_multimodal_info(img) assert objs == [img] assert spec is True def test_extract_multimodal_info_nested_schema(): """Nested schema with multimodal returns correct exclude spec shape""" class Inner(BaseIOSchema): """Inner schema.""" image: instructor.Image = Field(..., description="Image") label: str = Field(..., description="Label") class Outer(BaseIOSchema): """Outer schema.""" inner: Inner = Field(..., description="Inner") text: str = Field(..., description="Text") img = instructor.Image(source="url", media_type="image/jpeg", detail="low") objs, spec = ChatHistory._extract_multimodal_info(Outer(inner=Inner(image=img, label="test"), text="hello")) assert objs == [img] # Exclude spec should be {"inner": {"image": True}} assert spec == {"inner": {"image": True}} def test_extract_multimodal_info_list_all_multimodal(): """List where every item is multimodal collapses to True""" class Schema(BaseIOSchema): """Schema with all-multimodal list.""" images: List[instructor.Image] = Field(..., description="Images") img1 = instructor.Image(source="url1", media_type="image/jpeg", detail="low") img2 = instructor.Image(source="url2", media_type="image/jpeg", detail="low") objs, spec = ChatHistory._extract_multimodal_info(Schema(images=[img1, img2])) assert objs == [img1, img2] # All items multimodal → field excluded entirely assert spec == {"images": True} def test_extract_multimodal_info_list_partial_multimodal(): """List with mixed content returns index-based exclude spec""" class Schema(BaseIOSchema): """Schema with mixed list.""" items: List[Union[str, instructor.Image]] = Field(..., description="Items") img = instructor.Image(source="url", media_type="image/jpeg", detail="low") objs, spec = ChatHistory._extract_multimodal_info(Schema(items=["text", img])) assert objs == [img] # Only index 1 is multimodal assert spec == {"items": {1: True}} def test_extract_multimodal_info_tuple_support(): """Tuples are handled identically to lists""" img = instructor.Image(source="url", media_type="image/jpeg", detail="low") objs, spec = ChatHistory._extract_multimodal_info((img, "text")) assert objs == [img] assert spec == {0: True} def test_extract_multimodal_info_dict_with_multimodal(): """Dict values containing multimodal objects return correct exclude spec""" img = instructor.Image(source="url", media_type="image/jpeg", detail="low") objs, spec = ChatHistory._extract_multimodal_info({"key1": img, "key2": "text"}) assert objs == [img] assert spec == {"key1": True} def test_extract_multimodal_info_dict_all_multimodal(): """Dict where all values are multimodal collapses to True""" img1 = instructor.Image(source="url1", media_type="image/jpeg", detail="low") img2 = instructor.Image(source="url2", media_type="image/jpeg", detail="low") objs, spec = ChatHistory._extract_multimodal_info({"a": img1, "b": img2}) assert objs == [img1, img2] assert spec is True def test_extract_multimodal_info_dict_no_multimodal(): """Dict with no multimodal returns empty list and None""" objs, spec = ChatHistory._extract_multimodal_info({"key1": "text", "key2": 42}) assert objs == [] assert spec is None def test_get_history_dict_of_images(history): """Dict[str, Image] field exercises the dict code path in get_history""" class DictImageSchema(BaseIOSchema): """Schema with dict of images.""" image_map: Dict[str, instructor.Image] = Field(..., description="Named images") note: str = Field(..., description="A note") img_a = instructor.Image(source="url_a", media_type="image/jpeg", detail="low") img_b = instructor.Image(source="url_b", media_type="image/jpeg", detail="low") content = DictImageSchema(image_map={"front": img_a, "back": img_b}, note="Two views") history.add_message("user", content) result = history.get_history() assert len(result) == 1 assert isinstance(result[0]["content"], list) assert len(result[0]["content"]) == 3 # JSON + 2 images json_part = json.loads(result[0]["content"][0]) assert json_part["note"] == "Two views" assert "image_map" not in json_part assert img_a in result[0]["content"] assert img_b in result[0]["content"] class VideoMessageSchema(BaseIOSchema): """Test schema with a video field""" prompt: str = Field(..., description="Instruction for the video") video: VideoURL = Field(..., description="The video to analyze") def test_get_history_video_url_becomes_content_part(history): """VideoURL fields are separated from the JSON text as video_url content parts""" content = VideoMessageSchema( prompt="Summarize the video", video=VideoURL(url="https://example.com/clip.mp4", fps=1.0), ) history.add_message("user", content) result = history.get_history() assert len(result) == 1 assert json.loads(result[0]["content"][0]) == {"prompt": "Summarize the video"} assert result[0]["content"][1] == { "type": "video_url", "video_url": {"url": "https://example.com/clip.mp4", "fps": 1.0}, } def test_dump_and_load_video_url(history): """VideoURL round-trips through dump() and load() as a typed object""" history.add_message("user", VideoMessageSchema(prompt="Summarize", video=VideoURL(url="https://example.com/clip.mp4"))) loaded_history = ChatHistory() loaded_history.load(history.dump()) loaded_content = loaded_history.history[0].content assert isinstance(loaded_content.video, VideoURL) assert loaded_content.video.url == "https://example.com/clip.mp4" ``` ### File: atomic-agents/tests/context/test_system_prompt_generator.py ```python from typing import Dict, Optional import pytest from atomic_agents.context import ( SystemPromptGenerator, BaseDynamicContextProvider, BaseSystemPromptGenerator, ) class MockContextProvider(BaseDynamicContextProvider): def __init__(self, title: str, info: str): super().__init__(title) self._info = info def get_info(self) -> str: return self._info class MockSystemPromptGenerator(BaseSystemPromptGenerator): def __init__(self, system_prompt: str, context_providers: Optional[Dict[str, BaseDynamicContextProvider]] = None): super().__init__(context_providers) self.system_prompt = system_prompt def generate_prompt(self): return self.system_prompt 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 ) 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={ "provider1": MockContextProvider("Provider 1", "Info 1"), "provider2": MockContextProvider("Provider 2", "Info 2"), }, ) 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 def test_base_system_prompt_generator_repr(): mock_context_provider = MockContextProvider("Mock Provider", "Test") mock_generator = MockSystemPromptGenerator( context_providers={"mock_provider": mock_context_provider}, system_prompt="Test prompt" ) assert repr(mock_generator) == "MockSystemPromptGenerator (providers=['mock_provider'])" def test_custom_system_prompt_generator(): mock_context_provider = MockContextProvider("Mock Provider", "Test") mock_generator = MockSystemPromptGenerator( context_providers={"mock_provider": mock_context_provider}, system_prompt="Test prompt" ) assert mock_generator.context_providers == {"mock_provider": mock_context_provider} assert mock_generator.system_prompt == "Test prompt" def test_system_prompt_generator_with_no_generate_prompt(): with pytest.raises(TypeError): BaseSystemPromptGenerator() def test_base_system_prompt_generator_with_no_context_providers(): generator = MockSystemPromptGenerator(system_prompt="Test prompt") assert generator.context_providers == {} assert repr(generator) == "MockSystemPromptGenerator (providers=[])" ``` ### File: atomic-agents/tests/observability/instructor_transport.py ```python """Offline Instructor clients using real OpenAI request and validation paths.""" from dataclasses import dataclass import httpx import instructor import openai from openai.types import CompletionUsage from openai.types.completion_usage import CompletionTokensDetails, PromptTokensDetails @dataclass(frozen=True) class ScriptedResponse: arguments: str = '{"answer": "42"}' input_tokens: int = 3 output_tokens: int = 2 cached_tokens: int = 1 reasoning_tokens: int = 2 schema_name: str = "Answer" def chat_usage(self) -> CompletionUsage: return CompletionUsage( prompt_tokens=self.input_tokens, completion_tokens=self.output_tokens, total_tokens=self.input_tokens + self.output_tokens, prompt_tokens_details=PromptTokensDetails(cached_tokens=self.cached_tokens), completion_tokens_details=CompletionTokensDetails(reasoning_tokens=self.reasoning_tokens), ) def responses_body(self) -> dict[str, object]: return { "id": "resp_test", "object": "response", "created_at": 1, "status": "completed", "model": "response-model", "output": [ { "type": "function_call", "id": "fc_test", "call_id": "call_test", "name": self.schema_name, "arguments": self.arguments, "status": "completed", } ], "usage": { "input_tokens": self.input_tokens, "output_tokens": self.output_tokens, "total_tokens": self.input_tokens + self.output_tokens, "input_tokens_details": {"cached_tokens": self.cached_tokens}, "output_tokens_details": {"reasoning_tokens": self.reasoning_tokens}, }, } def chat_body(self, mode: instructor.Mode) -> dict[str, object]: message: dict[str, object] = {"role": "assistant", "content": self.arguments} if mode == instructor.Mode.TOOLS: message = { "role": "assistant", "content": None, "tool_calls": [ { "id": "call_test", "type": "function", "function": {"name": self.schema_name, "arguments": self.arguments}, } ], } return { "id": "chatcmpl_test", "object": "chat.completion", "created": 1, "model": "response-model", "choices": [{"index": 0, "message": message, "finish_reason": "stop"}], "usage": self.chat_usage().model_dump(), } def response_transport(mode: instructor.Mode, responses: list[ScriptedResponse]) -> httpx.MockTransport: scripted = iter(responses) def respond(request: httpx.Request) -> httpx.Response: response = next(scripted) body = response.responses_body() if mode.value == "responses_tools" else response.chat_body(mode) expected_path = "/v1/responses" if mode.value == "responses_tools" else "/v1/chat/completions" if request.url.path != expected_path: raise AssertionError(f"Unexpected transport path: {request.url.path}") return httpx.Response(200, json=body) return httpx.MockTransport(respond) def instructor_client(mode: instructor.Mode, responses: list[ScriptedResponse]) -> instructor.Instructor: return instructor.from_openai( openai.OpenAI(api_key="test", http_client=httpx.Client(transport=response_transport(mode, responses))), mode=mode ) def async_instructor_client(mode: instructor.Mode, responses: list[ScriptedResponse]) -> instructor.AsyncInstructor: return instructor.from_openai( openai.AsyncOpenAI(api_key="test", http_client=httpx.AsyncClient(transport=response_transport(mode, responses))), mode=mode, ) ``` ### File: atomic-agents/tests/observability/test_instructor_modes.py ```python """Real Instructor transport, validation and tracing across Ada's request modes.""" import asyncio import instructor import pytest from instructor.core.exceptions import InstructorRetryException from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import SpanKind, StatusCode from instructor_transport import ScriptedResponse, async_instructor_client, instructor_client from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema from atomic_agents.observability import TracingConfig from atomic_agents.observability import installation class Question(BaseIOSchema): """A question sent through the real client.""" question: str class Answer(BaseIOSchema): """The validated model answer.""" answer: str @pytest.fixture(params=["TOOLS", "JSON", "JSON_SCHEMA", "RESPONSES_TOOLS"]) def mode(request): selected = getattr(instructor.Mode, request.param, None) if selected is None: pytest.skip(f"Installed Instructor has no Mode.{request.param} for from_openai") return selected @pytest.fixture(params=[False, True], ids=["run", "run_async"]) def asynchronous(request): return request.param @pytest.fixture def exporter(): exported = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exported)) installation.install(tracer_provider=provider, config=TracingConfig(openinference_attributes=False)) yield exported installation.uninstall() provider.shutdown() def run_script(mode, asynchronous, responses): construct_client = async_instructor_client if asynchronous else instructor_client client = construct_client(mode, responses) configured = AtomicAgent[Question, Answer](AgentConfig(client=client, model="test-model")) configured.model_api_parameters["max_retries"] = 1 question = Question(question="What is the answer?") if asynchronous: return asyncio.run(configured.run_async(question)) return configured.run(question) def assert_usage(span, input_tokens, output_tokens, cached_tokens, reasoning_tokens): assert span.attributes["gen_ai.usage.input_tokens"] == input_tokens assert span.attributes["gen_ai.usage.output_tokens"] == output_tokens assert span.attributes["gen_ai.usage.cache_read.input_tokens"] == cached_tokens assert span.attributes["gen_ai.usage.reasoning.output_tokens"] == reasoning_tokens def assert_attempt_spans(exporter, responses): spans = exporter.get_finished_spans() parent = spans[-1] attempts = spans[:-1] assert parent.attributes["gen_ai.operation.name"] == "invoke_agent" assert len(attempts) == len(responses) assert len(spans) == len(responses) + 1 assert parent.attributes["atomic_agents.llm.attempts"] == len(responses) assert_models(parent) for number, (attempt, response) in enumerate(zip(attempts, responses), start=1): assert attempt.attributes["gen_ai.operation.name"] == "chat" assert attempt.name == "chat test-model" assert attempt.kind == SpanKind.CLIENT assert attempt.parent.span_id == parent.context.span_id assert attempt.attributes["atomic_agents.llm.attempt"] == number assert_models(attempt) assert_usage(attempt, response.input_tokens, response.output_tokens, response.cached_tokens, response.reasoning_tokens) return attempts, parent def assert_models(span): assert span.attributes["gen_ai.request.model"] == "test-model" assert span.attributes["gen_ai.response.model"] == "response-model" assert span.attributes["gen_ai.provider.name"] == "openai" def test_success_model_usage_and_attempt_count(mode, asynchronous, exporter): responses = [ScriptedResponse()] assert run_script(mode, asynchronous, responses).answer == "42" attempts, parent = assert_attempt_spans(exporter, responses) assert attempts[0].status.status_code == parent.status.status_code == StatusCode.OK assert_usage(parent, 3, 2, 1, 2) def test_validation_retry_preserves_each_attempt_and_total_usage(mode, asynchronous, exporter): responses = [ ScriptedResponse(arguments='{"answer": {"secret": "private output"}}'), ScriptedResponse(input_tokens=7, output_tokens=5, cached_tokens=3, reasoning_tokens=4), ] assert run_script(mode, asynchronous, responses).answer == "42" attempts, parent = assert_attempt_spans(exporter, responses) assert [attempt.status.status_code for attempt in attempts] == [StatusCode.ERROR, StatusCode.OK] assert parent.status.status_code == StatusCode.OK assert any(event.name == "validation_error" for event in attempts[0].events) assert any(event.name == "validation_error" for event in parent.events) assert_usage(parent, 10, 7, 4, 6) def test_hidden_retry_exhaustion_redacts_all_span_failures(mode, asynchronous, exporter): installed = installation._current_installation() installation.install( tracer_provider=installed.tracer_provider, config=TracingConfig(hide_outputs=True, openinference_attributes=False) ) responses = [ScriptedResponse(arguments='{"answer": {"secret": "private model output"}}')] * 2 with pytest.raises(InstructorRetryException): run_script(mode, asynchronous, responses) attempts, parent = assert_attempt_spans(exporter, responses) assert [attempt.status.status_code for attempt in attempts] == [StatusCode.ERROR, StatusCode.ERROR] assert parent.status.status_code == StatusCode.ERROR assert "__REDACTED__" in parent.attributes["atomic_agents.llm.last_validation_error"] assert_usage(parent, 6, 4, 2, 4) for span in (*attempts, parent): assert "__REDACTED__" in span.status.description assert "private model output" not in str(span.attributes) + str(span.events) + span.status.description for event in span.events: if event.name == "exception": assert "__REDACTED__" in event.attributes["exception.message"] assert "__REDACTED__" in event.attributes["exception.stacktrace"] ``` ### File: atomic-agents/tests/observability/test_llm_attempt_spans.py ```python import asyncio import json import threading from concurrent.futures import ThreadPoolExecutor, TimeoutError from types import SimpleNamespace from unittest.mock import Mock import instructor import pytest from instructor.core.hooks import Hooks from opentelemetry import context as context_api, trace as trace_api from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import SpanKind, StatusCode from pydantic import ValidationError from instructor.core.exceptions import InstructorRetryException from instructor_transport import ScriptedResponse, instructor_client from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema from atomic_agents.observability import TracingConfig, using_conversation from atomic_agents.observability import agent_spans, installation, llm_attempt_spans class Question(BaseIOSchema): """A test question.""" question: str class Answer(BaseIOSchema): """A test answer.""" answer: str @pytest.fixture def exporter(): exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) installation.install(tracer_provider=provider, config=TracingConfig(openinference_attributes=False)) yield exporter installation.uninstall() def usage(input_tokens=3, output_tokens=2): return ScriptedResponse(input_tokens=input_tokens, output_tokens=output_tokens).chat_usage() def finished_spans(exporter, count): spans = exporter.get_finished_spans() assert len(spans) == count return spans def response(): return SimpleNamespace(model="response-model", usage=usage()) def scripted_client(script): client = Mock(spec=instructor.Instructor) client.provider = instructor.Provider.OPENAI client.hooks = Hooks() client.on.side_effect = client.hooks.on client.off.side_effect = client.hooks.off client.chat = Mock() client.chat.completions = Mock() def create(**kwargs): script(client.hooks, kwargs) return Answer(answer="42") client.chat.completions.create.side_effect = create return client def agent(client): return AtomicAgent[Question, Answer](AgentConfig(client=client, model="test-model")) def retry_script(hooks, kwargs): hooks.emit_completion_arguments(**kwargs) hooks.emit_completion_response(response()) hooks.emit_parse_error(ValueError("invalid output")) hooks.emit_completion_arguments(**kwargs) hooks.emit_completion_response(response()) def run_agent(client): return agent(client).run(Question(question="test")) @pytest.mark.parametrize("mode", [instructor.Mode.TOOLS, instructor.Mode.JSON, instructor.Mode.JSON_SCHEMA]) def test_opt_in_real_transport_messages_follow_attempts(exporter, mode): installation.install(config=TracingConfig(capture_messages=True, max_attribute_length=20000)) client = instructor_client(mode, [ScriptedResponse(arguments='{"answer": {"secret": "retry"}}'), ScriptedResponse()]) assert run_agent(client).answer == "42" first, second, parent = finished_spans(exporter, 3) assert "gen_ai.output.messages" in first.attributes first_messages = json.loads(first.attributes["gen_ai.output.messages"]) second_messages = json.loads(second.attributes["gen_ai.output.messages"]) assert "retry" in str(first_messages) assert "retry" not in str(second_messages) assert "42" in str(second_messages) assert json.loads(first.attributes["gen_ai.input.messages"])[-1]["role"] == "user" assert first.parent.span_id == second.parent.span_id == parent.context.span_id assert parent.attributes["gen_ai.usage.input_tokens"] == 6 assert parent.attributes["atomic_agents.llm.attempts"] == 2 assert "gen_ai.output.messages" not in parent.attributes expected_part = "tool_call" if mode == instructor.Mode.TOOLS else "text" assert second_messages[0]["parts"][0]["type"] == expected_part @pytest.mark.parametrize("arguments", ['{"answer": "4', '{"answer": NaN}']) def test_unparseable_tool_call_arguments_stay_visible_across_the_retry(exporter, arguments): installation.install(config=TracingConfig(capture_messages=True, max_attribute_length=20000)) client = instructor_client(instructor.Mode.TOOLS, [ScriptedResponse(arguments=arguments), ScriptedResponse()]) assert run_agent(client).answer == "42" first, second, _ = finished_spans(exporter, 3) rejected_call = json.loads(first.attributes["gen_ai.output.messages"])[0]["parts"][0] replayed_calls = [ part for message in json.loads(second.attributes["gen_ai.input.messages"]) for part in message["parts"] if part["type"] == "tool_call" ] assert rejected_call["arguments"] == arguments assert [call["arguments"] for call in replayed_calls] == [arguments] def test_oversized_request_messages_leave_their_length_on_the_attempt_span(exporter): installation.install(config=TracingConfig(capture_messages=True, max_attribute_length=40)) run_agent(instructor_client(instructor.Mode.TOOLS, [ScriptedResponse()])) child, _ = finished_spans(exporter, 2) assert "gen_ai.input.messages" not in child.attributes assert child.attributes["gen_ai.input.messages.original_length"] > 40 @pytest.mark.parametrize("hide_inputs, hide_outputs", [(False, False), (True, False), (False, True), (True, True)]) def test_message_visibility_directions(exporter, hide_inputs, hide_outputs): installation.install(config=TracingConfig(capture_messages=True, hide_inputs=hide_inputs, hide_outputs=hide_outputs)) run_agent(instructor_client(instructor.Mode.TOOLS, [ScriptedResponse()])) child, parent = finished_spans(exporter, 2) assert ("gen_ai.input.messages" in child.attributes) is not (hide_inputs or hide_outputs) assert ("gen_ai.output.messages" in child.attributes) is not hide_outputs @pytest.mark.parametrize("mode", [instructor.Mode.TOOLS, instructor.Mode.JSON, instructor.Mode.JSON_SCHEMA]) def test_hidden_outputs_keep_rejected_answer_out_of_retry_spans(exporter, mode): installation.install(config=TracingConfig(capture_messages=True, hide_outputs=True, max_attribute_length=20000)) rejected_answer = ScriptedResponse(arguments='{"answer": {"secret": "rejected-draft"}}') assert run_agent(instructor_client(mode, [rejected_answer, ScriptedResponse()])).answer == "42" recorded = [ (dict(span.attributes), [dict(event.attributes) for event in span.events]) for span in finished_spans(exporter, 3) ] assert "rejected-draft" not in str(recorded) @pytest.mark.parametrize("enabled", [None, False]) def test_message_capture_default_and_explicit_disabled(exporter, enabled): if enabled is not None: installation.install(config=TracingConfig(capture_messages=enabled)) run_agent(instructor_client(instructor.Mode.JSON, [ScriptedResponse()])) child, parent = finished_spans(exporter, 2) assert "gen_ai.input.messages" not in child.attributes assert "gen_ai.output.messages" not in child.attributes def test_capture_with_attempt_spans_disabled_keeps_totals(exporter): installation.install(config=TracingConfig(capture_messages=True, llm_spans=False)) run_agent(instructor_client(instructor.Mode.TOOLS, [ScriptedResponse()])) (parent,) = finished_spans(exporter, 1) assert parent.attributes["gen_ai.usage.input_tokens"] == 3 assert "gen_ai.input.messages" not in parent.attributes assert "gen_ai.output.messages" not in parent.attributes def test_native_responses_capture_keeps_request_text_and_omits_response(exporter): installation.install(config=TracingConfig(capture_messages=True)) run_agent(instructor_client(instructor.Mode.RESPONSES_TOOLS, [ScriptedResponse()])) child, parent = finished_spans(exporter, 2) assert json.loads(child.attributes["gen_ai.input.messages"])[-1]["role"] == "user" assert "gen_ai.output.messages" not in child.attributes assert parent.attributes["gen_ai.usage.input_tokens"] == 3 def test_concurrent_message_capture_stays_with_own_run(exporter): installation.install(config=TracingConfig(capture_messages=True)) client = scripted_client(retry_script) async def complete(configured, text): client.hooks.emit_completion_arguments(model="test-model", messages=[{"role": "user", "content": text}]) await asyncio.sleep(0) client.hooks.emit_completion_response( SimpleNamespace( model="response-model", usage=usage(), choices=[SimpleNamespace(message={"role": "assistant", "content": text})], ) ) return Answer(answer=text) async def gather(): await asyncio.gather(*(agent_spans._trace_agent_async_call(complete)(agent(client), text) for text in ["one", "two"])) asyncio.run(gather()) children = [span for span in exporter.get_finished_spans() if span.name == "chat test-model"] assert len(children) == 2 for child in children: assert "gen_ai.input.messages" in child.attributes request = json.loads(child.attributes["gen_ai.input.messages"]) response_messages = json.loads(child.attributes["gen_ai.output.messages"]) assert request[0]["parts"] == response_messages[0]["parts"] assert {json.loads(child.attributes["gen_ai.input.messages"])[0]["parts"][0]["content"] for child in children} == { "one", "two", } def test_retry_spans_usage_validation_and_parentage(exporter): client = instructor_client( instructor.Mode.TOOLS, [ScriptedResponse(arguments='{"answer": {"secret": "private output"}}'), ScriptedResponse()] ) with using_conversation("conversation"): run_agent(client) first, second, parent = finished_spans(exporter, 3) assert [first.name, second.name] == ["chat test-model", "chat test-model"] assert [first.attributes["atomic_agents.llm.attempt"], second.attributes["atomic_agents.llm.attempt"]] == [1, 2] assert first.parent.span_id == second.parent.span_id == parent.context.span_id assert first.status.status_code == StatusCode.ERROR assert second.status.status_code == StatusCode.OK assert first.kind == second.kind == SpanKind.CLIENT assert first.events[0].name == "validation_error" assert first.attributes["error.type"] == "ValidationError" assert first.attributes["gen_ai.provider.name"] == "openai" assert first.attributes["gen_ai.conversation.id"] == "conversation" assert first.attributes["gen_ai.usage.input_tokens"] == second.attributes["gen_ai.usage.input_tokens"] == 3 assert first.attributes["gen_ai.usage.output_tokens"] == 2 assert first.attributes["gen_ai.usage.cache_read.input_tokens"] == 1 assert first.attributes["gen_ai.usage.reasoning.output_tokens"] == 2 assert parent.attributes["gen_ai.usage.input_tokens"] == 6 assert parent.attributes["gen_ai.usage.output_tokens"] == 4 assert parent.attributes.get("atomic_agents.llm.attempts") == 2 assert "private output" in parent.attributes["atomic_agents.llm.last_validation_error"] assert parent.attributes["gen_ai.response.model"] == first.attributes["gen_ai.response.model"] == "response-model" event = parent.events[0] assert event.name == "validation_error" assert event.attributes["atomic_agents.llm.attempt"] == 1 def test_disabled_attempt_spans_keep_totals(exporter): installed = installation._current_installation() installation.install(tracer_provider=installed.tracer_provider, config=TracingConfig(llm_spans=False)) run_agent(scripted_client(retry_script)) (parent,) = finished_spans(exporter, 1) assert parent.attributes.get("atomic_agents.llm.attempts") == 2 assert parent.attributes["gen_ai.usage.input_tokens"] == 6 def test_completion_error_is_bounded_and_ends_attempt(exporter): installed = installation._current_installation() installation.install(tracer_provider=installed.tracer_provider, config=TracingConfig(max_attribute_length=32)) error = ValueError("failed " * 20) expected = str(error)[:18] + "...[truncated]" def script(hooks, kwargs): hooks.emit_completion_arguments(**kwargs) hooks.emit_parse_error(error) hooks.emit_completion_error(error, attempt_number=1, max_attempts=1, is_last_attempt=True) run_agent(scripted_client(script)) child, parent = finished_spans(exporter, 2) assert child.status.status_code == StatusCode.ERROR assert child.attributes["error.type"] == "ValueError" assert child.events[-1].attributes["exception.message"] == expected assert len(child.events[-1].attributes["exception.stacktrace"]) <= 32 assert child.status.description == expected assert len(child.status.description) == 32 assert parent.attributes["atomic_agents.llm.last_validation_error"] == expected assert parent.attributes["atomic_agents.llm.last_validation_error.original_length"] == len(str(error)) assert parent.status.status_code == StatusCode.OK def test_registration_once_and_uninstall_preserves_application_hooks(exporter): client = scripted_client(retry_script) application_hook = Mock() client.hooks.on("completion:response", application_hook) for _ in range(3): run_agent(client) assert client.on.call_count == 4 assert application_hook.call_count == 6 installation.uninstall() assert client.off.call_count == 4 client.hooks.emit_completion_response(response()) assert application_hook.call_count == 7 assert not llm_attempt_spans._clients @pytest.mark.parametrize("asynchronous", [False, True]) def test_concurrent_shared_client_runs_are_isolated(exporter, asynchronous): client = scripted_client(retry_script) async def invoke(number): async def complete(agent, *args, **kwargs): client.hooks.emit_completion_arguments(model="test-model") await asyncio.sleep(0) client.hooks.emit_completion_response(SimpleNamespace(model="response-model", usage=usage(number))) return Answer(answer="ok") if asynchronous: return await agent_spans._trace_agent_async_call(complete)(agent(client)) return await asyncio.to_thread(run_agent, client) async def gather(): await asyncio.gather(invoke(10), invoke(20)) asyncio.run(gather()) parents = [span for span in exporter.get_finished_spans() if span.name.startswith("invoke_agent")] expected = [10, 20] if asynchronous else [6, 6] assert sorted(span.attributes.get("gen_ai.usage.input_tokens", -1) for span in parents) == expected assert all(span.attributes["atomic_agents.llm.attempts"] == (1 if asynchronous else 2) for span in parents) @pytest.mark.parametrize("kind", ["sync", "async", "stream", "async_stream"]) @pytest.mark.parametrize("fails", [False, True]) def test_open_attempt_ends_at_agent_end(exporter, kind, fails): client = scripted_client(retry_script) def start(): client.hooks.emit_completion_arguments(model="test-model") if fails: raise ValueError("agent failed") return Answer(answer="partial") def call(*args, **kwargs): return start() async def async_call(*args, **kwargs): return start() def stream(*args, **kwargs): yield start() yield Answer(answer="final") async def async_stream(*args, **kwargs): yield start() yield Answer(answer="final") def invoke(): if kind == "sync": return agent_spans._trace_agent_call(call)(agent(client)) if kind == "async": return asyncio.run(agent_spans._trace_agent_async_call(async_call)(agent(client))) if kind == "stream": iterator = agent_spans._trace_agent_stream(stream)(agent(client)) next(iterator) assert trace_api.get_current_span() is trace_api.INVALID_SPAN iterator.close() return None return asyncio.run(consume_async_stream()) async def consume_async_stream(): iterator = agent_spans._trace_agent_async_stream(async_stream)(agent(client)) await anext(iterator) assert trace_api.get_current_span() is trace_api.INVALID_SPAN await iterator.aclose() if fails: with pytest.raises(ValueError, match="agent failed"): invoke() else: invoke() child, parent = finished_spans(exporter, 2) assert child.status.status_code == StatusCode.UNSET assert child.parent.span_id == parent.context.span_id assert child.end_time <= parent.end_time def test_validation_text_redacted_and_anthropic_usage(exporter): installed = installation._current_installation() installation.install(tracer_provider=installed.tracer_provider, config=TracingConfig(hide_outputs=True)) def script(hooks, kwargs): hooks.emit_completion_arguments(**kwargs) with pytest.raises(ValidationError) as invalid: Answer.model_validate({"answer": {"secret": "private output"}}) hooks.emit_completion_response(SimpleNamespace(model="response-model", usage=usage(9, 4))) hooks.emit_parse_error(invalid.value) hooks.emit_completion_error(invalid.value, attempt_number=1, max_attempts=1, is_last_attempt=True) run_agent(scripted_client(script)) child, parent = finished_spans(exporter, 2) assert parent.attributes["atomic_agents.llm.last_validation_error"] == "__REDACTED__" assert parent.events[0].attributes["atomic_agents.llm.last_validation_error"] == "__REDACTED__" assert parent.attributes["gen_ai.usage.input_tokens"] == 9 assert child.status.description == "ValidationError: 1 errors (__REDACTED__)" assert child.events[-1].attributes["exception.message"] == child.status.description assert "ValidationError: 1 errors (__REDACTED__)" in child.events[-1].attributes["exception.stacktrace"] assert "private output" not in str(child.events) def test_hook_failure_warns_without_breaking_agent(exporter, caplog): class BrokenResponse: @property def model(self): raise RuntimeError("broken metadata") def script(hooks, kwargs): hooks.emit_completion_arguments(**kwargs) hooks.emit_completion_response(BrokenResponse()) assert run_agent(scripted_client(script)).answer == "42" child, parent = finished_spans(exporter, 2) assert child.status.status_code == StatusCode.UNSET assert "Instructor tracing hook failed" in caplog.text @pytest.mark.parametrize("asynchronous", [False, True]) def test_stream_attempt_context_survives_cursor_advances(exporter, asynchronous): client = scripted_client(retry_script) def begin(): client.hooks.emit_completion_arguments(model="test-model") return trace_api.get_current_span().get_span_context().span_id def complete(expected): assert trace_api.get_current_span().get_span_context().span_id == expected client.hooks.emit_completion_response(response()) return Answer(answer="final") def stream(*args, **kwargs): expected = begin() yield Answer(answer="partial") yield complete(expected) async def async_stream(*args, **kwargs): expected = begin() yield Answer(answer="partial") await asyncio.sleep(0) yield complete(expected) async def consume(): return [partial async for partial in agent_spans._trace_agent_async_stream(async_stream)(agent(client))] if asynchronous: assert len(asyncio.run(consume())) == 2 else: assert len(list(agent_spans._trace_agent_stream(stream)(agent(client)))) == 2 child, parent = finished_spans(exporter, 2) assert child.status.status_code == StatusCode.OK assert parent.status.status_code == StatusCode.OK def test_validation_text_truncation_and_extra_hook_keywords(exporter): installed = installation._current_installation() installation.install(tracer_provider=installed.tracer_provider, config=TracingConfig(max_attribute_length=4)) def script(hooks, kwargs): hooks.emit_completion_arguments(**kwargs) llm_attempt_spans._parse_error(ValueError("long error"), future_option=True) llm_attempt_spans._completion_response(response(), future_option=True) run_agent(scripted_client(script)) child, parent = finished_spans(exporter, 2) assert parent.attributes["atomic_agents.llm.last_validation_error"] == "...[" assert parent.attributes["atomic_agents.llm.last_validation_error.original_length"] == 10 def test_threaded_first_use_registers_once_and_uninstalls(exporter, monkeypatch): client = scripted_client(retry_script) first_registration = threading.Barrier(2, timeout=5) release_registration = threading.Event() original_on = Hooks.on def blocked_on(hooks, event, handler): if event == "completion:kwargs": first_registration.wait() assert release_registration.wait(timeout=5) return original_on(hooks, event, handler) monkeypatch.setattr(Hooks, "on", blocked_on) client.on.side_effect = client.hooks.on with ThreadPoolExecutor(max_workers=2) as executor: first = executor.submit(run_agent, client) first_registration.wait() second_started = threading.Event() def second_run(): second_started.set() return run_agent(client) second = executor.submit(second_run) assert second_started.wait(timeout=5) try: with pytest.raises(TimeoutError): second.result(timeout=0.1) assert client.on.call_count == 1 finally: release_registration.set() first.result(timeout=5) second.result(timeout=5) assert len(client.hooks._handlers) == 4 assert all(len(handlers) == 1 for handlers in client.hooks._handlers.values()) installation.uninstall() assert all(not handlers for handlers in client.hooks._handlers.values()) def test_partial_registration_is_rolled_back(exporter): client = scripted_client(retry_script) def fail_registration(event, handler): if event == "parse:error": raise RuntimeError("registration failed") client.hooks.on(event, handler) client.on.side_effect = fail_registration llm_attempt_spans.register_client(client) assert client not in llm_attempt_spans._clients assert all(not handlers for handlers in client.hooks._handlers.values()) client.on.side_effect = client.hooks.on llm_attempt_spans.register_client(client) assert all(len(handlers) == 1 for handlers in client.hooks._handlers.values()) @pytest.mark.parametrize("completion", ["response", "error"]) def test_suppressed_attempt_events_are_ignored_then_normal_attempt_works(exporter, completion): def script(hooks, kwargs): run = llm_attempt_spans._active_run.get() token = context_api.attach(context_api.set_value("suppress_instrumentation", True)) try: hooks.emit_completion_arguments(**kwargs) finally: context_api.detach(token) if completion == "response": hooks.emit_completion_response(response()) else: hooks.emit_completion_error(ValueError("suppressed"), attempt_number=1, max_attempts=2, is_last_attempt=False) hooks.emit_parse_error(ValueError("suppressed")) assert run.attempt_count == 0 assert run.open_attempt is None assert not exporter.get_finished_spans() hooks.emit_completion_arguments(**kwargs) hooks.emit_completion_response(response()) run_agent(scripted_client(script)) child, parent = finished_spans(exporter, 2) assert child.attributes["atomic_agents.llm.attempt"] == 1 assert parent.attributes["atomic_agents.llm.attempts"] == 1 assert parent.attributes["gen_ai.usage.input_tokens"] == 3 assert not parent.events @pytest.mark.parametrize( "mode", [instructor.Mode.TOOLS, instructor.Mode.JSON, instructor.Mode.JSON_SCHEMA, instructor.Mode.RESPONSES_TOOLS] ) def test_real_transport_usage_is_accumulated_per_response(exporter, mode): client = instructor_client( mode, [ ScriptedResponse(arguments='{"answer": {"private": "model text"}}'), ScriptedResponse(input_tokens=7, output_tokens=5, cached_tokens=3, reasoning_tokens=4), ], ) assert run_agent(client).answer == "42" first, second, parent = finished_spans(exporter, 3) assert first.status.status_code == StatusCode.ERROR assert first.events[0].name == "validation_error" assert second.status.status_code == StatusCode.OK assert first.attributes["gen_ai.usage.input_tokens"] == 3 assert second.attributes["gen_ai.usage.input_tokens"] == 7 assert parent.attributes["gen_ai.usage.input_tokens"] == 10 assert parent.attributes["gen_ai.usage.output_tokens"] == 7 assert parent.attributes["gen_ai.usage.cache_read.input_tokens"] == 4 assert parent.attributes["gen_ai.usage.reasoning.output_tokens"] == 6 assert parent.attributes["atomic_agents.llm.attempts"] == 2 assert "completion:usage" not in llm_attempt_spans._HANDLERS @pytest.mark.parametrize("hidden", [False, True]) def test_real_retry_exhaustion_redacts_or_bounds_every_exception(exporter, hidden): installed = installation._current_installation() installation.install( tracer_provider=installed.tracer_provider, config=TracingConfig(hide_outputs=hidden, max_attribute_length=256) ) secret = "private model text " * 100 client = instructor_client( instructor.Mode.TOOLS, [ScriptedResponse(arguments='{"answer": {"secret": "' + secret + '"}}')] * 3 ) configured = agent(client) configured.model_api_parameters["max_retries"] = 2 with pytest.raises(InstructorRetryException): configured.run(Question(question="test")) first, second, third, parent = finished_spans(exporter, 4) assert parent.status.status_code == StatusCode.ERROR assert parent.attributes["atomic_agents.llm.attempts"] == 3 assert parent.attributes["gen_ai.usage.input_tokens"] == 9 message = parent.events[-1].attributes["exception.message"] assert len(message) <= 256 assert len(parent.status.description) <= 256 assert len(parent.events[-1].attributes["exception.stacktrace"]) <= 256 if hidden: assert message == parent.status.description == "InstructorRetryException: 3 attempts (__REDACTED__)" assert "private model text" not in str(parent.events) assert "private model text" not in str(first.events) + str(second.events) + str(third.events) assert ( first.events[0].attributes["atomic_agents.llm.last_validation_error"] == "ValidationError: 1 errors (__REDACTED__)" ) else: assert message.endswith("...[truncated]") def test_direct_instructor_calls_inside_run_add_to_agent_totals(exporter): client = instructor_client(instructor.Mode.TOOLS, [ScriptedResponse(), ScriptedResponse(input_tokens=11)]) def call_twice(configured): for _ in range(2): client.chat.completions.create( model="test-model", response_model=Answer, messages=[{"role": "user", "content": "test"}] ) return Answer(answer="42") agent_spans._trace_agent_call(call_twice)(agent(client)) first, second, parent = finished_spans(exporter, 3) assert parent.attributes["gen_ai.usage.input_tokens"] == 14 assert parent.attributes["atomic_agents.llm.attempts"] == 2 assert first.parent.span_id == second.parent.span_id == parent.context.span_id def test_missing_model_and_final_completion_error_still_count(exporter): def script(hooks, kwargs): hooks.emit_completion_arguments() hooks.emit_completion_response(response()) hooks.emit_completion_arguments() hooks.emit_completion_error(RuntimeError("transport failed")) run_agent(scripted_client(script)) first, second, parent = finished_spans(exporter, 3) assert first.name == second.name == "chat" assert first.kind == second.kind == SpanKind.CLIENT assert second.status.status_code == StatusCode.ERROR assert parent.attributes["atomic_agents.llm.attempts"] == 2 assert parent.attributes["gen_ai.usage.input_tokens"] == 3 @pytest.mark.parametrize("chain", ["__cause__", "__context__"]) def test_wrapped_validation_error_and_tool_exception_are_redacted(exporter, chain): from atomic_agents.observability import tool_spans from atomic_agents.base.base_tool import BaseTool installed = installation._current_installation() installation.install(tracer_provider=installed.tracer_provider, config=TracingConfig(hide_outputs=True)) with pytest.raises(ValidationError) as invalid: Answer.model_validate({"answer": {"secret": "private model text"}}) wrapped = RuntimeError("private wrapper text") setattr(wrapped, chain, invalid.value) with pytest.raises(RuntimeError), tool_spans._tool_span(Mock(spec=BaseTool, tool_name="test"), (), {}): raise wrapped (span,) = finished_spans(exporter, 1) assert span.status.description == "RuntimeError: 1 errors (__REDACTED__)" assert span.events[0].attributes["exception.message"] == span.status.description stacktrace = span.events[0].attributes["exception.stacktrace"] assert "Traceback (most recent call last)" in stacktrace assert stacktrace.rstrip().endswith(span.status.description) assert "private model text" not in stacktrace assert "private wrapper text" not in stacktrace ``` ### File: atomic-agents/tests/observability/test_message_payloads.py ```python import json from dataclasses import replace from types import SimpleNamespace import pytest from opentelemetry import context as context_api from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from atomic_agents.observability import TracingConfig from atomic_agents.observability.message_payloads import write_messages class UnreadableContent: def __init__(self): self.reads = 0 @property def role(self): self.reads += 1 raise AssertionError("content was accessed") @property def choices(self): self.reads += 1 raise AssertionError("choices were accessed") class UnconsumableStream: def __iter__(self): raise AssertionError("stream was consumed") def captured(value, direction="input", **options): exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) with provider.get_tracer(__name__).start_as_current_span("chat") as span: write_messages(span, value, direction, replace(TracingConfig(capture_messages=True), **options)) return dict(exporter.get_finished_spans()[0].attributes) def decoded(value, direction="input"): return json.loads(captured(value, direction)[f"gen_ai.{direction}.messages"]) def test_message_and_part_order(): messages = [ {"role": "system", "content": "instructions"}, { "role": "user", "content": [ {"type": "image_url", "image_url": UnreadableContent()}, {"type": "text", "text": "first"}, {"type": "audio", "data": "private"}, {"type": "text", "text": "second"}, ], }, { "role": "assistant", "tool_calls": [ {"type": "function", "id": "call-1", "function": {"name": "lookup", "arguments": '{"city":"Ghent"}'}} ], }, {"role": "tool", "tool_call_id": "call-1", "content": "sunny"}, ] assert decoded(messages) == [ {"role": "system", "parts": [{"type": "text", "content": "instructions"}]}, {"role": "user", "parts": [{"type": "text", "content": "first"}, {"type": "text", "content": "second"}]}, { "role": "assistant", "parts": [{"type": "tool_call", "id": "call-1", "name": "lookup", "arguments": {"city": "Ghent"}}], }, {"role": "tool", "parts": [{"type": "tool_call_response", "id": "call-1", "response": "sunny"}]}, ] def test_output_choices_remain_separate_and_ordered(): response = SimpleNamespace( choices=[ SimpleNamespace(message=SimpleNamespace(role="assistant", content="first")), {"message": {"role": "assistant", "content": "second"}}, ] ) assert decoded(response, "output") == [ {"role": "assistant", "parts": [{"type": "text", "content": "first"}]}, {"role": "assistant", "parts": [{"type": "text", "content": "second"}]}, ] @pytest.mark.parametrize( "options", [ {"capture_messages": False}, {"hide_inputs": True}, {"hide_outputs": True}, {"llm_spans": False}, ], ) def test_input_gates_precede_content_access(options): content = UnreadableContent() assert captured([content], **options) == {} assert content.reads == 0 def test_output_gate_precedes_content_access(): content = UnreadableContent() assert captured(content, "output", hide_outputs=True) == {} assert content.reads == 0 @pytest.mark.parametrize("direction", ["input", "output"]) def test_suppression_precedes_content_access(direction): content = UnreadableContent() value = [content] if direction == "input" else content token = context_api.attach(context_api.set_value("suppress_instrumentation", True)) try: assert captured(value, direction) == {} finally: context_api.detach(token) assert content.reads == 0 def test_hidden_outputs_omit_request_messages_replaying_model_output(): request_messages = [ {"role": "user", "content": "question"}, {"role": "assistant", "content": "earlier model answer"}, {"role": "user", "content": "follow-up"}, ] assert captured(request_messages, hide_outputs=True) == {} @pytest.mark.parametrize( "value", [ None, UnconsumableStream(), UnreadableContent(), [UnreadableContent()], {}, [{"content": "missing role"}], [{"role": "user", "content": object()}], [{"role": "user", "content": ""}], [{"role": "user", "content": [{"type": "reasoning", "text": "secret"}]}], [{"role": "tool", "content": None}], [{"role": "assistant", "tool_calls": [{"type": "other"}]}], [{"role": "assistant", "tool_calls": [{"type": "function", "function": {"arguments": "{}"}}]}], ], ) def test_unrepresentable_input_is_omitted(value): assert captured(value) == {} @pytest.mark.parametrize("arguments", ['{"title": "Lease en', '{"title": NaN}', '{"ratio": -Infinity}', "broken"]) @pytest.mark.parametrize("direction", ["input", "output"]) def test_unparseable_tool_call_arguments_keep_their_raw_text(arguments, direction): rejected_answer = { "role": "assistant", "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "answer", "arguments": arguments}}], } value = ( [{"role": "user", "content": "question"}, rejected_answer] if direction == "input" else {"choices": [{"message": rejected_answer}]} ) tool_call = decoded(value, direction)[-1]["parts"][0] assert tool_call == {"type": "tool_call", "name": "answer", "arguments": arguments, "id": "call_1"} @pytest.mark.parametrize( "value", [ UnconsumableStream(), {"output": [{"type": "function_call"}]}, {"choices": UnconsumableStream()}, {"choices": [{"delta": {"content": "stream"}}]}, ], ) def test_unsupported_response_is_omitted(value): assert captured(value, "output") == {} @pytest.mark.parametrize("text", ['quote"\\\n', "Ghent ☀️ 漢字", "long" * 100]) def test_complete_json_length_bound(text): messages = [{"role": "user", "content": text}] encoded = captured(messages)["gen_ai.input.messages"] assert json.loads(encoded)[0]["parts"][0]["content"] == text assert captured(messages, max_attribute_length=len(encoded))["gen_ai.input.messages"] == encoded oversized = {"gen_ai.input.messages.original_length": len(encoded)} assert captured(messages, max_attribute_length=len(encoded) - 1) == oversized assert captured(messages, max_attribute_length=1) == oversized def test_oversized_output_messages_record_only_their_length(): response = {"choices": [{"message": {"role": "assistant", "content": "a long answer"}}]} encoded = captured(response, "output")["gen_ai.output.messages"] assert captured(response, "output", max_attribute_length=len(encoded) - 1) == { "gen_ai.output.messages.original_length": len(encoded) } @pytest.mark.parametrize("setting, expected", [("true", True), ("TRUE", True), ("1", True), ("false", False), ("yes", False)]) def test_environment_and_explicit_config_precedence(monkeypatch, setting, expected): monkeypatch.setenv("ATOMIC_AGENTS_TRACE_CAPTURE_MESSAGES", setting) assert TracingConfig.from_environment().capture_messages is expected assert TracingConfig.coerce(None).capture_messages is expected visibility = SimpleNamespace(hide_inputs=True, hide_outputs=False) assert TracingConfig.coerce(visibility).capture_messages is expected assert TracingConfig.coerce(TracingConfig()).capture_messages is False assert TracingConfig.coerce(TracingConfig(capture_messages=True)).capture_messages is True def test_capture_defaults_preserve_positional_fields(monkeypatch): monkeypatch.delenv("ATOMIC_AGENTS_TRACE_CAPTURE_MESSAGES", raising=False) assert TracingConfig(True, False, 50, False, False).capture_messages is False assert TracingConfig.from_environment().capture_messages is False ``` ### File: atomic-agents/tests/observability/test_opentelemetry_discovery.py ```python import logging import os from collections.abc import Iterator, Mapping from importlib.metadata import EntryPoint, distribution, entry_points from pathlib import Path import subprocess import sys import sysconfig from unittest.mock import Mock import pytest from openinference.instrumentation import TraceConfig from openinference.semconv.trace import SpanAttributes from opentelemetry.instrumentation.auto_instrumentation import _load from opentelemetry.instrumentation.distro import DefaultDistro from opentelemetry.sdk.trace import ReadableSpan, TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.util.types import AttributeValue import atomic_agents.observability as observability from atomic_agents.observability import AtomicAgentsInstrumentor from test_tracing import CalculationInputSchema, CalculatorTool, QuestionSchema, make_agent, sync_client @pytest.fixture def local_tracing() -> Iterator[tuple[TracerProvider, InMemorySpanExporter]]: exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) yield provider, exporter AtomicAgentsInstrumentor().uninstrument() provider.shutdown() def installed_entry_point(group: str) -> EntryPoint: points = [point for point in entry_points(group=group) if point.name == "atomic_agents"] assert len(points) == 1, f"Expected one installed Atomic Agents entry point in {group}" return points[0] def exported_attributes(span: ReadableSpan) -> Mapping[str, AttributeValue]: attributes = span.attributes assert attributes is not None return attributes def run_agent_and_tool() -> None: make_agent(sync_client()).run(QuestionSchema(question="private question")) CalculatorTool().run(CalculationInputSchema(expression="private expression")) def test_installed_opentelemetry_entry_point_activates_real_loader( monkeypatch: pytest.MonkeyPatch, local_tracing: tuple[TracerProvider, InMemorySpanExporter] ) -> None: provider, exporter = local_tracing instrumentor = installed_entry_point("opentelemetry_instrumentor") monkeypatch.setattr(_load, "entry_points", lambda group: (instrumentor,) if group == "opentelemetry_instrumentor" else ()) monkeypatch.setattr("opentelemetry.trace.get_tracer_provider", lambda: provider) monkeypatch.delenv("OTEL_PYTHON_DISABLED_INSTRUMENTATIONS", raising=False) monkeypatch.setenv("OPENINFERENCE_HIDE_INPUTS", "true") monkeypatch.setenv("OPENINFERENCE_HIDE_OUTPUTS", "true") _load._load_instrumentors(DefaultDistro()) run_agent_and_tool() spans = exporter.get_finished_spans() assert [exported_attributes(span)[SpanAttributes.OPENINFERENCE_SPAN_KIND] for span in spans] == ["AGENT", "TOOL"] assert [exported_attributes(span)[SpanAttributes.INPUT_VALUE] for span in spans] == ["__REDACTED__", "__REDACTED__"] assert [exported_attributes(span)[SpanAttributes.OUTPUT_VALUE] for span in spans] == ["__REDACTED__", "__REDACTED__"] def test_unknown_launcher_options_preserve_provider_and_privacy( local_tracing: tuple[TracerProvider, InMemorySpanExporter], caplog: pytest.LogCaptureFixture, ) -> None: provider, exporter = local_tracing AtomicAgentsInstrumentor().instrument( tracer_provider=provider, config=TraceConfig(hide_inputs=True, hide_outputs=True), skip_dep_check=True, custom_launcher_option=object(), ) run_agent_and_tool() assert "custom_launcher_option" in caplog.text assert any(record.levelno == logging.WARNING for record in caplog.records) assert "skip_dep_check" not in caplog.text spans = exporter.get_finished_spans() assert len(spans) == 2 assert [exported_attributes(span)[SpanAttributes.INPUT_VALUE] for span in spans] == ["__REDACTED__", "__REDACTED__"] assert [exported_attributes(span)[SpanAttributes.OUTPUT_VALUE] for span in spans] == ["__REDACTED__", "__REDACTED__"] def test_loader_skips_missing_opentelemetry_and_continues_without_error( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: instrumentor = installed_entry_point("opentelemetry_instrumentor") sentinel = Mock( name="sentinel", group="opentelemetry_instrumentor", value="sentinel:Instrumentor", dist=distribution("atomic-agents") ) sentinel.name = "sentinel" points = (instrumentor, sentinel) monkeypatch.setattr(_load, "entry_points", lambda group: points if group == "opentelemetry_instrumentor" else ()) monkeypatch.setattr(observability, "_opentelemetry_is_installed", lambda: False) monkeypatch.delenv("OTEL_PYTHON_DISABLED_INSTRUMENTATIONS", raising=False) with caplog.at_level(logging.DEBUG): _load._load_instrumentors(DefaultDistro()) sentinel.load.return_value.return_value.instrument.assert_called_once_with(skip_dep_check=True) assert ( 'Atomic Agents tracing needs opentelemetry-api. Install it with: pip install "atomic-agents[observability]" ' '(add "atomic-agents[phoenix]" for Arize Phoenix / OpenInference attributes)' ) in caplog.text assert not [record for record in caplog.records if record.levelno >= logging.ERROR] assert not AtomicAgentsInstrumentor().is_instrumented_by_opentelemetry @pytest.mark.parametrize("first_group", ["openinference_instrumentor", "opentelemetry_instrumentor"]) def test_both_entry_points_replace_provider_and_config( first_group: str, local_tracing: tuple[TracerProvider, InMemorySpanExporter], caplog: pytest.LogCaptureFixture ) -> None: provider, exporter = local_tracing second_group = { "openinference_instrumentor": "opentelemetry_instrumentor", "opentelemetry_instrumentor": "openinference_instrumentor", }[first_group] first = installed_entry_point(first_group).load()() second = installed_entry_point(second_group).load()() other_provider = TracerProvider() other_exporter = InMemorySpanExporter() other_provider.add_span_processor(SimpleSpanProcessor(other_exporter)) first.instrument(tracer_provider=provider, config=TraceConfig(hide_inputs=True)) try: with caplog.at_level(logging.WARNING): second.instrument(tracer_provider=other_provider, config=TraceConfig(hide_inputs=False), skip_dep_check=True) run_agent_and_tool() assert "tracer provider/config replaced" in caplog.text assert exporter.get_finished_spans() == () spans = other_exporter.get_finished_spans() assert len(spans) == 2 assert all(exported_attributes(span)[SpanAttributes.INPUT_VALUE] != "__REDACTED__" for span in spans) finally: second.uninstrument() other_provider.shutdown() @pytest.mark.parametrize("first_group", ["openinference_instrumentor", "opentelemetry_instrumentor"]) def test_either_entry_point_can_uninstrument_and_reinstall( first_group: str, local_tracing: tuple[TracerProvider, InMemorySpanExporter] ) -> None: provider, exporter = local_tracing second_group = { "openinference_instrumentor": "opentelemetry_instrumentor", "opentelemetry_instrumentor": "openinference_instrumentor", }[first_group] first = installed_entry_point(first_group).load()() second = installed_entry_point(second_group).load()() first.instrument(tracer_provider=provider) second.instrument(tracer_provider=provider) run_agent_and_tool() assert len(exporter.get_finished_spans()) == 2 second.uninstrument() first.uninstrument() run_agent_and_tool() assert len(exporter.get_finished_spans()) == 2 second.instrument(tracer_provider=provider, config=TraceConfig(hide_inputs=True), skip_dep_check=True) exporter.clear() run_agent_and_tool() assert [exported_attributes(span)[SpanAttributes.INPUT_VALUE] for span in exporter.get_finished_spans()] == [ "__REDACTED__", "__REDACTED__", ] @pytest.mark.parametrize("disabled_name,expected_kinds", [("", ["AGENT", "TOOL"]), ("atomic_agents", [])]) def test_real_launcher_respects_disabled_name(disabled_name: str, expected_kinds: list[str], tmp_path: Path) -> None: executable_name = "opentelemetry-instrument.exe" if os.name == "nt" else "opentelemetry-instrument" launcher = Path(sysconfig.get_path("scripts")) / executable_name environment = os.environ.copy() environment.update( OTEL_TRACES_EXPORTER="none", OTEL_METRICS_EXPORTER="none", OTEL_LOGS_EXPORTER="none", OTEL_PYTHON_DISABLED_INSTRUMENTATIONS=disabled_name, OPENINFERENCE_HIDE_INPUTS="true", OPENINFERENCE_HIDE_OUTPUTS="true", ) script = f""" import sys sys.path.insert(0, {str(Path(__file__).parent)!r}) from test_opentelemetry_discovery import exported_attributes, run_agent_and_tool from atomic_agents.observability import AtomicAgentsInstrumentor from opentelemetry import trace from opentelemetry.sdk.trace import ReadableSpan, TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter assert AtomicAgentsInstrumentor().is_instrumented_by_opentelemetry is {bool(expected_kinds)!r} provider = TracerProvider() exporter = InMemorySpanExporter() provider.add_span_processor(SimpleSpanProcessor(exporter)) trace.set_tracer_provider(provider) run_agent_and_tool() spans = exporter.get_finished_spans() assert len(spans) == {len(expected_kinds)} assert [exported_attributes(span)['openinference.span.kind'] for span in spans] == {expected_kinds!r} assert [exported_attributes(span)['input.value'] for span in spans] == {['__REDACTED__'] * len(expected_kinds)!r} AtomicAgentsInstrumentor().uninstrument() provider.shutdown() print('launcher assertions passed') """ script_path = tmp_path / "launcher_probe.py" script_path.write_text(script, encoding="utf-8") result = subprocess.run( [str(launcher), sys.executable, str(script_path)], env=environment, capture_output=True, text=True, timeout=45 ) assert result.returncode == 0, result.stdout + result.stderr assert "launcher assertions passed" in result.stdout assert "Traceback" not in result.stderr def test_normal_import_does_not_require_optional_observability_packages() -> None: script = """ import sys sys.modules['openinference'] = None sys.modules['opentelemetry'] = None import atomic_agents import atomic_agents.observability assert 'atomic_agents.observability.tracing' not in sys.modules print('optional-free import passed') """ result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, timeout=45) assert result.returncode == 0, result.stdout + result.stderr assert "optional-free import passed" in result.stdout ``` ### File: atomic-agents/tests/observability/test_tracing.py ```python import abc import asyncio import contextvars import gc import json import logging import types from concurrent.futures import ThreadPoolExecutor from importlib.metadata import entry_points from typing import Dict, List from unittest.mock import Mock, PropertyMock import instructor import pytest from openinference.instrumentation import OITracer, TraceConfig, using_metadata, using_session from openinference.semconv.trace import OpenInferenceSpanKindValues, SpanAttributes from opentelemetry import trace as trace_api from opentelemetry.sdk.trace import ReadableSpan, TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import StatusCode from pydantic import Field, ValidationError from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema, BaseTool from atomic_agents.context import BaseDynamicContextProvider, SystemPromptGenerator from atomic_agents.observability import AtomicAgentsInstrumentor import atomic_agents.observability as observability from atomic_agents.observability import agent_spans, installation, span_attributes, span_payloads, tracing as tracing class QuestionSchema(BaseIOSchema): """A question for the agent.""" question: str = Field(...) class AnswerSchema(BaseIOSchema): """An answer from the agent.""" answer: str = Field(...) class CalculationInputSchema(BaseIOSchema): """Expression to evaluate.""" expression: str = Field(...) class CalculationOutputSchema(BaseIOSchema): """Result of the evaluation.""" result: str = Field(...) class CalculatorTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): def run(self, params: CalculationInputSchema, precision: int = 2) -> CalculationOutputSchema: return CalculationOutputSchema(result=f"{len(params.expression):.{precision}f}") class AsyncCalculatorTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: return asyncio.run(self.arun(params)) async def arun(self, params: CalculationInputSchema) -> CalculationOutputSchema: return CalculationOutputSchema(result=str(len(params.expression))) class ThreadedCalculatorTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: return CalculationOutputSchema(result=str(len(params.expression))) async def run_async(self, params: CalculationInputSchema) -> CalculationOutputSchema: return await asyncio.to_thread(self.run, params) class RunMixin: def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: return CalculationOutputSchema(result="mixin") class MixinCalculatorTool(RunMixin, BaseTool[CalculationInputSchema, CalculationOutputSchema]): pass class StaticRunTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): @staticmethod def run(params: CalculationInputSchema) -> CalculationOutputSchema: return CalculationOutputSchema(result="static") class FailingTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: raise ValueError("boom") class AnswerStore: """Container object pydantic cannot serialize to JSON.""" def __init__(self): self.values: List[int] = [] class UnserializableInputSchema(BaseIOSchema): """Input carrying an arbitrary object.""" model_config = {"arbitrary_types_allowed": True} store: AnswerStore = Field(...) class EchoTool(BaseTool[UnserializableInputSchema, CalculationOutputSchema]): def run(self, params: UnserializableInputSchema) -> CalculationOutputSchema: return CalculationOutputSchema(result="echo") @pytest.fixture def exporter(): exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) exporter.provider = provider instrumentor = AtomicAgentsInstrumentor() instrumentor.instrument(tracer_provider=provider) yield exporter instrumentor.uninstrument() @pytest.fixture def provider_tracer(exporter): return exporter.provider.get_tracer("test") def sync_client(response=None, partials=None): client = Mock(spec=instructor.Instructor) client.chat = Mock() client.chat.completions = Mock() client.chat.completions.create = Mock(return_value=response or AnswerSchema(answer="42")) def create_partial(**kwargs): for partial in partials or [AnswerSchema(answer="4"), AnswerSchema(answer="42")]: yield partial client.chat.completions.create_partial = create_partial return client def async_client(response=None, partials=None): client = Mock(spec=instructor.AsyncInstructor) client.chat = Mock() client.chat.completions = Mock() async def create(**kwargs): return response or AnswerSchema(answer="42") async def create_partial(**kwargs): for partial in partials or [AnswerSchema(answer="4"), AnswerSchema(answer="42")]: yield partial client.chat.completions.create = create client.chat.completions.create_partial = create_partial return client def make_agent(client) -> AtomicAgent[QuestionSchema, AnswerSchema]: return AtomicAgent[QuestionSchema, AnswerSchema](AgentConfig(client=client, model="test-model")) def spans_of(exporter) -> List[ReadableSpan]: return list(exporter.get_finished_spans()) def kind_of(span: ReadableSpan) -> str: return span.attributes[SpanAttributes.OPENINFERENCE_SPAN_KIND] def input_json(span: ReadableSpan) -> Dict: return json.loads(span.attributes[SpanAttributes.INPUT_VALUE]) def output_json(span: ReadableSpan) -> Dict: return json.loads(span.attributes[SpanAttributes.OUTPUT_VALUE]) def test_run_emits_agent_span_with_io_and_framework_attributes(exporter): agent = make_agent(sync_client()) response = agent.run(QuestionSchema(question="6*7?")) assert response.answer == "42" (span,) = spans_of(exporter) assert span.name == "invoke_agent AtomicAgent" assert kind_of(span) == OpenInferenceSpanKindValues.AGENT.value assert span.attributes[SpanAttributes.AGENT_NAME] == "AtomicAgent" assert span.attributes[span_attributes.ATTR_AGENT_CLASS] == "AtomicAgent[QuestionSchema, AnswerSchema]" assert input_json(span) == {"question": "6*7?"} assert output_json(span) == {"answer": "42"} assert span.attributes[SpanAttributes.INPUT_MIME_TYPE] == "application/json" assert span.attributes[SpanAttributes.OUTPUT_MIME_TYPE] == "application/json" assert span.attributes[span_attributes.ATTR_MODEL] == "test-model" assert span.attributes[span_attributes.ATTR_MODE] == agent.mode.value assert span.attributes[span_attributes.ATTR_INPUT_SCHEMA] == "QuestionSchema" assert span.attributes[span_attributes.ATTR_OUTPUT_SCHEMA] == "AnswerSchema" assert span.attributes[span_attributes.ATTR_HISTORY_MESSAGE_COUNT] == 0 assert span.attributes[span_attributes.ATTR_CONTEXT_PROVIDERS] == () assert len(span.attributes[span_attributes.ATTR_SYSTEM_PROMPT_SHA256]) == 64 assert span.status.status_code == StatusCode.OK def test_second_turn_reports_history_size_entering_the_turn(exporter): agent = make_agent(sync_client()) agent.run(QuestionSchema(question="first")) agent.run(QuestionSchema(question="second")) first, second = spans_of(exporter) assert first.attributes[span_attributes.ATTR_HISTORY_MESSAGE_COUNT] == 0 assert second.attributes[span_attributes.ATTR_HISTORY_MESSAGE_COUNT] == 2 def test_agent_without_system_role_has_no_prompt_hash(exporter): agent = AtomicAgent[QuestionSchema, AnswerSchema](AgentConfig(client=sync_client(), model="m", system_role=None)) agent.run(QuestionSchema(question="?")) (span,) = spans_of(exporter) assert span_attributes.ATTR_SYSTEM_PROMPT_SHA256 not in span.attributes def test_abstract_intermediate_tool_base_is_not_wrapped(exporter): class AbstractCalculator(BaseTool[CalculationInputSchema, CalculationOutputSchema]): @abc.abstractmethod def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: raise NotImplementedError assert not hasattr(AbstractCalculator.__dict__["run"], installation.TRACED_MARKER) def test_run_without_input_sets_no_input_attribute(exporter): agent = make_agent(sync_client()) agent.run() (span,) = spans_of(exporter) assert SpanAttributes.INPUT_VALUE not in span.attributes assert output_json(span) == {"answer": "42"} def test_run_failure_propagates_and_marks_span_error(exporter): client = sync_client() client.chat.completions.create = Mock(side_effect=RuntimeError("provider down")) agent = make_agent(client) with pytest.raises(RuntimeError, match="provider down"): agent.run(QuestionSchema(question="?")) (span,) = spans_of(exporter) assert span.status.status_code == StatusCode.ERROR assert [event.name for event in span.events] == ["exception"] def test_subclassed_agent_is_named_after_the_subclass_and_emits_one_span(exporter): class ResearchAgent(AtomicAgent[QuestionSchema, AnswerSchema]): def run(self, user_input=None): return super().run(user_input) agent = ResearchAgent(AgentConfig(client=sync_client(), model="test-model")) agent.run(QuestionSchema(question="?")) (span,) = spans_of(exporter) assert span.name == "invoke_agent ResearchAgent" def test_agent_span_nests_under_a_user_chain_span(exporter, provider_tracer): agent = make_agent(sync_client()) with provider_tracer.start_as_current_span("turn") as turn: agent.run(QuestionSchema(question="?")) agent_span, turn_span = spans_of(exporter) assert agent_span.parent.span_id == turn.get_span_context().span_id assert turn_span.name == "turn" def test_manual_setup_from_the_guide_groups_agents_and_tools_under_one_chain_span(): exporter = InMemorySpanExporter() tracer_provider = TracerProvider() tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) trace_config = TraceConfig() instrumentor = AtomicAgentsInstrumentor() instrumentor.instrument(tracer_provider=tracer_provider, config=trace_config) tracer = OITracer(tracer_provider.get_tracer(__name__), config=trace_config) planner = make_agent(sync_client()) answerer = make_agent(sync_client()) calculator = CalculatorTool() try: with tracer.start_as_current_span("answer_question", openinference_span_kind="chain") as turn: turn.set_input("what is 2 + 2?") plan = planner.run(QuestionSchema(question="what is 2 + 2?")) calculation = calculator.run(CalculationInputSchema(expression=plan.answer)) answer = answerer.run(QuestionSchema(question=str(calculation.result))) turn.set_output(answer.answer) finally: instrumentor.uninstrument() spans = spans_of(exporter) chain_span = next(span for span in spans if kind_of(span) == OpenInferenceSpanKindValues.CHAIN.value) children = [span for span in spans if span is not chain_span] assert [kind_of(span) for span in children] == ["AGENT", "TOOL", "AGENT"] assert all(span.parent.span_id == chain_span.get_span_context().span_id for span in children) assert chain_span.attributes[SpanAttributes.INPUT_VALUE] == "what is 2 + 2?" assert chain_span.attributes[SpanAttributes.OUTPUT_VALUE] == answer.answer def test_llm_call_inside_run_becomes_child_of_agent_span(exporter, provider_tracer): def create(**kwargs): with provider_tracer.start_as_current_span("llm"): return AnswerSchema(answer="42") client = sync_client() client.chat.completions.create = create agent = make_agent(client) agent.run(QuestionSchema(question="?")) llm_span, agent_span = spans_of(exporter) assert llm_span.parent.span_id == agent_span.context.span_id def test_run_stream_records_final_output_and_keeps_child_context_per_step(exporter, provider_tracer): current_between_yields = [] def create_partial(**kwargs): yield AnswerSchema(answer="4") with provider_tracer.start_as_current_span("llm-chunk"): yield AnswerSchema(answer="42") client = sync_client() client.chat.completions.create_partial = create_partial agent = make_agent(client) partials = [] for partial in agent.run_stream(QuestionSchema(question="?")): partials.append(partial.answer) current_between_yields.append(trace_api.get_current_span()) assert partials == ["4", "42"] assert all(span is trace_api.INVALID_SPAN for span in current_between_yields) chunk_span, agent_span = spans_of(exporter) assert agent_span.name == "invoke_agent AtomicAgent" assert chunk_span.parent.span_id == agent_span.context.span_id assert output_json(agent_span) == {"answer": "42"} assert agent_span.status.status_code == StatusCode.OK def test_run_stream_preserves_the_generator_return_value(exporter): agent = make_agent(sync_client()) def consume(): final = yield from agent.run_stream(QuestionSchema(question="?")) return final generator = consume() with pytest.raises(StopIteration) as stop: while True: next(generator) assert stop.value.value.answer == "42" def test_run_stream_with_no_partials_ends_ok_without_output(exporter): client = sync_client() client.chat.completions.create_partial = lambda **kwargs: iter([]) agent = make_agent(client) assert list(agent.run_stream(QuestionSchema(question="?"))) == [] (span,) = spans_of(exporter) assert span.status.status_code == StatusCode.OK assert SpanAttributes.OUTPUT_VALUE not in span.attributes def test_run_stream_closed_early_ends_span_and_closes_inner_generator(exporter): inner_closed = [] def create_partial(**kwargs): try: yield AnswerSchema(answer="4") yield AnswerSchema(answer="42") finally: inner_closed.append(True) client = sync_client() client.chat.completions.create_partial = create_partial agent = make_agent(client) for partial in agent.run_stream(QuestionSchema(question="?")): break assert inner_closed == [True] (span,) = spans_of(exporter) assert span.attributes[span_attributes.ATTR_STREAM_CLOSED_EARLY] is True assert output_json(span) == {"answer": "4"} assert span.status.status_code == StatusCode.OK def test_run_stream_failure_mid_stream_marks_error(exporter): def create_partial(**kwargs): yield AnswerSchema(answer="4") raise RuntimeError("stream broke") client = sync_client() client.chat.completions.create_partial = create_partial agent = make_agent(client) with pytest.raises(RuntimeError, match="stream broke"): list(agent.run_stream(QuestionSchema(question="?"))) (span,) = spans_of(exporter) assert span.status.status_code == StatusCode.ERROR @pytest.mark.asyncio async def test_run_async_emits_agent_span(exporter): agent = make_agent(async_client()) response = await agent.run_async(QuestionSchema(question="?")) assert response.answer == "42" (span,) = spans_of(exporter) assert span.name == "invoke_agent AtomicAgent" assert output_json(span) == {"answer": "42"} assert span.status.status_code == StatusCode.OK @pytest.mark.asyncio async def test_run_async_stream_records_last_partial_and_nests_children(exporter, provider_tracer): async def create_partial(**kwargs): yield AnswerSchema(answer="4") with provider_tracer.start_as_current_span("llm-chunk"): yield AnswerSchema(answer="42") client = async_client() client.chat.completions.create_partial = create_partial agent = make_agent(client) partials = [partial.answer async for partial in agent.run_async_stream(QuestionSchema(question="?"))] assert partials == ["4", "42"] chunk_span, agent_span = spans_of(exporter) assert agent_span.name == "invoke_agent AtomicAgent" assert chunk_span.parent.span_id == agent_span.context.span_id assert output_json(agent_span) == {"answer": "42"} @pytest.mark.asyncio async def test_run_async_stream_closed_early_ends_span(exporter): agent = make_agent(async_client()) stream = agent.run_async_stream(QuestionSchema(question="?")) first = await stream.__anext__() await stream.aclose() assert first.answer == "4" (span,) = spans_of(exporter) assert span.attributes[span_attributes.ATTR_STREAM_CLOSED_EARLY] is True assert span.status.status_code == StatusCode.OK @pytest.mark.asyncio async def test_concurrent_async_runs_keep_separate_spans(exporter): first = make_agent(async_client(response=AnswerSchema(answer="one"))) second = make_agent(async_client(response=AnswerSchema(answer="two"))) await asyncio.gather(first.run_async(QuestionSchema(question="1")), second.run_async(QuestionSchema(question="2"))) outputs = {input_json(span)["question"]: output_json(span)["answer"] for span in spans_of(exporter)} assert outputs == {"1": "one", "2": "two"} assert all(span.parent is None for span in spans_of(exporter)) def test_tool_defined_before_instrumentation_emits_tool_span(exporter): tool = CalculatorTool() result = tool.run(CalculationInputSchema(expression="6*7"), precision=1) assert result.result == "3.0" (span,) = spans_of(exporter) assert span.name == "execute_tool CalculationInputSchema" assert kind_of(span) == OpenInferenceSpanKindValues.TOOL.value assert span.attributes[SpanAttributes.TOOL_NAME] == "CalculationInputSchema" assert span.attributes[SpanAttributes.TOOL_DESCRIPTION] == "Expression to evaluate." assert json.loads(span.attributes[SpanAttributes.TOOL_PARAMETERS])["title"] == "CalculationInputSchema" assert input_json(span) == {"expression": "6*7"} assert output_json(span) == {"result": "3.0"} assert span.status.status_code == StatusCode.OK def test_tool_accepts_keyword_params(exporter): result = CalculatorTool().run(params=CalculationInputSchema(expression="1+1")) assert result.result == "3.00" (span,) = spans_of(exporter) assert input_json(span) == {"expression": "1+1"} def test_tool_defined_after_instrumentation_emits_tool_span(exporter): class LateTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: return CalculationOutputSchema(result="late") LateTool().run(CalculationInputSchema(expression="x")) (span,) = spans_of(exporter) assert span.name == "execute_tool CalculationInputSchema" assert LateTool().input_schema is CalculationInputSchema def test_subclass_calling_super_run_emits_one_span(exporter): class LoudCalculatorTool(CalculatorTool): def run(self, params: CalculationInputSchema, precision: int = 2) -> CalculationOutputSchema: return super().run(params, precision) LoudCalculatorTool().run(CalculationInputSchema(expression="2*2")) (span,) = spans_of(exporter) assert span.name == "execute_tool CalculationInputSchema" def test_sync_run_delegating_to_arun_emits_one_span(exporter): AsyncCalculatorTool().run(CalculationInputSchema(expression="2*2")) (span,) = spans_of(exporter) assert span.name == "execute_tool CalculationInputSchema" assert output_json(span) == {"result": "3"} @pytest.mark.asyncio async def test_run_async_delegating_to_thread_emits_one_span(exporter): await ThreadedCalculatorTool().run_async(CalculationInputSchema(expression="3*3")) (span,) = spans_of(exporter) assert span.name == "execute_tool CalculationInputSchema" assert output_json(span) == {"result": "3"} def test_mixin_defined_run_is_traced(exporter): MixinCalculatorTool().run(CalculationInputSchema(expression="x")) (span,) = spans_of(exporter) assert span.name == "execute_tool CalculationInputSchema" def test_static_run_is_left_alone(exporter): assert StaticRunTool.run(CalculationInputSchema(expression="x")).result == "static" assert spans_of(exporter) == [] def test_tool_failure_propagates_same_exception_and_marks_error(exporter): with pytest.raises(ValueError, match="boom") as raised: FailingTool().run(CalculationInputSchema(expression="x")) (span,) = spans_of(exporter) assert span.status.status_code == StatusCode.ERROR assert raised.value.args == ("boom",) assert span.attributes["error.type"] == "ValueError" assert [event.name for event in span.events] == ["exception"] @pytest.mark.asyncio async def test_tool_run_async_cancellation_leaves_span_unset(exporter): started = asyncio.Event() release = asyncio.Event() class WaitingTool(CalculatorTool): async def run_async(self, params: CalculationInputSchema) -> CalculationOutputSchema: started.set() await release.wait() return CalculationOutputSchema(result="never") task = asyncio.create_task(WaitingTool().run_async(CalculationInputSchema(expression="x"))) await asyncio.wait_for(started.wait(), timeout=2) task.cancel() with pytest.raises(asyncio.CancelledError): await task (span,) = spans_of(exporter) assert span.attributes.get("atomic_agents.cancelled") is True assert span.status.status_code == StatusCode.UNSET assert span.events == () assert "error.type" not in span.attributes @pytest.mark.parametrize("error_type", [GeneratorExit, asyncio.CancelledError]) @pytest.mark.parametrize("method_name", ["run", "run_async"]) @pytest.mark.asyncio async def test_tool_cancellation_propagates_without_error(exporter, monkeypatch, error_type, method_name): def cancelled_result(*args, **kwargs): raise error_type monkeypatch.setattr(CalculationOutputSchema, "__init__", cancelled_result) tool = ThreadedCalculatorTool() with pytest.raises(error_type): result = getattr(tool, method_name)(CalculationInputSchema(expression="x")) if method_name == "run_async": await result (span,) = spans_of(exporter) assert span.attributes.get("atomic_agents.cancelled") is True assert span.status.status_code == StatusCode.UNSET assert span.events == () assert "error.type" not in span.attributes def test_unserializable_input_still_returns_result(exporter, caplog): with caplog.at_level(logging.WARNING): result = EchoTool().run(UnserializableInputSchema(store=AnswerStore())) assert result.result == "echo" (span,) = spans_of(exporter) assert span.status.status_code == StatusCode.OK def test_agent_run_inside_tool_nests_under_tool_span(exporter): agent = make_agent(sync_client()) class AgentTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: return CalculationOutputSchema(result=agent.run(QuestionSchema(question=params.expression)).answer) AgentTool().run(CalculationInputSchema(expression="?")) agent_span, tool_span = spans_of(exporter) assert agent_span.parent.span_id == tool_span.context.span_id def test_hide_inputs_and_outputs_redacts_values_and_drops_mime_types(): exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) instrumentor = AtomicAgentsInstrumentor() instrumentor.instrument(tracer_provider=provider, config=TraceConfig(hide_inputs=True, hide_outputs=True)) try: make_agent(sync_client()).run(QuestionSchema(question="secret")) CalculatorTool().run(CalculationInputSchema(expression="1+1")) finally: instrumentor.uninstrument() for span in spans_of(exporter): assert span.attributes[SpanAttributes.INPUT_VALUE] == "__REDACTED__" assert span.attributes[SpanAttributes.OUTPUT_VALUE] == "__REDACTED__" assert SpanAttributes.INPUT_MIME_TYPE not in span.attributes assert SpanAttributes.OUTPUT_MIME_TYPE not in span.attributes def test_session_and_user_metadata_survive_on_agent_span(exporter): agent = make_agent(sync_client()) with using_session("session-1"), using_metadata({"tenant": "acme"}): agent.run(QuestionSchema(question="?")) (span,) = spans_of(exporter) assert span.attributes[SpanAttributes.SESSION_ID] == "session-1" assert json.loads(span.attributes[SpanAttributes.METADATA]) == {"tenant": "acme"} def test_second_instrument_replaces_provider_and_warns(exporter, caplog): other_provider = TracerProvider() other_exporter = InMemorySpanExporter() other_provider.add_span_processor(SimpleSpanProcessor(other_exporter)) agent = make_agent(sync_client()) agent.run(QuestionSchema(question="first")) installed = installation._installation wrapped_run = AtomicAgent.run with caplog.at_level(logging.WARNING): AtomicAgentsInstrumentor().instrument(tracer_provider=other_provider) agent.run(QuestionSchema(question="second")) assert installation._installation is installed assert AtomicAgent.run is wrapped_run assert len(spans_of(exporter)) == 1 assert len(spans_of(other_exporter)) == 1 assert "tracer provider/config replaced" in caplog.text assert "in-flight spans finish on the provider they started with" in caplog.text assert caplog.records[-1].levelno == logging.WARNING other_provider.shutdown() def test_uninstrument_restores_originals_and_stops_tracing(): original_run = AtomicAgent.__dict__["run"] original_tool_run = CalculatorTool.__dict__["run"] original_init_subclass = BaseTool.__dict__["__init_subclass__"] exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) instrumentor = AtomicAgentsInstrumentor() instrumentor.instrument(tracer_provider=provider) assert instrumentor.is_instrumented_by_opentelemetry instrumentor.uninstrument() instrumentor.uninstrument() assert not instrumentor.is_instrumented_by_opentelemetry assert AtomicAgent.__dict__["run"] is original_run assert CalculatorTool.__dict__["run"] is original_tool_run assert BaseTool.__dict__["__init_subclass__"] is original_init_subclass assert "run" not in MixinCalculatorTool.__dict__ class AfterwardsTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: return CalculationOutputSchema(result="plain") assert AfterwardsTool().input_schema is CalculationInputSchema make_agent(sync_client()).run(QuestionSchema(question="?")) AfterwardsTool().run(CalculationInputSchema(expression="x")) assert spans_of(exporter) == [] def test_phoenix_entry_point_loads_and_instruments_with_only_a_tracer_provider(): (entry_point,) = [point for point in entry_points(group="openinference_instrumentor") if point.name == "atomic_agents"] exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) instrumentor = entry_point.load()() instrumentor.instrument(tracer_provider=provider) try: make_agent(sync_client()).run(QuestionSchema(question="?")) finally: instrumentor.uninstrument() assert isinstance(instrumentor, AtomicAgentsInstrumentor) assert len(spans_of(exporter)) == 1 @pytest.mark.parametrize("options", [{}, {"skip_dep_check": False}]) def test_instrument_without_opentelemetry_raises_install_hint(monkeypatch: pytest.MonkeyPatch, options: dict[str, object]): monkeypatch.setattr(observability, "_opentelemetry_is_installed", lambda: False) with pytest.raises(ImportError, match="atomic-agents\\[observability\\]"): AtomicAgentsInstrumentor().instrument(**options) def test_uninstrument_without_opentelemetry_is_a_no_op(monkeypatch): monkeypatch.setattr(observability, "_opentelemetry_is_installed", lambda: False) AtomicAgentsInstrumentor().uninstrument() assert not AtomicAgentsInstrumentor().is_instrumented_by_opentelemetry def test_openinference_below_the_floor_only_disables_context_with_a_warning(monkeypatch, caplog): monkeypatch.setattr(observability, "version", lambda distribution: "0.1.20") with caplog.at_level(logging.WARNING): AtomicAgentsInstrumentor().instrument() assert "0.1.38 or newer is required" in caplog.text assert AtomicAgentsInstrumentor().is_instrumented_by_opentelemetry assert span_attributes._context_attributes() == {} AtomicAgentsInstrumentor().uninstrument() class CountingProvider(BaseDynamicContextProvider): def __init__(self): super().__init__(title="Counter") self.calls = 0 def get_info(self) -> str: self.calls += 1 return f"call {self.calls}" def test_prompt_hash_identifies_static_sections_despite_dynamic_providers(exporter): generator = SystemPromptGenerator(background=["v1"], context_providers={"counter": CountingProvider()}) agent = AtomicAgent[QuestionSchema, AnswerSchema]( AgentConfig(client=sync_client(), model="m", system_prompt_generator=generator) ) agent.run(QuestionSchema(question="a")) agent.run(QuestionSchema(question="b")) first, second = spans_of(exporter) assert ( first.attributes[span_attributes.ATTR_SYSTEM_PROMPT_SHA256] == second.attributes[span_attributes.ATTR_SYSTEM_PROMPT_SHA256] ) assert first.attributes[span_attributes.ATTR_CONTEXT_PROVIDERS] == ("counter",) def test_prompt_hash_changes_when_static_sections_change(exporter): agent = make_agent(sync_client()) agent.run(QuestionSchema(question="a")) agent.system_prompt_generator.background = ["v2"] agent.run(QuestionSchema(question="b")) first, second = spans_of(exporter) assert ( first.attributes[span_attributes.ATTR_SYSTEM_PROMPT_SHA256] != second.attributes[span_attributes.ATTR_SYSTEM_PROMPT_SHA256] ) def test_run_stream_dropped_without_close_still_ends_span(exporter): agent = make_agent(sync_client()) stream = agent.run_stream(QuestionSchema(question="?")) next(stream) del stream gc.collect() (span,) = spans_of(exporter) assert span.attributes[span_attributes.ATTR_STREAM_CLOSED_EARLY] is True @pytest.mark.asyncio async def test_run_async_stream_cancelled_mid_await_is_not_an_error(exporter): release = asyncio.Event() async def create_partial(**kwargs): yield AnswerSchema(answer="4") await release.wait() yield AnswerSchema(answer="42") client = async_client() client.chat.completions.create_partial = create_partial agent = make_agent(client) async def consume(): async for _ in agent.run_async_stream(QuestionSchema(question="?")): pass task = asyncio.create_task(consume()) await asyncio.sleep(0.01) task.cancel() with pytest.raises(asyncio.CancelledError): await task (span,) = spans_of(exporter) assert span.attributes[span_attributes.ATTR_STREAM_CLOSED_EARLY] is True assert output_json(span) == {"answer": "4"} assert span.status.status_code != StatusCode.ERROR def test_forge_style_thread_pool_tool_emits_one_span_in_the_callers_trace(exporter, provider_tracer): class ForgeStyleTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: with ThreadPoolExecutor() as executor: return executor.submit(contextvars.copy_context().run, asyncio.run, self.run_async(params)).result() async def run_async(self, params: CalculationInputSchema) -> CalculationOutputSchema: return CalculationOutputSchema(result="pooled") with provider_tracer.start_as_current_span("root") as root: ForgeStyleTool().run(CalculationInputSchema(expression="x")) tool_span, root_span = spans_of(exporter) assert tool_span.name == "execute_tool CalculationInputSchema" assert tool_span.parent.span_id == root.get_span_context().span_id assert tool_span.context.trace_id == root_span.context.trace_id def test_mcp_style_generated_tool_class_emits_one_span(exporter): async def arun(self, params): return CalculationOutputSchema(result="mcp") def run(self, params): return asyncio.run(self.arun(params)) GeneratedTool = types.new_class( "GeneratedTool", (BaseTool[CalculationInputSchema, CalculationOutputSchema],), {}, lambda namespace: namespace.update({"run": run, "arun": arun}), ) GeneratedTool().run(CalculationInputSchema(expression="x")) (span,) = spans_of(exporter) assert span.name == "execute_tool CalculationInputSchema" assert output_json(span) == {"result": "mcp"} def test_async_generator_tool_method_is_left_alone(exporter): class StreamingTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: return CalculationOutputSchema(result="x") async def run_async(self, params: CalculationInputSchema): yield CalculationOutputSchema(result="chunk") assert not hasattr(StreamingTool.__dict__["run_async"], installation.TRACED_MARKER) assert hasattr(StreamingTool.__dict__["run"], installation.TRACED_MARKER) def test_instrument_uninstrument_instrument_cycle_traces_tools_defined_between_cycles(): exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) instrumentor = AtomicAgentsInstrumentor() instrumentor.instrument(tracer_provider=provider) instrumentor.uninstrument() class BetweenCyclesTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: return CalculationOutputSchema(result="between") original_run = BetweenCyclesTool.__dict__["run"] instrumentor.instrument(tracer_provider=provider) try: BetweenCyclesTool().run(CalculationInputSchema(expression="x")) finally: instrumentor.uninstrument() assert [span.name for span in spans_of(exporter)] == ["execute_tool CalculationInputSchema"] assert BetweenCyclesTool.__dict__["run"] is original_run @pytest.fixture def configured_tracing(): provider = TracerProvider() exporter = InMemorySpanExporter() provider.add_span_processor(SimpleSpanProcessor(exporter)) yield provider, exporter AtomicAgentsInstrumentor().uninstrument() provider.shutdown() @pytest.mark.parametrize("openinference_enabled", [False, True]) def test_genai_attributes_survive_without_openinference(monkeypatch, configured_tracing, openinference_enabled): provider, exporter = configured_tracing monkeypatch.setattr(span_attributes, "_import_context_attributes", lambda: None) config = tracing.TracingConfig(openinference_attributes=True if openinference_enabled else None) AtomicAgentsInstrumentor().instrument(tracer_provider=provider, config=config) client = sync_client() client.provider = types.SimpleNamespace(value="openai") make_agent(client).run(QuestionSchema(question="question")) CalculatorTool().run(CalculationInputSchema(expression="1+1")) agent_span, tool_span = spans_of(exporter) assert [agent_span.name, tool_span.name] == ["invoke_agent AtomicAgent", "execute_tool CalculationInputSchema"] assert agent_span.attributes["gen_ai.operation.name"] == "invoke_agent" assert agent_span.attributes["gen_ai.agent.name"] == "AtomicAgent" assert agent_span.attributes["gen_ai.request.model"] == "test-model" assert agent_span.attributes["gen_ai.provider.name"] == "openai" assert agent_span.attributes["atomic_agents.agent.class"] == "AtomicAgent[QuestionSchema, AnswerSchema]" assert json.loads(agent_span.attributes["atomic_agents.input"]) == {"question": "question"} assert json.loads(agent_span.attributes["atomic_agents.output"]) == {"answer": "42"} assert tool_span.attributes["gen_ai.operation.name"] == "execute_tool" assert tool_span.attributes["gen_ai.tool.name"] == "CalculationInputSchema" assert tool_span.attributes["atomic_agents.tool.class"] == "CalculatorTool" assert tool_span.attributes["gen_ai.tool.description"] == "Expression to evaluate." assert tool_span.attributes["gen_ai.tool.type"] == "function" assert json.loads(tool_span.attributes["gen_ai.tool.call.arguments"]) == {"expression": "1+1"} assert json.loads(tool_span.attributes["gen_ai.tool.call.result"]) == {"result": "3.00"} assert ("openinference.span.kind" in agent_span.attributes) is openinference_enabled assert ("openinference.span.kind" in tool_span.attributes) is openinference_enabled if openinference_enabled: assert agent_span.attributes["input.value"] == agent_span.attributes["atomic_agents.input"] assert tool_span.attributes["input.value"] == tool_span.attributes["gen_ai.tool.call.arguments"] assert tool_span.attributes["output.value"] == tool_span.attributes["gen_ai.tool.call.result"] else: assert not ({"input.value", "output.value", "agent.name"} & agent_span.attributes.keys()) assert not ({"input.value", "output.value", "tool.name", "tool.parameters"} & tool_span.attributes.keys()) @pytest.mark.parametrize("config_class", [tracing.TracingConfig, TraceConfig]) def test_both_visibility_configs_redact_primary_payloads(configured_tracing, config_class): provider, exporter = configured_tracing AtomicAgentsInstrumentor().instrument(tracer_provider=provider, config=config_class(hide_inputs=True, hide_outputs=True)) make_agent(sync_client()).run(QuestionSchema(question="secret")) CalculatorTool().run(CalculationInputSchema(expression="secret")) agent_span, tool_span = spans_of(exporter) assert agent_span.attributes["atomic_agents.input"] == "__REDACTED__" assert agent_span.attributes["atomic_agents.output"] == "__REDACTED__" assert tool_span.attributes["gen_ai.tool.call.arguments"] == "__REDACTED__" assert tool_span.attributes["gen_ai.tool.call.result"] == "__REDACTED__" @pytest.mark.parametrize("hide_outputs", [False, True]) def test_large_tool_result_is_truncated_unless_hidden(configured_tracing, hide_outputs): provider, exporter = configured_tracing AtomicAgentsInstrumentor().instrument( tracer_provider=provider, config=tracing.TracingConfig(hide_outputs=hide_outputs, openinference_attributes=True) ) class LargeResultTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: return CalculationOutputSchema(result="x" * 10240) result = LargeResultTool().run(CalculationInputSchema(expression="large")) (span,) = spans_of(exporter) serialized = result.model_dump_json() expected = "__REDACTED__" if hide_outputs else serialized[: 4096 - len("...[truncated]")] + "...[truncated]" assert span.attributes["gen_ai.tool.call.result"] == expected assert span.attributes["output.value"] == expected if hide_outputs: assert "gen_ai.tool.call.result.original_length" not in span.attributes assert "output.value.original_length" not in span.attributes else: assert span.attributes["gen_ai.tool.call.result.original_length"] == len(serialized) assert span.attributes["output.value.original_length"] == len(serialized) @pytest.mark.parametrize("agent_method", ["run", "run_async", "run_stream", "run_async_stream"]) @pytest.mark.asyncio async def test_suppression_skips_agent_and_tool_spans(exporter, agent_method): from opentelemetry import context client = async_client() if "async" in agent_method else sync_client() agent = make_agent(client) token = context.attach(context.set_value("suppress_instrumentation", True)) try: result = getattr(agent, agent_method)(QuestionSchema(question="suppressed")) assert await final_agent_answer(agent_method, result) == "42" assert CalculatorTool().run(CalculationInputSchema(expression="1+1")).result == "3.00" assert (await ThreadedCalculatorTool().run_async(CalculationInputSchema(expression="1+1"))).result == "3" finally: context.detach(token) assert spans_of(exporter) == [] def test_failures_record_qualified_error_type(exporter): client = sync_client() client.chat.completions.create.side_effect = ValueError("agent failure") with pytest.raises(ValueError, match="agent failure"): make_agent(client).run(QuestionSchema(question="fail")) with pytest.raises(ValueError, match="boom"): FailingTool().run(CalculationInputSchema(expression="fail")) for span in spans_of(exporter): assert span.attributes["error.type"] == "ValueError" assert span.status.status_code == StatusCode.ERROR assert len(span.events) == 1 def test_tracing_config_environment_and_coercion(monkeypatch): monkeypatch.setenv("OPENINFERENCE_HIDE_INPUTS", "TrUe") monkeypatch.setenv("OPENINFERENCE_HIDE_OUTPUTS", "1") monkeypatch.setenv("ATOMIC_AGENTS_TRACE_MAX_ATTRIBUTE_LENGTH", "128") from atomic_agents.observability import TracingConfig config = TracingConfig.from_environment() assert config == TracingConfig(hide_inputs=True, hide_outputs=True, max_attribute_length=128) assert TracingConfig.coerce(None) == config assert TracingConfig.coerce(config) is config assert TracingConfig.coerce(TraceConfig(hide_inputs=True)) == TracingConfig(hide_inputs=True, max_attribute_length=128) @pytest.mark.parametrize( "value, expected", [("false", False), ("FaLsE", False), ("0", False), ("true", True), ("1", True), ("", True)] ) def test_tracing_config_llm_spans_environment(monkeypatch, value, expected): from atomic_agents.observability import TracingConfig monkeypatch.setenv("ATOMIC_AGENTS_TRACE_LLM_SPANS", value) assert TracingConfig.from_environment().llm_spans is expected assert TracingConfig.coerce(None).llm_spans is expected def test_payload_visibility_preserves_environment_settings(monkeypatch): monkeypatch.setenv("ATOMIC_AGENTS_TRACE_LLM_SPANS", "false") monkeypatch.setenv("ATOMIC_AGENTS_TRACE_MAX_ATTRIBUTE_LENGTH", "128") monkeypatch.setenv("ATOMIC_AGENTS_TRACE_HIDE_OUTPUTS", "true") config = tracing.TracingConfig.coerce(TraceConfig(hide_inputs=True, hide_outputs=False)) assert config.llm_spans is False assert config.max_attribute_length == 128 assert config.hide_inputs is True assert config.hide_outputs is False def test_tracing_config_llm_spans_defaults_to_enabled(monkeypatch): from atomic_agents.observability import TracingConfig monkeypatch.delenv("ATOMIC_AGENTS_TRACE_LLM_SPANS", raising=False) assert TracingConfig.from_environment().llm_spans is True assert TracingConfig().llm_spans is True def test_session_is_mirrored_to_genai_conversation(exporter): with using_session("session-1"): make_agent(sync_client()).run(QuestionSchema(question="?")) CalculatorTool().run(CalculationInputSchema(expression="?")) for span in spans_of(exporter): assert span.attributes["gen_ai.conversation.id"] == "session-1" async def final_agent_answer(agent_method, result): if agent_method == "run_async": return (await result).answer elif agent_method == "run_stream": return list(result)[-1].answer elif agent_method == "run_async_stream": return [partial async for partial in result][-1].answer else: return result.answer @pytest.mark.parametrize( "name, expected", [(None, "AtomicAgent"), ("", "AtomicAgent"), ("orchestrator_planning", "orchestrator_planning")] ) def test_stable_agent_name_is_separate_from_schema_class(exporter, name, expected): agent = AtomicAgent[QuestionSchema, AnswerSchema](AgentConfig(client=sync_client(), name=name)) agent.run(QuestionSchema(question="?")) (span,) = spans_of(exporter) assert span.name == f"invoke_agent {expected}" assert span.attributes["gen_ai.agent.name"] == expected assert span.attributes["agent.name"] == expected assert span.attributes["atomic_agents.agent.class"] == "AtomicAgent[QuestionSchema, AnswerSchema]" @pytest.mark.parametrize("openinference_enabled", [False, True]) def test_conversation_context_overrides_session_on_agent_and_tool(exporter, openinference_enabled): config = tracing.TracingConfig(openinference_attributes=openinference_enabled) AtomicAgentsInstrumentor().instrument(tracer_provider=exporter.provider, config=config) with using_session("session-1"), observability.using_conversation("c-1"): make_agent(sync_client()).run(QuestionSchema(question="?")) CalculatorTool().run(CalculationInputSchema(expression="?")) agent_span, tool_span = spans_of(exporter) assert agent_span.attributes["gen_ai.conversation.id"] == "c-1" assert tool_span.attributes["gen_ai.conversation.id"] == "c-1" assert agent_span.attributes.get("session.id") == ("c-1" if openinference_enabled else None) assert tool_span.attributes.get("session.id") == ("c-1" if openinference_enabled else None) def test_session_is_mirrored_even_with_openinference_attributes_disabled(exporter): config = tracing.TracingConfig(openinference_attributes=False) AtomicAgentsInstrumentor().instrument(tracer_provider=exporter.provider, config=config) with using_session("session-1"): make_agent(sync_client()).run(QuestionSchema(question="?")) CalculatorTool().run(CalculationInputSchema(expression="?")) for span in spans_of(exporter): assert span.attributes["gen_ai.conversation.id"] == "session-1" assert "session.id" not in span.attributes def test_conversation_context_restores_after_nested_exception(exporter): with observability.using_conversation("outer"): with pytest.raises(ValueError), observability.using_conversation("inner"): CalculatorTool().run(CalculationInputSchema(expression="inner")) raise ValueError("reset context") CalculatorTool().run(CalculationInputSchema(expression="outer")) CalculatorTool().run(CalculationInputSchema(expression="unset")) inner, outer, unset = spans_of(exporter) assert inner.attributes["gen_ai.conversation.id"] == "inner" assert outer.attributes["gen_ai.conversation.id"] == "outer" assert "gen_ai.conversation.id" not in unset.attributes @pytest.mark.asyncio async def test_conversation_context_is_task_local(exporter): async def run_in_conversation(conversation_id): with observability.using_conversation(conversation_id): await asyncio.sleep(0) await make_agent(async_client()).run_async(QuestionSchema(question="?")) CalculatorTool().run(CalculationInputSchema(expression="?")) await asyncio.gather(run_in_conversation("first"), run_in_conversation("second")) assert [span.attributes["gen_ai.conversation.id"] for span in spans_of(exporter)] == ["first", "first", "second", "second"] def test_plain_parent_span_owns_agent_and_tool_trace(exporter): with exporter.provider.get_tracer("app").start_as_current_span("turn") as turn: make_agent(sync_client()).run(QuestionSchema(question="?")) CalculatorTool().run(CalculationInputSchema(expression="?")) agent_span, tool_span, parent = spans_of(exporter) assert agent_span.context.trace_id == turn.get_span_context().trace_id assert tool_span.context.trace_id == turn.get_span_context().trace_id assert agent_span.parent.span_id == parent.context.span_id assert tool_span.parent.span_id == parent.context.span_id def test_same_provider_and_config_is_a_silent_no_op(exporter, caplog): installed = installation._installation tracer = installed.tracer with caplog.at_level(logging.DEBUG, logger=tracing.__name__): AtomicAgentsInstrumentor().instrument(tracer_provider=exporter.provider, config=tracing.TracingConfig()) assert installation._installation is installed assert installed.tracer is tracer assert not caplog.records def test_reinstrument_replaces_config_without_rewrapping(exporter, caplog): wrapped_run = AtomicAgent.run agent = make_agent(sync_client()) config = tracing.TracingConfig(hide_inputs=True, openinference_attributes=False) with caplog.at_level(logging.WARNING): AtomicAgentsInstrumentor().instrument(tracer_provider=exporter.provider, config=config) agent.run(QuestionSchema(question="secret")) CalculatorTool().run(CalculationInputSchema(expression="secret")) agent_span, tool_span = spans_of(exporter) assert agent_span.attributes["atomic_agents.input"] == "__REDACTED__" assert tool_span.attributes["gen_ai.tool.call.arguments"] == "__REDACTED__" assert "agent.name" not in agent_span.attributes assert AtomicAgent.run is wrapped_run assert "tracer provider/config replaced" in caplog.text def test_in_flight_span_finishes_on_original_provider(exporter): other_provider = TracerProvider() other_exporter = InMemorySpanExporter() other_provider.add_span_processor(SimpleSpanProcessor(other_exporter)) client = sync_client() def replace_provider(**kwargs): AtomicAgentsInstrumentor().instrument(tracer_provider=other_provider) CalculatorTool().run(CalculationInputSchema(expression="?")) return AnswerSchema(answer="42") client.chat.completions.create.side_effect = replace_provider make_agent(client).run(QuestionSchema(question="?")) assert len(spans_of(exporter)) == 1 assert len(spans_of(other_exporter)) == 1 (agent_span,) = spans_of(exporter) (tool_span,) = spans_of(other_exporter) assert agent_span.name == "invoke_agent AtomicAgent" assert tool_span.parent.span_id == agent_span.context.span_id assert tool_span.context.trace_id == agent_span.context.trace_id other_provider.shutdown() @pytest.mark.parametrize("name", ["model", "mode"]) def test_none_descriptions_are_omitted_without_otel_warnings(exporter, caplog, name): agent = make_agent(sync_client()) setattr(agent, name, None) with agent_spans._agent_span(agent, (), {}): pass (span,) = spans_of(exporter) assert f"atomic_agents.{name}" not in span.attributes assert "gen_ai.request.model" not in span.attributes if name == "model" else True assert "Invalid type" not in caplog.text def test_broken_history_preserves_mandatory_agent_attributes(exporter, monkeypatch, caplog): agent = make_agent(sync_client()) monkeypatch.setattr(type(agent), "history", PropertyMock(side_effect=RuntimeError("history unavailable")), raising=False) with caplog.at_level(logging.DEBUG), agent_spans._agent_span(agent, (), {}): pass (span,) = spans_of(exporter) assert span.attributes["gen_ai.operation.name"] == "invoke_agent" assert span.attributes["gen_ai.agent.name"] == "AtomicAgent" assert span.attributes["atomic_agents.agent.class"] == "AtomicAgent[QuestionSchema, AnswerSchema]" assert span.attributes["gen_ai.request.model"] == "test-model" assert "atomic_agents.history.message_count" not in span.attributes assert sum("history.message_count" in record.message for record in caplog.records) == 1 def test_broken_tool_description_preserves_other_attributes(exporter, monkeypatch): monkeypatch.setattr(CalculatorTool, "tool_description", PropertyMock(side_effect=RuntimeError("unavailable"))) CalculatorTool().run(CalculationInputSchema(expression="1+1")) (span,) = spans_of(exporter) assert span.attributes["gen_ai.tool.name"] == "CalculationInputSchema" assert span.attributes["atomic_agents.tool.class"] == "CalculatorTool" assert span.attributes["tool.name"] == "CalculationInputSchema" assert "tool.parameters" in span.attributes assert "gen_ai.tool.description" not in span.attributes def test_configured_tool_name_is_the_span_join_key(exporter): tool = CalculatorTool() tool.config.title = "calculate" tool.run(CalculationInputSchema(expression="1+1")) (span,) = spans_of(exporter) assert span.name == "execute_tool calculate" assert span.attributes["gen_ai.tool.name"] == span.attributes["tool.name"] == "calculate" assert span.attributes["atomic_agents.tool.class"] == "CalculatorTool" @pytest.mark.parametrize("hide", [False, True]) def test_hidden_payload_is_never_serialized(exporter, caplog, hide): payload = UnserializableInputSchema(store=AnswerStore()) with exporter.provider.get_tracer("test").start_as_current_span("payload") as span: span_payloads._record_payload(span, (("payload", "payload.mime_type"),), payload, hide, tracing.TracingConfig()) (span,) = spans_of(exporter) assert ("Could not serialize tracing payload" in caplog.text) is not hide assert span.attributes.get("payload") == ("__REDACTED__" if hide else None) assert "payload.mime_type" not in span.attributes @pytest.mark.parametrize("hide", [False, True]) def test_hidden_non_model_payload_does_not_call_json_dumps(exporter, monkeypatch, caplog, hide): serialize = Mock(side_effect=ValueError("serialization failed")) monkeypatch.setattr(span_payloads.json, "dumps", serialize) with exporter.provider.get_tracer("test").start_as_current_span("payload") as span: span_payloads._record_payload(span, (("payload", None),), object(), hide, tracing.TracingConfig()) assert serialize.call_count == (0 if hide else 1) assert ("Could not serialize tracing payload" in caplog.text) is not hide @pytest.mark.parametrize( "provider,expected", [ ("azure_openai", "azure.ai.openai"), ("vertexai", "gcp.vertex_ai"), ("gemini", "gcp.gemini"), ("genai", "gcp.gemini"), ("bedrock", "aws.bedrock"), ("mistral", "mistral_ai"), ("xai", "x_ai"), ("openai", "openai"), ("anthropic", "anthropic"), ("cohere", "cohere"), ("groq", "groq"), ("deepseek", "deepseek"), ("perplexity", "perplexity"), ("custom", "custom"), ], ) def test_instructor_provider_names_follow_genai_conventions(exporter, provider, expected): client = sync_client() client.provider = types.SimpleNamespace(value=provider) make_agent(client).run(QuestionSchema(question="question")) (span,) = spans_of(exporter) assert span.attributes["gen_ai.provider.name"] == expected def test_none_context_attributes_are_filtered(exporter, monkeypatch, caplog): monkeypatch.setattr(span_attributes, "_context_attributes", lambda: {"optional": None}) make_agent(sync_client()).run(QuestionSchema(question="question")) (span,) = spans_of(exporter) assert "optional" not in span.attributes assert "gen_ai.provider.name" not in span.attributes assert "Invalid type" not in caplog.text @pytest.mark.parametrize("prefix", ["ATOMIC_AGENTS_TRACE", "OPENINFERENCE"]) def test_visibility_environment_supports_both_spellings(monkeypatch, prefix): monkeypatch.delenv("ATOMIC_AGENTS_TRACE_HIDE_INPUTS", raising=False) monkeypatch.delenv("ATOMIC_AGENTS_TRACE_HIDE_OUTPUTS", raising=False) monkeypatch.setenv(f"{prefix}_HIDE_INPUTS", "TRUE") monkeypatch.setenv(f"{prefix}_HIDE_OUTPUTS", "1") config = tracing.TracingConfig.from_environment() assert config.hide_inputs is True assert config.hide_outputs is True def test_atomic_visibility_environment_overrides_openinference(monkeypatch): monkeypatch.setenv("OPENINFERENCE_HIDE_INPUTS", "true") monkeypatch.setenv("OPENINFERENCE_HIDE_OUTPUTS", "true") monkeypatch.setenv("ATOMIC_AGENTS_TRACE_HIDE_INPUTS", "false") monkeypatch.setenv("ATOMIC_AGENTS_TRACE_HIDE_OUTPUTS", "") config = tracing.TracingConfig.from_environment() assert config.hide_inputs is False assert config.hide_outputs is False def test_error_status_and_recorded_message_are_bounded(configured_tracing): provider, exporter = configured_tracing AtomicAgentsInstrumentor().instrument(tracer_provider=provider, config=tracing.TracingConfig(max_attribute_length=12)) error = ValueError("private" * 100) agent = make_agent(sync_client()) with pytest.raises(ValueError) as raised, agent_spans._agent_span(agent, (), {}): raise error (span,) = spans_of(exporter) expected = "...[truncate" assert raised.value is error assert span.status.description == expected assert span.events[0].attributes["exception.message"] == expected assert len(span.events[0].attributes["exception.stacktrace"]) <= 12 def test_hidden_validation_errors_do_not_export_model_output(configured_tracing): provider, exporter = configured_tracing config = tracing.TracingConfig(hide_outputs=True) AtomicAgentsInstrumentor().instrument(tracer_provider=provider, config=config) try: QuestionSchema.model_validate({"question": ["secret model output"]}) except ValidationError as error: with provider.get_tracer("test").start_as_current_span("failure") as span: span_payloads._record_failure(span, error, config) (span,) = spans_of(exporter) assert span.status.description == "ValidationError: 1 errors (__REDACTED__)" assert span.events[0].attributes["exception.message"] == span.status.description assert "secret model output" not in str(span.events[0].attributes) def test_tool_parameters_use_the_payload_size_limit(configured_tracing): provider, exporter = configured_tracing AtomicAgentsInstrumentor().instrument( tracer_provider=provider, config=tracing.TracingConfig(max_attribute_length=12, openinference_attributes=True) ) CalculatorTool().run(CalculationInputSchema(expression="x")) (span,) = spans_of(exporter) parameters = json.dumps(CalculationInputSchema.model_json_schema()) assert span.attributes["tool.parameters"] == "...[truncate" assert span.attributes["tool.parameters.original_length"] == len(parameters) @pytest.mark.parametrize("value, expected", [("", 4096), ("4k", 4096), ("0", 4096), ("-5", 4096), ("512", 512)]) @pytest.mark.parametrize("visibility", [None, TraceConfig(hide_inputs=True)]) def test_attribute_length_environment_is_tolerant(monkeypatch, caplog, value, expected, visibility): monkeypatch.setenv("ATOMIC_AGENTS_TRACE_MAX_ATTRIBUTE_LENGTH", value) with caplog.at_level(logging.WARNING): config = tracing.TracingConfig.coerce(visibility) assert config.max_attribute_length == expected assert config.hide_inputs is (visibility is not None) warnings = [record for record in caplog.records if record.levelno == logging.WARNING] if expected == 4096: assert len(warnings) == 1 assert repr(value) in warnings[0].message else: assert warnings == [] def test_reinstall_without_config_preserves_redaction(exporter, caplog): config = tracing.TracingConfig(hide_inputs=True, openinference_attributes=False) AtomicAgentsInstrumentor().instrument(tracer_provider=exporter.provider, config=config) other_provider = TracerProvider() other_exporter = InMemorySpanExporter() other_provider.add_span_processor(SimpleSpanProcessor(other_exporter)) caplog.clear() try: with caplog.at_level(logging.WARNING): AtomicAgentsInstrumentor().instrument(tracer_provider=other_provider) assert installation._current_installation().config is config assert installation._current_installation().tracer_provider is other_provider make_agent(sync_client()).run(QuestionSchema(question="secret")) assert spans_of(other_exporter)[0].attributes["atomic_agents.input"] == "__REDACTED__" assert len(caplog.records) == 1 assert "replaced" in caplog.records[0].message finally: other_provider.shutdown() def test_reinstall_without_provider_preserves_provider(exporter, caplog): config = tracing.TracingConfig(hide_outputs=True) with caplog.at_level(logging.WARNING): AtomicAgentsInstrumentor().instrument(config=config) assert installation._current_installation().tracer_provider is exporter.provider caplog.clear() AtomicAgentsInstrumentor().instrument() assert caplog.records == [] make_agent(sync_client()).run(QuestionSchema(question="?")) assert spans_of(exporter)[0].attributes["atomic_agents.output"] == "__REDACTED__" @pytest.mark.asyncio async def test_run_async_cancellation_leaves_span_unset(exporter): started = asyncio.Event() release = asyncio.Event() async def create(**kwargs): started.set() await release.wait() return AnswerSchema(answer="never") client = async_client() client.chat.completions.create = create task = asyncio.create_task(make_agent(client).run_async(QuestionSchema(question="?"))) await asyncio.wait_for(started.wait(), timeout=2) task.cancel() with pytest.raises(asyncio.CancelledError): await task (span,) = spans_of(exporter) assert span.attributes.get("atomic_agents.cancelled") is True assert span.status.status_code == StatusCode.UNSET assert span.events == () assert "error.type" not in span.attributes @pytest.mark.parametrize("error_type", [GeneratorExit, asyncio.CancelledError]) def test_run_cancellation_leaves_span_unset(exporter, error_type): client = sync_client() client.chat.completions.create.side_effect = error_type with pytest.raises(error_type): make_agent(client).run(QuestionSchema(question="?")) (span,) = spans_of(exporter) assert span.attributes.get("atomic_agents.cancelled") is True assert span.status.status_code == StatusCode.UNSET assert span.events == () def test_run_stream_cancellation_marks_closed_early(exporter): def create_partial(**kwargs): yield AnswerSchema(answer="4") raise asyncio.CancelledError client = sync_client() client.chat.completions.create_partial = create_partial stream = make_agent(client).run_stream(QuestionSchema(question="?")) next(stream) with pytest.raises(asyncio.CancelledError): next(stream) (span,) = spans_of(exporter) assert span.attributes.get(span_attributes.ATTR_STREAM_CLOSED_EARLY) is True assert span.status.status_code == StatusCode.UNSET assert span.events == () @pytest.mark.parametrize("broken_name", [False, True]) def test_agent_name_falls_back_without_initialized_base(exporter, broken_name): class UninitializedAgent(AtomicAgent[QuestionSchema, AnswerSchema]): def __init__(self): self.client = sync_client() if broken_name: UninitializedAgent.name = property(lambda self: 1 / 0) agent = UninitializedAgent() try: display_name = span_attributes._agent_display_name(agent) except (AttributeError, ZeroDivisionError): display_name = None assert display_name == "UninitializedAgent" with agent_spans._agent_span(agent, (), {}): pass (span,) = spans_of(exporter) assert span.name == "invoke_agent UninitializedAgent" assert span.attributes["gen_ai.agent.name"] == "UninitializedAgent" @pytest.mark.parametrize("method", ["run", "run_async", "arun"]) @pytest.mark.parametrize("fail", [False, True]) @pytest.mark.asyncio async def test_uninstrument_in_flight_tool_preserves_result_or_error(exporter, method, fail): error = ValueError("original tool error") class LifecycleTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: AtomicAgentsInstrumentor().uninstrument() if fail: raise error return CalculationOutputSchema(result=params.expression) async def run_async(self, params: CalculationInputSchema) -> CalculationOutputSchema: return self.run(params) async def arun(self, params: CalculationInputSchema) -> CalculationOutputSchema: return self.run(params) operation = getattr(LifecycleTool(), method) params = CalculationInputSchema(expression="42") if fail: with pytest.raises(ValueError) as raised: if method == "run": operation(params) else: await operation(params) assert raised.value is error else: result = operation(params) if method == "run" else await operation(params) assert result.result == "42" (span,) = spans_of(exporter) assert span.status.status_code == (StatusCode.ERROR if fail else StatusCode.OK) @pytest.mark.parametrize("method", ["run", "run_async"]) @pytest.mark.parametrize("fail", [False, True]) @pytest.mark.asyncio async def test_uninstrument_in_flight_agent_preserves_result_or_error(exporter, method, fail): error = ValueError("original agent error") def complete(**kwargs): AtomicAgentsInstrumentor().uninstrument() if fail: raise error return AnswerSchema(answer="42") async def complete_async(**kwargs): return complete(**kwargs) client = sync_client() if method == "run" else async_client() client.chat.completions.create = complete if method == "run" else complete_async operation = getattr(make_agent(client), method) question = QuestionSchema(question="?") if fail: with pytest.raises(ValueError) as raised: if method == "run": operation(question) else: await operation(question) assert raised.value is error else: result = operation(question) if method == "run" else await operation(question) assert result.answer == "42" (span,) = spans_of(exporter) assert span.status.status_code == (StatusCode.ERROR if fail else StatusCode.OK) @pytest.mark.parametrize("ending", ["complete", "close", "error", "cancel"]) @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.asyncio async def test_uninstrument_in_flight_stream_preserves_lifecycle(exporter, ending, asynchronous): error = asyncio.CancelledError() if ending == "cancel" else ValueError("original stream error") def partials(**kwargs): yield AnswerSchema(answer="4") if ending in ("error", "cancel"): raise error yield AnswerSchema(answer="42") async def async_partials(**kwargs): for partial in partials(): yield partial client = async_client() if asynchronous else sync_client() client.chat.completions.create_partial = async_partials if asynchronous else partials agent = make_agent(client) question = QuestionSchema(question="?") stream = agent.run_async_stream(question) if asynchronous else agent.run_stream(question) first = await anext(stream) if asynchronous else next(stream) assert first.answer == "4" AtomicAgentsInstrumentor().uninstrument() if ending == "close": if asynchronous: await stream.aclose() else: stream.close() elif ending in ("error", "cancel"): with pytest.raises(type(error)) as raised: if asynchronous: await anext(stream) else: next(stream) assert raised.value is error else: remaining = [item async for item in stream] if asynchronous else list(stream) assert remaining[-1].answer == "42" (span,) = spans_of(exporter) expected_status = {"error": StatusCode.ERROR, "cancel": StatusCode.UNSET}.get(ending, StatusCode.OK) assert span.status.status_code == expected_status assert span.attributes.get("atomic_agents.stream.closed_early", False) == (ending in ("close", "cancel")) @pytest.mark.parametrize("component", ["agent", "tool"]) @pytest.mark.parametrize("fail", [False, True]) def test_in_flight_span_retains_output_privacy_config(exporter, component, fail): initial = tracing.TracingConfig(hide_outputs=True, openinference_attributes=False) AtomicAgentsInstrumentor().instrument(tracer_provider=exporter.provider, config=initial) error = ValueError("private failure") def reconfigure(): AtomicAgentsInstrumentor().instrument( tracer_provider=exporter.provider, config=tracing.TracingConfig(hide_outputs=False, openinference_attributes=True), ) if fail: raise error class ReconfiguringTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: reconfigure() return CalculationOutputSchema(result="private result") def complete(**kwargs): reconfigure() return AnswerSchema(answer="private result") client = sync_client() client.chat.completions.create = complete def operation(): if component == "agent": return make_agent(client).run(QuestionSchema(question="?")) return ReconfiguringTool().run(CalculationInputSchema(expression="?")) if fail: with pytest.raises(ValueError) as raised: operation() assert raised.value is error else: operation() (span,) = spans_of(exporter) assert "private" not in str(span.attributes) assert "output.value" not in span.attributes if fail: assert span.events[0].attributes["exception.message"] == "ValueError: 1 errors (__REDACTED__)" else: output_key = "atomic_agents.output" if component == "agent" else "gen_ai.tool.call.result" assert span.attributes[output_key] == "__REDACTED__" @pytest.mark.parametrize( "component,method", [("agent", "run"), ("agent", "run_async"), ("tool", "run"), ("tool", "run_async"), ("tool", "arun")] ) @pytest.mark.parametrize("fail", [False, True]) @pytest.mark.asyncio async def test_deferred_call_after_uninstrument_preserves_result_or_error(exporter, component, method, fail): error = ValueError("deferred operation failed") def complete(**kwargs): if fail: raise error return AnswerSchema(answer="42") async def complete_async(**kwargs): return complete(**kwargs) class DeferredTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: complete() return CalculationOutputSchema(result=params.expression) async def run_async(self, params: CalculationInputSchema) -> CalculationOutputSchema: return self.run(params) async def arun(self, params: CalculationInputSchema) -> CalculationOutputSchema: return self.run(params) if component == "agent": client = sync_client() if method == "run" else async_client() client.chat.completions.create = complete if method == "run" else complete_async operation = getattr(make_agent(client), method) params = QuestionSchema(question="?") else: operation = getattr(DeferredTool(), method) params = CalculationInputSchema(expression="42") pending = operation(params) if method != "run" else None AtomicAgentsInstrumentor().uninstrument() if fail: with pytest.raises(ValueError) as raised: if pending is None: operation(params) else: await pending assert raised.value is error else: result = operation(params) if pending is None else await pending assert (result.answer if component == "agent" else result.result) == "42" assert spans_of(exporter) == [] @pytest.mark.parametrize("ending", ["complete", "error", "cancel"]) @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.asyncio async def test_deferred_stream_after_uninstrument_preserves_lifecycle(exporter, ending, asynchronous): error = asyncio.CancelledError() if ending == "cancel" else ValueError("deferred stream failed") def partials(**kwargs): yield AnswerSchema(answer="4") if ending != "complete": raise error yield AnswerSchema(answer="42") async def async_partials(**kwargs): for partial in partials(): yield partial client = async_client() if asynchronous else sync_client() client.chat.completions.create_partial = async_partials if asynchronous else partials agent = make_agent(client) question = QuestionSchema(question="?") stream = agent.run_async_stream(question) if asynchronous else agent.run_stream(question) AtomicAgentsInstrumentor().uninstrument() first = await anext(stream) if asynchronous else next(stream) assert first.answer == "4" if ending != "complete": with pytest.raises(type(error)) as raised: if asynchronous: await anext(stream) else: next(stream) assert raised.value is error elif asynchronous: assert [item.answer async for item in stream] == ["42"] else: assert next(stream).answer == "42" with pytest.raises(StopIteration) as stopped: next(stream) assert stopped.value.value.answer == "42" assert spans_of(exporter) == [] @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.asyncio async def test_deferred_public_stream_after_uninstrument_closes_agent_stream(exporter, monkeypatch, asynchronous): instrumentor = AtomicAgentsInstrumentor() instrumentor.uninstrument() method = "run_async_stream" if asynchronous else "run_stream" original = getattr(AtomicAgent, method) inner_streams = [] def capture_stream(*args, **kwargs): stream = original(*args, **kwargs) inner_streams.append(stream) return stream monkeypatch.setattr(AtomicAgent, method, capture_stream) instrumentor.instrument(tracer_provider=exporter.provider) agent = make_agent(async_client() if asynchronous else sync_client()) stream = getattr(agent, method)(QuestionSchema(question="?")) instrumentor.uninstrument() first = await anext(stream) if asynchronous else next(stream) assert first.answer == "4" (inner,) = inner_streams frame = "ag_frame" if asynchronous else "gi_frame" assert getattr(inner, frame) is not None try: if asynchronous: await stream.aclose() else: stream.close() assert getattr(inner, frame) is None assert spans_of(exporter) == [] finally: if asynchronous: await inner.aclose() else: inner.close() ``` ### File: atomic-agents/tests/utils/test_format_tool_message.py ```python import uuid import json from datetime import date, datetime, timezone from decimal import Decimal from enum import Enum from pathlib import Path from pydantic import BaseModel, ConfigDict, field_serializer from instructor.processing.schema import generate_openai_schema import pytest from atomic_agents import BaseIOSchema from atomic_agents.utils import format_tool_message # Mock classes for testing class MockToolCall(BaseModel): """Mock class for testing""" param1: str param2: int def test_format_tool_message_with_provided_tool_id(): tool_call = MockToolCall(param1="test", param2=42) tool_id = "test-tool-id" result = format_tool_message(tool_call, tool_id) assert result == { "id": "test-tool-id", "type": "function", "function": {"name": "MockToolCall", "arguments": '{"param1": "test", "param2": 42}'}, } def test_format_tool_message_without_tool_id(): tool_call = MockToolCall(param1="test", param2=42) result = format_tool_message(tool_call) assert isinstance(result["id"], str) assert len(result["id"]) == 36 # UUID length assert result["type"] == "function" assert result["function"]["name"] == "MockToolCall" assert result["function"]["arguments"] == '{"param1": "test", "param2": 42}' def test_format_tool_message_with_different_tool(): class AnotherToolCall(BaseModel): """Another tool schema""" field1: bool field2: float tool_call = AnotherToolCall(field1=True, field2=3.14) result = format_tool_message(tool_call) assert result["type"] == "function" assert result["function"]["name"] == "AnotherToolCall" assert result["function"]["arguments"] == '{"field1": true, "field2": 3.14}' def test_format_tool_message_id_is_valid_uuid(): tool_call = MockToolCall(param1="test", param2=42) result = format_tool_message(tool_call) try: uuid.UUID(result["id"]) except ValueError: pytest.fail("The generated tool_id is not a valid UUID") def test_format_tool_message_consistent_output(): tool_call = MockToolCall(param1="test", param2=42) tool_id = "fixed-id" result1 = format_tool_message(tool_call, tool_id) result2 = format_tool_message(tool_call, tool_id) assert result1 == result2 def test_format_tool_message_with_complex_model(): class ComplexToolCall(BaseIOSchema): """Mock complex tool call schema""" nested: dict list_field: list tool_call = ComplexToolCall(nested={"key": "value"}, list_field=[1, 2, 3]) result = format_tool_message(tool_call) assert result["function"]["name"] == "ComplexToolCall" assert result["function"]["arguments"] == '{"nested": {"key": "value"}, "list_field": [1, 2, 3]}' def test_format_tool_message_serializes_pydantic_json_types(): class Priority(Enum): HIGH = "high" class Details(BaseModel): created: datetime due: date request_id: uuid.UUID amount: Decimal path: Path priority: Priority class RichToolCall(BaseModel): details: Details tool_call = RichToolCall( details=Details( created=datetime(2026, 1, 2, 3, 4, 5, tzinfo=timezone.utc), due=date(2026, 1, 3), request_id=uuid.UUID("12345678-1234-5678-1234-567812345678"), amount=Decimal("12.50"), path=Path("report.txt"), priority=Priority.HIGH, ) ) result = format_tool_message(tool_call, "request-1") assert result["function"]["name"] == "RichToolCall" assert json.loads(result["function"]["arguments"]) == { "details": { "created": "2026-01-02T03:04:05Z", "due": "2026-01-03", "request_id": "12345678-1234-5678-1234-567812345678", "amount": "12.50", "path": "report.txt", "priority": "high", } } def test_format_tool_message_honors_json_field_serializers(): class SerializedToolCall(BaseModel): value: int @field_serializer("value", when_used="json") def serialize_value(self, value): return f"item-{value}" result = format_tool_message(SerializedToolCall(value=42)) assert json.loads(result["function"]["arguments"]) == {"value": "item-42"} @pytest.mark.parametrize( "config", [ ConfigDict(title="search_documents"), ConfigDict(model_title_generator=lambda cls: "generated_search"), ConfigDict(json_schema_extra={"title": "custom_search"}), ], ) def test_format_tool_message_name_matches_advertised_schema(config): class SearchArguments(BaseIOSchema): """Search the available documents.""" model_config = config query: str call = SearchArguments(query="example") result = format_tool_message(call, "search-1") assert result["function"]["name"] == generate_openai_schema(SearchArguments)["name"] assert json.loads(result["function"]["arguments"]) == {"query": "example"} assert result["id"] == "search-1" if __name__ == "__main__": pytest.main() ``` ### File: atomic-agents/tests/utils/test_token_counter.py ```python import sys from unittest.mock import MagicMock, call, patch import pytest from litellm import token_counter as litellm_token_counter from atomic_agents.utils.token_counter import ( TokenCounter, TokenCountResult, TokenCountError, get_token_counter, ) class TestTokenCountResult: """Tests for TokenCountResult named tuple.""" def test_creation_with_all_fields(self): result = TokenCountResult( total=100, system_prompt=30, history=50, tools=20, model="gpt-4", max_tokens=8192, utilization=0.0122, ) assert result.total == 100 assert result.system_prompt == 30 assert result.history == 50 assert result.tools == 20 assert result.model == "gpt-4" assert result.max_tokens == 8192 assert result.utilization == 0.0122 def test_optional_fields_default_to_none(self): result = TokenCountResult( total=100, system_prompt=30, history=50, tools=20, model="gpt-4", ) assert result.max_tokens is None assert result.utilization is None def test_named_tuple_unpacking(self): result = TokenCountResult( total=100, system_prompt=30, history=50, tools=20, model="gpt-4", ) total, system_prompt, history, tools, model, max_tokens, utilization = result assert total == 100 assert system_prompt == 30 assert history == 50 assert tools == 20 assert model == "gpt-4" assert max_tokens is None assert utilization is None def test_access_by_index(self): result = TokenCountResult( total=100, system_prompt=30, history=50, tools=20, model="gpt-4", ) assert result[0] == 100 # total assert result[1] == 30 # system_prompt assert result[2] == 50 # history assert result[3] == 20 # tools assert result[4] == "gpt-4" # model assert result[5] is None # max_tokens assert result[6] is None # utilization class TestTokenCounter: """Tests for TokenCounter class.""" @patch("litellm.token_counter") def test_count_messages(self, mock_token_counter): mock_token_counter.return_value = 42 counter = TokenCounter() messages = [{"role": "user", "content": "Hello"}] result = counter.count_messages("gpt-4", messages) assert result == 42 mock_token_counter.assert_called_once_with(model="gpt-4", messages=messages) def test_count_messages_reports_missing_litellm(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setitem(sys.modules, "litellm", None) with pytest.raises( ImportError, match=r"litellm is required for token counting\. Install it with: pip install litellm" ) as error: TokenCounter().count_messages("gpt-4", [{"role": "user", "content": "Hello"}]) assert not isinstance(error.value, TokenCountError) assert isinstance(error.value.__cause__, ModuleNotFoundError) assert "litellm" in str(error.value.__cause__) @patch("litellm.token_counter") def test_count_messages_multiple(self, mock_token_counter): mock_token_counter.return_value = 100 counter = TokenCounter() messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}, ] result = counter.count_messages("gpt-4", messages) assert result == 100 mock_token_counter.assert_called_once() @patch("litellm.token_counter") def test_count_text(self, mock_token_counter): mock_token_counter.return_value = 5 counter = TokenCounter() result = counter.count_text("gpt-4", "Hello world") assert result == 5 # Should wrap text in a message mock_token_counter.assert_called_once_with(model="gpt-4", messages=[{"role": "user", "content": "Hello world"}]) @patch("litellm.get_model_info") def test_get_max_tokens(self, mock_get_model_info): mock_get_model_info.return_value = {"max_input_tokens": 128000, "max_tokens": 16384} counter = TokenCounter() result = counter.get_max_tokens("gpt-4") assert result == 128000 mock_get_model_info.assert_called_once_with("gpt-4") def test_get_max_tokens_reports_missing_litellm(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setitem(sys.modules, "litellm", None) with pytest.raises( ImportError, match=r"litellm is required for token counting\. Install it with: pip install litellm" ) as error: TokenCounter().get_max_tokens("gpt-4") assert not isinstance(error.value, TokenCountError) assert isinstance(error.value.__cause__, ModuleNotFoundError) assert "litellm" in str(error.value.__cause__) @pytest.mark.parametrize("model", [None, 123], ids=["none", "integer"]) @patch("litellm.get_model_info") def test_get_max_tokens_rejects_non_string_model_before_dependency_import( self, mock_get_model_info: MagicMock, monkeypatch: pytest.MonkeyPatch, model: int | None ) -> None: monkeypatch.setitem(sys.modules, "litellm", None) with pytest.raises(TypeError, match=f"model must be a string, got {type(model).__name__}"): TokenCounter().get_max_tokens(model) mock_get_model_info.assert_not_called() @patch("litellm.get_model_info") def test_get_max_tokens_falls_back_to_max_tokens(self, mock_get_model_info): mock_get_model_info.return_value = {"max_tokens": 8192} counter = TokenCounter() result = counter.get_max_tokens("gpt-4") assert result == 8192 mock_get_model_info.assert_called_once_with("gpt-4") @patch("litellm.get_model_info") def test_get_max_tokens_falls_back_when_max_input_tokens_is_none(self, mock_get_model_info): mock_get_model_info.return_value = {"max_input_tokens": None, "max_tokens": 8192} counter = TokenCounter() result = counter.get_max_tokens("gpt-4") assert result == 8192 @patch("litellm.get_model_info") def test_get_max_tokens_zero_input_tokens_returns_zero(self, mock_get_model_info): """Ensure max_input_tokens=0 is not confused with 'missing'.""" mock_get_model_info.return_value = {"max_input_tokens": 0, "max_tokens": 4096} counter = TokenCounter() result = counter.get_max_tokens("custom-model") assert result == 0 @patch("litellm.get_model_info") def test_get_max_tokens_both_keys_missing(self, mock_get_model_info): mock_get_model_info.return_value = {"model_name": "some-model"} counter = TokenCounter() result = counter.get_max_tokens("some-model") assert result is None @patch("litellm.get_model_info") def test_get_max_tokens_unknown_model(self, mock_get_model_info): mock_get_model_info.side_effect = Exception("Unknown model") counter = TokenCounter() result = counter.get_max_tokens("unknown-model") assert result is None @patch("litellm.token_counter") def test_count_messages_with_tools(self, mock_token_counter): mock_token_counter.return_value = 150 counter = TokenCounter() messages = [{"role": "user", "content": "Hello"}] tools = [{"type": "function", "function": {"name": "test_fn"}}] result = counter.count_messages("gpt-4", messages, tools=tools) assert result == 150 mock_token_counter.assert_called_once_with(model="gpt-4", messages=messages, tools=tools) @patch("litellm.token_counter") def test_count_messages_raises_token_count_error(self, mock_token_counter): mock_token_counter.side_effect = Exception("API error") counter = TokenCounter() with pytest.raises(TokenCountError) as exc_info: counter.count_messages("gpt-4", [{"role": "user", "content": "test"}]) assert "Failed to count tokens for model 'gpt-4'" in str(exc_info.value) def test_count_messages_raises_value_error_for_empty_model(self): counter = TokenCounter() with pytest.raises(ValueError) as exc_info: counter.count_messages("", [{"role": "user", "content": "test"}]) assert "model is required" in str(exc_info.value) @patch("litellm.get_model_info") @patch("litellm.token_counter") def test_count_context(self, mock_token_counter, mock_get_model_info): # Counting each group independently would report 30 + 70 = 100, but the # complete request only has 95 tokens because framing overhead is shared. mock_token_counter.side_effect = [30, 95] mock_get_model_info.return_value = {"max_input_tokens": 8192, "max_tokens": 4096} counter = TokenCounter() result = counter.count_context( model="gpt-4", system_messages=[{"role": "system", "content": "You are helpful"}], history_messages=[{"role": "user", "content": "Hello"}], ) assert result.total == 95 assert result.system_prompt == 30 assert result.history == 65 assert result.tools == 0 assert result.model == "gpt-4" assert result.max_tokens == 8192 assert result.utilization == pytest.approx(95 / 8192) @patch("litellm.get_model_info") @patch("litellm.token_counter") def test_count_context_with_tools(self, mock_token_counter, mock_get_model_info): # system=30, complete messages=95, complete request with tools=140 mock_token_counter.side_effect = [30, 95, 140] mock_get_model_info.return_value = {"max_input_tokens": 8192, "max_tokens": 4096} counter = TokenCounter() system_messages = [{"role": "system", "content": "You are helpful"}] history_messages = [{"role": "user", "content": "Hello"}] tools = [{"type": "function", "function": {"name": "test_fn"}}] result = counter.count_context( model="gpt-4", system_messages=system_messages, history_messages=history_messages, tools=tools, ) assert result.system_prompt == 30 assert result.history == 65 assert result.tools == 45 assert result.total == 140 assert result.model == "gpt-4" assert mock_token_counter.call_args_list == [ call(model="gpt-4", messages=system_messages), call(model="gpt-4", messages=system_messages + history_messages), call(model="gpt-4", messages=system_messages + history_messages, tools=tools), ] @patch("litellm.get_model_info") def test_count_context_matches_complete_litellm_request(self, mock_get_model_info: MagicMock) -> None: mock_get_model_info.return_value = {"max_input_tokens": 128000, "max_tokens": 16384} model = "gpt-4o-mini" system_messages = [{"role": "system", "content": "You are a concise research assistant."}] history_messages = [ {"role": "user", "content": "Summarize retrieval-augmented generation."}, {"role": "assistant", "content": "It grounds model responses in retrieved evidence."}, ] tools = [ { "type": "function", "function": { "name": "search_papers", "description": "Search for relevant papers", "parameters": { "type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"], }, }, } ] complete_request_tokens = litellm_token_counter( model=model, messages=system_messages + history_messages, tools=tools, ) result = TokenCounter().count_context( model=model, system_messages=system_messages, history_messages=history_messages, tools=tools, ) assert result.total == complete_request_tokens assert result.total == result.system_prompt + result.history + result.tools @patch("litellm.get_model_info") @patch("litellm.token_counter") def test_count_context_with_tools_and_no_messages( self, mock_token_counter: MagicMock, mock_get_model_info: MagicMock ) -> None: mock_token_counter.side_effect = [60, 10] mock_get_model_info.return_value = {"max_input_tokens": 8192, "max_tokens": 4096} counter = TokenCounter() tools = [{"type": "function", "function": {"name": "test_fn"}}] result = counter.count_context( model="gpt-4", system_messages=[], history_messages=[], tools=tools, ) placeholder = [{"role": "user", "content": ""}] assert result.system_prompt == 0 assert result.history == 0 assert result.tools == 50 assert result.total == 50 assert mock_token_counter.call_args_list == [ call(model="gpt-4", messages=placeholder, tools=tools), call(model="gpt-4", messages=placeholder), ] @patch("litellm.get_model_info") @patch("litellm.token_counter") def test_count_context_empty_system(self, mock_token_counter, mock_get_model_info): mock_token_counter.return_value = 50 mock_get_model_info.return_value = {"max_input_tokens": 4096, "max_tokens": 2048} counter = TokenCounter() result = counter.count_context( model="gpt-3.5-turbo", system_messages=[], # No system prompt history_messages=[{"role": "user", "content": "Hello"}], ) assert result.total == 50 assert result.system_prompt == 0 assert result.history == 50 assert result.model == "gpt-3.5-turbo" assert result.max_tokens == 4096 @patch("litellm.get_model_info") @patch("litellm.token_counter") def test_count_context_no_max_tokens(self, mock_token_counter, mock_get_model_info): mock_token_counter.side_effect = [20, 45] mock_get_model_info.side_effect = Exception("Unknown model") counter = TokenCounter() result = counter.count_context( model="custom-model", system_messages=[{"role": "system", "content": "Test"}], history_messages=[{"role": "user", "content": "Test"}], ) assert result.total == 45 assert result.max_tokens is None assert result.utilization is None @patch("litellm.token_counter") def test_count_messages_different_models(self, mock_token_counter): mock_token_counter.return_value = 10 counter = TokenCounter() # Test various model formats models = [ "gpt-4", "gpt-3.5-turbo", "claude-3-opus-20240229", "anthropic/claude-3-sonnet", "gemini-pro", "gemini/gemini-1.5-pro", ] for model in models: result = counter.count_messages(model, [{"role": "user", "content": "test"}]) assert result == 10 # Verify all models were called assert mock_token_counter.call_count == len(models) @patch("litellm.get_model_info") @patch("litellm.token_counter") def test_count_context_division_by_zero_prevention(self, mock_token_counter, mock_get_model_info): """Test that division by zero is prevented when max_tokens is 0.""" mock_token_counter.side_effect = [20, 45] mock_get_model_info.return_value = {"max_input_tokens": 0, "max_tokens": 0} # Edge case counter = TokenCounter() result = counter.count_context( model="custom-model", system_messages=[{"role": "system", "content": "Test"}], history_messages=[{"role": "user", "content": "Test"}], ) assert result.total == 45 assert result.max_tokens == 0 assert result.utilization is None # Should be None, not raise ZeroDivisionError @patch("litellm.get_model_info") @patch("litellm.token_counter") def test_count_context_empty_history(self, mock_token_counter, mock_get_model_info): """Test counting context with empty history messages.""" mock_token_counter.return_value = 30 # Only system mock_get_model_info.return_value = {"max_input_tokens": 4096, "max_tokens": 2048} counter = TokenCounter() result = counter.count_context( model="gpt-4", system_messages=[{"role": "system", "content": "You are helpful"}], history_messages=[], # Empty history ) assert result.total == 30 assert result.system_prompt == 30 assert result.history == 0 assert result.tools == 0 class TestGetTokenCounter: """Tests for the get_token_counter singleton function.""" def test_get_token_counter_returns_instance(self): """Test that get_token_counter returns a TokenCounter instance.""" counter = get_token_counter() assert isinstance(counter, TokenCounter) def test_get_token_counter_returns_same_instance(self): """Test that get_token_counter returns the same singleton instance.""" counter1 = get_token_counter() counter2 = get_token_counter() assert counter1 is counter2 class TestTokenCountError: """Tests for TokenCountError exception.""" def test_token_count_error_is_exception(self): """Test that TokenCountError is an Exception.""" error = TokenCountError("test error") assert isinstance(error, Exception) def test_token_count_error_message(self): """Test that TokenCountError preserves the error message.""" error = TokenCountError("Custom error message") assert str(error) == "Custom error message" class TestTokenCounterIntegration: """Integration tests that verify the module structure.""" def test_import_from_utils(self): """Test that all exports can be imported from utils.""" from atomic_agents.utils import ( TokenCounter, TokenCountResult, TokenCountError, get_token_counter, ) assert TokenCounter is not None assert TokenCountResult is not None assert TokenCountError is not None assert get_token_counter is not None def test_token_counter_instantiation(self): """Test that TokenCounter can be instantiated without arguments.""" counter = TokenCounter() assert counter is not None if __name__ == "__main__": pytest.main() ``` ================================================================================ END OF DOCUMENT ================================================================================