104 lines
3.2 KiB
Python
104 lines
3.2 KiB
Python
"""
|
|
Utility for logging traces from LLM calls.
|
|
|
|
This module provides structured JSONL logging of LLM inputs/outputs.
|
|
"""
|
|
|
|
import fcntl
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from src.config import (
|
|
ConfiguredModelSettings,
|
|
ModelConfig,
|
|
settings,
|
|
)
|
|
|
|
|
|
def get_reasoning_traces_file_path() -> Path | None:
|
|
"""Get the traces file path from settings."""
|
|
if settings.REASONING_TRACES_FILE:
|
|
return Path(settings.REASONING_TRACES_FILE)
|
|
return None
|
|
|
|
|
|
def log_reasoning_trace(
|
|
task_type: str,
|
|
model_config: ModelConfig | ConfiguredModelSettings,
|
|
prompt: str,
|
|
response: Any,
|
|
*,
|
|
max_tokens: int | None = None,
|
|
thinking_budget_tokens: int | None = None,
|
|
reasoning_effort: str | None = None,
|
|
json_mode: bool = False,
|
|
stop_seqs: list[str] | None = None,
|
|
messages: list[dict[str, Any]] | None = None,
|
|
) -> None:
|
|
"""
|
|
Log a trace to the configured JSONL file.
|
|
|
|
Args:
|
|
task_type: Type of task (e.g., "minimal_deriver", "dialectic_chat")
|
|
model_config: Model configuration used for the call
|
|
prompt: The full prompt text sent to the LLM (used if messages is None)
|
|
response: HonchoLLMCallResponse object with the LLM response
|
|
max_tokens: Max output tokens setting
|
|
thinking_budget_tokens: Anthropic thinking budget (if used)
|
|
reasoning_effort: OpenAI reasoning effort (if used)
|
|
json_mode: Whether JSON mode was enabled
|
|
stop_seqs: Stop sequences used (if any)
|
|
messages: Full conversation history for multi-turn/agentic calls
|
|
"""
|
|
traces_file = get_reasoning_traces_file_path()
|
|
if not traces_file:
|
|
return
|
|
|
|
# Serialize response content - handle Pydantic models
|
|
content = response.content
|
|
if isinstance(content, BaseModel):
|
|
content = content.model_dump()
|
|
|
|
trace_entry: dict[str, Any] = {
|
|
"timestamp": time.time(),
|
|
"task_type": task_type,
|
|
"provider": model_config.transport,
|
|
"model": model_config.model,
|
|
"settings": {
|
|
"max_tokens": max_tokens,
|
|
"thinking_budget_tokens": thinking_budget_tokens,
|
|
"reasoning_effort": reasoning_effort,
|
|
"json_mode": json_mode,
|
|
"stop_seqs": stop_seqs,
|
|
},
|
|
"input": {
|
|
"tokens": response.input_tokens,
|
|
},
|
|
"output": {
|
|
"content": content,
|
|
"tokens": response.output_tokens,
|
|
"finish_reasons": response.finish_reasons,
|
|
"thinking_content": response.thinking_content,
|
|
},
|
|
}
|
|
|
|
# Use messages for multi-turn/agentic calls, otherwise use prompt
|
|
if messages is not None:
|
|
trace_entry["input"]["messages"] = messages
|
|
else:
|
|
trace_entry["input"]["prompt"] = prompt
|
|
|
|
# Include tool calls if present
|
|
if hasattr(response, "tool_calls_made") and response.tool_calls_made:
|
|
trace_entry["output"]["tool_calls"] = response.tool_calls_made
|
|
|
|
# Use file locking to handle concurrent writes from multiple processes
|
|
with open(traces_file, "a") as f:
|
|
fcntl.flock(f.fileno(), fcntl.LOCK_EX)
|
|
f.write(json.dumps(trace_entry) + "\n")
|
|
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
|