Update app.py

app.py

@@ -8,10 +8,13 @@ import google.generativeai as genai
 import chromadb
 from langchain_chroma import Chroma
 
-# Load Gemini API key from environment variable
+# === 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"
 
-# Load and clean knowledge base
+# === Load and Clean Knowledge Base ===
 df_kb = pd.read_json("XENO_Uganda_KnowledgeBase_Advisory.json")
 df_kb.dropna(subset=['Content'], inplace=True)
 
@@ -33,10 +36,8 @@ def prepare_documents(data):
 xeno_data_list = df_kb.to_dict('records')
 documents, metadatas, ids = prepare_documents(xeno_data_list)
 
-# ChromaDB setup
+# === Setup ChromaDB ===
 client = chromadb.PersistentClient(path="./xeno_db")
-collection_name = "xeno_collection"
-
 try:
     collection = client.get_collection(name=collection_name)
 except:
@@ -46,27 +47,44 @@ except:
 vector_store = Chroma(client=client, collection_name=collection_name)
 retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 4})
 
-embedding_model = "models/embedding-001"
-llm_model_name = "models/gemma-3-4b-it"
-
-SYSTEM_PROMPT = """
-You are XENO Support Assistant, an AI-powered helpful and professional customer service representative.
+# === Prompt System ===
+SYSTEM_PROMPT = """You are 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.
-
-"""
-
+Only use context that is clearly relevant to the user's question.
+For greetings like “hi” or “hello”, respond politely without using the context."""
+
+# === 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)
-    user_prompt = f"""{SYSTEM_PROMPT}
-    ### CONTEXT ###
-    {context}
-    ### QUESTION ###
-    {question}"""
-    response = model.generate_content(user_prompt)
+    prompt = f"""{SYSTEM_PROMPT}
+
+### CONTEXT ###
+{context}
+
+### QUESTION ###
+{question}"""
+    response = model.generate_content(prompt)
     return response.text.strip()
 
+# === Main Interface Logic ===
 def get_context_and_answer(message, history):
+    if message.lower().strip() in {"hi", "hello", "hey"}:
+        return "Hello! How can I assist you with XENO services today?"
+
     queried_results = retriever.invoke(message)
     query_embedding = genai.embed_content(model=embedding_model,
                                           content=message,
@@ -78,13 +96,15 @@ def get_context_and_answer(message, history):
                                             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)
-    context = ""
-    top_docs = sorted(zip(queried_results, cosine_scores), key=lambda x: x[1], reverse=True)[:2]
-    for i, (result, score) in enumerate(top_docs, 1):
-        context += f"Knowledge Entry {i}:\nQ: {result.metadata['question']}\nA: {result.metadata['content']}\n{'-'*40}\n"
-    return generate_xeno_response(context, message)
 
+    # 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."
+
+    context = process_context(queried_results, cosine_scores)
+    return generate_xeno_response(context, message)
 
+# === Gradio UI ===
 iface = gr.ChatInterface(
     fn=get_context_and_answer,
     title="ASKXENO",
@@ -93,5 +113,4 @@ iface = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    iface.launch(share=True)
-
+    iface.launch()