59 lines
2.4 KiB
Python
59 lines
2.4 KiB
Python
from typing import Optional
|
|
|
|
from pydantic import BaseModel, Field, field_validator, model_validator
|
|
|
|
from mem0.llms.configs import LlmConfig
|
|
|
|
|
|
class Neo4jConfig(BaseModel):
|
|
url: Optional[str] = Field(None, description="Host address for the graph database")
|
|
username: Optional[str] = Field(None, description="Username for the graph database")
|
|
password: Optional[str] = Field(None, description="Password for the graph database")
|
|
database: Optional[str] = Field(None, description="Database for the graph database")
|
|
|
|
@model_validator(mode="before")
|
|
def check_host_port_or_path(cls, values):
|
|
url, username, password = (
|
|
values.get("url"),
|
|
values.get("username"),
|
|
values.get("password"),
|
|
)
|
|
if not url or not username or not password:
|
|
raise ValueError("Please provide 'url', 'username' and 'password'.")
|
|
return values
|
|
|
|
class MemgraphConfig(BaseModel):
|
|
url: Optional[str] = Field(None, description="Host address for the graph database")
|
|
username: Optional[str] = Field(None, description="Username for the graph database")
|
|
password: Optional[str] = Field(None, description="Password for the graph database")
|
|
|
|
@model_validator(mode="before")
|
|
def check_host_port_or_path(cls, values):
|
|
url, username, password = (
|
|
values.get("url"),
|
|
values.get("username"),
|
|
values.get("password"),
|
|
)
|
|
if not url or not username or not password:
|
|
raise ValueError("Please provide 'url', 'username' and 'password'.")
|
|
return values
|
|
|
|
|
|
class GraphStoreConfig(BaseModel):
|
|
provider: str = Field(description="Provider of the data store (e.g., 'neo4j')", default="neo4j")
|
|
config: Neo4jConfig = Field(description="Configuration for the specific data store", default=None)
|
|
llm: Optional[LlmConfig] = Field(description="LLM configuration for querying the graph store", default=None)
|
|
custom_prompt: Optional[str] = Field(
|
|
description="Custom prompt to fetch entities from the given text", default=None
|
|
)
|
|
|
|
@field_validator("config")
|
|
def validate_config(cls, v, values):
|
|
provider = values.data.get("provider")
|
|
if provider == "neo4j":
|
|
return Neo4jConfig(**v.model_dump())
|
|
elif provider == "memgraph":
|
|
return MemgraphConfig(**v.model_dump())
|
|
else:
|
|
raise ValueError(f"Unsupported graph store provider: {provider}")
|