app.py

#!/usr/bin/env python3
"""
Enhanced Budget Proposals Chatbot API using LangChain with Memory and Agentic RAG
"""

from flask import Flask, request, jsonify
from flask_cors import CORS
import os
import logging
import json
from datetime import datetime
from typing import Dict, List, Any

# LangChain imports
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.memory import ConversationBufferWindowMemory
from langchain.schema import HumanMessage, AIMessage
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.chains import LLMChain
from langchain_community.chat_message_histories import RedisChatMessageHistory
from langchain.tools import Tool
from langchain.agents import AgentExecutor, create_openai_functions_agent
from langchain.agents.openai_functions_agent.base import OpenAIFunctionsAgent
from langchain.schema import BaseMessage

# Vector database imports
from pinecone import Pinecone
from sentence_transformers import SentenceTransformer

# Language detection imports
import re
import requests
import json

app = Flask(__name__)
CORS(app)

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

# Configure Gemini
GEMINI_API_KEY = os.getenv('GEMINI_API_KEY')
if not GEMINI_API_KEY:
    logger.error("GEMINI_API_KEY not found in environment variables")
    raise ValueError("Please set GEMINI_API_KEY in your .env file")

# Configure Pinecone
PINECONE_API_KEY = os.getenv('PINECONE_API_KEY')
if not PINECONE_API_KEY:
    logger.error("PINECONE_API_KEY not found in environment variables")
    raise ValueError("Please set PINECONE_API_KEY in your .env file")

# Initialize Pinecone and embedding model
pc = Pinecone(api_key=PINECONE_API_KEY)
BUDGET_INDEX_NAME = "budget-proposals-index"
embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# Initialize LangChain components
llm = ChatGoogleGenerativeAI(
    model="gemini-2.5-flash",
    google_api_key=GEMINI_API_KEY,
    temperature=0.7,
    max_tokens=2000  # Increased for longer Sinhala responses
)

# Simplified initialization - Let Gemini handle everything
logger.info("Using Gemini for all language processing (transliteration, translation, responses)")

def detect_sinhala_content(text: str) -> bool:
    """Detect if text contains Sinhala characters"""
    # Sinhala Unicode range: U+0D80 to U+0DFF
    sinhala_pattern = re.compile(r'[\u0D80-\u0DFF]')
    return bool(sinhala_pattern.search(text))

def detect_tamil_content(text: str) -> bool:
    """Detect if text contains Tamil characters"""
    # Tamil Unicode range: U+0B80 to U+0BFF
    tamil_pattern = re.compile(r'[\u0B80-\u0BFF]')
    return bool(tamil_pattern.search(text))

def simple_detect_language(text: str) -> Dict[str, Any]:
    """Simplified language detection with Tamil support - let Gemini handle the complexity"""
    try:
        # Check for Sinhala Unicode first (most reliable)
        has_sinhala_unicode = detect_sinhala_content(text)
        if has_sinhala_unicode:
            return {
                'language': 'si',
                'confidence': 0.95,
                'is_sinhala_unicode': True,
                'is_tamil_unicode': False,
                'is_romanized_sinhala': False,
                'is_english': False,
                'detection_method': 'unicode_detection'
            }
        
        # Check for Tamil Unicode
        has_tamil_unicode = detect_tamil_content(text)
        if has_tamil_unicode:
            return {
                'language': 'ta',
                'confidence': 0.95,
                'is_sinhala_unicode': False,
                'is_tamil_unicode': True,
                'is_romanized_sinhala': False,
                'is_english': False,
                'detection_method': 'unicode_detection'
            }
        
        # Use enhanced rule-based detection for Singlish
        return enhanced_rule_based_detection(text)
        
    except Exception as e:
        logger.error(f"Language detection failed: {e}")
        return rule_based_language_detection(text)

