mirror of
https://github.com/imartinez/privateGPT.git
synced 2025-08-20 00:17:07 +00:00
add initial neo4j implementation
This commit is contained in:
parent
087cb0b7b7
commit
d69299b21e
0
private_gpt/components/graph_store/__init__.py
Normal file
0
private_gpt/components/graph_store/__init__.py
Normal file
77
private_gpt/components/graph_store/graph_store_component.py
Normal file
77
private_gpt/components/graph_store/graph_store_component.py
Normal file
@ -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()
|
@ -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
|
||||
|
||||
|
||||
"""
|
||||
|
@ -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"
|
||||
|
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user