diff --git a/lyra/llm.py b/lyra/llm.py index de02efe..080a6cc 100644 --- a/lyra/llm.py +++ b/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": mdl, "messages": messages, "stream": False}, + timeout=120, + ) + resp.raise_for_status() + out = resp.json()["message"]["content"] - resp = httpx.post( - f"{cfg.local_base_url}/api/chat", - json={"model": model or cfg.local_model, "messages": messages, "stream": False}, - timeout=120, - ) - resp.raise_for_status() - return 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 = {