app.py

import os
import gradio as gr
import pandas as pd
import torch
import numpy as np
from sentence_transformers import util
import google.generativeai as genai
import chromadb
from langchain_chroma import Chroma
import re
from typing import Dict, List, Tuple

# === Configuration ===
genai.configure(api_key=os.environ["GEMINI_API_KEY"])
embedding_model = "models/embedding-001"
llm_model_name = "models/gemma-3-4b-it"
collection_name = "xeno_collection"

# === Intent Classification System ===
class IntentClassifier:
    def __init__(self):
        # Define intent patterns and responses
        self.intent_patterns = {
            'greeting': {
                'patterns': [
                    r'\b(hi|hello|hey|good morning|good afternoon|good evening|greetings)\b',
                    r'^(hi|hello|hey)[\s!.]*$',
                    r'\b(how are you|how do you do)\b'
                ],
                'responses': [
                    "Hello! I'm XENO Assistant. How can I help you with XENO financial services today?",
                    "Hi there! I'm here to assist you with any questions about XENO services. What can I help you with?",
                    "Good day! Welcome to XENO Support. How may I assist you today?"
                ]
            },
            'thanks': {
                'patterns': [
                    r'\b(thank you|thanks|thank u|thx|appreciate|grateful)\b',
                    r'^(thanks|thank you)[\s!.]*$',
                    r'\b(much appreciated|thanks a lot|thank you so much)\b'
                ],
                'responses': [
                    "You're welcome! Is there anything else I can help you with regarding XENO services?",
                    "Happy to help! Feel free to ask if you have any other questions about XENO.",
                    "Glad I could assist you! Let me know if you need help with anything else."
                ]
            },
            'goodbye': {
                'patterns': [
                    r'\b(bye|goodbye|see you|farewell|take care|have a good day)\b',
                    r'^(bye|goodbye)[\s!.]*$',
                    r'\b(talk to you later|see you later|until next time)\b'
                ],
                'responses': [
                    "Goodbye! Thank you for using XENO services. Have a great day!",
                    "Take care! Feel free to return anytime you need help with XENO services.",
                    "Have a wonderful day! Don't hesitate to reach out if you need assistance with XENO."
                ]
            }
        }
    
    def classify_intent(self, message: str) -> Tuple[str, str]:
        """
        Classify the intent of a message and return appropriate response if it's a simple intent.
        Returns: (intent_name, response) - response is empty string if intent requires RAG
        """
        message_lower = message.lower().strip()
        
        # Check for each intent pattern
        for intent_name, intent_data in self.intent_patterns.items():
            for pattern in intent_data['patterns']:
                if re.search(pattern, message_lower, re.IGNORECASE):
                    # Return random response from available responses
                    import random
                    response = random.choice(intent_data['responses'])
                    return intent_name, response
        
        # If no simple intent found, it's a query that needs RAG
        return 'query', ''
    
    def is_simple_intent(self, intent: str) -> bool:
        """Check if intent can be handled without RAG"""
        simple_intents = ['greeting', 'thanks']
        return intent in simple_intents

# Initialize intent classifier
intent_classifier = IntentClassifier()

# === Load and Clean Knowledge Base ===
df_kb = pd.read_json("XENO_Uganda_KnowledgeBase_Advisory.json")
df_kb.dropna(subset=['Content'], inplace=True)

def prepare_documents(data):
    documents, metadatas, ids = [], [], []
    for item in data:
        documents.append(f"Question: {item['Question']}\nAnswer: {item['Content']}")
        metadatas.append({
            "question": item["Question"],
            "content": item["Content"],
            "section": item.get("Section", ""),
            "source": item.get("Source", ""),
            "owner": item.get("Owner", ""),
            "tag": item.get("Tag", "")
        })
        ids.append(item["ID"])
    return documents, metadatas, ids

xeno_data_list = df_kb.to_dict('records')
documents, metadatas, ids = prepare_documents(xeno_data_list)

# === Setup ChromaDB ===
client = chromadb.PersistentClient(path="./xeno_db")
try:
    collection = client.get_collection(name=collection_name)