def enhanced_rule_based_detection(text: str) -> Dict[str, Any]:
    """Enhanced rule-based detection with Singlish and Romanized Tamil recognition"""
    has_sinhala_unicode = detect_sinhala_content(text)
    has_tamil_unicode = detect_tamil_content(text)
    is_romanized_sinhala = detect_singlish(text) and not has_sinhala_unicode and not has_tamil_unicode
    is_romanized_tamil = detect_romanized_tamil(text) and not has_sinhala_unicode and not has_tamil_unicode and not is_romanized_sinhala
    
    # More sophisticated Singlish detection
    if not has_sinhala_unicode and not is_romanized_sinhala:
        # Check for common Sinhala sentence patterns in English letters
        sinhala_patterns = [
            r'\b(mokadda|kohomada|api|oya|mama)\b',
            r'\b(eka|meka|thiyenne|kiyala)\b',
            r'\b(gana|genna|danna|karanna)\b',
            r'\b(budget|proposal).*\b(gana|eka)\b'
        ]
        
        text_lower = text.lower()
        pattern_matches = sum(1 for pattern in sinhala_patterns if re.search(pattern, text_lower))
        
        if pattern_matches >= 1:  # Lower threshold for better detection
            is_romanized_sinhala = True
    
    if has_sinhala_unicode:
        language_code = 'si'
        confidence = 0.9
    elif has_tamil_unicode:
        language_code = 'ta'
        confidence = 0.9
    elif is_romanized_sinhala:
        language_code = 'singlish'
        confidence = 0.8
    elif is_romanized_tamil:
        language_code = 'romanized_tamil'
        confidence = 0.8
    else:
        language_code = 'en'
        confidence = 0.7
    
    return {
        'language': language_code,
        'confidence': confidence,
        'is_sinhala_unicode': has_sinhala_unicode,
        'is_tamil_unicode': has_tamil_unicode,
        'is_romanized_sinhala': is_romanized_sinhala,
        'is_romanized_tamil': is_romanized_tamil,
        'is_english': language_code == 'en',
        'detection_method': 'enhanced_rule_based'
    }

def rule_based_language_detection(text: str) -> Dict[str, Any]:
    """Fallback rule-based language detection with Tamil and Romanized Tamil support"""
    has_sinhala_unicode = detect_sinhala_content(text)
    has_tamil_unicode = detect_tamil_content(text)
    is_romanized_sinhala = detect_singlish(text) and not has_sinhala_unicode and not has_tamil_unicode
    is_romanized_tamil = detect_romanized_tamil(text) and not has_sinhala_unicode and not has_tamil_unicode and not is_romanized_sinhala
    is_english = not has_sinhala_unicode and not has_tamil_unicode and not is_romanized_sinhala and not is_romanized_tamil
    
    if has_sinhala_unicode:
        language_code = 'si'
    elif has_tamil_unicode:
        language_code = 'ta'
    elif is_romanized_sinhala:
        language_code = 'singlish'
    elif is_romanized_tamil:
        language_code = 'romanized_tamil'
    else:
        language_code = 'en'
    
    return {
        'language': language_code,
        'confidence': 0.8,  # Default confidence for rule-based
        'is_sinhala_unicode': has_sinhala_unicode,
        'is_tamil_unicode': has_tamil_unicode,
        'is_romanized_sinhala': is_romanized_sinhala,
        'is_romanized_tamil': is_romanized_tamil,
        'is_english': is_english,
        'detection_method': 'rule_based'
    }

def detect_singlish(text: str) -> bool:
    """Detect common Singlish patterns and words"""
    singlish_words = [
        'mokadda', 'kohomada', 'api', 'oya', 'mama', 'eka', 'meka', 'oya', 'dan', 'kiyala',
        'budget', 'proposal', 'karan', 'karanna', 'gana', 'genna', 'danna', 'ahala', 'denna',
        'mata', 'ape', 'wage', 'wenas', 'thiyenne', 'kiyanawa', 'balanawa', 'pennanna',
        'sampura', 'mudal', 'pasal', 'vyaparayak', 'rajaye', 'arthikaya', 'sammandala',
        'kara', 'karanna', 'giya', 'yanawa', 'enawa', 'gihin', 'awe', 'nane', 'inne',
        'danna', 'kiyanna', 'balanna', 'ganna', 'denna', 'yanna', 'enna'
    ]
    
    # Convert to lowercase and check for common Singlish words
    text_lower = text.lower()
    singlish_word_count = sum(1 for word in singlish_words if word in text_lower)
    
    # Consider it Singlish if it has 2 or more Singlish words
    return singlish_word_count >= 2

