From 7d4c047652b1e3a1f4d9ccfbf74d432e647ecac0 Mon Sep 17 00:00:00 2001 From: Adam Sorrenti Date: Mon, 3 Aug 2026 10:54:23 -0400 Subject: [PATCH] Trace LLMP completions in Langfuse Co-authored-by: Codex --- aieng-forecasting/README.md | 2 +- .../aieng/forecasting/methods/README.md | 2 +- .../methods/llm_processes/_client.py | 68 +++++++++++++++---- .../methods/llm_processes/test__client.py | 34 ++++++++++ 4 files changed, 92 insertions(+), 14 deletions(-) diff --git a/aieng-forecasting/README.md b/aieng-forecasting/README.md index 037c5735..f19edf22 100644 --- a/aieng-forecasting/README.md +++ b/aieng-forecasting/README.md @@ -32,7 +32,7 @@ pip install "aieng-forecasting[agentic]" Current extras: - `numerical` — Darts-based numerical predictors and related model dependencies -- `llm` — LiteLLM-based LLM-process predictors; Langfuse tracing via `langfuse_otel` +- `llm` — LiteLLM-based LLM-process predictors; each completion is recorded as a nested Langfuse Generation with its request messages and response - `agentic` — Google ADK runner (`AdkTextRunner`), generic agent factory (`build_adk_agent`), Track 1 predictor wrapper (`AgentPredictor`), structured agent output schemas, E2B code interpreter, and Langfuse / OpenInference tracing > **E2B setup:** the `agentic` extra requires a one-time sandbox image build. diff --git a/aieng-forecasting/aieng/forecasting/methods/README.md b/aieng-forecasting/aieng/forecasting/methods/README.md index d591f412..cda035a8 100644 --- a/aieng-forecasting/aieng/forecasting/methods/README.md +++ b/aieng-forecasting/aieng/forecasting/methods/README.md @@ -95,7 +95,7 @@ from aieng.forecasting.methods.agentic import ( | `llm_processes/sampled_trajectory.py` | `SampledTrajectoryLLMPredictor` | Samples full trajectories from an LLM, then computes empirical quantiles per horizon. Supports optional covariates: set `covariate_series_ids` to serialize labeled exogenous-series history into the prompt (Context-is-Key §5.4). | | `llm_processes/quantile_grid.py` | `QuantileGridLLMPredictor` | Asks an LLM for the standard quantile grid in one structured completion. | | `llm_processes/binary_probability.py` | `BinaryProbabilityLLMPredictor` | Direct elicitation of one calibrated event probability for binary tasks (Brier-scored), in one structured completion. | -| `llm_processes/categorical_probability.py` | `CategoricalProbabilityLLMPredictor` | Direct elicitation of a calibrated distribution over the task-declared ordered categories (RPS-scored); history serialized as category labels. | +| `llm_processes/categorical_probability.py` | `CategoricalProbabilityLLMPredictor` | Direct elicitation of a calibrated distribution over the task-declared ordered categories (RPS-scored); history serialized as category labels. Each completion appears as a nested Langfuse Generation, including its system/user messages and response. | | `llm_processes/point_intervals.py` | — | Placeholder for a compact point-plus-interval contract; may become configurable sparse quantile-grid elicitation. | ### Agentic diff --git a/aieng-forecasting/aieng/forecasting/methods/llm_processes/_client.py b/aieng-forecasting/aieng/forecasting/methods/llm_processes/_client.py index 24b91d72..562b1318 100644 --- a/aieng-forecasting/aieng/forecasting/methods/llm_processes/_client.py +++ b/aieng-forecasting/aieng/forecasting/methods/llm_processes/_client.py @@ -26,7 +26,8 @@ import os import warnings from concurrent.futures import ThreadPoolExecutor -from typing import Any, Callable, TypeVar +from contextlib import contextmanager +from typing import Any, Callable, Iterator, TypeVar from pydantic import BaseModel, ValidationError @@ -119,6 +120,40 @@ def trace_url_for(trace_id: str) -> str | None: return None +@contextmanager +def langfuse_generation(*, model: str, messages: list[dict[str, Any]]) -> Iterator[Any | None]: + """Create a best-effort Langfuse Generation for one LLM-process call. + + The LiteLLM OTEL callback is useful for standalone calls, but does not + reliably inherit the native Langfuse ``@observe`` context used by LLMP + predictors. Creating the Generation through the Langfuse client makes it + a child of that predictor trace and records the actual request messages. + """ + context: Any | None = None + try: + from langfuse import get_client # noqa: PLC0415 + + context = get_client().start_as_current_observation( + name="litellm.acompletion", + as_type="generation", + input=messages, + model=model, + ) + observation = context.__enter__() + except Exception: + logger.debug("Could not start Langfuse Generation for LLM-process call.", exc_info=True) + yield None + return + + try: + yield observation + finally: + try: + context.__exit__(None, None, None) + except Exception: + logger.debug("Could not close Langfuse Generation for LLM-process call.", exc_info=True) + + def set_current_trace_name(name: str) -> None: """Name the active Langfuse trace, if any, so it is identifiable in the UI. @@ -299,17 +334,26 @@ async def _one_completion_async( # models that don't support them (e.g. temperature on some o-series). kwargs["drop_params"] = True - resp = await litellm.acompletion(**kwargs) - cost = float(getattr(resp, "_hidden_params", {}).get("response_cost") or 0.0) - usage = getattr(resp, "usage", None) - in_tok = int(getattr(usage, "prompt_tokens", 0) or 0) if usage is not None else 0 - out_tok = int(getattr(usage, "completion_tokens", 0) or 0) if usage is not None else 0 - # Log full usage so we can see thinking-token breakdown when available. - # The proxy may populate completion_tokens_details.reasoning_tokens. - if usage is not None: - logger.debug("LLM usage: %s", vars(usage) if hasattr(usage, "__dict__") else usage) - raw = resp.choices[0].message.content - content = strip_markdown_fence(raw) if raw else raw + with langfuse_generation(model=str(kwargs["model"]), messages=messages) as generation: + resp = await litellm.acompletion(**kwargs) + cost = float(getattr(resp, "_hidden_params", {}).get("response_cost") or 0.0) + usage = getattr(resp, "usage", None) + in_tok = int(getattr(usage, "prompt_tokens", 0) or 0) if usage is not None else 0 + out_tok = int(getattr(usage, "completion_tokens", 0) or 0) if usage is not None else 0 + # Log full usage so we can see thinking-token breakdown when available. + # The proxy may populate completion_tokens_details.reasoning_tokens. + if usage is not None: + logger.debug("LLM usage: %s", vars(usage) if hasattr(usage, "__dict__") else usage) + raw = resp.choices[0].message.content + content = strip_markdown_fence(raw) if raw else raw + if generation is not None: + try: + generation.update( + output={"role": "assistant", "content": raw}, + usage_details={"input": in_tok, "output": out_tok}, + ) + except Exception: + logger.debug("Could not update Langfuse Generation for LLM-process call.", exc_info=True) return content, cost, in_tok, out_tok diff --git a/aieng-forecasting/tests/aieng/forecasting/methods/llm_processes/test__client.py b/aieng-forecasting/tests/aieng/forecasting/methods/llm_processes/test__client.py index d7154cc6..29335b1a 100644 --- a/aieng-forecasting/tests/aieng/forecasting/methods/llm_processes/test__client.py +++ b/aieng-forecasting/tests/aieng/forecasting/methods/llm_processes/test__client.py @@ -7,6 +7,7 @@ from __future__ import annotations +from contextlib import contextmanager from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -125,6 +126,39 @@ def _mock_litellm_response(content: str) -> MagicMock: _DUMMY_FORMAT = {"type": "json_schema", "json_schema": {"name": "x", "schema": {}, "strict": True}} +@pytest.mark.asyncio +async def test_completion_is_recorded_as_a_nested_langfuse_generation() -> None: + """LLMP calls explicitly record messages and output on the predictor trace.""" + generation = MagicMock() + + @contextmanager + def fake_generation(**kwargs): # type: ignore[no-untyped-def] + assert kwargs == {"model": "gemini-3-flash-preview", "messages": _DUMMY_MESSAGES} + yield generation + + with ( + patch( + "aieng.forecasting.methods.llm_processes._client.langfuse_generation", + side_effect=fake_generation, + ), + patch("litellm.acompletion", new=AsyncMock(return_value=_mock_litellm_response('{"cut": 0.2}'))), + ): + await _one_completion_async( + model="gemini-3-flash-preview", + messages=_DUMMY_MESSAGES, + response_format=_DUMMY_FORMAT, + temperature=1.0, + max_tokens=512, + timeout_s=30.0, + reasoning_effort=None, + ) + + generation.update.assert_called_once_with( + output={"role": "assistant", "content": '{"cut": 0.2}'}, + usage_details={"input": 0, "output": 0}, + ) + + @pytest.mark.asyncio async def test_proxy_path_prefixes_model_with_openai() -> None: """When api_base is set, the model is prefixed with 'openai/'.