Skip to content

Commit 9764d5e

Browse files
CaralHsifridayLtianxing02
authored
feat: add Nebula (#168)
* feat: one-click deployment with docker * feat: one-click deployment with docker * feat: one-click deployment with docker * feat: one-click deployment with docker * feat: one-click deployment with docker * feat: one-click deployment with docker * feat: one-click deployment with docker * feat: one-click deployment with docker * feat: one-click deployment with docker * feat: one-click deployment with docker * feat: docker settings modify * Update tree_config.json * Update simple_openapi_memos.py * feat: docker settings modify * feat: support nebular database * Update Dockerfile * Update docker-compose.yml * Update tree_config.json * Update simple_openapi_memos.py * Rename .env.example to docker/.env.example * feat: update requirement.txt * feat: update requirement.txt * feat: update requirement.txt * feat: support nebular database * feat: support nebular database * feat: support nebular database * feat: support nebular database * feat: support nebular database * feat: support nebular database * feat: support nebular database * feat: support nebular database * feat: support nebular database * feat: fix multi-tanant bug for nebula; add create index; fix some functions in nebula; fix example * fix: search bug; * fix index and multi-db-name * feat: finish get-by-metadata in nebula * feat: finish get_structure_optimization_candidates * fix: return value in nebula * feat: fix all nebula issue * feat: create pool connecting * fix: reorganize bug * add several edge types to nebula * feat: change edge name in nebula * fix: bug for graph not found * fix: fail to create db bug * delete test nebula in temp * fix: no node_not_exist bug --------- Co-authored-by: chunyu li <[email protected]> Co-authored-by: stx <[email protected]> Co-authored-by: Tianxing Shi <[email protected]>
1 parent 1d71884 commit 9764d5e

File tree

9 files changed

+1867
-55
lines changed

9 files changed

+1867
-55
lines changed

‎examples/basic_modules/nebular_example.py‎

Lines changed: 413 additions & 0 deletions
Large diffs are not rendered by default.

‎examples/basic_modules/neo4j_example.py‎

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import os
2+
13
from datetime import datetime
24

35
from memos.configs.embedder import EmbedderConfigFactory
@@ -8,7 +10,15 @@
810

911

1012
embedder_config = EmbedderConfigFactory.model_validate(
11-
{"backend": "ollama", "config": {"model_name_or_path": "nomic-embed-text:latest"}}
13+
{
14+
"backend": "universal_api",
15+
"config": {
16+
"provider": "openai",
17+
"api_key": os.getenv("OPENAI_API_KEY", "sk-xxxxx"),
18+
"model_name_or_path": "text-embedding-3-large",
19+
"base_url": os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1"),
20+
},
21+
}
1222
)
1323
embedder = EmbedderFactory.from_config(embedder_config)
1424

@@ -27,7 +37,7 @@ def example_multi_db(db_name: str = "paper"):
2737
"password": "12345678",
2838
"db_name": db_name,
2939
"auto_create": True,
30-
"embedding_dimension": 768,
40+
"embedding_dimension": 3072,
3141
"use_multi_db": True,
3242
},
3343
)
@@ -268,7 +278,7 @@ def example_shared_db(db_name: str = "shared-traval-group"):
268278
"user_name": user_name,
269279
"use_multi_db": False,
270280
"auto_create": True,
271-
"embedding_dimension": 768,
281+
"embedding_dimension": 3072,
272282
},
273283
)
274284
# Step 2: Instantiate graph store
@@ -331,7 +341,7 @@ def example_shared_db(db_name: str = "shared-traval-group"):
331341
"password": "12345678",
332342
"db_name": db_name,
333343
"user_name": user_list[0],
334-
"embedding_dimension": 768,
344+
"embedding_dimension": 3072,
335345
},
336346
)
337347
graph_alice = GraphStoreFactory.from_config(config_alice)
@@ -362,14 +372,14 @@ def run_user_session(
362372
"user_name": user_name,
363373
"use_multi_db": False,
364374
"auto_create": False, # Neo4j Community does not allow auto DB creation
365-
"embedding_dimension": 768,
375+
"embedding_dimension": 3072,
366376
"vec_config": {
367377
# Pass nested config to initialize external vector DB
368378
# If you use qdrant, please use Server instead of local mode.
369379
"backend": "qdrant",
370380
"config": {
371381
"collection_name": "neo4j_vec_db",
372-
"vector_dimension": 768,
382+
"vector_dimension": 3072,
373383
"distance_metric": "cosine",
374384
"host": "localhost",
375385
"port": 6333,
@@ -388,7 +398,7 @@ def run_user_session(
388398
"user_name": user_name,
389399
"use_multi_db": False,
390400
"auto_create": True,
391-
"embedding_dimension": 768,
401+
"embedding_dimension": 3072,
392402
},
393403
)
394404
graph = GraphStoreFactory.from_config(config)

‎src/memos/configs/graph_db.py‎

Lines changed: 49 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
class BaseGraphDBConfig(BaseConfig):
1010
"""Base class for all graph database configurations."""
1111

12-
uri: str
12+
uri: str | list
1313
user: str
1414
password: str
1515

@@ -103,13 +103,61 @@ def validate_community(self):
103103
return self
104104

105105

106+
class NebulaGraphDBConfig(BaseGraphDBConfig):
107+
"""
108+
NebulaGraph-specific configuration.
109+
110+
Key concepts:
111+
- `space`: Equivalent to a database or namespace. All tag/edge/schema live within a space.
112+
- `user_name`: Used for logical tenant isolation if needed.
113+
- `auto_create`: Whether to automatically create the target space if it does not exist.
114+
115+
Example:
116+
---
117+
hosts = ["127.0.0.1:9669"]
118+
user = "root"
119+
password = "nebula"
120+
space = "shared_graph"
121+
user_name = "alice"
122+
"""
123+
124+
space: str = Field(
125+
..., description="The name of the target NebulaGraph space (like a database)"
126+
)
127+
user_name: str | None = Field(
128+
default=None,
129+
description="Logical user or tenant ID for data isolation (optional, used in metadata tagging)",
130+
)
131+
auto_create: bool = Field(
132+
default=False,
133+
description="Whether to auto-create the space if it does not exist",
134+
)
135+
use_multi_db: bool = Field(
136+
default=True,
137+
description=(
138+
"If True: use Neo4j's multi-database feature for physical isolation; "
139+
"each user typically gets a separate database. "
140+
"If False: use a single shared database with logical isolation by user_name."
141+
),
142+
)
143+
embedding_dimension: int = Field(default=3072, description="Dimension of vector embedding")
144+
145+
@model_validator(mode="after")
146+
def validate_config(self):
147+
"""Validate config."""
148+
if not self.space:
149+
raise ValueError("`space` must be provided")
150+
return self
151+
152+
106153
class GraphDBConfigFactory(BaseModel):
107154
backend: str = Field(..., description="Backend for graph database")
108155
config: dict[str, Any] = Field(..., description="Configuration for the graph database backend")
109156

110157
backend_to_class: ClassVar[dict[str, Any]] = {
111158
"neo4j": Neo4jGraphDBConfig,
112159
"neo4j-community": Neo4jCommunityGraphDBConfig,
160+
"nebular": NebulaGraphDBConfig,
113161
}
114162

115163
@field_validator("backend")

‎src/memos/graph_dbs/base.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,7 @@ def search_by_embedding(self, vector: list[float], top_k: int = 5) -> list[dict]
146146
"""
147147

148148
@abstractmethod
149-
def get_by_metadata(self, filters: dict[str, Any]) -> list[str]:
149+
def get_by_metadata(self, filters: list[dict[str, Any]]) -> list[str]:
150150
"""
151151
Retrieve node IDs that match given metadata filters.
152152
@@ -162,6 +162,14 @@ def get_by_metadata(self, filters: dict[str, Any]) -> list[str]:
162162
- Can be used for faceted recall or prefiltering before embedding rerank.
163163
"""
164164

165+
@abstractmethod
166+
def get_structure_optimization_candidates(self, scope: str) -> list[dict]:
167+
"""
168+
Find nodes that are likely candidates for structure optimization:
169+
- Isolated nodes, nodes with empty background, or nodes with exactly one child.
170+
- Plus: the child of any parent node that has exactly one child.
171+
"""
172+
165173
# Structure Maintenance
166174
@abstractmethod
167175
def deduplicate_nodes(self) -> None:

‎src/memos/graph_dbs/factory.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
from memos.configs.graph_db import GraphDBConfigFactory
44
from memos.graph_dbs.base import BaseGraphDB
5+
from memos.graph_dbs.nebular import NebulaGraphDB
56
from memos.graph_dbs.neo4j import Neo4jGraphDB
67
from memos.graph_dbs.neo4j_community import Neo4jCommunityGraphDB
78

@@ -12,6 +13,7 @@ class GraphStoreFactory(BaseGraphDB):
1213
backend_to_class: ClassVar[dict[str, Any]] = {
1314
"neo4j": Neo4jGraphDB,
1415
"neo4j-community": Neo4jCommunityGraphDB,
16+
"nebular": NebulaGraphDB,
1517
}
1618

1719
@classmethod

0 commit comments

Comments
 (0)