ic

packages/qdrant/banortegpt/vector/qdrant/__init__.py

@@ -0,0 +1,4 @@
+from .blocking import Qdrant
+from .nonblocking import AsyncQdrant
+
+__all__ = ["Qdrant", "AsyncQdrant"]

packages/qdrant/banortegpt/vector/qdrant/base.py

@@ -0,0 +1,43 @@
+from typing import Protocol
+
+
+class BaseQdrant:
+    def __init__(
+        self, *, url: str, api_key: str | None, collection: str | None = None
+    ) -> None:
+        self.collection = collection
+
+    class Config(Protocol):
+        qdrant_url: str
+        qdrant_api_key: str | None
+
+    @classmethod
+    def from_config(cls, c: Config):
+        return cls(url=c.qdrant_url, api_key=c.qdrant_api_key)
+
+    @classmethod
+    def from_vault(
+        cls,
+        vault: str,
+        *,
+        collection: str | None = None,
+        url: str | None = None,
+        token: str | None = None,
+        mount_point: str = "secret",
+    ):
+        from hvac import Client
+
+        client = Client(url=url or "https://vault.ia-innovacion.work", token=token)
+
+        if not client.is_authenticated():
+            raise Exception("Vault authentication failed")
+
+        secret_map = client.secrets.kv.v2.read_secret_version(
+            path=vault, mount_point=mount_point
+        )["data"]["data"]
+
+        return cls(
+            url=secret_map["qdrant_api_url"],
+            api_key=secret_map["qdrant_api_key"],
+            collection=collection,
+        )

packages/qdrant/banortegpt/vector/qdrant/blocking.py

@@ -0,0 +1,99 @@
+from collections.abc import Sequence
+from typing import Any
+
+from langfuse.decorators import langfuse_context, observe
+from qdrant_client import QdrantClient, models
+
+from .base import BaseQdrant
+
+
+class Qdrant(BaseQdrant):
+    def __init__(
+        self, *, url: str, api_key: str | None, collection: str | None = None
+    ) -> None:
+        super().__init__(url=url, api_key=api_key, collection=collection)
+        self.client = QdrantClient(url=url, api_key=api_key)
+
+    def list_collections(self) -> Sequence[str]:
+        return [
+            collection.name for collection in self.client.get_collections().collections
+        ]
+
+    @observe(capture_input=False)
+    def semantic_search(
+        self,
+        embedding: Sequence[float] | models.NamedVector,
+        collection: str | None = None,
+        limit: int = 10,
+        conditions: Any | None = None,
+        threshold: float | None = None,
+        **kwargs,
+    ) -> Sequence[dict[str, Any]]:
+        if collection is None:
+            if self.collection is None:
+                raise ValueError(
+                    "No collection set; Please set a collection before calling 'semantic_search'"
+                )
+            collection = self.collection
+
+        langfuse_context.update_current_observation(
+            input={
+                "collection": collection,
+                "limit": limit,
+                "embedding": embedding,
+                "conditions": conditions,
+            }
+        )
+
+        points = self.client.search(
+            collection_name=collection,
+            query_vector=embedding,
+            query_filter=conditions,
+            limit=limit,
+            with_payload=True,
+            with_vectors=False,
+            score_threshold=threshold,
+            **kwargs,
+        )
+
+        return [point.payload for point in points if point.payload is not None]
+
+    def create_collection_if_not_exists(
+        self,
+        *,
+        collection: str | None = None,
+        vector_config: dict[str, models.VectorParams],
+    ):
+        if collection is None:
+            if self.collection is None:
+                raise ValueError(
+                    "No collection is set; Please set a collection before calling 'create_collection_if_not_exists'"
+                )
+            collection = self.collection
+
+        result = self.client.get_collections()
+        collection_names = [collection.name for collection in result.collections]
+
+        if collection not in collection_names:
+            return self.client.create_collection(
+                collection_name=collection,
+                vectors_config=vector_config,
+            )
+
+        return False
+
+    def upload_to_collection(
+        self,
+        *,
+        points: list[models.PointStruct],
+        collection: str | None = None,
+    ):
+        if collection is None:
+            if self.collection is None:
+                raise ValueError(
+                    "No collection is set; Please set a collection before calling 'create_collection_if_not_exists'"
+                )
+            collection = self.collection
+
+        for point in points:
+            self.client.upsert(collection_name=collection, points=[point], wait=True)

packages/qdrant/banortegpt/vector/qdrant/nonblocking.py

@@ -0,0 +1,96 @@
+from collections.abc import Sequence
+from typing import Any
+
+from langfuse.decorators import langfuse_context, observe
+from qdrant_client import AsyncQdrantClient, models
+
+from .base import BaseQdrant
+
+
+class AsyncQdrant(BaseQdrant):
+    def __init__(
+        self, *, url: str, api_key: str | None, collection: str | None = None
+    ) -> None:
+        super().__init__(url=url, api_key=api_key, collection=collection)
+        self.client = AsyncQdrantClient(url=url, api_key=api_key)
+
+    @observe(capture_input=False)
+    async def semantic_search(
+        self,
+        embedding: Sequence[float] | models.NamedVector,
+        collection: str | None = None,
+        limit: int = 10,
+        conditions: Any | None = None,
+        threshold: float | None = None,
+        **kwargs,
+    ) -> Sequence[dict[str, Any]]:
+        if collection is None:
+            if self.collection is None:
+                raise ValueError(
+                    "No collection set; Please set a collection before calling 'semantic_search'"
+                )
+            collection = self.collection
+
+        langfuse_context.update_current_observation(
+            input={
+                "collection": collection,
+                "limit": limit,
+                "embedding": embedding,
+                "conditions": conditions,
+            }
+        )
+
+        points = await self.client.search(
+            collection_name=collection,
+            query_vector=embedding,
+            query_filter=conditions,
+            limit=limit,
+            with_payload=True,
+            with_vectors=False,
+            score_threshold=threshold,
+            **kwargs,
+        )
+
+        return [point.payload for point in points if point.payload is not None]
+
+    async def create_collection_if_not_exists(
+        self,
+        *,
+        collection: str | None = None,
+        vector_config: dict[str, models.VectorParams],
+    ):
+        if collection is None:
+            if self.collection is None:
+                raise ValueError(
+                    "No collection is set; Please set a collection before calling 'create_collection_if_not_exists'"
+                )
+            collection = self.collection
+
+        result = await self.client.get_collections()
+        collection_names = [collection.name for collection in result.collections]
+
+        if collection not in collection_names:
+            return await self.client.create_collection(
+                collection_name=collection,
+                vectors_config=vector_config,
+            )
+
+        return False
+
+    async def upload_to_collection(
+        self,
+        *,
+        points: list[models.PointStruct],
+        collection: str | None = None,
+    ):
+        if collection is None:
+            if self.collection is None:
+                raise ValueError(
+                    "No collection is set; Please set a collection before calling 'create_collection_if_not_exists'"
+                )
+            collection = self.collection
+
+        for point in points:
+            await self.client.upsert(
+                collection_name=collection, points=[point], wait=True
+            )

packages/qdrant/pyproject.toml

@@ -0,0 +1,19 @@
+[project]
+name = "qdrant"
+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 = ["langfuse>=2.60.2", "qdrant-client>=1.12.2"]
+
+[tool.pyright]
+venvPath = "."
+venv = ".venv"
+
+[tool.hatch.build.targets.wheel]
+packages = ["banortegpt"]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"