Update app.py

app.py

@@ -356,7 +356,6 @@ if __name__ == "__main__":
 
 '''
 
-
 import os
 import gradio as gr
 import requests
@@ -503,6 +502,21 @@ def get_retrieved_context(query: str, top_k=3):
     return "
 ".join(context_parts)
 
+    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 "
+".join(context_parts)
+
 # ----------- 5. LLM Chain Setup (Lightning AI generator) -----------
 model = LitServeLLM(endpoint_url=LITSERVE_ENDPOINT)
 
@@ -764,4 +778,3 @@ if __name__ == "__main__":
         })
 
     demo.launch()
-