except:
    collection = client.create_collection(name=collection_name)
    collection.add(documents=documents, metadatas=metadatas, ids=ids)

vector_store = Chroma(client=client, collection_name=collection_name)
retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 4})

# === Prompt System ===
SYSTEM_PROMPT = """You are a friendly XENO Support Assistant, an AI-powered helpful and professional customer service representative.
Use only the information provided in the knowledge base context to answer user queries.
Do not hallucinate. If context doesn't contain relevant info, say so in a calm polite manner by saying I'm sorry, I can't assist with that.
Only use context that is clearly relevant to the user's question.
For greetings like “hi” or “hello”, respond politely without using the context.
remember previous conversations."""

# === Context Processing ===
def process_context(results, cosine_scores, max_results=2):
    sorted_indices = np.argsort(cosine_scores)[::-1][:max_results]
    formatted_context = ""
    for i, idx in enumerate(sorted_indices, 1):
        result = results[idx]
        score = cosine_scores[idx]
        formatted_context += f"Knowledge Entry {i}:\n"
        formatted_context += f"Q: {result.metadata.get('question', 'N/A')}\n"
        formatted_context += f"A: {result.metadata.get('content', 'N/A')}\n"
        formatted_context += "-" * 40 + "\n"
    return formatted_context

# === LLM Generation ===
def generate_xeno_response(context, question):
    model = genai.GenerativeModel(llm_model_name)
    prompt = f"""{SYSTEM_PROMPT}
### CONTEXT ###
{context}
### QUESTION ###
{question}"""
    response = model.generate_content(prompt)
    return response.text.strip()

# === Enhanced Main Interface Logic with Intent Classification ===
def get_context_and_answer(message, history):
    """
    Enhanced pipeline with intent classification
    """
    # Step 1: Intent Classification
    intent, direct_response = intent_classifier.classify_intent(message)
    
    # Step 2: Handle simple intents directly
    if intent_classifier.is_simple_intent(intent) and direct_response:
        return direct_response
    
    # Step 3: For queries that need RAG processing
    if intent == 'query':
        # Check if message is too short or unclear
        if len(message.strip()) < 3:
            return "I'd be happy to help! Could you please provide more details about what you'd like to know about XENO services?"
        
        # Retrieve relevant documents
        try:
            queried_results = retriever.invoke(message)
            query_embedding = genai.embed_content(
                model=embedding_model,
                content=message,
                task_type="retrieval_query"
            )['embedding']
            
            cosine_scores = []
            for doc in queried_results:
                doc_embedding = genai.embed_content(
                    model=embedding_model,
                    content=doc.page_content,
                    task_type="retrieval_document"
                )['embedding']
                cos_sim = util.cos_sim(
                    torch.tensor(query_embedding).float(), 
                    torch.tensor(doc_embedding).float()
                )[0][0].item()
                cosine_scores.append(cos_sim)

            # If none of the results have sufficient similarity, fallback
            if max(cosine_scores) < 0.4:
                return "I'm sorry, I couldn't find the specific information you're looking for in my knowledge base. Could you try rephrasing your question or contact XENO support directly for assistance?"

            context = process_context(queried_results, cosine_scores)
            return generate_xeno_response(context, message)
            
        except Exception as e:
            return "I apologize, but I'm experiencing a technical issue. Please contact XENO support directly for assistance with your query."
    
    # Fallback for any unhandled cases
    return "I'm here to help with XENO financial services. What would you like to know?"

# === Enhanced Gradio UI ===
def create_interface():
    """Create the Gradio interface with custom styling"""
    
    iface = gr.ChatInterface(
        fn=get_context_and_answer,
        title="🏦 ASKXENO - AI Support Assistant",
        description="""**Welcome to XENO AI Support!** 
I can help you with questions about XENO financial services including:
• Account management and setup
• Transaction processes and fees  
• Platform features and troubleshooting
• General service information
*Simply type your question below to get started!*""",
        theme="soft"
    )
    
    return iface

# === Main Execution ===
if __name__ == "__main__":
    iface = create_interface()
    iface.launch()