Update app.py

app.py

@@ -1,3 +1,111 @@
+'''
+import os
+import gradio as gr
+import requests
+from pinecone import Pinecone
+from langchain.prompts import PromptTemplate
+from langchain.chains.llm import LLMChain
+from langchain.llms.base import LLM
+from typing import Optional, List, Mapping, Any
+from langchain.embeddings import HuggingFaceEmbeddings
+
+# ----------- 1. Custom LLM to call your LitServe endpoint -----------
+class LitServeLLM(LLM):
+    endpoint_url: str
+
+    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+        payload = {"prompt": prompt}
+        response = requests.post(self.endpoint_url, json=payload)
+        if response.status_code == 200:
+            data = response.json()
+            return data.get("response", "").strip()
+        else:
+            raise ValueError(f"Request failed: {response.status_code} {response.text}")
+
+    @property
+    def _identifying_params(self) -> Mapping[str, Any]:
+        return {"endpoint_url": self.endpoint_url}
+
+    @property
+    def _llm_type(self) -> str:
+        return "litserve_llm"
+
+
+# ----------- 2. Connect to Pinecone -----------
+PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY")
+pc = Pinecone(api_key=PINECONE_API_KEY)
+index = pc.Index("rag-granite-index")
+
+# ----------- 3. Load embedding model -----------
+embeddings_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+
+# ----------- 4. Function to get top context from Pinecone -----------
+def get_retrieved_context(query: str, top_k=3):
+    query_embedding = embeddings_model.embed_query(query)
+    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. Create LLMChain with your model -----------
+model = LitServeLLM(
+    endpoint_url="https://8001-01k2h9d9mervcmgfn66ybkpwvq.cloudspaces.litng.ai/predict"
+)
+
+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:
+"""
+)
+
+llm_chain = LLMChain(llm=model, prompt=prompt)
+
+# ----------- 6. Main RAG Function -----------
+def rag_pipeline(question):
+    try:
+        retrieved_context = get_retrieved_context(question)
+        response = llm_chain.invoke({
+            "context": retrieved_context,
+            "question": question
+        })["text"].strip()
+
+        # Only keep what's after "Answer:"
+        if "Answer:" in response:
+            response = response.split("Answer:", 1)[-1].strip()
+
+        return response
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+
+# ----------- 7. Gradio UI -----------
+with gr.Blocks() as demo:
+    gr.Markdown("# 🧠 RAG Chatbot (Pinecone + LitServe)")
+    question_input = gr.Textbox(label="Ask your question here")
+    answer_output = gr.Textbox(label="Answer")
+    ask_button = gr.Button("Get Answer")
+    ask_button.click(rag_pipeline, inputs=question_input, outputs=answer_output)
+
+if _name_ == "_main_":
+    demo.launch()
+'''
+
+
+
+
 import os
 import gradio as gr
 import requests
@@ -8,81 +116,80 @@ from langchain.prompts import PromptTemplate
 from langchain.chains.llm import LLMChain
 from langchain.llms.base import LLM
 from typing import Optional, List, Mapping, Any
-import time
+import time  
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from dotenv import load_dotenv
 from datetime import datetime
 
-# DeepEval imports
-try:
-    from deepeval.test_case import LLMTestCase
-    from deepeval.metrics import AnswerRelevancyMetric, HallucinationMetric
-    from deepeval.metrics import BaseMetric
-    from deepeval.models.base_model import DeepEvalBaseLLM
-except Exception:
-    raise
-
-# Gemini imports (evaluation only)
-try:
-    import google.generativeai as genai
-    from langchain_google_genai import ChatGoogleGenerativeAI
-except Exception:
-    ChatGoogleGenerativeAI = None
-    genai = None
-
-# 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", "")
-
-# DagsHub + MLflow setup
-try:
-    dagshub.init(
-        repo_owner='prathamesh.khade20',
-        repo_name='Maintenance_AI_website',
-        mlflow=True
-    )
-except Exception:
-    pass
+# Load environment variables
+pinecone_api_key = os.environ["PINECONE_API_KEY"]
 
-mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)
-mlflow.set_experiment("Maintenance-RAG-Chatbot")
+mlflow_tracking_uri = os.environ["MLFLOW_TRACKING_URI"]
+
+# ----------- DagsHub & MLflow Setup -----------
 
-# ----------- Prompt template -----------
-prompt_template = """
+dagshub.init(
+    repo_owner='prathamesh.khade20',
+    repo_name='Maintenance_AI_website',
+    mlflow=True
+)
+
+mlflow.set_tracking_uri(mlflow_tracking_uri)
+mlflow.set_experiment("Maintenance-RAG-Chatbot")
+mlflow.langchain.autolog()
+
+
+
+# Initialize MLflow run for app configuration
+with mlflow.start_run(run_name=f"App-Config-{datetime.now().strftime('%Y%m%d-%H%M%S')}") as setup_run:
+    # Log environment configuration
+    mlflow.log_params({
+        "pinecone_index": "rag-granite-index",
+        "embedding_model": "all-MiniLM-L6-v2",
+        "namespace": "rag-ns",
+        "top_k": 3,
+        "llm_endpoint": "https://8001-01k2h9d9mervcmgfn66ybkpwvq.cloudspaces.litng.ai/predict"
+    })
+    
+    # Log important files as artifacts
+    
+    mlflow.log_text("""
 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 (Lightning AI generator) -----------