def detect_romanized_tamil(text: str) -> bool:
    """Detect common Romanized Tamil patterns and words (Tamil written in English letters)"""
    romanized_tamil_words = [
        # Common Tamil words in Roman script
        'enna', 'epdi', 'enga', 'yaar', 'naa', 'nee', 'avar', 'ivan', 'ival', 'ithu', 'athu',
        'vandhu', 'ponga', 'vanga', 'sollu', 'kelu', 'paaru', 'irukku', 'irukkanga', 'irundhu',
        'seiya', 'panna', 'mudiyum', 'mudiyathu', 'venum', 'vendam', 'puriyuthu', 'puriyala',
        'nalla', 'ketta', 'romba', 'konjam', 'neraya', 'kammi', 'adhikam', 'thaan', 'daan',
        # Budget/government related Tamil terms
        'budget', 'proposal', 'sarkar', 'arasaangam', 'vyavasai', 'panam', 'kaasu', 'thogai',
        'nilai', 'mari', 'maatram', 'thiruththam', 'yojana', 'thittam', 'mudhal', 'selavu',
        'varumanam', 'aayam', 'EPF', 'viduli', 'current', 'maternity', 'leave'
    ]
    
    # Convert to lowercase and check for common Romanized Tamil words
    text_lower = text.lower()
    tamil_word_count = sum(1 for word in romanized_tamil_words if word in text_lower)
    
    # Consider it Romanized Tamil if it has 2 or more Tamil words
    return tamil_word_count >= 2

# Removed: AI transliteration and Google Translate functions
# Gemini will handle all transliteration and translation needs

def simple_process_input(user_message: str) -> tuple:
    """
    Simplified input processing - let Gemini handle everything
    """
    # Step 1: Simple language detection
    language_info = simple_detect_language(user_message)
    original_language = language_info['language']
    confidence = language_info['confidence']
    detection_method = language_info['detection_method']
    
    logger.info(f"Language detection: {original_language} (confidence: {confidence:.2f}, method: {detection_method})")
    
    # Use original message for all processing - Gemini will handle the rest
    processed_message = user_message
    needs_translation = False  # Gemini handles translation internally
    transliteration_used = False  # Gemini handles transliteration internally
    ai_detection_used = detection_method == 'ai'
    
    logger.info(f"Input processing: keeping original '{user_message}' for Gemini to handle")
    
    return processed_message, original_language, needs_translation, transliteration_used, ai_detection_used, confidence

# Removed: translate_response_if_needed function
# Gemini handles all language responses automatically

def get_pinecone_index():
    """Get the budget proposals Pinecone index"""
    try:
        return pc.Index(BUDGET_INDEX_NAME)
    except Exception as e:
        logger.error(f"Error accessing Pinecone index: {e}")
        return None

def search_budget_proposals(query: str) -> str:
    """Search budget proposals using the semantic search API"""
    try:
        import requests
        
        # Use the deployed semantic search API
        response = requests.post(
            f"https://danulr05-budget-proposals-search-api.hf.space/api/search",
            json={"query": query, "top_k": 5},
            timeout=10
        )
        
        if response.status_code == 200:
            data = response.json()
            results = data.get("results", [])
            
            if not results:
                return "No relevant budget proposals found in the database."
            
            # Build context from search results
            context_parts = []
            for result in results[:3]:  # Limit to top 3 results
                file_path = result.get("file_path", "")
                category = result.get("category", "")
                summary = result.get("summary", "")
                cost = result.get("costLKR", "")
                title = result.get("title", "")
                content = result.get("content", "")  # Get the actual content
                
                context_parts.append(f"From {file_path} ({category}): {title}")
                if content:
                    context_parts.append(f"Content: {content}")
                elif summary:
                    context_parts.append(f"Summary: {summary}")
                if cost and cost != "No Costing Available":
                    context_parts.append(f"Cost: {cost}")
            
            return "\n\n".join(context_parts)
        else:
            return f"Error accessing semantic search API: {response.status_code}"
        
    except Exception as e:
        logger.error(f"Error searching budget proposals: {e}")
        return f"Error searching database: {str(e)}"

