344 lines
17 KiB
Python
344 lines
17 KiB
Python
import aiosqlite
|
|
import asyncio
|
|
import os
|
|
import time
|
|
import datetime
|
|
import re
|
|
from typing import Dict, List, Any, Optional, Tuple
|
|
import chromadb
|
|
from chromadb.utils import embedding_functions
|
|
from sentence_transformers import SentenceTransformer
|
|
import logging
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# --- Helper Function for Keyword Scoring ---
|
|
def calculate_keyword_score(text: str, context: str) -> int:
|
|
"""Calculates a simple keyword overlap score."""
|
|
if not context or not text:
|
|
return 0
|
|
context_words = set(re.findall(r'\b\w+\b', context.lower()))
|
|
text_words = set(re.findall(r'\b\w+\b', text.lower()))
|
|
# Ignore very common words (basic stopword list)
|
|
stopwords = {"the", "a", "is", "in", "it", "of", "and", "to", "for", "on", "with", "that", "this", "i", "you", "me", "my", "your"}
|
|
context_words -= stopwords
|
|
text_words -= stopwords
|
|
if not context_words: # Avoid division by zero if context is only stopwords
|
|
return 0
|
|
overlap = len(context_words.intersection(text_words))
|
|
# Normalize score slightly by context length (more overlap needed for longer context)
|
|
# score = overlap / (len(context_words) ** 0.5) # Example normalization
|
|
score = overlap # Simpler score for now
|
|
return score
|
|
|
|
class MemoryManager:
|
|
"""Handles database interactions for Gurt's memory (facts and semantic)."""
|
|
|
|
def __init__(self, db_path: str, max_user_facts: int = 20, max_general_facts: int = 100, semantic_model_name: str = 'all-MiniLM-L6-v2', chroma_path: str = "data/chroma_db"):
|
|
self.db_path = db_path
|
|
self.max_user_facts = max_user_facts
|
|
self.max_general_facts = max_general_facts
|
|
self.db_lock = asyncio.Lock() # Lock for SQLite operations
|
|
|
|
# Ensure data directories exist
|
|
os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
|
|
os.makedirs(chroma_path, exist_ok=True)
|
|
logger.info(f"MemoryManager initialized with db_path: {self.db_path}, chroma_path: {chroma_path}")
|
|
|
|
# --- Semantic Memory Setup ---
|
|
self.chroma_path = chroma_path
|
|
self.semantic_model_name = semantic_model_name
|
|
self.chroma_client = None
|
|
self.embedding_function = None
|
|
self.semantic_collection = None
|
|
self.transformer_model = None
|
|
self._initialize_semantic_memory_sync() # Initialize semantic components synchronously for simplicity during init
|
|
|
|
def _initialize_semantic_memory_sync(self):
|
|
"""Synchronously initializes ChromaDB client, model, and collection."""
|
|
try:
|
|
logger.info("Initializing ChromaDB client...")
|
|
# Use PersistentClient for saving data to disk
|
|
self.chroma_client = chromadb.PersistentClient(path=self.chroma_path)
|
|
|
|
logger.info(f"Loading Sentence Transformer model: {self.semantic_model_name}...")
|
|
# Load the model directly
|
|
self.transformer_model = SentenceTransformer(self.semantic_model_name)
|
|
|
|
# Create a custom embedding function using the loaded model
|
|
class CustomEmbeddingFunction(embedding_functions.EmbeddingFunction):
|
|
def __init__(self, model):
|
|
self.model = model
|
|
def __call__(self, input: chromadb.Documents) -> chromadb.Embeddings:
|
|
# Ensure input is a list of strings
|
|
if not isinstance(input, list):
|
|
input = [str(input)] # Convert single item to list
|
|
elif not all(isinstance(item, str) for item in input):
|
|
input = [str(item) for item in input] # Ensure all items are strings
|
|
|
|
logger.debug(f"Generating embeddings for {len(input)} documents.")
|
|
embeddings = self.model.encode(input, show_progress_bar=False).tolist()
|
|
logger.debug(f"Generated {len(embeddings)} embeddings.")
|
|
return embeddings
|
|
|
|
self.embedding_function = CustomEmbeddingFunction(self.transformer_model)
|
|
|
|
logger.info("Getting/Creating ChromaDB collection 'gurt_semantic_memory'...")
|
|
# Get or create the collection with the custom embedding function
|
|
self.semantic_collection = self.chroma_client.get_or_create_collection(
|
|
name="gurt_semantic_memory",
|
|
embedding_function=self.embedding_function,
|
|
metadata={"hnsw:space": "cosine"} # Use cosine distance for similarity
|
|
)
|
|
logger.info("ChromaDB collection initialized successfully.")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize semantic memory: {e}", exc_info=True)
|
|
# Set components to None to indicate failure
|
|
self.chroma_client = None
|
|
self.transformer_model = None
|
|
self.embedding_function = None
|
|
self.semantic_collection = None
|
|
|
|
async def initialize_sqlite_database(self):
|
|
"""Initializes the SQLite database and creates tables if they don't exist."""
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
await db.execute("PRAGMA journal_mode=WAL;")
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS user_facts (
|
|
user_id TEXT NOT NULL,
|
|
fact TEXT NOT NULL,
|
|
timestamp REAL DEFAULT (unixepoch('now')),
|
|
PRIMARY KEY (user_id, fact)
|
|
);
|
|
""")
|
|
await db.execute("CREATE INDEX IF NOT EXISTS idx_user_facts_user ON user_facts (user_id);")
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS general_facts (
|
|
fact TEXT PRIMARY KEY NOT NULL,
|
|
timestamp REAL DEFAULT (unixepoch('now'))
|
|
);
|
|
""")
|
|
# Removed channel/user state tables for brevity, can be added back if needed
|
|
await db.commit()
|
|
logger.info(f"SQLite database initialized/verified at {self.db_path}")
|
|
|
|
# --- SQLite Helper Methods ---
|
|
async def _db_execute(self, sql: str, params: tuple = ()):
|
|
async with self.db_lock:
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
await db.execute(sql, params)
|
|
await db.commit()
|
|
|
|
async def _db_fetchone(self, sql: str, params: tuple = ()) -> Optional[tuple]:
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
async with db.execute(sql, params) as cursor:
|
|
return await cursor.fetchone()
|
|
|
|
async def _db_fetchall(self, sql: str, params: tuple = ()) -> List[tuple]:
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
async with db.execute(sql, params) as cursor:
|
|
return await cursor.fetchall()
|
|
|
|
# --- User Fact Memory Methods (SQLite + Relevance) ---
|
|
|
|
async def add_user_fact(self, user_id: str, fact: str) -> Dict[str, Any]:
|
|
"""Stores a fact about a user in the SQLite database, enforcing limits."""
|
|
if not user_id or not fact:
|
|
return {"error": "user_id and fact are required."}
|
|
logger.info(f"Attempting to add user fact for {user_id}: '{fact}'")
|
|
try:
|
|
existing = await self._db_fetchone("SELECT 1 FROM user_facts WHERE user_id = ? AND fact = ?", (user_id, fact))
|
|
if existing:
|
|
logger.info(f"Fact already known for user {user_id}.")
|
|
return {"status": "duplicate", "user_id": user_id, "fact": fact}
|
|
|
|
count_result = await self._db_fetchone("SELECT COUNT(*) FROM user_facts WHERE user_id = ?", (user_id,))
|
|
current_count = count_result[0] if count_result else 0
|
|
|
|
status = "added"
|
|
if current_count >= self.max_user_facts:
|
|
logger.warning(f"User {user_id} fact limit ({self.max_user_facts}) reached. Deleting oldest.")
|
|
oldest_fact_row = await self._db_fetchone("SELECT fact FROM user_facts WHERE user_id = ? ORDER BY timestamp ASC LIMIT 1", (user_id,))
|
|
if oldest_fact_row:
|
|
await self._db_execute("DELETE FROM user_facts WHERE user_id = ? AND fact = ?", (user_id, oldest_fact_row[0]))
|
|
logger.info(f"Deleted oldest fact for user {user_id}: '{oldest_fact_row[0]}'")
|
|
status = "limit_reached" # Indicate limit was hit but fact was added
|
|
|
|
await self._db_execute("INSERT INTO user_facts (user_id, fact) VALUES (?, ?)", (user_id, fact))
|
|
logger.info(f"Fact added for user {user_id}.")
|
|
return {"status": status, "user_id": user_id, "fact_added": fact}
|
|
|
|
except Exception as e:
|
|
logger.error(f"SQLite error adding user fact for {user_id}: {e}", exc_info=True)
|
|
return {"error": f"Database error adding user fact: {str(e)}"}
|
|
|
|
async def get_user_facts(self, user_id: str, context: Optional[str] = None) -> List[str]:
|
|
"""Retrieves stored facts about a user, optionally scored by relevance to context."""
|
|
if not user_id:
|
|
logger.warning("get_user_facts called without user_id.")
|
|
return []
|
|
logger.info(f"Retrieving facts for user {user_id} (context provided: {bool(context)})")
|
|
try:
|
|
rows = await self._db_fetchall("SELECT fact FROM user_facts WHERE user_id = ?", (user_id,))
|
|
user_facts = [row[0] for row in rows]
|
|
|
|
if context and user_facts:
|
|
# Score facts based on context if provided
|
|
scored_facts = []
|
|
for fact in user_facts:
|
|
score = calculate_keyword_score(fact, context)
|
|
scored_facts.append({"fact": fact, "score": score})
|
|
|
|
# Sort by score (descending), then fallback to original order (implicitly newest first if DB returns that way)
|
|
scored_facts.sort(key=lambda x: x["score"], reverse=True)
|
|
# Return top N facts based on score
|
|
return [item["fact"] for item in scored_facts[:self.max_user_facts]]
|
|
else:
|
|
# No context or no facts, return newest N facts (assuming DB returns in insertion order or we add ORDER BY timestamp DESC)
|
|
# Let's add ORDER BY timestamp DESC to be explicit
|
|
rows_ordered = await self._db_fetchall(
|
|
"SELECT fact FROM user_facts WHERE user_id = ? ORDER BY timestamp DESC LIMIT ?",
|
|
(user_id, self.max_user_facts)
|
|
)
|
|
return [row[0] for row in rows_ordered]
|
|
|
|
except Exception as e:
|
|
logger.error(f"SQLite error retrieving user facts for {user_id}: {e}", exc_info=True)
|
|
return []
|
|
|
|
# --- General Fact Memory Methods (SQLite + Relevance) ---
|
|
|
|
async def add_general_fact(self, fact: str) -> Dict[str, Any]:
|
|
"""Stores a general fact in the SQLite database, enforcing limits."""
|
|
if not fact:
|
|
return {"error": "fact is required."}
|
|
logger.info(f"Attempting to add general fact: '{fact}'")
|
|
try:
|
|
existing = await self._db_fetchone("SELECT 1 FROM general_facts WHERE fact = ?", (fact,))
|
|
if existing:
|
|
logger.info(f"General fact already known: '{fact}'")
|
|
return {"status": "duplicate", "fact": fact}
|
|
|
|
count_result = await self._db_fetchone("SELECT COUNT(*) FROM general_facts", ())
|
|
current_count = count_result[0] if count_result else 0
|
|
|
|
status = "added"
|
|
if current_count >= self.max_general_facts:
|
|
logger.warning(f"General fact limit ({self.max_general_facts}) reached. Deleting oldest.")
|
|
oldest_fact_row = await self._db_fetchone("SELECT fact FROM general_facts ORDER BY timestamp ASC LIMIT 1", ())
|
|
if oldest_fact_row:
|
|
await self._db_execute("DELETE FROM general_facts WHERE fact = ?", (oldest_fact_row[0],))
|
|
logger.info(f"Deleted oldest general fact: '{oldest_fact_row[0]}'")
|
|
status = "limit_reached"
|
|
|
|
await self._db_execute("INSERT INTO general_facts (fact) VALUES (?)", (fact,))
|
|
logger.info(f"General fact added: '{fact}'")
|
|
return {"status": status, "fact_added": fact}
|
|
|
|
except Exception as e:
|
|
logger.error(f"SQLite error adding general fact: {e}", exc_info=True)
|
|
return {"error": f"Database error adding general fact: {str(e)}"}
|
|
|
|
async def get_general_facts(self, query: Optional[str] = None, limit: Optional[int] = 10, context: Optional[str] = None) -> List[str]:
|
|
"""Retrieves stored general facts, optionally filtering and scoring by relevance."""
|
|
logger.info(f"Retrieving general facts (query='{query}', limit={limit}, context provided: {bool(context)})")
|
|
limit = min(max(1, limit or 10), 50)
|
|
|
|
try:
|
|
sql = "SELECT fact FROM general_facts"
|
|
params = []
|
|
if query:
|
|
sql += " WHERE fact LIKE ?"
|
|
params.append(f"%{query}%")
|
|
|
|
# Fetch all matching facts first for scoring
|
|
rows = await self._db_fetchall(sql, tuple(params))
|
|
all_facts = [row[0] for row in rows]
|
|
|
|
if context and all_facts:
|
|
# Score facts based on context
|
|
scored_facts = []
|
|
for fact in all_facts:
|
|
score = calculate_keyword_score(fact, context)
|
|
scored_facts.append({"fact": fact, "score": score})
|
|
|
|
# Sort by score (descending)
|
|
scored_facts.sort(key=lambda x: x["score"], reverse=True)
|
|
# Return top N facts based on score
|
|
return [item["fact"] for item in scored_facts[:limit]]
|
|
else:
|
|
# No context or no facts, return newest N facts matching query (if any)
|
|
sql += " ORDER BY timestamp DESC LIMIT ?"
|
|
params.append(limit)
|
|
rows_ordered = await self._db_fetchall(sql, tuple(params))
|
|
return [row[0] for row in rows_ordered]
|
|
|
|
except Exception as e:
|
|
logger.error(f"SQLite error retrieving general facts: {e}", exc_info=True)
|
|
return []
|
|
|
|
# --- Semantic Memory Methods (ChromaDB) ---
|
|
|
|
async def add_message_embedding(self, message_id: str, text: str, metadata: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Generates embedding and stores a message in ChromaDB."""
|
|
if not self.semantic_collection:
|
|
return {"error": "Semantic memory (ChromaDB) is not initialized."}
|
|
if not text:
|
|
return {"error": "Cannot add empty text to semantic memory."}
|
|
|
|
logger.info(f"Adding message {message_id} to semantic memory.")
|
|
try:
|
|
# ChromaDB expects lists for inputs
|
|
await asyncio.to_thread(
|
|
self.semantic_collection.add,
|
|
documents=[text],
|
|
metadatas=[metadata],
|
|
ids=[message_id]
|
|
)
|
|
logger.info(f"Successfully added message {message_id} to ChromaDB.")
|
|
return {"status": "success", "message_id": message_id}
|
|
except Exception as e:
|
|
logger.error(f"ChromaDB error adding message {message_id}: {e}", exc_info=True)
|
|
return {"error": f"Semantic memory error adding message: {str(e)}"}
|
|
|
|
async def search_semantic_memory(self, query_text: str, n_results: int = 5, filter_metadata: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
|
|
"""Searches ChromaDB for messages semantically similar to the query text."""
|
|
if not self.semantic_collection:
|
|
logger.warning("Search semantic memory called, but ChromaDB is not initialized.")
|
|
return []
|
|
if not query_text:
|
|
logger.warning("Search semantic memory called with empty query text.")
|
|
return []
|
|
|
|
logger.info(f"Searching semantic memory (n_results={n_results}, filter={filter_metadata}) for query: '{query_text[:50]}...'")
|
|
try:
|
|
# Perform the query in a separate thread as ChromaDB operations can be blocking
|
|
results = await asyncio.to_thread(
|
|
self.semantic_collection.query,
|
|
query_texts=[query_text],
|
|
n_results=n_results,
|
|
where=filter_metadata, # Optional filter based on metadata
|
|
include=['metadatas', 'documents', 'distances'] # Include distance for relevance
|
|
)
|
|
logger.debug(f"ChromaDB query results: {results}")
|
|
|
|
# Process results
|
|
processed_results = []
|
|
if results and results.get('ids') and results['ids'][0]:
|
|
for i, doc_id in enumerate(results['ids'][0]):
|
|
processed_results.append({
|
|
"id": doc_id,
|
|
"document": results['documents'][0][i] if results.get('documents') else None,
|
|
"metadata": results['metadatas'][0][i] if results.get('metadatas') else None,
|
|
"distance": results['distances'][0][i] if results.get('distances') else None,
|
|
})
|
|
logger.info(f"Found {len(processed_results)} semantic results.")
|
|
return processed_results
|
|
|
|
except Exception as e:
|
|
logger.error(f"ChromaDB error searching memory for query '{query_text[:50]}...': {e}", exc_info=True)
|
|
return []
|