Update app.py

app.py

@@ -4,19 +4,15 @@ import gradio as gr
 from langfuse import Langfuse
 from langfuse.decorators import observe, langfuse_context
 import os 
+import faiss
+import pandas as pd
+from sentence_transformers import SentenceTransformer
+import datetime
 
 # Initialize Langfuse
-#langfuse = Langfuse(
-#    secret_key="sk-lf-229e10c5-6210-4a4b-a432-0f17bc66e56c",
-#    public_key="pk-lf-9f2c32d2-266f-421d-9b87-51377f0a268c", 
-#    host="https://chris4k-langfuse-template-space.hf.space"
-#)
-
-# Get keys for your project from the project settings page
-# https://cloud.langfuse.com
 os.environ["LANGFUSE_PUBLIC_KEY"] = "pk-lf-9f2c32d2-266f-421d-9b87-51377f0a268c"
 os.environ["LANGFUSE_SECRET_KEY"] = "sk-lf-229e10c5-6210-4a4b-a432-0f17bc66e56c"
-os.environ["LANGFUSE_HOST"] = "https://chris4k-langfuse-template-space.hf.space" # 🇪🇺 EU region
+os.environ["LANGFUSE_HOST"] = "https://chris4k-langfuse-template-space.hf.space"  # 🇪🇺 EU region
 
 langfuse = Langfuse()
 
@@ -26,10 +22,6 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name, device_map=None, torch_dtype=torch.float32)
 
 # Load FAISS and Embeddings
-import faiss
-import pandas as pd
-from sentence_transformers import SentenceTransformer
-
 embedder = SentenceTransformer('distiluse-base-multilingual-cased')
 url = 'https://www.bofrost.de/datafeed/DE/products.csv'
 data = pd.read_csv(url, sep='|')
@@ -67,6 +59,22 @@ def construct_prompt(user_input, context, chat_history, max_history_turns=1):
 # Main function to interact with the model
 @observe()
 def chat_with_model(user_input, chat_history=[]):
+    # Start trace for the entire chat process
+    trace = langfuse.trace(
+        name="ai-chat-execution",
+        user_id="user_12345",
+        metadata={"email": "[email protected]"},
+        tags=["chat", "product-query"],
+        release="v1.0.0"
+    )
+
+    # Span for product search
+    retrieval_span = trace.span(
+        name="product-retrieval",
+        metadata={"source": "faiss-index"},
+        input={"query": user_input}
+    )
+
     # Search for products
     search_results = search_products(user_input)
     if search_results:
@@ -76,18 +84,40 @@ def chat_with_model(user_input, chat_history=[]):
     else:
         context = "Das weiß ich nicht."
     
+    # End product search span with results
+    retrieval_span.end(
+        output={"search_results": search_results},
+        status_message=f"Found {len(search_results)} products"
+    )
+
+    # Update trace with search context
     langfuse_context.update_current_observation(
         input={"query": user_input},
         output={"context": context},
         metadata={"search_results_found": len(search_results)}
     )
 
-    # Generate prompt
+    # Generate prompt for Llama model
     prompt = construct_prompt(user_input, context, chat_history)
     input_ids = tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=4096)
+    
+    # Span for AI generation
+    generation_span = trace.span(
+        name="ai-response-generation",
+        metadata={"model": "Llama-3.2-3B-Instruct"},
+        input={"prompt": prompt}
+    )
+    
     outputs = model.generate(input_ids, max_new_tokens=1200, do_sample=True, top_k=50, temperature=0.7)
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # End model generation span
+    generation_span.end(
+        output={"response": response},
+        status_message="AI response generated"
+    )
 
+    # Update Langfuse context with usage details
     langfuse_context.update_current_observation(
         usage_details={
             "input_tokens": len(input_ids[0]),
@@ -95,7 +125,16 @@ def chat_with_model(user_input, chat_history=[]):
         }
     )
 
+    # Append the response to the chat history
     chat_history.append((user_input, response))
+
+    # Update trace final output
+    trace.update(
+        metadata={"final_status": "completed"},
+        output={"summary": response}
+    )
+
+    # Return the response
     return response, chat_history
 
 # Gradio interface
@@ -104,7 +143,7 @@ def gradio_interface(user_input, history):
     return response, updated_history
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🦙 Llama Instruct Chat with LangFuse Integration")
+    gr.Markdown("# 🦙 Llama Instruct Chat with LangFuse & Faiss Integration")
     user_input = gr.Textbox(label="Your Message", lines=2)
     submit_btn = gr.Button("Send")
     chat_history = gr.State([])