# Create the RAG tool
search_tool = Tool(
    name="search_budget_proposals",
    description="Search for relevant budget proposals in the vector database. Use this when you need specific information about budget proposals, costs, policies, or implementation details.",
    func=search_budget_proposals
)

# Create the prompt template for the agent
agent_prompt = ChatPromptTemplate.from_messages([
    ("system", """You are a helpful assistant for budget proposals in Sri Lanka. You have access to a vector database containing detailed information about various budget proposals. You can communicate in English, Sinhala, and understand Singlish (Sinhala written in English letters).

When a user asks about budget proposals, you should:
1. Use the search_budget_proposals tool to find relevant information
2. Provide accurate, detailed responses based on the retrieved information
3. Always cite the source documents when mentioning specific proposals
4. Be professional but approachable in any language
5. If the search doesn't return relevant results, acknowledge this and provide general guidance
6. Respond in the same language or style as the user's question when possible

Guidelines:
- Always use the search tool for specific questions about budget proposals
- Include source citations for any mention of proposals, costs, policies, revenue, or implementation
- Keep responses clear and informative in any language
- Use a balanced tone - helpful but not overly casual
- If asked about topics not covered, redirect to relevant topics professionally
- Be culturally sensitive when discussing Sri Lankan policies and economic matters
- When responding in Sinhala, use appropriate formal language for policy discussions"""),
    MessagesPlaceholder(variable_name="chat_history"),
    ("human", "{input}"),
    MessagesPlaceholder(variable_name="agent_scratchpad")
])

# Store conversation memories for different sessions
conversation_memories: Dict[str, ConversationBufferWindowMemory] = {}

def get_or_create_memory(session_id: str) -> ConversationBufferWindowMemory:
    """Get or create a memory instance for a session"""
    if session_id not in conversation_memories:
        # Create new memory with window of 10 messages (5 exchanges)
        conversation_memories[session_id] = ConversationBufferWindowMemory(
            k=10,  # Remember last 10 messages
            return_messages=True,
            memory_key="chat_history"
        )
        logger.info(f"Created new memory for session: {session_id}")
    
    return conversation_memories[session_id]

def create_agent(session_id: str) -> AgentExecutor:
    """Create a LangChain agent with memory and RAG capabilities"""
    memory = get_or_create_memory(session_id)
    
    # Create the agent
    agent = create_openai_functions_agent(
        llm=llm,
        tools=[search_tool],
        prompt=agent_prompt
    )
    
    # Create agent executor with memory
    agent_executor = AgentExecutor(
        agent=agent,
        tools=[search_tool],
        memory=memory,
        verbose=False,
        handle_parsing_errors=True
    )
    
    return agent_executor

def get_available_pdfs() -> List[str]:
    """Dynamically get list of available PDF files from assets directory"""
    try:
        import os
        pdf_dir = "assets/pdfs"
        if os.path.exists(pdf_dir):
            pdf_files = [f for f in os.listdir(pdf_dir) if f.lower().endswith('.pdf')]
            return pdf_files
        else:
            # Fallback to known PDFs if directory doesn't exist
            return ['MLB.pdf', 'Cigs.pdf', 'Elec.pdf', 'Audit_EPF.pdf', 'EPF.pdf', 'Discretion.pdf', '1750164001872.pdf']
    except Exception as e:
        logger.error(f"Error getting available PDFs: {e}")
        # Fallback to known PDFs
        return ['MLB.pdf', 'Cigs.pdf', 'Elec.pdf', 'Audit_EPF.pdf', 'EPF.pdf', 'Discretion.pdf', '1750164001872.pdf']

