open-swarm 0.1.1743070217__py3-none-any.whl

swarm/extensions/mcp/mcp_client.py

@@ -0,0 +1,341 @@
+"""
+MCP Client Module
+
+Manages connections and interactions with MCP servers using the MCP Python SDK.
+Redirects MCP server stderr to log files unless debug mode is enabled.
+"""
+
+import asyncio
+import logging
+import os
+from typing import Any, Dict, List, Callable
+from contextlib import contextmanager
+import sys
+import json # Added for result parsing
+
+# Attempt to import mcp types carefully
+try:
+    from mcp import ClientSession, StdioServerParameters # type: ignore
+    from mcp.client.stdio import stdio_client # type: ignore
+    MCP_AVAILABLE = True
+except ImportError:
+    MCP_AVAILABLE = False
+    # Define dummy classes if mcp is not installed
+    class ClientSession: pass
+    class StdioServerParameters: pass
+    def stdio_client(*args, **kwargs): raise ImportError("mcp library not installed")
+
+from ...types import Tool # Use Tool from swarm.types
+from .cache_utils import get_cache
+from ...settings import Settings # Import Swarm settings
+
+# Use Swarm's settings for logging configuration
+swarm_settings = Settings()
+logger = logging.getLogger(__name__)
+logger.setLevel(swarm_settings.log_level.upper()) # Use log level from settings
+# Ensure handler is added only if needed, respecting potential global config
+if not logger.handlers and not logging.getLogger('swarm').handlers:
+     handler = logging.StreamHandler()
+     # Use log format from settings
+     formatter = logging.Formatter(swarm_settings.log_format.value)
+     handler.setFormatter(formatter)
+     logger.addHandler(handler)
+
+class MCPClient:
+    """
+    Manages connections and interactions with MCP servers using the MCP Python SDK.
+    """
+
+    def __init__(self, server_config: Dict[str, Any], timeout: int = 15, debug: bool = False):
+        """
+        Initialize the MCPClient with server configuration.
+
+        Args:
+            server_config (dict): Configuration dictionary for the MCP server.
+            timeout (int): Timeout for operations in seconds.
+            debug (bool): If True, MCP server stderr goes to console; otherwise, suppressed.
+        """
+        if not MCP_AVAILABLE:
+            raise ImportError("The 'mcp-client' library is required for MCP functionality but is not installed.")
+
+        self.command = server_config.get("command", "npx")
+        self.args = server_config.get("args", [])
+        self.env = {**os.environ.copy(), **server_config.get("env", {})}
+        self.timeout = timeout
+        self.debug = debug or swarm_settings.debug # Use instance debug or global debug
+        self._tool_cache: Dict[str, Tool] = {}
+        self.cache = get_cache()
+
+        # Validate command and args types
+        if not isinstance(self.command, str):
+             raise TypeError(f"MCP server command must be a string, got {type(self.command)}")
+        if not isinstance(self.args, list) or not all(isinstance(a, str) for a in self.args):
+             raise TypeError(f"MCP server args must be a list of strings, got {self.args}")
+
+
+        logger.info(f"Initialized MCPClient with command={self.command}, args={self.args}, debug={self.debug}")
+
+    @contextmanager
+    def _redirect_stderr(self):
+        """Redirects stderr to /dev/null if not in debug mode."""
+        if not self.debug:
+            original_stderr = sys.stderr
+            devnull = None
+            try:
+                devnull = open(os.devnull, "w")
+                sys.stderr = devnull
+                yield
+            except Exception:
+                 # Restore stderr even if there was an error opening /dev/null or during yield
+                 if devnull: devnull.close()
+                 sys.stderr = original_stderr
+                 raise # Re-raise the exception
+            finally:
+                if devnull: devnull.close()
+                sys.stderr = original_stderr
+        else:
+            # If debug is True, don't redirect
+            yield
+
+    async def list_tools(self) -> List[Tool]:
+        """
+        Discover tools from the MCP server and cache their schemas.
+
+        Returns:
+            List[Tool]: A list of discovered tools with schemas.
+        """
+        logger.debug(f"Entering list_tools for command={self.command}, args={self.args}")
+
+        # Attempt to retrieve tools from cache
+        # Create a more robust cache key
+        args_string = json.dumps(self.args, sort_keys=True) # Serialize args consistently
+        cache_key = f"mcp_tools_{self.command}_{args_string}"
+        cached_tools_data = self.cache.get(cache_key)
+
+        if cached_tools_data:
+            logger.debug("Retrieved tools data from cache")
+            tools = []
+            for tool_data in cached_tools_data:
+                tool_name = tool_data["name"]
+                # Create Tool instance, ensuring func is a callable wrapper
+                tool = Tool(
+                    name=tool_name,
+                    description=tool_data["description"],
+                    input_schema=tool_data.get("input_schema", {"type": "object", "properties": {}}),
+                    func=self._create_tool_callable(tool_name), # Use the factory method
+                )
+                self._tool_cache[tool_name] = tool # Store in instance cache too
+                tools.append(tool)
+            logger.debug(f"Returning {len(tools)} cached tools")
+            return tools
+
+        # If not in cache, discover from server
+        server_params = StdioServerParameters(command=self.command, args=self.args, env=self.env)
+        logger.debug("Opening stdio_client connection")
+        try:
+            async with stdio_client(server_params) as (read, write):
+                logger.debug("Opening ClientSession")
+                async with ClientSession(read, write) as session:
+                    logger.info("Initializing session for tool discovery")
+                    await asyncio.wait_for(session.initialize(), timeout=self.timeout)
+                    logger.info("Requesting tool list from MCP server...")
+                    tools_response = await asyncio.wait_for(session.list_tools(), timeout=self.timeout)
+                    logger.debug(f"Tool list received: {tools_response}")
+
+                    if not hasattr(tools_response, 'tools') or not isinstance(tools_response.tools, list):
+                        logger.error(f"Invalid tool list response from MCP server: {tools_response}")
+                        return []
+
+                    serialized_tools = []
+                    tools = []
+                    for tool_proto in tools_response.tools:
+                         if not hasattr(tool_proto, 'name') or not tool_proto.name:
+                              logger.warning(f"Skipping tool with missing name in response: {tool_proto}")
+                              continue
+
+                         # Ensure inputSchema exists and is a dict, default if not
+                         input_schema = getattr(tool_proto, 'inputSchema', None)
+                         if not isinstance(input_schema, dict):
+                              input_schema = {"type": "object", "properties": {}}
+
+                         description = getattr(tool_proto, 'description', "") or "" # Ensure description is string
+
+                         serialized_tool_data = {
+                              'name': tool_proto.name,
+                              'description': description,
+                              'input_schema': input_schema,
+                         }
+                         serialized_tools.append(serialized_tool_data)
+
+                         # Create Tool instance for returning
+                         discovered_tool = Tool(
+                              name=tool_proto.name,
+                              description=description,
+                              input_schema=input_schema,
+                              func=self._create_tool_callable(tool_proto.name),
+                         )
+                         self._tool_cache[tool_proto.name] = discovered_tool # Cache instance
+                         tools.append(discovered_tool)
+                         logger.debug(f"Discovered tool: {tool_proto.name} with schema: {input_schema}")
+
+                    # Cache the serialized data
+                    self.cache.set(cache_key, serialized_tools, 3600)
+                    logger.debug(f"Cached {len(serialized_tools)} tools.")
+
+                    logger.debug(f"Returning {len(tools)} tools from MCP server")
+                    return tools
+
+        except asyncio.TimeoutError:
+            logger.error(f"Timeout after {self.timeout}s waiting for tool list")
+            raise RuntimeError("Tool list request timed out")
+        except Exception as e:
+            logger.error(f"Error listing tools: {e}", exc_info=True)
+            raise RuntimeError(f"Failed to list tools: {e}") from e
+
+    async def _do_list_resources(self) -> Any:
+        """Internal method to list resources with timeout."""
+        server_params = StdioServerParameters(command=self.command, args=self.args, env=self.env)
+        logger.debug("Opening stdio_client connection for resources")
+        try:
+            async with stdio_client(server_params) as (read, write):
+                logger.debug("Opening ClientSession for resources")
+                async with ClientSession(read, write) as session:
+                    with self._redirect_stderr(): # Suppress stderr if not debugging
+                        logger.debug("Initializing session before listing resources")
+                        await asyncio.wait_for(session.initialize(), timeout=self.timeout)
+                        logger.info("Requesting resource list from MCP server...")
+                        resources_response = await asyncio.wait_for(session.list_resources(), timeout=self.timeout)
+                    logger.debug("Resource list received from MCP server")
+                    return resources_response
+        except asyncio.TimeoutError:
+             logger.error(f"Timeout listing resources after {self.timeout}s")
+             raise RuntimeError("Resource list request timed out")
+        except Exception as e:
+             logger.error(f"Error listing resources: {e}", exc_info=True)
+             raise RuntimeError(f"Failed to list resources: {e}") from e
+
+    def _create_tool_callable(self, tool_name: str) -> Callable[..., Any]:
+        """
+        Dynamically create an async callable function for the specified tool.
+        This callable will establish a connection and execute the tool on demand.
+        """
+        async def dynamic_tool_func(**kwargs) -> Any:
+            logger.debug(f"Creating tool callable for '{tool_name}'")
+            server_params = StdioServerParameters(command=self.command, args=self.args, env=self.env)
+            try:
+                async with stdio_client(server_params) as (read, write):
+                    async with ClientSession(read, write) as session:
+                        # Initialize session first
+                        logger.debug(f"Initializing session for tool '{tool_name}'")
+                        await asyncio.wait_for(session.initialize(), timeout=self.timeout)
+
+                        # Validate input if schema is available in instance cache
+                        if tool_name in self._tool_cache:
+                            tool = self._tool_cache[tool_name]
+                            self._validate_input_schema(tool.input_schema, kwargs)
+                        else:
+                            logger.warning(f"Schema for tool '{tool_name}' not found in cache for validation.")
+
+                        logger.info(f"Calling tool '{tool_name}' with arguments: {kwargs}")
+                        # Execute the tool call
+                        result_proto = await asyncio.wait_for(session.call_tool(tool_name, kwargs), timeout=self.timeout)
+
+                        # Process result (assuming result_proto has a 'result' attribute)
+                        result_data = getattr(result_proto, 'result', None)
+                        if result_data is None:
+                            logger.warning(f"Tool '{tool_name}' executed but returned no result data.")
+                            return None # Or raise error?
+
+                        # Attempt to parse if it looks like JSON, otherwise return as is
+                        if isinstance(result_data, str):
+                            try:
+                                parsed_result = json.loads(result_data)
+                                logger.info(f"Tool '{tool_name}' executed successfully (result parsed as JSON).")
+                                return parsed_result
+                            except json.JSONDecodeError:
+                                logger.info(f"Tool '{tool_name}' executed successfully (result returned as string).")
+                                return result_data # Return raw string if not JSON
+                        else:
+                             logger.info(f"Tool '{tool_name}' executed successfully (result type: {type(result_data)}).")
+                             return result_data # Return non-string result directly
+
+            except asyncio.TimeoutError:
+                logger.error(f"Timeout after {self.timeout}s executing tool '{tool_name}'")
+                raise RuntimeError(f"Tool '{tool_name}' execution timed out")
+            except Exception as e:
+                logger.error(f"Failed to execute tool '{tool_name}': {e}", exc_info=True)
+                raise RuntimeError(f"Tool execution failed: {e}") from e
+
+        return dynamic_tool_func
+
+    def _validate_input_schema(self, schema: Dict[str, Any], kwargs: Dict[str, Any]):
+        """
+        Validate the provided arguments against the input schema.
+        """
+        # Ensure schema is a dictionary, default to no-op if not
+        if not isinstance(schema, dict):
+            logger.warning(f"Invalid schema format for validation: {type(schema)}. Skipping.")
+            return
+
+        required_params = schema.get("required", [])
+        # Ensure required_params is a list
+        if not isinstance(required_params, list):
+            logger.warning(f"Invalid 'required' list in schema: {type(required_params)}. Skipping requirement check.")
+            required_params = []
+
+        for param in required_params:
+            if param not in kwargs:
+                raise ValueError(f"Missing required parameter: '{param}'")
+
+        # Optional: Add type validation based on schema['properties'][param]['type']
+        properties = schema.get("properties", {})
+        if isinstance(properties, dict):
+            for key, value in kwargs.items():
+                 if key in properties:
+                      expected_type = properties[key].get("type")
+                      # Basic type mapping (add more as needed)
+                      type_map = {"string": str, "integer": int, "number": (int, float), "boolean": bool, "array": list, "object": dict}
+                      if expected_type in type_map:
+                           if not isinstance(value, type_map[expected_type]):
+                                logger.warning(f"Type mismatch for parameter '{key}'. Expected '{expected_type}', got '{type(value).__name__}'. Attempting to proceed.")
+                                # Allow proceeding but log warning, or raise ValueError for strict validation
+
+        logger.debug(f"Validated input against schema: {schema} with arguments: {kwargs}")
+
+    async def list_resources(self) -> Any:
+        """
+        Discover resources from the MCP server using the internal method with enforced timeout.
+        """
+        return await self._do_list_resources() # Timeout handled in _do_list_resources
+
+    async def get_resource(self, resource_uri: str) -> Any:
+        """
+        Retrieve a specific resource from the MCP server.
+
+        Args:
+            resource_uri (str): The URI of the resource to retrieve.
+
+        Returns:
+            Any: The resource retrieval response.
+        """
+        server_params = StdioServerParameters(command=self.command, args=self.args, env=self.env)
+        logger.debug("Opening stdio_client connection for resource retrieval")
+        try:
+            async with stdio_client(server_params) as (read, write):
+                logger.debug("Opening ClientSession for resource retrieval")
+                async with ClientSession(read, write) as session:
+                    with self._redirect_stderr(): # Suppress stderr if not debugging
+                        logger.debug(f"Initializing session for resource retrieval of {resource_uri}")
+                        await asyncio.wait_for(session.initialize(), timeout=self.timeout)
+                        logger.info(f"Retrieving resource '{resource_uri}' from MCP server")
+                        response = await asyncio.wait_for(session.read_resource(resource_uri), timeout=self.timeout)
+                    logger.info(f"Resource '{resource_uri}' retrieved successfully")
+                    # Process response if needed (e.g., getattr(response, 'content', None))
+                    return response
+        except asyncio.TimeoutError:
+            logger.error(f"Timeout retrieving resource '{resource_uri}' after {self.timeout}s")
+            raise RuntimeError(f"Resource '{resource_uri}' retrieval timed out")
+        except Exception as e:
+            logger.error(f"Failed to retrieve resource '{resource_uri}': {e}", exc_info=True)
+            raise RuntimeError(f"Resource retrieval failed: {e}") from e
+

