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