Trace LLMP completions in Langfuse

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.

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

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
 
 

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/'.