swarm/extensions/mcp/mcp_constants.py

@@ -0,0 +1,7 @@
+"""
+Constants specific to MCP interactions.
+"""
+
+# Separator used in tool results to signal agent handoff
+MCP_SEPARATOR = ":::"
+

swarm/extensions/mcp/mcp_tool_provider.py

@@ -0,0 +1,110 @@
+"""
+MCPToolProvider Module for Open-Swarm
+
+This module is responsible for discovering tools from MCP (Model Context Protocol) servers
+and integrating them into the Open-Swarm framework as `Tool` instances.
+"""
+
+import logging
+import json
+import re # Standard library for regular expressions
+from typing import List, Dict, Any
+
+from ...settings import Settings # Use Swarm settings
+from ...types import Tool, Agent
+from .mcp_client import MCPClient
+from .cache_utils import get_cache
+
+# Use Swarm's settings for logging configuration
+swarm_settings = Settings()
+logger = logging.getLogger(__name__)
+logger.setLevel(swarm_settings.log_level.upper())
+# Ensure handler is added only if needed
+if not logger.handlers and not logging.getLogger('swarm').handlers:
+     handler = logging.StreamHandler()
+     formatter = logging.Formatter(swarm_settings.log_format.value)
+     handler.setFormatter(formatter)
+     logger.addHandler(handler)
+
+
+class MCPToolProvider:
+    """
+    MCPToolProvider discovers tools from an MCP server and converts them into `Tool` instances.
+    Uses caching to avoid repeated discovery.
+    """
+    _instances: Dict[str, "MCPToolProvider"] = {}
+
+    @classmethod
+    def get_instance(cls, server_name: str, server_config: Dict[str, Any], timeout: int = 15, debug: bool = False) -> "MCPToolProvider":
+        """Get or create an instance for the given server name."""
+        config_key = json.dumps(server_config, sort_keys=True)
+        instance_key = f"{server_name}_{config_key}_{timeout}_{debug}"
+
+        if instance_key not in cls._instances:
+            logger.debug(f"Creating new MCPToolProvider instance for key: {instance_key}")
+            cls._instances[instance_key] = cls(server_name, server_config, timeout, debug)
+        else:
+             logger.debug(f"Reusing existing MCPToolProvider instance for key: {instance_key}")
+        return cls._instances[instance_key]
+
+    def __init__(self, server_name: str, server_config: Dict[str, Any], timeout: int = 15, debug: bool = False):
+        """
+        Initialize an MCPToolProvider instance. Use get_instance() for shared instances.
+        """
+        self.server_name = server_name
+        effective_debug = debug or swarm_settings.debug
+        try:
+            self.client = MCPClient(server_config=server_config, timeout=timeout, debug=effective_debug)
+        except ImportError as e:
+             logger.error(f"Failed to initialize MCPClient for '{server_name}': {e}. MCP features will be unavailable.")
+             self.client = None
+        except Exception as e:
+             logger.error(f"Error initializing MCPClient for '{server_name}': {e}", exc_info=True)
+             self.client = None
+
+        self.cache = get_cache()
+        logger.debug(f"Initialized MCPToolProvider for server '{self.server_name}' with timeout {timeout}s.")
+
+    async def discover_tools(self, agent: Agent) -> List[Tool]:
+        """
+        Discover tools from the MCP server using the MCPClient.
+
+        Args:
+            agent (Agent): The agent for which tools are being discovered.
+
+        Returns:
+            List[Tool]: A list of discovered `Tool` instances with prefixed names.
+        """
+        if not self.client:
+             logger.warning(f"MCPClient for '{self.server_name}' not initialized. Cannot discover tools.")
+             return []
+
+        logger.debug(f"Starting tool discovery via MCPClient for server '{self.server_name}'.")
+        try:
+            tools = await self.client.list_tools()
+            logger.debug(f"Discovered {len(tools)} tools from MCP server '{self.server_name}'.")
+
+            separator = "__"
+            prefixed_tools = []
+            for tool in tools:
+                 prefixed_name = f"{self.server_name}{separator}{tool.name}"
+                 # Validate prefixed name against OpenAI pattern
+                 if not re.match(r"^[a-zA-Z0-9_-]{1,64}$", prefixed_name):
+                      logger.warning(f"Generated MCP tool name '{prefixed_name}' might violate OpenAI pattern. Skipping.")
+                      continue
+
+                 prefixed_tool = Tool(
+                      name=prefixed_name,
+                      description=tool.description,
+                      input_schema=tool.input_schema,
+                      func=tool.func # Callable already targets this client/tool
+                 )
+                 prefixed_tools.append(prefixed_tool)
+                 logger.debug(f"Added prefixed tool: {prefixed_tool.name}")
+
+            return prefixed_tools
+
+        except Exception as e:
+            logger.error(f"Failed to discover tools from MCP server '{self.server_name}': {e}", exc_info=True)
+            return []
+

