92 lines
2.6 KiB
Python
92 lines
2.6 KiB
Python
"""
|
|
Lab 5.9: Set Up Agent Observability
|
|
|
|
Trace every tool call + LLM completion to a local SQLite database.
|
|
"""
|
|
|
|
import sqlite3
|
|
import json
|
|
from datetime import datetime, timezone
|
|
from typing import Any
|
|
|
|
|
|
# TODO 1: Define the database schema
|
|
DB_SCHEMA = """
|
|
CREATE TABLE IF NOT EXISTS tool_calls (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
session_id TEXT NOT NULL,
|
|
turn_number INTEGER NOT NULL,
|
|
tool_name TEXT NOT NULL,
|
|
tool_params TEXT NOT NULL,
|
|
tool_result TEXT,
|
|
duration_ms INTEGER,
|
|
timestamp TEXT NOT NULL
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS llm_completions (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
session_id TEXT NOT NULL,
|
|
turn_number INTEGER NOT NULL,
|
|
input_tokens INTEGER,
|
|
output_tokens INTEGER,
|
|
cached_tokens INTEGER,
|
|
cost_cents REAL,
|
|
timestamp TEXT NOT NULL
|
|
);
|
|
|
|
-- TODO: Add a sessions table with total cost, total tokens, status
|
|
"""
|
|
|
|
|
|
class ObservabilityTracker:
|
|
"""Tracks agent observability data to SQLite."""
|
|
|
|
def __init__(self, db_path: str = "agent_observability.db"):
|
|
self.db_path = db_path
|
|
self.session_id = f"session_{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}"
|
|
self._init_db()
|
|
|
|
def _init_db(self):
|
|
"""Initialize database with schema."""
|
|
# TODO: Connect to SQLite and create tables
|
|
pass
|
|
|
|
def log_tool_call(self, turn: int, tool_name: str, params: dict,
|
|
result: str, duration_ms: int):
|
|
"""Log a tool call."""
|
|
# TODO: Insert tool call record
|
|
pass
|
|
|
|
def log_llm_completion(self, turn: int, input_tokens: int,
|
|
output_tokens: int, cost_cents: float):
|
|
"""Log an LLM completion with token counts."""
|
|
# TODO: Insert LLM completion record
|
|
pass
|
|
|
|
def get_session_summary(self) -> dict:
|
|
"""Get total cost, tool calls, tokens for this session."""
|
|
# TODO: Query and return summary
|
|
pass
|
|
|
|
|
|
# TODO: Decorate your agent with observability
|
|
def with_observability(agent_func):
|
|
"""Decorator that wraps an agent with observability tracking."""
|
|
def wrapper(prompt: str):
|
|
tracker = ObservabilityTracker()
|
|
# TODO: Wrap the agent execution with tracking
|
|
result = agent_func(prompt)
|
|
return result
|
|
return wrapper
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Test the tracker
|
|
tracker = ObservabilityTracker()
|
|
tracker.log_tool_call(1, "read_file", {"path": "test.txt"}, "file contents", 150)
|
|
tracker.log_llm_completion(1, 500, 200, 0.003)
|
|
|
|
summary = tracker.get_session_summary()
|
|
print(f"Session: {tracker.session_id}")
|
|
print(f"Summary: {json.dumps(summary, indent=2)}")
|