def extract_sources_from_response(response: str) -> List[str]:
    """Extract source documents mentioned in the response"""
    sources = []
    
    # Get dynamically available PDF files
    available_pdfs = get_available_pdfs()
    
    # Look for source patterns like "(Source: MLB.pdf)" or "(Sources: MLB.pdf, EPF.pdf)"
    for pdf in available_pdfs:
        if pdf in response:
            sources.append(pdf)
    
    return list(set(sources))  # Remove duplicates

def generate_response_with_rag(user_message: str, session_id: str) -> Dict[str, Any]:
    """Generate response using RAG with memory and multilingual support"""
    try:
        # Process multilingual input
        processed_message, original_language, needs_translation, transliteration_used, ai_detection_used, confidence = simple_process_input(user_message)
        logger.info(f"Input processing: original='{user_message}', processed='{processed_message}', lang='{original_language}', transliteration='{transliteration_used}', ai_detection='{ai_detection_used}', confidence='{confidence:.2f}'")
        
        # Get or create memory for this session
        memory = get_or_create_memory(session_id)
        
        # Let Gemini handle both specific and general questions intelligently
        # Always search with the user's actual query - Gemini will handle vague questions
        search_context = search_budget_proposals(processed_message)
        
        # Get conversation history for context
        chat_history = memory.chat_memory.messages
        conversation_context = ""
        if chat_history:
            # Get last few messages for context
            recent_messages = chat_history[-6:]  # Last 3 exchanges
            conversation_parts = []
            for msg in recent_messages:
                if isinstance(msg, HumanMessage):
                    conversation_parts.append(f"User: {msg.content}")
                elif isinstance(msg, AIMessage):
                    conversation_parts.append(f"Assistant: {msg.content}")
            conversation_context = "\n".join(conversation_parts)
        
        # Create a prompt with conversation history and retrieved context
        language_instruction = ""
        if original_language == 'si':
            language_instruction = "\n\nIMPORTANT: The user asked in Sinhala. Please respond in the same language (Sinhala) using proper Sinhala script and formal language appropriate for policy discussions. The question was: '{}'".format(user_message)
        elif original_language == 'ta':
            language_instruction = "\n\nIMPORTANT: The user asked in Tamil. Please respond in the same language (Tamil) using proper Tamil script and formal language appropriate for policy discussions. Use Sri Lankan Tamil terminology and context. The question was: '{}'".format(user_message)
        elif original_language == 'singlish':
            language_instruction = "\n\nIMPORTANT: The user asked in Singlish (Romanized Sinhala - Sinhala words written in English letters). Please respond in proper Sinhala script using formal language appropriate for policy discussions. Translate their question and provide a comprehensive answer in Sinhala. The original question was: '{}'".format(user_message)
        elif original_language == 'romanized_tamil':
            language_instruction = "\n\nIMPORTANT: The user asked in Romanized Tamil (Tamil words written in English letters). Please respond in proper Tamil script using formal language appropriate for policy discussions. Use Sri Lankan Tamil terminology and context. Translate their question and provide a comprehensive answer in Tamil. The original question was: '{}'".format(user_message)
        
        prompt = f"""You are a helpful assistant for budget proposals in Sri Lanka. You can communicate in English, Sinhala, Tamil (Sri Lankan Tamil), and understand Singlish and Romanized Tamil.

FORMATTING RULES:
- DO NOT use asterisks (*) for formatting or emphasis
- DO NOT use markdown formatting like **bold** or *italic*
- Use plain text without any special formatting characters
- Keep responses clean and readable without formatting symbols

IMPORTANT: This website contains various budget proposals for Sri Lanka including:
- Maternity leave benefits proposals
- Cigarette tax reform proposals  
- EPF (Employee Provident Fund) changes
- Electricity tariff reforms
- Tax policy changes
- Economic growth initiatives
- Social protection measures

Based on the following information from the budget proposals database:

{search_context}

{conversation_context}

Current user question: {processed_message}
Original user input: {user_message}
{language_instruction}

Guidelines:
- For general questions like "monada meh" (what is this), "help", or vague inquiries, provide a helpful overview of available budget proposals
- Never say "I couldn't process your request" - always provide useful information about budget proposals  
- Be professional but approachable in any language
- Include specific details from the retrieved information when available
- Cite the source documents when mentioning specific proposals
- If the search doesn't return relevant results, provide an overview of available proposals with examples
- For vague questions, proactively explain what's available and guide users to specific topics (EPF, electricity, maternity leave, cigarette taxes, etc.)
- Keep responses clear and informative
- Reference previous conversation context when relevant
- Maintain conversation continuity
- Be culturally sensitive when discussing Sri Lankan policies
- When responding in Sinhala, use appropriate formal language for policy discussions
- When responding in Tamil, use Sri Lankan Tamil dialect and formal language appropriate for policy discussions
- Always be helpful - turn any question into an opportunity to inform about budget proposals

Please provide a helpful response:"""

        # Generate response using the LLM directly
        response = llm.invoke(prompt)
        response_text = response.content.strip()
        
        # No need to translate response - Gemini handles language matching automatically
        
        # Extract sources from response
        sources = extract_sources_from_response(response_text)
        
        # Add messages to memory (store original user message for context)
        memory.chat_memory.add_user_message(user_message)
        memory.chat_memory.add_ai_message(response_text)
        
        # Get updated conversation history for context
        chat_history = memory.chat_memory.messages
        
        return {
            "response": response_text,
            "confidence": "high",
            "session_id": session_id,
            "conversation_length": len(chat_history),
            "memory_used": True,
            "rag_used": True,
            "sources": sources,
            "language_detected": original_language,
            "translation_used": needs_translation,
            "transliteration_used": transliteration_used,
            "ai_detection_used": ai_detection_used,
            "detection_confidence": confidence
        }
        
    except Exception as e:
        logger.error(f"Error generating response with RAG: {e}")
        # Provide error message in appropriate language
        error_message = "I'm sorry, I'm having trouble processing your request right now. Please try again later."
        if original_language == 'si':
            try:
                error_message = translate_text(error_message, 'si')
            except:
                pass  # Keep English if translation fails
        
        return {
            "response": error_message,
            "confidence": "error",
            "session_id": session_id,
            "memory_used": False,
            "rag_used": False,
            "sources": [],
            "language_detected": original_language if 'original_language' in locals() else 'en',
            "translation_used": False,
            "transliteration_used": False,
            "ai_detection_used": False,
            "detection_confidence": 0.0
        }

