client/anthropic_bridge.py

'''Classes to connect to Anthropic inference endpoint'''

import abc
from typing import Dict, List, Any, Optional
import anthropic
from client.mcp_client import MCPClientWrapper, ToolDef, ToolParameter, ToolInvocationResult

DEFAULT_ANTHROPIC_MODEL = 'claude-3-haiku-20240307'

# Type mapping from Python/MCP types to JSON Schema types
TYPE_MAPPING = {
    'int': 'integer',
    'bool': 'boolean',
    'str': 'string',
    'float': 'number',
    'list': 'array',
    'dict': 'object',
    'boolean': 'boolean',
    'string': 'string',
    'integer': 'integer',
    'number': 'number',
    'array': 'array',
    'object': 'object'
}


class LLMBridge(abc.ABC):
    '''Abstract base class for LLM bridge implementations.'''

    def __init__(self, mcp_client: MCPClientWrapper):
        '''Initialize the LLM bridge with an MCPClient instance.

        Args:
            mcp_client: An initialized MCPClient instance
        '''
        self.mcp_client = mcp_client
        self.tools = None


    async def fetch_tools(self) -> List[ToolDef]:
        '''Fetch available tools from the MCP endpoint.

        
        Returns:
            List of ToolDef objects
        '''
        self.tools = await self.mcp_client.list_tools()
        return self.tools


    @abc.abstractmethod
    async def format_tools(self, tools: List[ToolDef]) -> Any:
        '''Format tools for the specific LLM provider.

        Args:
            tools: List of ToolDef objects
            
        Returns:
            Formatted tools in the LLM-specific format
        '''
        pass


    @abc.abstractmethod
    async def submit_query(self, system_prompt: str, query: List[Dict], formatted_tools: Any) -> Dict[str, Any]:
        '''Submit a query to the LLM with the formatted tools.

        Args:
            query: User query string            # messages=[
            #     {'role': 'user', 'content': query}
            # ],
            formatted_tools: Tools in the LLM-specific format
            
        Returns:
            LLM response
        '''
        pass


    @abc.abstractmethod
    async def parse_tool_call(self, llm_response: Any) -> Optional[Dict[str, Any]]:
        '''Parse the LLM response to extract tool calls.

        Args:
            llm_response: Response from the LLM
            
        Returns:
            Dictionary with tool name and parameters, or None if no tool call
        '''
        pass


    async def execute_tool(self, tool_name: str, kwargs: Dict[str, Any]) -> ToolInvocationResult:
        '''Execute a tool with the given parameters.

        Args:
            tool_name: Name of the tool to invoke
            kwargs: Dictionary of parameters to pass to the tool
            
        Returns:
            ToolInvocationResult containing the tool's response
        '''
        return await self.mcp_client.invoke_tool(tool_name, kwargs)


    async def process_query(self, system_prompt: str, query: List[Dict]) -> Dict[str, Any]:
        '''Process a user query through the LLM and execute any tool calls.

        This method handles the full flow:
        1. Fetch tools if not already fetched
        2. Format tools for the LLM
        3. Submit query to LLM
        4. Parse tool calls from LLM response
        5. Execute tool if needed
        
        Args:
            query: User query string
            
        Returns:
            Dictionary containing the LLM response, tool call, and tool result
        '''
        # 1. Fetch tools if not already fetched
        if self.tools is None:
            await self.fetch_tools()

        # 2. Format tools for the LLM
        formatted_tools = await self.format_tools(self.tools)

        # 3. Submit query to LLM
        llm_response = await self.submit_query(system_prompt, query, formatted_tools)

        # 4. Parse tool calls from LLM response
        tool_call = await self.parse_tool_call(llm_response)

        result = {
            'llm_response': llm_response,
            'tool_call': tool_call,
            'tool_result': None
        }

        # 5. Execute tool if needed
        if tool_call:
            tool_name = tool_call.get('name')
            kwargs = tool_call.get('parameters', {})
            tool_result = await self.execute_tool(tool_name, kwargs)
            result['tool_result'] = tool_result

        return result


