diff --git a/private_gpt/components/graph_store/__init__.py b/private_gpt/components/graph_store/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/private_gpt/components/graph_store/graph_store_component.py b/private_gpt/components/graph_store/graph_store_component.py new file mode 100644 index 00000000..20543d92 --- /dev/null +++ b/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() diff --git a/private_gpt/settings/settings.py b/private_gpt/settings/settings.py index 5896f00d..4211634c 100644 --- a/private_gpt/settings/settings.py +++ b/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 """ diff --git a/pyproject.toml b/pyproject.toml index d5689998..342193dd 100644 --- a/pyproject.toml +++ b/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" diff --git a/settings-ollama.yaml b/settings-ollama.yaml index d7e1a12c..3037a69b 100644 --- a/settings-ollama.yaml +++ b/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