diff --git a/.gitignore b/.gitignore index 9beb3c165..76345f956 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,12 @@ # MemOS home .memos/ +# FuXi CLI local state +.fuxi/ + +# Node dependencies +node_modules/ + # Temporary files tmp/ **/tmp_data/ diff --git a/apps/memos-local-plugin/adapters/deepseek-harness/index.ts b/apps/memos-local-plugin/adapters/deepseek-harness/index.ts index bc63ba5d6..803db522e 100644 --- a/apps/memos-local-plugin/adapters/deepseek-harness/index.ts +++ b/apps/memos-local-plugin/adapters/deepseek-harness/index.ts @@ -413,10 +413,21 @@ export async function apply( })); registrations.push(ctx.on("session/event", (session: Session, event: SessionEvent): void => { - bridge!.onSessionEvent( - session as unknown as DshSessionLike, - event as unknown as DshSessionEventLike, - ); + // Memory is an optional enhancement. A malformed/unexpected event shape + // must not break the host agent's event loop, so swallow with a warning. + try { + bridge!.onSessionEvent( + session as unknown as DshSessionLike, + event as unknown as DshSessionEventLike, + ); + } catch (error) { + const message = + error instanceof Error ? error.message : String(error); + const stack = error instanceof Error && error.stack ? `\n${error.stack}` : ""; + ctx.logger.warn( + `memos-local-memory: session event ignored (${event.type}): ${message}${stack}`, + ); + } })); registrations.push(ctx.on("session/disposed", (session: Session): void => { diff --git a/docker/requirements-full.txt b/docker/requirements-full.txt index b148d43d6..4e6e4affb 100644 --- a/docker/requirements-full.txt +++ b/docker/requirements-full.txt @@ -169,7 +169,6 @@ zstandard==0.23.0 prometheus_client==0.23.1 beartype==0.22.5 diskcache==5.6.3 -iniconfig==2.3.0 jaraco.classes==3.4.0 jaraco.context==6.0.1 jaraco.functools==4.3.0 @@ -178,11 +177,9 @@ more-itertools==10.8.0 pathable==0.4.4 pathvalidate==3.3.1 platformdirs==4.5.0 -pluggy==1.6.0 psycopg2-binary==2.9.11 py-key-value-aio==0.2.8 py-key-value-shared==0.2.8 PyJWT==2.10.1 -pytest==9.0.2 alibabacloud-oss-v2==1.2.2 tavily-python==0.5.0 diff --git a/docker/requirements.txt b/docker/requirements.txt index 988e64b83..ea77de45b 100644 --- a/docker/requirements.txt +++ b/docker/requirements.txt @@ -34,7 +34,6 @@ httpx==0.28.1 httpx-sse==0.4.3 huggingface-hub==0.36.0 idna==3.11 -iniconfig==2.3.0 itsdangerous==2.2.0 jaraco.classes==3.4.0 jaraco.context==6.0.1 @@ -64,7 +63,6 @@ pathable==0.4.4 pathvalidate==3.3.1 pika==1.3.2 platformdirs==4.5.0 -pluggy==1.6.0 portalocker==2.8.0 prometheus_client==0.23.1 protobuf==6.33.1 @@ -81,7 +79,6 @@ PyJWT==2.10.1 pymilvus==2.6.5 PyMySQL==1.1.2 pyperclip==1.11.0 -pytest==9.0.2 python-dateutil==2.9.0.post0 python-dotenv==1.2.1 python-multipart==0.0.20 diff --git a/src/memos/api/config.py b/src/memos/api/config.py index 0247dfa42..e983aa260 100644 --- a/src/memos/api/config.py +++ b/src/memos/api/config.py @@ -240,7 +240,7 @@ def _auth_headers(): return try: data_props = cls.parse_properties(content) - logger.info("nacos config:", data_props) + logger.info("nacos config: %s", data_props) _update_env_from_dict(data_props) logger.info("✅ parse Nacos setting is Properties ") except Exception as e: @@ -365,7 +365,7 @@ def minimax_config() -> dict[str, Any]: @staticmethod def vllm_config() -> dict[str, Any]: - """Get Qwen configuration.""" + """Get vLLM configuration.""" return { "model_name_or_path": os.getenv("MOS_CHAT_MODEL", "Qwen/Qwen3-1.7B"), "temperature": float(os.getenv("MOS_CHAT_TEMPERATURE", "0.8")), @@ -378,7 +378,7 @@ def vllm_config() -> dict[str, Any]: @staticmethod def get_activation_config() -> dict[str, Any]: - """Get Ollama configuration.""" + """Get activation (kv_cache) configuration.""" return { "backend": "kv_cache", "config": { @@ -547,7 +547,7 @@ def get_feedback_llm_config() -> dict[str, Any]: @staticmethod def get_activation_vllm_config() -> dict[str, Any]: - """Get Ollama configuration.""" + """Get activation (vLLM kv_cache) configuration.""" return { "backend": "vllm_kv_cache", "config": { @@ -580,7 +580,7 @@ def get_preference_memory_config() -> dict[str, Any]: @staticmethod def get_reranker_config() -> dict[str, Any]: - """Get embedder configuration.""" + """Get reranker configuration.""" embedder_backend = os.getenv("MOS_RERANKER_BACKEND", "http_bge") if embedder_backend in ["http_bge", "http_bge_strategy"]: @@ -606,7 +606,7 @@ def get_reranker_config() -> dict[str, Any]: @staticmethod def get_feedback_reranker_config() -> dict[str, Any]: - """Get embedder configuration.""" + """Get feedback reranker configuration.""" embedder_backend = os.getenv("MOS_FEEDBACK_RERANKER_BACKEND", "http_bge") if embedder_backend in ["http_bge", "http_bge_strategy"]: @@ -827,7 +827,7 @@ def get_neo4j_config(user_id: str | None = None) -> dict[str, Any]: @staticmethod def get_noshared_neo4j_config(user_id) -> dict[str, Any]: - """Get Neo4j configuration.""" + """Get per-user non-shared Neo4j configuration.""" return { "uri": os.getenv("NEO4J_URI", "bolt://localhost:7687"), "user": os.getenv("NEO4J_USER", "neo4j"), @@ -840,7 +840,7 @@ def get_noshared_neo4j_config(user_id) -> dict[str, Any]: @staticmethod def get_neo4j_shared_config(user_id: str | None = None) -> dict[str, Any]: - """Get Neo4j configuration.""" + """Get shared Neo4j configuration (multi-tenant).""" return { "uri": os.getenv("NEO4J_URI", "bolt://localhost:7687"), "user": os.getenv("NEO4J_USER", "neo4j"), @@ -852,7 +852,9 @@ def get_neo4j_shared_config(user_id: str | None = None) -> dict[str, Any]: "embedding_dimension": int(os.getenv("EMBEDDING_DIMENSION", 3072)), } - def get_milvus_config(): + @staticmethod + def get_milvus_config() -> dict[str, Any]: + """Get Milvus vector database configuration.""" return { "collection_name": [ "explicit_preference", diff --git a/src/memos/api/handlers/chat_handler.py b/src/memos/api/handlers/chat_handler.py index 58a96cd75..b29bd7670 100644 --- a/src/memos/api/handlers/chat_handler.py +++ b/src/memos/api/handlers/chat_handler.py @@ -7,6 +7,7 @@ import asyncio import json +import logging import os import re import time @@ -48,6 +49,51 @@ from memos.types import MessageList +# Fields safe to log for chat requests. Sensitive/bulk content (query, +# history, system_prompt, filter) and credentials (business_key) are excluded +# or masked so request logging never leaks memory content or auth keys. +# frozenset for O(1) membership checks on the hot path. +_CHAT_REQ_LOG_WHITELIST = frozenset(( + "user_id", + "manager_user_id", + "project_id", + "mem_cube_id", + "readable_cube_ids", + "writable_cube_ids", + "session_id", + "mode", + "top_k", + "threshold", + "relativity", + "pref_top_k", + "max_tokens", + "temperature", + "top_p", + "internet_search", + "include_preference", + "add_message_on_answer", +)) + + +def _safe_chat_req_log(chat_req: Any, prefix: str) -> str: + """Build a loggable summary of a chat request without sensitive content.""" + try: + data = chat_req.model_dump() + except AttributeError: + data = getattr(chat_req, "__dict__", {}) + safe = {k: v for k, v in data.items() if k in _CHAT_REQ_LOG_WHITELIST} + if "business_key" in data: + safe["business_key"] = "***" if data.get("business_key") else None + return f"{prefix} Chat Req: {safe}" + + +def _log_chat_req(logger: Any, chat_req: Any, prefix: str) -> None: + """Emit the safe chat-request log line, skipping expensive serialization + entirely when INFO logging is disabled.""" + if logger.isEnabledFor(logging.INFO): + logger.info(_safe_chat_req_log(chat_req, prefix)) + + class ChatHandler(BaseHandler): """ Handler for chat operations. @@ -116,7 +162,7 @@ def handle_chat_complete(self, chat_req: APIChatCompleteRequest) -> dict[str, An Raises: HTTPException: If chat fails """ - self.logger.info(f"[ChatHandler] Chat Req is: {chat_req}") + _log_chat_req(self.logger, chat_req, "[ChatHandler]") try: # Resolve readable cube IDs (for search) readable_cube_ids = chat_req.readable_cube_ids or [chat_req.user_id] @@ -251,7 +297,7 @@ def handle_chat_stream(self, chat_req: ChatRequest) -> StreamingResponse: Raises: HTTPException: If stream initialization fails """ - self.logger.info(f"[ChatHandler] Chat Req is: {chat_req}") + _log_chat_req(self.logger, chat_req, "[ChatHandler]") try: def generate_chat_response() -> Generator[str, None, None]: @@ -436,7 +482,7 @@ def handle_chat_stream_playground(self, chat_req: ChatPlaygroundRequest) -> Stre Raises: HTTPException: If stream initialization fails """ - self.logger.info(f"[ChatHandler] Chat Req is: {chat_req}") + _log_chat_req(self.logger, chat_req, "[ChatHandler]") try: def generate_chat_response() -> Generator[str, None, None]: @@ -780,7 +826,7 @@ def handle_chat_stream_for_business_user( self, chat_req: ChatBusinessRequest ) -> StreamingResponse: """Chat API for business user.""" - self.logger.info(f"[ChatBusinessHandler] Chat Req is: {chat_req}") + _log_chat_req(self.logger, chat_req, "[ChatBusinessHandler]") # Validate business_key permission business_chat_keys = os.environ.get("BUSINESS_CHAT_KEYS", "[]") diff --git a/src/memos/graph_dbs/neo4j.py b/src/memos/graph_dbs/neo4j.py index 56c3e08a0..5a4db9337 100644 --- a/src/memos/graph_dbs/neo4j.py +++ b/src/memos/graph_dbs/neo4j.py @@ -13,6 +13,35 @@ logger = get_logger(__name__) +# Relationship types used across the codebase (see tree_text_memory/organize/*, +# mem_scheduler handlers). Neo4j does not support parameterized relationship +# types in MATCH patterns, so any dynamic `type` value must pass this allowlist +# before being interpolated into Cypher. +ALLOWED_EDGE_TYPES = ( + "FOLLOWS", + "PARENT", + "MERGED_TO", + "RELATE", + "RELATED", + "RELATE_TO", + "INFERS", + "AGGREGATE_TO", + "CAUSE", + "CONDITION", + "CONFLICT", +) + + +def _validate_edge_type(type: str) -> str: + """Validate a relationship type against the allowlist; return it.""" + if type != "ANY" and type not in ALLOWED_EDGE_TYPES: + raise ValueError( + f"Invalid relationship type: {type!r}. " + f"Must be one of {ALLOWED_EDGE_TYPES} or 'ANY'." + ) + return type + + def _compose_node(item: dict[str, Any]) -> tuple[str, str, dict[str, Any]]: node_id = item["id"] memory = item["memory"] @@ -660,14 +689,63 @@ def get_neighbors( ) -> list[str]: """ Get connected node IDs in a specific direction and relationship type. + Args: id: Source node ID. - type: Relationship type. + type: Relationship type to match, or 'ANY' to match all. direction: Edge direction to follow ('out', 'in', or 'both'). + - 'out': nodes connected by edges leaving `id` + - 'in': nodes connected by edges entering `id` + - 'both': nodes connected either way (deduplicated) + user_name (str, optional): User name for filtering in non-multi-db mode + Returns: List of neighboring node IDs. """ - raise NotImplementedError + if direction not in ("in", "out", "both"): + raise ValueError("Invalid direction. Must be 'in', 'out', or 'both'.") + _validate_edge_type(type) + + user_name = user_name if user_name else self.config.user_name + rel_type = "" if type == "ANY" else f":{type}" + + # 'both' uses an undirected pattern; Neo4j matches each edge in both + # orientations, so it traverses roughly 2x the edges of a directed + # query (correctness is preserved by the DISTINCT + b.id <> $id guard). + if direction == "out": + pattern = f"(a:Memory)-[r{rel_type}]->(b:Memory)" + where_clause = "a.id = $id" + elif direction == "in": + pattern = f"(b:Memory)-[r{rel_type}]->(a:Memory)" + where_clause = "a.id = $id" + else: # both + pattern = f"(a:Memory)-[r{rel_type}]-(b:Memory)" + where_clause = "a.id = $id AND b.id <> $id" + + params = {"id": id} + if not self.config.use_multi_db: + if not user_name: + raise ValueError("user_name is required in non-multi-db mode") + where_clause += " AND a.user_name = $user_name AND b.user_name = $user_name" + params["user_name"] = user_name + else: + # Contract: in multi-db mode each database is a single tenant, so + # no user filter is applied. This is a security assumption — if a + # database ever holds more than one user, tenant isolation breaks. + logger.debug( + "get_neighbors: use_multi_db=True, assuming per-tenant databases; " + "no user_name scoping applied" + ) + + query = f""" + MATCH {pattern} + WHERE {where_clause} + RETURN DISTINCT b.id AS neighbor_id + """ + + with self.driver.session(database=self.db_name) as session: + result = session.run(query, params) + return [record["neighbor_id"] for record in result] def get_neighbors_by_tag( self, @@ -748,14 +826,56 @@ def get_path( ) -> list[str]: """ Get the path of nodes from source to target within a limited depth. + + Tenant isolation: in multi-db mode each database is expected to hold a + single tenant's graph, so no user filter is applied. In non-multi-db + mode a `user_name` filter is always applied (and required). + Args: source_id: Starting node ID. target_id: Target node ID. max_depth: Maximum path length to traverse. + user_name (str, optional): User name for filtering in non-multi-db mode Returns: Ordered list of node IDs along the path. """ - raise NotImplementedError + if not isinstance(max_depth, int) or isinstance(max_depth, bool): + raise TypeError(f"max_depth must be an int, got {type(max_depth).__name__!r}") + user_name = user_name if user_name else self.config.user_name + + user_filter = "" + params = {"source_id": source_id, "target_id": target_id} + if not self.config.use_multi_db: + if not user_name: + raise ValueError("user_name is required in non-multi-db mode") + user_filter = ( + "WHERE n.user_name = $user_name AND m.user_name = $user_name " + "AND all(x IN nodes(p) WHERE x.user_name = $user_name)" + ) + params["user_name"] = user_name + else: + # Contract: in multi-db mode each database is a single tenant, so + # no user filter is applied. This is a security assumption — if a + # database ever holds more than one user, paths can cross tenants. + logger.debug( + "get_path: use_multi_db=True, assuming per-tenant databases; " + "no user_name scoping applied" + ) + + # Neo4j does not allow parameters in variable-length hop bounds; the + # literal must be inlined. Cap it to avoid runaway traversal. + hops = max(1, min(max_depth, 10)) + query = f""" + MATCH p = shortestPath((n:Memory {{id: $source_id}})-[*1..{hops}]-(m:Memory {{id: $target_id}})) + {user_filter} + RETURN [x IN nodes(p) | x.id] AS path_ids + LIMIT 1 + """ + + with self.driver.session(database=self.db_name) as session: + result = session.run(query, params) + record = result.single() + return record["path_ids"] if record else [] def get_subgraph( self, @@ -834,7 +954,7 @@ def get_context_chain(self, id: str, type: str = "FOLLOWS") -> list[str]: Returns: List of ordered node IDs in the chain. """ - raise NotImplementedError + return self.get_neighbors(id, type, "out") # Search / recall operations def search_by_embedding( diff --git a/src/memos/graph_dbs/polardb.py b/src/memos/graph_dbs/polardb.py index bf74fbb8b..0dfbf7821 100644 --- a/src/memos/graph_dbs/polardb.py +++ b/src/memos/graph_dbs/polardb.py @@ -1,6 +1,7 @@ import json import os import random +import re import textwrap import threading import time @@ -21,6 +22,53 @@ logger = get_logger(__name__) +# Relationship types used across the codebase (see tree_text_memory/organize/*, +# mem_scheduler handlers). Hardcoded allowlist so dynamic `type` values can never +# be interpolated into Cypher patterns. +ALLOWED_EDGE_TYPES = ( + "FOLLOWS", + "PARENT", + "MERGED_TO", + "RELATE", + "RELATED", + "RELATE_TO", + "INFERS", + "AGGREGATE_TO", + "CAUSE", + "CONDITION", + "CONFLICT", +) + + +def _validate_edge_type(type: str) -> str: + """Validate a relationship type against the allowlist; return it.""" + if type != "ANY" and type not in ALLOWED_EDGE_TYPES: + raise ValueError( + f"Invalid relationship type: {type!r}. " + f"Must be one of {ALLOWED_EDGE_TYPES} or 'ANY'." + ) + return type + + +def _cypher_safe_id(value: Any) -> str: + """Coerce a node id / user_name to a safe Cypher literal. + + Only allow characters that can appear in real ids / user names + (UUID hex, letters, digits, '-', '_', '.'). Anything else — quotes, + `$`, backslashes, spaces, `;`, etc. — is rejected outright, so a value + can never break out of the dollar-quoted Cypher body or its string + literal. Empty values are allowed (they simply never match any node). + + Apache AGE does not treat SQL double-quote-escaping (`''`) as a valid + escape inside a Cypher single-quoted string, so we defend by strict + character allowlisting instead of escaping. + """ + text = str(value or "") + if not re.fullmatch(r"[0-9A-Za-z_.\-]*", text): + raise ValueError(f"Invalid identifier for Cypher embedding: {value!r}") + return text + + def _build_lightweight_return_columns(return_fields: list[str]) -> str: columns = [] for field in return_fields: @@ -744,63 +792,6 @@ def delete_edge(self, source_id: str, target_id: str, type: str) -> None: cursor.execute(query, (source_id, target_id, type)) logger.info(f"Edge deleted: {source_id} -[{type}]-> {target_id}") - @timed - def edge_exists_old( - self, source_id: str, target_id: str, type: str = "ANY", direction: str = "OUTGOING" - ) -> bool: - """ - Check if an edge exists between two nodes. - Args: - source_id: ID of the source node. - target_id: ID of the target node. - type: Relationship type. Use "ANY" to match any relationship type. - direction: Direction of the edge. - Use "OUTGOING" (default), "INCOMING", or "ANY". - Returns: - True if the edge exists, otherwise False. - """ - where_clauses = [] - params = [] - # SELECT * FROM - # cypher('memtensor_memos_graph', $$ - # MATCH(a: Memory - # {id: "13bb9df6-0609-4442-8bed-bba77dadac92"})-[r] - (b:Memory {id: "2dd03a5b-5d5f-49c9-9e0a-9a2a2899b98d"}) - # RETURN - # r - # $$) AS(r - # agtype); - - if direction == "OUTGOING": - where_clauses.append("source_id = %s AND target_id = %s") - params.extend([source_id, target_id]) - elif direction == "INCOMING": - where_clauses.append("source_id = %s AND target_id = %s") - params.extend([target_id, source_id]) - elif direction == "ANY": - where_clauses.append( - "((source_id = %s AND target_id = %s) OR (source_id = %s AND target_id = %s))" - ) - params.extend([source_id, target_id, target_id, source_id]) - else: - raise ValueError( - f"Invalid direction: {direction}. Must be 'OUTGOING', 'INCOMING', or 'ANY'." - ) - - if type != "ANY": - where_clauses.append("edge_type = %s") - params.append(type) - - where_clause = " AND ".join(where_clauses) - - query = f""" - SELECT 1 FROM "{self.db_name}_graph"."Edges" - WHERE {where_clause} - LIMIT 1 - """ - with self._get_connection() as conn, conn.cursor() as cursor: - cursor.execute(query, params) - result = cursor.fetchone() - return result is not None @timed def edge_exists( @@ -999,199 +990,67 @@ def get_nodes(self, ids: list[str], user_name: str, **kwargs) -> list[dict[str, nodes.append(self._parse_node(properties)) return nodes - @timed - def get_edges_old( - self, id: str, type: str = "ANY", direction: str = "ANY" - ) -> list[dict[str, str]]: - """ - Get edges connected to a node, with optional type and direction filter. - - Args: - id: Node ID to retrieve edges for. - type: Relationship type to match, or 'ANY' to match all. - direction: 'OUTGOING', 'INCOMING', or 'ANY'. - - Returns: - List of edges: - [ - {"from": "source_id", "to": "target_id", "type": "RELATE"}, - ... - ] - """ - - # Create a simple edge table to store relationships (if not exists) - try: - with self.connection.cursor() as cursor: - # Create edge table - cursor.execute(f""" - CREATE TABLE IF NOT EXISTS "{self.db_name}_graph"."Edges" ( - id SERIAL PRIMARY KEY, - source_id TEXT NOT NULL, - target_id TEXT NOT NULL, - edge_type TEXT NOT NULL, - properties JSONB, - created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - FOREIGN KEY (source_id) REFERENCES "{self.db_name}_graph"."Memory"(id), - FOREIGN KEY (target_id) REFERENCES "{self.db_name}_graph"."Memory"(id) - ); - """) - - # Create indexes - cursor.execute(f""" - CREATE INDEX IF NOT EXISTS idx_edges_source - ON "{self.db_name}_graph"."Edges" (source_id); - """) - cursor.execute(f""" - CREATE INDEX IF NOT EXISTS idx_edges_target - ON "{self.db_name}_graph"."Edges" (target_id); - """) - cursor.execute(f""" - CREATE INDEX IF NOT EXISTS idx_edges_type - ON "{self.db_name}_graph"."Edges" (edge_type); - """) - except Exception as e: - logger.warning(f"Failed to create edges table: {e}") - - # Query edges - where_clauses = [] - params = [id] - - if type != "ANY": - where_clauses.append("edge_type = %s") - params.append(type) - - if direction == "OUTGOING": - where_clauses.append("source_id = %s") - elif direction == "INCOMING": - where_clauses.append("target_id = %s") - else: # ANY - where_clauses.append("(source_id = %s OR target_id = %s)") - params.append(id) # Add second parameter for ANY direction - - where_clause = " AND ".join(where_clauses) - - query = f""" - SELECT source_id, target_id, edge_type - FROM "{self.db_name}_graph"."Edges" - WHERE {where_clause} - """ - - with self.connection.cursor() as cursor: - cursor.execute(query, params) - results = cursor.fetchall() - - edges = [] - for row in results: - source_id, target_id, edge_type = row - edges.append({"from": source_id, "to": target_id, "type": edge_type}) - return edges def get_neighbors( self, id: str, type: str, direction: Literal["in", "out", "both"] = "out" ) -> list[str]: """Get connected node IDs in a specific direction and relationship type.""" - raise NotImplementedError - - @timed - def get_neighbors_by_tag_old( - self, - tags: list[str], - exclude_ids: list[str], - top_k: int = 5, - min_overlap: int = 1, - ) -> list[dict[str, Any]]: - """ - Find top-K neighbor nodes with maximum tag overlap. - - Args: - tags: The list of tags to match. - exclude_ids: Node IDs to exclude (e.g., local cluster). - top_k: Max number of neighbors to return. - min_overlap: Minimum number of overlapping tags required. - - Returns: - List of dicts with node details and overlap count. - """ - # Build query conditions - where_clauses = [] - params = [] - - # Exclude specified IDs - if exclude_ids: - placeholders = ",".join(["%s"] * len(exclude_ids)) - where_clauses.append(f"id NOT IN ({placeholders})") - params.extend(exclude_ids) - - # Status filter - where_clauses.append("properties->>'status' = %s") - params.append("activated") - - # Type filter - where_clauses.append("properties->>'type' != %s") - params.append("reasoning") - - where_clauses.append("properties->>'memory_type' != %s") - params.append("WorkingMemory") - - # User filter - if not self._get_config_value("use_multi_db", True) and self._get_config_value("user_name"): - where_clauses.append("properties->>'user_name' = %s") - params.append(self._get_config_value("user_name")) - - where_clause = " AND ".join(where_clauses) + if direction not in ("in", "out", "both"): + raise ValueError("Invalid direction. Must be 'in', 'out', or 'both'.") + _validate_edge_type(type) + + # In multi-db mode user_name may be None (single-tenant databases), so + # only apply the tenant filter when one is configured. + user_name = self._get_config_value("user_name") + id_safe = _cypher_safe_id(id) + user_clause = "" + if user_name: + user_safe = _cypher_safe_id(user_name) + user_clause = f" AND a.user_name = '{user_safe}'" + type_filter = f":{type}" if type != "ANY" else "" - # Get all candidate nodes + if direction == "out": + cypher_body = f""" + MATCH (a:Memory)-[r{type_filter}]->(b:Memory) + WHERE a.id = '{id_safe}'{user_clause} + RETURN DISTINCT b.id AS neighbor_id + """ + elif direction == "in": + cypher_body = f""" + MATCH (b:Memory)-[r{type_filter}]->(a:Memory) + WHERE a.id = '{id_safe}'{user_clause} + RETURN DISTINCT b.id AS neighbor_id + """ + else: # both + cypher_body = f""" + MATCH (a:Memory)-[r{type_filter}]-(b:Memory) + WHERE a.id = '{id_safe}'{user_clause} + RETURN DISTINCT b.id AS neighbor_id + """ query = f""" - SELECT id, properties, embedding - FROM "{self.db_name}_graph"."Memory" - WHERE {where_clause} + SELECT * FROM cypher('{self.db_name}_graph', $$ + {cypher_body.strip()} + $$) AS (neighbor_id agtype) """ + try: + with self._get_connection() as conn, conn.cursor() as cursor: + cursor.execute(query) + results = cursor.fetchall() - with self.connection.cursor() as cursor: - cursor.execute(query, params) - results = cursor.fetchall() - - nodes_with_overlap = [] - for row in results: - node_id, properties_json, embedding_json = row - properties = properties_json if properties_json else {} - - # Parse embedding - if embedding_json is not None: - try: - embedding = ( - json.loads(embedding_json) - if isinstance(embedding_json, str) - else embedding_json - ) - properties["embedding"] = embedding - except (json.JSONDecodeError, TypeError): - logger.warning(f"Failed to parse embedding for node {node_id}") - - # Compute tag overlap - node_tags = properties.get("tags", []) - if isinstance(node_tags, str): - try: - node_tags = json.loads(node_tags) - except (json.JSONDecodeError, TypeError): - node_tags = [] - - overlap_tags = [tag for tag in tags if tag in node_tags] - overlap_count = len(overlap_tags) - - if overlap_count >= min_overlap: - node_data = self._parse_node( - { - "id": properties.get("id", node_id), - "memory": properties.get("memory", ""), - "metadata": properties, - } - ) - nodes_with_overlap.append((node_data, overlap_count)) + neighbors = [] + for row in results: + # Guard against NULL results (e.g. optional match / graph inconsistency) + if row[0] is None: + continue + raw = row[0].value if hasattr(row[0], "value") else row[0] + if isinstance(raw, str) and raw.startswith('"') and raw.endswith('"'): + raw = raw[1:-1] + neighbors.append(str(raw)) + return neighbors + except Exception as e: + logger.error(f"Failed to get neighbors: {e}", exc_info=True) + raise - # Sort by overlap count and return top_k - nodes_with_overlap.sort(key=lambda x: x[1], reverse=True) - return [node for node, _ in nodes_with_overlap[:top_k]] @timed def get_children_with_embeddings( @@ -1274,7 +1133,50 @@ def get_children_with_embeddings( def get_path(self, source_id: str, target_id: str, max_depth: int = 3) -> list[str]: """Get the path of nodes from source to target within a limited depth.""" - raise NotImplementedError + if not isinstance(max_depth, int) or isinstance(max_depth, bool): + raise TypeError(f"max_depth must be an int, got {type(max_depth).__name__!r}") + + user_name = self._get_config_value("user_name") + source_safe = _cypher_safe_id(source_id) + target_safe = _cypher_safe_id(target_id) + user_clause = "" + if user_name: + user_safe = _cypher_safe_id(user_name) + user_clause = f"WHERE n.user_name = '{user_safe}' AND m.user_name = '{user_safe}'" + + # Variable-length path [*1..N] counts edges; cap at a sane upper bound + # to avoid unbounded traversal in AGE. + hops = max(1, min(max_depth, 6)) + query = f""" + SELECT * FROM cypher('{self.db_name}_graph', $cypher$ + MATCH p = (n:Memory {{id: '{source_safe}'}})-[*1..{hops}]-(m:Memory {{id: '{target_safe}'}}) + {user_clause} + RETURN [x IN nodes(p) | x.id] AS path_ids + ORDER BY length(p) ASC + LIMIT 1 + $cypher$) AS (path_ids agtype) + """ + try: + with self._get_connection() as conn, conn.cursor() as cursor: + cursor.execute(query) + row = cursor.fetchone() + if row is None: + return [] + raw = row[0].value if hasattr(row[0], "value") else row[0] + if isinstance(raw, list): + result = [] + for x in raw: + if x is None: + continue + val = x.value if hasattr(x, "value") else x + if isinstance(val, str) and val.startswith('"') and val.endswith('"'): + val = val[1:-1] + result.append(str(val)) + return result + return [] + except Exception as e: + logger.error(f"Failed to get path: {e}", exc_info=True) + raise @timed def get_subgraph( @@ -1482,7 +1384,7 @@ def get_subgraph( def get_context_chain(self, id: str, type: str = "FOLLOWS") -> list[str]: """Get the ordered context chain starting from a node.""" - raise NotImplementedError + return self.get_neighbors(id, type, "out") def _extract_fields_from_properties( self, properties: Any, return_fields: list[str] @@ -2109,81 +2011,6 @@ def get_by_metadata( logger.info("get_by_metadata internal took %.1f ms", elapsed) return ids - @timed - def get_grouped_counts1( - self, - group_fields: list[str], - where_clause: str = "", - params: dict[str, Any] | None = None, - user_name: str | None = None, - ) -> list[dict[str, Any]]: - """ - Count nodes grouped by any fields. - - Args: - group_fields (list[str]): Fields to group by, e.g., ["memory_type", "status"] - where_clause (str, optional): Extra WHERE condition. E.g., - "WHERE n.status = 'activated'" - params (dict, optional): Parameters for WHERE clause. - - Returns: - list[dict]: e.g., [{ 'memory_type': 'WorkingMemory', 'status': 'active', 'count': 10 }, ...] - """ - user_name = user_name if user_name else self.config.user_name - if not group_fields: - raise ValueError("group_fields cannot be empty") - - final_params = params.copy() if params else {} - if not self.config.use_multi_db and (self.config.user_name or user_name): - user_clause = "n.user_name = $user_name" - final_params["user_name"] = user_name - if where_clause: - where_clause = where_clause.strip() - if where_clause.upper().startswith("WHERE"): - where_clause += f" AND {user_clause}" - else: - where_clause = f"WHERE {where_clause} AND {user_clause}" - else: - where_clause = f"WHERE {user_clause}" - # Force RETURN field AS field to guarantee key match - group_fields_cypher = ", ".join([f"n.{field} AS {field}" for field in group_fields]) - """ - # group_fields_cypher_polardb = "agtype, ".join([f"{field}" for field in group_fields]) - """ - group_fields_cypher_polardb = ", ".join([f"{field} agtype" for field in group_fields]) - query = f""" - SELECT * FROM cypher('{self.db_name}_graph', $$ - MATCH (n:Memory) - {where_clause} - RETURN {group_fields_cypher}, COUNT(n) AS count1 - $$ ) as ({group_fields_cypher_polardb}, count1 agtype); - """ - try: - with self.connection.cursor() as cursor: - # Handle parameterized query - if params and isinstance(params, list): - cursor.execute(query, final_params) - else: - cursor.execute(query) - results = cursor.fetchall() - - output = [] - for row in results: - group_values = {} - for i, field in enumerate(group_fields): - value = row[i] - if hasattr(value, "value"): - group_values[field] = value.value - else: - group_values[field] = str(value) - count_value = row[-1] # Last column is count - output.append({**group_values, "count": count_value}) - - return output - - except Exception as e: - logger.error(f"Failed to get grouped counts: {e}", exc_info=True) - return [] @timed def get_grouped_counts( @@ -2704,111 +2531,6 @@ def get_all_memory_items( return nodes - def get_all_memory_items_old( - self, scope: str, include_embedding: bool = False, user_name: str | None = None - ) -> list[dict]: - """ - Retrieve all memory items of a specific memory_type. - - Args: - scope (str): Must be one of 'WorkingMemory', 'LongTermMemory', or 'UserMemory'. - include_embedding: with/without embedding - user_name (str, optional): User name for filtering in non-multi-db mode - - Returns: - list[dict]: Full list of memory items under this scope. - """ - user_name = user_name if user_name else self._get_config_value("user_name") - if scope not in {"WorkingMemory", "LongTermMemory", "UserMemory", "OuterMemory"}: - raise ValueError(f"Unsupported memory type scope: {scope}") - - # Use cypher query to retrieve memory items - if include_embedding: - cypher_query = f""" - WITH t as ( - SELECT * FROM cypher('{self.db_name}_graph', $$ - MATCH (n:Memory) - WHERE n.memory_type = '{scope}' AND n.user_name = '{user_name}' - RETURN id(n) as id1,n - LIMIT 100 - $$) AS (id1 agtype,n agtype) - ) - SELECT - m.embedding, - t.n - FROM t, - {self.db_name}_graph."Memory" m - WHERE t.id1 = m.id; - """ - else: - cypher_query = f""" - SELECT * FROM cypher('{self.db_name}_graph', $$ - MATCH (n:Memory) - WHERE n.memory_type = '{scope}' AND n.user_name = '{user_name}' - RETURN properties(n) as props - LIMIT 100 - $$) AS (nprops agtype) - """ - - nodes = [] - try: - with self.connection.cursor() as cursor: - cursor.execute(cypher_query) - results = cursor.fetchall() - - for row in results: - node_agtype = row[0] - - # Handle string-formatted data - if isinstance(node_agtype, str): - try: - # Remove ::vertex suffix - json_str = node_agtype.replace("::vertex", "") - node_data = json.loads(json_str) - - if isinstance(node_data, dict) and "properties" in node_data: - properties = node_data["properties"] - # Build node data - parsed_node_data = { - "id": properties.get("id", ""), - "memory": properties.get("memory", ""), - "metadata": properties, - } - - if include_embedding and "embedding" in properties: - parsed_node_data["embedding"] = properties["embedding"] - - nodes.append(self._parse_node(parsed_node_data)) - logger.debug( - f"[get_all_memory_items] Parsed node successfully: {properties.get('id', '')}" - ) - else: - logger.warning(f"Invalid node data format: {node_data}") - - except (json.JSONDecodeError, TypeError) as e: - logger.error(f"JSON parsing failed: {e}") - elif node_agtype and hasattr(node_agtype, "value"): - # Handle agtype object - node_props = node_agtype.value - if isinstance(node_props, dict): - # Parse node properties - node_data = { - "id": node_props.get("id", ""), - "memory": node_props.get("memory", ""), - "metadata": node_props, - } - - if include_embedding and "embedding" in node_props: - node_data["embedding"] = node_props["embedding"] - - nodes.append(self._parse_node(node_data)) - else: - logger.warning(f"Unknown data format: {type(node_agtype)}") - - except Exception as e: - logger.error(f"Failed to get memories: {e}", exc_info=True) - - return nodes @timed def get_structure_optimization_candidates( @@ -2985,17 +2707,14 @@ def get_structure_optimization_candidates( return candidates def drop_database(self) -> None: - """Permanently delete the entire graph this instance is using.""" - return - if self._get_config_value("use_multi_db", True): - with self.connection.cursor() as cursor: - cursor.execute(f"SELECT drop_graph('{self.db_name}_graph', true)") - logger.info(f"Graph '{self.db_name}_graph' has been dropped.") - else: - raise ValueError( - f"Refusing to drop graph '{self.db_name}_graph' in " - f"Shared Database Multi-Tenant mode" - ) + """ + Permanently delete the entire graph this instance is using. + + Intentionally a no-op: dropping a PolarDB graph is handled by the + operator (e.g. dropping the whole database/schema), not by this + process, to avoid accidental destructive actions on shared deployment. + """ + pass def _parse_node(self, node_data: dict[str, Any]) -> dict[str, Any]: """Parse node data from database format to standard format.""" @@ -3088,8 +2807,11 @@ def _strip_wrapping_quotes(value: Any) -> Any: def __del__(self): """Close database connection when object is destroyed.""" - if hasattr(self, "connection") and self.connection: - self.connection.close() + if hasattr(self, "connection_pool"): + try: + self.connection_pool.closeall() + except Exception as e: + logger.warning(f"Failed to close connection pool in __del__: {e}") @timed def add_node( @@ -3551,164 +3273,6 @@ def get_neighbors_by_tag( logger.error(f"Failed to get neighbors by tag: {e}", exc_info=True) return [] - def get_neighbors_by_tag_ccl( - self, - tags: list[str], - exclude_ids: list[str], - top_k: int = 5, - min_overlap: int = 1, - include_embedding: bool = False, - user_name: str | None = None, - ) -> list[dict[str, Any]]: - """ - Find top-K neighbor nodes with maximum tag overlap. - - Args: - tags: The list of tags to match. - exclude_ids: Node IDs to exclude (e.g., local cluster). - top_k: Max number of neighbors to return. - min_overlap: Minimum number of overlapping tags required. - include_embedding: with/without embedding - user_name (str, optional): User name for filtering in non-multi-db mode - - Returns: - List of dicts with node details and overlap count. - """ - if not tags: - return [] - - user_name = user_name if user_name else self._get_config_value("user_name") - - # Build query conditions shared with other graph backends - where_clauses = [ - 'n.status = "activated"', - 'NOT (n.node_type = "reasoning")', - 'NOT (n.memory_type = "WorkingMemory")', - ] - where_clauses = [ - 'n.status = "activated"', - 'NOT (n.memory_type = "WorkingMemory")', - ] - - if exclude_ids: - exclude_ids_str = "[" + ", ".join(f'"{id}"' for id in exclude_ids) + "]" - where_clauses.append(f"NOT (n.id IN {exclude_ids_str})") - - where_clauses.append(f'n.user_name = "{user_name}"') - - where_clause = " AND ".join(where_clauses) - tag_list_literal = "[" + ", ".join(f'"{t}"' for t in tags) + "]" - - return_fields = [ - "n.id AS id", - "n.memory AS memory", - "n.user_name AS user_name", - "n.user_id AS user_id", - "n.session_id AS session_id", - "n.status AS status", - "n.key AS key", - "n.confidence AS confidence", - "n.tags AS tags", - "n.created_at AS created_at", - "n.updated_at AS updated_at", - "n.memory_type AS memory_type", - "n.sources AS sources", - "n.source AS source", - "n.node_type AS node_type", - "n.visibility AS visibility", - "n.background AS background", - ] - - if include_embedding: - return_fields.append("n.embedding AS embedding") - - return_fields_str = ", ".join(return_fields) - result_fields = [] - for field in return_fields: - # Extract field name 'id' from 'n.id AS id' - field_name = field.split(" AS ")[-1] - result_fields.append(f"{field_name} agtype") - - # Add overlap_count - result_fields.append("overlap_count agtype") - result_fields_str = ", ".join(result_fields) - # Use Cypher query to keep the graph query path aligned - query = f""" - SELECT * FROM ( - SELECT * FROM cypher('{self.db_name}_graph', $$ - WITH {tag_list_literal} AS tag_list - MATCH (n:Memory) - WHERE {where_clause} - RETURN {return_fields_str}, - size([tag IN n.tags WHERE tag IN tag_list]) AS overlap_count - $$) AS ({result_fields_str}) - ) AS subquery - ORDER BY (overlap_count::integer) DESC - LIMIT {top_k} - """ - logger.debug(f"get_neighbors_by_tag: {query}") - try: - with self.connection.cursor() as cursor: - cursor.execute(query) - results = cursor.fetchall() - - neighbors = [] - for row in results: - # Parse results - props = {} - overlap_count = None - - # Manually parse each field - field_names = [ - "id", - "memory", - "user_name", - "user_id", - "session_id", - "status", - "key", - "confidence", - "tags", - "created_at", - "updated_at", - "memory_type", - "sources", - "source", - "node_type", - "visibility", - "background", - ] - - if include_embedding: - field_names.append("embedding") - field_names.append("overlap_count") - - for i, field in enumerate(field_names): - if field == "overlap_count": - overlap_count = row[i].value if hasattr(row[i], "value") else row[i] - else: - props[field] = row[i].value if hasattr(row[i], "value") else row[i] - overlap_int = int(overlap_count) - if overlap_count is not None and overlap_int >= min_overlap: - parsed = self._parse_node(props) - parsed["overlap_count"] = overlap_int - neighbors.append(parsed) - - # Sort by overlap count - neighbors.sort(key=lambda x: x["overlap_count"], reverse=True) - neighbors = neighbors[:top_k] - - # Remove overlap_count field - result = [] - for neighbor in neighbors: - neighbor.pop("overlap_count", None) - result.append(neighbor) - - return result - - except Exception as e: - logger.error(f"Failed to get neighbors by tag: {e}", exc_info=True) - return [] @timed def import_graph(self, data: dict[str, Any], user_name: str | None = None) -> None: