Adds langfuse ec2 and api key to code

code/utils/langgraph_agent.py

@@ -14,9 +14,21 @@ from langgraph.graph.message import add_messages
 from langchain_aws import ChatBedrock
 from langchain_core.messages import HumanMessage, AIMessage, ToolMessage, SystemMessage
 from langchain_core.tools import tool
+import os
+from langfuse import Langfuse
+from langfuse.langchain import CallbackHandler 
+from utils.secrets_manager import SECRETS
 
 CODE=""
 
+langfuse = Langfuse(
+    secret_key=SECRETS["LANGFUSE-SECRET-KEY"],
+    public_key=SECRETS["LANGFUSE-PUBLIC-KEY"],
+    host=os.environ.get("LANGFUSE_HOST", ""),
+)
+
+
+
 # Base paths
 SCRIPTS_DIR = Path(__file__).parent
 JSON_OUTPUT_DIR = SCRIPTS_DIR / "json_output"
@@ -564,13 +576,26 @@ Start your answer with either:
     print("-" * 50)
 
     # Run the agent
-    final_state = agent.invoke(initial_state)
+    langfuse_handler = CallbackHandler()
+    config = {"callbacks": [langfuse_handler]}
+    final_state = agent.invoke(initial_state, config=config)
 
     # Get the final response
     final_message = final_state["messages"][-1]
     response = final_message.content if hasattr(final_message, "content") else str(final_message)
 
+    # Count tokens from all AI messages
+    input_tokens = 0
+    output_tokens = 0
+    for msg in final_state["messages"]:
+        usage = getattr(msg, "usage_metadata", None)
+        if usage:
+            input_tokens += usage.get("input_tokens", 0)
+            output_tokens += usage.get("output_tokens", 0)
+
+    langfuse.flush()
     print(f"Agent: {response}")
-    return response
+    print(f"Tokens - input: {input_tokens}, output: {output_tokens}")
+    return {"response": response, "input_tokens": input_tokens, "output_tokens": output_tokens}