Fix/langfuse llmp generation

aieng-forecasting/aieng/forecasting/langfuse_tracing.py

@@ -24,12 +24,11 @@ def _langfuse_credentials_present() -> bool:
 
 
 class _LangfuseTracingBootstrap:
-    """Registers LiteLLM + ADK exporters at most once per process."""
+    """Registers the Langfuse client and ADK instrumentation once per process."""
 
-    __slots__ = ("_google_adk_instrumented", "_langfuse_client_initialized", "_litellm_instrumented")
+    __slots__ = ("_google_adk_instrumented", "_langfuse_client_initialized")
 
     def __init__(self) -> None:
-        self._litellm_instrumented = False
         self._google_adk_instrumented = False
         self._langfuse_client_initialized = False
 
@@ -46,7 +45,6 @@ def init(self) -> None:
         # this, ADK spans are emitted into a no-op provider and never reach Langfuse.
         self._ensure_langfuse_client()
 
-        self._register_litellm_langfuse_otel()
         self._instrument_google_adk()
 
     def _ensure_langfuse_client(self) -> None:
@@ -64,21 +62,6 @@ def _ensure_langfuse_client(self) -> None:
             return
         self._langfuse_client_initialized = True
 
-    def _register_litellm_langfuse_otel(self) -> None:
-        """Register LiteLLM Langfuse callback."""
-        if self._litellm_instrumented:
-            return
-        try:
-            import litellm  # noqa: PLC0415
-        except ImportError:
-            logger.debug("litellm not installed; skipping LiteLLM Langfuse callback.")
-            return
-
-        existing = list(getattr(litellm, "callbacks", None) or [])
-        if "langfuse_otel" not in existing:
-            litellm.callbacks = [*existing, "langfuse_otel"]
-        self._litellm_instrumented = True
-
     def _instrument_google_adk(self) -> None:
         """Instrument Google ADK."""
         if self._google_adk_instrumented:
@@ -120,11 +103,13 @@ def init_langfuse_tracing() -> None:
        ``TracerProvider`` receives Langfuse's span processor.  This is required
        for ADK spans emitted via ``openinference-instrumentation-google-adk``
        to reach Langfuse.
