project-lyra/cortex/autonomy/tools/adapters/ollama_adapter.py

"""
Ollama adapter for tool calling using XML-structured prompts.

Since Ollama doesn't have native function calling, this adapter uses
XML-based prompts to instruct the model how to call tools.
"""

import json
import re
from typing import Dict, List, Optional
from .base import ToolAdapter


class OllamaAdapter(ToolAdapter):
    """Ollama adapter using XML-structured prompts for tool calling.

    This adapter injects tool descriptions into the system prompt and
    teaches the model to respond with XML when it wants to use a tool.
    """

    SYSTEM_PROMPT = """You have access to the following tools:

{tool_descriptions}

To use a tool, respond with XML in this exact format:
<tool_call>
  <name>tool_name</name>
  <arguments>
    <arg_name>value</arg_name>
  </arguments>
  <reason>why you're using this tool</reason>
</tool_call>

You can call multiple tools by including multiple <tool_call> blocks.
If you don't need to use any tools, respond normally without XML.
After tools are executed, you'll receive results and can continue the conversation."""

    async def prepare_request(
        self,
        messages: List[Dict],
        tools: List[Dict],
        tool_choice: Optional[str] = None
    ) -> Dict:
        """Inject tool descriptions into system prompt.

        Args:
            messages: Conversation history
            tools: Lyra tool definitions
            tool_choice: Ignored for Ollama (no native support)

        Returns:
            dict: Request payload with modified messages
        """
        # Format tool descriptions
        tool_desc = "\n".join([
            f"- {t['name']}: {t['description']}\n  Parameters: {self._format_parameters(t['parameters'], t.get('required', []))}"
            for t in tools
        ])

        system_msg = self.SYSTEM_PROMPT.format(tool_descriptions=tool_desc)

        # Check if first message is already a system message
        modified_messages = messages.copy()
        if modified_messages and modified_messages[0].get("role") == "system":
            # Prepend tool instructions to existing system message
            modified_messages[0]["content"] = system_msg + "\n\n" + modified_messages[0]["content"]
        else:
            # Add new system message at the beginning
            modified_messages.insert(0, {"role": "system", "content": system_msg})

        return {"messages": modified_messages}

    def _format_parameters(self, parameters: Dict, required: List[str]) -> str:
        """Format parameters for tool description.

        Args:
            parameters: Parameter definitions
            required: List of required parameter names

        Returns:
            str: Human-readable parameter description
        """
        param_strs = []
        for name, spec in parameters.items():
            req_marker = "(required)" if name in required else "(optional)"
            param_strs.append(f"{name} {req_marker}: {spec.get('description', '')}")
        return ", ".join(param_strs)

    async def parse_response(self, response) -> Dict:
        """Extract tool calls from XML in response.

        Args:
            response: String response from Ollama

        Returns:
            dict: Standardized Lyra format with content and tool_calls
        """
        import logging
        logger = logging.getLogger(__name__)

        # Ollama returns a string
        if isinstance(response, dict):
            content = response.get("message", {}).get("content", "")
        else:
            content = str(response)

        logger.info(f"🔍 OllamaAdapter.parse_response: content length={len(content)}, has <tool_call>={('<tool_call>' in content)}")
        logger.debug(f"🔍 Content preview: {content[:500]}")

        # Parse XML tool calls
        tool_calls = []
        if "<tool_call>" in content:
            # Split content by <tool_call> to get each block
            blocks = content.split('<tool_call>')
            logger.info(f"🔍 Split into {len(blocks)} blocks")

            # First block is content before any tool calls
            clean_parts = [blocks[0]]

            for idx, block in enumerate(blocks[1:]):  # Skip first block (pre-tool content)
                # Extract tool name
                name_match = re.search(r'<name>(.*?)</name>', block)
                if not name_match:
                    logger.warning(f"Block {idx} has no <name> tag, skipping")
                    continue

                name = name_match.group(1).strip()
                arguments = {}

                # Extract arguments
                args_match = re.search(r'<arguments>(.*?)</arguments>', block, re.DOTALL)
                if args_match:
                    args_xml = args_match.group(1)
                    # Parse <key>value</key> pairs
                    arg_pairs = re.findall(r'<(\w+)>(.*?)</\1>', args_xml, re.DOTALL)
                    arguments = {k: v.strip() for k, v in arg_pairs}

                tool_calls.append({
                    "id": f"call_{idx}",
                    "name": name,
                    "arguments": arguments
                })

                # For clean content, find what comes AFTER the tool call block
                # Look for the last closing tag (</tool_call> or malformed </xxx>) and keep what's after
                # Split by any closing tag at the END of the tool block
                remaining = block
                # Remove everything up to and including a standalone closing tag
                # Pattern: find </something> that's not followed by more XML
                end_match = re.search(r'</[a-z_]+>\s*(.*)$', remaining, re.DOTALL)
                if end_match:
                    after_content = end_match.group(1).strip()
                    if after_content and not after_content.startswith('<'):
                        # Only keep if it's actual text content, not more XML
                        clean_parts.append(after_content)

            clean_content = ''.join(clean_parts).strip()
        else:
            clean_content = content

        return {
            "content": clean_content,
            "tool_calls": tool_calls if tool_calls else None
        }

    def format_tool_result(
        self,
        tool_call_id: str,
        tool_name: str,
        result: Dict
    ) -> Dict:
        """Format tool result as XML for next prompt.

        Args:
            tool_call_id: ID from the original tool call
            tool_name: Name of the executed tool
            result: Tool execution result

        Returns:
            dict: Message in user role with XML-formatted result
        """
        # Format result as XML
        result_xml = f"""<tool_result>
  <tool>{tool_name}</tool>
  <result>{json.dumps(result, ensure_ascii=False)}</result>
</tool_result>"""

        return {
            "role": "user",
            "content": result_xml
        }