Raiff1982 commited on
Commit
01a0455
·
verified ·
1 Parent(s): 7863970

Rename database.py to codette_memory_sql.py

Browse files
Files changed (2) hide show
  1. codette_memory_sql.py +92 -0
  2. database.py +0 -49
codette_memory_sql.py ADDED
@@ -0,0 +1,92 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import logging
3
+ import pyodbc
4
+ import sqlite3
5
+ from typing import Optional, List, Tuple
6
+
7
+
8
+ def connect_to_memory_database() -> Optional[pyodbc.Connection]:
9
+ """Connect to Azure SQL Database using HF secret environment variables."""
10
+ try:
11
+ server = os.getenv("AZURE_SQL_SERVER")
12
+ database = os.getenv("AZURE_SQL_DATABASE")
13
+ username = os.getenv("AZURE_SQL_USERNAME")
14
+ password = os.getenv("AZURE_SQL_PASSWORD")
15
+ driver = "{ODBC Driver 17 for SQL Server}"
16
+
17
+ if not all([server, database, username, password]):
18
+ logging.warning("Azure SQL credentials missing. Falling back to SQLite.")
19
+ return None
20
+
21
+ conn_str = (
22
+ f"DRIVER={driver};SERVER={server};DATABASE={database};"
23
+ f"UID={username};PWD={password}"
24
+ )
25
+ conn = pyodbc.connect(conn_str)
26
+ logging.info("Connected to Azure SQL Database.")
27
+ return conn
28
+ except pyodbc.Error as e:
29
+ logging.error(f"Azure SQL connection error: {e}")
30
+ return None
31
+
32
+
33
+ def log_memory(conn, query: str, response: str) -> None:
34
+ try:
35
+ cursor = conn.cursor()
36
+ cursor.execute(
37
+ "CREATE TABLE IF NOT EXISTS codette_memory (query TEXT, response TEXT)"
38
+ )
39
+ cursor.execute(
40
+ "INSERT INTO codette_memory (query, response) VALUES (?, ?)", (query, response)
41
+ )
42
+ conn.commit()
43
+ logging.info("Memory logged successfully.")
44
+ except Exception as e:
45
+ logging.error(f"Failed to log memory: {e}")
46
+
47
+
48
+ def get_recent_entries(conn, limit: int = 5) -> List[Tuple[str, str]]:
49
+ try:
50
+ cursor = conn.cursor()
51
+ cursor.execute(
52
+ "SELECT TOP (?) query, response FROM codette_memory ORDER BY 1 DESC", (limit,)
53
+ )
54
+ return cursor.fetchall()
55
+ except Exception as e:
56
+ logging.warning(f"Failed to fetch memory: {e}")
57
+ return []
58
+
59
+
60
+ def ping_database(conn) -> bool:
61
+ try:
62
+ cursor = conn.cursor()
63
+ cursor.execute("SELECT 1")
64
+ return True
65
+ except Exception as e:
66
+ logging.warning(f"Ping failed: {e}")
67
+ return False
68
+
69
+
70
+ # SQLite fallback (in-memory or file-backed)
71
+ def get_sqlite_memory(fallback_path: Optional[str] = ":memory:") -> sqlite3.Connection:
72
+ conn = sqlite3.connect(fallback_path)
73
+ conn.execute("CREATE TABLE IF NOT EXISTS codette_memory (query TEXT, response TEXT)")
74
+ return conn
75
+
76
+
77
+ # Context manager for auto-closing
78
+ class MemoryConnection:
79
+ def __init__(self):
80
+ self.conn = connect_to_memory_database()
81
+ if self.conn is None:
82
+ self.conn = get_sqlite_memory()
83
+
84
+ def __enter__(self):
85
+ return self.conn
86
+
87
+ def __exit__(self, exc_type, exc_value, traceback):
88
+ try:
89
+ self.conn.close()
90
+ logging.info("Memory DB connection closed.")
91
+ except Exception as e:
92
+ logging.warning(f"Failed to close DB: {e}")
database.py DELETED
@@ -1,49 +0,0 @@
1
- import pyodbc
2
- import logging
3
- import os
4
- from dotenv import load_dotenv
5
-
6
- # Load environment variables from .env file
7
- load_dotenv()
8
-
9
- def connect_to_database() -> pyodbc.Connection:
10
- """Connect to the Azure SQL Database."""
11
- try:
12
- server = os.getenv('AZURE_SQL_SERVER')
13
- database = os.getenv('AZURE_SQL_DATABASE')
14
- username = os.getenv('AZURE_SQL_USERNAME')
15
- password = os.getenv('AZURE_SQL_PASSWORD')
16
- driver = '{ODBC Driver 17 for SQL Server}'
17
-
18
- conn_str = f'DRIVER={driver};SERVER={server};DATABASE={database};UID={username};PWD={password}'
19
- conn = pyodbc.connect(conn_str)
20
- logging.info(f"Connected to the Azure SQL Database at {server}")
21
- return conn
22
- except pyodbc.Error as e:
23
- logging.error(f"Error connecting to the Azure SQL Database: {e}")
24
- return None
25
-
26
- def close_database_connection(conn: pyodbc.Connection) -> None:
27
- """Close the database connection."""
28
- try:
29
- if conn:
30
- conn.close()
31
- logging.info("Database connection closed.")
32
- except pyodbc.Error as e:
33
- logging.error(f"Error closing the database connection: {e}")
34
-
35
- # Example usage with context management
36
- class DatabaseConnection:
37
- def __enter__(self) -> pyodbc.Connection:
38
- self.conn = connect_to_database()
39
- return self.conn
40
-
41
- def __exit__(self, exc_type, exc_value, traceback) -> None:
42
- close_database_connection(self.conn)
43
-
44
- # Example usage
45
- if __name__ == "__main__":
46
- with DatabaseConnection() as conn:
47
- if conn:
48
- # Perform database operations
49
- pass