Skip to content

Commit 5a6ce9f

Browse files
authored
feat: llm tracing link for langsmith (#700)
# Motivation <!-- Why is this change necessary? --> # Content <!-- Please include a summary of the change --> # Testing <!-- How was the change tested? --> # Please check the following before marking your PR as ready for review - [x] I have added tests for my changes - [x] I have updated the documentation or added new documentation as needed --------- Co-authored-by: kopekC <28070492+kopekC@users.noreply.github.com>
1 parent e9e58e3 commit 5a6ce9f

File tree

1 file changed

+107
-3
lines changed

1 file changed

+107
-3
lines changed

src/codegen/agents/code_agent.py

Lines changed: 107 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,19 @@
1+
import os
12
from typing import TYPE_CHECKING, Optional
23
from uuid import uuid4
34

45
from langchain.tools import BaseTool
56
from langchain_core.messages import AIMessage
7+
from langsmith import Client
68

79
from codegen.extensions.langchain.agent import create_codebase_agent
810

911
if TYPE_CHECKING:
1012
from codegen import Codebase
1113

14+
# Remove logger configuration
15+
# logger = logging.getLogger(__name__)
16+
1217

1318
class CodeAgent:
1419
"""Agent for interacting with a codebase."""
@@ -30,6 +35,11 @@ def __init__(self, codebase: "Codebase", model_provider: str = "anthropic", mode
3035
"""
3136
self.codebase = codebase
3237
self.agent = create_codebase_agent(self.codebase, model_provider=model_provider, model_name=model_name, memory=memory, additional_tools=tools, **kwargs)
38+
self.langsmith_client = Client()
39+
40+
# Get project name from environment variable or use a default
41+
self.project_name = os.environ.get("LANGCHAIN_PROJECT", "RELACE")
42+
print(f"Using LangSmith project: {self.project_name}")
3343

3444
def run(self, prompt: str, thread_id: Optional[str] = None) -> str:
3545
"""Run the agent with a prompt.
@@ -49,7 +59,10 @@ def run(self, prompt: str, thread_id: Optional[str] = None) -> str:
4959
input = {"messages": [("user", prompt)]}
5060

5161
# we stream the steps instead of invoke because it allows us to access intermediate nodes
52-
stream = self.agent.stream(input, config={"configurable": {"thread_id": thread_id}, "recursion_limit": 100}, stream_mode="values")
62+
stream = self.agent.stream(input, config={"configurable": {"thread_id": thread_id, "metadata": {"project": self.project_name}}, "recursion_limit": 100}, stream_mode="values")
63+
64+
# Keep track of run IDs from the stream
65+
run_ids = []
5366

5467
for s in stream:
5568
message = s["messages"][-1]
@@ -61,5 +74,96 @@ def run(self, prompt: str, thread_id: Optional[str] = None) -> str:
6174
else:
6275
message.pretty_print()
6376

64-
# last stream object contains all messages. message[-1] is the last message
65-
return s["messages"][-1].content
77+
# Try to extract run ID if available in metadata
78+
if hasattr(message, "additional_kwargs") and "run_id" in message.additional_kwargs:
79+
run_ids.append(message.additional_kwargs["run_id"])
80+
81+
# Get the last message content
82+
result = s["messages"][-1].content
83+
84+
# Try to find run IDs in the LangSmith client's recent runs
85+
try:
86+
# Get the most recent runs with proper filter parameters
87+
# We need to provide at least one filter parameter as required by the API
88+
recent_runs = list(
89+
self.langsmith_client.list_runs(
90+
# Use the project name from environment variable
91+
project_name=self.project_name,
92+
# Limit to just the most recent run
93+
limit=1,
94+
)
95+
)
96+
97+
if recent_runs and len(recent_runs) > 0:
98+
# Make sure we have a valid run object with an id attribute
99+
if hasattr(recent_runs[0], "id"):
100+
# Convert the ID to string to ensure it's in the right format
101+
run_id = str(recent_runs[0].id)
102+
103+
# Get the run URL using the run_id parameter
104+
run_url = self.get_langsmith_url(run_id=run_id)
105+
106+
separator = "=" * 60
107+
print(f"\n{separator}\n🔍 LangSmith Run URL: {run_url}\n{separator}")
108+
else:
109+
separator = "=" * 60
110+
print(f"\n{separator}\nRun object has no 'id' attribute: {recent_runs[0]}\n{separator}")
111+
else:
112+
# If no runs found with project name, try a more general approach
113+
# Use a timestamp filter to get recent runs (last 10 minutes)
114+
import datetime
115+
116+
ten_minutes_ago = datetime.datetime.now() - datetime.timedelta(minutes=10)
117+
118+
recent_runs = list(self.langsmith_client.list_runs(start_time=ten_minutes_ago.isoformat(), limit=1))
119+
120+
if recent_runs and len(recent_runs) > 0 and hasattr(recent_runs[0], "id"):
121+
# Convert the ID to string to ensure it's in the right format
122+
run_id = str(recent_runs[0].id)
123+
124+
# Get the run URL using the run_id parameter
125+
run_url = self.get_langsmith_url(run_id=run_id)
126+
127+
separator = "=" * 60
128+
print(f"\n{separator}\n🔍 LangSmith Run URL: {run_url}\n{separator}")
129+
else:
130+
separator = "=" * 60
131+
print(f"\n{separator}\nNo valid runs found\n{separator}")
132+
except Exception as e:
133+
separator = "=" * 60
134+
print(f"\n{separator}\nCould not retrieve LangSmith URL: {e}")
135+
import traceback
136+
137+
print(traceback.format_exc())
138+
print(separator)
139+
140+
return result
141+
142+
def get_langsmith_url(self, run_id: str, project_name: Optional[str] = None) -> str:
143+
"""Get the URL for a run in LangSmith.
144+
145+
Args:
146+
run_id: The ID of the run
147+
project_name: Optional name of the project
148+
149+
Returns:
150+
The URL for the run in LangSmith
151+
"""
152+
# Construct the URL directly using the host URL and run ID
153+
# This avoids the issue with the client's get_run_url method expecting a run object
154+
host_url = self.langsmith_client._host_url
155+
tenant_id = self.langsmith_client._get_tenant_id()
156+
157+
# If project_name is not provided, use the default one
158+
if project_name is None:
159+
project_name = self.project_name
160+
161+
try:
162+
# Get the project ID from the project name
163+
project_id = self.langsmith_client.read_project(project_name=project_name).id
164+
# Construct the URL
165+
return f"{host_url}/o/{tenant_id}/projects/p/{project_id}/r/{run_id}?poll=true"
166+
except Exception as e:
167+
# If we can't get the project ID, construct a URL without it
168+
print(f"Could not get project ID for {project_name}: {e}")
169+
return f"{host_url}/o/{tenant_id}/r/{run_id}?poll=true"

0 commit comments

Comments
 (0)