lyra/llm.py

@@ -2,11 +2,13 @@
 from __future__ import annotations
 
 import json
+import time
 from typing import Iterator, Literal, TypedDict
 
 import httpx
 from openai import OpenAI
 
+from lyra import logbus
 from lyra.config import load
 
 
@@ -18,30 +20,54 @@ class Message(TypedDict):
 Backend = Literal["local", "cloud", "mi50"]
 
 
+def _approx_tok(messages: list) -> int:
+    """Rough prompt size (chars/4) — enough to see what's loading a backend."""
+    total = 0
+    for m in messages or []:
+        if isinstance(m, dict) and isinstance(m.get("content"), str):
+            total += len(m["content"])
+    return total // 4
+
+
+def _resolved_model(cfg, backend: Backend, model: str | None) -> str:
+    if backend == "cloud":
+        return model or cfg.cloud_model
+    if backend == "mi50":
+        return model or cfg.mi50_model
+    return model or cfg.local_model
+
+
 def complete(messages: list[Message], backend: Backend = "local", model: str | None = None) -> str:
     """Generate a completion. `model` overrides the backend's default model
     (used so live chat can run a stronger cloud model than bulk consolidation)."""
     cfg = load()
+    mdl = _resolved_model(cfg, backend, model)
+    logbus.log("info", "llm call", kind="complete", backend=backend, model=mdl, tok=_approx_tok(messages))
+    t0 = time.monotonic()
+
     if backend == "cloud":
         if not cfg.openai_api_key:
             raise RuntimeError("OPENAI_API_KEY is not set")
         client = OpenAI(api_key=cfg.openai_api_key)
-        resp = client.chat.completions.create(model=model or cfg.cloud_model, messages=messages)
-        return resp.choices[0].message.content or ""
-
-    if backend == "mi50":
+        resp = client.chat.completions.create(model=mdl, messages=messages)
+        out = resp.choices[0].message.content or ""
+    elif backend == "mi50":
         # MI50 box runs an OpenAI-compatible llama.cpp server; key is unused.
         client = OpenAI(api_key="not-needed", base_url=cfg.mi50_base_url)
-        resp = client.chat.completions.create(model=model or cfg.mi50_model, messages=messages)
-        return resp.choices[0].message.content or ""
-
+        resp = client.chat.completions.create(model=mdl, messages=messages)
+        out = resp.choices[0].message.content or ""
+    else:
         resp = httpx.post(
             f"{cfg.local_base_url}/api/chat",
-        json={"model": model or cfg.local_model, "messages": messages, "stream": False},
+            json={"model": mdl, "messages": messages, "stream": False},
             timeout=120,
         )
         resp.raise_for_status()
-    return resp.json()["message"]["content"]
+        out = resp.json()["message"]["content"]
+
+    logbus.log("info", "llm done", kind="complete", backend=backend,
+               ms=int((time.monotonic() - t0) * 1000), out=len(out))
+    return out
 
 
 def chat_call(
@@ -68,6 +94,8 @@ def chat_call(
         kwargs: dict = {"model": mdl, "messages": messages}
         if tools:
             kwargs["tools"] = tools
+        logbus.log("info", "llm call", kind="chat", backend=backend, model=mdl, tok=_approx_tok(messages))
+        t0 = time.monotonic()
         msg = client.chat.completions.create(**kwargs).choices[0].message
         tcs = None
         if getattr(msg, "tool_calls", None):
@@ -75,6 +103,9 @@ def chat_call(
                 {"id": tc.id, "name": tc.function.name, "arguments": tc.function.arguments}
                 for tc in msg.tool_calls
             ]
+        logbus.log("info", "llm done", kind="chat", backend=backend,
+                   ms=int((time.monotonic() - t0) * 1000), out=len(msg.content or ""),
+                   tools=[t["name"] for t in tcs] if tcs else None)
         return msg.model_dump(), tcs
 
     # local (Ollama): no tool-calling here — return plain content.
@@ -105,6 +136,8 @@ def chat_call_stream(
         kwargs: dict = {"model": mdl, "messages": messages, "stream": True}
         if tools:
             kwargs["tools"] = tools
+        logbus.log("info", "llm call", kind="chat-stream", backend=backend, model=mdl, tok=_approx_tok(messages))
+        t0 = time.monotonic()
         parts: list[str] = []
         frags: dict[int, dict] = {}  # tool-call fragments accumulated by index
         for chunk in client.chat.completions.create(**kwargs):
@@ -123,6 +156,9 @@ def chat_call_stream(
                 if tc.function and tc.function.arguments:
                     slot["arguments"] += tc.function.arguments
         content = "".join(parts)
+        logbus.log("info", "llm done", kind="chat-stream", backend=backend,
+                   ms=int((time.monotonic() - t0) * 1000), out=len(content),
+                   tools=[frags[i]["name"] for i in sorted(frags)] if frags else None)
         if frags:
             calls = [frags[i] for i in sorted(frags)]
             assistant = {