class AnthropicBridge(LLMBridge):
    '''Anthropic-specific implementation of the LLM Bridge.'''

    def __init__(self, mcp_client, api_key, model=DEFAULT_ANTHROPIC_MODEL): # Use imported default
        '''Initialize Anthropic bridge with API key and model.
        
        Args:
            mcp_client: An initialized MCPClient instance
            api_key: Anthropic API key
            model: Anthropic model to use (default: from models.py)
        '''
        super().__init__(mcp_client)
        self.llm_client = anthropic.Anthropic(api_key=api_key)
        self.model = model


    async def format_tools(self, tools: List[ToolDef]) -> List[Dict[str, Any]]:
        '''Format tools for Anthropic.
        
        Args:
            tools: List of ToolDef objects
            
        Returns:
            List of tools in Anthropic format
        '''
        return to_anthropic_format(tools)


    async def submit_query(
            self,
            system_prompt: str,
            query: List[Dict],
            formatted_tools: List[Dict[str, Any]]
    ) -> Dict[str, Any]:
        '''Submit a query to Anthropic with the formatted tools.
        
        Args:
            query: User query string
            formatted_tools: Tools in Anthropic format
            
        Returns:
            Anthropic API response
        '''
        response = self.llm_client.messages.create(
            model=self.model,
            max_tokens=4096,
            system=system_prompt,
            messages=query,
            tools=formatted_tools
        )

        return response


    async def parse_tool_call(self, llm_response: Any) -> Optional[Dict[str, Any]]:
        '''Parse the Anthropic response to extract tool calls.
        
        Args:
            llm_response: Response from Anthropic
            
        Returns:
            Dictionary with tool name and parameters, or None if no tool call
        '''
        for content in llm_response.content:
            if content.type == 'tool_use':
                return {
                    'name': content.name,
                    'parameters': content.input
                }

        return None



def to_anthropic_format(tools: List[ToolDef]) -> List[Dict[str, Any]]:
    '''Convert ToolDef objects to Anthropic tool format.
    
    Args:
        tools: List of ToolDef objects to convert
        
    Returns:
        List of dictionaries in Anthropic tool format
    '''

    anthropic_tools = []
    for tool in tools:
        anthropic_tool = {
            'name': tool.name,
            'description': tool.description,
            'input_schema': {
                'type': 'object',
                'properties': {},
                'required': []
            }
        }

        # Add properties
        for param in tool.parameters:
            # Map the type or use the original if no mapping exists
            schema_type = TYPE_MAPPING.get(param.parameter_type, param.parameter_type)

            param_schema = {
                'type': schema_type,  # Use mapped type
                'description': param.description
            }

            # For arrays, we need to specify the items type
            if schema_type == 'array':
                item_type = _infer_array_item_type(param)
                param_schema['items'] = {'type': item_type}

            anthropic_tool['input_schema']['properties'][param.name] = param_schema

            # Add default value if provided
            if param.default is not None:
                anthropic_tool['input_schema']['properties'][param.name]['default'] = param.default

            # Add to required list if required
            if param.required:
                anthropic_tool['input_schema']['required'].append(param.name)

        anthropic_tools.append(anthropic_tool)
    return anthropic_tools


def _infer_array_item_type(param: ToolParameter) -> str:
    '''Infer the item type for an array parameter based on its name and description.
    
    Args:
        param: The ToolParameter object
        
    Returns:
        The inferred JSON Schema type for array items
    '''
    # Default to string items
    item_type = 'string'

    # Check if parameter name contains hints about item type
    param_name_lower = param.name.lower()
    if any(hint in param_name_lower for hint in ['language', 'code', 'tag', 'name', 'id']):
        item_type = 'string'
    elif any(hint in param_name_lower for hint in ['number', 'count', 'amount', 'index']):
        item_type = 'integer'

    # Also check the description for hints
    if param.description:
        desc_lower = param.description.lower()
        if 'string' in desc_lower or 'text' in desc_lower or 'language' in desc_lower:
            item_type = 'string'
        elif 'number' in desc_lower or 'integer' in desc_lower or 'int' in desc_lower:
            item_type = 'integer'

    return item_type