microsoft--graphrag/packages/graphrag-vectors/example_notebooks/custom_vector_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "971701cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) 2026 Microsoft Corporation.\n",
    "# Licensed under the MIT License."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bff6557e",
   "metadata": {},
   "source": [
    "## Custom vector example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d76748c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from graphrag_vectors import (\n",
    "    IndexSchema,\n",
    "    VectorStore,\n",
    "    VectorStoreConfig,\n",
    "    VectorStoreDocument,\n",
    "    create_vector_store,\n",
    "    register_vector_store,\n",
    ")\n",
    "\n",
    "\n",
    "class MyCustomVectorStore(VectorStore):\n",
    "    \"\"\"Custom vector store implementation.\"\"\"\n",
    "\n",
    "    def __init__(self, my_param, **kwargs):\n",
    "        self.my_param = my_param\n",
    "\n",
    "    def connect(self):\n",
    "        \"\"\"Connect to the vector store.\"\"\"\n",
    "\n",
    "    def create_index(self):\n",
    "        \"\"\"Create an index in the vector store.\"\"\"\n",
    "\n",
    "    def load_documents(self, documents, overwrite=False):\n",
    "        \"\"\"Load documents into the vector store.\"\"\"\n",
    "\n",
    "    def search_by_id(self, id) -> VectorStoreDocument:\n",
    "        \"\"\"Search for a document by ID.\"\"\"\n",
    "        msg = \"search_by_id not implemented\"\n",
    "        raise NotImplementedError(msg)\n",
    "\n",
    "    def similarity_search_by_vector(self, query_embedding, k=10, **kwargs):\n",
    "        \"\"\"Search for similar documents by vector.\"\"\"\n",
    "        msg = \"similarity_search_by_vector not implemented\"\n",
    "        raise NotImplementedError(msg)\n",
    "\n",
    "\n",
    "# Register your custom implementation\n",
    "register_vector_store(\"my_custom_store\", MyCustomVectorStore)\n",
    "\n",
    "# Define an index schema\n",
    "schema_config = IndexSchema(\n",
    "    index_name=\"my_index\",\n",
    "    vector_size=1536,\n",
    ")\n",
    "\n",
    "# Use your custom vector store\n",
    "config = VectorStoreConfig(\n",
    "    type=\"my_custom_store\",\n",
    "    my_param=\"something\",  # type: ignore\n",
    ")\n",
    "custom_store = create_vector_store(\n",
    "    config=config,\n",
    "    index_schema=schema_config,\n",
    ")\n",
    "\n",
    "custom_store.connect()\n",
    "custom_store.create_index()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}