swarm/llm/chat_completion.py

@@ -0,0 +1,195 @@
+"""
+Chat Completion Module
+
+This module handles chat completion logic for the Swarm framework, including message preparation,
+tool call repair, and interaction with the OpenAI API. Located in llm/ for LLM-specific functionality.
+"""
+
+import os
+import json
+import logging
+from typing import List, Optional, Dict, Any, Union, AsyncGenerator # Added AsyncGenerator
+from collections import defaultdict
+
+import asyncio
+from openai import AsyncOpenAI, OpenAIError
+# Make sure ChatCompletionMessage is correctly imported if it's defined elsewhere
+# Assuming it might be part of the base model or a common types module
+# For now, let's assume it's implicitly handled or use a dict directly
+# from ..types import ChatCompletionMessage, Agent # If defined in types
+from ..types import Agent # Import Agent
+from ..utils.redact import redact_sensitive_data
+from ..utils.general_utils import serialize_datetime
+from ..utils.message_utils import filter_duplicate_system_messages, update_null_content
+from ..utils.context_utils import get_token_count, truncate_message_history
+from ..utils.message_sequence import repair_message_payload
+
+# Configure module-level logging
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+if not logger.handlers:
+    stream_handler = logging.StreamHandler()
+    formatter = logging.Formatter("[%(levelname)s] %(asctime)s - %(name)s - %(message)s")
+    stream_handler.setFormatter(formatter)
+    logger.addHandler(stream_handler)
+
+
+async def get_chat_completion(
+    client: AsyncOpenAI,
+    agent: Agent,
+    history: List[Dict[str, Any]],
+    context_variables: dict,
+    current_llm_config: Dict[str, Any],
+    max_context_tokens: int,
+    max_context_messages: int,
+    tools: Optional[List[Dict[str, Any]]] = None, # <-- Added tools parameter
+    tool_choice: Optional[str] = "auto",         # <-- Added tool_choice parameter
+    model_override: Optional[str] = None,
+    stream: bool = False,
+    debug: bool = False
+) -> Union[Dict[str, Any], AsyncGenerator[Any, None]]: # Adjusted return type hint
+    """
+    Retrieve a chat completion from the LLM for the given agent and history.
+
+    Args:
+        client: AsyncOpenAI client instance.
+        agent: The agent processing the completion.
+        history: List of previous messages in the conversation.
+        context_variables: Variables to include in the agent's context.
+        current_llm_config: Current LLM configuration dictionary.
+        max_context_tokens: Maximum token limit for context.
+        max_context_messages: Maximum message limit for context.
+        tools: Optional list of tools in OpenAI format.
+        tool_choice: Tool choice mode (e.g., "auto", "none").
+        model_override: Optional model to use instead of default.
+        stream: If True, stream the response; otherwise, return complete.
+        debug: If True, log detailed debugging information.
+
+    Returns:
+        Union[Dict[str, Any], AsyncGenerator[Any, None]]: The LLM's response message (as dict) or stream.
+    """
+    if not agent:
+        logger.error("Cannot generate chat completion: Agent is None")
+        raise ValueError("Agent is required")
+
+    logger.debug(f"Generating chat completion for agent '{agent.name}'")
+    active_model = model_override or current_llm_config.get("model", "default")
+    client_kwargs = {
+        "api_key": current_llm_config.get("api_key"),
+        "base_url": current_llm_config.get("base_url")
+    }
+    client_kwargs = {k: v for k, v in client_kwargs.items() if v is not None}
+    redacted_kwargs = redact_sensitive_data(client_kwargs, sensitive_keys=["api_key"])
+    logger.debug(f"Using client with model='{active_model}', base_url='{client_kwargs.get('base_url', 'default')}', api_key={redacted_kwargs['api_key']}")
+
+    context_variables = defaultdict(str, context_variables)
+    instructions = agent.instructions(context_variables) if callable(agent.instructions) else agent.instructions
+    if not isinstance(instructions, str):
+        logger.warning(f"Invalid instructions type for '{agent.name}': {type(instructions)}. Converting to string.")
+        instructions = str(instructions)
+    messages = repair_message_payload([{"role": "system", "content": instructions}], debug=debug)
+
+    if not isinstance(history, list):
+        logger.error(f"Invalid history type for '{agent.name}': {type(history)}. Expected list.")
+        history = []
+    seen_ids = set()
+    for msg in history:
+        msg_id = msg.get("id", hash(json.dumps(msg, sort_keys=True, default=serialize_datetime)))
+        if msg_id not in seen_ids:
+            seen_ids.add(msg_id)
+            if "tool_calls" in msg and msg["tool_calls"] is not None and not isinstance(msg["tool_calls"], list):
+                logger.warning(f"Invalid tool_calls in history for '{msg.get('sender', 'unknown')}': {msg['tool_calls']}. Setting to None.")
+                msg["tool_calls"] = None
+            # Ensure content: None becomes content: "" for API compatibility
+            if "content" in msg and msg["content"] is None:
+                 msg["content"] = ""
+            messages.append(msg)
+    messages = filter_duplicate_system_messages(messages)
+    messages = truncate_message_history(messages, active_model, max_context_tokens, max_context_messages)
+    messages = repair_message_payload(messages, debug=debug)  # Ensure tool calls are paired post-truncation
+    # Final content None -> "" check after repair
+    messages = update_null_content(messages)
+
+    logger.debug(f"Prepared {len(messages)} messages for '{agent.name}'")
+    if debug:
+        logger.debug(f"Messages: {json.dumps(messages, indent=2, default=str)}")
+
+    create_params = {
+        "model": active_model,
+        "messages": messages,
+        "stream": stream,
+        "temperature": current_llm_config.get("temperature", 0.7),
+        # --- Pass tools and tool_choice ---
+        "tools": tools if tools else None,
+        "tool_choice": tool_choice if tools else None, # Only set tool_choice if tools are provided
+    }
+    if getattr(agent, "response_format", None):
+        create_params["response_format"] = agent.response_format
+    create_params = {k: v for k, v in create_params.items() if v is not None} # Clean None values
+
+    tool_info_log = f", tools_count={len(tools)}" if tools else ", tools=None"
+    logger.debug(f"Chat completion params: model='{active_model}', messages_count={len(messages)}, stream={stream}{tool_info_log}, tool_choice={create_params.get('tool_choice')}")
+
+    try:
+        logger.debug(f"Calling OpenAI API for '{agent.name}' with model='{active_model}'")
+        # Temporary workaround for potential env var conflicts if client doesn't isolate well
+        prev_openai_api_key = os.environ.pop("OPENAI_API_KEY", None)
+        try:
+            completion = await client.chat.completions.create(**create_params)
+            if stream:
+                return completion # Return stream object directly
+
+            # --- Handle Non-Streaming Response ---
+            if completion.choices and len(completion.choices) > 0 and completion.choices[0].message:
+                message_dict = completion.choices[0].message.model_dump(exclude_none=True)
+                log_msg = message_dict.get("content", "No content")[:50] if message_dict.get("content") else "No content"
+                if message_dict.get("tool_calls"): log_msg += f" (+{len(message_dict['tool_calls'])} tool calls)"
+                logger.debug(f"OpenAI completion received for '{agent.name}': {log_msg}...")
+                return message_dict # Return the message dictionary
+            else:
+                logger.warning(f"No valid message in completion for '{agent.name}'")
+                return {"role": "assistant", "content": "No response generated"} # Return dict
+        finally:
+            if prev_openai_api_key is not None:
+                os.environ["OPENAI_API_KEY"] = prev_openai_api_key
+    except OpenAIError as e:
+        logger.error(f"Chat completion failed for '{agent.name}': {e}")
+        raise
+    except Exception as e: # Catch broader errors during API call
+        logger.error(f"Unexpected error during chat completion for '{agent.name}': {e}", exc_info=True)
+        raise # Re-raise
+
+
+async def get_chat_completion_message(
+    client: AsyncOpenAI,
+    agent: Agent,
+    history: List[Dict[str, Any]],
+    context_variables: dict,
+    current_llm_config: Dict[str, Any],
+    max_context_tokens: int,
+    max_context_messages: int,
+    tools: Optional[List[Dict[str, Any]]] = None, # <-- Added tools
+    tool_choice: Optional[str] = "auto",        # <-- Added tool_choice
+    model_override: Optional[str] = None,
+    stream: bool = False,
+    debug: bool = False
+) -> Union[Dict[str, Any], AsyncGenerator[Any, None]]: # Return dict or stream
+    """
+    Wrapper to retrieve and validate a chat completion message (returns dict or stream).
+
+    Args:
+        Same as get_chat_completion.
+
+    Returns:
+        Union[Dict[str, Any], AsyncGenerator[Any, None]]: Validated LLM response message as dict or the stream.
+    """
+    logger.debug(f"Fetching chat completion message for '{agent.name}'")
+    completion_result = await get_chat_completion(
+        client, agent, history, context_variables, current_llm_config,
+        max_context_tokens, max_context_messages,
+        tools=tools, tool_choice=tool_choice, # Pass through
+        model_override=model_override, stream=stream, debug=debug
+    )
+    # If streaming, completion_result is already the generator
+    # If not streaming, it's the message dictionary
+    return completion_result