ai-microcore 4.0.0.dev22__tar.gz → 4.0.0.dev23__tar.gz

{ai_microcore-4.0.0.dev22 → ai_microcore-4.0.0.dev23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-microcore
-Version: 4.0.0.dev22
+Version: 4.0.0.dev23
 Summary: # Minimalistic Foundation for AI Applications
 Keywords: llm,large language models,ai,similarity search,ai search,gpt,openai,framework,adapter
 Author-email: Vitalii Stepanenko <mail@vitalii.in>

{ai_microcore-4.0.0.dev22 → ai_microcore-4.0.0.dev23}/microcore/__init__.py

@@ -186,4 +186,4 @@ __all__ = [
     # "wrappers",
 ]
 
-__version__ = "4.0.0-dev22"
+__version__ = "4.0.0-dev23"

{ai_microcore-4.0.0.dev22 → ai_microcore-4.0.0.dev23}/microcore/_env.py

@@ -151,6 +151,17 @@ class Env:
             from .embedding_db.chromadb import ChromaEmbeddingDB
 
             self.texts = ChromaEmbeddingDB(self.config)
+            return
+
+        if self.config.EMBEDDING_DB_TYPE == EmbeddingDbType.QDRANT:
+            if find_spec("qdrant_client") is None:
+                raise ModuleNotFoundError(
+                    "To use Qdrant, install the `qdrant-client` package. "
+                    "Run `pip install qdrant-client`."
+                )
+            from .embedding_db.qdrant import QdrantEmbeddingDB
+            self.texts = QdrantEmbeddingDB(self.config)
+            return
 
 
 @dataclass

{ai_microcore-4.0.0.dev22 → ai_microcore-4.0.0.dev23}/microcore/configuration.py

@@ -84,6 +84,7 @@ class ApiType(str, Enum):
 
 class EmbeddingDbType(str, Enum):
     CHROMA = "chroma"
+    QDRANT = "qdrant"
     NONE = ""
 
     def __str__(self):
@@ -392,6 +393,9 @@ class Config(LLMConfig):
 
     EMBEDDING_DB_TYPE: str = from_env(EmbeddingDbType.CHROMA)
 
+    EMBEDDING_DB_SIZE: int = from_env(default=0)
+    """Used with Qdrant"""
+
     DEFAULT_ENCODING: str = from_env("utf-8")
     """Used in file system operations, utf-8 by default"""
 
@@ -435,3 +439,19 @@ class Config(LLMConfig):
                 raise e
         if self.TEXT_TO_SPEECH_PATH is None:
             self.TEXT_TO_SPEECH_PATH = Path(self.STORAGE_PATH) / "voicing"
+
+    def validate(self):
+        super().validate()
+        if self.EMBEDDING_DB_TYPE == EmbeddingDbType.QDRANT:
+            if not self.EMBEDDING_DB_SIZE:
+                raise LLMConfigError(
+                    "EMBEDDING_DB_SIZE is required configuration parameter for Qdrant"
+                )
+            if not self.EMBEDDING_DB_HOST:
+                raise LLMConfigError(
+                    "EMBEDDING_DB_HOST is required configuration parameter for Qdrant"
+                )
+            if not self.EMBEDDING_DB_FUNCTION:
+                raise LLMConfigError(
+                    "EMBEDDING_DB_FUNCTION is required configuration parameter for Qdrant"
+                )

{ai_microcore-4.0.0.dev22 → ai_microcore-4.0.0.dev23}/microcore/embedding_db/chromadb.py

@@ -107,6 +107,7 @@ class ChromaEmbeddingDB(AbstractEmbeddingDB):
         return chroma_collection.count() if chroma_collection else 0
 
     def delete(self, collection: str, what: str | list[str] | dict):
+        # pylint: disable=R0801, duplicate-code
         if not self.collection_exists(collection):
             return
 

ai_microcore-4.0.0.dev23/microcore/embedding_db/qdrant.py

@@ -0,0 +1,227 @@
+import hashlib
+import logging
+import sys
+from dataclasses import dataclass
+import uuid
+
+from qdrant_client import QdrantClient
+from qdrant_client.http.models import (
+    PointIdsList,
+    CollectionInfo,
+    Record,
+    FieldCondition,
+    Filter,
+    MatchValue,
+    MatchText
+)
+from qdrant_client.models import VectorParams, Distance, PointStruct, ScoredPoint
+
+from ..configuration import Config
+from .. import SearchResult, SearchResults, AbstractEmbeddingDB
+
+
+def is_sentence_transformer(fn):
+    return fn.__class__.__name__ == 'SentenceTransformer'
+
+
+def prepare_embedding_function(fn):
+    if is_sentence_transformer(fn):
+        return lambda x: fn.encode(x).tolist()
+    return fn
+
+
+@dataclass
+class QdrantEmbeddingDB(AbstractEmbeddingDB):
+    config: Config
+    embedding_function: callable = None
+    client: QdrantClient = None
+
+    def __post_init__(self):
+
+        logging.info(
+            "Connecting to Qdrant at %s:%s",
+            self.config.EMBEDDING_DB_HOST,
+            self.config.EMBEDDING_DB_PORT
+        )
+        self.client = QdrantClient(
+            host=self.config.EMBEDDING_DB_HOST,
+            port=self.config.EMBEDDING_DB_PORT or 6333
+        )
+        self.embedding_function = prepare_embedding_function(self.config.EMBEDDING_DB_FUNCTION)
+
+    @classmethod
+    def _wrap_results(cls, points: list[ScoredPoint | Record]) -> list[str | SearchResult]:
+        return SearchResults(
+            [
+                SearchResult(
+                    i.payload["_text"],
+                    dict(
+                        id=i.id,
+                        distance=i.score if hasattr(i, "score") else 0,
+                        metadata={k: v for k, v in i.payload.items() if k != "_text"},
+                    ),
+
+                )
+                for i in points
+            ]
+        )
+
+    @classmethod
+    def _convert_where(cls, where: dict | None, kwargs=None) -> Filter | None:
+        conditions = []
+        # ChromaDB format
+        if kwargs and "where_document" in kwargs and kwargs["where_document"]:
+            conditions.append(
+                FieldCondition(
+                    key="_text",
+                    match=MatchText(text=kwargs["where_document"]["$contains"])
+                )
+            )
+        if where:
+            for k, v in where.items():
+                conditions.append(FieldCondition(key=k, match=MatchValue(value=v)))
+
+        return Filter(must=conditions) if conditions else None
+
+    def search(
+        self,
+        collection: str,
+        query: str | list,
+        n_results: int = 5,
+        where: dict = None,
+        **kwargs,
+    ) -> list[str | SearchResult]:
+        if not self.collection_exists(collection):
+            return SearchResults([])
+
+        if not isinstance(query, str):
+            raise ValueError("`query` must be a string")
+
+        query_vector = self.embedding_function(query)
+        where = self._convert_where(where, kwargs)
+        kwargs.pop("where_document", None)
+        hits = self.client.query_points(
+            collection_name=collection,
+            query=query_vector,
+            limit=n_results,
+            query_filter=where,
+            **kwargs,
+        )
+        return self._wrap_results(hits.points)
+
+    def save_many(self, collection: str, items: list[tuple[str, dict] | str]):
+        if not self.collection_exists(collection):
+            self._create_collection(collection)
+        point_structs = []
+        ids = set()
+        unique = not self.config.EMBEDDING_DB_ALLOW_DUPLICATES
+        for i in items:
+            if isinstance(i, str):
+                text = i
+                metadata = dict()
+            else:
+                text = i[0]
+                metadata = i[1] or {}
+            metadata["_text"] = text
+            if unique:
+                new_id = str(uuid.UUID(hashlib.md5(text.encode()).hexdigest()))
+                if new_id in ids:
+                    continue
+                ids.add(new_id)
+            else:
+                new_id = str(uuid.uuid4())
+            point_structs.append(
+                PointStruct(
+                    id=new_id,
+                    vector=self.embedding_function(text),
+                    payload=metadata
+                )
+            )
+
+        operation_info = self.client.upsert(
+            collection_name=collection,
+            wait=True,
+            points=point_structs,
+        )
+        return operation_info
+
+    def clear(self, collection: str):
+        if self.collection_exists(collection):
+            self.client.delete_collection(collection_name=collection)
+
+    def count(self, collection: str) -> int:
+        if self.collection_exists(collection):
+            return self._get_collection(collection).points_count
+        return 0
+
+    def delete(self, collection: str, what: str | list[str] | dict):
+        # pylint: disable=R0801, duplicate-code
+        if not self.collection_exists(collection):
+            return
+
+        if isinstance(what, str):
+            ids, where = [what], None
+        elif isinstance(what, list):
+            ids, where = what, None
+        elif isinstance(what, dict):
+            ids, where = None, what
+        else:
+            raise ValueError("Invalid `what` argument")
+        if ids is not None:
+            points_selector = PointIdsList(points=ids)
+        else:
+            points_selector = self._convert_where(where)
+        self.client.delete(collection_name=collection, points_selector=points_selector)
+
+    def get(
+        self,
+        collection: str,
+        ids: list[str] | str = None,
+        limit: int = None,
+        offset: int = None,
+        where: dict = None,
+        **kwargs,
+    ) -> list[str | SearchResult] | str | SearchResult | None:
+        if not self.collection_exists(collection):
+            return SearchResults([]) if not isinstance(ids, str) else None
+        if ids:
+            raise NotImplementedError("Getting by ids is not supported for Qdrant")
+        where = self._convert_where(where, kwargs)
+        kwargs.pop("where_document", None)
+        search_results = self._wrap_results(
+            self.client.scroll(
+                collection,
+                limit=limit or sys.maxsize - 1,
+                offset=offset or 0,
+                scroll_filter=where,
+                **kwargs
+            )[0]
+        )
+        if isinstance(ids, str):
+            return search_results[0] if search_results else None
+        return search_results
+
+    def get_all(self, collection: str) -> list[str | SearchResult]:
+        if not self.collection_exists(collection):
+            return SearchResults([])
+        return self._wrap_results(self.client.scroll(collection, limit=sys.maxsize - 1)[0])
+
+    def collection_exists(self, collection: str) -> bool:
+        return self.client.collection_exists(collection)
+
+    def _create_collection(self, name: str):
+        assert self.config.EMBEDDING_DB_SIZE > 0
+        size = self.config.EMBEDDING_DB_SIZE
+        distance = Distance.COSINE
+        self.client.create_collection(
+            collection_name=name,
+            vectors_config=VectorParams(size=size, distance=distance),
+        )
+
+    def _get_collection(self, name: str, create: bool = False) -> CollectionInfo | None:
+        if not self.collection_exists(name):
+            if create:
+                self._create_collection(name)
+            else:
+                return None
+        return self.client.get_collection(name)

{ai_microcore-4.0.0.dev22 → ai_microcore-4.0.0.dev23}/microcore/ui.py

@@ -54,7 +54,7 @@ def ask_choose(msg: str, variants: list):
             continue
         break
 
-    item = variants[int(i)]
+    item = variants[i]
     return item