open-swarm 0.1.1745274976__py3-none-any.whl → 0.1.1748636259__py3-none-any.whl

swarm/extensions/launchers/swarm_api.py

@@ -1,12 +1,68 @@
-"""
-Swarm API entry point for installation via PyPI or local dev.
-"""
+#!/usr/bin/env python3
+import argparse
+import subprocess
 import sys
-import os
-from swarm.core.swarm_api import main
+from os import path, listdir, makedirs
 
-def main_entry():
-    main()
+def main():
+    parser = argparse.ArgumentParser(description="Swarm REST Launcher")
+    parser.add_argument("--blueprint", required=True, help="Comma-separated blueprint file paths or names for configuration purposes")
+    parser.add_argument("--port", type=int, default=8000, help="Port to run the REST server")
+    parser.add_argument("--config", default="~/.swarm/swarm_config.json", help="Configuration file path")
+    parser.add_argument("--daemon", action="store_true", help="Run in daemon mode and print process id")
+    args = parser.parse_args()
+
+    # Split blueprints by comma and strip whitespace
+    bp_list = [bp.strip() for bp in args.blueprint.split(",") if bp.strip()]
+    blueprint_paths = []
+    for bp_arg in bp_list:
+        resolved = None
+        if path.exists(bp_arg):
+            if path.isdir(bp_arg):
+                resolved = bp_arg
+                print(f"Using blueprint directory: {resolved}")
+            else:
+                resolved = bp_arg
+                print(f"Using blueprint file: {resolved}")
+        else:
+            managed_path = path.expanduser("~/.swarm/blueprints/" + bp_arg)
+            if path.isdir(managed_path):
+                matches = [f for f in listdir(managed_path) if f.startswith("blueprint_") and f.endswith(".py")]
+                if not matches:
+                    print("Error: No blueprint file found in managed directory:", managed_path)
+                    sys.exit(1)
+                resolved = path.join(managed_path, matches[0])
+                print(f"Using managed blueprint: {resolved}")
+            else:
+                print("Warning: Blueprint not found:", bp_arg, "- skipping.")
+                continue
+        if resolved:
+            blueprint_paths.append(resolved)
+
+    if not blueprint_paths:
+        print("Error: No valid blueprints found.")
+        sys.exit(1)
+    print("Blueprints to be configured:")
+    for bp in blueprint_paths:
+        print(" -", bp)
+
+    config_path = path.expanduser(args.config)
+    if not path.exists(config_path):
+        makedirs(path.dirname(config_path), exist_ok=True)
+        with open(config_path, 'w') as f:
+            f.write("{}")
+        print("Default config file created at:", config_path)
+
+    print("Launching Django server on port 0.0.0.0:{}".format(args.port))
+    try:
+        if args.daemon:
+            proc = subprocess.Popen(["python", "manage.py", "runserver", f"0.0.0.0:{args.port}"])
+            print("Running in daemon mode. Process ID:", proc.pid)
+        else:
+            subprocess.run(["python", "manage.py", "runserver", f"0.0.0.0:{args.port}"], check=True)
+    except subprocess.CalledProcessError as e:
+        print("Error launching Django server:", e)
+        sys.exit(1)
 
 if __name__ == "__main__":
-    main_entry()
+    main()

swarm/extensions/launchers/swarm_cli.py

@@ -1,12 +1,304 @@
-"""
-Swarm CLI entry point for installation via PyPI or local dev.
-"""
-import sys
+#!/usr/bin/env python3
+import argparse
+import importlib.util
 import os
-from swarm.extensions.cli.main import main
+import sys
+import subprocess
+import shutil
+import json
+import PyInstaller.__main__
 
