# src/memory.py
import sqlite3
from datetime import datetime, timedelta
import json
from typing import List, Dict, Any, Tuple
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import logging
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
class MemoryManager:
def __init__(self, db_path: str):
self.conn = sqlite3.connect(db_path)
self.cursor = self.conn.cursor()
self.create_tables()
self.vectorizer = TfidfVectorizer(stop_words='english')
logging.info("MemoryManager initialized and tables created.")
def create_tables(self):
# Create tables if they don't exist
self.cursor.execute('''CREATE TABLE IF NOT EXISTS semantic_memory
(id INTEGER PRIMARY KEY, concept TEXT, description TEXT, last_accessed DATETIME, tags TEXT, importance REAL DEFAULT 0.5)''')
# Add tags and importance columns if they don't exist
self.cursor.execute("PRAGMA table_info(semantic_memory)")
columns = [column[1] for column in self.cursor.fetchall()]
if 'tags' not in columns:
self.cursor.execute("ALTER TABLE semantic_memory ADD COLUMN tags TEXT")
if 'importance' not in columns:
self.cursor.execute("ALTER TABLE semantic_memory ADD COLUMN importance REAL DEFAULT 0.5")
self.cursor.execute('''CREATE INDEX IF NOT EXISTS idx_semantic_concept ON semantic_memory (concept)''')
self.cursor.execute('''CREATE INDEX IF NOT EXISTS idx_semantic_last_accessed ON semantic_memory (last_accessed)''')
self.cursor.execute('''CREATE INDEX IF NOT EXISTS idx_semantic_tags ON semantic_memory (tags)''')
# Create table for user interactions
self.cursor.execute('''CREATE TABLE IF NOT EXISTS user_interactions
(user_id TEXT, query TEXT, response TEXT, timestamp DATETIME)''')
self.cursor.execute('''CREATE INDEX IF NOT EXISTS idx_user_interactions_timestamp ON user_interactions (timestamp)''')
self.conn.commit()
logging.info("Tables and indexes created successfully.")
def add_semantic_memory(self, concept: str, description: str, tags: List[str] = None):
if tags is None:
tags = []
tags_str = json.dumps(tags)
self.cursor.execute("INSERT INTO semantic_memory (concept, description, last_accessed, tags) VALUES (?, ?, ?, ?)",
(concept, description, datetime.now().isoformat(), tags_str))
self.conn.commit()
logging.info("Semantic memory added.")
def retrieve_relevant_memories(self, query: str, limit: int = 30) -> List[Dict[str, Any]]:
all_memories = self._get_all_memories()
# Handle empty or stop-word-only query
if not query.strip() or self.vectorizer.stop_words and all(word in self.vectorizer.stop_words for word in query.split()):
return []
scored_memories = self._score_memories(query, all_memories)
return [memory for memory, score in sorted(scored_memories, key=lambda x: x[1], reverse=True)[:limit]]
def _get_all_memories(self) -> List[Tuple[Dict[str, Any], datetime]]:
self.cursor.execute("SELECT concept, description, importance, last_accessed, tags FROM semantic_memory ORDER BY importance DESC, last_accessed DESC")
semantic_memories = self.cursor.fetchall()
all_memories = [({"concept": concept, "description": description, "importance": importance},
datetime.fromisoformat(last_accessed), json.loads(tags) if tags else None) for concept, description, importance, last_accessed, tags in semantic_memories]
return all_memories
def _score_memories(self, query: str, memories: List[Tuple[Dict[str, Any], datetime, List[str]]]) -> List[Tuple[Dict[str, Any], float]]:
query_vector = self.vectorizer.fit_transform([query])
scored_memories = []
for memory, timestamp, tags in memories:
text = f"{memory['concept']} {memory['description']}"
importance = memory.get('importance', 0.5)
memory_vector = self.vectorizer.transform([text])
similarity = cosine_similarity(query_vector, memory_vector)[0][0]
if timestamp:
recency = 1 / (1 + (datetime.now() - timestamp).total_seconds() / 60) # Favor recent memories
else:
recency = 0.5 # Neutral recency for semantic memories
score = (similarity + importance + recency) / 3
scored_memories.append((memory, score))
return scored_memories
def section_exists(self, concept: str) -> bool:
# Normalize the concept to lowercase
concept = concept.lower()
self.cursor.execute("SELECT COUNT(*) FROM semantic_memory WHERE concept LIKE ?", (f"{concept}%",))
count = self.cursor.fetchone()[0]
return count > 0
def add_user_interaction(self, user_id: str, query: str, response: str):
self.cursor.execute("INSERT INTO user_interactions (user_id, query, response, timestamp) VALUES (?, ?, ?, ?)",
(user_id, query, response, datetime.now().isoformat()))
self.conn.commit()
logging.info(f"User interaction added: User ID: {user_id}, Query: {query}, Response: {response}")
def get_user_interactions(self, user_id: str) -> List[Dict[str, Any]]:
self.cursor.execute("SELECT query, response, timestamp FROM user_interactions WHERE user_id = ?", (user_id,))
interactions = self.cursor.fetchall()
return [{"query": query, "response": response, "timestamp": timestamp} for query, response, timestamp in interactions]
def cleanup_expired_interactions(self):
cutoff_time = datetime.now() - timedelta(minutes=5)
self.cursor.execute("DELETE FROM user_interactions WHERE timestamp < ?", (cutoff_time.isoformat(),))
self.conn.commit()
logging.info(f"Expired user interactions cleaned up. Cutoff time: {cutoff_time}")
def get_section_description(self, section_name: str) -> str:
# Normalize the section name to lowercase
section_name = section_name.lower()
# Retrieve the specific section from the database
self.cursor.execute("SELECT description FROM semantic_memory WHERE concept LIKE ?", (f"{section_name}%",))
result = self.cursor.fetchone()
if result:
logging.info(f"Found section: {section_name}")
return result[0]
else:
logging.warning(f"Section not found: {section_name}")
return ""
def count_chroniques(self) -> int:
# Count the number of chroniques in the database
self.cursor.execute("SELECT COUNT(*) FROM semantic_memory WHERE concept LIKE 'chronique #%'")
count = self.cursor.fetchone()[0]
logging.info(f"Number of chroniques: {count}")
return count
def count_flash_infos(self) -> int:
# Count the number of flash infos in the database
self.cursor.execute("SELECT COUNT(*) FROM semantic_memory WHERE concept LIKE 'flash info fl-%'")
count = self.cursor.fetchone()[0]
logging.info(f"Number of flash infos: {count}")
return count
def count_chronique_faqs(self) -> int:
# Count the number of chronique-faqs in the database
self.cursor.execute("SELECT COUNT(*) FROM semantic_memory WHERE concept LIKE 'chronique-faq #%'")
count = self.cursor.fetchone()[0]
logging.info(f"Number of chronique-faqs: {count}")
return count
if __name__ == "__main__":
db_path = "agent.db"
memory_manager = MemoryManager(db_path)
memory_manager.cleanup_expired_interactions()