From 2662d1b2531724791207a1aeef1d4c4d3c117dbc Mon Sep 17 00:00:00 2001 From: zphinx Date: Mon, 4 May 2026 06:03:39 +0200 Subject: [PATCH] feat(cli): add structured JSONL session logging for AI output --- src/tai/cli.py | 100 +++++++++++++++++++++++++++++++++++--- src/tai/session_log.py | 34 +++++++++++++ tests/test_session_log.py | 22 +++++++++ 3 files changed, 149 insertions(+), 7 deletions(-) create mode 100644 src/tai/session_log.py create mode 100644 tests/test_session_log.py diff --git a/src/tai/cli.py b/src/tai/cli.py index 1b3a0a7..16952b3 100644 --- a/src/tai/cli.py +++ b/src/tai/cli.py @@ -15,6 +15,7 @@ from tai.input_parser import InputValidationError, build_request from tai.models import TroubleshootRequest from tai.plan import plan_from_request from tai.prompt_builder import build_system_prompt, build_user_message +from tai.session_log import SessionLogger from tai.ssh_client import SSHClient, SSHCommandResult, SSHConnectionConfig, SSHSession app = typer.Typer(no_args_is_help=True, add_completion=False) @@ -85,6 +86,13 @@ def run( str, typer.Option("--ai-key", help="API key for the AI backend (not needed for Ollama)."), ] = "ollama", + log_file: Annotated[ + str | None, + typer.Option( + "--log-file", + help="Optional JSONL file path to log AI and session output.", + ), + ] = None, ) -> None: """Start an interactive troubleshooting session scaffold.""" try: @@ -120,6 +128,7 @@ def run( return # nothing SSH-related requested ai_config = AIConfig(host=ai_host, model=model, api_key=ai_key) + logger = SessionLogger.create(log_file) if log_file else None if analyze or interactive: console.print(f"[cyan]AI:[/cyan] {AIClient(ai_config).summary()}") @@ -133,6 +142,7 @@ def run( analyze=analyze, interactive=interactive, ai_config=ai_config, + logger=logger, ) ) except typer.Exit: @@ -154,13 +164,36 @@ async def _async_main( analyze: bool, interactive: bool, ai_config: AIConfig, + logger: SessionLogger | None, ) -> None: """Open a single SSH session and run probe / collection / analysis through it.""" client = SSHClient(config) + if logger is not None: + logger.log_event( + "session_start", + { + "host": req.host, + "port": req.port, + "issue": req.issue, + "probe": probe, + "collect": collect, + "analyze": analyze, + "interactive": interactive, + }, + ) async with client.connect() as session: if probe: result = await session.probe() _handle_probe_result(result) + if logger is not None: + logger.log_event( + "probe_result", + { + "exit_code": result.exit_code, + "stdout": result.stdout, + "stderr": result.stderr, + }, + ) report: CollectionReport | None = None if collect or analyze: @@ -168,12 +201,20 @@ async def _async_main( console.print(f"[cyan]Collecting diagnostics:[/cyan] {len(plan)} commands") report = await collect_from_plan(session, plan) _handle_collection_report(report) + if logger is not None: + logger.log_event( + "collection_summary", + { + "total": report.total, + "failed": report.failed, + }, + ) if analyze and report is not None: - _run_analysis(ai_config, req.issue, report) + _run_analysis(ai_config, req.issue, report, logger=logger) if interactive: - await _interactive_loop(session, req, ai_config, report) + await _interactive_loop(session, req, ai_config, report, logger=logger) async def _interactive_loop( @@ -181,6 +222,7 @@ async def _interactive_loop( req: TroubleshootRequest, ai_config: AIConfig, report: CollectionReport | None, + logger: SessionLogger | None, ) -> None: """Run a follow-up loop for collecting and conversational analysis.""" console.print( @@ -210,6 +252,8 @@ async def _interactive_loop( command = input("tai> ").strip() except (EOFError, KeyboardInterrupt): console.print("\n[yellow]Exiting interactive mode.[/yellow]") + if logger is not None: + logger.log_event("interactive_exit", {"reason": "signal_or_eof"}) return if not command: @@ -217,6 +261,8 @@ async def _interactive_loop( if command in {"/quit", "/exit"}: console.print("[green]Bye.[/green]") + if logger is not None: + logger.log_event("interactive_exit", {"reason": "user_quit"}) return if command == "/help": @@ -230,6 +276,14 @@ async def _interactive_loop( report = await collect_from_plan(session, plan) _handle_collection_report(report) messages = _reset_messages(report) + if logger is not None: + logger.log_event( + "collection_summary", + { + "total": report.total, + "failed": report.failed, + }, + ) continue if command == "/analyze": @@ -249,7 +303,7 @@ async def _interactive_loop( "content": "Provide an updated diagnosis from the current diagnostics.", } ) - response = _stream_conversation(ai, messages) + response = _stream_conversation(ai, messages, logger=logger) messages.append({"role": "assistant", "content": response}) continue @@ -265,7 +319,7 @@ async def _interactive_loop( continue messages.append({"role": "user", "content": command}) - response = _stream_conversation(ai, messages) + response = _stream_conversation(ai, messages, logger=logger) messages.append({"role": "assistant", "content": response}) @@ -293,7 +347,13 @@ def _handle_collection_report(report: CollectionReport) -> None: console.print(f"- {item.name}: {status}{trunc}") -def _run_analysis(ai_config: AIConfig, issue: str, report: CollectionReport) -> None: +def _run_analysis( + ai_config: AIConfig, + issue: str, + report: CollectionReport, + *, + logger: SessionLogger | None, +) -> None: """Send collected data to the AI and stream the analysis to stdout.""" console.print("[cyan]Analyzing...[/cyan]\n") ai = AIClient(ai_config) @@ -303,13 +363,29 @@ def _run_analysis(ai_config: AIConfig, issue: str, report: CollectionReport) -> chunks: list[str] = [] for chunk in ai.stream(system_prompt, user_message): chunks.append(chunk) - console.print(Markdown("".join(chunks))) + response = "".join(chunks) + console.print(Markdown(response)) + if logger is not None: + logger.log_event( + "analysis_response", + { + "issue": issue, + "response": response, + }, + ) except Exception as exc: # noqa: BLE001 console.print(f"[red]AI analysis failed:[/red] {exc}") + if logger is not None: + logger.log_event("analysis_error", {"error": str(exc)}) raise typer.Exit(code=1) from exc -def _stream_conversation(ai: AIClient, messages: list[dict[str, str]]) -> str: +def _stream_conversation( + ai: AIClient, + messages: list[dict[str, str]], + *, + logger: SessionLogger | None, +) -> str: """Stream a multi-turn AI response and return the final text.""" console.print("[cyan]Analyzing...[/cyan]\n") try: @@ -318,9 +394,19 @@ def _stream_conversation(ai: AIClient, messages: list[dict[str, str]]) -> str: chunks.append(chunk) response = "".join(chunks) console.print(Markdown(response)) + if logger is not None and messages: + logger.log_event( + "analysis_response", + { + "last_user_message": messages[-1].get("content", ""), + "response": response, + }, + ) return response except Exception as exc: # noqa: BLE001 console.print(f"[red]AI analysis failed:[/red] {exc}") + if logger is not None: + logger.log_event("analysis_error", {"error": str(exc)}) raise typer.Exit(code=1) from exc diff --git a/src/tai/session_log.py b/src/tai/session_log.py new file mode 100644 index 0000000..c3773ac --- /dev/null +++ b/src/tai/session_log.py @@ -0,0 +1,34 @@ +"""Structured session logging helpers for troubleshooting runs.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + + +@dataclass(slots=True) +class SessionLogger: + """Append JSONL events to a log file for post-run analysis.""" + + path: Path + + @classmethod + def create(cls, file_path: str) -> SessionLogger: + """Create a logger for *file_path*, ensuring parent directories exist.""" + path = Path(file_path).expanduser() + path.parent.mkdir(parents=True, exist_ok=True) + return cls(path=path) + + def log_event(self, event: str, payload: dict[str, Any]) -> None: + """Write one timestamped event row to the JSONL log.""" + row = { + "ts": datetime.now(UTC).isoformat(), + "event": event, + "payload": payload, + } + with self.path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(row, ensure_ascii=True)) + handle.write("\n") \ No newline at end of file diff --git a/tests/test_session_log.py b/tests/test_session_log.py new file mode 100644 index 0000000..a812cfe --- /dev/null +++ b/tests/test_session_log.py @@ -0,0 +1,22 @@ +"""Tests for structured session logging.""" + +from __future__ import annotations + +import json + +from tai.session_log import SessionLogger + + +def test_session_logger_writes_jsonl_row(tmp_path) -> None: # type: ignore[no-untyped-def] + log_path = tmp_path / "logs" / "session.jsonl" + logger = SessionLogger.create(str(log_path)) + + logger.log_event("analysis_response", {"response": "Root cause is X"}) + + lines = log_path.read_text(encoding="utf-8").splitlines() + assert len(lines) == 1 + + row = json.loads(lines[0]) + assert row["event"] == "analysis_response" + assert row["payload"]["response"] == "Root cause is X" + assert "ts" in row \ No newline at end of file