Skip to content

Commit 61a49e3

Browse files
vprashrexvprashrex
committed
feat(langfuse): link LLM call traces to Kaapi IDs with shared telemetry keys
- Trace metadata carries job.id, llm.call_id, llm.config.id/version, tenant ids, gen_ai.provider.name and gen_ai.request.model, the same keys as the Sentry spans (_CallContext.identity_attributes is the single source). - Generation metadata adds llm.input/output type, format and language code, and gen_ai.response.id. - billable_output_tokens is the one reasoning-token rule for Sentry and Langfuse. - LLM_* keys and GEN_AI_RESPONSE_ID move to telemetry constants.
1 parent 357e775 commit 61a49e3

4 files changed

Lines changed: 117 additions & 35 deletions

File tree

‎backend/app/core/langfuse/langfuse.py‎

Lines changed: 56 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@
1414
from opentelemetry.sdk.trace import TracerProvider
1515

1616
from app.core.config import settings
17+
from app.core.telemetry.constants import GEN_AI_RESPONSE_ID, LLM_COMPLETION_TYPE
18+
from app.core.telemetry.llm import billable_output_tokens
1719
from app.models.llm import (
1820
AudioContent,
1921
AudioOutput,
@@ -24,7 +26,6 @@
2426
QueryParams,
2527
TextInput,
2628
TextOutput,
27-
Usage,
2829
)
2930
from app.models.llm.request import QueryInput
3031

@@ -179,6 +180,46 @@ def extract_query_input(
179180
return _query_input_item(query_input)
180181

181182

183+
def _input_metadata(trace_input: TraceInput) -> dict[str, str]:
184+
"""Input type, format and language code as filterable metadata."""
185+
items = trace_input if isinstance(trace_input, list) else [trace_input]
186+
types: dict[str, None] = {}
187+
formats: dict[str, None] = {}
188+
languages: dict[str, None] = {}
189+
for item in items:
190+
types[str(item["type"])] = None
191+
parts = item.get("items")
192+
for part in parts if isinstance(parts, list) else [item]:
193+
if part.get("format"):
194+
formats[str(part["format"])] = None
195+
if item.get("language_code"):
196+
languages[str(item["language_code"])] = None
197+
metadata = {
198+
"llm.input.type": ",".join(types),
199+
"llm.input.format": ",".join(formats),
200+
}
201+
if languages:
202+
metadata["llm.input.language_code"] = ",".join(languages)
203+
return metadata
204+
205+
206+
def _output_metadata(llm_output: TextOutput | AudioOutput | None) -> dict[str, str]:
207+
if isinstance(llm_output, TextOutput):
208+
metadata = {
209+
"llm.output.type": "text",
210+
"llm.output.format": llm_output.content.format,
211+
}
212+
if llm_output.content.language_code:
213+
metadata["llm.output.language_code"] = llm_output.content.language_code
214+
return metadata
215+
if isinstance(llm_output, AudioOutput):
216+
return {
217+
"llm.output.type": "audio",
218+
"llm.output.format": llm_output.content.format,
219+
}
220+
return {}
221+
222+
182223
def _model_parameters(params: dict[str, object]) -> dict[str, ModelParameterValue]:
183224
"""Scalar config params, one nested level flattened (``reasoning.effort``)."""
184225
scalars = (str, int, float, bool)
@@ -206,14 +247,6 @@ def _generation_input(params: dict[str, object], trace_input: TraceInput) -> obj
206247
]
207248

208249

