From 61494362a3e92803eb7a0ffce05f8d2dd773db0c Mon Sep 17 00:00:00 2001 From: buenosairesam Date: Wed, 3 Jun 2026 10:33:22 -0300 Subject: [PATCH] langfuse generation spans with token usage; switch llm to vllm (qwen2.5-coder-7b on mcrndeb) --- api/analyses/compare_periods.py | 6 ++- api/analyses/drill_down/helpers.py | 8 ++-- api/langfuse_client.py | 17 +++++-- api/llm.py | 72 +++++++++++++++++++++++++----- api/runtime/runner.py | 3 +- ctrl/k8s/base/api.yaml | 14 ++++-- ctrl/k8s/base/configmap.yaml | 8 +++- 7 files changed, 103 insertions(+), 25 deletions(-) diff --git a/api/analyses/compare_periods.py b/api/analyses/compare_periods.py index b1cd401..fbe6b04 100644 --- a/api/analyses/compare_periods.py +++ b/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"): diff --git a/api/analyses/drill_down/helpers.py b/api/analyses/drill_down/helpers.py index 984a3de..47e3e8f 100644 --- a/api/analyses/drill_down/helpers.py +++ 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 ── diff --git a/api/langfuse_client.py b/api/langfuse_client.py index dacc936..ee54c36 100644 --- a/api/langfuse_client.py +++ b/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=""` 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() diff --git a/api/llm.py b/api/llm.py index 3c91d33..76ae7d4 100644 --- a/api/llm.py +++ b/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 diff --git a/api/runtime/runner.py b/api/runtime/runner.py index 2ce8feb..06f9be6 100644 --- a/api/runtime/runner.py +++ b/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} diff --git a/ctrl/k8s/base/api.yaml b/ctrl/k8s/base/api.yaml index 6fa025a..adf945d 100644 --- a/ctrl/k8s/base/api.yaml +++ b/ctrl/k8s/base/api.yaml @@ -13,14 +13,20 @@ spec: labels: app: api spec: - # lng.local.ar resolves to ::1 via the host's dnsmasq → coredns chain, - # which is the pod's own loopback (wrong). Override to the kind bridge - # gateway so traffic from this pod reaches the host's system Caddy and - # gets reverse-proxied to Langfuse on :3000. hostAliases: + # lng.local.ar resolves to ::1 via the host's dnsmasq → coredns chain, + # which is the pod's own loopback (wrong). Override to the kind bridge + # gateway so traffic from this pod reaches the host's system Caddy and + # gets reverse-proxied to Langfuse on :3000. - ip: "172.19.0.1" hostnames: - "lng.local.ar" + # mcrndeb is the GPU box on the LAN (vLLM lives there). kind-nvi + # pods reach it via routing through the host; we just need the name + # to resolve since coredns doesn't know about LAN hostnames. + - ip: "192.168.1.3" + hostnames: + - "mcrndeb" containers: - name: api image: nvi-api diff --git a/ctrl/k8s/base/configmap.yaml b/ctrl/k8s/base/configmap.yaml index 30d3fff..a18b689 100644 --- a/ctrl/k8s/base/configmap.yaml +++ b/ctrl/k8s/base/configmap.yaml @@ -13,7 +13,13 @@ data: # LLM. Supported providers: groq, anthropic, openai. # Secrets (*_API_KEY) live in the Secret built by kustomize from .env. - LLM_PROVIDER: "groq" + # Flip LLM_PROVIDER to switch — the other configs stay so we can A/B easily. + LLM_PROVIDER: "openai" + + # Local vLLM serving Qwen2.5-Coder-7B-Instruct-AWQ on mcrndeb. + OPENAI_BASE_URL: "http://mcrndeb:11000/v1" + OPENAI_MODEL: "Qwen/Qwen2.5-Coder-7B-Instruct-AWQ" + GROQ_MODEL: "llama-3.3-70b-versatile" ANTHROPIC_MODEL: "claude-sonnet-4-6"