RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Track user and session ID on Langfuse traces (#814) 

* telemetry: use session and user IDs in langfuse

* test: update old span test

* fix: disable langfuse in unit tests

* fix: add post-loop synthesis span

* refactor: address PR review feedback on langfuse tracing

- Consolidate track_name onto LLMTelemetryContext as the sole home;
  remove the honcho_llm_call kwarg and update 4 callers to set it on
  telemetry directly. Sentry ai_track now reads telemetry.track_name.
- Decouple escaped-stream self-stamping from run-context exit ordering:
  stream_final_response now resets _in_agent_run explicitly around drain.
- Narrow langfuse_agent_step wrap in the tool loop — between-turn
  bookkeeping (iteration_callback, choice switch, increment) lifted
  outside the span so it scopes only the LLM call + tools.
- Reword test conftest comment to behavior-only language.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* refactor: switch langfuse spans to imperative handles

Replaces the context-manager-based langfuse_agent_run/step with imperative
LangfuseAgentRun/Step handles so the run span can outlive the function that
opens it. Streaming responses now own the run handle from construction and
close it after drain, stamping the accumulated streamed text as trace output
(previously blank). Multi-turn generations always stamp provider/model and
step metadata, fixing the regression where only the first turn was annotated.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* fix(llm): record effective prompt-only input on run span

The run-level Langfuse span recorded the raw messages parameter, which is
None for prompt-only calls. Mirror execute_tool_loop's handling and record
the synthesized user message so the trace input isn't blank.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* fix(llm): drop StreamingResponseWithMetadata.__anext__ to prevent span leak

The standalone __anext__ delegated straight to the inner stream, bypassing
the token-folding and Langfuse run-handle close that live only in the
__aiter__ generator. Any caller driving the wrapper via anext() instead of
`async for` would leak the run span and lose final-stream token accounting.
Latent today (all callers use `async for`), removed to close the footgun.

Add tests covering the run-handle drain path: full drain stamps the
accumulated streamed text as the span output and closes once; an abandoned
stream still closes via the finally rather than leaking.

* chore(llm): document intentional empty-body propagate_attributes block

The `with propagate_attributes(...): pass` stamps the active @observe trace
root via the context manager's __enter__ side effect; the empty body reads
as deletable dead code. Add a comment so it isn't removed. Addresses PR review.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fix(llm): restore api.py types after __anext__ removal

Dropping StreamingResponseWithMetadata.__anext__ made it stop satisfying
the AsyncIterator protocol, breaking the result annotation and the
isinstance narrowing in honcho_llm_call. Widen the tool-less result
annotation to include StreamingResponseWithMetadata and narrow positively
to HonchoLLMCallResponse before reading .content.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
This commit is contained in:
Aakash Kattelu 2026-06-23 00:03:21 -04:00 committed by GitHub
parent e2ff106f28
commit e8ef1a06e5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
17 changed files with 1501 additions and 294 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.env.template
View File

@ -25,6 +25,7 @@ LOG_LEVEL=INFO
# LANGFUSE_HOST=
# LANGFUSE_PUBLIC_KEY=
# LANGFUSE_SECRET_KEY=
# COLLECT_METRICS_LOCAL=false
# LOCAL_METRICS_FILE=metrics.jsonl

2
src/deriver/deriver.py
View File

@ -147,7 +147,6 @@ async def process_representation_tasks_batch(
model_config=model_config,
prompt=prompt,
max_tokens=max_tokens,
track_name="Minimal Deriver",
response_model=PromptRepresentation,
json_mode=True,
max_input_tokens=settings.DERIVER.MAX_INPUT_TOKENS,
@ -159,6 +158,7 @@ async def process_representation_tasks_batch(
call_purpose=CallPurpose.DERIVER_REPRESENTATION.value,
parent_category="representation",
observed=observed,
track_name="Minimal Deriver",
),
)
llm_duration = (time.perf_counter() - llm_start) * 1000

12
src/dialectic/core.py
View File

@ -301,13 +301,14 @@ class DialecticAgent:
return tool_executor, task_name, run_id, start_time
def _telemetry_context(self) -> LLMTelemetryContext:
def _telemetry_context(self, track_name: str | None = None) -> LLMTelemetryContext:
"""Build the LLMTelemetryContext shared by answer() and answer_stream().
Carries the instance's `_run_id` (always set in __init__) + workspace +
peer identifiers so LLMCallCompletedEvent and 's
AgentIterationEvent can attribute every per-iteration LLM call back to
this dialectic invocation.
this dialectic invocation. `track_name` names the Langfuse trace/step
(e.g. "Dialectic Agent" vs "Dialectic Agent Stream").
"""
return LLMTelemetryContext(
workspace_name=self.workspace_name,
@ -316,6 +317,7 @@ class DialecticAgent:
agent_type="dialectic",
run_id=self._run_id,
peer_name=self.observed,
track_name=track_name,
)
def _log_response_metrics(
@ -446,10 +448,9 @@ class DialecticAgent:
tool_executor=tool_executor,
max_tool_iterations=level_settings.MAX_TOOL_ITERATIONS,
messages=self.messages,
track_name="Dialectic Agent",
max_input_tokens=settings.DIALECTIC.MAX_INPUT_TOKENS,
trace_name="dialectic_chat",
telemetry=self._telemetry_context(),
telemetry=self._telemetry_context(track_name="Dialectic Agent"),
)
self._log_response_metrics(
@ -515,10 +516,9 @@ class DialecticAgent:
tool_executor=tool_executor,
max_tool_iterations=level_settings.MAX_TOOL_ITERATIONS,
messages=self.messages,
track_name="Dialectic Agent Stream",
max_input_tokens=settings.DIALECTIC.MAX_INPUT_TOKENS,
trace_name="dialectic_chat",
telemetry=self._telemetry_context(),
telemetry=self._telemetry_context(track_name="Dialectic Agent Stream"),
),
)

2
src/dreamer/specialists.py
View File

@ -286,7 +286,6 @@ If you update it, send the full deduplicated list and remove stale entries.
tool_executor=tool_executor,
max_tool_iterations=self.get_max_iterations(),
messages=messages,
track_name=f"Dreamer/{self.name}",
telemetry=LLMTelemetryContext(
workspace_name=workspace_name,
call_purpose=call_purpose_slug,
@ -295,6 +294,7 @@ If you update it, send the full deduplicated list and remove stale entries.
run_id=run_id,
observer=observer,
observed=observed,
track_name=f"Dreamer/{self.name}",
),
)

103
src/llm/api.py
View File

