langfuse generation spans with token usage; switch llm to vllm (qwen2.5-coder-7b on mcrndeb)

api/analyses/compare_periods.py

@@ -83,7 +83,10 @@ class ComparePeriods(Analysis):
                     system_len=len(interpret_system),
                     user_len=len(interpret_user),
                 ))
-                summary = chat(system=interpret_system, user=interpret_user, max_tokens=512)
+                summary = chat(
+                    system=interpret_system, user=interpret_user,
+                    max_tokens=512, span_name="compare_periods.interpret",
+                )
             except Exception as e:
                 logger.exception("compare_periods failed")
                 await events.publish_current(events.tool_call_end("compare_periods", error=str(e)))
@@ -122,6 +125,7 @@ def _generate_pair(question: str, period_a: str, period_b: str) -> dict[str, dic
             metrics_block=schema.render_metrics(),
         ),
         max_tokens=1024,
+        span_name="compare_periods.pair",
     )
     obj = json.loads(_extract_json(text))
     for k in ("a", "b"):

api/analyses/drill_down/helpers.py

@@ -43,11 +43,13 @@ async def decide_next(question: str, metric: str, dimensions: list[str],
         history=format_history(slices),
         budget=budget,
     )
-    with lf.span("drill_down.next", as_type="generation", input={"budget": budget}) as span:
+    with lf.span("drill_down.next", input={"budget": budget}) as span:
         await events.publish_current(events.llm_call(
             "drill_down.next", system_len=len(system), user_len=len(user),
         ))
-        decision = _parse_json(chat(system=system, user=user, max_tokens=256))
+        decision = _parse_json(chat(
+            system=system, user=user, max_tokens=256, span_name="drill_down.next.gen",
+        ))
 
         # Hard guard: the chosen dimension MUST be in the candidate list.
         if decision.get("action") == "drill":
@@ -142,7 +144,7 @@ async def interpret(question: str, metric: str, slices: list[Slice]) -> str:
     await events.publish_current(events.llm_call(
         "drill_down.interpret", system_len=len(system), user_len=len(user),
     ))
-    return chat(system=system, user=user, max_tokens=512)
+    return chat(system=system, user=user, max_tokens=512, span_name="drill_down.interpret")
 
 
 # ── Prompt-context formatters ──

api/langfuse_client.py

@@ -55,9 +55,14 @@ def span(
     as_type: str = "span",
     input: Any = None,
     metadata: dict[str, Any] | None = None,
+    model: str | None = None,
 ) -> Iterator[Any]:
     """Open a Langfuse observation; yields the span object or a no-op.
 
+    Pass `model="<id>"` when `as_type="generation"` so Langfuse can label
+    the trace and compute cost when it knows the model's pricing. Token
+    usage is reported by the body via `span.update(usage_details=...)`.
+
     Exceptions raised by the *body* of the `with` block propagate up
     untouched. Only Langfuse's own setup failures fall through to a no-op
     span — we never want to mask a real error.
@@ -66,10 +71,16 @@ def span(
     if lf is None:
         yield _NullSpan()
         return
+    kwargs: dict[str, Any] = {
+        "name": name,
+        "as_type": as_type,
+        "input": input,
+        "metadata": metadata or {},
+    }
+    if model is not None:
+        kwargs["model"] = model
     try:
-        cm = lf.start_as_current_observation(
-            name=name, as_type=as_type, input=input, metadata=metadata or {}
-        )
+        cm = lf.start_as_current_observation(**kwargs)
     except Exception as e:
         logger.warning("langfuse setup for %s failed: %s", name, e)
         yield _NullSpan()

api/llm.py

@@ -3,20 +3,25 @@
 Provider selection is via `settings.llm_provider`. Supported:
 - "groq"      — hits Groq via the openai SDK (groq_base_url + groq_api_key).
 - "anthropic" — native Anthropic SDK.
-- "openai"    — openai SDK with the standard base_url.
+- "openai"    — openai SDK with the standard base_url (also handles vLLM /
+                any OpenAI-compatible endpoint).
 
-The Anthropic and OpenAI/Groq clients are cached so HTTP connection pools
-are reused across calls.
+Every call opens a Langfuse `generation` span automatically, tagged with the
+model name and populated with token usage from the provider's response.
+Spans nest under whatever observation the caller has open, so the trace
+hierarchy shows `text_to_sql > llm.chat` with usage on the inner node.
 """
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from functools import lru_cache
 from typing import Callable
 
 from anthropic import Anthropic
 from openai import OpenAI
 
+from api import langfuse_client as lf
 from api.config import get_settings
 
 
@@ -41,7 +46,13 @@ def _openai() -> OpenAI:
 
 # ── Per-provider chat impls ──
 
