Test MongoDB code

src/ctp_slack_bot/containers.py

@@ -4,6 +4,7 @@ from slack_bolt.async_app import AsyncApp
 
 from ctp_slack_bot.core.config import Settings
 from ctp_slack_bot.db.mongo_db import MongoDB
+from ctp_slack_bot.db.repositories import MongoVectorizedChunkRepository
 from ctp_slack_bot.services.answer_retrieval_service import AnswerRetrievalService
 from ctp_slack_bot.services.content_ingestion_service import ContentIngestionService
 from ctp_slack_bot.services.context_retrieval_service import ContextRetrievalService
@@ -20,12 +21,11 @@ class Container(DeclarativeContainer):
     settings = Singleton(Settings)
     event_brokerage_service = Singleton(EventBrokerageService)
     slack_bolt_app = Factory(AsyncApp, token=settings.provided.SLACK_BOT_TOKEN().get_secret_value())
-    mongo_db = Singleton(MongoDB, settings=settings) # TODO: we could really use less commitment to MongoDB.
-    # Repositories
-    # transcript_repository = Factory(
-    #     # Your transcript repository class
-    #     mongo_db=mongo_db
-    # )
+    mongo_db = Singleton(MongoDB, settings=settings) # TODO: generalize to any database.
+    vectorized_chunk_repository = Singleton(
+        MongoVectorizedChunkRepository,
+        mongo_db=mongo_db
+    )
     vector_database_service = Singleton(VectorDatabaseService, settings=settings, mongo_db=mongo_db)
     embeddings_model_service = Singleton(EmbeddingsModelService, settings=settings)
     vectorization_service = Singleton(VectorizationService, settings=settings, embeddings_model_service=embeddings_model_service)
@@ -35,4 +35,4 @@ class Container(DeclarativeContainer):
     answer_retrieval_service = Singleton(AnswerRetrievalService, settings=settings, event_brokerage_service=event_brokerage_service, language_model_service=language_model_service)
     question_dispatch_service = Singleton(QuestionDispatchService, settings=settings, event_brokerage_service=event_brokerage_service, content_ingestion_service=content_ingestion_service, context_retrieval_service=context_retrieval_service, answer_retrieval_service=answer_retrieval_service)
     slack_service = Singleton(SlackService, event_brokerage_service=event_brokerage_service, slack_bolt_app=slack_bolt_app)
-    primordial_services = List(settings, event_brokerage_service, slack_bolt_app, slack_service, question_dispatch_service, content_ingestion_service)
+    primordial_services = List(settings, event_brokerage_service, slack_bolt_app, slack_service, vectorized_chunk_repository, question_dispatch_service, content_ingestion_service)

src/ctp_slack_bot/db/mongo_db.py

@@ -1,131 +1,102 @@
-#from motor.motor_asyncio import AsyncIOMotorClient
+from motor.motor_asyncio import AsyncIOMotorClient
+from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError
 from loguru import logger
 from pydantic import BaseModel, model_validator, PrivateAttr
-#from pymongo import IndexModel, ASCENDING
-from typing import Optional, Self
+from typing import Any, Dict, Optional, Self
+import asyncio
 
 from ctp_slack_bot.core.config import Settings
 
 class MongoDB(BaseModel):
     """
-    MongoDB connection and initialization class.
-    Handles connection to MongoDB, database selection, and index creation.
+    MongoDB connection manager using Motor for async operations.
     """
-
     settings: Settings
     _client: PrivateAttr = PrivateAttr()
     _db: PrivateAttr = PrivateAttr()
-    _vector_collection: PrivateAttr = PrivateAttr()
-    _initialized: PrivateAttr = PrivateAttr()
-
-    def __init__(self: Self, **data) -> None:
-        super().__init__(**data)
-        self._client = None
-        self._db = None
-        self._vector_collection = None
-        self._initialized = False
-
+    
+    class Config:
+        arbitrary_types_allowed = True
+    
     @model_validator(mode='after')
     def post_init(self: Self) -> Self:
+        """Initialize MongoDB connection after model creation."""
+        self._initialize_client()
         logger.debug("Created {}", self.__class__.__name__)
         return self
