add initial neo4j implementation

This commit is contained in:
Javier Martínez Álvarez 2024-03-26 23:18:01 +01:00
parent 087cb0b7b7
commit d69299b21e
No known key found for this signature in database
5 changed files with 112 additions and 0 deletions

View 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()

View File

@ -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
"""

View File

@ -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"

View File

@ -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