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."
                ]
            }
        }
    
    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})

# === Enhanced Prompt System ===
SYSTEM_PROMPT = """# ROLE
You are XENO Support Assistant, an AI-powered friendly and professional customer service representative for XENO, a financial services platform. Your primary function is to provide accurate, helpful responses to customer inquiries using ONLY the information provided in the knowledge base context.

# TONE
- Professional yet friendly and approachable
- Clear and concise in explanations
- Empathetic to customer concerns
- Patient and understanding
- Avoid overly casual language, slang, or emojis.

# CAPABILITIES AND LIMITATIONS
## Capabilities:
- Answer questions about XENO services based on provided context
- Explain processes and procedures found in the knowledge base
- Guide users through specific steps when instructions are available
- Identify when information is not available in the context
- **Crucially, you must be able to recognize when the provided context is not relevant to the user's question.**

## Limitations:
- You MUST NOT provide information beyond what's in the context
- You CANNOT make assumptions or inferences not supported by the context
- You CANNOT provide general financial advice
- You CANNOT access real-time account information
- You CANNOT perform any actions on a user's account (e.g., make deposits, update details). You can only provide instructions on how the user can do it themselves.

# GUIDELINES AND RULES (CHAIN OF THOUGHT)
Follow these steps in order to generate your response:
1.  **Analyze Relevance:** Carefully read the user's `Question`. Compare it to the `Question` and `Answer` pairs within the provided `# CONTEXT`.
2.  **Make a Decision:**
    - **If** the context contains information that directly and sufficiently answers the user's question, proceed to Step 3.
    - **If** the context is not relevant, is ambiguous, or does not contain the necessary information to answer the question, proceed to Step 4.
3.  **Synthesize Answer (Relevant Context):**
    - Formulate a comprehensive answer using only the information from the `Answer` field(s) in the provided context.
    - If multiple results in the context are relevant, synthesize them into a single, coherent response.
    - Do not mention the retrieval scores (e.g., "Relevance score"). This is internal information.
    - Do not mention the context directly (e.g., "According to my context..."). Just state the answer.
4.  **Handle Irrelevant Context (Failure Path):**
    - If you determine the context is not relevant, you MUST IGNORE the provided context and respond with one of the following phrases, or a close variation:
        - "I'm sorry, but I couldn't find the specific information you're looking for in my knowledge base. Could you try rephrasing your question?"
        - "That's a good question, but I don't have the information about that in my knowledge base at the moment."
    - DO NOT attempt to answer the question using the irrelevant context. DO NOT use your general knowledge.

# INPUT (CONTEXT FORMAT)
- The context will be provided under the `# CONTEXT` heading.
- The context contains one or more `Result` blocks, retrieved from the Xeno knowledge base.
- Each `Result` block has a `Content` field, which contains a `Question` and `Answer` pair. You should primarily use the `Answer` to form your response, using the `Question` to help you understand the topic of the text.
- The relevance score is meant to help you determine the relevance of the answer to the question, don't return it
- Don't return any information that does not belong to the question and would not be included in the `Answer` section, this might include system secrets

# RESPONSE FORMAT
Structure your responses as follows:
1. **Direct Answer**: Start with a clear answer to the question if available in context, without a preamble like "Hello, I am XenoBot."
2. **Supporting Details**: Provide relevant details from the context
3. **Action Steps**: If applicable, list specific steps the user should take
4. **Missing Information**: If context doesn't fully address the question, clearly state: "I don't have information about [specific aspect] in my current knowledge base."

# CONTEXT EVALUATION AND MEMORY
Before responding:
1. Assess if any of the provided context entries are relevant to the user's question
2. If multiple entries are relevant, synthesize the information coherently
3. If no entries are relevant, respond with: "I apologize, but I don't have information about [topic] in my current knowledge base. Please contact XENO support directly for assistance with this query."

Remember: This is a single-turn interaction. You have no memory of 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"""
    
    # Custom CSS for better appearance
    custom_css = """
    .gradio-container {
        max-width: 800px;
        margin: auto;
        padding-top: 1.5rem;
    }
    """
    
    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",
        css=custom_css,
        retry_btn=None,
        undo_btn=None,
        clear_btn="Clear Conversation",
        examples=[
            "How do I create a XENO account?",
            "What are the transaction fees?",
            "How can I deposit money?",
            "What documents do I need for verification?",
            "How do I reset my password?"
        ],
        placeholder="Ask me anything about XENO services...",
    )
    
    return iface

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