main.py

import os
import json
import logging
import traceback
from typing import Dict, Any
import cohere
from agents.email_agent import process_email
from agents.json_agent import process_json
from agents.pdf_agent import process_pdf
from memory.memory_store import init_db, log_input, update_input_status
from utils.pdf_loader import read_pdf
from action_router import ActionRouter

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize components
init_db()
action_router = ActionRouter()

# Setup Cohere client
try:
    co = cohere.Client(os.getenv("COHERE_API_KEY"))
except Exception as e:
    logger.error(f"Failed to initialize Cohere client: {e}")
    co = None

def detect_format(content: str) -> str:
    """Detect input format with improved validation"""
    if not content:
        return "Unknown"
    
    if isinstance(content, str) and content.lower().endswith('.pdf'):
        return "PDF"
    
    # Check content patterns
    content_lower = content.lower()
    
    try:
        if content.strip().startswith("{") or '"json"' in content_lower:
            json.loads(content)  # Validate JSON
            return "JSON"
    except json.JSONDecodeError:
        pass
        
    if any(x in content_lower for x in ["@gmail.com", "@yahoo.com", "subject:", "from:"]):
        return "Email"
    elif any(x in content_lower for x in ["invoice", "policy", "statement"]):
        return "PDF"
    
    return "Unknown"

def detect_intent(content: str) -> str:
    """Detect intent with fallback mechanisms"""
    if not content:
        return "Unknown"
    
    content_lower = content.lower()
    
    # Pattern matching first
    if any(x in content_lower for x in ["invoice", "bill", "payment"]):
        return "Invoice"
    elif any(x in content_lower for x in ["complaint", "angry", "dissatisfied"]):
        return "Complaint"
    elif any(x in content_lower for x in ["rfq", "request for quote", "quotation"]):
        return "RFQ"
    elif any(x in content_lower for x in ["regulation", "compliance", "gdpr"]):
        return "Regulation"
    
    # LLM fallback
    if not co:
        logger.warning("Cohere client not available for intent detection")
        return "Unknown"
    
    prompt = f"""Classify the business intent of this message. Respond ONLY with one of:
- RFQ
- Complaint
- Invoice
- Regulation
- Fraud Risk
- Unknown\n\n{content[:1000]}"""
    
    try:
        response = co.generate(
            model='command',
            prompt=prompt,
            max_tokens=10,
            temperature=0.3
        )
        intent = response.generations[0].text.strip().split('\n')[0]
        return intent if intent in ["RFQ", "Complaint", "Invoice", "Regulation", "Fraud Risk"] else "Unknown"
    except Exception as e:
        logger.error(f"LLM intent detection failed: {e}")
        return "Unknown"

def classify_input(content: str) -> Dict[str, str]:
    """Classify input with error handling"""
    try:
        return {
            "format": detect_format(content),
            "intent": detect_intent(content)
        }
    except Exception as e:
        logger.error(f"Classification failed: {e}")
        return {"format": "Unknown", "intent": "Unknown"}

def route_input(content: str) -> Dict[str, Any]:
    """Route input to appropriate agent with full tracing"""
    classification = classify_input(content)
    logger.info(f"Classification: {classification}")
    
    input_id = log_input(
        source_type=classification["format"],
        intent=classification["intent"],
        format_=classification["format"],
        raw_content=content[:1000]  # Store sample
    )
    
    try:
        result = None
        
        if classification["format"] == "PDF":
            pdf_text = read_pdf(content) if os.path.isfile(content) else content
            result = process_pdf(pdf_text)  # ✅ pdf_text is now actual text, not a path

        elif classification["format"] == "Email":
            result = process_email(content)
            
        elif classification["format"] == "JSON":
            try:
                json_data = json.loads(content)
                result = process_json(json_data)
            except json.JSONDecodeError as e:
                result = {"error": f"Invalid JSON: {str(e)}"}
                
        else:
            result = {"error": f"Unsupported format: {classification['format']}"}
        
        # Update status and route actions
        if "error" in result:
            update_input_status(input_id, "failed")
            logger.error(f"Processing failed: {result['error']}")
        else:
            update_input_status(input_id, "processed")
            action_router.handle_agent_output({
                **result,
                "input_id": input_id,
                "classification": classification
            })
            
        return {**result, "input_id": input_id}
        
    except Exception as e:
        error_msg = f"Processing failed: {str(e)}"
        update_input_status(input_id, "failed")
        logger.error(error_msg)
        return {"error": error_msg, "input_id": input_id}

if __name__ == "__main__":
    test_cases = [
        r"C:\Users\omlap\OneDrive\Desktop\multi_agent_ai_system\inputs\sample_invoice.pdf",
        r"C:\Users\omlap\OneDrive\Desktop\multi_agent_ai_system\inputs\sample_email.txt",
        r"C:\Users\omlap\OneDrive\Desktop\multi_agent_ai_system\inputs\fake_pdf.pdf",
    ]

    for idx, test_case in enumerate(test_cases):
        try:
            print("\n" + "=" * 70)
            print(f"🧪 Processing Test Case {idx + 1}")
            
            # Show path if it's a file
            if isinstance(test_case, str) and os.path.exists(test_case):
                print(f"📄 File Input: {os.path.basename(test_case)}")
            elif len(test_case) < 100:
                print(f"📝 Raw Input Preview: {test_case.strip()}")
            else:
                print(f"📝 Raw Input (Truncated): {test_case[:100]}...")
            if test_case.lower().endswith('.pdf'):
                content = read_pdf(test_case)  # ✅ Extract actual text
            elif test_case.lower().endswith('.txt'):
                with open(test_case, 'r', encoding='utf-8') as f:
                    content = f.read()
            else:
                content = test_case  # Raw text or unknown case

            result = route_input(content)

            # Output result
            print("📤 Final Result:")
            print(json.dumps(result, indent=2))

        except Exception as e:
            print(f"❌ Test Case {idx + 1} Failed: {str(e)}")
            print(traceback.format_exc())