zphinx/tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

feat(cli): add structured JSONL session logging for AI output

Browse Source
This commit is contained in:
zphinx
2026-05-04 06:03:39 +02:00
parent fdcde37e46
commit 2662d1b253
3 changed files with 149 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

100
src/tai/cli.py
View File

@@ -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

34
src/tai/session_log.py Normal file
View File

@@ -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")

22
tests/test_session_log.py Normal file
View File

@@ -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