ic

packages/mongo-search/banortegpt/vector/mongo_search/__init__.py

@@ -0,0 +1,49 @@
+from pymongo import AsyncMongoClient
+
+
+class AsyncMongo:
+    def __init__(
+        self, mongo_uri: str, database_name: str = "MayaContigo", mode: str = "cosmosdb"
+    ) -> None:
+        self.mode = mode
+        self.client: AsyncMongoClient = AsyncMongoClient(mongo_uri)
+        self.db = self.client.get_database(database_name)
+
+    def build_pipeline(self, embedding, limit):
+        if self.mode == "native":
+            return {
+                "$vectorSearch": {
+                    "index": "vector_index",
+                    "queryVector": embedding,
+                    "path": "embedding",
+                    "limit": limit,
+                    "numCandidates": 3 * limit,
+                }
+            }
+        elif self.mode == "cosmosdb":
+            return {
+                "$search": {
+                    "cosmosSearch": {
+                        "path": "vector",
+                        "vector": embedding[:2000],
+                        "k": limit,
+                    }
+                }
+            }
+        else:
+            raise ValueError("Invalid mode")
+
+    async def semantic_search(
+        self,
+        collection: str,
+        embedding: list[float],
+        limit: int = 10,
+        conditions: dict | None = None,
+        threshold: float | None = None,
+        **kwargs,
+    ):
+        db_collection = self.db[collection]
+
+        pipeline = self.build_pipeline(embedding, limit)
+        aggregate = await db_collection.aggregate([pipeline])
+        return [result async for result in aggregate]