add initial neo4j implementation

private_gpt/components/graph_store/graph_store_component.py

@@ -0,0 +1,77 @@
+import logging
+import typing
+
+from injector import inject, singleton
+from llama_index.core.graph_stores.types import (
+    GraphStore,
+)
+from llama_index.core.indices.knowledge_graph import (
+    KnowledgeGraphRAGRetriever,
+)
+from llama_index.core.llms.llm import LLM
+from llama_index.core.storage import StorageContext
+
+from private_gpt.settings.settings import Settings
+
+logger = logging.getLogger(__name__)
+
+
+@singleton
+class GraphStoreComponent:
+    settings: Settings
+    graph_store: GraphStore
+
+    @inject
+    def __init__(self, settings: Settings) -> None:
+        self.settings = settings
+
+        # If no graphstore is defined, return, making the graphstore optional
+        if settings.graphstore is None:
+            return
+
+        match settings.graphstore.database:
+            case "neo4j":
+                try:
+                    from llama_index.graph_stores.neo4j import (  # type: ignore
+                        Neo4jGraphStore,
+                    )
+                except ImportError as e:
+                    raise ImportError(
+                        "Neo4j dependencies not found, install with `poetry install --extras graph-stores-neo4j`"
+                    ) from e
+
+                if settings.neo4j is None:
+                    raise ValueError(
+                        "Neo4j settings not found. Please provide settings."
+                    )
+
+                self.graph_store = typing.cast(
+                    GraphStore,
+                    Neo4jGraphStore(
+                        **settings.neo4j.model_dump(exclude_none=True),
+                    ),  # TODO
+                )
+            case _:
+                # Should be unreachable
+                # The settings validator should have caught this
+                raise ValueError(
+                    f"Vectorstore database {settings.vectorstore.database} not supported"
+                )
+
+    def get_knowledge_graph(
+        self,
+        llm: LLM,
+    ) -> KnowledgeGraphRAGRetriever:
+        if self.graph_store is None:
+            raise ValueError("GraphStore not defined in settings")
+
+        storage_context = StorageContext.from_defaults(graph_store=self.graph_store)
+        return KnowledgeGraphRAGRetriever(
+            storage_context=storage_context,
+            llm=llm,
+            verbose=True,
+        )
+
+    def close(self) -> None:
+        if hasattr(self.graph_store.client, "close"):
+            self.graph_store.client.close()

private_gpt/settings/settings.py

@@ -114,6 +114,10 @@ class NodeStoreSettings(BaseModel):
     database: Literal["simple", "postgres"]
 
 
+class GraphStoreSettings(BaseModel):
+    database: Literal["neo4j"]
+
+
 class LlamaCPPSettings(BaseModel):
     llm_hf_repo_id: str
     llm_hf_model_file: str
@@ -376,6 +380,25 @@ class QdrantSettings(BaseModel):
     )
 
 
+class Neo4jSettings(BaseModel):
+    url: str | None = Field(
+        "bolt://localhost:7687",
+        description="URL of the Neo4j database.",
+    )
+    username: str | None = Field(
+        "neo4j",
+        description="Username to connect to the Neo4j database.",
+    )
+    password: str | None = Field(
+        "password",
+        description="Password to connect to the Neo4j database.",
+    )
+    database: str | None = Field(
+        "neo4j",
+        description="Database name to connect to the Neo4j database.",
+    )
+
+
 class Settings(BaseModel):
     server: ServerSettings
     data: DataSettings
@@ -389,10 +412,12 @@ class Settings(BaseModel):
     ollama: OllamaSettings
     azopenai: AzureOpenAISettings
     vectorstore: VectorstoreSettings
+    graphstore: GraphStoreSettings
     nodestore: NodeStoreSettings
     rag: RagSettings
     qdrant: QdrantSettings | None = None
     postgres: PostgresSettings | None = None
+    neo4j: Neo4jSettings | None = None
 
 
 """

pyproject.toml

@@ -57,6 +57,7 @@ vector-stores-qdrant = ["llama-index-vector-stores-qdrant"]
 vector-stores-chroma = ["llama-index-vector-stores-chroma"]
 vector-stores-postgres = ["llama-index-vector-stores-postgres"]
 storage-nodestore-postgres = ["llama-index-storage-docstore-postgres","llama-index-storage-index-store-postgres","psycopg2-binary","asyncpg"]
+graph-stores-neo4j = ["llama-index-graph-stores-neo4j"]
 
 [tool.poetry.group.dev.dependencies]
 black = "^22"

settings-ollama.yaml

@@ -24,5 +24,14 @@ ollama:
 vectorstore:
   database: qdrant
 
+graphstore:
+  database: neo4j
+
 qdrant:
   path: local_data/private_gpt/qdrant
+
+neo4j:
+  url: neo4j://localhost:7687
+  username: neo4j
+  password: password
+  database: neo4j