-def app():
-    main()
+def resolve_env_vars(data):
+    if isinstance(data, dict):
+        return {k: resolve_env_vars(v) for k, v in data.items()}
+    elif isinstance(data, list):
+        return [resolve_env_vars(item) for item in data]
+    elif isinstance(data, str):
+        return os.path.expandvars(data)
+    else:
+        return data
+
+MANAGED_DIR = os.path.expanduser("~/.swarm/blueprints")
+BIN_DIR = os.path.expanduser("~/.swarm/bin")
+
+def ensure_managed_dir():
+    if not os.path.exists(MANAGED_DIR):
+        os.makedirs(MANAGED_DIR, exist_ok=True)
+    if not os.path.exists(BIN_DIR):
+        os.makedirs(BIN_DIR, exist_ok=True)
+
+def add_blueprint(source_path, blueprint_name=None):
+    source_path = os.path.normpath(source_path)
+    if not os.path.exists(source_path):
+        print("Error: source file/directory does not exist:", source_path)
+        sys.exit(1)
+    if os.path.isdir(source_path):
+        if not blueprint_name:
+            blueprint_name = os.path.basename(os.path.normpath(source_path))
+        target_dir = os.path.join(MANAGED_DIR, blueprint_name)
+        if os.path.exists(target_dir):
+            shutil.rmtree(target_dir)
+        os.makedirs(target_dir, exist_ok=True)
+        for root, dirs, files in os.walk(source_path):
+            rel_path = os.path.relpath(root, source_path)
+            dest_root = os.path.join(target_dir, rel_path) if rel_path != '.' else target_dir
+            os.makedirs(dest_root, exist_ok=True)
+            for file in files:
+                shutil.copy2(os.path.join(root, file), os.path.join(dest_root, file))
+        print(f"Blueprint '{blueprint_name}' added successfully to {target_dir}.")
+    else:
+        blueprint_file = source_path
+        if not blueprint_name:
+            base = os.path.basename(blueprint_file)
+            if base.startswith("blueprint_") and base.endswith(".py"):
+                blueprint_name = base[len("blueprint_"):-3]
+            else:
+                blueprint_name = os.path.splitext(base)[0]
+        target_dir = os.path.join(MANAGED_DIR, blueprint_name)
+        os.makedirs(target_dir, exist_ok=True)
+        target_file = os.path.join(target_dir, f"blueprint_{blueprint_name}.py")
+        shutil.copy2(blueprint_file, target_file)
+        print(f"Blueprint '{blueprint_name}' added successfully to {target_dir}.")
+
+def list_blueprints():
+    ensure_managed_dir()
+    entries = os.listdir(MANAGED_DIR)
+    blueprints = [d for d in entries if os.path.isdir(os.path.join(MANAGED_DIR, d))]
+    if blueprints:
+        print("Registered blueprints:")
+        for bp in blueprints:
+            print(" -", bp)
+    else:
+        print("No blueprints registered.")
+
+def delete_blueprint(blueprint_name):
+    target_dir = os.path.join(MANAGED_DIR, blueprint_name)
+    if os.path.exists(target_dir) and os.path.isdir(target_dir):
+        shutil.rmtree(target_dir)
+        print(f"Blueprint '{blueprint_name}' deleted successfully.")
+    else:
+        print(f"Error: Blueprint '{blueprint_name}' does not exist.")
+        sys.exit(1)
+
+def run_blueprint(blueprint_name):
+    target_dir = os.path.join(MANAGED_DIR, blueprint_name)
+    blueprint_file = os.path.join(target_dir, f"blueprint_{blueprint_name}.py")
+    if not os.path.exists(blueprint_file):
+        print(f"Error: Blueprint file not found for '{blueprint_name}'. Install it using 'swarm-cli add <path>'.")
+        sys.exit(1)
+    spec = importlib.util.spec_from_file_location("blueprint_module", blueprint_file)
+    if spec is None or spec.loader is None:
+        print("Error: Failed to load blueprint module from:", blueprint_file)
+        sys.exit(1)
+    blueprint = importlib.util.module_from_spec(spec)
+    loader = spec.loader
+    src_path = os.path.join(os.getcwd(), "src")
+    if src_path not in sys.path:
+        sys.path.insert(0, src_path)
+    loader.exec_module(blueprint)
+    if hasattr(blueprint, "main"):
+        blueprint.main()
+    else:
+        print("Error: The blueprint does not have a main() function.")
+        sys.exit(1)
+
+def install_blueprint(blueprint_name):
+    target_dir = os.path.join(MANAGED_DIR, blueprint_name)
+    blueprint_file = os.path.join(target_dir, f"blueprint_{blueprint_name}.py")
+    if not os.path.exists(blueprint_file):
+        print(f"Error: Blueprint '{blueprint_name}' is not registered. Add it using 'swarm-cli add <path>'.")
+        sys.exit(1)
+    cli_name = blueprint_name  # Use blueprint_name as default cli_name for simplicity
+    try:
+        PyInstaller.__main__.run([
+            blueprint_file,
+            "--onefile",
+            "--name", cli_name,
+            "--distpath", BIN_DIR,
+            "--workpath", os.path.join(target_dir, "build"),
+            "--specpath", target_dir
+        ])
+    except KeyboardInterrupt:
+        print("Installation aborted by user request.")
+        sys.exit(1)
+    print(f"Blueprint '{blueprint_name}' installed as CLI utility '{cli_name}' at: {os.path.join(BIN_DIR, cli_name)}")
+
+def uninstall_blueprint(blueprint_name, blueprint_only=False, wrapper_only=False):
+    target_dir = os.path.join(MANAGED_DIR, blueprint_name)
+    blueprint_file = os.path.join(target_dir, f"blueprint_{blueprint_name}.py")
+    cli_name = blueprint_name  # Default to blueprint_name for uninstall
+    cli_path = os.path.join(BIN_DIR, cli_name)
+    removed = False
+    
+    if not blueprint_only and not wrapper_only:  # Remove both by default
+        if os.path.exists(target_dir):
+            shutil.rmtree(target_dir)
+            print(f"Blueprint '{blueprint_name}' removed from {MANAGED_DIR}.")
+            removed = True
+        if os.path.exists(cli_path):
+            os.remove(cli_path)
+            print(f"Wrapper '{cli_name}' removed from {BIN_DIR}.")
+            removed = True
+    elif blueprint_only:
+        if os.path.exists(target_dir):
+            shutil.rmtree(target_dir)
+            print(f"Blueprint '{blueprint_name}' removed from {MANAGED_DIR}.")
+            removed = True
+    elif wrapper_only:
+        if os.path.exists(cli_path):
+            os.remove(cli_path)
+            print(f"Wrapper '{cli_name}' removed from {BIN_DIR}.")
+            removed = True
+    
+    if not removed:
+        print(f"Error: Nothing to uninstall for '{blueprint_name}' with specified options.")
+        sys.exit(1)
+
+def main():
+    os.environ.pop("SWARM_BLUEPRINTS", None)
+    parser = argparse.ArgumentParser(
+        description="Swarm CLI Launcher\n\nSubcommands:\n"
+                    "  add     : Add a blueprint to the managed directory.\n"
+                    "  list    : List registered blueprints.\n"
+                    "  delete  : Delete a registered blueprint.\n"
+                    "  run     : Run a blueprint by name.\n"
+                    "  install : Install a blueprint as a CLI utility with PyInstaller.\n"
+                    "  uninstall : Uninstall a blueprint and/or its CLI wrapper.\n"
+                    "  migrate : Apply Django database migrations.\n"
+                    "  config  : Manage swarm configuration (LLM and MCP servers).",
+        formatter_class=argparse.RawTextHelpFormatter)
+    subparsers = parser.add_subparsers(dest="command", required=True, help="Available subcommands")
+    
+    parser_add = subparsers.add_parser("add", help="Add a blueprint from a file or directory.")
+    parser_add.add_argument("source", help="Source blueprint file or directory.")
+    parser_add.add_argument("--name", help="Optional blueprint name. If not provided, inferred from filename.")
+    
+    parser_list = subparsers.add_parser("list", help="List registered blueprints.")
+    
+    parser_delete = subparsers.add_parser("delete", help="Delete a registered blueprint by name.")
+    parser_delete.add_argument("name", help="Blueprint name to delete.")
+    
+    parser_run = subparsers.add_parser("run", help="Run a blueprint by name.")
+    parser_run.add_argument("name", help="Blueprint name to run.")
+    parser_run.add_argument("--config", default="~/.swarm/swarm_config.json", help="Path to configuration file.")
+    
+    parser_install = subparsers.add_parser("install", help="Install a blueprint as a CLI utility with PyInstaller.")
+    parser_install.add_argument("name", help="Blueprint name to install as a CLI utility.")
+    
+    parser_uninstall = subparsers.add_parser("uninstall", help="Uninstall a blueprint and/or its CLI wrapper.")
+    parser_uninstall.add_argument("name", help="Blueprint name to uninstall.")
+    parser_uninstall.add_argument("--blueprint-only", action="store_true", help="Remove only the blueprint directory.")
+    parser_uninstall.add_argument("--wrapper-only", action="store_true", help="Remove only the CLI wrapper.")
+    
+    parser_migrate = subparsers.add_parser("migrate", help="Apply Django database migrations.")
+    
+    parser_config = subparsers.add_parser("config", help="Manage swarm configuration (LLM and MCP servers).")
+    parser_config.add_argument("action", choices=["add", "list", "remove"], help="Action to perform on configuration")
+    parser_config.add_argument("--section", required=True, choices=["llm", "mcpServers"], help="Configuration section to manage")
+    parser_config.add_argument("--name", help="Name of the configuration entry (required for add and remove)")
+    parser_config.add_argument("--json", help="JSON string for configuration entry (required for add)")
+    parser_config.add_argument("--config", default="~/.swarm/swarm_config.json", help="Path to configuration file")
+    
+    args = parser.parse_args()
+    ensure_managed_dir()
+    
+    if args.command == "add":
+        add_blueprint(args.source, args.name)
+    elif args.command == "list":
+        list_blueprints()
+    elif args.command == "delete":
+        delete_blueprint(args.name)
+    elif args.command == "run":
+        config_path = os.path.expanduser(args.config)
+        if not os.path.exists(config_path):
+            os.makedirs(os.path.dirname(config_path), exist_ok=True)
+            default_config = {"llm": {}, "mcpServers": {}}
+            with open(config_path, 'w') as f:
+                json.dump(default_config, f, indent=4)
+            print("Default config file created at:", config_path)
+        run_blueprint(args.name)
+    elif args.command == "install":
+        install_blueprint(args.name)
+    elif args.command == "uninstall":
+        uninstall_blueprint(args.name, args.blueprint_only, args.wrapper_only)
+    elif args.command == "migrate":
+        try:
+            subprocess.run(["python", "manage.py", "migrate"], check=True)
+            print("Migrations applied successfully.")
+        except subprocess.CalledProcessError as e:
+            print("Error applying migrations:", e)
+            sys.exit(1)
+    elif args.command == "config":
+        config_path = os.path.expanduser(args.config)
+        if not os.path.exists(config_path):
+            default_conf = {"llm": {}, "mcpServers": {}}
+            os.makedirs(os.path.dirname(config_path), exist_ok=True)
+            with open(config_path, "w") as f:
+                json.dump(default_conf, f, indent=4)
+            print("Default config file created at:", config_path)
+            config = default_conf
+        else:
+            try:
+                with open(config_path, "r") as f:
+                    config = json.load(f)
+            except json.JSONDecodeError:
+                print("Error: Invalid configuration file.")
+                sys.exit(1)
+        section = args.section
+        if args.action == "list":
+            entries = config.get(section, {})
+            if entries:
+                print(f"Entries in {section}:")
+                for key, value in entries.items():
+                    print(f" - {key}: {json.dumps(value, indent=4)}")
+            else:
+                print(f"No entries found in {section}.")
+        elif args.action == "add":
+            if args.section == "mcpServers" and not args.name:
+                if not args.json:
+                    print("Error: --json is required for adding an mcpServers block when --name is omitted.")
+                    sys.exit(1)
+                try:
+                    update_data = json.loads(args.json)
+                except json.JSONDecodeError:
+                    print("Error: --json must be a valid JSON string.")
+                    sys.exit(1)
+                if "mcpServers" not in update_data:
+                    print("Error: JSON block must contain 'mcpServers' key for merging.")
+                    sys.exit(1)
+                config.setdefault("mcpServers", {})
+                config["mcpServers"].update(update_data["mcpServers"])
+                with open(config_path, "w") as f:
+                    json.dump(config, f, indent=4)
+                print("MCP servers updated in configuration.")
+            else:
+                if not args.name or not args.json:
+                    print("Error: --name and --json are required for adding an entry.")
+                    sys.exit(1)
+                try:
+                    entry_data = json.loads(args.json)
+                except json.JSONDecodeError:
+                    print("Error: --json must be a valid JSON string.")
+                    sys.exit(1)
+                config.setdefault(section, {})[args.name] = entry_data
+                with open(config_path, "w") as f:
+                    json.dump(config, f, indent=4)
+                print(f"Entry '{args.name}' added to {section} in configuration.")
+        elif args.action == "remove":
+            if not args.name:
+                print("Error: --name is required for removing an entry.")
+                sys.exit(1)
+            if args.name in config.get(section, {}):
+                del config[section][args.name]
+                with open(config_path, "w") as f:
+                    json.dump(config, f, indent=4)
+                print(f"Entry '{args.name}' removed from {section} in configuration.")
+            else:
+                print(f"Error: Entry '{args.name}' not found in {section}.")
+                sys.exit(1)
+    else:
+        parser.print_help()
+        sys.exit(1)
 
 if __name__ == "__main__":
-    app()
+    main()

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