356 lines
15 KiB
Python
356 lines
15 KiB
Python
"""PostgreSQL Database Persistence Manager for GCP Solution Architecture Agent.
|
|
|
|
Manages connection pooling, schema initialization, and database persistence
|
|
for ADK sessions, workflow executions, artifacts, evaluations, and orchestrator review logs.
|
|
Supports external PostgreSQL database with SQLite fallback for offline development/testing.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import sqlite3
|
|
from datetime import datetime, timezone
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
try:
|
|
import psycopg2
|
|
from psycopg2.extras import RealDictCursor
|
|
HAS_PSYCOPG2 = True
|
|
except ImportError:
|
|
HAS_PSYCOPG2 = False
|
|
|
|
from app.config import get_settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DatabaseManager:
|
|
"""Handles persistence of sessions, execution runs, artifacts, and review logs."""
|
|
|
|
def __init__(self, db_url: Optional[str] = None) -> None:
|
|
self.settings = get_settings()
|
|
self.db_url = db_url or self.settings.get_database_url()
|
|
self.is_postgres = self.db_url.startswith("postgresql://") or self.db_url.startswith("postgres://")
|
|
self._init_db()
|
|
|
|
def _get_connection(self):
|
|
"""Get database connection (PostgreSQL if available/accessible, else SQLite)."""
|
|
if self.is_postgres and HAS_PSYCOPG2:
|
|
try:
|
|
conn = psycopg2.connect(self.db_url, connect_timeout=3)
|
|
return conn
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"PostgreSQL connection to %s failed (%s); falling back to local SQLite database %s",
|
|
self.db_url,
|
|
exc,
|
|
self.settings.SQLITE_FALLBACK_DB,
|
|
)
|
|
|
|
# SQLite fallback connection
|
|
conn = sqlite3.connect(self.settings.SQLITE_FALLBACK_DB)
|
|
conn.row_factory = sqlite3.Row
|
|
return conn
|
|
|
|
def _init_db(self) -> None:
|
|
"""Initialize database schema tables if they do not exist."""
|
|
conn = self._get_connection()
|
|
try:
|
|
cursor = conn.cursor()
|
|
|
|
# 1. Sessions Table
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS adk_sessions (
|
|
session_id VARCHAR(128) PRIMARY KEY,
|
|
agent_name VARCHAR(128) NOT NULL,
|
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
state_data TEXT NOT NULL,
|
|
metadata TEXT
|
|
)
|
|
""")
|
|
|
|
# 2. Workflow Executions Table
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS workflow_executions (
|
|
execution_id VARCHAR(128) PRIMARY KEY,
|
|
session_id VARCHAR(128) NOT NULL,
|
|
status VARCHAR(32) NOT NULL,
|
|
current_phase VARCHAR(64) NOT NULL,
|
|
loop_count INTEGER DEFAULT 0,
|
|
started_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
completed_at TIMESTAMP WITH TIME ZONE,
|
|
request_summary TEXT,
|
|
results_json TEXT
|
|
)
|
|
""")
|
|
|
|
# 3. Artifacts Table
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS adk_artifacts (
|
|
artifact_id VARCHAR(128) PRIMARY KEY,
|
|
execution_id VARCHAR(128) NOT NULL,
|
|
artifact_type VARCHAR(64) NOT NULL,
|
|
file_path VARCHAR(256) NOT NULL,
|
|
content TEXT NOT NULL,
|
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
|
|
# 4. Review & Orchestration Loop Logs Table
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS orchestrator_review_logs (
|
|
log_id VARCHAR(128) PRIMARY KEY,
|
|
execution_id VARCHAR(128) NOT NULL,
|
|
iteration INTEGER NOT NULL,
|
|
review_status VARCHAR(32) NOT NULL,
|
|
reviewer_agent VARCHAR(64) NOT NULL,
|
|
feedback TEXT,
|
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
|
|
# 5. Evaluations Table
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS adk_evaluations (
|
|
eval_id VARCHAR(128) PRIMARY KEY,
|
|
case_id VARCHAR(128) NOT NULL,
|
|
total_score FLOAT NOT NULL,
|
|
max_score FLOAT NOT NULL,
|
|
pass_rate FLOAT NOT NULL,
|
|
passed BOOLEAN NOT NULL,
|
|
details_json TEXT NOT NULL,
|
|
evaluated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
|
|
conn.commit()
|
|
logger.info("Database schema initialized successfully.")
|
|
except Exception as exc:
|
|
logger.error("Failed to initialize database schema: %s", exc)
|
|
finally:
|
|
conn.close()
|
|
|
|
def save_session(self, session_id: str, agent_name: str, state_data: Dict[str, Any], metadata: Optional[Dict[str, Any]] = None) -> bool:
|
|
"""Persist or update ADK session state in PostgreSQL."""
|
|
conn = self._get_connection()
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
try:
|
|
cursor = conn.cursor()
|
|
state_json = json.dumps(state_data)
|
|
meta_json = json.dumps(metadata or {})
|
|
|
|
if isinstance(conn, sqlite3.Connection):
|
|
cursor.execute("""
|
|
INSERT INTO adk_sessions (session_id, agent_name, created_at, updated_at, state_data, metadata)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
ON CONFLICT(session_id) DO UPDATE SET
|
|
updated_at = excluded.updated_at,
|
|
state_data = excluded.state_data,
|
|
metadata = excluded.metadata
|
|
""", (session_id, agent_name, now, now, state_json, meta_json))
|
|
else:
|
|
cursor.execute("""
|
|
INSERT INTO adk_sessions (session_id, agent_name, created_at, updated_at, state_data, metadata)
|
|
VALUES (%s, %s, %s, %s, %s, %s)
|
|
ON CONFLICT(session_id) DO UPDATE SET
|
|
updated_at = EXCLUDED.updated_at,
|
|
state_data = EXCLUDED.state_data,
|
|
metadata = EXCLUDED.metadata
|
|
""", (session_id, agent_name, now, now, state_json, meta_json))
|
|
|
|
conn.commit()
|
|
return True
|
|
except Exception as exc:
|
|
logger.error("Error saving session %s: %s", session_id, exc)
|
|
return False
|
|
finally:
|
|
conn.close()
|
|
|
|
def get_session(self, session_id: str) -> Optional[Dict[str, Any]]:
|
|
"""Retrieve ADK session state by session_id."""
|
|
conn = self._get_connection()
|
|
try:
|
|
cursor = conn.cursor()
|
|
if isinstance(conn, sqlite3.Connection):
|
|
cursor.execute("SELECT * FROM adk_sessions WHERE session_id = ?", (session_id,))
|
|
else:
|
|
cursor.execute("SELECT * FROM adk_sessions WHERE session_id = %s", (session_id,))
|
|
|
|
row = cursor.fetchone()
|
|
if not row:
|
|
return None
|
|
|
|
if isinstance(conn, sqlite3.Connection):
|
|
return {
|
|
"session_id": row["session_id"],
|
|
"agent_name": row["agent_name"],
|
|
"state_data": json.loads(row["state_data"]),
|
|
"metadata": json.loads(row["metadata"] or "{}"),
|
|
"updated_at": row["updated_at"],
|
|
}
|
|
else:
|
|
return {
|
|
"session_id": row[0],
|
|
"agent_name": row[1],
|
|
"created_at": str(row[2]),
|
|
"updated_at": str(row[3]),
|
|
"state_data": json.loads(row[4]),
|
|
"metadata": json.loads(row[5] or "{}"),
|
|
}
|
|
except Exception as exc:
|
|
logger.error("Error reading session %s: %s", session_id, exc)
|
|
return None
|
|
finally:
|
|
conn.close()
|
|
|
|
def save_workflow_execution(
|
|
self,
|
|
execution_id: str,
|
|
session_id: str,
|
|
status: str,
|
|
current_phase: str,
|
|
loop_count: int,
|
|
request_summary: str,
|
|
results: Dict[str, Any],
|
|
) -> bool:
|
|
"""Save workflow execution record to database."""
|
|
conn = self._get_connection()
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
try:
|
|
cursor = conn.cursor()
|
|
results_json = json.dumps(results)
|
|
|
|
if isinstance(conn, sqlite3.Connection):
|
|
cursor.execute("""
|
|
INSERT INTO workflow_executions
|
|
(execution_id, session_id, status, current_phase, loop_count, started_at, completed_at, request_summary, results_json)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
ON CONFLICT(execution_id) DO UPDATE SET
|
|
status = excluded.status,
|
|
current_phase = excluded.current_phase,
|
|
loop_count = excluded.loop_count,
|
|
completed_at = excluded.completed_at,
|
|
results_json = excluded.results_json
|
|
""", (execution_id, session_id, status, current_phase, loop_count, now, now, request_summary, results_json))
|
|
else:
|
|
cursor.execute("""
|
|
INSERT INTO workflow_executions
|
|
(execution_id, session_id, status, current_phase, loop_count, started_at, completed_at, request_summary, results_json)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
ON CONFLICT(execution_id) DO UPDATE SET
|
|
status = EXCLUDED.status,
|
|
current_phase = EXCLUDED.current_phase,
|
|
loop_count = EXCLUDED.loop_count,
|
|
completed_at = EXCLUDED.completed_at,
|
|
results_json = EXCLUDED.results_json
|
|
""", (execution_id, session_id, status, current_phase, loop_count, now, now, request_summary, results_json))
|
|
|
|
conn.commit()
|
|
return True
|
|
except Exception as exc:
|
|
logger.error("Error saving workflow execution %s: %s", execution_id, exc)
|
|
return False
|
|
finally:
|
|
conn.close()
|
|
|
|
def save_artifact(self, artifact_id: str, execution_id: str, artifact_type: str, file_path: str, content: str) -> bool:
|
|
"""Write ADK artifact metadata & content to database."""
|
|
conn = self._get_connection()
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
try:
|
|
cursor = conn.cursor()
|
|
if isinstance(conn, sqlite3.Connection):
|
|
cursor.execute("""
|
|
INSERT INTO adk_artifacts (artifact_id, execution_id, artifact_type, file_path, content, created_at)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
ON CONFLICT(artifact_id) DO UPDATE SET
|
|
content = excluded.content,
|
|
file_path = excluded.file_path
|
|
""", (artifact_id, execution_id, artifact_type, file_path, content, now))
|
|
else:
|
|
cursor.execute("""
|
|
INSERT INTO adk_artifacts (artifact_id, execution_id, artifact_type, file_path, content, created_at)
|
|
VALUES (%s, %s, %s, %s, %s, %s)
|
|
ON CONFLICT(artifact_id) DO UPDATE SET
|
|
content = EXCLUDED.content,
|
|
file_path = EXCLUDED.file_path
|
|
""", (artifact_id, execution_id, artifact_type, file_path, content, now))
|
|
|
|
conn.commit()
|
|
return True
|
|
except Exception as exc:
|
|
logger.error("Error saving artifact %s: %s", artifact_id, exc)
|
|
return False
|
|
finally:
|
|
conn.close()
|
|
|
|
def record_orchestrator_log(
|
|
self,
|
|
log_id: str,
|
|
execution_id: str,
|
|
iteration: int,
|
|
review_status: str,
|
|
reviewer_agent: str,
|
|
feedback: str,
|
|
) -> bool:
|
|
"""Record orchestrator loop iteration review event into PostgreSQL log table."""
|
|
conn = self._get_connection()
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
try:
|
|
cursor = conn.cursor()
|
|
if isinstance(conn, sqlite3.Connection):
|
|
cursor.execute("""
|
|
INSERT INTO orchestrator_review_logs (log_id, execution_id, iteration, review_status, reviewer_agent, feedback, created_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?)
|
|
""", (log_id, execution_id, iteration, review_status, reviewer_agent, feedback, now))
|
|
else:
|
|
cursor.execute("""
|
|
INSERT INTO orchestrator_review_logs (log_id, execution_id, iteration, review_status, reviewer_agent, feedback, created_at)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s)
|
|
""", (log_id, execution_id, iteration, review_status, reviewer_agent, feedback, now))
|
|
|
|
conn.commit()
|
|
return True
|
|
except Exception as exc:
|
|
logger.error("Error recording orchestrator review log %s: %s", log_id, exc)
|
|
return False
|
|
finally:
|
|
conn.close()
|
|
|
|
def save_evaluation(self, eval_id: str, case_id: str, total_score: float, max_score: float, pass_rate: float, passed: bool, details: Dict[str, Any]) -> bool:
|
|
"""Save ADK evaluation benchmark result to database."""
|
|
conn = self._get_connection()
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
try:
|
|
cursor = conn.cursor()
|
|
details_json = json.dumps(details)
|
|
if isinstance(conn, sqlite3.Connection):
|
|
cursor.execute("""
|
|
INSERT INTO adk_evaluations (eval_id, case_id, total_score, max_score, pass_rate, passed, details_json, evaluated_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
|
""", (eval_id, case_id, total_score, max_score, pass_rate, passed, details_json, now))
|
|
else:
|
|
cursor.execute("""
|
|
INSERT INTO adk_evaluations (eval_id, case_id, total_score, max_score, pass_rate, passed, details_json, evaluated_at)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
|
|
""", (eval_id, case_id, total_score, max_score, pass_rate, passed, details_json, now))
|
|
|
|
conn.commit()
|
|
return True
|
|
except Exception as exc:
|
|
logger.error("Error saving evaluation %s: %s", eval_id, exc)
|
|
return False
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
_db_manager: Optional[DatabaseManager] = None
|
|
|
|
|
|
def get_db_manager() -> DatabaseManager:
|
|
"""Get singleton DatabaseManager instance."""
|
|
global _db_manager
|
|
if _db_manager is None:
|
|
_db_manager = DatabaseManager()
|
|
return _db_manager
|