-
-    # async def connect(self):
-    #     """
-    #     Connect to MongoDB using connection string from settings.
-    #     """
-    #     if self.client is not None:
-    #         return
-
-    #     if not settings.MONGODB_URI:
-    #         raise ValueError("MONGODB_URI is not set in environment variables")
-
-    #     try:
-    #         # Create MongoDB connection
-    #         self.client = AsyncIOMotorClient(settings.MONGODB_URI.get_secret_value())
-    #         self.db = self.client[settings.MONGODB_DB_NAME]
-    #         self.vector_collection = self.db["vector_store"]
-    #         logger.info(f"Connected to MongoDB: {settings.MONGODB_DB_NAME}")
-    #     except Exception as e:
-    #         logger.error(f"Error connecting to MongoDB: {str(e)}")
-    #         raise
-
-    # async def initialize(self):
-    #     """
-    #     Initialize MongoDB with required collections and indexes.
-    #     """
-    #     if self.initialized:
-    #         return
-        
-    #     if not self.client:
-    #         await self.connect()
-        
-    #     try:
-    #         # Create vector index for similarity search
-    #         await self.create_vector_index()
-    #         self.initialized = True
-    #         logger.info("MongoDB initialized successfully")
-    #     except Exception as e:
-    #         logger.error(f"Error initializing MongoDB: {str(e)}")
-    #         raise
-
-    # async def create_vector_index(self):
-    #     """
-    #     Create vector index for similarity search using MongoDB Atlas Vector Search.
-    #     """
-    #     try:
-    #         # Check if index already exists
-    #         existing_indexes = await self.vector_collection.list_indexes().to_list(length=None)
-    #         index_names = [index.get('name') for index in existing_indexes]
+    
+    def _initialize_client(self: Self) -> None:
+        """Initialize MongoDB client with settings."""
+        try:
+            connection_string = self.settings.MONGODB_URI.get_secret_value()
             
-    #         if "vector_index" not in index_names:
-    #             # Create vector search index
-    #             index_definition = {
-    #                 "mappings": {
-    #                     "dynamic": True,
-    #                     "fields": {
-    #                         "embedding": {
-    #                             "dimensions": settings.VECTOR_DIMENSION,
-    #                             "similarity": "cosine",
-    #                             "type": "knnVector"
-    #                         }
-    #                     }
-    #                 }
-    #             }
-                
-    #             # Create the index
-    #             await self.db.command({
-    #                 "createIndexes": self.vector_collection.name,
-    #                 "indexes": [
-    #                     {
-    #                         "name": "vector_index",
-    #                         "key": {"embedding": "vector"},
-    #                         "weights": {"embedding": 1},
-    #                         "vectorSearchOptions": index_definition
-    #                     }
-    #                 ]
-    #             })
-                
-    #             # Create additional metadata indexes for filtering
-    #             await self.vector_collection.create_index([("metadata.source", ASCENDING)])
-    #             await self.vector_collection.create_index([("metadata.timestamp", ASCENDING)])
-                
-    #             logger.info("Vector search index created")
-    #         else:
-    #             logger.info("Vector search index already exists")
-                
-    #     except Exception as e:
-    #         logger.error(f"Error creating vector index: {str(e)}")
-    #         raise
-
-    # async def close(self):
-    #     """
-    #     Close MongoDB connection.
-    #     """
-    #     if self.client:
-    #         self.client.close()
-    #         self.client = None
-    #         self.db = None
-    #         self.vector_collection = None
-    #         self.initialized = False
-    #         logger.info("MongoDB connection closed")
+            # Create client with appropriate settings
+            self._client = AsyncIOMotorClient(
+                connection_string,
+                serverSelectionTimeoutMS=5000,
+                connectTimeoutMS=10000,
+                socketTimeoutMS=45000,
+                maxPoolSize=100,
+                retryWrites=True,
+                w="majority"
+            )
+            
+            # Set database
+            db_name = self.settings.MONGODB_NAME
+            
+            self._db = self._client[db_name]
+            logger.debug("MongoDB client initialized for database: {}", db_name)
+        
+        except Exception as e:
+            logger.error("Failed to initialize MongoDB client: {}", e)
+            self._client = None
+            self._db = None
+    
+    @property
+    def client(self: Self) -> AsyncIOMotorClient:
+        """Get the MongoDB client instance."""
+        if not self._client:
+            logger.warning("MongoDB client not initialized. Attempting to initialize.")
+            self._initialize_client()
+            if not self._client:
+                raise ConnectionError("Failed to initialize MongoDB client")
+        return self._client
+    
+    @property
+    def db(self: Self) -> Any:
+        """Get the MongoDB database instance."""
+        if not self._db:
+            logger.warning("MongoDB database not initialized. Attempting to initialize client.")
+            self._initialize_client()
+            if not self._db:
+                raise ConnectionError("Failed to initialize MongoDB database")
+        return self._db
+    
+    async def ping(self: Self) -> bool:
+        """Check if MongoDB connection is alive."""
+        try:
+            await self.client.admin.command('ping')
+            return True
+        except (ConnectionFailure, ServerSelectionTimeoutError) as e:
+            logger.error("MongoDB connection failed: {}", e)
+            return False
+    
+    async def get_collection(self: Self, name: str) -> Any:
+        """Get a collection by name with validation."""
+        if not await self.ping():
+            raise ConnectionError("MongoDB connection is not available")
+        return self.db[name]
+    
+    async def create_indexes(self: Self, collection_name: str, indexes: list) -> None:
+        """Create indexes on a collection."""
+        collection = await self.get_collection(collection_name)
+        await collection.create_indexes(indexes)
+        logger.info("Created indexes for collection {}: {}", collection_name, indexes)
+    
+    async def close(self: Self) -> None:
+        """Close MongoDB connection."""
+        if self._client:
+            self._client.close()
+            logger.info("MongoDB connection closed")
+            self._client = None
+            self._db = None

