Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import os
 import gradio as gr
 import requests
@@ -23,23 +22,22 @@ try:
 except Exception:
     raise
 
-# Optional LangChain Google generative integration (Gemini)
+# Gemini imports (evaluation only)
 try:
     import google.generativeai as genai
-    from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
+    from langchain_google_genai import ChatGoogleGenerativeAI
 except Exception:
     ChatGoogleGenerativeAI = None
     genai = None
 
-# Load environment variables
+# Load env vars
 load_dotenv()
-
 PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY", "")
 MLFLOW_TRACKING_URI = os.environ.get("MLFLOW_TRACKING_URI", "http://localhost:5000")
 GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY", "")
-LITSERVE_ENDPOINT = os.environ.get("LITSERVE_ENDPOINT", "http://localhost:8000/predict")
+LITSERVE_ENDPOINT = os.environ.get("LITSERVE_ENDPOINT", "")
 
-# DagsHub & MLflow Setup (guarded)
+# DagsHub + MLflow setup
 try:
     dagshub.init(
         repo_owner='prathamesh.khade20',
@@ -52,25 +50,17 @@ except Exception:
 mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)
 mlflow.set_experiment("Maintenance-RAG-Chatbot")
 
-# ----------- App configuration logging -----------
-with mlflow.start_run(run_name=f"App-Config-{datetime.now().strftime('%Y%m%d-%H%M%S')}") as setup_run:
-    mlflow.log_params({
-        "pinecone_index": "rag-granite-index",
-        "embedding_model": "all-MiniLM-L6-v2",
-        "namespace": "rag-ns",
-        "top_k": 3,
-        "llm_endpoint": LITSERVE_ENDPOINT
-    })
-    mlflow.log_text("""
+# ----------- Prompt template -----------
+prompt_template = """
 You are a smart assistant. Based on the provided context, answer the question in 1–2 lines only.
 If the context has more details, summarize it concisely.
 Context:
 {context}
 Question: {question}
 Answer:
-""", "artifacts/prompt_template.txt")
+"""
 
-# ----------- 1. Custom LLM for LitServe endpoint -----------
+# ----------- 1. Custom LLM for LitServe (Lightning AI generator) -----------
 class LitServeLLM(LLM):
     endpoint_url: str
 
@@ -117,53 +107,58 @@ embeddings_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
 
 # ----------- 4. Context Retrieval -----------
 def get_retrieved_context(query: str, top_k=3):
+    start_time = time.time()
     query_embedding = embeddings_model.embed_query(query)
+    mlflow.log_metric("embedding_latency", time.time() - start_time)
+
     if index is None:
         return ""
+
+    start_time = time.time()
     results = index.query(
         namespace="rag-ns",
         vector=query_embedding,
         top_k=top_k,
         include_metadata=True
     )
+    mlflow.log_metric("pinecone_latency", time.time() - start_time)
+    mlflow.log_metric("retrieved_chunks", len(results['matches']))
+
     context_parts = [match['metadata']['text'] for match in results['matches']]
     return "\n".join(context_parts)
 
 # ----------- 5. LLM Chain Setup -----------
 model = LitServeLLM(endpoint_url=LITSERVE_ENDPOINT)
-
-prompt = PromptTemplate(
-    input_variables=["context", "question"],
-    template="""
-You are a smart assistant. Based on the provided context, answer the question in 1–2 lines only.
-If the context has more details, summarize it concisely.
-Context:
-{context}
-Question: {question}
-Answer:
-"""
-)
-
+prompt = PromptTemplate(input_variables=["context", "question"], template=prompt_template)
 llm_chain = LLMChain(llm=model, prompt=prompt)
 
-# ----------- 6. RAG Pipeline -----------
+# ----------- 6. RAG Pipeline (Lightning AI) -----------
 def rag_pipeline(question):
     try:
-        retrieved_context = get_retrieved_context(question)
-        mlflow.log_text(retrieved_context, "artifacts/retrieved_context.txt")
-        response_obj = llm_chain.invoke({"context": retrieved_context, "question": question})
-        response = response_obj.get("text") if isinstance(response_obj, dict) else getattr(response_obj, "text", str(response_obj))
-        response = response.strip()
-        if "Answer:" in response:
-            response = response.split("Answer:", 1)[-1].strip()
-        mlflow.log_text(response, "artifacts/response.txt")
-        return response
+        with mlflow.start_run(run_name=f"Query-{datetime.now().strftime('%H%M%S')}", nested=True):
+            mlflow.log_param("user_question", question)
+            retrieved_context = get_retrieved_context(question)
+            mlflow.log_text(retrieved_context, "artifacts/retrieved_context.txt")
+
+            start_time = time.time()
+            response_obj = llm_chain.invoke({"context": retrieved_context, "question": question})
+            response = response_obj.get("text") if isinstance(response_obj, dict) else getattr(response_obj, "text", str(response_obj))
+            response = response.strip()
+
+            if "Answer:" in response:
+                response = response.split("Answer:", 1)[-1].strip()
+
+            mlflow.log_metric("response_latency", time.time() - start_time)
+            mlflow.log_metric("response_length", len(response))
+            mlflow.log_text(response, "artifacts/response.txt")
+            return response
     except Exception as e:
+        mlflow.log_metric("pipeline_errors", 1)
         error_info = {"error": str(e), "question": question, "timestamp": datetime.now().isoformat()}
         mlflow.log_dict(error_info, "artifacts/pipeline_errors.json")
         return f"Error: {str(e)}"
 
-# ----------- 7. DeepEval Wrappers -----------
+# ----------- 7. DeepEval Wrappers (Gemini evaluator only) -----------
 class GoogleVertexAI(DeepEvalBaseLLM):
     def __init__(self, model):
         self.model = model
@@ -172,31 +167,21 @@ class GoogleVertexAI(DeepEvalBaseLLM):
         return self.model
 
     def generate(self, prompt: str) -> str:
-        chat_model = self.load_model()
-        res = chat_model.invoke(prompt)
+        res = self.model.invoke(prompt)
         if hasattr(res, 'content'):
             return res.content
         if isinstance(res, dict):
             return res.get('content') or res.get('text') or str(res)
         return str(res)
 
-    def get_model_name(self):
-        return "Vertex AI Model"
-
-class LitServeWrapper(DeepEvalBaseLLM):
-    def __init__(self, lit_llm: LitServeLLM):
-        self.lit_llm = lit_llm
-
-    def load_model(self):
-        return self.lit_llm
-
-    def generate(self, prompt: str) -> str:
-        return self.lit_llm._call(prompt)
+    async def a_generate(self, prompt: str) -> str:
+        res = await self.model.ainvoke(prompt)
+        return getattr(res, 'content', str(res))
 
     def get_model_name(self):
-        return "LitServeModel"
+        return "Gemini Evaluator"
 
-# ----------- 8. Custom Metric -----------
+# Length-based utility metric
 class LengthMetric(BaseMetric):
     def __init__(self, min_tokens: int = 1, max_tokens: int = 200):
         self.min_tokens = min_tokens
@@ -214,6 +199,9 @@ class LengthMetric(BaseMetric):
         self.success = (self.min_tokens <= tokens <= self.max_tokens)
         return self.score
 
+    async def a_measure(self, test_case: LLMTestCase):
+        return self.measure(test_case)
+
     def is_successful(self):
         return self.success
 
@@ -221,73 +209,113 @@ class LengthMetric(BaseMetric):
     def name(self):
         return "Length Metric"
 
-# ----------- 9. Evaluation Setup -----------
-def get_deepeval_model(choice: str = 'gemini'):
-    if choice == 'gemini' and ChatGoogleGenerativeAI is not None and GOOGLE_API_KEY:
-        genai.configure(api_key=GOOGLE_API_KEY)
-        chat_model = ChatGoogleGenerativeAI(model="gemini-pro", google_api_key=GOOGLE_API_KEY)
-        return GoogleVertexAI(model=chat_model)
-    else:
-        return LitServeWrapper(lit_llm=model)
-
-def run_deepeval_tests(test_cases: List[LLMTestCase], eval_model_choice: str = 'gemini'):
-    model_wrapper = get_deepeval_model(eval_model_choice)
+# ----------- 8. Run DeepEval Tests (Gemini only) -----------
+def run_deepeval_tests(test_cases: List[LLMTestCase]):
+    if ChatGoogleGenerativeAI is None or not GOOGLE_API_KEY:
+        raise RuntimeError("Gemini API not available — set GOOGLE_API_KEY")
+
+    genai.configure(api_key=GOOGLE_API_KEY)
+    chat_model = ChatGoogleGenerativeAI(model="gemini-pro", google_api_key=GOOGLE_API_KEY)
+    model_wrapper = GoogleVertexAI(model=chat_model)
+
     answer_relevancy_metric = AnswerRelevancyMetric(threshold=0.5, model=model_wrapper)
     hallucination_metric = HallucinationMetric(threshold=0.5, model=model_wrapper)
     length_metric = LengthMetric(min_tokens=3, max_tokens=200)
 
     results = []
-    for i, tc in enumerate(test_cases):
-        if tc.context:
-            mlflow.log_text("\n".join(tc.context), f"artifacts/tc_{i}_context.txt")
-        answer_relevancy_metric.measure(tc)
-        hallucination_metric.measure(tc)
-        length_metric.measure(tc)
-        entry = {
-            "input": tc.input,
-            "actual_output": tc.actual_output,
-            "context": tc.context,
-            "answer_relevancy_score": answer_relevancy_metric.score,
-            "hallucination_score": hallucination_metric.score,
-            "length_score": length_metric.score
-        }
-        results.append(entry)
+    with mlflow.start_run(run_name=f"DeepEval-{datetime.now().strftime('%H%M%S')}", nested=True):
+        for i, tc in enumerate(test_cases):
+            mlflow.log_param(f"tc_{i}_input", tc.input)
+            mlflow.log_param(f"tc_{i}_actual", tc.actual_output)
+            if tc.context:
+                mlflow.log_text("\n".join(tc.context), f"artifacts/tc_{i}_context.txt")
+
+            answer_relevancy_metric.measure(tc)
+            hallucination_metric.measure(tc)
+            length_metric.measure(tc)
+
+            entry = {
+                "input": tc.input,
+                "actual_output": tc.actual_output,
+                "context": tc.context,
+                "answer_relevancy_score": answer_relevancy_metric.score,
+                "hallucination_score": hallucination_metric.score,
+                "length_score": length_metric.score
+            }
+
+            mlflow.log_metric(f"tc_{i}_answer_relevancy", answer_relevancy_metric.score)
+            mlflow.log_metric(f"tc_{i}_hallucination", hallucination_metric.score)
+            mlflow.log_metric(f"tc_{i}_length", length_metric.score)
+
+            results.append(entry)
     return results
 
-# ----------- 10. Gradio App -----------
+# ----------- 9. Gradio UI -----------
 with gr.Blocks() as demo:
-    gr.Markdown("# 🛠️ Maintenance AI Assistant + DeepEval")
+    gr.Markdown("# 🛠️ Maintenance AI Assistant (Lightning AI Generator + Gemini Evaluator)")
 
     with gr.Tabs():
         with gr.TabItem("Chat (RAG)"):
+            usage_counter = gr.State(value=0)
+            session_start = gr.State(value=datetime.now().isoformat())
+
             question_input = gr.Textbox(label="Ask your maintenance question")
             answer_output = gr.Textbox(label="AI Response")
             ask_button = gr.Button("Get Answer")
 
-            def handle_question(question):
-                return rag_pipeline(question)
-
-            ask_button.click(handle_question, inputs=[question_input], outputs=[answer_output])
+            def track_usage(question, count, session_start):
+                count += 1
+                with mlflow.start_run(run_name=f"User-Interaction-{count}", nested=True):
+                    mlflow.log_param("question", question)
+                    mlflow.log_param("session_start", session_start)
+                    response = rag_pipeline(question)
+                    mlflow.log_metric("total_queries", count)
+                    return response, count, session_start
+
+            ask_button.click(
+                track_usage,
+                inputs=[question_input, usage_counter, session_start],
+                outputs=[answer_output, usage_counter, session_start]
+            )
 
         with gr.TabItem("DeepEval — Model Tests"):
+            gr.Markdown("### Evaluate with Gemini (no expected output needed)")
+
             tc_input = gr.Textbox(label="Test Input (prompt)")
-            tc_actual = gr.Textbox(label="Actual Output (leave empty to auto-generate)")
+            tc_actual = gr.Textbox(label="Actual Output (leave empty to auto-generate via Lightning AI)")
             tc_context = gr.Textbox(label="Context (optional)")
-            auto_generate = gr.Checkbox(label="Auto-generate actual output", value=True)
-            model_choice = gr.Radio(["gemini", "litserve"], value="gemini", label="Evaluation backend")
+
+            auto_generate = gr.Checkbox(label="Auto-generate actual output from RAG", value=True)
             run_button = gr.Button("Run DeepEval")
             eval_output = gr.JSON(label="Evaluation Results")
 
-            def run_single_eval(inp, actual, context, autogen, eval_backend):
+            def run_single_eval(inp, actual, context, autogen):
                 if autogen or not actual.strip():
                     actual_output = rag_pipeline(inp)
                 else:
                     actual_output = actual
-                tc = LLMTestCase(input=inp, actual_output=actual_output, expected_output=None, context=[context] if context else None)
-                results = run_deepeval_tests([tc], eval_model_choice=eval_backend)
+
+                tc = LLMTestCase(
+                    input=inp,
+                    actual_output=actual_output,
+                    expected_output=None,
+                    context=[context] if context else None
+                )
+                results = run_deepeval_tests([tc])
                 return results
 
-            run_button.click(run_single_eval, inputs=[tc_input, tc_actual, tc_context, auto_generate, model_choice], outputs=[eval_output])
+            run_button.click(
+                run_single_eval,
+                inputs=[tc_input, tc_actual, tc_context, auto_generate],
+                outputs=[eval_output]
+            )
 
 if __name__ == "__main__":
+    with mlflow.start_run(run_name="Deployment-Info"):
+        mlflow.log_params({
+            "app_version": "1.4.0",
+            "deployment_platform": "Lightning AI / HuggingFace Space",
+            "deployment_time": datetime.now().isoformat(),
+            "code_version": os.getenv("GIT_COMMIT", "dev")
+        })
     demo.launch()