llm_client.py

"""
LLM Client for natural language to MCP tool calling
"""
import ollama
import json
import logging
from typing import Dict, Any, Optional, List
import asyncio

logger = logging.getLogger(__name__)

class LLMToolCaller:
    """Lightweight LLM that can naturally call MCP tools"""
    
    def __init__(self, model_name: str = "llama3.2:1b"):
        self.model_name = model_name
        self.tools = [
            {
                "name": "search_products",
                "description": "Search for products in the PrestaShop catalog using keywords",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "Search terms for finding products"
                        }
                    },
                    "required": ["query"]
                }
            },
            {
                "name": "get_product_details",
                "description": "Get detailed information about a specific product by ID",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "product_id": {
                            "type": "string",
                            "description": "The ID of the product to get details for"
                        }
                    },
                    "required": ["product_id"]
                }
            },
            {
                "name": "get_product_features",
                "description": "Get the features and specifications of a specific product",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "product_id": {
                            "type": "string",
                            "description": "The ID of the product to get features for"
                        }
                    },
                    "required": ["product_id"]
                }
            },
            {
                "name": "get_product_images",
                "description": "Get images/photos of a specific product",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "product_id": {
                            "type": "string",
                            "description": "The ID of the product to get images for"
                        }
                    },
                    "required": ["product_id"]
                }
            }
        ]

    async def parse_message_for_tool_call(self, message: str) -> tuple[Optional[str], Dict[str, Any]]:
        """Use LLM to determine which tool to call and with what parameters"""
        
        # Create a prompt that helps the LLM understand tool selection
        system_prompt = f"""You are a PrestaShop assistant. Analyze the user message and determine which tool to call.

Available tools:
{json.dumps(self.tools, indent=2)}

User message: "{message}"

Respond with ONLY a JSON object in this format:
{{"tool_name": "tool_name", "parameters": {{"param": "value"}}}}

Examples:
- "tell me about product 12" → {{"tool_name": "get_product_details", "parameters": {{"product_id": "12"}}}}
- "search for shoes" → {{"tool_name": "search_products", "parameters": {{"query": "shoes"}}}}
- "features of product 5" → {{"tool_name": "get_product_features", "parameters": {{"product_id": "5"}}}}
- "images of product 3" → {{"tool_name": "get_product_images", "parameters": {{"product_id": "3"}}}}

If the message is conversational and doesn't need tools, respond with: {{"tool_name": null, "parameters": {{}}}}
"""

        try:
            # Call the local LLM
            response = ollama.chat(
                model=self.model_name,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": message}
                ]
            )
            
            # Parse the response
            content = response['message']['content'].strip()
            logger.info(f"🤖 LLM response: {content}")
            
            # Try to parse as JSON
            try:
                parsed = json.loads(content)
                tool_name = parsed.get("tool_name")
                parameters = parsed.get("parameters", {})
                
                if tool_name and tool_name != "null":
                    logger.info(f"🔧 LLM selected tool: {tool_name} with params: {parameters}")
                    return tool_name, parameters
                else:
                    logger.info("💬 LLM determined this is conversational, no tool needed")
                    return None, {}
                    
            except json.JSONDecodeError:
                logger.warning(f"Failed to parse LLM response as JSON: {content}")
                # Fallback to simple pattern matching
                return self._fallback_tool_selection(message)
                
        except Exception as e:
            logger.error(f"Error calling LLM: {e}")
            # Fallback to simple pattern matching
            return self._fallback_tool_selection(message)

    def _fallback_tool_selection(self, message: str) -> tuple[Optional[str], Dict[str, Any]]:
        """Fallback tool selection if LLM fails"""
        message_lower = message.lower()
        words = message.split()
        
        # Look for product ID
        product_id = None
        for word in words:
            if word.isdigit():
                product_id = word
                break
        
        # Simple pattern matching
        if any(pattern in message_lower for pattern in ['details', 'about product', 'tell me about']):
            if product_id:
                return "get_product_details", {"product_id": product_id}
        
        if any(pattern in message_lower for pattern in ['features', 'specifications']):
            if product_id:
                return "get_product_features", {"product_id": product_id}
        
        if any(pattern in message_lower for pattern in ['images', 'photos', 'pictures']):
            if product_id:
                return "get_product_images", {"product_id": product_id}
        
        # Default to search
        if any(pattern in message_lower for pattern in ['search', 'find', 'look for']):
            return "search_products", {"query": message}
        
        return None, {}