refactored mongo init, context and vectordb services to be synchronous

src/ctp_slack_bot/db/mongo_db.py

@@ -1,8 +1,7 @@
-#from motor.motor_asyncio import AsyncIOMotorClient
+from pymongo import MongoClient, ASCENDING
 from loguru import logger
-from pydantic import BaseModel, model_validator
-#from pymongo import IndexModel, ASCENDING
-from typing import Optional, Self
+from pydantic import BaseModel, model_validator, ConfigDict
+from typing import Optional, Self, Any
 
 from ctp_slack_bot.core.config import Settings
 
@@ -13,115 +12,114 @@ class MongoDB(BaseModel):
     """
 
     settings: Settings
+    client: Optional[MongoClient] = None
+    db: Optional[Any] = None
+    vector_collection: Optional[Any] = None
+    initialized: bool = False
+    
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     @model_validator(mode='after')
     def post_init(self: Self) -> Self:
         logger.debug("Created {}", self.__class__.__name__)
         return self
 
-    def __init__(self: Self, settings: Settings) -> Self:
-        #self.client: Optional[AsyncIOMotorClient] = None
-        #self.db = None
-        #self.vector_collection = None
-        #self.initialized = False
-        pass # The above initialization needs to be done some other way.
+    def connect(self):
+        """
+        Connect to MongoDB using connection string from settings.
+        """
+        if self.client is not None:
+            return
 
-    # async def connect(self):
-    #     """
-    #     Connect to MongoDB using connection string from settings.
-    #     """
-    #     if self.client is not None:
-    #         return
+        if not self.settings.MONGODB_URI:
+            raise ValueError("MONGODB_URI is not set in environment variables")
 
-    #     if not settings.MONGODB_URI:
-    #         raise ValueError("MONGODB_URI is not set in environment variables")
+        try:
+            # Create MongoDB connection
+            self.client = MongoClient(self.settings.MONGODB_URI.get_secret_value())
+            self.db = self.client[self.settings.MONGODB_NAME]
+            self.vector_collection = self.db["vector_store"]
+            logger.info(f"Connected to MongoDB: {self.settings.MONGODB_NAME}")
+        except Exception as e:
+            logger.error(f"Error connecting to MongoDB: {str(e)}")
+            raise
 
-    #     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
+    def initialize(self):
+        """
+        Initialize MongoDB with required collections and indexes.
+        """
+        if self.initialized:
+            return
         
-    #     if not self.client:
-    #         await self.connect()
+        if not self.client:
+            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
+        try:
+            # Create vector index for similarity search
+            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 create_vector_index(self):
+        """
+        Create vector index for similarity search using MongoDB Atlas Vector Search.
+        """
+        try:
+            # Check if index already exists
+            existing_indexes = list(self.vector_collection.list_indexes())
+            index_names = [index.get('name') for index in existing_indexes]
             
