app.py

from dotenv import load_dotenv
import gradio as gr
import cohere
from typing import Dict, List, Optional
import json
from dataclasses import dataclass
import os
from datetime import datetime

# Load environment variables
load_dotenv()


@dataclass
class IntentResponse:
    intent: str
    confidence: float
    entities: Dict
    suggested_action: str
    explanation: str


class EcommerceLLMIntentRecognizer:
    def __init__(self):
        # Get API key from environment variable
        api_key = os.getenv('COHERE_API_KEY')
        if not api_key:
            raise ValueError("Please set COHERE_API_KEY environment variable")

        self.co = cohere.Client(api_key)

        # Define our intent taxonomy
        self.valid_intents = {
            'product_search': 'SEARCH_CATALOG',
            'price_inquiry': 'FETCH_PRICE',
            'order_status': 'CHECK_ORDER_STATUS',
            'return_request': 'INITIATE_RETURN',
            'cart_management': 'MODIFY_CART',
            'availability_check': 'CHECK_INVENTORY',
            'checkout_help': 'ASSIST_CHECKOUT',
            'shipping_info': 'PROVIDE_SHIPPING_INFO',
            'product_comparison': 'COMPARE_PRODUCTS',
            'size_guide': 'SHOW_SIZE_GUIDE',
            'warranty_info': 'PROVIDE_WARRANTY_INFO',
            'cancel_order': 'PROCESS_CANCELLATION'
        }

    def _generate_prompt(self, query: str) -> str:
        return f"""As an e-commerce AI assistant, analyze the following customer query and extract the shopping intent, relevant entities, and determine the appropriate action.

Valid intents are: {', '.join(self.valid_intents.keys())}

Customer Query: "{query}"

Provide your analysis in the following JSON format:
{{
    "intent": "the_identified_intent",
    "confidence": 0.XX,
    "entities": {{
        "product": "identified_product",
        "category": "product_category",
        "specifications": ["any", "relevant", "specs"],
        "quantity": "if_mentioned",
        "order_id": "if_mentioned",
        "price_range": {{
            "min": "if_mentioned",
            "max": "if_mentioned"
        }}
    }},
    "explanation": "Brief explanation of why this intent was chosen"
}}

JSON Response:"""

    def recognize_intent(self, query: str) -> IntentResponse:
        # Generate LLM response
        response = self.co.generate(
            model='command',
            prompt=self._generate_prompt(query),
            max_tokens=500,
            temperature=0.2,
            k=0,
            stop_sequences=["\n\n"],
            return_likelihoods='NONE'
        )

        try:
            # Parse the LLM's response
            result = json.loads(response.generations[0].text)

            # Map to our action system
            suggested_action = self.valid_intents.get(
                result['intent'],
                'UNKNOWN_ACTION'
            )

            return IntentResponse(
                intent=result['intent'],
                confidence=result['confidence'],
                entities=result['entities'],
                suggested_action=suggested_action,
                explanation=result['explanation']
            )

        except json.JSONDecodeError:
            return IntentResponse(
                intent='parse_error',
                confidence=0.0,
                entities={},
                suggested_action='HANDLE_ERROR',
                explanation='Failed to parse LLM response'
            )


def process_query(user_query: str) -> str:
    try:
        recognizer = EcommerceLLMIntentRecognizer()
        response = recognizer.recognize_intent(user_query)

        return json.dumps({
            'timestamp': datetime.now().isoformat(),
            'query': user_query,
            'intent': response.intent,
            'confidence': response.confidence,
            'entities': response.entities,
            'suggested_action': response.suggested_action,
            'explanation': response.explanation
        }, indent=2)
    except ValueError as e:
        return json.dumps({
            'error': str(e),
            'hint': 'Please ensure COHERE_API_KEY is set in your .env file'
        }, indent=2)


# Create Gradio interface
iface = gr.Interface(
    fn=process_query,
    inputs=gr.Textbox(label="Enter customer query"),
    outputs=gr.JSON(label="Intent Analysis"),
    title="E-commerce LLM Intent Recognition System",
    description="""This system uses Cohere's Command model to understand customer intentions in an e-commerce context.
                  Enter your query to see the detailed intent analysis.""",
    examples=[
        ["I'm looking for a waterproof smart watch under $300"],
        ["Can you compare the iPhone 13 and iPhone 14 Pro?"],
        ["Need to return my order #ABC123, it's the wrong size"],
        ["Do you have this dress in size medium and in red?"],
        ["What's your shipping time to California?"]
    ]
)

if __name__ == "__main__":
    iface.launch()