SupplyAgent/supply_agent/logging/logger.py

from __future__ import annotations

import json
import logging
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from supply_agent.types import AgentResult, Message, ToolDefinition


def _serialize(obj: Any) -> str:
    """Serialize any object to a complete, non-truncated JSON string."""

    def default(o: Any) -> Any:
        if hasattr(o, "model_dump"):
            return o.model_dump()
        if hasattr(o, "to_api_dict"):
            return o.to_api_dict()
        return str(o)

    return json.dumps(obj, ensure_ascii=False, indent=2, default=default)


def _try_parse_json(text: str) -> Any:
    """Best-effort JSON parse; return original string on failure."""
    try:
        return json.loads(text)
    except (json.JSONDecodeError, TypeError):
        return text


def _extract_usage(raw_response: Any) -> dict[str, Any] | None:
    if raw_response is None:
        return None
    usage = getattr(raw_response, "usage", None)
    if usage is None and isinstance(raw_response, dict):
        usage = raw_response.get("usage")
    if usage is None:
        return None
    if hasattr(usage, "model_dump"):
        return usage.model_dump()
    if isinstance(usage, dict):
        return usage
    return {"raw": str(usage)}


def _extract_skill_name(skill_name_or_args: str) -> str:
    parsed = _try_parse_json(skill_name_or_args)
    if isinstance(parsed, dict):
        return str(parsed.get("name") or skill_name_or_args)
    return skill_name_or_args


class _FullContentFormatter(logging.Formatter):
    """Formatter that never truncates message content."""

    def format(self, record: logging.LogRecord) -> str:
        record.message = record.getMessage()
        return f"[{self.formatTime(record, '%Y-%m-%d %H:%M:%S')}] {record.message}"


