src/transaction_clarifier.py

import openai
import os
import json
from typing import Dict, Any, Optional, List, Tuple
from enum import Enum
from pydantic import BaseModel
from models import EntityExtraction

class ClarificationStatus(str, Enum):
    COMPLETE = "complete"
    NEEDS_CLARIFICATION = "needs_clarification"
    CANCELLED = "cancelled"

class ClarificationRequest(BaseModel):
    missing_fields: List[str]
    questions: List[str]
    suggested_values: Dict[str, Any] = {}
    explanation: str

class TransactionClarifier:
    def __init__(self, api_key: Optional[str] = None):
        """Initialize OpenAI client for transaction clarification"""
        self.client = openai.OpenAI(
            api_key=api_key or os.getenv('OPENAI_API_KEY')
        )
    
    def analyze_transaction_completeness(self, entities: EntityExtraction) -> Tuple[ClarificationStatus, Optional[ClarificationRequest]]:
        """
        Analyze if a transaction has all necessary information
        
        Args:
            entities: Extracted entities from user input
            
        Returns:
            Tuple of (status, clarification_request)
        """
        
        # Define required and optional fields based on transaction type
        if entities.transaction_type == "purchase":
            required_fields = ["product", "quantity", "supplier", "unit_price"]
            optional_fields = ["total_amount"]
        elif entities.transaction_type == "sale":
            required_fields = ["product", "quantity", "customer", "unit_price"]
            optional_fields = ["total_amount"]
        else:
            return ClarificationStatus.COMPLETE, None
        
        # Check for missing required fields
        missing_fields = []
        entity_dict = entities.dict()
        
        for field in required_fields:
            if not entity_dict.get(field):
                missing_fields.append(field)
        
        # If all required fields are present, transaction is complete
        if not missing_fields:
            return ClarificationStatus.COMPLETE, None
        
        # Generate intelligent clarification request
        clarification = self._generate_clarification_request(entities, missing_fields)
        
        return ClarificationStatus.NEEDS_CLARIFICATION, clarification
    
    def _generate_clarification_request(self, entities: EntityExtraction, missing_fields: List[str]) -> ClarificationRequest:
        """Generate intelligent questions for missing information"""
        
        # Prepare context about what we already know
        known_info = {}
        entity_dict = entities.dict()
        
        for field, value in entity_dict.items():
            if value is not None and field != "notes":
                known_info[field] = value
        
        system_prompt = f"""You are a helpful business assistant helping complete a {entities.transaction_type} transaction.

Generate natural, conversational questions to gather missing information. The user should be able to:
1. Provide the missing information
2. Say "N/A" or "skip" if the information is not available/applicable
3. Ask for suggestions if they're unsure

Create personalized questions based on the context of what we already know.

Return your response in this exact JSON format:
{{
    "questions": ["question1", "question2", ...],
    "suggested_values": {{"field": "suggested_value", ...}},
    "explanation": "Brief explanation of why we need this information"
}}

Missing fields to ask about: {missing_fields}
Transaction type: {entities.transaction_type}
"""

        user_prompt = f"""We're processing a {entities.transaction_type} transaction and need to gather some missing information.

What we already know:
{json.dumps(known_info, indent=2)}

Missing fields: {missing_fields}

Generate friendly, specific questions to gather the missing information. Make suggestions when appropriate."""

        try:
            response = self.client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=0.3,
                max_tokens=400
            )
            
            response_text = response.choices[0].message.content.strip()
            
            try:
                result_dict = json.loads(response_text)
                return ClarificationRequest(
                    missing_fields=missing_fields,
                    questions=result_dict.get("questions", []),
                    suggested_values=result_dict.get("suggested_values", {}),
                    explanation=result_dict.get("explanation", "I need some additional information to complete this transaction.")
                )
            except (json.JSONDecodeError, KeyError) as e:
                # Fallback to simple questions
                return self._generate_fallback_questions(entities, missing_fields)
            
        except Exception as e:
            print(f"Error generating clarification: {e}")
            return self._generate_fallback_questions(entities, missing_fields)
    
    def _generate_fallback_questions(self, entities: EntityExtraction, missing_fields: List[str]) -> ClarificationRequest:
        """Generate fallback questions when LLM fails"""
        
        question_templates = {
            "product": "What product or item is involved in this transaction?",
            "quantity": f"How many units {'were purchased' if entities.transaction_type == 'purchase' else 'were sold'}?",
            "supplier": "Which supplier or vendor is this purchase from?",
            "customer": "Who is the customer for this sale?",
            "unit_price": "What is the price per unit?",
            "total_amount": "What is the total amount for this transaction?"
        }
        
        questions = []
        for field in missing_fields:
            questions.append(question_templates.get(field, f"What is the {field.replace('_', ' ')}?"))
        
        return ClarificationRequest(
            missing_fields=missing_fields,
            questions=questions,
            suggested_values={},
            explanation="I need some additional information to complete this transaction."
        )
    
    def process_clarification_response(self, original_entities: EntityExtraction, 
                                     missing_fields: List[str], 
                                     user_response: str) -> Tuple[EntityExtraction, bool]:
        """
        Process user's response to clarification questions
        
        Args:
            original_entities: Original extracted entities
            missing_fields: Fields we asked about
            user_response: User's response to our questions
            
        Returns:
            Tuple of (updated_entities, is_complete)
        """
        
        system_prompt = f"""You are processing a user's response to clarification questions about a {original_entities.transaction_type} transaction.

Extract the missing information from the user's response. The user may:
1. Provide specific values for the missing fields
2. Say "N/A", "skip", "not applicable", or similar to indicate the field should be null
3. Ask for help or say they don't know

Missing fields we asked about: {missing_fields}

Return a JSON object with the extracted values. Use null for fields that are N/A or skipped.

Example response format:
{{
    "product": "extracted product name",
    "quantity": 10,
    "supplier": null,
    "unit_price": 5.99,
    "interpretation": "Brief explanation of what you extracted"
}}"""

        user_prompt = f"""Original transaction: {original_entities.transaction_type}
Missing fields: {missing_fields}
User's response: "{user_response}"

Extract the values for the missing fields from the user's response."""

        try:
            response = self.client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=0.1,
                max_tokens=300
            )
            
            response_text = response.choices[0].message.content.strip()
            
            try:
                extracted_values = json.loads(response_text)
                
                # Update original entities with extracted values
                updated_entities = self._update_entities(original_entities, extracted_values, missing_fields)
                
                # Check if transaction is now complete
                status, _ = self.analyze_transaction_completeness(updated_entities)
                is_complete = (status == ClarificationStatus.COMPLETE)
                
                return updated_entities, is_complete
                
            except (json.JSONDecodeError, KeyError) as e:
                print(f"Error parsing clarification response: {e}")
                return original_entities, False
                
        except Exception as e:
            print(f"Error processing clarification: {e}")
            return original_entities, False
    
    def _update_entities(self, original_entities: EntityExtraction, 
                        extracted_values: Dict[str, Any], 
                        missing_fields: List[str]) -> EntityExtraction:
        """Update entities with extracted clarification values"""
        
        # Convert to dict for easier manipulation
        entity_dict = original_entities.dict()
        
        # Update with extracted values
        for field in missing_fields:
            if field in extracted_values:
                value = extracted_values[field]
                
                # Handle type conversions
                if field in ["quantity"] and value is not None:
                    try:
                        entity_dict[field] = int(value)
                    except (ValueError, TypeError):
                        entity_dict[field] = None
                elif field in ["unit_price", "total_amount"] and value is not None:
                    try:
                        entity_dict[field] = float(value)
                    except (ValueError, TypeError):
                        entity_dict[field] = None
                else:
                    entity_dict[field] = value
        
        # Recalculate total if we have quantity and unit_price
        if entity_dict.get("quantity") and entity_dict.get("unit_price"):
            entity_dict["total_amount"] = entity_dict["quantity"] * entity_dict["unit_price"]
        
        return EntityExtraction(**entity_dict)
    
    def format_clarification_message(self, clarification: ClarificationRequest) -> str:
        """Format clarification request as a user-friendly message"""
        
        message = f"📝 {clarification.explanation}\n\n"
        
        for i, question in enumerate(clarification.questions, 1):
            message += f"{i}. {question}\n"
        
        # Add suggestions if available
        if clarification.suggested_values:
            message += "\n💡 Suggestions:\n"
            for field, suggestion in clarification.suggested_values.items():
                message += f"   • {field.replace('_', ' ').title()}: {suggestion}\n"
        
        message += "\n✨ You can say 'N/A' or 'skip' for any information that's not available."
        message += "\n📞 Please provide the missing information in your next message."
        
        return message