graphrag-vectors 3.0.0__tar.gz

graphrag_vectors-3.0.0/.gitignore

@@ -0,0 +1,65 @@
+# Python Artifacts
+python/*/lib/
+dist/
+build/
+*.egg-info/
+
+# Test Output
+.coverage
+coverage/
+licenses.txt
+examples_notebooks/*/data
+tests/fixtures/cache
+tests/fixtures/*/cache
+tests/fixtures/*/output
+output/lancedb
+
+
+# Random
+.DS_Store
+*.log*
+.venv
+venv/
+.conda
+.tmp
+packages/graphrag-llm/notebooks/metrics
+packages/graphrag-llm/notebooks/cache
+
+.env
+build.zip
+
+.turbo
+
+__pycache__
+
+.pipeline
+
+# Azurite
+temp_azurite/
+__azurite*.json
+__blobstorage*.json
+__blobstorage__/
+
+# Getting started example
+ragtest/
+.ragtest/
+.pipelines
+.pipeline
+
+
+# mkdocs
+site/
+
+# Docs migration
+docsite/
+.yarn/
+.pnp*
+
+# PyCharm
+.idea/
+
+# Jupyter notebook
+.ipynb_checkpoints/
+
+# Root build assets
+packages/*/LICENSE

graphrag_vectors-3.0.0/LICENSE

@@ -0,0 +1,21 @@
+    MIT License
+
+    Copyright (c) Microsoft Corporation.
+
+    Permission is hereby granted, free of charge, to any person obtaining a copy
+    of this software and associated documentation files (the "Software"), to deal
+    in the Software without restriction, including without limitation the rights
+    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+    copies of the Software, and to permit persons to whom the Software is
+    furnished to do so, subject to the following conditions:
+
+    The above copyright notice and this permission notice shall be included in all
+    copies or substantial portions of the Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    SOFTWARE

graphrag_vectors-3.0.0/PKG-INFO

@@ -0,0 +1,136 @@
+Metadata-Version: 2.4
+Name: graphrag-vectors
+Version: 3.0.0
+Summary: GraphRAG vector store package.
+Project-URL: Source, https://github.com/microsoft/graphrag
+Author: Mónica Carvajal
+Author-email: Alonso Guevara Fernández <alonsog@microsoft.com>, Andrés Morales Esquivel <andresmor@microsoft.com>, Chris Trevino <chtrevin@microsoft.com>, David Tittsworth <datittsw@microsoft.com>, Dayenne de Souza <ddesouza@microsoft.com>, Derek Worthen <deworthe@microsoft.com>, Gaudy Blanco Meneses <gaudyb@microsoft.com>, Ha Trinh <trinhha@microsoft.com>, Jonathan Larson <jolarso@microsoft.com>, Josh Bradley <joshbradley@microsoft.com>, Kate Lytvynets <kalytv@microsoft.com>, Kenny Zhang <zhangken@microsoft.com>, Nathan Evans <naevans@microsoft.com>, Rodrigo Racanicci <rracanicci@microsoft.com>, Sarah Smith <smithsarah@microsoft.com>
+License: MIT
+License-File: LICENSE
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: <3.14,>=3.11
+Requires-Dist: azure-core~=1.32
+Requires-Dist: azure-cosmos~=4.9
+Requires-Dist: azure-identity~=1.19
+Requires-Dist: azure-search-documents~=11.6
+Requires-Dist: graphrag-common==3.0.0
+Requires-Dist: lancedb~=0.24.1
+Requires-Dist: numpy~=2.1
+Requires-Dist: pyarrow~=22.0
+Requires-Dist: pydantic~=2.10
+Description-Content-Type: text/markdown
+
+# GraphRAG Vectors
+
+This package provides vector store implementations for GraphRAG with support for multiple backends including LanceDB, Azure AI Search, and Azure Cosmos DB. It offers both a convenient configuration-driven API and direct factory access for creating and managing vector stores with flexible index schema definitions.
+
+## Basic usage with the utility function (recommended)
+
+This demonstrates the recommended approach to create a vector store using the create_vector_store convenience function with configuration objects that specify the store type and index schema. The example shows setting up a LanceDB vector store with a defined index configuration, then connecting to it and creating the index for vector operations.
+
+```python
+from graphrag_vectors import (
+    create_vector_store,
+    VectorStoreType,
+    IndexSchema,
+)
+
+# Create a vector store using the convenience function
+store_config = VectorStoreConfig(
+    type="lancedb",
+    db_uri="lance"
+)
+
+schema_config = IndexSchema(
+    index_name="my_index",
+    vector_size=1536,
+)
+
+vector_store = create_vector_store(
+    config=store_config
+    index_schema=schema_config,
+)
+
+vector_store.connect()
+vector_store.create_index()
+```
+
+## Basic usage implementing the factory directly
+
+This example shows a different approach to create vector stores by directly using the vector_store_factory with enum types and dictionary-based initialization arguments. This method provides more direct control over the factory creation process while bypassing the convenience function layer.
+
+```python
+from graphrag_vectors import (
+    VectorStoreFactory,
+    vector_store_factory,
+    VectorStoreType,
+    IndexSchema,
+)
+
+# Create a vector store using the factory
+schema_config = IndexSchema(
+    index_name="my_index",
+    vector_size=1536,
+)
+
+vector_store = vector_store_factory.create(
+    VectorStoreType.LanceDB,
+    {
+        "index_schema": schema_config,
+        "db_uri": "./lancedb"
+    }
+)
+
+vector_store.connect()
+vector_store.create_index()
+```
+
+## Supported Vector Stores
+
+- **LanceDB**: Local vector database
+- **Azure AI Search**: Azure's managed search service with vector capabilities
+- **Azure Cosmos DB**: Azure's NoSQL database with vector search support
+
+## Custom Vector Store
+
+You can register custom vector store implementations:
+
+```python
+from graphrag_vectors import VectorStore, register_vector_store, create_vector_store
+
+class MyCustomVectorStore(VectorStore):
+    def __init__(self, my_param):
+        self.my_param = my_param
+
+    def connect(self):
+        # Implementation
+        pass
+    
+    def create_index(self):
+        # Implementation
+        pass
+    
+    # ... implement other required methods
+
+# Register your custom implementation
+register_vector_store("my_custom_store", MyCustomVectorStore)
+
+# Use your custom vector store
+config = VectorStoreConfig(
+    type="my_custom_store",
+    my_param="something"
+)
+custom_store = create_vector_store(
+    config=config,
+    index_schema=schema_config,
+)
+```
+
+## Configuration
+
+Vector stores are configured using:
+- `VectorStoreConfig`: baseline parameters for the store
+- `IndexSchema`: Schema configuration for the specific index to create/connect to (index name, field names, vector size)

graphrag_vectors-3.0.0/README.md

@@ -0,0 +1,111 @@
+# GraphRAG Vectors
+
+This package provides vector store implementations for GraphRAG with support for multiple backends including LanceDB, Azure AI Search, and Azure Cosmos DB. It offers both a convenient configuration-driven API and direct factory access for creating and managing vector stores with flexible index schema definitions.
+
+## Basic usage with the utility function (recommended)
+
+This demonstrates the recommended approach to create a vector store using the create_vector_store convenience function with configuration objects that specify the store type and index schema. The example shows setting up a LanceDB vector store with a defined index configuration, then connecting to it and creating the index for vector operations.
+
+```python
+from graphrag_vectors import (
+    create_vector_store,
+    VectorStoreType,
+    IndexSchema,
+)
+
+# Create a vector store using the convenience function
+store_config = VectorStoreConfig(
+    type="lancedb",
+    db_uri="lance"
+)
+
+schema_config = IndexSchema(
+    index_name="my_index",
+    vector_size=1536,
+)
+
+vector_store = create_vector_store(
+    config=store_config
+    index_schema=schema_config,
+)
+
+vector_store.connect()
+vector_store.create_index()
+```
+
+## Basic usage implementing the factory directly
+
+This example shows a different approach to create vector stores by directly using the vector_store_factory with enum types and dictionary-based initialization arguments. This method provides more direct control over the factory creation process while bypassing the convenience function layer.
+
+```python
+from graphrag_vectors import (
+    VectorStoreFactory,
+    vector_store_factory,
+    VectorStoreType,
+    IndexSchema,
+)
+
+# Create a vector store using the factory
+schema_config = IndexSchema(
+    index_name="my_index",
+    vector_size=1536,
+)
+
+vector_store = vector_store_factory.create(
+    VectorStoreType.LanceDB,
+    {
+        "index_schema": schema_config,
+        "db_uri": "./lancedb"
+    }
+)
+
+vector_store.connect()
+vector_store.create_index()
+```
+
+## Supported Vector Stores
+
+- **LanceDB**: Local vector database
+- **Azure AI Search**: Azure's managed search service with vector capabilities
+- **Azure Cosmos DB**: Azure's NoSQL database with vector search support
+
+## Custom Vector Store
+
+You can register custom vector store implementations:
+
+```python
+from graphrag_vectors import VectorStore, register_vector_store, create_vector_store
+
+class MyCustomVectorStore(VectorStore):
+    def __init__(self, my_param):
+        self.my_param = my_param
+
+    def connect(self):
+        # Implementation
+        pass
+    
+    def create_index(self):
+        # Implementation
+        pass
+    
+    # ... implement other required methods
+
+# Register your custom implementation
+register_vector_store("my_custom_store", MyCustomVectorStore)
+
+# Use your custom vector store
+config = VectorStoreConfig(
+    type="my_custom_store",
+    my_param="something"
+)
+custom_store = create_vector_store(
+    config=config,
+    index_schema=schema_config,
+)
+```
+
+## Configuration
+
+Vector stores are configured using:
+- `VectorStoreConfig`: baseline parameters for the store
+- `IndexSchema`: Schema configuration for the specific index to create/connect to (index name, field names, vector size)

graphrag_vectors-3.0.0/graphrag_vectors/__init__.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""GraphRAG vector store implementations."""
+
+from graphrag_vectors.index_schema import IndexSchema
+from graphrag_vectors.types import TextEmbedder
+from graphrag_vectors.vector_store import (
+    VectorStore,
+    VectorStoreDocument,
+    VectorStoreSearchResult,
+)
+from graphrag_vectors.vector_store_config import VectorStoreConfig
+from graphrag_vectors.vector_store_factory import (
+    VectorStoreFactory,
+    create_vector_store,
+    register_vector_store,
+    vector_store_factory,
+)
+from graphrag_vectors.vector_store_type import VectorStoreType
+
+__all__ = [
+    "IndexSchema",
+    "TextEmbedder",
+    "VectorStore",
+    "VectorStoreConfig",
+    "VectorStoreDocument",
+    "VectorStoreFactory",
+    "VectorStoreSearchResult",
+    "VectorStoreType",
+    "create_vector_store",
+    "register_vector_store",
+    "vector_store_factory",
+]

graphrag_vectors-3.0.0/graphrag_vectors/azure_ai_search.py

@@ -0,0 +1,173 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A package containing the Azure AI Search  vector store implementation."""
+
+from typing import Any
+
+from azure.core.credentials import AzureKeyCredential
+from azure.identity import DefaultAzureCredential
+from azure.search.documents import SearchClient
+from azure.search.documents.indexes import SearchIndexClient
+from azure.search.documents.indexes.models import (
+    HnswAlgorithmConfiguration,
+    HnswParameters,
+    SearchField,
+    SearchFieldDataType,
+    SearchIndex,
+    SimpleField,
+    VectorSearch,
+    VectorSearchAlgorithmMetric,
+    VectorSearchProfile,
+)
+from azure.search.documents.models import VectorizedQuery
+
+from graphrag_vectors.vector_store import (
+    VectorStore,
+    VectorStoreDocument,
+    VectorStoreSearchResult,
+)
+
+
+class AzureAISearchVectorStore(VectorStore):
+    """Azure AI Search vector storage implementation."""
+
+    index_client: SearchIndexClient
+
+    def __init__(
+        self,
+        url: str,
+        api_key: str | None = None,
+        audience: str | None = None,
+        vector_search_profile_name: str = "vectorSearchProfile",
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        if not url:
+            msg = "url must be provided for Azure AI Search."
+            raise ValueError(msg)
+        self.url = url
+        self.api_key = api_key
+        self.audience = audience
+        self.vector_search_profile_name = vector_search_profile_name
+
+    def connect(self) -> Any:
+        """Connect to AI search vector storage."""
+        audience_arg = (
+            {"audience": self.audience} if self.audience and not self.api_key else {}
+        )
+        self.db_connection = SearchClient(
+            endpoint=self.url,
+            index_name=self.index_name,
+            credential=(
+                AzureKeyCredential(self.api_key)
+                if self.api_key
+                else DefaultAzureCredential()
+            ),
+            **audience_arg,
+        )
+        self.index_client = SearchIndexClient(
+            endpoint=self.url,
+            credential=(
+                AzureKeyCredential(self.api_key)
+                if self.api_key
+                else DefaultAzureCredential()
+            ),
+            **audience_arg,
+        )
+
+    def create_index(self) -> None:
+        """Load documents into an Azure AI Search index."""
+        if (
+            self.index_name is not None
+            and self.index_name in self.index_client.list_index_names()
+        ):
+            self.index_client.delete_index(self.index_name)
+
+        # Configure vector search profile
+        vector_search = VectorSearch(
+            algorithms=[
+                HnswAlgorithmConfiguration(
+                    name="HnswAlg",
+                    parameters=HnswParameters(
+                        metric=VectorSearchAlgorithmMetric.COSINE
+                    ),
+                )
+            ],
+            profiles=[
+                VectorSearchProfile(
+                    name=self.vector_search_profile_name,
+                    algorithm_configuration_name="HnswAlg",
+                )
+            ],
+        )
+        # Configure the index
+        index = SearchIndex(
+            name=self.index_name,
+            fields=[
+                SimpleField(
+                    name=self.id_field,
+                    type=SearchFieldDataType.String,
+                    key=True,
+                ),
+                SearchField(
+                    name=self.vector_field,
+                    type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
+                    searchable=True,
+                    hidden=False,  # DRIFT needs to return the vector for client-side similarity
+                    vector_search_dimensions=self.vector_size,
+                    vector_search_profile_name=self.vector_search_profile_name,
+                ),
+            ],
+            vector_search=vector_search,
+        )
+        self.index_client.create_or_update_index(
+            index,
+        )
+
+    def load_documents(self, documents: list[VectorStoreDocument]) -> None:
+        """Load documents into an Azure AI Search index."""
+        batch = [
+            {
+                self.id_field: doc.id,
+                self.vector_field: doc.vector,
+            }
+            for doc in documents
+            if doc.vector is not None
+        ]
+
+        if len(batch) > 0:
+            self.db_connection.upload_documents(batch)
+
+    def similarity_search_by_vector(
+        self, query_embedding: list[float], k: int = 10
+    ) -> list[VectorStoreSearchResult]:
+        """Perform a vector-based similarity search."""
+        vectorized_query = VectorizedQuery(
+            vector=query_embedding, k_nearest_neighbors=k, fields=self.vector_field
+        )
+
+        response = self.db_connection.search(
+            vector_queries=[vectorized_query],
+        )
+
+        return [
+            VectorStoreSearchResult(
+                document=VectorStoreDocument(
+                    id=doc.get(self.id_field, ""),
+                    vector=doc.get(self.vector_field, []),
+                ),
+                # Cosine similarity between 0.333 and 1.000
+                # https://learn.microsoft.com/en-us/azure/search/hybrid-search-ranking#scores-in-a-hybrid-search-results
+                score=doc["@search.score"],
+            )
+            for doc in response
+        ]
+
+    def search_by_id(self, id: str) -> VectorStoreDocument:
+        """Search for a document by id."""
+        response = self.db_connection.get_document(id)
+        return VectorStoreDocument(
+            id=response.get(self.id_field, ""),
+            vector=response.get(self.vector_field, []),
+        )

graphrag_vectors-3.0.0/graphrag_vectors/cosmosdb.py

@@ -0,0 +1,244 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A package containing the CosmosDB vector store implementation."""
+
+from typing import Any
+
+from azure.cosmos import ContainerProxy, CosmosClient, DatabaseProxy
+from azure.cosmos.exceptions import CosmosHttpResponseError
+from azure.cosmos.partition_key import PartitionKey
+from azure.identity import DefaultAzureCredential
+
+from graphrag_vectors.vector_store import (
+    VectorStore,
+    VectorStoreDocument,
+    VectorStoreSearchResult,
+)
+
+
+class CosmosDBVectorStore(VectorStore):
+    """Azure CosmosDB vector storage implementation."""
+
+    _cosmos_client: CosmosClient
+    _database_client: DatabaseProxy
+    _container_client: ContainerProxy
+
+    def __init__(
+        self,
+        database_name: str,
+        connection_string: str | None = None,
+        url: str | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        if self.id_field != "id":
+            msg = "CosmosDB requires the id_field to be 'id'."
+            raise ValueError(msg)
+        if not connection_string and not url:
+            msg = "Either connection_string or url must be provided for CosmosDB."
+            raise ValueError(msg)
+
+        self.database_name = database_name
+        self.connection_string = connection_string
+        self.url = url
+
+    def connect(self) -> Any:
+        """Connect to CosmosDB vector storage."""
+        if self.connection_string:
+            self._cosmos_client = CosmosClient.from_connection_string(
+                self.connection_string
+            )
+        else:
+            self._cosmos_client = CosmosClient(
+                url=self.url, credential=DefaultAzureCredential()
+            )
+
+        self._create_database()
+        self._create_container()
+
+    def _create_database(self) -> None:
+        """Create the database if it doesn't exist."""
+        self._cosmos_client.create_database_if_not_exists(id=self.database_name)
+        self._database_client = self._cosmos_client.get_database_client(
+            self.database_name
+        )
+
+    def _delete_database(self) -> None:
+        """Delete the database if it exists."""
+        if self._database_exists():
+            self._cosmos_client.delete_database(self.database_name)
+
+    def _database_exists(self) -> bool:
+        """Check if the database exists."""
+        existing_database_names = [
+            database["id"] for database in self._cosmos_client.list_databases()
+        ]
+        return self.database_name in existing_database_names
+
+    def _create_container(self) -> None:
+        """Create the container if it doesn't exist."""
+        partition_key = PartitionKey(path=f"/{self.id_field}", kind="Hash")
+
+        # Define the container vector policy
+        vector_embedding_policy = {
+            "vectorEmbeddings": [
+                {
+                    "path": f"/{self.vector_field}",
+                    "dataType": "float32",
+                    "distanceFunction": "cosine",
+                    "dimensions": self.vector_size,
+                }
+            ]
+        }
+
+        # Define the vector indexing policy
+        indexing_policy = {
+            "indexingMode": "consistent",
+            "automatic": True,
+            "includedPaths": [{"path": "/*"}],
+            "excludedPaths": [
+                {"path": "/_etag/?"},
+                {"path": f"/{self.vector_field}/*"},
+            ],
+        }
+
+        # Currently, the CosmosDB emulator does not support the diskANN policy.
+        try:
+            # First try with the standard diskANN policy
+            indexing_policy["vectorIndexes"] = [
+                {"path": f"/{self.vector_field}", "type": "diskANN"}
+            ]
+
+            # Create the container and container client
+            self._database_client.create_container_if_not_exists(
+                id=self.index_name,
+                partition_key=partition_key,
+                indexing_policy=indexing_policy,
+                vector_embedding_policy=vector_embedding_policy,
+            )
+        except CosmosHttpResponseError:
+            # If diskANN fails (likely in emulator), retry without vector indexes
+            indexing_policy.pop("vectorIndexes", None)
+
+            # Create the container with compatible indexing policy
+            self._database_client.create_container_if_not_exists(
+                id=self.index_name,
+                partition_key=partition_key,
+                indexing_policy=indexing_policy,
+                vector_embedding_policy=vector_embedding_policy,
+            )
+
+        self._container_client = self._database_client.get_container_client(
+            self.index_name
+        )
+
+    def _delete_container(self) -> None:
+        """Delete the vector store container in the database if it exists."""
+        if self._container_exists():
+            self._database_client.delete_container(self.index_name)
+
+    def _container_exists(self) -> bool:
+        """Check if the container name exists in the database."""
+        existing_container_names = [
+            container["id"] for container in self._database_client.list_containers()
+        ]
+        return self.index_name in existing_container_names
+
+    def create_index(self) -> None:
+        """Load documents into CosmosDB."""
+        # Create a CosmosDB container on overwrite
+        self._delete_container()
+        self._create_container()
+
+        if self._container_client is None:
+            msg = "Container client is not initialized."
+            raise ValueError(msg)
+
+    def load_documents(self, documents: list[VectorStoreDocument]) -> None:
+        """Load documents into CosmosDB."""
+        # Upload documents to CosmosDB
+        for doc in documents:
+            if doc.vector is not None:
+                doc_json = {
+                    self.id_field: doc.id,
+                    self.vector_field: doc.vector,
+                }
+                self._container_client.upsert_item(doc_json)
+
+    def similarity_search_by_vector(
+        self, query_embedding: list[float], k: int = 10
+    ) -> list[VectorStoreSearchResult]:
+        """Perform a vector-based similarity search."""
+        if self._container_client is None:
+            msg = "Container client is not initialized."
+            raise ValueError(msg)
+
+        try:
+            query = f"SELECT TOP {k} c.{self.id_field}, c.{self.vector_field}, VectorDistance(c.{self.vector_field}, @embedding) AS SimilarityScore FROM c ORDER BY VectorDistance(c.{self.vector_field}, @embedding)"  # noqa: S608
+            query_params = [{"name": "@embedding", "value": query_embedding}]
+            items = list(
+                self._container_client.query_items(
+                    query=query,
+                    parameters=query_params,
+                    enable_cross_partition_query=True,
+                )
+            )
+        except (CosmosHttpResponseError, ValueError):
+            # Currently, the CosmosDB emulator does not support the VectorDistance function.
+            # For emulator or test environments - fetch all items and calculate distance locally
+            query = f"SELECT c.{self.id_field}, c.{self.vector_field} FROM c"  # noqa: S608
+            items = list(
+                self._container_client.query_items(
+                    query=query,
+                    enable_cross_partition_query=True,
+                )
+            )
+
+            # Calculate cosine similarity locally (1 - cosine distance)
+            from numpy import dot
+            from numpy.linalg import norm
+
+            def cosine_similarity(a, b):
+                if norm(a) * norm(b) == 0:
+                    return 0.0
+                return dot(a, b) / (norm(a) * norm(b))
+
+            # Calculate scores for all items
+            for item in items:
+                item_vector = item.get(self.vector_field, [])
+                similarity = cosine_similarity(query_embedding, item_vector)
+                item["SimilarityScore"] = similarity
+
+            # Sort by similarity score (higher is better) and take top k
+            items = sorted(
+                items, key=lambda x: x.get("SimilarityScore", 0.0), reverse=True
+            )[:k]
+
+        return [
+            VectorStoreSearchResult(
+                document=VectorStoreDocument(
+                    id=item.get(self.id_field, ""),
+                    vector=item.get(self.vector_field, []),
+                ),
+                score=item.get("SimilarityScore", 0.0),
+            )
+            for item in items
+        ]
+
+    def search_by_id(self, id: str) -> VectorStoreDocument:
+        """Search for a document by id."""
+        if self._container_client is None:
+            msg = "Container client is not initialized."
+            raise ValueError(msg)
+
+        item = self._container_client.read_item(item=id, partition_key=id)
+        return VectorStoreDocument(
+            id=item.get(self.id_field, ""),
+            vector=item.get(self.vector_field, []),
+        )
+
+    def clear(self) -> None:
+        """Clear the vector store."""
+        self._delete_container()
+        self._delete_database()

graphrag_vectors-3.0.0/graphrag_vectors/index_schema.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Parameterization settings for the default configuration."""
+
+import re
+
+from pydantic import BaseModel, Field, model_validator
+
+DEFAULT_VECTOR_SIZE: int = 3072
+
+VALID_IDENTIFIER_REGEX = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+
+
+def is_valid_field_name(field: str) -> bool:
+    """Check if a field name is valid for CosmosDB."""
+    return bool(VALID_IDENTIFIER_REGEX.match(field))
+
+
+class IndexSchema(BaseModel):
+    """The default configuration section for Vector Store Schema."""
+
+    index_name: str = Field(
+        description="The index name to use.", default="vector_index"
+    )
+
+    id_field: str = Field(
+        description="The ID field to use.",
+        default="id",
+    )
+
+    vector_field: str = Field(
+        description="The vector field to use.",
+        default="vector",
+    )
+
+    vector_size: int = Field(
+        description="The vector size to use.",
+        default=DEFAULT_VECTOR_SIZE,
+    )
+
+    def _validate_schema(self) -> None:
+        """Validate the schema."""
+        for field in [
+            self.id_field,
+            self.vector_field,
+        ]:
+            if not is_valid_field_name(field):
+                msg = f"Unsafe or invalid field name: {field}"
+                raise ValueError(msg)
+
+    @model_validator(mode="after")
+    def _validate_model(self):
+        """Validate the model."""
+        self._validate_schema()
+        return self

graphrag_vectors-3.0.0/graphrag_vectors/lancedb.py

@@ -0,0 +1,128 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""The LanceDB vector storage implementation package."""
+
+from typing import Any
+
+import lancedb
+import numpy as np
+import pyarrow as pa
+
+from graphrag_vectors.vector_store import (
+    VectorStore,
+    VectorStoreDocument,
+    VectorStoreSearchResult,
+)
+
+
+class LanceDBVectorStore(VectorStore):
+    """LanceDB vector storage implementation."""
+
+    def __init__(self, db_uri: str = "lancedb", **kwargs: Any):
+        super().__init__(**kwargs)
+        self.db_uri = db_uri
+
+    def connect(self) -> Any:
+        """Connect to the vector storage."""
+        self.db_connection = lancedb.connect(self.db_uri)
+
+        if self.index_name and self.index_name in self.db_connection.table_names():
+            self.document_collection = self.db_connection.open_table(self.index_name)
+
+    def create_index(self) -> None:
+        """Create index."""
+        dummy_vector = np.zeros(self.vector_size, dtype=np.float32)
+        flat_array = pa.array(dummy_vector, type=pa.float32())
+        vector_column = pa.FixedSizeListArray.from_arrays(flat_array, self.vector_size)
+
+        data = pa.table({
+            self.id_field: pa.array(["__DUMMY__"], type=pa.string()),
+            self.vector_field: vector_column,
+        })
+
+        self.document_collection = self.db_connection.create_table(
+            self.index_name if self.index_name else "",
+            data=data,
+            mode="overwrite",
+            schema=data.schema,
+        )
+
+        # Step 5: Create index now that schema exists
+        self.document_collection.create_index(
+            vector_column_name=self.vector_field, index_type="IVF_FLAT"
+        )
+
+    def load_documents(self, documents: list[VectorStoreDocument]) -> None:
+        """Load documents into vector storage."""
+        self.document_collection.delete(f"{self.id_field} = '__DUMMY__'")
+
+        # Step 1: Prepare data columns manually
+        ids = []
+        vectors = []
+
+        for document in documents:
+            self.vector_size = (
+                len(document.vector) if document.vector else self.vector_size
+            )
+            if document.vector is not None and len(document.vector) == self.vector_size:
+                ids.append(document.id)
+                vectors.append(np.array(document.vector, dtype=np.float32))
+
+        # Step 2: Handle empty case
+        if len(ids) == 0:
+            data = None
+        else:
+            # Step 3: Flatten the vectors and build FixedSizeListArray manually
+            flat_vector = np.concatenate(vectors).astype(np.float32)
+            flat_array = pa.array(flat_vector, type=pa.float32())
+            vector_column = pa.FixedSizeListArray.from_arrays(
+                flat_array, self.vector_size
+            )
+
+            # Step 4: Create PyArrow table (let schema be inferred)
+            data = pa.table({
+                self.id_field: pa.array(ids, type=pa.string()),
+                self.vector_field: vector_column,
+            })
+
+            if data:
+                self.document_collection.add(data)
+
+    def similarity_search_by_vector(
+        self, query_embedding: list[float] | np.ndarray, k: int = 10
+    ) -> list[VectorStoreSearchResult]:
+        """Perform a vector-based similarity search."""
+        query_embedding = np.array(query_embedding, dtype=np.float32)
+
+        docs = (
+            self.document_collection
+            .search(query=query_embedding, vector_column_name=self.vector_field)
+            .limit(k)
+            .to_list()
+        )
+        return [
+            VectorStoreSearchResult(
+                document=VectorStoreDocument(
+                    id=doc[self.id_field],
+                    vector=doc[self.vector_field],
+                ),
+                score=1 - abs(float(doc["_distance"])),
+            )
+            for doc in docs
+        ]
+
+    def search_by_id(self, id: str) -> VectorStoreDocument:
+        """Search for a document by id."""
+        doc = (
+            self.document_collection
+            .search()
+            .where(f"{self.id_field} == '{id}'", prefilter=True)
+            .to_list()
+        )
+        if doc:
+            return VectorStoreDocument(
+                id=doc[0][self.id_field],
+                vector=doc[0][self.vector_field],
+            )
+        return VectorStoreDocument(id=id, vector=None)

graphrag_vectors-3.0.0/graphrag_vectors/types.py

@@ -0,0 +1,8 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Common types for vector stores."""
+
+from collections.abc import Callable
+
+TextEmbedder = Callable[[str], list[float]]

graphrag_vectors-3.0.0/graphrag_vectors/vector_store.py

@@ -0,0 +1,81 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Base classes for vector stores."""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
+from graphrag_vectors.types import TextEmbedder
+
+
+@dataclass
+class VectorStoreDocument:
+    """A document that is stored in vector storage."""
+
+    id: str | int
+    """unique id for the document"""
+
+    vector: list[float] | None
+
+
+@dataclass
+class VectorStoreSearchResult:
+    """A vector storage search result."""
+
+    document: VectorStoreDocument
+    """Document that was found."""
+
+    score: float
+    """Similarity score between -1 and 1. Higher is more similar."""
+
+
+class VectorStore(ABC):
+    """The base class for vector storage data-access classes."""
+
+    def __init__(
+        self,
+        index_name: str = "vector_index",
+        id_field: str = "id",
+        vector_field: str = "vector",
+        vector_size: int = 3072,
+        **kwargs: Any,
+    ):
+        self.index_name = index_name
+        self.id_field = id_field
+        self.vector_field = vector_field
+        self.vector_size = vector_size
+
+    @abstractmethod
+    def connect(self) -> None:
+        """Connect to vector storage."""
+
+    @abstractmethod
+    def create_index(self) -> None:
+        """Create index."""
+
+    @abstractmethod
+    def load_documents(self, documents: list[VectorStoreDocument]) -> None:
+        """Load documents into the vector-store."""
+
+    @abstractmethod
+    def similarity_search_by_vector(
+        self, query_embedding: list[float], k: int = 10
+    ) -> list[VectorStoreSearchResult]:
+        """Perform ANN search by vector."""
+
+    def similarity_search_by_text(
+        self, text: str, text_embedder: TextEmbedder, k: int = 10
+    ) -> list[VectorStoreSearchResult]:
+        """Perform a text-based similarity search."""
+        query_embedding = text_embedder(text)
+        if query_embedding:
+            return self.similarity_search_by_vector(
+                query_embedding=query_embedding, k=k
+            )
+        return []
+
+    @abstractmethod
+    def search_by_id(self, id: str) -> VectorStoreDocument:
+        """Search for a document by id."""

graphrag_vectors-3.0.0/graphrag_vectors/vector_store_config.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Parameterization settings for the default configuration."""
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from graphrag_vectors.index_schema import IndexSchema
+from graphrag_vectors.vector_store_type import VectorStoreType
+
+
+class VectorStoreConfig(BaseModel):
+    """The default configuration section for Vector Store."""
+
+    model_config = ConfigDict(extra="allow")
+    """Allow extra fields to support custom vector implementations."""
+
+    type: str = Field(
+        description="The vector store type to use.",
+        default=VectorStoreType.LanceDB,
+    )
+
+    db_uri: str | None = Field(
+        description="The database URI to use (only used by lancedb for built-in stores).",
+        default=None,
+    )
+
+    url: str | None = Field(
+        description="The database URL when type == azure_ai_search or cosmosdb.",
+        default=None,
+    )
+
+    api_key: str | None = Field(
+        description="The database API key when type == azure_ai_search.",
+        default=None,
+    )
+
+    audience: str | None = Field(
+        description="The database audience when type == azure_ai_search.",
+        default=None,
+    )
+
+    connection_string: str | None = Field(
+        description="The connection string when type == cosmosdb.",
+        default=None,
+    )
+
+    database_name: str | None = Field(
+        description="The database name to use when type == cosmosdb.",
+        default=None,
+    )
+
+    index_schema: dict[str, IndexSchema] = {}

graphrag_vectors-3.0.0/graphrag_vectors/vector_store_factory.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Factory functions for creating a vector store."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from graphrag_common.factory import Factory, ServiceScope
+
+from graphrag_vectors.vector_store import VectorStore
+from graphrag_vectors.vector_store_type import VectorStoreType
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from graphrag_vectors.index_schema import IndexSchema
+    from graphrag_vectors.vector_store_config import VectorStoreConfig
+
+
+class VectorStoreFactory(Factory[VectorStore]):
+    """A factory for vector stores.
+
+    Includes a method for users to register a custom vector store implementation.
+
+    Configuration arguments are passed to each vector store implementation as kwargs
+    for individual enforcement of required/optional arguments.
+    """
+
+
+vector_store_factory = VectorStoreFactory()
+
+
+def register_vector_store(
+    vector_store_type: str,
+    vector_store_initializer: Callable[..., VectorStore],
+    scope: ServiceScope = "transient",
+) -> None:
+    """Register a custom vector store implementation.
+
+    Args
+    ----
+        - vector_store_type: str
+            The vector store id to register.
+        - vector_store_initializer: Callable[..., VectorStore]
+            The vector store initializer to register.
+        - scope: ServiceScope
+            The service scope for the vector store (default: "transient").
+    """
+    vector_store_factory.register(vector_store_type, vector_store_initializer, scope)
+
+
+def create_vector_store(
+    config: VectorStoreConfig, index_schema: IndexSchema
+) -> VectorStore:
+    """Create a vector store implementation based on the given type and configuration.
+
+    Args
+    ----
+        - config: VectorStoreConfig
+            The base vector store configuration.
+        - index_schema: IndexSchema
+            The index schema configuration for the vector store instance - i.e., for the specific table we are reading/writing.
+
+    Returns
+    -------
+        VectorStore
+            The created vector store implementation.
+    """
+    strategy = config.type
+
+    # Lazy load built-in implementations
+    if strategy not in vector_store_factory:
+        match strategy:
+            case VectorStoreType.LanceDB:
+                from graphrag_vectors.lancedb import LanceDBVectorStore
+
+                register_vector_store(VectorStoreType.LanceDB, LanceDBVectorStore)
+            case VectorStoreType.AzureAISearch:
+                from graphrag_vectors.azure_ai_search import AzureAISearchVectorStore
+
+                register_vector_store(
+                    VectorStoreType.AzureAISearch, AzureAISearchVectorStore
+                )
+            case VectorStoreType.CosmosDB:
+                from graphrag_vectors.cosmosdb import CosmosDBVectorStore
+
+                register_vector_store(VectorStoreType.CosmosDB, CosmosDBVectorStore)
+            case _:
+                msg = f"Vector store type '{strategy}' is not registered in the VectorStoreFactory. Registered types: {', '.join(vector_store_factory.keys())}."
+                raise ValueError(msg)
+
+    # collapse the base config and specific index config into a single dict for the initializer
+    config_model = config.model_dump()
+    index_model = index_schema.model_dump()
+    return vector_store_factory.create(
+        strategy, init_args={**config_model, **index_model}
+    )

graphrag_vectors-3.0.0/graphrag_vectors/vector_store_type.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Vector store type enum."""
+
+from enum import StrEnum
+
+
+class VectorStoreType(StrEnum):
+    """The supported vector store types."""
+
+    LanceDB = "lancedb"
+    AzureAISearch = "azure_ai_search"
+    CosmosDB = "cosmosdb"

graphrag_vectors-3.0.0/pyproject.toml

@@ -0,0 +1,49 @@
+[project]
+name = "graphrag-vectors"
+version = "3.0.0"
+description = "GraphRAG vector store package."
+authors = [
+    {name = "Alonso Guevara Fernández", email = "alonsog@microsoft.com"},
+    {name = "Andrés Morales Esquivel", email = "andresmor@microsoft.com"},
+    {name = "Chris Trevino", email = "chtrevin@microsoft.com"},
+    {name = "David Tittsworth", email = "datittsw@microsoft.com"},
+    {name = "Dayenne de Souza", email = "ddesouza@microsoft.com"},
+    {name = "Derek Worthen", email = "deworthe@microsoft.com"},
+    {name = "Gaudy Blanco Meneses", email = "gaudyb@microsoft.com"},
+    {name = "Ha Trinh", email = "trinhha@microsoft.com"},
+    {name = "Jonathan Larson", email = "jolarso@microsoft.com"},
+    {name = "Josh Bradley", email = "joshbradley@microsoft.com"},
+    {name = "Kate Lytvynets", email = "kalytv@microsoft.com"},
+    {name = "Kenny Zhang", email = "zhangken@microsoft.com"},
+    {name = "Mónica Carvajal"},
+    {name = "Nathan Evans", email = "naevans@microsoft.com"},
+    {name = "Rodrigo Racanicci", email = "rracanicci@microsoft.com"},
+    {name = "Sarah Smith", email = "smithsarah@microsoft.com"},
+]
+license = {text = "MIT"}
+readme = "README.md"
+requires-python = ">=3.11,<3.14"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+]
+dependencies = [
+    "azure-core~=1.32",
+    "azure-cosmos~=4.9",
+    "azure-identity~=1.19",
+    "azure-search-documents~=11.6",
+    "graphrag-common==3.0.0",
+    "lancedb~=0.24.1",
+    "numpy~=2.1",
+    "pyarrow~=22.0",
+    "pydantic~=2.10",
+]
+
+[project.urls]
+Source = "https://github.com/microsoft/graphrag"
+
+[build-system]
+requires = ["hatchling>=1.27.0,<2.0.0"]
+build-backend = "hatchling.build"