my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

my-project-public — gitlawb

#1	from typing import Optional, Union
#2
#3	from pydantic import BaseModel, Field, field_validator, model_validator
#4
#5	from mem0.llms.configs import LlmConfig
#6
#7
#8	class Neo4jConfig(BaseModel):
#9	url: Optional[str] = Field(None, description="Host address for the graph database")
#10	username: Optional[str] = Field(None, description="Username for the graph database")
#11	password: Optional[str] = Field(None, description="Password for the graph database")
#12	database: Optional[str] = Field(None, description="Database for the graph database")
#13	base_label: Optional[bool] = Field(None, description="Whether to use base node label __Entity__ for all entities")
#14
#15	@model_validator(mode="before")
#16	def check_host_port_or_path(cls, values):
#17	url, username, password = (
#18	values.get("url"),
#19	values.get("username"),
#20	values.get("password"),
#21	)
#22	if not url or not username or not password:
#23	raise ValueError("Please provide 'url', 'username' and 'password'.")
#24	return values
#25
#26
#27	class MemgraphConfig(BaseModel):
#28	url: Optional[str] = Field(None, description="Host address for the graph database")
#29	username: Optional[str] = Field(None, description="Username for the graph database")
#30	password: Optional[str] = Field(None, description="Password for the graph database")
#31
#32	@model_validator(mode="before")
#33	def check_host_port_or_path(cls, values):
#34	url, username, password = (
#35	values.get("url"),
#36	values.get("username"),
#37	values.get("password"),
#38	)
#39	if not url or not username or not password:
#40	raise ValueError("Please provide 'url', 'username' and 'password'.")
#41	return values
#42
#43
#44	class NeptuneConfig(BaseModel):
#45	app_id: Optional[str] = Field("Mem0", description="APP_ID for the connection")
#46	endpoint: Optional[str] = (
#47	Field(
#48	None,
#49	description="Endpoint to connect to a Neptune-DB Cluster as 'neptune-db://<host>' or Neptune Analytics Server as 'neptune-graph://<graphid>'",
#50	),
#51	)
#52	base_label: Optional[bool] = Field(None, description="Whether to use base node label __Entity__ for all entities")
#53	collection_name: Optional[str] = Field(None, description="vector_store collection name to store vectors when using Neptune-DB Clusters")
#54
#55	@model_validator(mode="before")
#56	def check_host_port_or_path(cls, values):
#57	endpoint = values.get("endpoint")
#58	if not endpoint:
#59	raise ValueError("Please provide 'endpoint' with the format as 'neptune-db://<endpoint>' or 'neptune-graph://<graphid>'.")
#60	if endpoint.startswith("neptune-db://"):
#61	# This is a Neptune DB Graph
#62	return values
#63	elif endpoint.startswith("neptune-graph://"):
#64	# This is a Neptune Analytics Graph
#65	graph_identifier = endpoint.replace("neptune-graph://", "")
#66	if not graph_identifier.startswith("g-"):
#67	raise ValueError("Provide a valid 'graph_identifier'.")
#68	values["graph_identifier"] = graph_identifier
#69	return values
#70	else:
#71	raise ValueError(
#72	"You must provide an endpoint to create a NeptuneServer as either neptune-db://<endpoint> or neptune-graph://<graphid>"
#73	)
#74
#75
#76	class KuzuConfig(BaseModel):
#77	db: Optional[str] = Field(":memory:", description="Path to a Kuzu database file")
#78
#79
#80	class GraphStoreConfig(BaseModel):
#81	provider: str = Field(
#82	description="Provider of the data store (e.g., 'neo4j', 'memgraph', 'neptune', 'kuzu')",
#83	default="neo4j",
#84	)
#85	config: Union[Neo4jConfig, MemgraphConfig, NeptuneConfig, KuzuConfig] = Field(
#86	description="Configuration for the specific data store", default=None
#87	)
#88	llm: Optional[LlmConfig] = Field(description="LLM configuration for querying the graph store", default=None)
#89	custom_prompt: Optional[str] = Field(
#90	description="Custom prompt to fetch entities from the given text", default=None
#91	)
#92	threshold: float = Field(
#93	description="Threshold for embedding similarity when matching nodes during graph ingestion. "
#94	"Range: 0.0 to 1.0. Higher values require closer matches. "
#95	"Use lower values (e.g., 0.5-0.7) for distinct entities with similar embeddings. "
#96	"Use higher values (e.g., 0.9+) when you want stricter matching.",
#97	default=0.7,
#98	ge=0.0,
#99	le=1.0,
#100	)
#101
#102	@field_validator("config")
#103	def validate_config(cls, v, values):
#104	provider = values.data.get("provider")
#105	if provider == "neo4j":
#106	return Neo4jConfig(**v.model_dump())
#107	elif provider == "memgraph":
#108	return MemgraphConfig(**v.model_dump())
#109	elif provider == "neptune" or provider == "neptunedb":
#110	return NeptuneConfig(**v.model_dump())
#111	elif provider == "kuzu":
#112	return KuzuConfig(**v.model_dump())
#113	else:
#114	raise ValueError(f"Unsupported graph store provider: {provider}")
#115

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public