diff --git a/src/fabledassistant/services/generation_task.py b/src/fabledassistant/services/generation_task.py index 4f26c51..2fe4699 100644 --- a/src/fabledassistant/services/generation_task.py +++ b/src/fabledassistant/services/generation_task.py @@ -212,7 +212,11 @@ async def run_generation( t_start = time.monotonic() timing: dict = { + "think": think, "tools": [], + "rounds": 0, + "prompt_tokens": None, + "output_tokens": None, "ttft_ms": None, "generation_ms": None, "total_ms": None, @@ -228,6 +232,7 @@ async def run_generation( research_completed = False for _round in range(MAX_TOOL_ROUNDS + 1): + timing["rounds"] = _round + 1 round_tool_calls: list[dict] = [] logger.info("Generation round %d started for conv %d (model=%s)", _round, conv_id, model) @@ -261,6 +266,12 @@ async def run_generation( logger.warning("Failed periodic flush for message %d", msg_id, exc_info=True) last_flush = now + elif chunk.type == "done": + if chunk.prompt_tokens is not None: + timing["prompt_tokens"] = (timing["prompt_tokens"] or 0) + chunk.prompt_tokens + if chunk.output_tokens is not None: + timing["output_tokens"] = (timing["output_tokens"] or 0) + chunk.output_tokens + elif chunk.type == "tool_calls" and chunk.tool_calls: logger.info("Round %d: model returned %d tool call(s)", _round, len(chunk.tool_calls)) for tc in chunk.tool_calls: diff --git a/src/fabledassistant/services/llm.py b/src/fabledassistant/services/llm.py index 9a327e4..ce3f02c 100644 --- a/src/fabledassistant/services/llm.py +++ b/src/fabledassistant/services/llm.py @@ -149,6 +149,9 @@ class ChatChunk: type: Literal["content", "thinking", "tool_calls", "done"] content: str = "" tool_calls: list[dict] | None = None + # Token counts from the Ollama done event (only set on type="done") + prompt_tokens: int | None = None + output_tokens: int | None = None async def stream_chat_with_tools( @@ -221,7 +224,11 @@ async def stream_chat_with_tools( yield ChatChunk(type="tool_calls", tool_calls=accumulated_tool_calls) else: logger.debug("Ollama done with no tool calls") - yield ChatChunk(type="done") + yield ChatChunk( + type="done", + prompt_tokens=data.get("prompt_eval_count"), + output_tokens=data.get("eval_count"), + ) break diff --git a/src/fabledassistant/services/logging.py b/src/fabledassistant/services/logging.py index e7a07c0..0c70f6e 100644 --- a/src/fabledassistant/services/logging.py +++ b/src/fabledassistant/services/logging.py @@ -68,12 +68,12 @@ async def log_generation( """Persist per-generation timing breakdown to app_logs for benchmarking.""" async with async_session() as session: log = AppLog( - category="usage", + category="generation", user_id=user_id, action="generation", endpoint=f"/chat/conversations/{conv_id}", duration_ms=timing.get("total_ms"), - details=json.dumps({"model": model, **timing}), + details=json.dumps({"model": model, "conv_id": conv_id, **timing}), ) session.add(log) await session.commit()