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]

packages/mongo-search/pyproject.toml

@@ -0,0 +1,15 @@
+[project]
+name = "mongo-search"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+authors = [{ name = "ajac-zero", email = "ajcardoza2000@gmail.com" }]
+requires-python = ">=3.12"
+dependencies = ["pymongo>=4.10.1"]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["banortegpt"]