Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import os
 import gradio as gr
 import requests
@@ -37,7 +36,7 @@ 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", "https://8001-01k2h9d9mervcmgfn66ybkpwvq.cloudspaces.litng.ai/predict")
+LITSERVE_ENDPOINT = os.environ.get("LITSERVE_ENDPOINT", "http://localhost:8000/predict")
 
 # DagsHub & MLflow Setup (guarded)
 try:
@@ -70,31 +69,29 @@ Question: {question}
 Answer:
 """, "artifacts/prompt_template.txt")
 
-# ----------- 1. Custom LLM for LitServe endpoint (Lightning AI) -----------
+# ----------- 1. Custom LLM for LitServe endpoint -----------
 class LitServeLLM(LLM):
     endpoint_url: str
 
-    @mlflow.trace
     def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
         payload = {"prompt": prompt}
-        with mlflow.start_span("lit_serve_request"):
-            start_time = time.time()
-            response = requests.post(self.endpoint_url, json=payload)
-            latency = time.time() - start_time
-            mlflow.log_metric("lit_serve_latency", latency)
-            if response.status_code == 200:
-                data = response.json()
-                mlflow.log_metric("response_tokens", len(data.get("response", "").split()))
-                return data.get("response", "").strip()
-            else:
-                mlflow.log_metric("request_errors", 1)
-                error_info = {
-                    "status_code": response.status_code,
-                    "error": response.text,
-                    "timestamp": datetime.now().isoformat()
-                }
-                mlflow.log_dict(error_info, "artifacts/error_log.json")
-                raise ValueError(f"Request failed: {response.status_code}")
+        start_time = time.time()
+        response = requests.post(self.endpoint_url, json=payload)
+        latency = time.time() - start_time
+        mlflow.log_metric("lit_serve_latency", latency)
+        if response.status_code == 200:
+            data = response.json()
+            mlflow.log_metric("response_tokens", len(data.get("response", "").split()))
+            return data.get("response", "").strip()
+        else:
+            mlflow.log_metric("request_errors", 1)
+            error_info = {
+                "status_code": response.status_code,
+                "error": response.text,
+                "timestamp": datetime.now().isoformat()
+            }
+            mlflow.log_dict(error_info, "artifacts/error_log.json")
+            raise ValueError(f"Request failed: {response.status_code}")
 
     @property
     def _identifying_params(self) -> Mapping[str, Any]:
@@ -105,9 +102,7 @@ class LitServeLLM(LLM):
         return "litserve_llm"
 
 # ----------- 2. Pinecone Connection -----------
-@mlflow.trace
 def init_pinecone():
-    PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY")
     pc = Pinecone(api_key=PINECONE_API_KEY)
     return pc.Index("rag-granite-index")
 
@@ -119,33 +114,21 @@ except Exception:
 # ----------- 3. Embedding Model -----------
 embeddings_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
 
-# ----------- 4. Context Retrieval with Tracing -----------
-@mlflow.trace
+# ----------- 4. Context Retrieval -----------
 def get_retrieved_context(query: str, top_k=3):
-    with mlflow.start_span("embedding_generation"):
-        start_time = time.time()
-        query_embedding = embeddings_model.embed_query(query)
-        mlflow.log_metric("embedding_latency", time.time() - start_time)
-
+    query_embedding = embeddings_model.embed_query(query)
     if index is None:
         return ""
-
-    with mlflow.start_span("pinecone_query"):
-        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']))
-
+    results = index.query(
+        namespace="rag-ns",
+        vector=query_embedding,
+        top_k=top_k,
+        include_metadata=True
+    )
     context_parts = [match['metadata']['text'] for match in results['matches']]
     return "\n".join(context_parts)
 
-
-# ----------- 5. LLM Chain Setup (Lightning AI generator) -----------
+# ----------- 5. LLM Chain Setup -----------
 model = LitServeLLM(endpoint_url=LITSERVE_ENDPOINT)
 
 prompt = PromptTemplate(
@@ -162,42 +145,24 @@ Answer:
 
 llm_chain = LLMChain(llm=model, prompt=prompt)
 
-# ----------- 6. RAG Pipeline with Full Tracing (uses Lightning AI) -----------
-@mlflow.trace
+# ----------- 6. RAG Pipeline -----------
 def rag_pipeline(question):
     try:
-        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
+        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
     except Exception as e:
-        mlflow.log_metric("pipeline_errors", 1)
-        error_info = {
-            "error": str(e),
-            "question": question,
-            "timestamp": datetime.now().isoformat()
-        }
+        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 Wrapper(s) and Metrics Integration (Gemini evaluation) -----------
+# ----------- 7. DeepEval Wrappers -----------
 class GoogleVertexAI(DeepEvalBaseLLM):
     def __init__(self, model):
         self.model = model
@@ -214,11 +179,6 @@ class GoogleVertexAI(DeepEvalBaseLLM):
             return res.get('content') or res.get('text') or str(res)
         return str(res)
 
-    async def a_generate(self, prompt: str) -> str:
-        chat_model = self.load_model()
-        res = await chat_model.ainvoke(prompt)
-        return getattr(res, 'content', str(res))
-
     def get_model_name(self):
         return "Vertex AI Model"
 
@@ -232,13 +192,10 @@ class LitServeWrapper(DeepEvalBaseLLM):
     def generate(self, prompt: str) -> str:
         return self.lit_llm._call(prompt)
 
-    async def a_generate(self, prompt: str) -> str:
-        return self.generate(prompt)
-
     def get_model_name(self):
         return "LitServeModel"
 
-# Custom metric that DOES NOT require expected_output: Length-based utility metric
+# ----------- 8. Custom Metric -----------
 class LengthMetric(BaseMetric):
     def __init__(self, min_tokens: int = 1, max_tokens: int = 200):
         self.min_tokens = min_tokens
@@ -256,9 +213,6 @@ 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
 
@@ -266,142 +220,73 @@ class LengthMetric(BaseMetric):
     def name(self):
         return "Length Metric"
 
-# Helper to get eval model: GEMINI will be used as evaluator by default
+# ----------- 9. Evaluation Setup -----------
 def get_deepeval_model(choice: str = 'gemini'):
     if choice == 'gemini' and ChatGoogleGenerativeAI is not None and GOOGLE_API_KEY:
-        try:
-            genai.configure(api_key=GOOGLE_API_KEY)
-        except Exception:
-            pass
+        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:
-        # fallback to litserve wrapper if gemini isn't available
         return LitServeWrapper(lit_llm=model)
 
-# Function to run Deepeval tests and log to mlflow (only metrics that don't need expected_output)
-@mlflow.trace
 def run_deepeval_tests(test_cases: List[LLMTestCase], eval_model_choice: str = 'gemini'):
     model_wrapper = get_deepeval_model(eval_model_choice)
-
-    # Only metrics that do not require expected output
     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 = []
-    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")
-
-            # Measure metrics
-            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
-            }
-
-            # Log metrics to mlflow
-            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)
-
+    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)
     return results
 
-# ----------- 8. Gradio UI with Evaluation Tab (Auto-generate actual output from Lightning AI) -----------
+# ----------- 10. Gradio App -----------
 with gr.Blocks() as demo:
-    gr.Markdown("# 🛠️ Maintenance AI Assistant + DeepEval (Lightning AI generator, Gemini evaluator)")
+    gr.Markdown("# 🛠️ Maintenance AI Assistant + DeepEval")
 
     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")
-            feedback = gr.Radio(["Helpful", "Not Helpful"], label="Was this response helpful?")
-
-            def track_usage(question, count, session_start, feedback=None):
-                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)
-                    if feedback:
-                        mlflow.log_param("user_feedback", feedback)
-                        mlflow.log_metric("helpful_responses", 1 if feedback == "Helpful" else 0)
-                    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]
-            )
-
-            feedback.change(
-                track_usage,
-                inputs=[question_input, usage_counter, session_start, feedback],
-                outputs=[answer_output, usage_counter, session_start]
-            )
 
-        with gr.TabItem("DeepEval — Model Tests"):
-            gr.Markdown("### Run DeepEval metrics (no expected output needed). Provide input; optionally auto-generate the model response (Lightning AI). Gemini will evaluate by default.")
+            def handle_question(question):
+                return rag_pipeline(question)
+
+            ask_button.click(handle_question, inputs=[question_input], outputs=[answer_output])
 
+        with gr.TabItem("DeepEval — Model Tests"):
             tc_input = gr.Textbox(label="Test Input (prompt)")
-            tc_actual = gr.Textbox(label="Actual Output (paste model response or leave empty to auto-generate)")
+            tc_actual = gr.Textbox(label="Actual Output (leave empty to auto-generate)")
             tc_context = gr.Textbox(label="Context (optional)")
-
-            auto_generate = gr.Checkbox(label="Auto-generate actual output from RAG (Lightning AI)", value=True)
-            model_choice = gr.Radio(["gemini", "litserve"], value="gemini", label="Evaluation backend (Gemini recommended)")
+            auto_generate = gr.Checkbox(label="Auto-generate actual output", value=True)
+            model_choice = gr.Radio(["gemini", "litserve"], value="gemini", label="Evaluation backend")
             run_button = gr.Button("Run DeepEval")
             eval_output = gr.JSON(label="Evaluation Results")
 
             def run_single_eval(inp, actual, context, autogen, eval_backend):
-                # If autogen is True, generate actual output via RAG pipeline (Lightning AI)
-                if autogen or (actual is None or actual.strip() == ""):
-                    generated = rag_pipeline(inp)
-                    actual_output = generated
+                if autogen or not actual.strip():
+                    actual_output = rag_pipeline(inp)
                 else:
                     actual_output = actual
-
-                # Log that actual was autogenerated
-                with mlflow.start_run(run_name=f"DE-Run-{datetime.now().strftime('%H%M%S')}", nested=True):
-                    mlflow.log_param("input", inp)
-                    mlflow.log_param("autogenerated_actual", autogen)
-                    if context:
-                        mlflow.log_text(context, "artifacts/eval_context.txt")
-
                 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)
                 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, model_choice], outputs=[eval_output])
 
 if __name__ == "__main__":
-    with mlflow.start_run(run_name="Deployment-Info"):
-        mlflow.log_params({
-            "app_version": "1.3.0",
-            "deployment_platform": "Lightning AI / HuggingFace Space",
-            "deployment_time": datetime.now().isoformat(),
-            "code_version": os.getenv("GIT_COMMIT", "dev")
-        })
-
     demo.launch()