Skip to content

Commit 79bb4fe

Browse files
dinmukhamedmnirga
andauthored
fix(langchain): tools in message history (#2939)
Co-authored-by: Nir Gazit <[email protected]>
1 parent 3a97238 commit 79bb4fe

File tree

10 files changed

+847
-101
lines changed

10 files changed

+847
-101
lines changed

packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py

Lines changed: 108 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
from opentelemetry.context.context import Context
2323
from opentelemetry.trace import SpanKind, set_span_in_context, Tracer
2424
from opentelemetry.trace.span import Span
25+
from opentelemetry.util.types import AttributeValue
2526

2627
from opentelemetry import context as context_api
2728
from opentelemetry.instrumentation.langchain.utils import (
@@ -53,12 +54,14 @@ def _message_type_to_role(message_type: str) -> str:
5354
return "system"
5455
elif message_type == "ai":
5556
return "assistant"
57+
elif message_type == "tool":
58+
return "tool"
5659
else:
5760
return "unknown"
5861

5962

60-
def _set_span_attribute(span, name, value):
61-
if value is not None:
63+
def _set_span_attribute(span: Span, name: str, value: AttributeValue):
64+
if value is not None and value != "":
6265
span.set_attribute(name, value)
6366

6467

@@ -75,9 +78,9 @@ def _set_request_params(span, kwargs, span_holder: SpanHolder):
7578
else:
7679
model = "unknown"
7780

78-
span.set_attribute(SpanAttributes.LLM_REQUEST_MODEL, model)
81+
_set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, model)
7982
# response is not available for LLM requests (as opposed to chat)
80-
span.set_attribute(SpanAttributes.LLM_RESPONSE_MODEL, model)
83+
_set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, model)
8184