class AgentLogger:
    """
    Encapsulated logger for agent runs.

    Each run writes:
    - ``logs/run_<id>.log``   — human-readable full dump
    - ``logs/run_<id>.jsonl`` — structured event stream for visualization
    """

    def __init__(self, logs_dir: Path | str = "logs", *, enabled: bool = True) -> None:
        self.logs_dir = Path(logs_dir)
        self.enabled = enabled
        self._run_id: str | None = None
        self._log_file: Path | None = None
        self._jsonl_file: Path | None = None
        self._logger: logging.Logger | None = None
        self._seq: int = 0
        self._jsonl_fh: Any | None = None

    @property
    def run_id(self) -> str | None:
        return self._run_id

    @property
    def log_file(self) -> Path | None:
        return self._log_file

    @property
    def jsonl_file(self) -> Path | None:
        return self._jsonl_file

    def start_run(self, user_input: str, *, model: str) -> str:
        """Start a new run log file. Returns the run id."""
        if not self.enabled:
            return ""

        self.logs_dir.mkdir(parents=True, exist_ok=True)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self._run_id = f"{timestamp}_{uuid.uuid4().hex[:8]}"
        self._log_file = self.logs_dir / f"run_{self._run_id}.log"
        self._jsonl_file = self.logs_dir / f"run_{self._run_id}.jsonl"
        self._seq = 0

        self._logger = logging.getLogger(f"supply_agent.run.{self._run_id}")
        self._logger.setLevel(logging.DEBUG)
        self._logger.handlers.clear()
        self._logger.propagate = False

        formatter = _FullContentFormatter()

        file_handler = logging.FileHandler(self._log_file, encoding="utf-8")
        file_handler.setFormatter(formatter)
        self._logger.addHandler(file_handler)

        console_handler = logging.StreamHandler()
        console_handler.setFormatter(formatter)
        self._logger.addHandler(console_handler)

        self._jsonl_fh = open(self._jsonl_file, "w", encoding="utf-8")

        self._emit(
            "run_start",
            {
                "run_id": self._run_id,
                "model": model,
                "user_input": user_input,
                "log_file": str(self._log_file),
                "jsonl_file": str(self._jsonl_file),
            },
            title="RUN START",
        )
        return self._run_id

    def log_llm_input(
        self,
        iteration: int,
        model: str,
        messages: list[Message],
        tools: list[ToolDefinition] | None,
        temperature: float,
    ) -> None:
        """Log the complete LLM request payload."""
        if not self.enabled or not self._logger:
            return

        payload = {
            "iteration": iteration,
            "model": model,
            "temperature": temperature,
            "messages": [m.to_api_dict() for m in messages],
            "tools": [t.to_api_dict() for t in tools] if tools else None,
            "message_count": len(messages),
            "tool_count": len(tools) if tools else 0,
        }
        self._emit(
            "llm_input",
            payload,
            iteration=iteration,
            title=f"LLM INPUT | step={self._seq + 1} | iteration={iteration}",
        )

    def log_llm_output(
        self,
        iteration: int,
        response: Message,
        raw_response: Any | None = None,
    ) -> None:
        """Log the complete LLM response, highlighting reasoning and tool calls."""
        if not self.enabled or not self._logger:
            return

        parsed = response.model_dump()
        tool_calls = parsed.get("tool_calls") or []
        usage = _extract_usage(raw_response)

        payload: dict[str, Any] = {
            "iteration": iteration,
            "content": parsed.get("content"),
            "reasoning": parsed.get("reasoning"),
            "tool_calls": tool_calls,
            "has_reasoning": bool(parsed.get("reasoning")),
            "tool_call_count": len(tool_calls),
            "usage": usage,
            "parsed": parsed,
        }
        if raw_response is not None:
            if hasattr(raw_response, "model_dump"):
                payload["raw"] = raw_response.model_dump()
            else:
                payload["raw"] = raw_response

        self._emit(
            "llm_output",
            payload,
            iteration=iteration,
            title=f"LLM OUTPUT | step={self._seq + 1} | iteration={iteration}",
        )

    def log_tool_call(
        self,
        iteration: int,
        name: str,
        arguments: str,
        result: str,
        is_error: bool = False,
        *,
        tool_call_id: str | None = None,
    ) -> None:
        """Log a tool execution with full arguments and result."""
        if not self.enabled or not self._logger:
            return

        args_parsed = _try_parse_json(arguments)
        result_parsed = _try_parse_json(result)

        payload = {
            "iteration": iteration,
            "tool": name,
            "tool_call_id": tool_call_id,
            "arguments": arguments,
            "arguments_parsed": args_parsed,
            "result": result,
            "result_parsed": result_parsed,
            "is_error": is_error,
        }
        status = "ERROR" if is_error else "OK"
        self._emit(
            "tool_call",
            payload,
            iteration=iteration,
            title=f"TOOL CALL | step={self._seq + 1} | iteration={iteration} | {name} [{status}]",
        )

    def log_skill_loaded(self, iteration: int, skill_name: str) -> None:
        """Log when a skill is loaded."""
        if not self.enabled or not self._logger:
            return

        name = _extract_skill_name(skill_name)
        self._emit(
            "skill_loaded",
            {"iteration": iteration, "skill": name},
            iteration=iteration,
            title=f"SKILL LOADED | step={self._seq + 1} | iteration={iteration} | {name}",
        )

    def end_run(self, result: AgentResult) -> None:
        """Log run summary and close the run log."""
        if not self.enabled or not self._logger:
            return

        self._emit(
            "run_end",
            {
                "run_id": self._run_id,
                "iterations": result.iterations,
                "tool_calls_made": result.tool_calls_made,
                "skills_used": result.skills_used,
                "final_content": result.content,
            },
            title="RUN END",
        )

        for handler in self._logger.handlers:
            handler.close()
        self._logger.handlers.clear()

        if self._jsonl_fh is not None:
            self._jsonl_fh.close()
            self._jsonl_fh = None

    def _emit(
        self,
        event: str,
        data: dict[str, Any],
        *,
        title: str,
        iteration: int | None = None,
    ) -> None:
        assert self._logger is not None
        self._seq += 1
        ts = datetime.now(timezone.utc).astimezone().isoformat(timespec="seconds")

        record: dict[str, Any] = {
            "event": event,
            "ts": ts,
            "seq": self._seq,
            "run_id": self._run_id,
            "data": data,
        }
        if iteration is not None:
            record["iteration"] = iteration

        if self._jsonl_fh is not None:
            self._jsonl_fh.write(json.dumps(record, ensure_ascii=False, default=str) + "\n")
            self._jsonl_fh.flush()

        separator = "=" * 80
        body = _serialize(data)
        self._logger.info("%s\n%s\n%s\n%s", separator, title, separator, body)


# Module-level default logger instance
_default_logger: AgentLogger | None = None


def get_agent_logger(
    logs_dir: Path | str = "logs",
    *,
    enabled: bool = True,
) -> AgentLogger:
    """Get or create the default AgentLogger instance."""
    global _default_logger
    if _default_logger is None:
        _default_logger = AgentLogger(logs_dir, enabled=enabled)
    return _default_logger