def clear_session_memory(session_id: str) -> bool:
    """Clear memory for a specific session"""
    try:
        if session_id in conversation_memories:
            del conversation_memories[session_id]
            logger.info(f"Cleared memory for session: {session_id}")
            return True
        return False
    except Exception as e:
        logger.error(f"Error clearing memory: {e}")
        return False

@app.route('/api/chat', methods=['POST'])
def chat():
    """Enhanced chat endpoint with memory"""
    try:
        data = request.get_json()
        user_message = data.get('message', '').strip()
        session_id = data.get('session_id', 'default')
        
        if not user_message:
            return jsonify({
                "error": "Message is required"
            }), 400
        
        # Generate response with memory
        result = generate_response_with_rag(user_message, session_id)
        
        return jsonify({
            "response": result["response"],
            "confidence": result["confidence"],
            "session_id": session_id,
            "conversation_length": result.get("conversation_length", 0),
            "memory_used": result.get("memory_used", False),
            "rag_used": result.get("rag_used", False),
            "sources": result.get("sources", []),
            "user_message": user_message,
            "language_detected": result.get("language_detected", "en"),
            "translation_used": result.get("translation_used", False),
            "transliteration_used": result.get("transliteration_used", False),
            "ai_detection_used": result.get("ai_detection_used", False),
            "detection_confidence": result.get("detection_confidence", 0.0)
        })
    
    except Exception as e:
        logger.error(f"Chat API error: {e}")
        return jsonify({"error": str(e)}), 500