src/ctp_slack_bot/db/repositories/__init__.py

@@ -0,0 +1,2 @@
+from ctp_slack_bot.db.repositories.mongo_db_vectorized_chunk_repository import MongoVectorizedChunkRepository
+from ctp_slack_bot.db.repositories.vectorized_chunk_repository import VectorizedChunkRepository

src/ctp_slack_bot/db/repositories/mongo_db_vectorized_chunk_repository.py

@@ -0,0 +1,65 @@
+from typing import List, Optional, Dict, Any
+import pymongo
+from bson import ObjectId
+
+from ctp_slack_bot.db import MongoDB
+from ctp_slack_bot.db.repositories.vectorized_chunk_repository import VectorizedChunkRepository
+from ctp_slack_bot.models.base import VectorizedChunk
+
+class MongoVectorizedChunkRepository(VectorizedChunkRepository):
+    """MongoDB implementation of VectorizedChunkRepository."""
+    
+    def __init__(self, mongo_db: MongoDB):
+        self.mongo_db = mongo_db
+        self.collection = self.mongo_db.client.get_database().get_collection("vectorized_chunks")
+        
+        # Create indexes for efficient queries
+        self.collection.create_index("chunk_id")
+        self.collection.create_index("parent_id")
+        
+    async def find_by_id(self, id: str) -> Optional[VectorizedChunk]:
+        doc = await self.collection.find_one({"_id": ObjectId(id)})
+        return self._map_to_entity(doc) if doc else None
+    
+    async def find_all(self) -> List[VectorizedChunk]:
+        cursor = self.collection.find({})
+        return [self._map_to_entity(doc) async for doc in cursor]
+    
+    async def find_by_parent_id(self, parent_id: str) -> List[VectorizedChunk]:
+        cursor = self.collection.find({"parent_id": parent_id})
+        return [self._map_to_entity(doc) async for doc in cursor]
+    
+    async def save(self, chunk: VectorizedChunk) -> VectorizedChunk:
+        doc = self._map_to_document(chunk)
+        
+        if "_id" in doc and doc["_id"]:
+            # Update existing document
+            await self.collection.replace_one({"_id": doc["_id"]}, doc)
+        else:
+            # Insert new document
+            result = await self.collection.insert_one(doc)
+            doc["_id"] = result.inserted_id
+            
+        return self._map_to_entity(doc)
+    
+    async def delete(self, id: str) -> bool:
+        result = await self.collection.delete_one({"_id": ObjectId(id)})
+        return result.deleted_count > 0
+    
+    async def find_by_metadata(self, metadata_query: Dict[str, Any]) -> List[VectorizedChunk]:
+        # Convert the metadata query to MongoDB query format
+        query = {f"metadata.{k}": v for k, v in metadata_query.items()}
+        cursor = self.collection.find(query)
+        return [self._map_to_entity(doc) async for doc in cursor]
+    
+    def _map_to_document(self, chunk: VectorizedChunk) -> Dict[str, Any]:
+        """Convert a VectorizedChunk to a MongoDB document."""
+        doc = chunk.model_dump()
+        # Handle any special conversions needed
+        return doc
+    
+    def _map_to_entity(self, doc: Dict[str, Any]) -> VectorizedChunk:
+        """Convert a MongoDB document to a VectorizedChunk."""
+        if "_id" in doc:
+            doc["id"] = str(doc.pop("_id"))
+        return VectorizedChunk(**doc)

src/ctp_slack_bot/db/repositories/vectorized_chunk_repository.py

@@ -0,0 +1,30 @@
+from typing import List, Optional, Dict, Any
+
+from ctp_slack_bot.models.base import VectorizedChunk
+
+class VectorizedChunkRepository:
+    """Repository interface for VectorizedChunk entities."""
+
+    async def find_by_id(self, id: str) -> Optional[VectorizedChunk]:
+        """Find a chunk by its ID."""
+        pass
+
+    async def find_all(self) -> List[VectorizedChunk]:
+        """Find all chunks."""
+        pass
+
+    async def find_by_parent_id(self, parent_id: str) -> List[VectorizedChunk]:
+        """Find chunks by parent document ID."""
+        pass
+
+    async def save(self, chunk: VectorizedChunk) -> VectorizedChunk:
+        """Save a chunk to the database."""
+        pass
+
+    async def delete(self, id: str) -> bool:
+        """Delete a chunk by its ID."""
+        pass
+
+    async def find_by_metadata(self, metadata_query: Dict[str, Any]) -> List[VectorizedChunk]:
+        """Find chunks by metadata criteria."""
+        pass