Updated memory and connected.

The memory module has been updated to work with the new database files, some of those were fixed in this commit. Modules that rely on memory were also updated in suit.

src/api/routes/chat.py

@@ -30,7 +30,10 @@ async def chat_endpoint(
 		if not user_profile:
 			state.memory_system.create_user(request.user_id, request.user_role or "Anonymous")
 			if request.user_specialty:
-				state.memory_system.set_user_preference(request.user_id, "specialty", request.user_specialty)
+				state.memory_system.set_user_preferences(
+					request.user_id,
+					{"specialty": request.user_specialty}
+				)
 
 		# Get or create session
 		session = state.memory_system.get_session(request.session_id)
@@ -43,8 +46,8 @@ async def chat_endpoint(
 		# Get medical context from memory
 		medical_context = state.history_manager.get_conversation_context(
 			request.user_id,
-			request.session_id,
-			request.message
+			#request.session_id,
+			#request.message
 		)
 
 		# Generate response using Gemini AI

src/core/history.py

@@ -59,7 +59,7 @@ class MedicalHistoryManager:
 		user_id: str,
 	) -> str:
 		"""Retrieves relevant conversation context for a new question."""
-		return self.memory.get_medical_context(user_id, session_id, question)
+		return self.memory.get_medical_context(user_id=user_id)
 
 	def get_user_medical_history(
 		self,

src/core/memory.py

@@ -6,132 +6,110 @@ from typing import Any
 
 from src.core.profile import UserProfile
 from src.core.session import ChatSession
-from src.data import mongodb
+from src.data.repositories import account as account_repo
+from src.data.repositories import medical as medical_repo
+from src.data.repositories import session as chat_repo
 from src.utils.logger import logger
 
 
 class MemoryLRU:
 	"""
-	Memory system using MongoDB for persistence, supporting:
-	- Multiple users with profiles
-	- Multiple chat sessions per user
-	- Chat history and continuity
-	- Medical context summaries
+	A memory system that orchestrates data access between the application core
+	and the data repositories, managing users, sessions, and medical context.
 	"""
+
 	def __init__(self, max_sessions_per_user: int = 10):
 		self.max_sessions_per_user = max_sessions_per_user
 
 	def create_user(self, user_id: str, name: str = "Anonymous") -> UserProfile:
-		"""Create a new user profile"""
-		user = UserProfile(user_id, name)
-		mongodb.create_account({
-			"_id": user_id,
-			"name": name,
-			"created_at": datetime.now(timezone.utc),
-			"last_seen": datetime.now(timezone.utc),
-			"preferences": {}
-		})
-		return user
+		"""Creates a new user profile."""
+		account_repo.create_account(name=name, user_id=user_id)
+		return UserProfile(user_id, name)
 
 	def get_user(self, user_id: str) -> UserProfile | None:
-		"""Get user profile by ID"""
-		data = mongodb.get_user_profile(user_id)
+		"""Retrieves a user profile by its ID."""
+		data = account_repo.get_user_profile(user_id)
 		return UserProfile.from_dict(data) if data else None
 
 	def create_session(self, user_id: str, title: str = "New Chat") -> str:
-		"""Create a new chat session"""
-		session_id = str(uuid.uuid4())
-		mongodb.create_chat_session({
-			"_id": session_id,
-			"user_id": user_id,
-			"title": title,
-			"messages": []
-		})
-		return session_id
+		"""Creates a new chat session for a user."""
+		return chat_repo.create_session(user_id, title)
 
 	def get_session(self, session_id: str) -> ChatSession | None:
-		"""Get chat session by ID"""
+		"""Retrieves a single chat session by its ID."""
 		try:
-			data = mongodb.get_session(session_id)
-			if not data:
-				logger().info(f"Session not found: {session_id}")
-				return None
-
-			logger().debug(f"Retrieved session data: {data}")
-			return ChatSession.from_dict(data)
+			data = chat_repo.get_session(session_id)
+			return ChatSession.from_dict(data) if data else None
 		except Exception as e:
 			logger().error(f"Error retrieving session {session_id}: {e}")
-			logger().error(f"Stack trace:", exc_info=True)
 			raise
 
 	def get_user_sessions(self, user_id: str) -> list[ChatSession]:
-		"""Get all sessions for a user"""
-		sessions_data = mongodb.get_user_sessions(user_id, limit=self.max_sessions_per_user)
+		"""Retrieves all sessions for a specific user."""
+		sessions_data = chat_repo.get_user_sessions(user_id, limit=self.max_sessions_per_user)
 		return [ChatSession.from_dict(data) for data in sessions_data]
 
-	def add_message_to_session(self, session_id: str, role: str, content: str, metadata: dict | None = None):
-		"""Add a message to a session"""
+	def add_message_to_session(
+		self,
+		session_id: str,
+		role: str,
+		content: str,
+		metadata: dict = {}
+	):
+		"""Adds a message to a chat session."""
 		message = {
 			"id": str(uuid.uuid4()),
 			"role": role,
 			"content": content,
 			"timestamp": datetime.now(timezone.utc),
-			"metadata": metadata or {}
+			"metadata": metadata
 		}
-		mongodb.add_message(session_id, message)
+		chat_repo.add_message(session_id, message)
 
 	def update_session_title(self, session_id: str, title: str):
-		"""Update session title"""
-		mongodb.update_session_title(session_id, title)
+		"""Updates the title of a session."""
+		chat_repo.update_session_title(session_id, title)
 
 	def delete_session(self, session_id: str):
-		"""Delete a chat session"""
-		mongodb.delete_chat_session(session_id)
+		"""Deletes a chat session."""
+		chat_repo.delete_chat_session(session_id)
 
-	def set_user_preference(self, user_id: str, key: str, value: Any):
-		"""Set user preference"""
-		mongodb.set_user_preference(user_id, key, value)
+	def set_user_preferences(
+		self,
+		user_id: str,
+		update_data: dict[str, Any]
+	):
+		"""Sets a preference for a user."""
+		account_repo.set_user_preferences(user_id, update_data)
 
-	# Medical context methods
 	def add(self, user_id: str, summary: str):
-		"""Add a medical context summary"""
-		mongodb.add_medical_context(user_id, summary)
+		"""Adds a medical context summary for a user."""
+		medical_repo.add_medical_context(user_id, summary)
 
 	def all(self, user_id: str) -> list[str]:
-		"""Get all medical context summaries for a user"""
-		contexts = mongodb.get_medical_context(user_id)
+		"""Retrieves all medical context summaries for a user."""
+		contexts = medical_repo.get_medical_context(user_id)
 		return [ctx["summary"] for ctx in contexts]
 
 	def recent(self, user_id: str, n: int) -> list[str]:
-		"""Get n most recent medical context summaries"""
-		contexts = mongodb.get_medical_context(user_id, limit=n)
+		"""Retrieves the N most recent medical context summaries."""
+		contexts = medical_repo.get_medical_context(user_id, limit=n)
 		return [ctx["summary"] for ctx in contexts]
 
 	def rest(self, user_id: str, skip: int) -> list[str]:
-		"""Get all summaries except the most recent n"""
-		contexts = mongodb.get_medical_context(user_id)
-		return [ctx["summary"] for ctx in contexts[skip:]]
-
-	def get_medical_context(self, user_id: str, session_id: str, question: str) -> str:
-		"""Get relevant medical context for a question"""
+		"""Retrieves all summaries except for the N most recent ones."""
+		all_contexts = self.all(user_id)
+		return all_contexts[skip:]
+
+	def get_medical_context(
+		self,
+		user_id: str,
+		limit: int = 5
+	) -> str:
+		"""Retrieves and formats recent medical context into a single string."""
 		try:
-			# Get recent contexts
-			contexts = mongodb.get_medical_context(user_id, limit=5)
-			if not contexts:
-				return ""
-
-			# Format contexts into a string
-			context_texts = []
-			for ctx in contexts:
-				summary = ctx.get("summary")
-				if summary:
-					context_texts.append(summary)
-
-			if not context_texts:
-				return ""
-
-			return "\n\n".join(context_texts)
+			contexts = self.recent(user_id, limit)
+			return "\n\n".join(contexts)
 		except Exception as e:
-			logger().error(f"Error getting medical context: {e}")
-			logger().error("Stack trace:", exc_info=True)
+			logger().error(f"Error getting medical context for user {user_id}: {e}")
 			return ""

src/core/state.py

@@ -7,30 +7,33 @@ from src.utils.rotator import APIKeyRotator
 
 
 class MedicalState:
-	"""Global state management for Medical AI system"""
-	_instance: 'MedicalState | None' = None
+	"""Manages the global state of the application using a Singleton pattern."""
+	_instance = None
+
+	def __new__(cls):
+		if cls._instance is None:
+			cls._instance = super(MedicalState, cls).__new__(cls)
+			cls._instance._initialized = False
+		return cls._instance
 
 	def __init__(self):
+		if self._initialized:
+			return
 		self.memory_system: MemoryLRU
 		self.embedding_client: EmbeddingClient
 		self.history_manager: MedicalHistoryManager
 		self.gemini_rotator: APIKeyRotator
 		self.nvidia_rotator: APIKeyRotator
+		self._initialized = True
 
 	def initialize(self):
-		"""Initialize all core components"""
+		"""Initializes all core application components."""
 		self.memory_system = MemoryLRU(max_sessions_per_user=20)
-		self.embedding_client = EmbeddingClient("all-MiniLM-L6-v2", 384)
+		self.embedding_client = EmbeddingClient(model_name="all-MiniLM-L6-v2", dimension=384)
 		self.history_manager = MedicalHistoryManager(self.memory_system, self.embedding_client)
 		self.gemini_rotator = APIKeyRotator("GEMINI_API_", max_slots=5)
 		self.nvidia_rotator = APIKeyRotator("NVIDIA_API_", max_slots=5)
 
-	@classmethod
-	def get_instance(cls) -> 'MedicalState':
-		if cls._instance is None:
-			cls._instance = MedicalState()
-		return cls._instance
-
 def get_state() -> MedicalState:
-	"""FastAPI dependency for getting application state"""
-	return MedicalState.get_instance()
+	"""Provides access to the application state, for use as a dependency."""
+	return MedicalState()

src/data/repositories/account.py

@@ -16,16 +16,28 @@ class UserNotFound(Exception):
 	pass
 
 def create_account(
-	user_data: dict[str, Any],
+	name: str,
+	preferences: dict[str, Any] = {},
 	*,
+	user_id: str,
 	collection_name: str = ACCOUNTS_COLLECTION
 ) -> str:
-	"""Creates a new user account."""
+	"""
+		Creates a new user account.
+
+		@TODO Revise if the user id should be passed in or generated by the database.
+	"""
 	collection = get_collection(collection_name)
 	now = datetime.now(timezone.utc)
-	user_data.update({"created_at": now, "updated_at": now})
+	user_data: dict[str, Any] = {
+		"_id": user_id,
+		"name": name,
+		"created_at": now,
+		"updated_at": now
+	}
 	try:
 		result = collection.insert_one(user_data)
+		set_user_preferences(user_id, preferences)
 		logger().info(f"Created new account: {result.inserted_id}")
 		return str(result.inserted_id)
 	except DuplicateKeyError as e:
@@ -72,19 +84,19 @@ def get_user_profile(
 def set_user_preferences(
 	user_id: str,
 	/,
-	update_data: dict[str, Any],
+	preferences: dict[str, Any],
 	*,
 	collection_name: str = ACCOUNTS_COLLECTION
 ) -> bool:
 	"""Sets a preference for a user."""
 	try:
 		collection = get_collection(collection_name)
-		update_data = {f"preferences.{key}": value for key, value in update_data}
-		update_data["updated_at"] = datetime.now(timezone.utc)
+		preferences = {f"preferences.{key}": value for key, value in preferences}
+		preferences["updated_at"] = datetime.now(timezone.utc)
 		result = collection.update_one(
 			{"_id": user_id},
 			{
-				"$set": update_data
+				"$set": preferences
 			}
 		)
 		if result.matched_count == 0:

src/data/repositories/session.py

@@ -1,9 +1,9 @@
 # data/repositories/chat.py
 
+import uuid
 from datetime import datetime, timezone
 from typing import Any
 
-from bson import ObjectId
 from pymongo import DESCENDING
 from pymongo.errors import (ConnectionFailure, DuplicateKeyError,
                             OperationFailure, PyMongoError)
@@ -13,23 +13,26 @@ from src.utils.logger import logger
 
 CHAT_SESSIONS_COLLECTION = "chat_sessions"
 
-def create_chat_session(
-	session_data: dict[str, Any],
+def create_session(
+	user_id: str,
+	title: str,
 	*,
 	collection_name: str = CHAT_SESSIONS_COLLECTION
 ) -> str:
-	"""Creates a new chat session."""
+	"""
+		Creates a new chat session.
+
+		@TODO Revise if the session id should be passed in or generated by the database.
+	"""
 	collection = get_collection(collection_name)
 	now = datetime.now(timezone.utc)
-	session_data.update({
+	session_data: dict[str, Any] = {
+		"_id": str(uuid.uuid4()),
+		"user_id": user_id,
+		"title": title,
 		"created_at": now,
-		"updated_at": now,
-		"messages": session_data.get("messages", []),
-		"title": str(session_data.get("title", "New Chat")),
-		"user_id": str(session_data["user_id"])
-	})
-	#if "_id" not in session_data:
-	#	session_data["_id"] = str(ObjectId())
+		"updated_at": now
+	}
 	try:
 		result = collection.insert_one(session_data)
 		return str(result.inserted_id)

src/main.py

@@ -24,7 +24,7 @@ except Exception as e:
 
 # Import project modules after trying to load environment variables
 from src.api.routes import chat, session, static, system, user
-from src.core.state import MedicalState
+from src.core.state import MedicalState, get_state
 
 
 def startup_event(state: MedicalState):
@@ -73,7 +73,7 @@ def shutdown_event():
 @asynccontextmanager
 async def lifespan(app: FastAPI):
 	# Initialize state
-	state = MedicalState.get_instance()
+	state = get_state()
 	state.initialize()
 
 	# Startup code here