@app.route('/api/chat/clear', methods=['POST'])
def clear_chat():
    """Clear chat memory for a session"""
    try:
        data = request.get_json()
        session_id = data.get('session_id', 'default')
        
        success = clear_session_memory(session_id)
        
        return jsonify({
            "success": success,
            "session_id": session_id,
            "message": "Chat memory cleared successfully" if success else "Session not found"
        })
    
    except Exception as e:
        logger.error(f"Clear chat error: {e}")
        return jsonify({"error": str(e)}), 500

@app.route('/api/chat/sessions', methods=['GET'])
def list_sessions():
    """List all active chat sessions"""
    try:
        sessions = []
        for session_id, memory in conversation_memories.items():
            messages = memory.chat_memory.messages
            sessions.append({
                "session_id": session_id,
                "message_count": len(messages),
                "last_activity": datetime.now().isoformat()  # Simplified for now
            })
        
        return jsonify({
            "sessions": sessions,
            "total_sessions": len(sessions)
        })
    
    except Exception as e:
        logger.error(f"List sessions error: {e}")
        return jsonify({"error": str(e)}), 500

@app.route('/api/chat/history/<session_id>', methods=['GET'])
def get_chat_history(session_id: str):
    """Get chat history for a specific session"""
    try:
        if session_id not in conversation_memories:
            return jsonify({
                "session_id": session_id,
                "history": [],
                "message_count": 0
            })
        
        memory = conversation_memories[session_id]
        messages = memory.chat_memory.messages
        
        history = []
        for msg in messages:
            if isinstance(msg, HumanMessage):
                history.append({
                    "type": "human",
                    "content": msg.content,
                    "timestamp": datetime.now().isoformat()
                })
            elif isinstance(msg, AIMessage):
                history.append({
                    "type": "ai",
                    "content": msg.content,
                    "timestamp": datetime.now().isoformat()
                })
        
        return jsonify({
            "session_id": session_id,
            "history": history,
            "message_count": len(history)
        })
    
    except Exception as e:
        logger.error(f"Get chat history error: {e}")
        return jsonify({"error": str(e)}), 500

@app.route('/api/chat/health', methods=['GET'])
def chat_health():
    """Health check for the enhanced chatbot"""
    try:
        # Test LangChain connection and vector database
        test_agent = create_agent("health_check")
        test_response = test_agent.invoke({"input": "Hello"})
        
        # Test vector database connection
        pc_index = get_pinecone_index()
        vector_db_status = "connected" if pc_index else "disconnected"
        
        return jsonify({
            "status": "healthy",
            "message": "Enhanced budget proposals chatbot with RAG is running",
            "langchain_status": "connected" if test_response else "disconnected",
            "vector_db_status": vector_db_status,
            "rag_enabled": True,
            "active_sessions": len(conversation_memories),
            "memory_enabled": True
        })
    except Exception as e:
        return jsonify({
            "status": "unhealthy",
            "message": f"Error: {str(e)}"
        }), 500

@app.route('/api/chat/debug/<session_id>', methods=['GET'])
def debug_session(session_id: str):
    """Debug endpoint to check session memory"""
    try:
        memory_exists = session_id in conversation_memories
        memory_info = {
            "session_id": session_id,
            "memory_exists": memory_exists,
            "total_sessions": len(conversation_memories),
            "session_keys": list(conversation_memories.keys())
        }
        
        if memory_exists:
            memory = conversation_memories[session_id]
            messages = memory.chat_memory.messages
            memory_info.update({
                "message_count": len(messages),
                "messages": [
                    {
                        "type": getattr(msg, 'type', 'unknown'),
                        "content": getattr(msg, 'content', '')[:100] + "..." if len(getattr(msg, 'content', '')) > 100 else getattr(msg, 'content', '')
                    }
                    for msg in messages
                ]
            })
        
        return jsonify(memory_info)
    
    except Exception as e:
        logger.error(f"Debug session error: {e}")
        return jsonify({"error": str(e)}), 500

