diff --git a/examples/semantic_kernel_vector_store.py b/examples/semantic_kernel_vector_store.py new file mode 100644 index 0000000..82d876a --- /dev/null +++ b/examples/semantic_kernel_vector_store.py @@ -0,0 +1,89 @@ +"""Use dynavec as a Semantic Kernel VectorStore. + +pip install "dynavec[sentence-transformers,semantic-kernel]" +python examples/semantic_kernel_vector_store.py +""" + +import asyncio + +from semantic_kernel.data.vector import ( + FieldTypes, + VectorStoreCollectionDefinition, + VectorStoreField, +) + +from dynavec import Dynavec, DynavecConfig +from dynavec.embeddings import SentenceTransformerEmbedder +from dynavec.integrations.semantic_kernel import DynavecStore + + +async def main(): + embedder = SentenceTransformerEmbedder(model="all-MiniLM-L6-v2") + + cfg = DynavecConfig( + vector_bucket="dynavec-demo", + index="sk-vector-store", + table="dynavec_sk_vector_store", + dimension=embedder.dimension, + region="us-east-1", + auto_provision=True, + ) + + db = Dynavec(cfg, embedder=embedder) + store = DynavecStore(db) + + definition = VectorStoreCollectionDefinition( + fields=[ + VectorStoreField(field_type=FieldTypes.KEY, name="id"), + VectorStoreField(field_type=FieldTypes.DATA, name="text"), + VectorStoreField(field_type=FieldTypes.DATA, name="category"), + VectorStoreField( + field_type=FieldTypes.VECTOR, + name="embedding", + dimensions=embedder.dimension, + ), + ], + ) + + collection = store.get_collection( + record_type=dict, + definition=definition, + collection_name="kb", + ) + + texts = [ + "Semantic Kernel is a framework for building AI applications.", + "Retrieval-augmented generation grounds answers in your documents.", + ] + vectors = embedder.embed_documents(texts) + + await collection.upsert( + [ + { + "id": "sk-1", + "text": texts[0], + "category": "ai", + "embedding": vectors[0], + }, + { + "id": "sk-2", + "text": texts[1], + "category": "rag", + "embedding": vectors[1], + }, + ] + ) + + query_vector = embedder.embed_query("How can I build an AI application?") + + results = await collection.search( + vector=query_vector, + top=2, + ) + + async for result in results.results: + print(result.score, "-", result.record["text"]) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/pyproject.toml b/pyproject.toml index c85b5ef..5acc3c1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,7 @@ ingest = ["pypdf>=4.0", "python-docx>=1.1", "python-pptx>=0.6", "openpyxl>=3.1", langchain = ["langchain-core>=0.3"] llamaindex = ["llama-index-core>=0.11"] crewai = ["crewai>=0.70,!=1.14.0; python_version >= '3.10'"] +semantic-kernel = ["semantic-kernel>=1.44"] haystack = ["haystack-ai>=2.0"] dspy = ["dspy>=3.3; python_version >= '3.10'"] all = [ @@ -67,6 +68,7 @@ all = [ "llama-index-core>=0.11", "haystack-ai>=2.0", "crewai>=0.70,!=1.14.0; python_version >= '3.10'", + "semantic-kernel>=1.44", "dspy>=3.3; python_version >= '3.10'", ] diff --git a/src/dynavec/integrations/semantic_kernel.py b/src/dynavec/integrations/semantic_kernel.py new file mode 100644 index 0000000..7ba4614 --- /dev/null +++ b/src/dynavec/integrations/semantic_kernel.py @@ -0,0 +1,308 @@ +"""Semantic Kernel vector store adapter for dynavec.""" + +from __future__ import annotations + +import ast +from collections.abc import Sequence +from typing import Any, Generic + +from ..client import Dynavec +from ..exceptions import MissingDependencyError +from ..models import Document as DVDocument + +try: + from semantic_kernel.data.vector import ( + KernelSearchResults, + SearchType, + TKey, + TModel, + VectorSearch, + VectorSearchOptions, + VectorSearchResult, + VectorStore, + VectorStoreCollection, + VectorStoreCollectionDefinition, + ) +except ImportError as exc: # pragma: no cover - import guard + raise MissingDependencyError( + "DynavecSemanticKernelStore", + "semantic-kernel", + "semantic-kernel", + ) from exc + + +class DynavecCollection( + VectorStoreCollection[TKey, TModel], + VectorSearch[TKey, TModel], + Generic[TKey, TModel], +): + """Semantic Kernel collection backed by Dynavec.""" + + supported_search_types = {SearchType.VECTOR} + + def __init__( + self, + client: Dynavec, + record_type: type[TModel], + definition: VectorStoreCollectionDefinition | None = None, + collection_name: str | None = None, + **kwargs: Any, + ) -> None: + super().__init__( + record_type=record_type, + definition=definition, + collection_name=collection_name, + **kwargs, + ) + self._client = client + self._namespace = collection_name or "default" + + async def _inner_upsert( + self, + records: Sequence[Any], + **kwargs: Any, + ) -> Sequence[TKey]: + result = self._client.upsert( + list(records), + namespace=self._namespace, + ) + return result.ids + + async def _inner_get( + self, + keys: Sequence[TKey] | None = None, + options=None, + **kwargs: Any, + ): + if keys is None: + raise NotImplementedError("Dynavec connector currently supports get by key only.") + + results = self._client.get( + [str(key) for key in keys], + namespace=self._namespace, + ) + + return [ + DVDocument( + id=result.id, + text=result.text, + vector=result.vector, + metadata=result.metadata, + ) + for result in results + ] + + async def _inner_delete( + self, + keys: Sequence[TKey], + **kwargs: Any, + ) -> None: + self._client.delete( + [str(key) for key in keys], + namespace=self._namespace, + ) + + async def ensure_collection_exists(self, **kwargs: Any) -> None: + return None + + async def collection_exists(self, **kwargs: Any) -> bool: + return True + + async def ensure_collection_deleted(self, **kwargs: Any) -> None: + return None + + async def _inner_search( + self, + search_type: SearchType, + options: VectorSearchOptions, + values: Any | None = None, + vector: Sequence[float | int] | None = None, + **kwargs: Any, + ) -> KernelSearchResults[VectorSearchResult[TModel]]: + search_filter = self._build_filter(options.filter) + + if vector is not None: + results = self._client.search( + vector=list(vector), + top_k=options.skip + options.top, + namespace=self._namespace, + filter=search_filter, + include_vectors=options.include_vectors, + ) + else: + results = self._client.search( + query=values, + top_k=options.skip + options.top, + namespace=self._namespace, + filter=search_filter, + include_vectors=options.include_vectors, + ) + + results = results[options.skip : options.skip + options.top] + + return KernelSearchResults( + results=self._get_vector_search_results_from_results( + results, + options, + ), + ) + + def _get_record_from_result(self, result: Any) -> DVDocument: + return DVDocument( + id=result.id, + text=result.text, + vector=result.vector, + metadata=result.metadata, + ) + + def _get_score_from_result(self, result: Any) -> float | None: + return result.score + + def _serialize_dicts_to_store_models( + self, records: list[dict[str, Any]], **kwargs: Any + ) -> list[DVDocument]: + key_name = self._key_field_storage_name + + vector_field = self.definition.vector_fields[0] if self.definition.vector_fields else None + vector_name = vector_field.storage_name or vector_field.name if vector_field else None + + documents = [] + + for record in records: + record_id = str(record[key_name]) + vector = record.get(vector_name) if vector_name else None + + text = None + metadata = {} + + for field in self.definition.data_fields: + field_name = field.storage_name or field.name + value = record.get(field_name) + + if field.name in {"text", "content"}: + text = value + elif value is not None: + metadata[field_name] = value + + documents.append( + DVDocument( + id=record_id, + text=text, + vector=vector, + metadata=metadata, + ) + ) + + return documents + + def _deserialize_store_models_to_dicts( + self, records: list[DVDocument], **kwargs: Any + ) -> list[dict[str, Any]]: + result = [] + + for record in records: + data = { + self._key_field_storage_name: record.id, + } + + for field in self.definition.vector_fields: + field_name = field.storage_name or field.name + data[field_name] = record.vector + + for field in self.definition.data_fields: + field_name = field.storage_name or field.name + + if field.name in {"text", "content"}: + data[field_name] = record.text + else: + data[field_name] = record.metadata.get(field_name) + + result.append(data) + + return result + + def _lambda_parser(self, node: ast.AST) -> Any: + if isinstance(node, ast.Name): + return node.id + + if isinstance(node, ast.Constant): + return node.value + + if isinstance(node, (ast.List, ast.Tuple, ast.Set)): + return [self._lambda_parser(element) for element in node.elts] + + if isinstance(node, ast.Subscript): + return self._lambda_parser(node.slice) + + if isinstance(node, ast.Attribute): + return node.attr + + if isinstance(node, ast.Compare): + field = self._lambda_parser(node.left) + filters = [] + + for operator, comparator in zip(node.ops, node.comparators): + value = self._lambda_parser(comparator) + + if isinstance(operator, ast.Eq): + filters.append({field: value}) + elif isinstance(operator, ast.NotEq): + filters.append({field: {"$ne": value}}) + elif isinstance(operator, ast.Gt): + filters.append({field: {"$gt": value}}) + elif isinstance(operator, ast.GtE): + filters.append({field: {"$gte": value}}) + elif isinstance(operator, ast.Lt): + filters.append({field: {"$lt": value}}) + elif isinstance(operator, ast.LtE): + filters.append({field: {"$lte": value}}) + elif isinstance(operator, ast.In): + filters.append({field: {"$in": value}}) + elif isinstance(operator, ast.NotIn): + filters.append({field: {"$nin": value}}) + else: + raise ValueError(f"Unsupported comparison operator: {type(operator).__name__}") + + return filters[0] if len(filters) == 1 else {"$and": filters} + + if isinstance(node, ast.BoolOp): + values = [self._lambda_parser(value) for value in node.values] + + if isinstance(node.op, ast.And): + return {"$and": values} + + if isinstance(node.op, ast.Or): + return {"$or": values} + + raise ValueError(f"Unsupported filter expression: {type(node).__name__}") + + +class DynavecStore(VectorStore): + """Semantic Kernel VectorStore backed by Dynavec.""" + + def __init__( + self, + client: Dynavec, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self._client = client + + def get_collection( + self, + record_type: type[TModel], + *, + definition: VectorStoreCollectionDefinition | None = None, + collection_name: str | None = None, + **kwargs: Any, + ) -> DynavecCollection: + return DynavecCollection( + client=self._client, + record_type=record_type, + definition=definition, + collection_name=collection_name, + **kwargs, + ) + + async def list_collection_names(self, **kwargs: Any) -> Sequence[str]: + return [] diff --git a/tests/test_semantic_kernel.py b/tests/test_semantic_kernel.py new file mode 100644 index 0000000..4028e4f --- /dev/null +++ b/tests/test_semantic_kernel.py @@ -0,0 +1,271 @@ +"""Tests for the Semantic Kernel vector store integration.""" + +import pytest + +pytest.importorskip("semantic_kernel") + +from semantic_kernel.data.vector import ( # noqa: E402 + FieldTypes, + VectorStoreCollectionDefinition, + VectorStoreField, +) + +from dynavec.integrations.semantic_kernel import ( # noqa: E402 + DynavecCollection, + DynavecStore, +) +from dynavec.models import SearchResult # noqa: E402 + + +class _FakeResult: + ids = ["doc-1"] + + +class _FakeClient: + def __init__(self): + self.calls = [] + + def upsert(self, documents, **kwargs): + self.calls.append(("upsert", documents, kwargs)) + return _FakeResult() + + def search(self, **kwargs): + self.calls.append(("search", kwargs)) + return [ + SearchResult( + id="doc-1", + score=0.95, + text="Banking knowledge", + metadata={"category": "banking"}, + vector=[0.1, 0.2, 0.3], + ), + SearchResult( + id="doc-2", + score=0.85, + text="Finance knowledge", + metadata={"category": "finance"}, + vector=[0.4, 0.5, 0.6], + ), + ] + + def get(self, ids, **kwargs): + self.calls.append(("get", ids, kwargs)) + return [ + SearchResult( + id="doc-1", + score=1.0, + text="Banking knowledge", + metadata={"category": "banking"}, + vector=[0.1, 0.2, 0.3], + ) + ] + + def delete(self, ids, **kwargs): + self.calls.append(("delete", ids, kwargs)) + + +def _definition(): + return VectorStoreCollectionDefinition( + fields=[ + VectorStoreField( + field_type=FieldTypes.KEY, + name="id", + ), + VectorStoreField( + field_type=FieldTypes.DATA, + name="text", + ), + VectorStoreField( + field_type=FieldTypes.DATA, + name="category", + ), + VectorStoreField( + field_type=FieldTypes.VECTOR, + name="embedding", + dimensions=3, + ), + ], + ) + + +@pytest.mark.asyncio +async def test_upsert(): + client = _FakeClient() + collection = DynavecCollection( + client=client, + record_type=dict, + definition=_definition(), + collection_name="kb", + ) + + ids = await collection.upsert( + [ + { + "id": "doc-1", + "text": "Banking knowledge", + "embedding": [0.1, 0.2, 0.3], + } + ] + ) + + assert ids == ["doc-1"] + + _, documents, kwargs = client.calls[0] + assert len(documents) == 1 + assert documents[0].id == "doc-1" + assert documents[0].text == "Banking knowledge" + assert documents[0].vector == [0.1, 0.2, 0.3] + assert kwargs == {"namespace": "kb"} + + +@pytest.mark.asyncio +async def test_search(): + client = _FakeClient() + collection = DynavecCollection( + client=client, + record_type=dict, + definition=_definition(), + collection_name="kb", + ) + + results = await collection.search( + values="banking", + top=2, + ) + + search_results = [result async for result in results.results] + + assert len(search_results) == 2 + + first = search_results[0] + assert first.record["id"] == "doc-1" + assert first.record["text"] == "Banking knowledge" + assert first.record["category"] == "banking" + assert first.score == 0.95 + + +@pytest.mark.asyncio +async def test_search_with_filter(): + client = _FakeClient() + collection = DynavecCollection( + client=client, + record_type=dict, + definition=_definition(), + collection_name="kb", + ) + + results = await collection.search( + values="banking", + filter=lambda x: x["category"] == "banking", + top=2, + ) + + search_results = [result async for result in results.results] + + assert len(search_results) == 2 + assert client.calls[-1] == ( + "search", + { + "query": "banking", + "top_k": 2, + "namespace": "kb", + "filter": {"category": "banking"}, + "include_vectors": False, + }, + ) + + +@pytest.mark.asyncio +async def test_get(): + client = _FakeClient() + collection = DynavecCollection( + client=client, + record_type=dict, + definition=_definition(), + collection_name="kb", + ) + + results = await collection.get(["doc-1"]) + + assert len(results) == 1 + assert results[0]["id"] == "doc-1" + assert results[0]["text"] == "Banking knowledge" + assert results[0]["category"] == "banking" + assert client.calls[-1] == ( + "get", + ["doc-1"], + {"namespace": "kb"}, + ) + + +@pytest.mark.asyncio +async def test_delete(): + client = _FakeClient() + collection = DynavecCollection( + client=client, + record_type=dict, + definition=_definition(), + collection_name="kb", + ) + + await collection.delete(["doc-1"]) + + assert client.calls[-1] == ( + "delete", + ["doc-1"], + {"namespace": "kb"}, + ) + + +@pytest.mark.asyncio +async def test_search_with_gte_filter(): + client = _FakeClient() + collection = DynavecCollection( + client=client, + record_type=dict, + definition=_definition(), + collection_name="kb", + ) + + results = await collection.search( + values="banking", + filter=lambda x: x["score"] >= 80, + top=2, + ) + + [result async for result in results.results] + + assert client.calls[-1] == ( + "search", + { + "query": "banking", + "top_k": 2, + "namespace": "kb", + "filter": {"score": {"$gte": 80}}, + "include_vectors": False, + }, + ) + + +def test_store_get_collection(): + client = _FakeClient() + store = DynavecStore(client=client) + + collection = store.get_collection( + record_type=dict, + definition=_definition(), + collection_name="kb", + ) + + assert isinstance(collection, DynavecCollection) + assert collection._client is client + assert collection._namespace == "kb" + + +@pytest.mark.asyncio +async def test_store_list_collection_names(): + store = DynavecStore(client=_FakeClient()) + + names = await store.list_collection_names() + + assert names == []