+# ----------- 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}
-        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}")
+        
+        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}")
 
     @property
     def _identifying_params(self) -> Mapping[str, Any]:
@@ -93,229 +200,157 @@ 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")
 
-try:
-    index = init_pinecone()
-except Exception:
-    index = None
+index = init_pinecone()
 
 # ----------- 3. Embedding Model -----------
 embeddings_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
 
-# ----------- 4. Context Retrieval -----------
+# ----------- 4. Context Retrieval with Tracing -----------
+@mlflow.trace
 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']))
-
+    """Retrieve context from Pinecone with performance tracing"""
+    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)
+    
+    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']))
+    
     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=prompt_template)
+model = LitServeLLM(
+    endpoint_url="https://8001-01k2h9d9mervcmgfn66ybkpwvq.cloudspaces.litng.ai/predict"
+)
+
+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:
+"""
+)
+
 llm_chain = LLMChain(llm=model, prompt=prompt)
 
-# ----------- 6. RAG Pipeline (Lightning AI) -----------
+# ----------- 6. RAG Pipeline with Full Tracing -----------
+@mlflow.trace
 def rag_pipeline(question):
+    """End-to-end RAG pipeline with MLflow tracing"""
     try:
+        # Start a new nested run for each query
         with mlflow.start_run(run_name=f"Query-{datetime.now().strftime('%H%M%S')}", nested=True):
             mlflow.log_param("user_question", question)
+            
+            # Retrieve context
             retrieved_context = get_retrieved_context(question)
             mlflow.log_text(retrieved_context, "artifacts/retrieved_context.txt")
-
+            
+            # Generate response
             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()
-
+            response = llm_chain.invoke({
+                "context": retrieved_context,
+                "question": question
+            })["text"].strip()
+            
+            # Clean response
             if "Answer:" in response:
                 response = response.split("Answer:", 1)[-1].strip()
-
+            
+            # Log metrics
             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()}
+        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 (Gemini evaluator only) -----------
-class GoogleVertexAI(DeepEvalBaseLLM):
-    def __init__(self, model):
-        self.model = model
-
-    def load_model(self):
-        return self.model
-
-    def generate(self, prompt: str) -> str:
-        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)
-
-    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 "Gemini Evaluator"
-
-# Length-based utility metric
-class LengthMetric(BaseMetric):
-    def __init__(self, min_tokens: int = 1, max_tokens: int = 200):
-        self.min_tokens = min_tokens
-        self.max_tokens = max_tokens
-        self.score = 0.0
-        self.success = False
-
-    def measure(self, test_case: LLMTestCase):
-        text = (test_case.actual_output or "")
-        tokens = len(text.split())
-        mid = (self.min_tokens + self.max_tokens) / 2
-        dist = abs(tokens - mid)
-        max_dist = max(mid - self.min_tokens, self.max_tokens - mid)
-        self.score = max(0.0, 1.0 - (dist / max_dist))
-        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
-
-    @property
-    def name(self):
-        return "Length Metric"
-
-# ----------- 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 = []
-    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
-
-# ----------- 9. Gradio UI -----------
+# ----------- 7. Gradio UI with Enhanced Tracking -----------
 with gr.Blocks() as demo:
-    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 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 via Lightning AI)")
-            tc_context = gr.Textbox(label="Context (optional)")
-
-            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):
-                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])
-                return results
-
-            run_button.click(
-                run_single_eval,
-                inputs=[tc_input, tc_actual, tc_context, auto_generate],
-                outputs=[eval_output]
-            )
-
-if __name__ == "__main__":
+    gr.Markdown("# 🛠 Maintenance AI Assistant")
+    
+    # Track additional UI metrics
+    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):
+        """Wrapper to track usage metrics with feedback"""
+        count += 1
+        
+        # Start tracking context
+        with mlflow.start_run(run_name=f"User-Interaction-{count}", nested=True):
+            mlflow.log_param("question", question)
+            mlflow.log_param("session_start", session_start)
+            
+            # Get response
+            response = rag_pipeline(question)
+            
+            # Log feedback if provided
+            if feedback:
+                mlflow.log_param("user_feedback", feedback)
+                mlflow.log_metric("helpful_responses", 1 if feedback == "Helpful" else 0)
+            
+            # Update metrics
+            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]
+    )
+
+if _name_ == "_main_":
+    # Log deployment information
     with mlflow.start_run(run_name="Deployment-Info"):
         mlflow.log_params({
-            "app_version": "1.4.0",
-            "deployment_platform": "Lightning AI / HuggingFace Space",
+            "app_version": "1.0.0",
+            "deployment_platform": "Lightning AI",
             "deployment_time": datetime.now().isoformat(),
             "code_version": os.getenv("GIT_COMMIT", "dev")
         })
+    
+    # Start Gradio app
     demo.launch()
+