refactored: to utilise the google adk and production grade agent
Some checks failed
validation / verify (push) Failing after 10s
Some checks failed
validation / verify (push) Failing after 10s
This commit is contained in:
355
app/database.py
Normal file
355
app/database.py
Normal file
@@ -0,0 +1,355 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user