209-
def _usage_details(usage: Usage) -> dict[str, int]:
210-
"""Gemini's output_tokens exclude reasoning_tokens; fold them in so cost is right."""
211-
return {
212-
"input": usage.input_tokens,
213-
"output": usage.output_tokens + (usage.reasoning_tokens or 0),
214-
}
215-
216-
217250
def extract_response_output(
218251
llm_output: TextOutput | AudioOutput | None,
219252
) -> str | dict[str, Any]:
@@ -399,12 +432,14 @@ def flush(self) -> None:
399432
def observe_llm_execution(
400433
session_id: str | None = None,
401434
credentials: dict | None = None,
435+
trace_metadata: dict[str, str] | None = None,
402436
) -> Callable:
403437
"""Decorator to add Langfuse observability to LLM provider execute methods.
404438
405439
Args:
406440
credentials: Langfuse credentials with public_key, secret_key, and host
407441
session_id: Session ID for grouping traces (conversation_id)
442+
trace_metadata: Call IDs, same keys as the Sentry spans (_CallContext.identity_attributes)
408443
409444
Usage:
410445
decorated_execute = observe_llm_execution(
@@ -462,6 +497,7 @@ def langfuse_call(fn, *args, **kwargs):
462497
as_type="span",
463498
name="unified-llm-call",
464499
input=trace_input,
500+
metadata=trace_metadata,
465501
)
466502
if trace:
467503
langfuse_call(
@@ -472,6 +508,7 @@ def langfuse_call(fn, *args, **kwargs):
472508
session_id=session_id,
473509
# OTel drops the whole list if any item is an Enum, not a plain str.
474510
tags=[completion_config.provider, str(completion_config.type)],
511+
metadata=trace_metadata,
475512
)
476513

477514
generation = None
@@ -483,7 +520,10 @@ def langfuse_call(fn, *args, **kwargs):
483520
input=_generation_input(params, trace_input),
484521
model=params.get("model"),
485522
model_parameters=_model_parameters(params),
486-
metadata={"completion_type": completion_config.type},
523+
metadata={
524+
LLM_COMPLETION_TYPE: str(completion_config.type),
525+
**_input_metadata(trace_input),
526+
},
487527
)
488528

489529
def finish_failure(error_msg: str) -> None:
@@ -528,10 +568,14 @@ def finish_failure(error_msg: str) -> None:
528568
langfuse_call(
529569
generation.update,
530570
output=output,
531-
usage_details=_usage_details(response.usage),
571+
usage_details={
572+
"input": response.usage.input_tokens,
573+
"output": billable_output_tokens(response.usage),
574+
},
532575
model=response.response.model,
533576
metadata={
534-
"provider_response_id": response.response.provider_response_id
577+
GEN_AI_RESPONSE_ID: response.response.provider_response_id,
578+
**_output_metadata(response.response.output),
535579
},
536580
)
537581
langfuse_call(generation.end)

‎backend/app/core/telemetry/constants.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
GEN_AI_REQUEST_PRESENCE_PENALTY,
1414
GEN_AI_REQUEST_TEMPERATURE,
1515
GEN_AI_REQUEST_TOP_P,
16+
GEN_AI_RESPONSE_ID,
1617
GEN_AI_RESPONSE_MODEL,
1718
GEN_AI_SYSTEM,
1819
GEN_AI_USAGE_INPUT_TOKENS,
@@ -38,6 +39,13 @@
3839
GEN_AI_REQUEST_AVAILABLE_TOOLS = "gen_ai.request.available_tools"
3940
GEN_AI_USAGE_TOTAL_TOKENS = "gen_ai.usage.total_tokens"
4041

42+
# --- LLM call: same keys on Sentry spans and Langfuse trace metadata ---
43+
LLM_CALL_ID = "llm.call_id"
44+
LLM_CHAIN_ID = "llm.chain_id"
45+
LLM_CONFIG_ID = "llm.config.id"
46+
LLM_CONFIG_VERSION = "llm.config.version"
47+
LLM_COMPLETION_TYPE = "llm.completion_type"
48+
4149
# --- HTTP ---
4250
HTTP_METHOD = "http.method"
4351
HTTP_ROUTE = "http.route"

‎backend/app/core/telemetry/llm.py‎

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,18 @@
3232
from app.core.telemetry.metrics import emit_sentry_metric
3333

3434
if TYPE_CHECKING:
35-
from app.models.llm.response import LLMCallResponse
35+
from app.models.llm.response import LLMCallResponse, Usage
36+
37+
38+
def billable_output_tokens(usage: "Usage") -> int:
39+
"""Output plus reasoning when the provider reports them separately (Gemini), so cost is right."""
40+
reasoning_tokens = usage.reasoning_tokens or 0
41+
if (
42+
reasoning_tokens
43+
and usage.input_tokens + usage.output_tokens < usage.total_tokens
44+
):
45+
return usage.output_tokens + reasoning_tokens
46+
return usage.output_tokens
3647

3748

3849
def _llm_call_attrs(
@@ -144,15 +155,13 @@ def set_gen_ai_response_attributes(
144155
"""Set OTel GenAI response attributes (usage, model) on `span`."""
145156
usage = response.usage
146157
if usage:
147-
reasoning_tokens = usage.reasoning_tokens or 0
148-
output_tokens = usage.output_tokens
149-
if reasoning_tokens and usage.input_tokens + output_tokens < usage.total_tokens:
150-
output_tokens += reasoning_tokens
151158
span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens)
152-
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens)
159+
span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, billable_output_tokens(usage))
153160
span.set_attribute(GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens)
154-
if reasoning_tokens:
155-
span.set_attribute(GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, reasoning_tokens)
161+
if usage.reasoning_tokens:
162+
span.set_attribute(
163+
GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, usage.reasoning_tokens
164+
)
156165

157166
if response.response and response.response.model:
158167
span.set_attribute(GEN_AI_RESPONSE_MODEL, response.response.model)

‎backend/app/services/llm/jobs.py‎

Lines changed: 36 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,11 @@
2525
GEN_AI_PROVIDER_NAME,
2626
GEN_AI_REQUEST_MODEL,
2727
JOB_ID,
28+
LLM_CALL_ID,
29+
LLM_CHAIN_ID,
30+
LLM_COMPLETION_TYPE,
31+
LLM_CONFIG_ID,
32+
LLM_CONFIG_VERSION,
2833
TENANT_ORG_ID,
2934
TENANT_PROJECT_ID,
3035
)
@@ -116,9 +121,9 @@ def _set_traceability_attributes(
116121
if job_id is not None:
117122
span.set_attribute(JOB_ID, str(job_id))
118123
if llm_call_id is not None:
119-
span.set_attribute("llm.call_id", str(llm_call_id))
124+
span.set_attribute(LLM_CALL_ID, str(llm_call_id))
120125
if chain_id is not None:
121-
span.set_attribute("llm.chain_id", str(chain_id))
126+
span.set_attribute(LLM_CHAIN_ID, str(chain_id))
122127
if trace_id is not None:
123128
span.set_attribute(CORRELATION_ID, trace_id)
124129
if project_id is not None:
@@ -136,6 +141,7 @@ def _execute_provider_call(
136141
query: QueryParams,
137142
credentials: dict | None,
138143
session_id: str | None,
144+
trace_metadata: dict[str, str] | None = None,
139145
**kwargs: Any,
140146
) -> tuple[Any, Any]:
141147
kwargs.pop("organization_id", None)
@@ -145,6 +151,7 @@ def _execute_provider_call(
145151
decorated = observe_llm_execution(
146152
session_id=session_id,
147153
credentials=credentials,
154+
trace_metadata=trace_metadata,
148155
)(func)
149156

150157
with suppress_http_instrumentation():
@@ -604,21 +611,29 @@ class _CallContext:
604611
trace_id: str | None
605612
record_call: bool
606613
llm_call_id: UUID | None = None
614+
config_id: UUID | None = None
615+
config_version: int | None = None
607616
started_at: float = 0.0
608617

618+
def identity_attributes(self) -> dict[str, str]:
619+
"""IDs of this call, with the same keys on Sentry spans and Langfuse trace metadata."""
620+
ids = {
621+
JOB_ID: self.job_id,
622+
LLM_CALL_ID: self.llm_call_id,
623+
LLM_CHAIN_ID: self.chain_id,
624+
LLM_CONFIG_ID: self.config_id,
625+
LLM_CONFIG_VERSION: self.config_version,
626+
CORRELATION_ID: self.trace_id,
627+
TENANT_ORG_ID: self.organization_id,
628+
TENANT_PROJECT_ID: self.project_id,
629+
}
630+
return {key: str(value) for key, value in ids.items() if value is not None}
631+
609632
@contextmanager
610633
def span(self, name: str) -> Iterator[trace.Span]:
611634
span_tracer = tracer if self.record_call else trace.NoOpTracer()
612635
with span_tracer.start_as_current_span(name) as span:
613-
_set_traceability_attributes(
614-
span,
615-
job_id=self.job_id,
616-
llm_call_id=self.llm_call_id,
617-
chain_id=self.chain_id,
618-
trace_id=self.trace_id,
619-
project_id=self.project_id,
620-
organization_id=self.organization_id,
621-
)
636+
span.set_attributes(self.identity_attributes())
622637
yield span
623638

624639
def metrics_started(self, *, provider: str, model: str) -> None:
@@ -663,8 +678,9 @@ def _load_config_blob(
663678
with ctx.span("llm.resolve_config") as cfg_span:
664679
cfg_span.set_attribute("llm.config.is_stored", config.is_stored_config)
665680
if config.is_stored_config:
666-
cfg_span.set_attribute("llm.config.id", str(config.id))
667-
cfg_span.set_attribute("llm.config.version", str(config.version))
681+
ctx.config_id, ctx.config_version = config.id, config.version
682+
cfg_span.set_attribute(LLM_CONFIG_ID, str(config.id))
683+
cfg_span.set_attribute(LLM_CONFIG_VERSION, str(config.version))
668684
config_crud = ConfigVersionCrud(
669685
session=session, project_id=ctx.project_id, config_id=config.id
670686
)
@@ -1354,7 +1370,7 @@ def execute_llm_call(
13541370
with ctx.span(ai_span_name) as ai_span:
13551371
ai_span.set_attribute("sentry.op", "gen_ai.chat")
13561372
if completion_type:
1357-
ai_span.set_attribute("llm.completion_type", completion_type)
1373+
ai_span.set_attribute(LLM_COMPLETION_TYPE, completion_type)
13581374
set_gen_ai_request_attributes(
13591375
ai_span,
13601376
provider=provider_name,
@@ -1376,7 +1392,7 @@ def execute_llm_call(
13761392
)
13771393
if completion_type:
13781394
provider_span.set_attribute(
1379-
"llm.completion_type", completion_type
1395+
LLM_COMPLETION_TYPE, completion_type
13801396
)
13811397
if model_name:
13821398
provider_span.set_attribute(GEN_AI_REQUEST_MODEL, model_name)
@@ -1387,6 +1403,11 @@ def execute_llm_call(
13871403
query=query,
13881404
credentials=langfuse_credentials,
13891405
session_id=conversation_id,
1406+
trace_metadata={
1407+
**ctx.identity_attributes(),
1408+
GEN_AI_PROVIDER_NAME: provider_name,
1409+
GEN_AI_REQUEST_MODEL: model_name,
1410+
},
13901411
organization_id=organization_id,
13911412
project_id=project_id,
13921413
telemetry_span=provider_span,

0 commit comments

Comments
 (0)