-    #         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"
-    #                         }
-    #                     }
-    #                 }
-    #             }
+            if "vector_index" not in index_names:
+                # Create vector search index
+                index_definition = {
+                    "mappings": {
+                        "dynamic": True,
+                        "fields": {
+                            "embedding": {
+                                "dimensions": self.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 the index
+                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)])
+                # Create additional metadata indexes for filtering
+                self.vector_collection.create_index([("metadata.source", ASCENDING)])
+                self.vector_collection.create_index([("metadata.timestamp", ASCENDING)])
                 
-    #             logger.info("Vector search index created")
-    #         else:
-    #             logger.info("Vector search index already exists")
+                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
+        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")
+    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")

src/ctp_slack_bot/services/context_retrieval_service.py

@@ -1,6 +1,6 @@
 from loguru import logger
 from pydantic import BaseModel, model_validator
-from typing import Any, Dict, List, Optional, Self
+from typing import Self, List
 
 from ctp_slack_bot.core.config import Settings
 from ctp_slack_bot.models import RetreivedContext, SlackMessage, VectorQuery
@@ -21,13 +21,7 @@ class ContextRetrievalService(BaseModel):
         logger.debug("Created {}", self.__class__.__name__)
         return self
     
-    async def initialize(self):
-        """
-        Initialize the required services.
-        """
-        await self.vector_database_service.initialize()
-    
-    async def get_context(self, message: SlackMessage) -> List[RetreivedContext]:
+    def get_context(self, message: SlackMessage) -> List[RetreivedContext]:
         """
         Retrieve relevant context for a given Slack message.
         
@@ -59,12 +53,12 @@ class ContextRetrievalService(BaseModel):
             # Create vector query
             vector_query = VectorQuery(
                 query_text=message.text,
-                k=settings.TOP_K_MATCHES,
+                k=self.settings.TOP_K_MATCHES,
                 score_threshold=0.7  # Minimum similarity threshold
             )
             
             # Search for similar content in vector database
-            context_results = await self.vector_database_service.search_by_similarity(
+            context_results = self.vector_database_service.search_by_similarity(
                 query=vector_query,
                 query_embedding=query_embedding
             )

src/ctp_slack_bot/services/vector_database_service.py

@@ -1,6 +1,6 @@
 from loguru import logger
 from pydantic import BaseModel, model_validator
-from typing import Any, Dict, List, Optional, Self
+from typing import Any, Dict, List, Self
 
 from ctp_slack_bot.core import Settings
 from ctp_slack_bot.db import MongoDB
@@ -10,7 +10,7 @@ class VectorDatabaseService(BaseModel): # TODO: this should not rely specificall
     """
     Service for storing and retrieving vector embeddings from MongoDB.
     """
-
+    
     settings: Settings
     mongo_db: MongoDB
 
@@ -19,13 +19,7 @@ class VectorDatabaseService(BaseModel): # TODO: this should not rely specificall
         logger.debug("Created {}", self.__class__.__name__)
         return self
 
-    async def initialize(self):
-        """
-        Initialize the database connection.
-        """
-        await mongodb.initialize()
-    
-    async def store(self, text: str, embedding: List[float], metadata: Dict[str, Any]) -> str:
+    def store(self, text: str, embedding: List[float], metadata: Dict[str, Any]) -> str:
         """
         Store text and its embedding vector in the database.
         
@@ -37,8 +31,8 @@ class VectorDatabaseService(BaseModel): # TODO: this should not rely specificall
         Returns:
             str: The ID of the stored document
         """
-        if not mongodb.initialized:
-            await mongodb.initialize()
+        if not self.mongo_db.initialized:
+            self.mongo_db.initialize()
             
         try:
             # Create document to store
@@ -49,7 +43,7 @@ class VectorDatabaseService(BaseModel): # TODO: this should not rely specificall
             }
             
             # Insert into collection
-            result = await mongodb.vector_collection.insert_one(document)
+            result = self.mongo_db.vector_collection.insert_one(document)
             logger.debug(f"Stored document with ID: {result.inserted_id}")
             
             return str(result.inserted_id)
@@ -57,7 +51,7 @@ class VectorDatabaseService(BaseModel): # TODO: this should not rely specificall
             logger.error(f"Error storing embedding: {str(e)}")
             raise
     
-    async def search_by_similarity(self, query: VectorQuery, query_embedding: List[float]) -> List[RetreivedContext]:
+    def search_by_similarity(self, query: VectorQuery, query_embedding: List[float]) -> List[RetreivedContext]:
         """
         Query the vector database for similar documents.
         
@@ -68,8 +62,8 @@ class VectorDatabaseService(BaseModel): # TODO: this should not rely specificall
         Returns:
             List[RetreivedContext]: List of similar documents with similarity scores
         """
-        if not mongodb.initialized:
-            await mongodb.initialize()
+        if not self.mongo_db.initialized:
+            self.mongo_db.initialize()
             
         try:
             # Build aggregation pipeline for vector search
@@ -100,7 +94,7 @@ class VectorDatabaseService(BaseModel): # TODO: this should not rely specificall
                 pipeline.insert(1, {"$match": metadata_filter})
             
             # Execute the pipeline
-            results = await mongodb.vector_collection.aggregate(pipeline).to_list(length=query.k)
+            results = list(self.mongo_db.vector_collection.aggregate(pipeline, maxTimeMS=30000))
             
             # Convert to RetreivedContext objects directly
             context_results = []