-    2. Appends ``"langfuse_otel"`` to ``litellm.callbacks`` once (if
-       ``litellm`` is importable).
-    3. Runs ``GoogleADKInstrumentor().instrument()`` once (if
+    2. Runs ``GoogleADKInstrumentor().instrument()`` once (if
        ``openinference-instrumentation-google-adk`` is importable).
 
+    LiteLLM's ``langfuse_otel`` callback is deliberately not registered: it is
+    unusable against the Langfuse v4 SDK and stamps a zero ``llm.cost.total``
+    on the active span, which suppresses Langfuse's own cost calculation.
+
     Set ``LANGFUSE_HOST`` or ``LANGFUSE_BASE_URL`` for non-default regions.
     For short-lived processes, call ``langfuse.get_client().flush()`` before
     exit so pending spans are exported.

aieng-forecasting/aieng/forecasting/methods/agentic/agent_factory.py

@@ -669,12 +669,16 @@ def build_adk_agent(
     if isinstance(model, str) and config.openai_base_url:
         from google.adk.models.lite_llm import LiteLlm  # noqa: PLC0415
 
-        # Prefix with "openai/" so LiteLLM uses the OpenAI-compatible path.
-        # LiteLLM strips the prefix before sending, so the proxy receives the
-        # bare model name.
-        litellm_model = model if model.startswith("openai/") else f"openai/{model}"
+        # Route via LiteLLM's OpenAI-compatible path with ``custom_llm_provider``
+        # rather than an ``openai/`` model prefix. ADK stamps ``LlmRequest.model``
+        # from this name and OpenInference reports it to Langfuse, which matches
+        # its per-model price table on the bare name. A prefixed name matches
+        # nothing, so the generation is logged at zero cost. Both forms route
+        # identically.
+        bare_model = model[len("openai/") :] if model.startswith("openai/") else model
         model = LiteLlm(
-            model=litellm_model,
+            model=bare_model,
+            custom_llm_provider="openai",
             api_base=config.openai_base_url,
             api_key=config.openai_api_key,
         )

aieng-forecasting/aieng/forecasting/methods/llm_processes/_client.py

@@ -20,13 +20,13 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import contextvars
 import json
 import logging
-import os
 import warnings
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, Callable, TypeVar
+from typing import Any, Callable, Iterator, TypeVar
 
 from pydantic import BaseModel, ValidationError
 
@@ -39,21 +39,22 @@
 
 
 def bootstrap_litellm() -> None:
-    """One-time wiring of LiteLLM callbacks.
+    """Suppress LiteLLM and OpenTelemetry logging noise, once per process.
 
     Lazy and idempotent so non-LLM predictors do not require Langfuse env vars.
-    The Langfuse OTEL callback is registered only when ``LANGFUSE_PUBLIC_KEY``
-    is set in the environment.
+
+    LiteLLM's ``langfuse_otel`` callback is deliberately not registered. It is
+    unusable against the Langfuse v4 SDK this repo depends on, and it stamps
+    ``llm.cost.total`` on the active span from LiteLLM's own ``response_cost``,
+    which is ``0`` for every proxy-routed model. Langfuse honours a supplied
+    cost instead of deriving one from usage, so the callback forced agent-path
+    generations to $0. Instead, :func:`langfuse_generation` creates LLM-process
+    generations directly and OpenInference covers the agent path, so both price
+    correctly from ``usage_details``.
     """
     global _BOOTSTRAP_DONE  # noqa: PLW0603
     if _BOOTSTRAP_DONE:
         return
-    import litellm  # noqa: PLC0415
-
-    if os.environ.get("LANGFUSE_PUBLIC_KEY"):
-        existing = list(getattr(litellm, "callbacks", []) or [])
-        if "langfuse_otel" not in existing:
-            litellm.callbacks = [*existing, "langfuse_otel"]
 
     # Suppress LiteLLM startup and OTEL noise (mirrors agent_factory.py filter).
     # Bedrock/SageMaker "no botocore" and OTEL proxy-server notices are harmless.
@@ -91,6 +92,61 @@ def _noop(fn: Any) -> Any:
         return _noop
 
 
+class _NoopGeneration:
+    """Stand-in used when Langfuse is unavailable, so callers need no branching."""
+
+    def update(self, **kwargs: Any) -> None:
+        """Discard the update."""
+        return
+
+
+@contextlib.contextmanager
+def langfuse_generation(*, name: str, model: str, input_messages: Any) -> Iterator[Any]:
+    """Create a Langfuse ``generation`` around one LLM call.
+
+    LiteLLM's ``langfuse_otel`` callback emits no generation when the call runs
+    inside an already-active Langfuse span, which is every LLM-process
+    ``predict`` because they are wrapped in :func:`langfuse_observe`. Token
+    usage, and so cost, never reached Langfuse for those runs. Creating the
+    generation here works when nested and keeps the model, usage, and payload
+    under this module's control.
+
+    Cost is deliberately not set. Langfuse derives it from ``usage_details``
+    against its own per-model prices, which match the Vector proxy's published
+    rates.
+
+    ``start_as_current_observation(as_type="generation")`` is a first-class
+    Langfuse v4 instrumentation API. LiteLLM's Langfuse bridge targets the v2
+    SDK, pinning ``langfuse = ^2.45.0``, while this repo requires
+    ``langfuse>=4.5.1``. Support for v4 is BerriAI/litellm#24123, open and
+    unanswered since 2026-03-19. Retire this helper in favour of the callback
+    once that issue is closed and ``langfuse_otel`` is confirmed to emit a
+    generation under an active Langfuse span.
+
+    Yields a handle exposing ``update(**kwargs)``. That handle is a no-op
+    stand-in when Langfuse is not installed or a generation cannot be started,
+    so predictors remain usable without the ``agentic`` and ``llm`` extras.
+    """
+    manager = None
+    try:
+        from langfuse import get_client  # noqa: PLC0415
+
+        manager = get_client().start_as_current_observation(
+            as_type="generation",
+            name=name,
+            model=model,
+            input=input_messages,
+        )
+    except Exception:  # pragma: no cover - depends on optional dependency
+        logger.debug("Langfuse generation unavailable; usage will not be traced.", exc_info=True)
+
+    if manager is None:
+        yield _NoopGeneration()
+        return
+    with manager as generation:
+        yield generation
+
+
 def current_trace_info() -> tuple[str | None, str | None]:
     """Return ``(trace_id, trace_url)`` from the active Langfuse client, if any."""
     try:
@@ -299,17 +355,21 @@ 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
+    # ``model`` is the bare name as configured, before any "openai/" prefixing
+    # above; that is what Langfuse's price table matches on.
+    with langfuse_generation(name="llm_completion", model=model, input_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
+        # Full usage exposes the thinking-token breakdown when the proxy
+        # populates 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
+        generation.update(output=content, usage_details={"input": in_tok, "output": out_tok})
     return content, cost, in_tok, out_tok
 
 

aieng-forecasting/tests/aieng/forecasting/methods/agentic/test_agent_factory.py

@@ -108,10 +108,12 @@ def test_string_model_wrapped_in_litellm_when_proxy_set(self) -> None:
         agent = build_adk_agent(config)
 
         assert isinstance(agent.model, LiteLlm)
-        # LiteLlm receives the "openai/" prefix so LiteLLM routes via the
-        # OpenAI-compatible proxy path; the prefix is stripped before the
-        # proxy sees the model name.
-        assert agent.model.model == "openai/gemini-3.1-flash-lite-preview"
+        # The bare model name is kept and the OpenAI-compatible proxy route is
+        # selected via custom_llm_provider instead of an "openai/" prefix:
+        # OpenInference reports this name to Langfuse, which matches its price
+        # table on the bare name (a prefixed name logs zero cost).
+        assert agent.model.model == "gemini-3.1-flash-lite-preview"
+        assert agent.model._additional_args["custom_llm_provider"] == "openai"
 
     def test_string_model_kept_as_string_without_proxy(self, monkeypatch: pytest.MonkeyPatch) -> None:
         """Without a proxy URL the model is passed as a plain string to LlmAgent."""

aieng-forecasting/tests/aieng/forecasting/methods/llm_processes/test__client.py

@@ -7,6 +7,8 @@
 
 from __future__ import annotations
 
+import contextlib
+from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -121,6 +123,7 @@ def _mock_litellm_response(content: str) -> MagicMock:
     return resp
 
 
+_CLIENT = "aieng.forecasting.methods.llm_processes._client"
 _DUMMY_MESSAGES = [{"role": "user", "content": "forecast"}]
 _DUMMY_FORMAT = {"type": "json_schema", "json_schema": {"name": "x", "schema": {}, "strict": True}}
 
@@ -257,3 +260,75 @@ async def fake_acompletion(**kwargs):  # type: ignore[override]
     assert kw["reasoning_effort"] == "low"
     assert "extra_body" not in kw or "reasoning_effort" not in kw.get("extra_body", {})
     assert kw.get("drop_params") is True
+
+
+# ---------------------------------------------------------------------------
+# langfuse_generation: usage and cost reporting for the Langfuse trace
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_generation_receives_bare_model_usage_and_output() -> None:
+    """The Langfuse generation gets the bare model name, token usage, and output.
+
+    The bare (un-prefixed) name matters: Langfuse matches its per-model price
+    table on it, so sending ``openai/<model>`` would yield no cost.
+    """
+    captured: dict = {}
+
+    @contextlib.contextmanager
+    def _fake_generation(*, name: str, model: str, input_messages: object):
+        captured.update(name=name, model=model, input_messages=input_messages)
+
+        class _Gen:
+            def update(self, **kwargs: object) -> None:
+                captured["update"] = kwargs
+
+        yield _Gen()
+
+    resp = _mock_litellm_response('{"ok": 1}')
+    resp.usage = SimpleNamespace(prompt_tokens=11, completion_tokens=7)
+
+    with (
+        patch(f"{_CLIENT}.langfuse_generation", _fake_generation),
+        patch("litellm.acompletion", new=AsyncMock(return_value=resp)),
+    ):
+        await _one_completion_async(
+            model="gemini-3.5-flash",
+            messages=_DUMMY_MESSAGES,
+            response_format=_DUMMY_FORMAT,
+            temperature=1.0,
+            max_tokens=512,
+            timeout_s=30.0,
+            reasoning_effort=None,
+            api_base="https://proxy.example.com/v1",
+        )
+
+    assert captured["model"] == "gemini-3.5-flash"  # not "openai/gemini-3.5-flash"
+    assert captured["input_messages"] == _DUMMY_MESSAGES
+    assert captured["update"]["usage_details"] == {"input": 11, "output": 7}
+    assert captured["update"]["output"] == '{"ok": 1}'
+
+
+@pytest.mark.asyncio
+async def test_completion_succeeds_when_langfuse_is_unavailable() -> None:
+    """A failing Langfuse client degrades to a no-op and the completion returns."""
+    resp = _mock_litellm_response('{"ok": 1}')
+    resp.usage = SimpleNamespace(prompt_tokens=3, completion_tokens=4)
+
+    with (
+        patch(f"{_CLIENT}.get_client", side_effect=RuntimeError("no langfuse"), create=True),
+        patch("litellm.acompletion", new=AsyncMock(return_value=resp)),
+    ):
+        content, _cost, in_tok, out_tok = await _one_completion_async(
+            model="gemini-3.5-flash",
+            messages=_DUMMY_MESSAGES,
+            response_format=_DUMMY_FORMAT,
+            temperature=1.0,
+            max_tokens=512,
+            timeout_s=30.0,
+            reasoning_effort=None,
+        )
+
+    assert content == '{"ok": 1}'
+    assert (in_tok, out_tok) == (3, 4)