import os import logging import pyodbc import sqlite3 from typing import Optional, List, Tuple def connect_to_memory_database() -> Optional[pyodbc.Connection]: """Connect to Azure SQL Database using HF secret environment variables.""" try: server = os.getenv("AZURE_SQL_SERVER") database = os.getenv("AZURE_SQL_DATABASE") username = os.getenv("AZURE_SQL_USERNAME") password = os.getenv("AZURE_SQL_PASSWORD") driver = "{ODBC Driver 17 for SQL Server}" if not all([server, database, username, password]): logging.warning("Azure SQL credentials missing. Falling back to SQLite.") return None conn_str = ( f"DRIVER={driver};SERVER={server};DATABASE={database};" f"UID={username};PWD={password}" ) conn = pyodbc.connect(conn_str) logging.info("Connected to Azure SQL Database.") return conn except pyodbc.Error as e: logging.error(f"Azure SQL connection error: {e}") return None def log_memory(conn, query: str, response: str) -> None: try: cursor = conn.cursor() cursor.execute( "CREATE TABLE IF NOT EXISTS codette_memory (query TEXT, response TEXT)" ) cursor.execute( "INSERT INTO codette_memory (query, response) VALUES (?, ?)", (query, response) ) conn.commit() logging.info("Memory logged successfully.") except Exception as e: logging.error(f"Failed to log memory: {e}") def get_recent_entries(conn, limit: int = 5) -> List[Tuple[str, str]]: try: cursor = conn.cursor() cursor.execute( "SELECT TOP (?) query, response FROM codette_memory ORDER BY 1 DESC", (limit,) ) return cursor.fetchall() except Exception as e: logging.warning(f"Failed to fetch memory: {e}") return [] def ping_database(conn) -> bool: try: cursor = conn.cursor() cursor.execute("SELECT 1") return True except Exception as e: logging.warning(f"Ping failed: {e}") return False # SQLite fallback (in-memory or file-backed) def get_sqlite_memory(fallback_path: Optional[str] = ":memory:") -> sqlite3.Connection: conn = sqlite3.connect(fallback_path) conn.execute("CREATE TABLE IF NOT EXISTS codette_memory (query TEXT, response TEXT)") return conn # Context manager for auto-closing class MemoryConnection: def __init__(self): self.conn = connect_to_memory_database() if self.conn is None: self.conn = get_sqlite_memory() def __enter__(self): return self.conn def __exit__(self, exc_type, exc_value, traceback): try: self.conn.close() logging.info("Memory DB connection closed.") except Exception as e: logging.warning(f"Failed to close DB: {e}")