8285
if "invocation_params" in kwargs:
8386
params = (
@@ -108,11 +111,13 @@ def _set_llm_request(
108111

109112
if should_send_prompts():
110113
for i, msg in enumerate(prompts):
111-
span.set_attribute(
114+
_set_span_attribute(
115+
span,
112116
f"{SpanAttributes.LLM_PROMPTS}.{i}.role",
113117
"user",
114118
)
115-
span.set_attribute(
119+
_set_span_attribute(
120+
span,
116121
f"{SpanAttributes.LLM_PROMPTS}.{i}.content",
117122
msg,
118123
)
@@ -144,20 +149,30 @@ def _set_chat_request(
144149
i = 0
145150
for message in messages:
146151
for msg in message:
147-
span.set_attribute(
152+
_set_span_attribute(
153+
span,
148154
f"{SpanAttributes.LLM_PROMPTS}.{i}.role",
149155
_message_type_to_role(msg.type),
150156
)
151-
# if msg.content is string
152-
if isinstance(msg.content, str):
153-
span.set_attribute(
154-
f"{SpanAttributes.LLM_PROMPTS}.{i}.content",
155-
msg.content,
156-
)
157+
tool_calls = (
158+
msg.tool_calls
159+
if hasattr(msg, "tool_calls")
160+
else msg.additional_kwargs.get("tool_calls")
161+
)
162+
163+
if tool_calls:
164+
_set_chat_tool_calls(span, f"{SpanAttributes.LLM_PROMPTS}.{i}", tool_calls)
165+
157166
else:
158-
span.set_attribute(
167+
content = (
168+
msg.content
169+
if isinstance(msg.content, str)
170+
else json.dumps(msg.content, cls=CallbackFilteredJSONEncoder)
171+
)
172+
_set_span_attribute(
173+
span,
159174
f"{SpanAttributes.LLM_PROMPTS}.{i}.content",
160-
json.dumps(msg.content, cls=CallbackFilteredJSONEncoder),
175+
content,
161176
)
162177
i += 1
163178

@@ -195,86 +210,117 @@ def _set_chat_response(span: Span, response: LLMResult) -> None:
195210
if should_send_prompts():
196211
prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{i}"
197212
if hasattr(generation, "text") and generation.text != "":
198-
span.set_attribute(
213+
_set_span_attribute(
214+
span,
199215
f"{prefix}.content",
200216
generation.text,
201217
)
202-
span.set_attribute(f"{prefix}.role", "assistant")
218+
_set_span_attribute(span, f"{prefix}.role", "assistant")
203219
else:
204-
span.set_attribute(
220+
_set_span_attribute(
221+
span,
205222
f"{prefix}.role",
206223
_message_type_to_role(generation.type),
207224
)
208225
if generation.message.content is str:
209-
span.set_attribute(
226+
_set_span_attribute(
227+
span,
210228
f"{prefix}.content",
211229
generation.message.content,
212230
)
213231
else:
214-
span.set_attribute(
232+
_set_span_attribute(
233+
span,
215234
f"{prefix}.content",
216235
json.dumps(
217236
generation.message.content, cls=CallbackFilteredJSONEncoder
218237
),
219238
)
220239
if generation.generation_info.get("finish_reason"):
221-
span.set_attribute(
240+
_set_span_attribute(
241+
span,
222242
f"{prefix}.finish_reason",
223243
generation.generation_info.get("finish_reason"),
224244
)
225245

226246
if generation.message.additional_kwargs.get("function_call"):
227-
span.set_attribute(
247+
_set_span_attribute(
248+
span,
228249
f"{prefix}.tool_calls.0.name",
229250
generation.message.additional_kwargs.get("function_call").get(
230251
"name"
231252
),
232253
)
233-
span.set_attribute(
254+
_set_span_attribute(
255+
span,
234256
f"{prefix}.tool_calls.0.arguments",
235257
generation.message.additional_kwargs.get("function_call").get(
236258
"arguments"
237259
),
238260
)
239261

240-
if generation.message.additional_kwargs.get("tool_calls"):
241-
for idx, tool_call in enumerate(
242-
generation.message.additional_kwargs.get("tool_calls")
243-
):
244-
tool_call_prefix = f"{prefix}.tool_calls.{idx}"
245-
246-
span.set_attribute(
247-
f"{tool_call_prefix}.id", tool_call.get("id")
248-
)
249-
span.set_attribute(
250-
f"{tool_call_prefix}.name",
251-
tool_call.get("function").get("name"),
252-
)
253-
span.set_attribute(
254-
f"{tool_call_prefix}.arguments",
255-
tool_call.get("function").get("arguments"),
256-
)
257-
i += 1
262+
if hasattr(generation, "message"):
263+
tool_calls = (
264+
generation.message.tool_calls
265+
if hasattr(generation.message, "tool_calls")
266+
else generation.message.additional_kwargs.get("tool_calls")
267+
)
268+
if tool_calls and isinstance(tool_calls, list):
269+
_set_span_attribute(
270+
span,
271+
f"{prefix}.role",
272+
"assistant",
273+
)
274+
_set_chat_tool_calls(span, prefix, tool_calls)
275+
i += 1
258276

259277
if input_tokens > 0 or output_tokens > 0 or total_tokens > 0 or cache_read_tokens > 0:
260-
span.set_attribute(
278+
_set_span_attribute(
279+
span,
261280
SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
262281
input_tokens,
263282
)
264-
span.set_attribute(
283+
_set_span_attribute(
284+
span,
265285
SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
266286
output_tokens,
267287
)
268-
span.set_attribute(
288+
_set_span_attribute(
289+
span,
269290
SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
270291
total_tokens,
271292
)
272-
span.set_attribute(
293+
_set_span_attribute(
294+
span,
273295
SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS,
274296
cache_read_tokens,
275297
)
276298

277299

300+
def _set_chat_tool_calls(span: Span, prefix: str, tool_calls: list[dict[str, Any]]) -> None:
301+
for idx, tool_call in enumerate(tool_calls):
302+
tool_call_prefix = f"{prefix}.tool_calls.{idx}"
303+
tool_call_dict = dict(tool_call)
304+
tool_id = tool_call_dict.get("id")
305+
tool_name = tool_call_dict.get("name", tool_call_dict.get("function", {}).get("name"))
306+
tool_args = tool_call_dict.get("args", tool_call_dict.get("function", {}).get("arguments"))
307+
308+
_set_span_attribute(
309+
span,
310+
f"{tool_call_prefix}.id", tool_id
311+
)
312+
_set_span_attribute(
313+
span,
314+
f"{tool_call_prefix}.name",
315+
tool_name,
316+
)
317+
_set_span_attribute(
318+
span,
319+
f"{tool_call_prefix}.arguments",
320+
json.dumps(tool_args, cls=CallbackFilteredJSONEncoder),
321+
)
322+
323+
278324
def _sanitize_metadata_value(value: Any) -> Any:
279325
"""Convert metadata values to OpenTelemetry-compatible types."""
280326
if value is None:
@@ -364,8 +410,8 @@ def _create_span(
364410
else:
365411
span = self.tracer.start_span(span_name, kind=kind)
366412

367-
span.set_attribute(SpanAttributes.TRACELOOP_WORKFLOW_NAME, workflow_name)
368-
span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_PATH, entity_path)
413+
_set_span_attribute(span, SpanAttributes.TRACELOOP_WORKFLOW_NAME, workflow_name)
414+
_set_span_attribute(span, SpanAttributes.TRACELOOP_ENTITY_PATH, entity_path)
369415

370416
token = context_api.attach(
371417
context_api.set_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, True)
@@ -402,8 +448,8 @@ def _create_task_span(
402448
metadata=metadata,
403449
)
404450

405-
span.set_attribute(SpanAttributes.TRACELOOP_SPAN_KIND, kind.value)
406-
span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, entity_name)
451+
_set_span_attribute(span, SpanAttributes.TRACELOOP_SPAN_KIND, kind.value)
452+
_set_span_attribute(span, SpanAttributes.TRACELOOP_ENTITY_NAME, entity_name)
407453

408454
return span
409455

@@ -427,8 +473,8 @@ def _create_llm_span(
427473
entity_path=entity_path,
428474
metadata=metadata,
429475
)
430-
span.set_attribute(SpanAttributes.LLM_SYSTEM, "Langchain")
431-
span.set_attribute(SpanAttributes.LLM_REQUEST_TYPE, request_type.value)
476+
_set_span_attribute(span, SpanAttributes.LLM_SYSTEM, "Langchain")
477+
_set_span_attribute(span, SpanAttributes.LLM_REQUEST_TYPE, request_type.value)
432478

433479
return span
434480

@@ -475,7 +521,8 @@ def on_chain_start(
475521
metadata,
476522
)
477523
if should_send_prompts():
478-
span.set_attribute(
524+
_set_span_attribute(
525+
span,
479526
SpanAttributes.TRACELOOP_ENTITY_INPUT,
480527
json.dumps(
481528
{
@@ -506,7 +553,8 @@ def on_chain_end(
506553
span_holder = self.spans[run_id]
507554
span = span_holder.span
508555
if should_send_prompts():
509-
span.set_attribute(
556+
_set_span_attribute(
557+
span,
510558
SpanAttributes.TRACELOOP_ENTITY_OUTPUT,
511559
json.dumps(
512560
{"outputs": outputs, "kwargs": kwargs},
@@ -586,13 +634,13 @@ def on_llm_end(
586634
"model_name"
587635
) or response.llm_output.get("model_id")
588636
if model_name is not None:
589-
span.set_attribute(SpanAttributes.LLM_RESPONSE_MODEL, model_name)
637+
_set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, model_name)
590638

591639
if self.spans[run_id].request_model is None:
592-
span.set_attribute(SpanAttributes.LLM_REQUEST_MODEL, model_name)
640+
_set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, model_name)
593641
id = response.llm_output.get("id")
594642
if id is not None and id != "":
595-
span.set_attribute(GEN_AI_RESPONSE_ID, id)
643+
_set_span_attribute(span, GEN_AI_RESPONSE_ID, id)
596644

597645
token_usage = (response.llm_output or {}).get("token_usage") or (
598646
response.llm_output or {}
@@ -687,7 +735,8 @@ def on_tool_start(
687735
entity_path,
688736
)
689737
if should_send_prompts():
690-
span.set_attribute(
738+
_set_span_attribute(
739+
span,
691740
SpanAttributes.TRACELOOP_ENTITY_INPUT,
692741
json.dumps(
693742
{
@@ -716,7 +765,8 @@ def on_tool_end(
716765

717766
span = self._get_span(run_id)
718767
if should_send_prompts():
719-
span.set_attribute(
768+
_set_span_attribute(
769+
span,
720770
SpanAttributes.TRACELOOP_ENTITY_OUTPUT,
721771
json.dumps(
722772
{"output": output, "kwargs": kwargs},

packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/utils.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import logging
55
import os
66
import traceback
7+
78
from opentelemetry import context as context_api
89
from opentelemetry.instrumentation.langchain.config import Config
910
from pydantic import BaseModel

0 commit comments

Comments
 (0)