langufse fix

agent_gradio_chat.py → app.py

@@ -8,7 +8,7 @@ import xml.etree.ElementTree as ET
 from typing import Any, Dict, List, Optional
 from openai import OpenAI
 from dotenv import load_dotenv
-from langfuse import Langfuse
+from langfuse import observe, get_client
 
 load_dotenv()
 
@@ -30,11 +30,6 @@ BASE_URL = "https://router.huggingface.co/v1"
 
 client = OpenAI(base_url=BASE_URL, api_key=HF_TOKEN)
 
-langfuse = Langfuse(
-    public_key=os.getenv("LANGFUSE_PUBLIC_KEY"),
-    secret_key=os.getenv("LANGFUSE_SECRET_KEY")
-)
-
 # ---------- Tools ----------
 def fetch_google_news_rss(num: int = 10) -> List[Dict[str, Any]]:
     """Fetch general news from Google News RSS feed."""
@@ -317,52 +312,39 @@ def run_agent(user_prompt: str, site_limit: Optional[str] = None, model: str = D
     return "I could not complete the task within the step limit. Try refining your query."
 
 # ---------- Gradio Interface ----------
+@observe()
 def chat_with_agent(message, history, model):
     """Handle chat messages and return agent responses."""
     if not message.strip():
         return history
     
-    # Create a trace for this interaction
-    trace = langfuse.trace(
-        name="chat_interaction",
+    lf = get_client()
+    lf.update_current_trace(
         input={"user_message": message, "model": model, "history_length": len(history)}
     )
-    
+
     try:
         response = run_agent(message, None, model)
-        
-        # Update trace with success info and output
-        trace.update(
+
+        lf.update_current_trace(
             output={"agent_response": response},
             metadata={
                 "model": model,
                 "message_length": len(message),
                 "response_length": len(response),
-                "success": True
-            }
+                "success": True,
+            },
         )
-        
-        # Flush the trace to send it to Langfuse
-        langfuse.flush()
-        
+
         history.append({"role": "user", "content": message})
         history.append({"role": "assistant", "content": response})
         return history
         
     except Exception as e:
-        # Update trace with error info
-        trace.update(
+        lf.update_current_trace(
             output={"error": str(e)},
-            level="ERROR",
-            metadata={
-                "error": str(e),
-                "success": False
-            }
+            metadata={"success": False, "error": str(e)},
         )
-        
-        # Flush the trace to send it to Langfuse
-        langfuse.flush()
-        
         error_msg = f"Sorry, I encountered an error: {str(e)}"
         history.append({"role": "user", "content": message})
         history.append({"role": "assistant", "content": error_msg})