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

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