@app.route('/api/chat/suggestions', methods=['GET'])
def get_chat_suggestions():
    """Get suggested questions for the chatbot with multilingual support"""
    suggestions = [
        "What are the maternity leave benefits proposed? 🤱",
        "How do the cigarette tax proposals work? 💰",
        "What changes are proposed for electricity tariffs? ⚡",
        "Tell me about the EPF audit proposals 📊",
        "What tax reforms are being suggested? 🏛️",
        "How will these proposals affect the economy? 📈",
        "What is the cost of implementing these proposals? 💵",
        "Can you compare the costs of different proposals? ⚖️",
        "What are the main benefits of these proposals? ✨",
        "Budget proposals gana kiyanna 📋",
        "EPF eka gana mokadda thiyenne? 💰",
        "Electricity bill eka wenas wenawada? ⚡",
        "Maternity leave benefits kiyannako 🤱",
        "මේ budget proposals වල cost එක කීයද? 💵",
        "රජයේ ආර්థික ප්‍රතිපත්ති ගැන කියන්න 🏛️"
    ]
    
    return jsonify({
        "suggestions": suggestions,
        "supported_languages": ["English", "Sinhala", "Singlish"]
    })

@app.route('/api/chat/available-pdfs', methods=['GET'])
def get_available_pdfs_endpoint():
    """Get list of available PDF files for debugging"""
    try:
        available_pdfs = get_available_pdfs()
        return jsonify({
            "available_pdfs": available_pdfs,
            "count": len(available_pdfs),
            "pdf_directory": "assets/pdfs"
        })
    except Exception as e:
        logger.error(f"Error getting available PDFs: {e}")
        return jsonify({"error": str(e)}), 500

@app.route('/api/chat/detect-language', methods=['POST'])
def detect_language():
    """Test language detection functionality"""
    try:
        data = request.get_json()
        text = data.get('text', '').strip()
        
        if not text:
            return jsonify({
                "error": "Text is required"
            }), 400
        
        processed_message, original_language, needs_translation, transliteration_used, ai_detection_used, confidence = simple_process_input(text)
        
        return jsonify({
            "original_text": text,
            "processed_text": processed_message,
            "language_detected": original_language,
            "translation_needed": needs_translation,
            "transliteration_used": transliteration_used,
            "ai_detection_used": ai_detection_used,
            "detection_confidence": confidence,
            "contains_sinhala": detect_sinhala_content(text),
            "is_singlish": detect_singlish(text)
        })
    
    except Exception as e:
        logger.error(f"Language detection error: {e}")
        return jsonify({"error": str(e)}), 500

@app.route('/', methods=['GET'])
def home():
    """Home endpoint with API documentation"""
    return jsonify({
        "message": "Multilingual Budget Proposals Chatbot API with Swabhasha Pipeline",
        "version": "2.1.0",
        "supported_languages": ["English", "Sinhala", "Tamil (Sri Lankan)", "Romanized Sinhala (Singlish)", "Romanized Tamil"],
        "features": ["RAG", "Memory", "Swabhasha Transliteration", "Google Translation", "FAISS Vector Store"],
        "pipeline": "Romanized Sinhala → Swabhasha → Sinhala Script → Google Translate → English → LLM → Response",
        "endpoints": {
            "POST /api/chat": "Chat with memory, RAG, and multilingual support",
            "POST /api/chat/clear": "Clear chat memory",
            "GET /api/chat/sessions": "List active sessions",
            "GET /api/chat/history/<session_id>": "Get chat history",
            "GET /api/chat/health": "Health check",
            "GET /api/chat/suggestions": "Get suggested questions (multilingual)",
            "GET /api/chat/available-pdfs": "Get available PDF files",
            "POST /api/chat/detect-language": "Test language detection"
        },
        "status": "running"
    })

if __name__ == '__main__':
    app.run(debug=False, host='0.0.0.0', port=7860)