Copyright (c) 2026 Microsoft Corporation.

# Copyright (c) 2026 Microsoft Corporation.
# Licensed under the MIT License.

from graphrag_vectors import (
    IndexSchema,
    VectorStore,
    VectorStoreConfig,
    VectorStoreDocument,
    create_vector_store,
    register_vector_store,
)


class MyCustomVectorStore(VectorStore):
    """Custom vector store implementation."""

    def __init__(self, my_param, **kwargs):
        self.my_param = my_param

    def connect(self):
        """Connect to the vector store."""

    def create_index(self):
        """Create an index in the vector store."""

    def load_documents(self, documents, overwrite=False):
        """Load documents into the vector store."""

    def search_by_id(self, id) -> VectorStoreDocument:
        """Search for a document by ID."""
        msg = "search_by_id not implemented"
        raise NotImplementedError(msg)

    def similarity_search_by_vector(self, query_embedding, k=10, **kwargs):
        """Search for similar documents by vector."""
        msg = "similarity_search_by_vector not implemented"
        raise NotImplementedError(msg)


# Register your custom implementation
register_vector_store("my_custom_store", MyCustomVectorStore)

# Define an index schema
schema_config = IndexSchema(
    index_name="my_index",
    vector_size=1536,
)

# Use your custom vector store
config = VectorStoreConfig(
    type="my_custom_store",
    my_param="something",  # type: ignore
)
custom_store = create_vector_store(
    config=config,
    index_schema=schema_config,
)

custom_store.connect()
custom_store.create_index()