genkit-plugin-dev-local-vectorstore 0.3.0.dev1__tar.gz → 0.3.1__tar.gz

{genkit_plugin_dev_local_vectorstore-0.3.0.dev1 → genkit_plugin_dev_local_vectorstore-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: genkit-plugin-dev-local-vectorstore
-Version: 0.3.0.dev1
+Version: 0.3.1
 Summary: Genkit Local Vector Store Plugin
 Author: Google
 License: Apache-2.0

{genkit_plugin_dev_local_vectorstore-0.3.0.dev1 → genkit_plugin_dev_local_vectorstore-0.3.1}/pyproject.toml

@@ -27,7 +27,7 @@ license = { text = "Apache-2.0" }
 name = "genkit-plugin-dev-local-vectorstore"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "0.3.0.dev1"
+version = "0.3.1"
 
 [build-system]
 build-backend = "hatchling.build"

{genkit_plugin_dev_local_vectorstore-0.3.0.dev1/src/genkit/plugins/dev_local_vector_store → genkit_plugin_dev_local_vectorstore-0.3.1/src/genkit/plugins/dev_local_vectorstore}/__init__.py

@@ -13,3 +13,9 @@
 # limitations under the License.
 #
 # SPDX-License-Identifier: Apache-2.0
+
+from .plugin_api import DevLocalVectorStore
+
+__all__ = [
+    DevLocalVectorStore.__name__,
+]

{genkit_plugin_dev_local_vectorstore-0.3.0.dev1/src/genkit/plugins/dev_local_vector_store → genkit_plugin_dev_local_vectorstore-0.3.1/src/genkit/plugins/dev_local_vectorstore}/constant.py

@@ -15,19 +15,11 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-from typing import Any
-
 from pydantic import BaseModel
 
 from genkit.types import DocumentData, Embedding
 
 
-class Params(BaseModel):
-    index_name: str
-    embedder: str
-    embedder_options: dict[str, Any] | None = None
-
-
 class DbValue(BaseModel):
     doc: DocumentData
     embedding: Embedding

{genkit_plugin_dev_local_vectorstore-0.3.0.dev1/src/genkit/plugins/dev_local_vector_store → genkit_plugin_dev_local_vectorstore-0.3.1/src/genkit/plugins/dev_local_vectorstore}/indexer.py

@@ -20,19 +20,21 @@ import json
 from hashlib import md5
 
 from genkit.blocks.document import Document
-from genkit.plugins.dev_local_vector_store.constant import DbValue
-from genkit.plugins.dev_local_vector_store.local_vector_store_api import (
+from genkit.codec import dump_json
+from genkit.types import DocumentData, Embedding
+
+from .constant import DbValue
+from .local_vector_store_api import (
     LocalVectorStoreAPI,
 )
-from genkit.types import Docs, Embedding
 
 
 class DevLocalVectorStoreIndexer(LocalVectorStoreAPI):
-    async def index(self, docs: Docs) -> None:
+    async def index(self, docs: list[DocumentData]) -> None:
         data = self._load_filestore()
         tasks = []
 
-        for doc_data in docs.root:
+        for doc_data in docs:
             tasks.append(
                 self.process_document(
                     document=Document.from_document_data(document_data=doc_data),
@@ -43,17 +45,17 @@ class DevLocalVectorStoreIndexer(LocalVectorStoreAPI):
         await asyncio.gather(*tasks)
 
         with open(self.index_file_name, 'w', encoding='utf-8') as f:
-            json.dump(data, f, indent=2)
+            f.write(dump_json(self._serialize_data(data=data), indent=2))
 
     async def process_document(self, document: Document, data: dict[str, DbValue]) -> None:
         embeddings = await self.ai.embed(
-            embedder=self.params.embedder,
+            embedder=self.embedder,
             documents=[document],
-            options=self.params.embedder_options,
+            options=self.embedder_options,
         )
         embedding_docs = document.get_embedding_documents(embeddings.embeddings)
 
-        for embedding, emb_doc in zip(embeddings, embedding_docs, strict=False):
+        for embedding, emb_doc in zip(embeddings.embeddings, embedding_docs, strict=False):
             self._add_document(data=data, embedding=embedding, doc=emb_doc)
 
     def _add_document(

{genkit_plugin_dev_local_vectorstore-0.3.0.dev1/src/genkit/plugins/dev_local_vector_store → genkit_plugin_dev_local_vectorstore-0.3.1/src/genkit/plugins/dev_local_vectorstore}/local_vector_store_api.py

@@ -23,19 +23,22 @@ from functools import cached_property
 from typing import Any
 
 from genkit.ai import Genkit
-from genkit.plugins.dev_local_vector_store.constant import DbValue, Params
+
+from .constant import DbValue
 
 
 class LocalVectorStoreAPI(ABC):
     _LOCAL_FILESTORE_TEMPLATE = '__db_{index_name}.json'
 
-    def __init__(self, ai: Genkit, params: Params):
+    def __init__(self, ai: Genkit, index_name: str, embedder: str, embedder_options: dict[str, Any] | None = None):
         self.ai = ai
-        self.params = params
+        self.index_name = index_name
+        self.embedder = embedder
+        self.embedder_options = embedder_options
 
     @cached_property
     def index_file_name(self):
-        return self._LOCAL_FILESTORE_TEMPLATE.format(index_name=self.params.index_name)
+        return self._LOCAL_FILESTORE_TEMPLATE.format(index_name=self.index_name)
 
     def _load_filestore(self) -> dict[str, DbValue]:
         data = {}
@@ -53,7 +56,7 @@ class LocalVectorStoreAPI(ABC):
     def _serialize_data(data: dict[str, DbValue]) -> dict[str, Any]:
         data = copy.deepcopy(data)
         for k in data:
-            data[k] = DbValue.model_dump(data[k])
+            data[k] = DbValue.model_dump(data[k], exclude_none=True)
         return data
 
     @staticmethod

{genkit_plugin_dev_local_vectorstore-0.3.0.dev1/src/genkit/plugins/dev_local_vector_store → genkit_plugin_dev_local_vectorstore-0.3.1/src/genkit/plugins/dev_local_vectorstore}/plugin_api.py

@@ -16,29 +16,19 @@
 
 """Local file-based vectorstore plugin that provides retriever and indexer for Genkit."""
 
+from typing import Any
+
 from genkit.ai import GenkitRegistry, Plugin
 from genkit.core.action import Action
-from genkit.plugins.dev_local_vector_store.constant import Params
-from genkit.plugins.dev_local_vector_store.indexer import (
+from genkit.types import Docs
+
+from .indexer import (
     DevLocalVectorStoreIndexer,
 )
-from genkit.plugins.dev_local_vector_store.retriever import (
+from .retriever import (
     DevLocalVectorStoreRetriever,
     RetrieverOptionsSchema,
 )
-from genkit.types import Docs
-
-
-def dev_local_vectorstore_name(name: str) -> str:
-    """Create a Dev Local Vector Store action name.
-
-    Args:
-        name: Base name for the action.
-
-    Returns:
-        The fully qualified Dev Local Vector Store action name.
-    """
-    return f'devLocalVectorstore/{name}'
 
 
 class DevLocalVectorStore(Plugin):
@@ -50,8 +40,10 @@ class DevLocalVectorStore(Plugin):
     name = 'devLocalVectorstore'
     _indexers: dict[str, DevLocalVectorStoreIndexer] = {}
 
-    def __init__(self, params: list[Params]):
-        self.params = params
+    def __init__(self, name: str, embedder: str, embedder_options: dict[str, Any] | None = None):
+        self.index_name = name
+        self.embedder = embedder
+        self.embedder_options = embedder_options
 
     def initialize(self, ai: GenkitRegistry) -> None:
         """Initialize the plugin by registering actions with the registry.
@@ -65,12 +57,10 @@ class DevLocalVectorStore(Plugin):
         Returns:
             None
         """
-        for params in self.params:
-            self._configure_dev_local_retriever(ai=ai, params=params)
-            self._configure_dev_local_indexer(ai=ai, params=params)
+        self._configure_dev_local_retriever(ai=ai)
+        self._configure_dev_local_indexer(ai=ai)
 
-    @classmethod
-    def _configure_dev_local_retriever(cls, ai: GenkitRegistry, params: Params) -> Action:
+    def _configure_dev_local_retriever(self, ai: GenkitRegistry) -> Action:
         """Registers Local Vector Store retriever for provided parameters.
 
         Args:
@@ -82,17 +72,18 @@ class DevLocalVectorStore(Plugin):
         """
         retriever = DevLocalVectorStoreRetriever(
             ai=ai,
-            params=params,
+            index_name=self.index_name,
+            embedder=self.embedder,
+            embedder_options=self.embedder_options,
         )
 
         return ai.define_retriever(
-            name=dev_local_vectorstore_name(params.index_name),
+            name=self.index_name,
             config_schema=RetrieverOptionsSchema,
             fn=retriever.retrieve,
         )
 
-    @classmethod
-    def _configure_dev_local_indexer(cls, ai: GenkitRegistry, params: Params) -> Action:
+    def _configure_dev_local_indexer(self, ai: GenkitRegistry) -> Action:
         """Registers Local Vector Store indexer for provided parameters.
 
         Args:
@@ -104,10 +95,12 @@ class DevLocalVectorStore(Plugin):
         """
         indexer = DevLocalVectorStoreIndexer(
             ai=ai,
-            params=params,
+            index_name=self.index_name,
+            embedder=self.embedder,
+            embedder_options=self.embedder_options,
         )
 
-        cls._indexers[params.index_name] = indexer
+        DevLocalVectorStore._indexers[self.index_name] = indexer
 
     @classmethod
     async def index(cls, index_name: str, documents: Docs) -> None:

{genkit_plugin_dev_local_vectorstore-0.3.0.dev1/src/genkit/plugins/dev_local_vector_store → genkit_plugin_dev_local_vectorstore-0.3.1/src/genkit/plugins/dev_local_vectorstore}/retriever.py

@@ -18,10 +18,9 @@
 from pydantic import BaseModel, Field
 
 from genkit.ai import ActionRunContext, Document
-from genkit.plugins.dev_local_vector_store.local_vector_store_api import (
-    LocalVectorStoreAPI,
-)
-from genkit.types import Embedding, RetrieverRequest
+from genkit.types import Embedding, RetrieverRequest, RetrieverResponse
+
+from .local_vector_store_api import LocalVectorStoreAPI
 
 
 class ScoredDocument(BaseModel):
@@ -34,29 +33,31 @@ class RetrieverOptionsSchema(BaseModel):
 
 
 class DevLocalVectorStoreRetriever(LocalVectorStoreAPI):
-    async def retrieve(self, request: RetrieverRequest, _: ActionRunContext):
+    async def retrieve(self, request: RetrieverRequest, _: ActionRunContext) -> RetrieverResponse:
         document = Document.from_document_data(document_data=request.query)
         embeddings = await self.ai.embed(
-            embedder=self.params.embedder,
+            embedder=self.embedder,
             documents=[document],
-            options=self.params.embedder_options,
+            options=self.embedder_options,
         )
-        if self.params.embedder_options:
-            k = self.params.embedder_options.get('limit') or 3
+        if self.embedder_options:
+            k = self.embedder_options.get('limit') or 3
         else:
             k = 3
-        return self._get_closest_documents(
+        docs = self._get_closest_documents(
             k=k,
             query_embeddings=embeddings.embeddings[0],
         )
 
+        return RetrieverResponse(documents=[d.document for d in docs])
+
     def _get_closest_documents(self, k: int, query_embeddings: Embedding) -> list[ScoredDocument]:
         db = self._load_filestore()
         scored_documents = []
 
         for val in db.values():
             this_embedding = val.embedding.embedding
-            score = self.cosine_similarity(query_embeddings, this_embedding)
+            score = self.cosine_similarity(query_embeddings.embedding, this_embedding)
             scored_documents.append(
                 ScoredDocument(
                     score=score,
@@ -68,9 +69,9 @@ class DevLocalVectorStoreRetriever(LocalVectorStoreAPI):
         return scored_documents[:k]
 
     @classmethod
-    def cosine_similarity(cls, a: list[int], b: list[int]) -> float:
+    def cosine_similarity(cls, a: list[float], b: list[float]) -> float:
         return cls.dot(a, b) / ((cls.dot(a, a) ** 0.5) * (cls.dot(b, b) ** 0.5))
 
     @staticmethod
-    def dot(a: list[int], b: list[int]) -> float:
-        return sum(a * b for a, b in zip(a, b, strict=False))
+    def dot(a: list[float], b: list[float]) -> float:
+        return sum(av * bv for av, bv in zip(a, b, strict=False))