-def _chat_anthropic(system: str, user: str, max_tokens: int) -> str:
+@dataclass
+class _Call:
+    text: str
+    usage: dict[str, int]   # {"input": N, "output": M, "total": N+M}
+
+
+def _chat_anthropic(system: str, user: str, max_tokens: int) -> _Call:
     s = get_settings()
     msg = _anthropic().messages.create(
         model=s.anthropic_model,
@@ -49,10 +60,18 @@ def _chat_anthropic(system: str, user: str, max_tokens: int) -> str:
         system=system,
         messages=[{"role": "user", "content": user}],
     )
-    return "".join(b.text for b in msg.content if getattr(b, "type", None) == "text").strip()
+    text = "".join(b.text for b in msg.content if getattr(b, "type", None) == "text").strip()
+    return _Call(
+        text=text,
+        usage={
+            "input": int(msg.usage.input_tokens),
+            "output": int(msg.usage.output_tokens),
+            "total": int(msg.usage.input_tokens + msg.usage.output_tokens),
+        },
+    )
 
 
-def _chat_openai_compat(client: OpenAI, model: str, system: str, user: str, max_tokens: int) -> str:
+def _chat_openai_compat(client: OpenAI, model: str, system: str, user: str, max_tokens: int) -> _Call:
     resp = client.chat.completions.create(
         model=model,
         max_tokens=max_tokens,
@@ -61,31 +80,60 @@ def _chat_openai_compat(client: OpenAI, model: str, system: str, user: str, max_
             {"role": "user", "content": user},
         ],
     )
-    return (resp.choices[0].message.content or "").strip()
+    text = (resp.choices[0].message.content or "").strip()
+    u = resp.usage
+    usage = {
+        "input": int(u.prompt_tokens) if u else 0,
+        "output": int(u.completion_tokens) if u else 0,
+        "total": int(u.total_tokens) if u else 0,
+    }
+    return _Call(text=text, usage=usage)
 
 
-def _chat_groq(system: str, user: str, max_tokens: int) -> str:
+def _chat_groq(system: str, user: str, max_tokens: int) -> _Call:
     return _chat_openai_compat(_groq(), get_settings().groq_model, system, user, max_tokens)
 
 
-def _chat_openai(system: str, user: str, max_tokens: int) -> str:
+def _chat_openai(system: str, user: str, max_tokens: int) -> _Call:
     return _chat_openai_compat(_openai(), get_settings().openai_model, system, user, max_tokens)
 
 
-_PROVIDERS: dict[str, Callable[[str, str, int], str]] = {
+_PROVIDERS: dict[str, Callable[[str, str, int], _Call]] = {
     "groq": _chat_groq,
     "anthropic": _chat_anthropic,
     "openai": _chat_openai,
 }
 
 
+def _active_model() -> str:
+    s = get_settings()
+    return {
+        "groq": s.groq_model,
+        "anthropic": s.anthropic_model,
+        "openai": s.openai_model,
+    }.get(s.llm_provider.lower(), s.llm_provider)
+
+
 # ── Public surface ──
 
-def chat(*, system: str, user: str, max_tokens: int = 1024) -> str:
+def chat(*, system: str, user: str, max_tokens: int = 1024, span_name: str = "llm.chat") -> str:
+    """Run a single chat completion. Opens a Langfuse generation span around
+    the call with the model name + token usage. `span_name` is the label
+    shown in Langfuse — pass something descriptive (e.g. "text_to_sql.gen")
+    so traces are scannable."""
     provider = get_settings().llm_provider.lower()
     impl = _PROVIDERS.get(provider)
     if impl is None:
         raise ValueError(
             f"unknown llm_provider {provider!r}; supported: {sorted(_PROVIDERS)}"
         )
-    return impl(system, user, max_tokens)
+
+    with lf.span(
+        span_name,
+        as_type="generation",
+        model=_active_model(),
+        input={"system": system, "user": user},
+    ) as gen:
+        call = impl(system, user, max_tokens)
+        gen.update(output=call.text, usage_details=call.usage)
+        return call.text

api/runtime/runner.py

@@ -128,7 +128,7 @@ async def _synthesize_node(state: RunState) -> dict[str, Any]:
         for f in findings
     )
 
-    with lf.span("synthesize", as_type="generation", input={"findings": len(findings)}) as span:
+    with lf.span("synthesize", input={"findings": len(findings)}) as span:
         answer = chat(
             system=load("synthesize.system"),
             user=render(
@@ -137,6 +137,7 @@ async def _synthesize_node(state: RunState) -> dict[str, Any]:
                 findings_block=findings_block,
             ),
             max_tokens=512,
+            span_name="synthesize.gen",
         )
         span.update(output={"answer": answer})
     return {"answer": answer}