@ -21,7 +21,6 @@ from tenacity import retry, stop_after_attempt, wait_exponential
from src.config import ConfiguredModelSettings, ModelConfig
from src.exceptions import ValidationException
from src.telemetry.logging import conditional_observe
from src.telemetry.reasoning_traces import log_reasoning_trace
from .executor import honcho_llm_call_inner
@ -31,7 +30,7 @@ from .runtime import (
effective_temperature,
plan_attempt,
resolve_runtime_model_config,
update_current_langfuse_observation,
start_langfuse_agent_run,
)
from .tool_loop import execute_tool_loop
from .types import (
@ -54,7 +53,6 @@ async def honcho_llm_call(
model_config: ModelConfig | ConfiguredModelSettings,
prompt: str,
max_tokens: int,
track_name: str | None = None,
response_model: type[M],
json_mode: bool = False,
temperature: float | None = None,
@ -84,7 +82,6 @@ async def honcho_llm_call(
model_config: ModelConfig | ConfiguredModelSettings,
prompt: str,
max_tokens: int,
track_name: str | None = None,
response_model: None = None,
json_mode: bool = False,
temperature: float | None = None,
@ -114,7 +111,6 @@ async def honcho_llm_call(
model_config: ModelConfig | ConfiguredModelSettings,
prompt: str,
max_tokens: int,
track_name: str | None = None,
response_model: type[BaseModel] | None = None,
json_mode: bool = False,
temperature: float | None = None,
@ -138,13 +134,11 @@ async def honcho_llm_call(
) -> AsyncIterator[HonchoLLMCallStreamChunk] | StreamingResponseWithMetadata: ...
@conditional_observe(name="LLM Call")
async def honcho_llm_call(
*,
model_config: ModelConfig | ConfiguredModelSettings,
prompt: str,
max_tokens: int,
track_name: str | None = None,
response_model: type[BaseModel] | None = None,
json_mode: bool = False,
temperature: float | None = None,
@ -206,11 +200,6 @@ async def honcho_llm_call(
call_thinking_budget_tokens=thinking_budget_tokens,
call_reasoning_effort=reasoning_effort,
)
update_current_langfuse_observation(
plan.provider,
plan.model,
name=track_name,
)
return plan
async def _call_with_provider_selection() -> (
@ -267,8 +256,9 @@ async def honcho_llm_call(
decorated = _call_with_provider_selection
if track_name:
decorated = ai_track(track_name)(decorated)
sentry_track_name = telemetry.track_name if telemetry is not None else None
if sentry_track_name:
decorated = ai_track(sentry_track_name)(decorated)
def before_retry_callback(retry_state: Any) -> None:
"""Update attempt counter before each retry + log transient failures.
@ -397,8 +387,8 @@ async def honcho_llm_call(
)
wrapped = _toolless_call
if track_name:
wrapped = ai_track(track_name)(wrapped)
if sentry_track_name:
wrapped = ai_track(sentry_track_name)(wrapped)
if enable_retry:
wrapped = retry(
stop=stop_after_attempt(retry_attempts),
@ -406,7 +396,9 @@ async def honcho_llm_call(
before_sleep=before_retry_callback,
)(wrapped)
result: (
HonchoLLMCallResponse[Any] | AsyncIterator[HonchoLLMCallStreamChunk]
HonchoLLMCallResponse[Any]
| AsyncIterator[HonchoLLMCallStreamChunk]
| StreamingResponseWithMetadata
) = await wrapped()
else:
result = await decorated()
@ -429,30 +421,59 @@ async def honcho_llm_call(
)
return result
# execute_tool_loop raises ValidationException on out-of-range
# max_tool_iterations; fail-fast is cheaper than silent clamping here.
result = await execute_tool_loop(
prompt=prompt,
max_tokens=max_tokens,
messages=messages,
tools=tools,
tool_choice=tool_choice,
tool_executor=tool_executor,
max_tool_iterations=max_tool_iterations,
response_model=response_model,
json_mode=json_mode,
temperature=temperature,
stop_seqs=stop_seqs,
verbosity=verbosity,
enable_retry=enable_retry,
retry_attempts=retry_attempts,
max_input_tokens=max_input_tokens,
get_attempt_plan=_get_attempt_plan,
before_retry_callback=before_retry_callback,
stream_final=stream_final_only,
iteration_callback=iteration_callback,
telemetry=telemetry,
)
# One run-level Langfuse trace wraps the whole run; step/LLM/tool spans
# nest under it (the run handle keeps `start_as_current_observation` open
# via ExitStack, so the run span stays current OTel-wise even though we
# never use a `with` block here). The handle is passed into
# `execute_tool_loop` so streaming results own it from construction and
# close the span after drain — that's how the streamed text shows up as
# the trace's output instead of blank. Non-streaming results: we end in
# the `finally`.
run_label = (telemetry.track_name if telemetry else None) or "Agent"
run_handle = start_langfuse_agent_run(run_label, telemetry)
if run_handle is not None:
# Mirror execute_tool_loop's prompt-only handling: when messages is
# omitted it seeds the conversation with a single user message built
# from prompt. Record that same effective input so the run span isn't
# blank for prompt-only calls.
run_handle.update(
input=messages if messages else [{"role": "user", "content": prompt}]
)
try:
# execute_tool_loop raises ValidationException on out-of-range
# max_tool_iterations; fail-fast is cheaper than silent clamping here.
result = await execute_tool_loop(
prompt=prompt,
max_tokens=max_tokens,
messages=messages,
tools=tools,
tool_choice=tool_choice,
tool_executor=tool_executor,
max_tool_iterations=max_tool_iterations,
response_model=response_model,
json_mode=json_mode,
temperature=temperature,
stop_seqs=stop_seqs,
verbosity=verbosity,
enable_retry=enable_retry,
retry_attempts=retry_attempts,
max_input_tokens=max_input_tokens,
get_attempt_plan=_get_attempt_plan,
before_retry_callback=before_retry_callback,
stream_final=stream_final_only,
iteration_callback=iteration_callback,
telemetry=telemetry,
langfuse_run_handle=run_handle,
)
except BaseException:
if run_handle is not None:
run_handle.end()
raise
# Streaming wrapper owns the handle and closes it after drain;
# non-streaming paths (always a HonchoLLMCallResponse here) close it now
# with the final content as output.
if run_handle is not None and isinstance(result, HonchoLLMCallResponse):
run_handle.end(output=result.content)
if trace_name and isinstance(result, HonchoLLMCallResponse):
log_reasoning_trace(
task_type=trace_name,

19
src/llm/executor.py
View File

@ -20,13 +20,18 @@ from typing import Any, Literal, TypeVar, overload
from pydantic import BaseModel
from src.config import ModelConfig, ModelTransport
from src.telemetry.logging import conditional_observe
from .backend import CompletionResult as BackendCompletionResult
from .backend import StreamChunk as BackendStreamChunk
from .backend import ToolCallResult
from .registry import CLIENTS, backend_for_provider
from .request_builder import execute_completion, execute_stream
from .runtime import AttemptPlan, effective_config_for_call
from .runtime import (
AttemptPlan,
annotate_current_langfuse_trace,
effective_config_for_call,
)
from .types import (
HonchoLLMCallResponse,
HonchoLLMCallStreamChunk,
@ -261,6 +266,7 @@ async def honcho_llm_call_inner(
) -> AsyncIterator[HonchoLLMCallStreamChunk]: ...
@conditional_observe(name="LLM Call", as_type="generation")
async def honcho_llm_call_inner(
provider: ModelTransport,
model: str,
@ -284,6 +290,11 @@ async def honcho_llm_call_inner(
) -> HonchoLLMCallResponse[Any] | AsyncIterator[HonchoLLMCallStreamChunk]:
"""One backend call. No retry, no fallback, no tool loop.
This is the Langfuse trace boundary (``@conditional_observe``): every
provider call is its own trace. Multi-turn agents thread a shared
``run_id`` through ``telemetry`` so their per-iteration traces roll up into
one Langfuse session (see ``annotate_current_langfuse_trace``).
The outer src/llm/api.py `honcho_llm_call` handles retry + fallback +
tool orchestration on top of this.
@ -300,6 +311,12 @@ async def honcho_llm_call_inner(
if client is None:
raise ValueError(f"Missing client for {provider}")
# Stamp this trace (user_id/session_id/metadata) now that the @observe
# span is open and the resolved provider/model are known. Set early so the
# annotation lands even on the stream path, where the span closes once the
# generator is returned (before chunks drain).
annotate_current_langfuse_trace(provider, model, telemetry=telemetry)
if messages is None:
messages = [{"role": "user", "content": prompt}]

284
src/llm/runtime.py
View File

@ -12,8 +12,9 @@ Owns:
from __future__ import annotations
import logging
from contextlib import ExitStack
from contextvars import ContextVar
from dataclasses import dataclass
from dataclasses import dataclass, field
from typing import Any
from src.config import (
@ -25,39 +26,284 @@ from src.config import (
)
from .registry import backend_for_provider, client_for_model_config
from .types import ProviderClient, ReasoningEffortType
from .types import LLMTelemetryContext, ProviderClient, ReasoningEffortType
logger = logging.getLogger(__name__)
# ContextVar tracking the current retry attempt for provider switching.
current_attempt: ContextVar[int] = ContextVar("current_attempt", default=0)
# True while a `LangfuseAgentRun` handle is live (start → end). Set by
# `start_langfuse_agent_run`, reset by `LangfuseAgentRun.end`. Used by
# `annotate_current_langfuse_trace` to decide whether the current generation
# is the trace root (single-shot callers like the deriver — stamp trace attrs)
# or nested under an active run (multi-turn / streaming — skip trace attrs;
# the run span already carries them via `propagate_attributes`).
_in_agent_run: ContextVar[bool] = ContextVar("_in_agent_run", default=False)
def update_current_langfuse_observation(
def annotate_current_langfuse_trace(
provider: ModelTransport,
model: str,
*,
name: str | None = None,
telemetry: LLMTelemetryContext | None = None,
) -> None:
"""Best-effort annotation of the current Langfuse span with LLM routing."""
"""Stamp provider/model + step metadata on the current Langfuse generation.
Inside an active agent run, `propagate_attributes` already stamped
user_id/session_id/trace_name on the run span; this call only needs to
decorate the per-iteration generation. Outside a run (single-shot
callers — deriver, summarizer), this generation IS the trace root, so we
also stamp the trace attrs.
Note: `model`/`metadata` are set on every call regardless of `inside_run`
so multi-turn iterations no longer lose provider/model attribution.
"""
if not settings.LANGFUSE_PUBLIC_KEY:
return
try:
from langfuse import get_client, propagate_attributes
inside_run = _in_agent_run.get()
gen_metadata = _step_metadata(telemetry) if telemetry is not None else {}
gen_metadata["provider"] = str(provider)
gen_metadata["model"] = str(model)
gen_name = (
f"{telemetry.track_name} LLM call"
if telemetry is not None and telemetry.track_name
else None
)
if not inside_run:
run_id = telemetry.run_id if telemetry is not None else None
trace_name = telemetry.track_name if telemetry is not None else None
trace_metadata: dict[str, str] = dict(gen_metadata)
if telemetry is None:
trace_metadata.setdefault("namespace", str(settings.NAMESPACE))
# Empty body is intentional: propagate_attributes stamps the active
# @observe generation (this trace root, for single-shot callers) at
# __enter__; there are no child spans to scope here. Don't delete as
# dead code — the enter-time side effect is the point.
with propagate_attributes(
user_id=str(settings.NAMESPACE),
session_id=run_id,
trace_name=trace_name,
metadata=trace_metadata,
):
pass
get_client().update_current_generation(
name=gen_name,
model=str(model),
metadata=gen_metadata,
)
except Exception as exc: # pragma: no cover - best-effort telemetry
logger.debug("Failed to update Langfuse trace metadata: %s", exc)
def _base_metadata(telemetry: LLMTelemetryContext) -> dict[str, str]:
"""Static routing/attribution metadata (everything except ``iteration``).
Rebuilt per run (cheap); callers that need ``iteration`` copy and add it.
"""
metadata: dict[str, str] = {"namespace": str(settings.NAMESPACE)}
for key, value in (
("workspace_name", telemetry.workspace_name),
("call_purpose", telemetry.call_purpose),
("agent_type", telemetry.agent_type),
("observer", telemetry.observer),
("observed", telemetry.observed),
("peer_name", telemetry.peer_name),
):
if value is not None:
metadata[key] = str(value)
return metadata
def _step_metadata(
telemetry: LLMTelemetryContext,
base: dict[str, str] | None = None,
) -> dict[str, str]:
"""Per-step metadata: ``base`` (or freshly computed) plus ``iteration``."""
metadata = dict(base) if base is not None else _base_metadata(telemetry)
if telemetry.iteration is not None:
metadata["iteration"] = str(telemetry.iteration)
return metadata
@dataclass
class LangfuseAgentRun:
"""Imperative handle for the run-level Langfuse span.
Owns an ``ExitStack`` that keeps ``start_as_current_observation`` and
``propagate_attributes`` open until ``.end()``. This lets the run span
outlive the function that created it — streaming flows transfer the
handle to the response wrapper, which calls ``.end(output=...)`` after
the stream drains. While the handle is alive, the run span is the
current OTel observation, so step spans and auto-instrumented LLM
generations nest under it without any ContextVar choreography.
Use ``start_langfuse_agent_run`` to construct; never instantiate directly.
"""
span: Any # LangfuseSpan; opaque to keep src/llm/ free of langfuse imports.
_stack: ExitStack
_run_token: Any
_ended: bool = field(default=False)
def update(self, **kwargs: Any) -> None:
"""Set input/output/metadata on the run span (best-effort, no-op if ended)."""
if self._ended or self.span is None:
return
try:
self.span.update(**kwargs)
except Exception as exc: # pragma: no cover - best-effort telemetry
logger.debug("Failed to update Langfuse run span: %s", exc)
def end(self, *, output: Any = None) -> None:
"""Stamp final output (optional) and close the run span. Idempotent."""
if self._ended:
return
self._ended = True
try:
if self.span is not None and output is not None:
self.span.update(output=output)
except Exception as exc: # pragma: no cover - best-effort telemetry
logger.debug("Failed to set Langfuse run output: %s", exc)
try:
self._stack.close()
except Exception as exc: # pragma: no cover - best-effort telemetry
logger.debug("Failed to close Langfuse run span: %s", exc)
try:
_in_agent_run.reset(self._run_token)
except (ValueError, LookupError) as exc: # pragma: no cover
# ContextVar.reset can raise if end() runs in a different async
# context than start(); telemetry must not fail user code.
logger.debug("Failed to reset _in_agent_run: %s", exc)
def start_langfuse_agent_run(
name: str, telemetry: LLMTelemetryContext | None
) -> LangfuseAgentRun | None:
"""Open the one run-level Langfuse trace per agentic run, imperatively.
Returns ``None`` when Langfuse is disabled or there's no ``run_id``
(single-shot callers — those self-stamp via
``annotate_current_langfuse_trace``). When non-None, the caller MUST
eventually call ``.end()`` — typically in a ``finally`` block, or by
transferring ownership to the streaming wrapper.
"""
if not settings.LANGFUSE_PUBLIC_KEY or telemetry is None or not telemetry.run_id:
return None
stack = ExitStack()
try:
from langfuse import get_client, propagate_attributes
span = stack.enter_context(
get_client().start_as_current_observation(as_type="span", name=name)
)
stack.enter_context(
propagate_attributes(
user_id=str(settings.NAMESPACE),
session_id=telemetry.run_id,
trace_name=name,
metadata=_base_metadata(telemetry),
)
)
except Exception as exc: # pragma: no cover - best-effort telemetry
logger.debug("Failed to open Langfuse agent run: %s", exc)
stack.close()
return None
run_token = _in_agent_run.set(True)
return LangfuseAgentRun(span=span, _stack=stack, _run_token=run_token)
@dataclass
class LangfuseAgentStep:
"""Imperative handle for a per-iteration step span under the run root.
Owns an ``ExitStack`` holding ``start_as_current_observation`` open until
``.end()``. While alive the step span is the current OTel observation,
so the LLM generation (auto-instrumented or otherwise) nests under it.
No trace attrs (the run root carries them); just the per-step
``iteration`` metadata.
"""
span: Any
_stack: ExitStack
_ended: bool = field(default=False)
def update(self, **kwargs: Any) -> None:
"""Set input/output/metadata on the step span (best-effort, no-op if ended)."""
if self._ended or self.span is None:
return
try:
self.span.update(**kwargs)
except Exception as exc: # pragma: no cover - best-effort telemetry
logger.debug("Failed to update Langfuse step span: %s", exc)
def annotate_io(
self,
messages: list[dict[str, Any]],
content: Any,
tool_calls: list[dict[str, Any]],
) -> None:
"""Stamp this turn's messages-in / content-or-tool-summary-out.
On a tool-calling turn the model returns no text yet, so we summarize
the tool calls for the step output preview; otherwise the assistant
text is used.
"""
if self._ended or self.span is None:
return
if isinstance(content, str) and content.strip():
output: Any = content
elif tool_calls:
output = {"tool_calls": [tc.get("name") for tc in tool_calls]}
else:
output = content
self.update(input=messages, output=output)
def end(self, *, output: Any = None) -> None:
"""Stamp final output (optional) and close the step span. Idempotent."""
if self._ended:
return
self._ended = True
try:
if self.span is not None and output is not None:
self.span.update(output=output)
except Exception as exc: # pragma: no cover - best-effort telemetry
logger.debug("Failed to set Langfuse step output: %s", exc)
try:
self._stack.close()
except Exception as exc: # pragma: no cover - best-effort telemetry
logger.debug("Failed to close Langfuse step span: %s", exc)
def start_langfuse_agent_step(
name: str, telemetry: LLMTelemetryContext | None
) -> LangfuseAgentStep | None:
"""Open a per-iteration step span, imperatively. Returns ``None`` when
Langfuse is disabled or there's no ``run_id`` (no agent run to nest under).
"""
if not settings.LANGFUSE_PUBLIC_KEY or telemetry is None or not telemetry.run_id:
return None
stack = ExitStack()
try:
from langfuse import get_client
update_kwargs: dict[str, Any] = {
"metadata": {
"namespace": settings.NAMESPACE,
"provider": provider,
"model": model,
}
}
if name is not None:
update_kwargs["name"] = name
get_client().update_current_span(**update_kwargs)
span = stack.enter_context(
get_client().start_as_current_observation(
as_type="span", name=name, metadata=_step_metadata(telemetry)
)
)
except Exception as exc: # pragma: no cover - best-effort telemetry
logger.debug("Failed to update Langfuse span metadata: %s", exc)
logger.debug("Failed to open Langfuse agent step: %s", exc)
stack.close()
return None
return LangfuseAgentStep(span=span, _stack=stack)
@dataclass(frozen=True)
@ -231,6 +477,9 @@ def resolve_backend_for_plan(plan: AttemptPlan) -> Any:
__all__ = [
"AttemptPlan",
"LangfuseAgentRun",
"LangfuseAgentStep",
"annotate_current_langfuse_trace",
"current_attempt",
"effective_config_for_call",
"effective_temperature",
@ -238,5 +487,6 @@ __all__ = [
"resolve_backend_for_plan",
"resolve_runtime_model_config",
"select_model_config_for_attempt",
"update_current_langfuse_observation",
"start_langfuse_agent_run",
"start_langfuse_agent_step",
]

439
src/llm/tool_loop.py
View File

@ -36,6 +36,7 @@ from .runtime import (
AttemptPlan,
current_attempt,
effective_temperature,
start_langfuse_agent_step,
)
from .types import (
HonchoLLMCallResponse,
@ -68,6 +69,17 @@ def _with_iteration_scope(
return wrapper
def _step_label(base: LLMTelemetryContext | None) -> str:
"""Stable per-agent step-span name, e.g. "Dialectic Agent step".
No step number — Langfuse aggregates by name; the index rides on the
``iteration`` metadata. The " step" suffix distinguishes it from the bare
agent name, which names the enclosing run trace (see
`start_langfuse_agent_run`).
"""
return f"{(base.track_name if base else None) or 'Agent'} step"
def _telemetry_for_iteration(
base: LLMTelemetryContext | None, iteration: int
) -> LLMTelemetryContext | None:
@ -79,17 +91,7 @@ def _telemetry_for_iteration(
"""
if base is None:
return None
return LLMTelemetryContext(
workspace_name=base.workspace_name,
call_purpose=base.call_purpose,
parent_category=base.parent_category,
run_id=base.run_id,
iteration=iteration,
observer=base.observer,
observed=base.observed,
peer_name=base.peer_name,
agent_type=base.agent_type,
)
return dataclasses.replace(base, iteration=iteration)
def _emit_agent_iteration(
@ -221,6 +223,13 @@ async def stream_final_response(
# value — telemetry can't tell the retry sequence apart.
stream_attempt = 0
# No ContextVar gymnastics around `_in_agent_run` here: the run handle
# is alive for the lifetime of the stream (owned by
# `StreamingResponseWithMetadata` and closed on drain), so this streamed
# generation correctly nests under the run span as the current OTel
# observation. The previous code had to flip `_in_agent_run` to escape
# the run; with imperative handles the run isn't going anywhere.
async def _setup_stream() -> AsyncIterator[HonchoLLMCallStreamChunk]:
nonlocal stream_attempt
stream_attempt += 1
@ -292,6 +301,7 @@ async def execute_tool_loop(
stream_final: bool = False,
iteration_callback: IterationCallback | None = None,
telemetry: LLMTelemetryContext | None = None,
langfuse_run_handle: Any | None = None,
) -> HonchoLLMCallResponse[Any] | StreamingResponseWithMetadata:
"""Run the iterative tool calling loop for agentic LLM interactions.
@ -337,195 +347,220 @@ async def execute_tool_loop(
effective_tool_choice = tool_choice
while iteration < max_tool_iterations:
# Reset attempt counter so each iteration starts with the primary provider.
current_attempt.set(1)
logger.debug(f"Tool execution iteration {iteration + 1}/{max_tool_iterations}")
if max_input_tokens is not None:
if count_message_tokens(conversation_messages) > max_input_tokens:
hit_input_token_cap = True
conversation_messages = truncate_messages_to_fit(
conversation_messages, max_input_tokens
)
async def _call_with_messages(
effective_tool_choice: str | dict[str, Any] | None = effective_tool_choice,
conversation_messages: list[dict[str, Any]] = conversation_messages,
iteration_for_call: int = iteration + 1,
) -> HonchoLLMCallResponse[Any]:
plan = get_attempt_plan()
return await honcho_llm_call_inner(
plan.provider,
plan.model,
prompt, # ignored when messages is passed
max_tokens,
response_model,
json_mode,
effective_temperature(temperature),
stop_seqs,
plan.reasoning_effort,
verbosity,
plan.thinking_budget_tokens,
stream=False,
client_override=plan.client,
tools=tools,
tool_choice=effective_tool_choice,
messages=conversation_messages,
selected_config=plan.selected_config,
plan=plan,
telemetry=_telemetry_for_iteration(telemetry, iteration_for_call),
)
if enable_retry:
call_func = retry(
stop=stop_after_attempt(retry_attempts),
wait=wait_exponential(multiplier=1, min=4, max=10),
before_sleep=before_retry_callback,
)(_call_with_messages)
else:
call_func = _call_with_messages
response = await call_func()
total_input_tokens += response.input_tokens
total_output_tokens += response.output_tokens
total_cache_creation_tokens += response.cache_creation_input_tokens
total_cache_read_tokens += response.cache_read_input_tokens
# emit one AgentIterationEvent per LLM response BEFORE the
# no-tool early return. The terminating iteration counts too — it has
# an empty tool_calls list and is essential for cost calibration.
_emit_agent_iteration(telemetry, iteration + 1, response)
if not response.tool_calls_made:
logger.debug("No tool calls in response, finishing")
if (
isinstance(response.content, str)
and not response.content.strip()
and empty_response_retries < 1
and iteration < max_tool_iterations - 1
):
empty_response_retries += 1
conversation_messages.append(
{
"role": "user",
"content": (
"Your last response was empty. Provide a concise answer "
"to the original query using the available context."
),
}
)
iteration += 1
continue
if stream_final:
# Snapshot the plan that just succeeded — streaming retries
# pin to this exact client/model so we don't bounce back to
# primary after the tool loop settled on fallback.
winning_plan = get_attempt_plan()
stream = stream_final_response(
winning_plan=winning_plan,
prompt=prompt,
max_tokens=max_tokens,
conversation_messages=conversation_messages,
response_model=response_model,
json_mode=json_mode,
temperature=temperature,
stop_seqs=stop_seqs,
verbosity=verbosity,
enable_retry=enable_retry,
retry_attempts=retry_attempts,
before_retry_callback=before_retry_callback,
telemetry=_telemetry_for_iteration(telemetry, iteration + 1),
)
return StreamingResponseWithMetadata(
stream=stream,
tool_calls_made=all_tool_calls,
input_tokens=total_input_tokens,
output_tokens=total_output_tokens,
cache_creation_input_tokens=total_cache_creation_tokens,
cache_read_input_tokens=total_cache_read_tokens,
thinking_content=response.thinking_content,
iterations=iteration + 1,
hit_input_token_cap=hit_input_token_cap,
)
response.tool_calls_made = all_tool_calls
response.input_tokens = total_input_tokens
response.output_tokens = total_output_tokens
response.cache_creation_input_tokens = total_cache_creation_tokens
response.cache_read_input_tokens = total_cache_read_tokens
response.iterations = iteration + 1
response.hit_input_token_cap = (
response.hit_input_token_cap or hit_input_token_cap
)
return response
current_provider = get_attempt_plan().provider
assistant_message = format_assistant_tool_message(
current_provider,
response.content,
response.tool_calls_made,
response.thinking_blocks,
response.reasoning_details,
step = start_langfuse_agent_step(
_step_label(telemetry),
_telemetry_for_iteration(telemetry, iteration + 1),
)
conversation_messages.append(assistant_message)
try:
# Reset attempt counter so each iteration starts with the primary provider.
current_attempt.set(1)
logger.debug(
f"Tool execution iteration {iteration + 1}/{max_tool_iterations}"
)
# Telemetry context — 1-indexed iteration.
set_current_iteration(iteration + 1)
tool_results: list[dict[str, Any]] = []
for seq, tool_call in enumerate(response.tool_calls_made):
tool_name = tool_call["name"]
tool_input = tool_call["input"]
tool_id = tool_call.get("id", "")
logger.debug(f"Executing tool: {tool_name}")
# the executor closure reads these from
# ContextVars to populate AgentToolCallCompletedEvent. Set BEFORE
# the executor call so two calls to the same tool in one iteration
# get distinct seq values. Reset last-tool metadata so we never
# observe stale state from a prior call.
set_current_tool_call_seq(seq, tool_id or None)
set_last_tool_metadata({})
try:
tool_result = await tool_executor(tool_name, tool_input)
# Stash ToolResult.metadata on all_tool_calls so
# specialist rollups can read created/deleted observation
# counts without round-tripping through the event store.
tool_result_metadata = get_last_tool_metadata()
tool_results.append(
{
"tool_id": tool_id,
"tool_name": tool_name,
"result": tool_result,
}
)
all_tool_calls.append(
{
"tool_name": tool_name,
"tool_input": tool_input,
"tool_result": tool_result,
"tool_result_metadata": tool_result_metadata,
}
)
except Exception as e:
logger.error(f"Tool execution failed for {tool_name}: {e}")
tool_results.append(
{
"tool_id": tool_id,
"tool_name": tool_name,
"result": f"Error: {str(e)}",
"is_error": True,
}
if max_input_tokens is not None:
if count_message_tokens(conversation_messages) > max_input_tokens:
hit_input_token_cap = True
conversation_messages = truncate_messages_to_fit(
conversation_messages, max_input_tokens
)
append_tool_results(current_provider, tool_results, conversation_messages)
async def _call_with_messages(
tool_choice_for_call: str
| dict[str, Any]
| None = effective_tool_choice,
captured_messages: list[dict[str, Any]] = conversation_messages,
iteration_for_call: int = iteration + 1,
) -> HonchoLLMCallResponse[Any]:
plan = get_attempt_plan()
return await honcho_llm_call_inner(
plan.provider,
plan.model,
prompt, # ignored when messages is passed
max_tokens,
response_model,
json_mode,
effective_temperature(temperature),
stop_seqs,
plan.reasoning_effort,
verbosity,
plan.thinking_budget_tokens,
stream=False,
client_override=plan.client,
tools=tools,
tool_choice=tool_choice_for_call,
messages=captured_messages,
selected_config=plan.selected_config,
plan=plan,
telemetry=_telemetry_for_iteration(telemetry, iteration_for_call),
)
call_func: Callable[[], Awaitable[HonchoLLMCallResponse[Any]]]
if enable_retry:
call_func = retry(
stop=stop_after_attempt(retry_attempts),
wait=wait_exponential(multiplier=1, min=4, max=10),
before_sleep=before_retry_callback,
)(_call_with_messages)
else:
call_func = _call_with_messages # pyright: ignore[reportGeneralTypeIssues]
response = await call_func()
total_input_tokens += response.input_tokens
total_output_tokens += response.output_tokens
total_cache_creation_tokens += response.cache_creation_input_tokens
total_cache_read_tokens += response.cache_read_input_tokens
# emit one AgentIterationEvent per LLM response BEFORE the
# no-tool early return. The terminating iteration counts too — it has
# an empty tool_calls list and is essential for cost calibration.
_emit_agent_iteration(telemetry, iteration + 1, response)
# Step span is current again (the generation closed); stamp this
# turn's I/O so it isn't blank.
if step is not None:
step.annotate_io(
conversation_messages,
response.content,
response.tool_calls_made,
)
if not response.tool_calls_made:
logger.debug("No tool calls in response, finishing")
if (
isinstance(response.content, str)
and not response.content.strip()
and empty_response_retries < 1
and iteration < max_tool_iterations - 1
):
empty_response_retries += 1
conversation_messages.append(
{
"role": "user",
"content": (
"Your last response was empty. Provide a concise answer "
"to the original query using the available context."
),
}
)
iteration += 1
continue
if stream_final:
# Snapshot the plan that just succeeded — streaming retries
# pin to this exact client/model so we don't bounce back to
# primary after the tool loop settled on fallback.
winning_plan = get_attempt_plan()
stream = stream_final_response(
winning_plan=winning_plan,
prompt=prompt,
max_tokens=max_tokens,
conversation_messages=conversation_messages,
response_model=response_model,
json_mode=json_mode,
temperature=temperature,
stop_seqs=stop_seqs,
verbosity=verbosity,
enable_retry=enable_retry,
retry_attempts=retry_attempts,
before_retry_callback=before_retry_callback,
telemetry=_telemetry_for_iteration(telemetry, iteration + 1),
)
return StreamingResponseWithMetadata(
stream=stream,
tool_calls_made=all_tool_calls,
input_tokens=total_input_tokens,
output_tokens=total_output_tokens,
cache_creation_input_tokens=total_cache_creation_tokens,
cache_read_input_tokens=total_cache_read_tokens,
thinking_content=response.thinking_content,
iterations=iteration + 1,
hit_input_token_cap=hit_input_token_cap,
langfuse_run_handle=langfuse_run_handle,
)
response.tool_calls_made = all_tool_calls
response.input_tokens = total_input_tokens
response.output_tokens = total_output_tokens
response.cache_creation_input_tokens = total_cache_creation_tokens
response.cache_read_input_tokens = total_cache_read_tokens
response.iterations = iteration + 1
response.hit_input_token_cap = (
response.hit_input_token_cap or hit_input_token_cap
)
return response
current_provider = get_attempt_plan().provider
assistant_message = format_assistant_tool_message(
current_provider,
response.content,
response.tool_calls_made,
response.thinking_blocks,
response.reasoning_details,
)
conversation_messages.append(assistant_message)
# Telemetry context — 1-indexed iteration.
set_current_iteration(iteration + 1)
tool_results: list[dict[str, Any]] = []
for seq, tool_call in enumerate(response.tool_calls_made):
tool_name = tool_call["name"]
tool_input = tool_call["input"]
tool_id = tool_call.get("id", "")
logger.debug(f"Executing tool: {tool_name}")
# the executor closure reads these from
# ContextVars to populate AgentToolCallCompletedEvent. Set BEFORE
# the executor call so two calls to the same tool in one iteration
# get distinct seq values. Reset last-tool metadata so we never
# observe stale state from a prior call.
set_current_tool_call_seq(seq, tool_id or None)
set_last_tool_metadata({})
try:
tool_result = await tool_executor(tool_name, tool_input)
# Stash ToolResult.metadata on all_tool_calls so
# specialist rollups can read created/deleted observation
# counts without round-tripping through the event store.
tool_result_metadata = get_last_tool_metadata()
tool_results.append(
{
"tool_id": tool_id,
"tool_name": tool_name,
"result": tool_result,
}
)
all_tool_calls.append(
{
"tool_name": tool_name,
"tool_input": tool_input,
"tool_result": tool_result,
"tool_result_metadata": tool_result_metadata,
}
)
except Exception as e:
logger.error(f"Tool execution failed for {tool_name}: {e}")
tool_results.append(
{
"tool_id": tool_id,
"tool_name": tool_name,
"result": f"Error: {str(e)}",
"is_error": True,
}
)
append_tool_results(current_provider, tool_results, conversation_messages)
finally:
if step is not None:
step.end()
# Between-turn bookkeeping lives outside the step span — the span
# scopes the LLM call + its tools, not the iteration accounting.
if iteration_callback is not None:
try:
iteration_data = IterationData(
@ -602,6 +637,7 @@ async def execute_tool_loop(
thinking_content=None,
iterations=iteration + 1,
hit_input_token_cap=hit_input_token_cap,
langfuse_run_handle=langfuse_run_handle,
)
current_attempt.set(1)
@ -639,7 +675,18 @@ async def execute_tool_loop(
else:
final_call_func = _final_call
final_response = await final_call_func()
# Step span around the synthesis call — same shape as in-loop iterations
# so the generation nests under the run root instead of dangling at the
# trace. Imperative pair with a try/finally for the .end().
synthesis_step = start_langfuse_agent_step(
_step_label(telemetry),
_telemetry_for_iteration(telemetry, synthesis_iteration),
)
try:
final_response = await final_call_func()
finally:
if synthesis_step is not None:
synthesis_step.end()
# emit the synthesis-call iteration event BEFORE merging cumulative
# totals onto final_response below — otherwise the event's per-iteration

56
src/llm/types.py
View File

@ -75,6 +75,12 @@ class LLMTelemetryContext:
# agent — dialectic/deduction/induction. Used by agent iteration
# event and tool call event.
agent_type: str | None = None
# Human-readable name for the Langfuse trace + per-call generation
# (e.g. "Dialectic Agent", "Minimal Deriver"). Sole home for this name —
# callers set it here; `honcho_llm_call` no longer takes a separate kwarg.
# Also used to label the sentry `ai_track` decorator and as the source for
# the run-level `langfuse_agent_run` label.
track_name: str | None = None
IterationCallback = Callable[[IterationData], None]
@ -134,6 +140,13 @@ class StreamingResponseWithMetadata:
reflects tool-loop output + final-stream output. Callers that read
`output_tokens` AFTER fully iterating the stream get the true total;
callers that read it before drain see only the tool-loop portion.
`langfuse_run_handle` (optional) is the run-level Langfuse span handle
transferred from `honcho_llm_call` when streaming. The wrapper owns it
after construction: on drain, the accumulated streamed text is stamped
as the run span's output and the span is closed. Without this transfer,
streaming traces would show blank output because the synchronous return
happens before any chunks arrive.
"""
_stream: AsyncIterator[HonchoLLMCallStreamChunk]
@ -145,6 +158,7 @@ class StreamingResponseWithMetadata:
thinking_content: str | None
iterations: int
hit_input_token_cap: bool
_langfuse_run_handle: Any | None
def __init__(
self,
@ -157,6 +171,7 @@ class StreamingResponseWithMetadata:
thinking_content: str | None = None,
iterations: int = 0,
hit_input_token_cap: bool = False,
langfuse_run_handle: Any | None = None,
):
self._stream = stream
self.tool_calls_made = tool_calls_made
@ -167,6 +182,7 @@ class StreamingResponseWithMetadata:
self.thinking_content = thinking_content
self.iterations = iterations
self.hit_input_token_cap = hit_input_token_cap
self._langfuse_run_handle = langfuse_run_handle
def __aiter__(self) -> AsyncIterator[HonchoLLMCallStreamChunk]:
# Wrap the underlying iterator to capture final-stream output_tokens
@ -180,20 +196,32 @@ class StreamingResponseWithMetadata:
self,
) -> AsyncIterator[HonchoLLMCallStreamChunk]:
final_stream_output_tokens = 0
async for chunk in self._stream:
if chunk.output_tokens is not None:
# Take the LATEST value, not the sum — providers report
# the cumulative usage in the final chunk, not deltas.
final_stream_output_tokens = chunk.output_tokens
yield chunk
# Stream drained — fold the final-stream output tokens into the
# tool-loop totals so DialecticCompletedEvent / downstream readers
# see the true cost.
if final_stream_output_tokens > 0:
self.output_tokens += final_stream_output_tokens
async def __anext__(self) -> HonchoLLMCallStreamChunk:
return await self._stream.__anext__()
# Only accumulate when a Langfuse run handle is attached — for non-
# traced streams the buffer is dead weight.
accumulate = self._langfuse_run_handle is not None
accumulated_text: list[str] = []
try:
async for chunk in self._stream:
if chunk.output_tokens is not None:
# Take the LATEST value, not the sum — providers report
# the cumulative usage in the final chunk, not deltas.
final_stream_output_tokens = chunk.output_tokens
if accumulate and chunk.content:
accumulated_text.append(chunk.content)
yield chunk
# Stream drained — fold the final-stream output tokens into the
# tool-loop totals so DialecticCompletedEvent / downstream readers
# see the true cost.
if final_stream_output_tokens > 0:
self.output_tokens += final_stream_output_tokens
finally:
# Close the run span once, stamping the streamed text as its
# output. In `finally` so an early-exit caller still closes
# the span rather than leaking it.
handle = self._langfuse_run_handle
if handle is not None:
self._langfuse_run_handle = None
handle.end(output="".join(accumulated_text) or None)
__all__ = [

13
src/telemetry/__init__.py
View File

@ -52,8 +52,17 @@ async def shutdown_telemetry() -> None:
This should be called during application shutdown to ensure:
- CloudEvents buffer is flushed
- Langfuse's buffered observations are flushed (its background exporter and
atexit hook don't fire reliably on SIGTERM)
"""
from src.telemetry.events import shutdown_telemetry_events
from src.telemetry.logging import flush_langfuse
# Shutdown CloudEvents emitter (flushes buffer)
await shutdown_telemetry_events()
# Flush Langfuse even if the CloudEvents shutdown raises, so the final
# batch of spans isn't dropped on a noisy shutdown.
try:
# Shutdown CloudEvents emitter (flushes buffer)
await shutdown_telemetry_events()
finally:
# Flush any buffered Langfuse spans before the process exits.
flush_langfuse()

46
src/telemetry/logging.py
View File

@ -5,9 +5,10 @@ and a conditional observe decorator that only applies when Langfuse is configure
"""
import datetime
import logging
from collections import OrderedDict
from collections.abc import Callable
from typing import ParamSpec, TypeVar, overload
from typing import Literal, ParamSpec, TypeVar, overload
from fastapi import Request
from langfuse import observe
@ -24,6 +25,8 @@ from src.utils.representation import (
Representation,
)
logger = logging.getLogger(__name__)
# Global console instance for consistent formatting
console = Console(markup=True)
@ -32,6 +35,21 @@ COLLECT_METRICS_LOCAL = settings.COLLECT_METRICS_LOCAL
P = ParamSpec("P")
R = TypeVar("R")
# Langfuse observation types accepted by `@observe(as_type=...)`. Mirrors the
# literal union the SDK exposes; kept local so callers don't import langfuse
# internals just to name an observation type.
ObserveAsType = Literal[
"generation",
"embedding",
"span",
"agent",
"tool",
"chain",
"retriever",
"evaluator",
"guardrail",
]
@overload
def conditional_observe(
@ -42,7 +60,8 @@ def conditional_observe(
@overload
def conditional_observe(
*,
name: str,
name: str | None = None,
as_type: ObserveAsType | None = None,
) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
@ -50,17 +69,20 @@ def conditional_observe(
func: Callable[P, R] | None = None,
*,
name: str | None = None,
as_type: ObserveAsType | None = None,
) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
"""
Conditionally apply the @observe decorator only when LANGFUSE_PUBLIC_KEY is present.
Can be used in two ways:
1. As a decorator: @conditional_observe
2. As a decorator factory: @conditional_observe(name="...")
2. As a decorator factory: @conditional_observe(name="...", as_type="generation")
Args:
func: The function to potentially decorate (when used as @conditional_observe)
name: Optional name for the observation (when used as @conditional_observe(name="..."))
as_type: Optional Langfuse observation type (e.g. "generation", "tool"). When
omitted, Langfuse infers a default span.
Returns:
The decorated function if Langfuse is configured, otherwise the original function
@ -69,6 +91,8 @@ def conditional_observe(
def decorator(f: Callable[P, R]) -> Callable[P, R]:
if settings.LANGFUSE_PUBLIC_KEY:
observe_name = name if name is not None else f.__name__
if as_type is not None:
return observe(name=observe_name, as_type=as_type)(f)
return observe(name=observe_name)(f)
else:
return f
@ -81,6 +105,22 @@ def conditional_observe(
return decorator
def flush_langfuse() -> None:
"""Flush buffered Langfuse spans on shutdown.
The SDK's background timer/atexit hook don't fire reliably on SIGTERM, so
the final batch is dropped without this. No-op when Langfuse is unconfigured.
"""
if not settings.LANGFUSE_PUBLIC_KEY:
return
try:
from langfuse import get_client
get_client().flush()
except Exception:
logger.debug("Failed to flush Langfuse on shutdown", exc_info=True)
# Bounded OrderedDict for accumulated metrics to prevent memory leaks.
# If an exception occurs between accumulate_metric() and log_performance_metrics(),
# the metrics would stay in memory forever. Using OrderedDict allows us to evict

37
src/utils/agent_tools.py
View File

@ -2403,6 +2403,10 @@ async def create_tool_executor(
metadata: dict[str, Any] = {}
is_error: bool = False
# Langfuse tool observation; auto-parents under the active step span.
# Closed in the finally below with output + level.
tool_obs = _begin_tool_observation(tool_name, tool_input)
try:
handler = _TOOL_HANDLERS.get(tool_name)
if handler:
@ -2481,11 +2485,44 @@ async def create_tool_executor(
provider_tool_call_id=get_current_provider_tool_call_id(),
)
_finish_tool_observation(tool_obs, result_str, is_error)
return result_str
return execute_tool
def _begin_tool_observation(tool_name: str, tool_input: dict[str, Any]) -> Any:
"""Open a non-current Langfuse "tool" observation for one tool execution.
Auto-parents under the active step span (else standalone). Returns a handle
(closed by `_finish_tool_observation`) or None when disabled/setup fails.
All tools are ``as_type="tool"`` — they share one generic dispatcher.
"""
if not settings.LANGFUSE_PUBLIC_KEY:
return None
try:
from langfuse import get_client
return get_client().start_observation(
as_type="tool", name=tool_name, input=tool_input
)
except Exception: # pragma: no cover - best-effort telemetry
logger.debug("Failed to open Langfuse tool observation", exc_info=True)
return None
def _finish_tool_observation(tool_obs: Any, result_str: str, is_error: bool) -> None:
"""Close a Langfuse tool observation opened by `_begin_tool_observation`."""
if tool_obs is None:
return
try:
tool_obs.update(output=result_str, level="ERROR" if is_error else None)
tool_obs.end()
except Exception: # pragma: no cover - best-effort telemetry
logger.debug("Failed to close Langfuse tool observation", exc_info=True)
def _emit_agent_tool_call_completed(
*,
ctx: "ToolContext",

9
tests/conftest.py
View File

@ -34,10 +34,17 @@ from src.config import settings
from src.db import Base
from src.dependencies import get_db, get_read_db
from src.exceptions import HonchoException
from src.main import app
from src.models import Peer, Workspace
from src.security import JWTParams, create_admin_jwt, create_jwt
# Disable Langfuse for the whole suite before importing src.main: @conditional_observe
# binds to settings.LANGFUSE_PUBLIC_KEY at import time, so blanking it here keeps mocked
# test calls from emitting traces to a configured Langfuse backend. Tests that exercise
# Langfuse patch settings.LANGFUSE_PUBLIC_KEY themselves.
settings.LANGFUSE_PUBLIC_KEY = None
from src.main import app # noqa: E402
# Create a custom handler that doesn't get closed prematurely
class TestHandler(logging.Handler):

499
tests/llm/test_langfuse_trace_annotation.py Normal file
View File

@ -0,0 +1,499 @@
# pyright: reportPrivateUsage=false, reportUnusedParameter=false
"""Tests for the Langfuse session/trace wiring in `src/llm/runtime.py`.
One agentic run = one trace; `session_id = run_id` (globally unique, so it's
conflict-free across tenants — unlike the Honcho session name). The run handle
(`start_langfuse_agent_run`) opens an `as_type="span"` root and keeps it
current via an ``ExitStack`` until `.end()`. Step spans + nested generations
nest under the run while it's open. A single-call agent (deriver, summarizer)
gets no run handle and self-stamps its lone generation as the trace root.
Disabled (no LANGFUSE_PUBLIC_KEY) → no Langfuse calls at all.
"""
from __future__ import annotations
import contextlib
from typing import Any
import pytest
from src.config import settings
from src.llm import runtime
from src.llm.types import LLMTelemetryContext
@pytest.fixture
def capture_propagate(monkeypatch: pytest.MonkeyPatch):
"""Stub `langfuse.propagate_attributes`, capturing the kwargs it's called with.
Returns a dict that's empty until propagate_attributes is invoked.
"""
captured: dict[str, Any] = {}
@contextlib.contextmanager
def fake_propagate(**kwargs: Any):
captured.clear()
captured.update(kwargs)
yield
import langfuse
monkeypatch.setattr(langfuse, "propagate_attributes", fake_propagate)
return captured
@pytest.fixture
def langfuse_client(monkeypatch: pytest.MonkeyPatch):
"""Stub `langfuse.get_client()`, capturing observation/generation/span calls.
Returns ``{"observation", "generation", "span", "run_span"}`` — each sub-dict
stays empty until the corresponding call is made:
- ``observation``: the span opened via `start_as_current_observation`
(run root or step span).
- ``generation``: the rename via `update_current_generation`.
- ``span``: the step I/O via `update_current_span`.
- ``run_span``: I/O written onto the span object handed back by
`start_as_current_observation` (merged, mirroring Langfuse's update
semantics).
"""
captured: dict[str, dict[str, Any]] = {
"observation": {},
"generation": {},
"span": {},
"run_span": {},
}
class FakeSpan:
def update(self, **kwargs: Any) -> None:
# Merge (don't clear): Langfuse's span.update accumulates, so input
# set at run start and output set at run end coexist.
captured["run_span"].update(kwargs)
@contextlib.contextmanager
def fake_observation(**kwargs: Any):
captured["observation"].clear()
captured["observation"].update(kwargs)
yield FakeSpan()
class FakeClient:
def start_as_current_observation(self, **kwargs: Any):
return fake_observation(**kwargs)
def update_current_generation(self, **kwargs: Any) -> None:
captured["generation"].clear()
captured["generation"].update(kwargs)
def update_current_span(self, **kwargs: Any) -> None:
captured["span"].clear()
captured["span"].update(kwargs)
import langfuse
monkeypatch.setattr(langfuse, "get_client", lambda: FakeClient())
return captured
@pytest.fixture
def langfuse_enabled(monkeypatch: pytest.MonkeyPatch):
"""Turn the integration on with a known NAMESPACE (the tenant / user_id)."""
monkeypatch.setattr(settings, "LANGFUSE_PUBLIC_KEY", "pk-test")
monkeypatch.setattr(settings, "NAMESPACE", "acme-tenant")
@contextlib.contextmanager
def _inside_agent_run():
"""Set the `_in_agent_run` ContextVar for the body, resetting it after.
Simulates execution nested inside a run handle without opening one (which
would itself call propagate_attributes and pollute the capture).
"""
token = runtime._in_agent_run.set(True)
try:
yield
finally:
runtime._in_agent_run.reset(token)
class TestAnnotateDisabled:
def test_noop_when_key_unset(
self, monkeypatch: pytest.MonkeyPatch, capture_propagate: dict[str, Any]
):
monkeypatch.setattr(settings, "LANGFUSE_PUBLIC_KEY", None)
runtime.annotate_current_langfuse_trace(
"anthropic",
"claude-x",
telemetry=LLMTelemetryContext(run_id="run-abc"),
)
assert capture_propagate == {}
class TestAnnotateInsideRun:
"""A generation nested inside an active run handle: the run owns the trace
attrs, so this call must NOT propagate. It still stamps model + per-step
metadata + name on the generation (the multi-turn regression fix —
provider/model used to be dropped on every iteration after the first)."""
def test_nested_generation_does_not_propagate(
self,
langfuse_enabled: None,
capture_propagate: dict[str, Any],
langfuse_client: dict[str, dict[str, Any]],
):
telemetry = LLMTelemetryContext(
workspace_name="ws1",
call_purpose="dialectic.answer",
agent_type="dialectic",
run_id="run-abc",
iteration=2,
peer_name="alice",
track_name="Dialectic Agent",
)
with _inside_agent_run():
runtime.annotate_current_langfuse_trace(
"anthropic", "claude-x", telemetry=telemetry
)
# Run handle owns user_id/session_id/trace_name — re-propagating here
# would clobber the run's session, so we don't propagate at all.
assert capture_propagate == {}
# Per-call generation: name + model + step metadata stamped every
# iteration (formerly only name was stamped, dropping provider/model).
gen = langfuse_client["generation"]
assert gen["name"] == "Dialectic Agent LLM call"
assert gen["model"] == "claude-x"
assert gen["metadata"]["provider"] == "anthropic"
assert gen["metadata"]["model"] == "claude-x"
assert gen["metadata"]["iteration"] == "2"
assert gen["metadata"]["agent_type"] == "dialectic"
class TestAnnotateOwnTraceRoot:
"""A generation that IS its own trace root stamps the trace attributes:
single calls (no run_id → no session) and the no-telemetry case."""
def test_single_call_has_no_session_and_names_trace(
self,
langfuse_enabled: None,
capture_propagate: dict[str, Any],
langfuse_client: dict[str, dict[str, Any]],
):
telemetry = LLMTelemetryContext(
workspace_name="ws1",
call_purpose="deriver.representation",
observed="bob",
track_name="Minimal Deriver",
)
runtime.annotate_current_langfuse_trace(
"gemini", "gemini-x", telemetry=telemetry
)
assert capture_propagate["session_id"] is None
assert capture_propagate["user_id"] == "acme-tenant"
# Single-call: this generation IS the trace root, so it names the trace.
assert capture_propagate["trace_name"] == "Minimal Deriver"
assert capture_propagate["metadata"]["observed"] == "bob"
# The generation observation is still named per agent+action.
assert langfuse_client["generation"]["name"] == "Minimal Deriver LLM call"
def test_no_telemetry_still_stamps_user_id(
self,
langfuse_enabled: None,
capture_propagate: dict[str, Any],
langfuse_client: dict[str, dict[str, Any]],
):
runtime.annotate_current_langfuse_trace("openai", "gpt-x", telemetry=None)
assert capture_propagate["user_id"] == "acme-tenant"
assert capture_propagate["session_id"] is None
assert capture_propagate["trace_name"] is None
assert capture_propagate["metadata"]["provider"] == "openai"
# No telemetry → no per-agent generation name, but provider/model still set.
gen = langfuse_client["generation"]
assert gen["name"] is None
assert gen["model"] == "gpt-x"
class TestAgentRun:
"""`start_langfuse_agent_run` returns an imperative handle: opens an
``as_type="span"`` root, stamps ``session_id = run_id`` via
``propagate_attributes``, and keeps the span open until ``.end()``.
Only fires for multi-turn runs (run_id present)."""
def test_noop_when_disabled(
self,
monkeypatch: pytest.MonkeyPatch,
langfuse_client: dict[str, dict[str, Any]],
capture_propagate: dict[str, Any],
):
monkeypatch.setattr(settings, "LANGFUSE_PUBLIC_KEY", None)
handle = runtime.start_langfuse_agent_run(
"Dialectic Agent", LLMTelemetryContext(run_id="r1")
)
assert handle is None
assert langfuse_client["observation"] == {}
assert capture_propagate == {}
def test_noop_without_run_id(
self,
langfuse_enabled: None,
langfuse_client: dict[str, dict[str, Any]],
capture_propagate: dict[str, Any],
):
# Single-call agents (deriver/summarizer) have no run_id → no run root,
# so their LLM calls stay standalone, sessionless traces.
handle = runtime.start_langfuse_agent_run(
"Minimal Deriver", LLMTelemetryContext(workspace_name="ws1")
)
assert handle is None
assert langfuse_client["observation"] == {}
assert capture_propagate == {}
def test_noop_without_telemetry(
self,
langfuse_enabled: None,
langfuse_client: dict[str, dict[str, Any]],
):
handle = runtime.start_langfuse_agent_run("anything", None)
assert handle is None
assert langfuse_client["observation"] == {}
def test_opens_run_root_and_owns_trace_attrs(
self,
langfuse_enabled: None,
langfuse_client: dict[str, dict[str, Any]],
capture_propagate: dict[str, Any],
):
tele = LLMTelemetryContext(
workspace_name="ws1",
call_purpose="dialectic.answer",
agent_type="dialectic",
run_id="run-abc",
observed="bob",
track_name="Dialectic Agent",
)
handle = runtime.start_langfuse_agent_run("Dialectic Agent", tele)
assert handle is not None
try:
# The run IS the trace root: an as_type="span" observation whose
# name is STABLE (no step number) so Langfuse aggregates by name.
observation = langfuse_client["observation"]
assert observation["as_type"] == "span"
assert observation["name"] == "Dialectic Agent"
# Trace grouping: one Langfuse session per run, drillable per tenant.
assert capture_propagate["session_id"] == "run-abc"
assert capture_propagate["user_id"] == "acme-tenant"
assert capture_propagate["trace_name"] == "Dialectic Agent"
md = capture_propagate["metadata"]
assert md["workspace_name"] == "ws1"
assert md["agent_type"] == "dialectic"
assert md["observed"] == "bob"
# Honcho's Session is deliberately NOT the grouping key.
assert "honcho_session_id" not in md
finally:
handle.end()
def test_marks_in_agent_run_for_nested_calls(
self,
langfuse_enabled: None,
langfuse_client: dict[str, dict[str, Any]],
capture_propagate: dict[str, Any],
):
# While the run handle is live, nested generations see _in_agent_run
# set so they stay silent (the run owns the trace attrs); end() resets it.
assert runtime._in_agent_run.get() is False
handle = runtime.start_langfuse_agent_run(
"Dialectic Agent",
LLMTelemetryContext(run_id="r1", track_name="Dialectic Agent"),
)
assert handle is not None
assert runtime._in_agent_run.get() is True
handle.end()
assert runtime._in_agent_run.get() is False
def test_end_is_idempotent(
self,
langfuse_enabled: None,
langfuse_client: dict[str, dict[str, Any]],
capture_propagate: dict[str, Any],
):
# The streaming wrapper may call .end() after the api.py finally already
# called it (or vice-versa); the handle has to tolerate that.
handle = runtime.start_langfuse_agent_run(
"Dialectic Agent", LLMTelemetryContext(run_id="r1")
)
assert handle is not None
handle.end()
handle.end() # must not raise
class TestAgentRunIO:
"""The run handle exposes `.update(input=..., output=...)` for stamping
the run-root span — the trace's input/output preview in the Langfuse UI.
A second call merges into the first (Langfuse's update semantics)."""
def test_sets_input_then_output_on_handle(
self,
langfuse_enabled: None,
langfuse_client: dict[str, dict[str, Any]],
capture_propagate: dict[str, Any],
):
messages = [{"role": "user", "content": "How many coffees?"}]
handle = runtime.start_langfuse_agent_run(
"Dialectic Agent", LLMTelemetryContext(run_id="run-abc")
)
assert handle is not None
try:
handle.update(input=messages)
finally:
handle.end(output="You bought 4 coffees.")
assert langfuse_client["run_span"]["input"] == messages
assert langfuse_client["run_span"]["output"] == "You bought 4 coffees."
def test_end_without_output_leaves_output_unset(
self,
langfuse_enabled: None,
langfuse_client: dict[str, dict[str, Any]],
capture_propagate: dict[str, Any],
):
handle = runtime.start_langfuse_agent_run(
"Dialectic Agent", LLMTelemetryContext(run_id="run-abc")
)
assert handle is not None
handle.update(input=[{"role": "user"}])
handle.end()
assert "input" in langfuse_client["run_span"]
# Only input was passed → output is not written (so it isn't blanked).
assert "output" not in langfuse_client["run_span"]
class TestAgentStep:
"""`start_langfuse_agent_step` opens a per-iteration child span under the
run root (one reasoning turn). Unlike the run handle, it does NOT touch
trace attributes."""
def test_noop_when_disabled(
self,
monkeypatch: pytest.MonkeyPatch,
langfuse_client: dict[str, dict[str, Any]],
capture_propagate: dict[str, Any],
):
monkeypatch.setattr(settings, "LANGFUSE_PUBLIC_KEY", None)
step = runtime.start_langfuse_agent_step(
"Dialectic Agent step", LLMTelemetryContext(run_id="r1")
)
assert step is None
assert langfuse_client["observation"] == {}
assert capture_propagate == {}
def test_noop_without_run_id(
self,
langfuse_enabled: None,
langfuse_client: dict[str, dict[str, Any]],
capture_propagate: dict[str, Any],
):
step = runtime.start_langfuse_agent_step(
"Minimal Deriver step", LLMTelemetryContext(workspace_name="ws1")
)
assert step is None
assert langfuse_client["observation"] == {}
assert capture_propagate == {}
def test_noop_without_telemetry(
self,
langfuse_enabled: None,
langfuse_client: dict[str, dict[str, Any]],
):
step = runtime.start_langfuse_agent_step("anything", None)
assert step is None
assert langfuse_client["observation"] == {}
def test_opens_child_span_without_propagating(
self,
langfuse_enabled: None,
langfuse_client: dict[str, dict[str, Any]],
capture_propagate: dict[str, Any],
):
tele = LLMTelemetryContext(
workspace_name="ws1",
agent_type="dialectic",
run_id="run-abc",
iteration=2,
observed="bob",
track_name="Dialectic Agent",
)
step = runtime.start_langfuse_agent_step("Dialectic Agent step", tele)
assert step is not None
try:
observation = langfuse_client["observation"]
assert observation["as_type"] == "span"
assert observation["name"] == "Dialectic Agent step"
# The per-step index rides on the span's metadata (str-coerced).
assert observation["metadata"]["iteration"] == "2"
assert observation["metadata"]["observed"] == "bob"
# The run root owns the trace attrs — the step must NOT propagate.
assert capture_propagate == {}
finally:
step.end()
class TestStepIO:
"""`step.annotate_io` stamps this turn's I/O on the step span — without it,
only the nested generation would carry I/O and the step would show blank."""
def test_text_answer_sets_input_and_output(
self,
langfuse_enabled: None,
langfuse_client: dict[str, dict[str, Any]],
):
messages = [{"role": "user", "content": "What is the user's name?"}]
step = runtime.start_langfuse_agent_step(
"Dialectic Agent step", LLMTelemetryContext(run_id="run-abc")
)
assert step is not None
try:
step.annotate_io(messages, "The user's name is Jordan.", [])
finally:
step.end()
# The step span's input/output is stamped via the handle's underlying
# span.update() — captured on the run_span fixture key.
assert langfuse_client["run_span"]["input"] == messages
assert langfuse_client["run_span"]["output"] == "The user's name is Jordan."
def test_tool_calling_turn_summarizes_tools_as_output(
self,
langfuse_enabled: None,
langfuse_client: dict[str, dict[str, Any]],
):
# A tool-calling turn has no text yet — the step's "output" is the set
# of tools it chose, by name (per-tool I/O lives on the tool children).
step = runtime.start_langfuse_agent_step(
"Dialectic Agent step", LLMTelemetryContext(run_id="run-abc")
)
assert step is not None
try:
step.annotate_io(
[{"role": "user", "content": "how many coffees?"}],
"",
[{"name": "grep_messages"}, {"name": "search_memory"}],
)
finally:
step.end()
assert langfuse_client["run_span"]["output"] == {
"tool_calls": ["grep_messages", "search_memory"]
}

83
tests/llm/test_telemetry_llm_call.py
View File

@ -13,7 +13,7 @@ Targets:
from __future__ import annotations
from typing import Any
from typing import Any, cast
from unittest.mock import AsyncMock, patch
import pytest
@ -684,3 +684,84 @@ class TestStreamingResponseTokenWriteBack:
# And we yielded every chunk to the caller — the wrapper is a
# passthrough, not a sink.
assert len(chunks) == 3
class TestStreamingResponseRunHandleClose:
"""When a `langfuse_run_handle` is transferred to the streaming wrapper,
the wrapper owns it: the accumulated streamed text is stamped as the run
span's output and the span is closed exactly once when the stream drains.
The close lives in a `finally`, so an early-exit caller still closes the
span rather than leaking it.
"""
class _FakeRunHandle:
def __init__(self) -> None:
self.end_calls: list[Any] = []
def end(self, *, output: Any = None) -> None:
self.end_calls.append(output)
@staticmethod
async def _fake_stream() -> Any:
from src.llm.types import HonchoLLMCallStreamChunk
yield HonchoLLMCallStreamChunk(content="hel", output_tokens=None)
yield HonchoLLMCallStreamChunk(content="lo", output_tokens=None)
yield HonchoLLMCallStreamChunk(content="", is_done=True, output_tokens=7)
@pytest.mark.asyncio
async def test_full_drain_stamps_output_and_closes_once(self):
from src.llm.types import StreamingResponseWithMetadata
handle = self._FakeRunHandle()
wrapper = StreamingResponseWithMetadata(
stream=self._fake_stream(),
tool_calls_made=[],
input_tokens=0,
output_tokens=0,
cache_creation_input_tokens=0,
cache_read_input_tokens=0,
langfuse_run_handle=handle,
)
async for _ in wrapper:
pass
# Closed exactly once, with the concatenated streamed text as output.
assert handle.end_calls == ["hello"]
# Ownership released so a second drain can't double-close.
assert wrapper._langfuse_run_handle is None
@pytest.mark.asyncio
async def test_abandoned_stream_still_closes_via_finally(self):
from src.llm.types import StreamingResponseWithMetadata
handle = self._FakeRunHandle()
wrapper = StreamingResponseWithMetadata(
stream=self._fake_stream(),
tool_calls_made=[],
input_tokens=0,
output_tokens=0,
cache_creation_input_tokens=0,
cache_read_input_tokens=0,
langfuse_run_handle=handle,
)
# Consume one chunk, then abandon the stream. `aclose()` is what the
# runtime/GC drives when a caller stops iterating early; it throws
# GeneratorExit at the suspended `yield`, firing the `finally`.
from collections.abc import AsyncGenerator
from src.llm.types import HonchoLLMCallStreamChunk
agen = cast(
"AsyncGenerator[HonchoLLMCallStreamChunk, None]", wrapper.__aiter__()
)
first = await agen.__anext__()
assert first.content == "hel"
await agen.aclose()
# Span closed once with only the text accumulated before abandonment —
# the span is closed, not leaked.
assert handle.end_calls == ["hel"]
assert wrapper._langfuse_run_handle is None

103
tests/unified/test_cases/dialectic_tool_calls.json Normal file
View File

@ -0,0 +1,103 @@
{
"description": "Force the dialectic agent to actually invoke tools (grep/search) by asking an enumeration+aggregation question that the prefetched first turn cannot answer in one shot. Used to verify Langfuse tool-call observations nest under the step span.",
"workspace_config": {},
"steps": [
{
"step_type": "create_session",
"session_id": "tool_calls_test",
"config": {
"reasoning": {
"enabled": false
}
},
"peer_configs": {
"user": {
"observe_me": true,
"observe_others": false
},
"assistant": {
"observe_me": false,
"observe_others": true
}
}
},
{
"step_type": "add_messages",
"session_id": "tool_calls_test",
"messages": [
{
"peer_id": "user",
"content": "Monday I grabbed a $5 latte at Starbucks before my standup.",
"created_at": "2024-03-04T08:30:00"
},
{
"peer_id": "assistant",
"content": "Nice, a classic way to start the week.",
"created_at": "2024-03-04T08:31:00"
},
{
"peer_id": "user",
"content": "Tuesday I tried a $4 cold brew from Blue Bottle, really smooth.",
"created_at": "2024-03-05T09:15:00"
},
{
"peer_id": "assistant",
"content": "Blue Bottle makes a solid cold brew.",
"created_at": "2024-03-05T09:16:00"
},
{
"peer_id": "user",
"content": "Wednesday was a $6 oat-milk mocha at a little place downtown.",
"created_at": "2024-03-06T08:45:00"
},
{
"peer_id": "assistant",
"content": "Oat milk mochas are underrated.",
"created_at": "2024-03-06T08:46:00"
},
{
"peer_id": "user",
"content": "Thursday I skipped coffee and just had tea at home.",
"created_at": "2024-03-07T08:20:00"
},
{
"peer_id": "assistant",
"content": "A calm morning, sounds good.",
"created_at": "2024-03-07T08:21:00"
},
{
"peer_id": "user",
"content": "Friday I splurged on a $7 pour-over at the roastery near the office.",
"created_at": "2024-03-08T08:50:00"
},
{
"peer_id": "assistant",
"content": "Ending the week strong!",
"created_at": "2024-03-08T08:51:00"
}
]
},
{
"step_type": "wait",
"target": "queue_empty",
"timeout": 180,
"flush": true
},
{
"step_type": "query",
"description": "Global query (no session history) + empty prefetch (reasoning off) forces grep/search tool calls",
"target": "chat",
"observer_peer_id": "assistant",
"observed_peer_id": "user",
"reasoning_level": "max",
"input": "How many separate coffees did I buy this week, and exactly how much did I spend in total across all of them?",
"assertions": [
{
"assertion_type": "llm_judge",
"prompt": "Does the response state that the user bought 4 coffees totaling $22 (or correctly enumerate the $5, $4, $6, and $7 purchases)?",
"pass_if": true
}
]
}
]
}

87
tests/utils/test_clients.py
View File

@ -9,6 +9,7 @@ Tests cover:
- Provider-specific features
"""
import contextlib
from typing import Any
from unittest.mock import AsyncMock, Mock, patch
@ -37,6 +38,7 @@ from src.llm import (
honcho_llm_call,
honcho_llm_call_inner,
)
from src.llm.types import LLMTelemetryContext
class SampleTestModel(BaseModel):
@ -910,8 +912,10 @@ class TestMainLLMCallFunction:
assert response.content == "No retry response"
async def test_track_name_updates_langfuse_span_name(self):
"""track_name should rename the top-level Langfuse span."""
async def test_track_name_on_telemetry_names_langfuse_trace_and_generation(self):
"""track_name on telemetry should name the Langfuse trace + generation
per agent and stamp provider/model metadata (via propagate_attributes +
update_current_generation; see annotate_current_langfuse_trace)."""
mock_llm_client = AsyncMock(spec=AsyncAnthropic)
mock_response = Mock()
@ -921,11 +925,18 @@ class TestMainLLMCallFunction:
mock_llm_client.messages.create = AsyncMock(return_value=mock_response)
mock_langfuse_client = Mock()
captured: dict[str, Any] = {}
@contextlib.contextmanager
def fake_propagate(**kwargs: Any):
captured.update(kwargs)
yield
with (
patch.dict(CLIENTS, {"anthropic": mock_llm_client}),
patch.object(settings, "LANGFUSE_PUBLIC_KEY", "test-public-key"),
patch("langfuse.get_client", return_value=mock_langfuse_client),
patch("langfuse.propagate_attributes", fake_propagate),
):
response = await honcho_llm_call(
model_config=ConfiguredModelSettings(
@ -935,19 +946,75 @@ class TestMainLLMCallFunction:
prompt="Hello",
max_tokens=100,
enable_retry=False,
track_name="Dialectic Agent",
telemetry=LLMTelemetryContext(
workspace_name="ws1", track_name="Dialectic Agent"
),
)
assert response.content == "Named response"
mock_langfuse_client.update_current_span.assert_called_once_with(
name="Dialectic Agent",
metadata={
"namespace": settings.NAMESPACE,
"provider": "anthropic",
"model": "claude-4-sonnet",
},
# No run_id → this single call IS the trace root: it names the trace
# and stamps metadata via propagate_attributes...
assert captured["trace_name"] == "Dialectic Agent"
assert captured["session_id"] is None
assert captured["metadata"]["namespace"] == settings.NAMESPACE
assert captured["metadata"]["provider"] == "anthropic"
assert captured["metadata"]["model"] == "claude-4-sonnet"
# ...and the generation is named + carries per-call model/metadata.
mock_langfuse_client.update_current_generation.assert_called_once()
gen_kwargs = mock_langfuse_client.update_current_generation.call_args.kwargs
assert gen_kwargs["name"] == "Dialectic Agent LLM call"
assert gen_kwargs["model"] == "claude-4-sonnet"
assert gen_kwargs["metadata"]["provider"] == "anthropic"
async def test_no_telemetry_still_stamps_trace_without_name(self):
"""Without telemetry, propagate_attributes still fires with namespace
metadata, but the trace stays unnamed — track_name lives exclusively
on telemetry now. The per-call generation still gets model/metadata
stamped (the multi-turn-regression fix means we always stamp these)."""
mock_llm_client = AsyncMock(spec=AsyncAnthropic)
mock_response = Mock()
mock_response.content = [TextBlock(text="Unnamed response", type="text")]
mock_response.usage = Usage(input_tokens=5, output_tokens=5)
mock_response.stop_reason = "stop"
mock_llm_client.messages.create = AsyncMock(return_value=mock_response)
mock_langfuse_client = Mock()
captured: dict[str, Any] = {}
@contextlib.contextmanager
def fake_propagate(**kwargs: Any):
captured.update(kwargs)
yield
with (
patch.dict(CLIENTS, {"anthropic": mock_llm_client}),
patch.object(settings, "LANGFUSE_PUBLIC_KEY", "test-public-key"),
patch("langfuse.get_client", return_value=mock_langfuse_client),
patch("langfuse.propagate_attributes", fake_propagate),
):
response = await honcho_llm_call(
model_config=ConfiguredModelSettings(
model="claude-4-sonnet",
transport="anthropic",
),
prompt="Hello",
max_tokens=100,
enable_retry=False,
)
assert response.content == "Unnamed response"
assert captured["user_id"] == str(settings.NAMESPACE)
assert captured["trace_name"] is None
assert captured["metadata"]["namespace"] == settings.NAMESPACE
assert captured["metadata"]["provider"] == "anthropic"
# Generation gets model + metadata even without a track_name — only
# the name kwarg stays None.
mock_langfuse_client.update_current_generation.assert_called_once()
gen_kwargs = mock_langfuse_client.update_current_generation.call_args.kwargs
assert gen_kwargs["name"] is None
assert gen_kwargs["model"] == "claude-4-sonnet"
class TestEdgeCases:
"""Tests for edge cases and boundary conditions"""