Add traces

Author: w-k-s

README.md

@@ -19,19 +19,19 @@
 4. Publish the hr_agent prompts to langfuse prompt registry
 
     ```shell
-    uv run main.py prompts --publish hr_agent
+    uv run langfuse_experiment prompts --publish hr_agent
     ```
 
 6. Before you can runs evals, you need to publish the eval dataset to langfuse dataset registry
 
     ```shell
-    uv run main.py datasets --publish ci
+    uv run langfuse_experiment datasets --publish ci
     ```
 
 5. Chat with the agent
 
     ```
-    uv run main.py chat
+    uv run langfuse_experiment chat
     ```
 
 ![Screenshot](./docs/Screenshot.png)
@@ -42,8 +42,8 @@
 - [x] Prompt gemma, falcon from langchain
 - [x] Setup chat memory
 - [x] Get the system prompt from SaaS langfuse
-- [ ] Setup offline Evals to run on the CI/CD pipeline 
-- [ ] Setup traces
+- [x] Setup offline Evals to run on the CI/CD pipeline 
+- [x] Setup traces
 - [ ] Generate a generic hr handbook 
 - [ ] Use LlamaIndex to do rag over the handbook. 
 - [ ] Add testing for rag (ragas?)

src/langfuse_experiment/graph.py

@@ -1,23 +1,50 @@
+from langfuse import Langfuse, propagate_attributes
 from dataclasses import dataclass
 from langfuse import Langfuse
 from langchain.chat_models import BaseChatModel
 from langgraph.runtime import Runtime
 from langgraph.graph import END, START, MessagesState, StateGraph
 from langgraph.checkpoint.memory import InMemorySaver
 from langchain_ollama import ChatOllama
+from langchain_core.runnables import RunnableConfig
 
 
 @dataclass
 class ContextSchema:
-    prompt_registry: Langfuse
+    langfuse: Langfuse
     llm: BaseChatModel
 
 
-def call_model(state: MessagesState, runtime: Runtime[ContextSchema]):
-    prompt_client = runtime.context.prompt_registry.get_prompt("hr_agent")
-    system_prompt = prompt_client.compile()
-    ai_msg = runtime.context.llm.invoke([system_prompt] + state["messages"])
-    return {"messages": [ai_msg]}
+def call_model(
+    state: MessagesState, runtime: Runtime[ContextSchema], config: RunnableConfig
+):
+    langfuse = runtime.context.langfuse
+    llm = runtime.context.llm
+    thread_id = config["configurable"]["thread_id"]
+
+    with propagate_attributes(session_id=thread_id):
+
+        with langfuse.start_as_current_observation(
+            as_type="span", name="call-model"
+        ) as root_span:
+
+            prompt_client = langfuse.get_prompt("hr_agent")
+            system_prompt = prompt_client.compile()
+            messages = [system_prompt] + state["messages"]
+
+            with root_span.start_as_current_observation(
+                as_type="generation",
+                name="generate-response",
+                model=runtime.context.llm.model,
+                input=messages,
+                prompt=prompt_client,
+                model_parameters={"temperature": llm.temperature},
+            ) as gen:
+                ai_msg = llm.invoke(messages)
+
+                gen.update(output=ai_msg, usage_details=ai_msg.usage_metadata)
+
+                return {"messages": [ai_msg]}
 
 
 def build_model():

src/langfuse_experiment/main.py

@@ -1,5 +1,6 @@
 import sys
 import argparse
+from uuid import uuid4
 from dotenv import load_dotenv
 from langfuse import get_client
 from langchain_core.runnables import RunnableConfig
@@ -13,7 +14,7 @@ def chat(langfuse):
 
     graph = build_graph()
 
-    config: RunnableConfig = {"configurable": {"thread_id": "1"}}
+    config: RunnableConfig = {"configurable": {"thread_id": str(uuid4())}}
 
     while True:
         user_input = input("Human: ")