mariano/nvi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
buenosairesam
2026-06-03 10:33:22 -03:00
parent bd34c8c5ce
commit 61494362a3
7 changed files with 103 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
api/analyses/compare_periods.py
View File

@@ -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"):

8
api/analyses/drill_down/helpers.py
View File

@@ -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 ──

17
api/langfuse_client.py
View File

@@ -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()

72
api/llm.py
View File

@@ -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

3
api/runtime/runner.py
View File

@@ -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}

8
ctrl/k8s/base/api.yaml
View File

@@ -13,14 +13,20 @@ spec:
labels:
app: api
spec:
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.
hostAliases:
- 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

8
ctrl/k8s/base/configmap.yaml
View File

@@ -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"