diff --git a/src/memos/graph_dbs/base.py b/src/memos/graph_dbs/base.py index 0bc4a54f8..dab4fe59a 100644 --- a/src/memos/graph_dbs/base.py +++ b/src/memos/graph_dbs/base.py @@ -80,6 +80,27 @@ def delete_edge(self, source_id: str, target_id: str, type: str) -> None: type: Relationship type to remove. """ + @abstractmethod + def get_edges( + self, id: str, type: str = "ANY", direction: str = "ANY", user_name: str | None = None + ) -> 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: 'OUT'/'OUTGOING', 'IN'/'INCOMING', or 'ANY'. + user_name: Optional user/tenant scope filter. + + Returns: + List of edges: + [ + {"from": "source_id", "to": "target_id", "type": "RELATE"}, + ... + ] + """ + @abstractmethod def edge_exists(self, source_id: str, target_id: str, type: str) -> bool: """ diff --git a/src/memos/graph_dbs/postgres.py b/src/memos/graph_dbs/postgres.py index 594f7e695..2f501f3ff 100644 --- a/src/memos/graph_dbs/postgres.py +++ b/src/memos/graph_dbs/postgres.py @@ -184,6 +184,48 @@ def _init_schema(self): # Node Management # ========================================================================= + def node_not_exist(self, scope: str, user_name: str | None = None) -> bool: + """Return True when no memory of the given scope exists for the user. + + Used by the tree text memory reorganizer to skip optimization when a + scope has no nodes yet. + """ + user_name = user_name or self.user_name + conn = self._get_conn() + try: + with conn.cursor() as cur: + cur.execute( + f""" + SELECT 1 FROM {self.schema}.memories + WHERE properties->>'memory_type' = %s + AND user_name = %s + LIMIT 1 + """, + (scope, user_name), + ) + return cur.fetchone() is None + finally: + self._put_conn(conn) + + def get_memory_count(self, memory_type: str, user_name: str | None = None) -> int: + """Count memory nodes of a given type for the user.""" + user_name = user_name or self.user_name + conn = self._get_conn() + try: + with conn.cursor() as cur: + cur.execute( + f""" + SELECT COUNT(*) FROM {self.schema}.memories + WHERE properties->>'memory_type' = %s + AND user_name = %s + """, + (memory_type, user_name), + ) + row = cur.fetchone() + return int(row[0]) if row and row[0] is not None else 0 + finally: + self._put_conn(conn) + def remove_oldest_memory( self, memory_type: str, keep_latest: int, user_name: str | None = None ) -> None: @@ -674,23 +716,133 @@ def delete_edge( finally: self._put_conn(conn) - def edge_exists(self, source_id: str, target_id: str, type: str) -> bool: - """Check if edge exists.""" + def edge_exists( + self, + source_id: str, + target_id: str, + type: str = "ANY", + direction: str = "ANY", + user_name: str | None = None, + ) -> 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: "OUTGOING"/"OUT", "INCOMING"/"IN", or "ANY". + Defaults to "ANY" so existing two-node checks match either way. + user_name: Optional user/tenant scope filter (unused by this backend). + + Returns: + True if the edge exists, otherwise False. + """ + conditions: list[str] = [] + params: list[Any] = [] + + if direction in ("OUTGOING", "OUT"): + conditions = ["source_id = %s", "target_id = %s"] + params = [source_id, target_id] + if type != "ANY": + conditions.append("edge_type = %s") + params.append(type) + elif direction in ("INCOMING", "IN"): + # Incoming edge: source_id is the edge's target, and vice versa. + conditions = ["source_id = %s", "target_id = %s"] + params = [target_id, source_id] + if type != "ANY": + conditions.append("edge_type = %s") + params.append(type) + elif direction == "ANY": + # Match the pair in either orientation. + conditions = [ + "(source_id = %s AND target_id = %s) OR (source_id = %s AND target_id = %s)" + ] + params = [source_id, target_id, target_id, source_id] + if type != "ANY": + conditions[0] = f"({conditions[0]})" + conditions.append("edge_type = %s") + params.append(type) + else: + raise ValueError( + f"Invalid direction: {direction}. Must be 'OUTGOING', 'INCOMING', or 'ANY'." + ) + + where_clause = " AND ".join(conditions) + conn = self._get_conn() try: with conn.cursor() as cur: cur.execute( f""" SELECT 1 FROM {self.schema}.edges - WHERE source_id = %s AND target_id = %s AND edge_type = %s + WHERE {where_clause} LIMIT 1 """, - (source_id, target_id, type), + params, ) return cur.fetchone() is not None finally: self._put_conn(conn) + def get_edges( + self, id: str, type: str = "ANY", direction: str = "ANY", user_name: str | None = None + ) -> 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: 'OUT'/'OUTGOING', 'IN'/'INCOMING', or 'ANY'. + user_name: Optional user/tenant scope filter. + + Returns: + List of edges: + [ + {"from": "source_id", "to": "target_id", "type": "RELATE"}, + ... + ] + """ + if direction in ("OUT", "OUTGOING"): + where_clause = "source_id = %s" + params: list[Any] = [id] + elif direction in ("IN", "INCOMING"): + where_clause = "target_id = %s" + params = [id] + elif direction == "ANY": + where_clause = "(source_id = %s OR target_id = %s)" + params = [id, id] + else: + raise ValueError( + f"Invalid direction: {direction}. Must be 'OUTGOING', 'INCOMING', or 'ANY'." + ) + + if type != "ANY": + where_clause += " AND edge_type = %s" + params.append(type) + + if user_name: + where_clause += ( + f" AND (source_id IN (SELECT id FROM {self.schema}.memories WHERE user_name = %s)" + f" AND target_id IN (SELECT id FROM {self.schema}.memories WHERE user_name = %s))" + ) + params.extend([user_name, user_name]) + + conn = self._get_conn() + try: + with conn.cursor() as cur: + cur.execute( + f""" + SELECT source_id, target_id, edge_type + FROM {self.schema}.edges + WHERE {where_clause} + """, + params, + ) + return [{"from": row[0], "to": row[1], "type": row[2]} for row in cur.fetchall()] + finally: + self._put_conn(conn) + # ========================================================================= # Graph Queries # ========================================================================= @@ -957,10 +1109,10 @@ def get_all_memory_items( self._put_conn(conn) def get_structure_optimization_candidates( - self, scope: str, include_embedding: bool = False + self, scope: str, include_embedding: bool = False, **kwargs ) -> list[dict]: - """Find isolated nodes (no edges).""" - user_name = self.user_name + """Find isolated nodes (no edges) for a given scope and user.""" + user_name = kwargs.get("user_name") or self.user_name conn = self._get_conn() try: with conn.cursor() as cur: @@ -983,6 +1135,117 @@ def get_structure_optimization_candidates( finally: self._put_conn(conn) + def search_by_fulltext( + self, + query_words: list[str], + top_k: int = 10, + scope: str | None = None, + status: str | None = None, + threshold: float | None = None, + search_filter: dict | None = None, + user_name: str | None = None, + filter: dict | None = None, + knowledgebase_ids: list[str] | None = None, + tsquery_config: str = "simple", + **kwargs, + ) -> list[dict]: + """Fulltext search over memory content using PostgreSQL text search. + + Compatible with the graph-store interface used by the tree text memory + keyword/fulltext recall path. Words are safely quoted and OR-joined + into a ``to_tsquery`` expression; the query is scored with ``ts_rank``. + + Args: + query_words: List of query words (already-quoted terms are handled). + top_k: Max number of results. + scope: Filter by ``properties->>'memory_type'``. + status: Filter by ``properties->>'status'``; defaults to 'activated'. + threshold: Minimum score to keep a result. + search_filter: Simple dict of property equality filters. + user_name: User/tenant scope. + filter: Rich filter dict (and/or) built via ``_build_filter_where_clause``. + knowledgebase_ids: Accepted for interface compatibility. + tsquery_config: PostgreSQL text-search configuration name. + Defaults to 'simple', which is always available; callers with the + jiebacfg extension may pass their own config. + + Returns: + list of {"id": ..., "score": ...} dicts ordered by descending score. + """ + user_name = user_name or self.user_name + + # Build the tsquery from quoted, OR-joined terms so that user input can + # never inject operators into to_tsquery(...). + def _quote_term(word: str) -> str: + word = str(word).strip().strip("'\"") + return "'" + word.replace("'", "''") + "'" + + terms = [ + _quote_term(w) + for w in query_words + if w and str(w).strip() and str(w).strip().strip("'\"") + ] + if not terms: + return [] + tsquery_string = " | ".join(terms) + + if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", tsquery_config): + raise ValueError(f"Invalid tsquery_config: {tsquery_config}") + + conditions = ["user_name = %s", "memory IS NOT NULL"] + params: list[Any] = [user_name] + + if scope: + conditions.append("properties->>'memory_type' = %s") + params.append(scope) + + if status: + conditions.append("properties->>'status' = %s") + params.append(status) + else: + conditions.append( + "(properties->>'status' = 'activated' OR properties->>'status' IS NULL)" + ) + + if search_filter: + for key, value in search_filter.items(): + if not self._is_safe_field_name(str(key)): + raise ValueError(f"Invalid search_filter key: {key}") + conditions.append(f"properties->>'{key}' = %s") + params.append(str(value)) + + if filter: + filter_where = self._build_filter_where_clause(filter, params) + if filter_where: + conditions.append(f"({filter_where})") + + where_clause = " AND ".join(conditions) + + conn = self._get_conn() + try: + with conn.cursor() as cur: + cur.execute( + f""" + SELECT id, ts_rank(to_tsvector('{tsquery_config}', memory), + to_tsquery('{tsquery_config}', %s)) AS score + FROM {self.schema}.memories + WHERE {where_clause} + AND to_tsvector('{tsquery_config}', memory) @@ to_tsquery('{tsquery_config}', %s) + ORDER BY score DESC + LIMIT %s + """, + (*params, tsquery_string, tsquery_string, top_k), + ) + + results = [] + for row in cur.fetchall(): + score = float(row[1]) + if threshold is None or score >= threshold: + results.append({"id": row[0], "score": score}) + return results + finally: + self._put_conn(conn) + # ========================================================================= # Maintenance # ========================================================================= diff --git a/tests/graph_dbs/test_postgres_graph_db.py b/tests/graph_dbs/test_postgres_graph_db.py new file mode 100644 index 000000000..9c5d6fd68 --- /dev/null +++ b/tests/graph_dbs/test_postgres_graph_db.py @@ -0,0 +1,361 @@ +"""Unit tests for PostgresGraphDB reorganizer/handler-facing methods. + +These tests exercise the graph-store interface used by the tree text memory +reorganizer and handler without requiring a live PostgreSQL server. The +psycopg2 module is patched with a lightweight fake so the SQL that would run +against PostgreSQL can be asserted directly. +""" + +import sys +import types + +from datetime import datetime +from unittest.mock import MagicMock, patch + +import pytest + + +_PSYCOPG2 = sys.modules.setdefault("psycopg2", types.ModuleType("psycopg2")) +_PSYCOPG2_POOL = sys.modules.setdefault("psycopg2.pool", types.ModuleType("psycopg2.pool")) +_PSYCOPG2.pool = _PSYCOPG2_POOL + +from memos.configs.graph_db import PostgresGraphDBConfig # noqa: E402 (stubs above) +from memos.graph_dbs.postgres import PostgresGraphDB # noqa: E402 (stubs above) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +def _make_fake_pool(): + """Return (fake_pool, conn_mock, cursor_mock) wired like ThreadedConnectionPool.""" + pool = MagicMock() + conn = MagicMock() + conn.closed = 0 + cursor = MagicMock() + cursor.fetchone.return_value = None + cursor.fetchall.return_value = [] + # `_get_conn` health check calls conn.cursor() without a context manager, + # while query methods use `with conn.cursor() as cur`. Make both return the + # same cursor object so the configured results are seen everywhere. + conn.cursor.return_value = cursor + cursor.__enter__.return_value = cursor + cursor.__exit__.return_value = False + pool.getconn.return_value = conn + return pool, conn, cursor + + +def _make_graph_db(): + config = PostgresGraphDBConfig( + host="localhost", + port=5432, + user="test", + password="test", + db_name="test_db", + schema_name="memos", + user_name="alice", + embedding_dimension=3, + maxconn=5, + ) + return PostgresGraphDB(config) + + +@pytest.fixture +def graph_db(): + with patch.object(_PSYCOPG2_POOL, "ThreadedConnectionPool", create=True) as mock_pool_cls: + pool, _conn, cursor = _make_fake_pool() + mock_pool_cls.return_value = pool + yield _make_graph_db(), cursor + + +# --------------------------------------------------------------------------- +# node_not_exist / get_memory_count +# --------------------------------------------------------------------------- + + +def test_node_not_exist_returns_true_when_scope_missing(graph_db): + db, cursor = graph_db + cursor.fetchone.return_value = None + + assert db.node_not_exist("WorkingMemory", user_name="alice") is True + sql = cursor.execute.call_args[0][0] + assert "FROM memos.memories" in sql + assert "memory_type" in sql + + +def test_node_not_exist_returns_false_when_scope_has_nodes(graph_db): + db, cursor = graph_db + cursor.fetchone.return_value = ("some-id",) + + assert db.node_not_exist("WorkingMemory", user_name="alice") is False + + +def test_node_not_exist_falls_back_to_config_user(graph_db): + db, cursor = graph_db + cursor.fetchone.return_value = None + + assert db.node_not_exist("WorkingMemory") is True + params = cursor.execute.call_args[0][1] + assert "alice" in params + + +def test_get_memory_count(graph_db): + db, cursor = graph_db + cursor.fetchone.return_value = (7,) + + assert db.get_memory_count("WorkingMemory", user_name="alice") == 7 + sql = cursor.execute.call_args[0][0] + assert "COUNT" in sql.upper() + params = cursor.execute.call_args[0][1] + assert "WorkingMemory" in params + assert "alice" in params + + +def test_get_memory_count_zero_when_no_rows(graph_db): + db, cursor = graph_db + cursor.fetchone.return_value = None + + assert db.get_memory_count("WorkingMemory", user_name="alice") == 0 + + +# --------------------------------------------------------------------------- +# get_edges +# --------------------------------------------------------------------------- + + +def test_get_edges_any_direction(graph_db): + db, cursor = graph_db + cursor.fetchall.return_value = [("a", "b", "PARENT"), ("c", "a", "MERGED_TO")] + + edges = db.get_edges("a", type="ANY", direction="ANY", user_name="alice") + + assert edges == [ + {"from": "a", "to": "b", "type": "PARENT"}, + {"from": "c", "to": "a", "type": "MERGED_TO"}, + ] + sql = cursor.execute.call_args[0][0] + assert "FROM memos.edges" in sql + assert "source_id = %s OR target_id = %s" in sql + + +def test_get_edges_outgoing(graph_db): + db, cursor = graph_db + cursor.fetchall.return_value = [("a", "b", "PARENT")] + + edges = db.get_edges("a", type="PARENT", direction="OUTGOING", user_name="alice") + + assert edges == [{"from": "a", "to": "b", "type": "PARENT"}] + sql = cursor.execute.call_args[0][0] + assert "source_id = %s" in sql + params = cursor.execute.call_args[0][1] + assert "a" in params + + +def test_get_edges_accepts_out_alias(graph_db): + """The scheduler handler calls get_edges(..., direction='OUT').""" + db, cursor = graph_db + cursor.fetchall.return_value = [("a", "b", "MERGED_TO")] + + edges = db.get_edges("a", type="MERGED_TO", direction="OUT") + + assert edges == [{"from": "a", "to": "b", "type": "MERGED_TO"}] + + +def test_get_edges_filters_by_type(graph_db): + db, cursor = graph_db + cursor.fetchall.return_value = [] + + db.get_edges("a", type="FOLLOWS", direction="ANY", user_name="alice") + + sql = cursor.execute.call_args[0][0] + assert "edge_type" in sql + params = cursor.execute.call_args[0][1] + assert "FOLLOWS" in params + + +def test_get_edges_raises_on_invalid_direction(graph_db): + _db, _cursor = graph_db + with pytest.raises(ValueError): + _db.get_edges("a", type="ANY", direction="SIDEWAYS") + + +# --------------------------------------------------------------------------- +# edge_exists +# --------------------------------------------------------------------------- + + +def test_edge_exists_directed(graph_db): + db, cursor = graph_db + cursor.fetchone.return_value = (1,) + + assert db.edge_exists("a", "b", "PARENT", direction="OUTGOING", user_name="alice") is True + sql = cursor.execute.call_args[0][0] + assert "source_id = %s" in sql + assert "target_id = %s" in sql + assert "edge_type = %s" in sql + + +def test_edge_exists_any_direction(graph_db): + db, cursor = graph_db + cursor.fetchone.return_value = (1,) + + assert db.edge_exists("a", "b", "PARENT", direction="ANY", user_name="alice") is True + sql = cursor.execute.call_args[0][0] + assert "source_id = %s AND target_id = %s" in sql + assert "edge_type = %s" in sql + params = cursor.execute.call_args[0][1] + assert params == ["a", "b", "b", "a", "PARENT"] + + +def test_edge_exists_false_when_missing(graph_db): + db, cursor = graph_db + cursor.fetchone.return_value = None + + assert db.edge_exists("a", "b", "PARENT", direction="ANY", user_name="alice") is False + + +def test_edge_exists_default_any_type(graph_db): + db, cursor = graph_db + cursor.fetchone.return_value = (1,) + + assert db.edge_exists("a", "b") is True + sql = cursor.execute.call_args[0][0] + assert "edge_type" not in sql + + +def test_edge_exists_keeps_old_three_arg_call(graph_db): + """Backwards compatibility: old callers pass type positionally.""" + db, cursor = graph_db + cursor.fetchone.return_value = (1,) + + assert db.edge_exists("a", "b", "PARENT") is True + + +# --------------------------------------------------------------------------- +# get_structure_optimization_candidates +# --------------------------------------------------------------------------- + + +def test_get_structure_optimization_candidates_respects_user_name(graph_db): + db, cursor = graph_db + cursor.fetchall.return_value = [ + ( + "n1", + "mem one", + {"memory_type": "LongTermMemory"}, + datetime(2026, 1, 1), + datetime(2026, 1, 1), + ), + ] + + nodes = db.get_structure_optimization_candidates( + "LongTermMemory", user_name="bob", include_embedding=False + ) + + assert len(nodes) == 1 + assert nodes[0]["id"] == "n1" + assert nodes[0]["memory"] == "mem one" + assert nodes[0]["metadata"]["memory_type"] == "LongTermMemory" + # per-user scope must be honored, not the config default + params = cursor.execute.call_args[0][1] + # params layout: (scope, user_name) — pin both the positive and negative + # expectations so a config-default leak (alice) cannot slip through + assert params[1] == "bob" + assert "alice" not in params + + +# --------------------------------------------------------------------------- +# search_by_fulltext +# --------------------------------------------------------------------------- + + +def test_search_by_fulltext_basic(graph_db): + db, cursor = graph_db + cursor.fetchall.return_value = [("n1", 0.85), ("n2", 0.6)] + + results = db.search_by_fulltext( + query_words=["hello", "world"], + top_k=10, + scope="LongTermMemory", + status="activated", + user_name="alice", + ) + + assert results == [{"id": "n1", "score": 0.85}, {"id": "n2", "score": 0.6}] + sql = cursor.execute.call_args[0][0] + assert "ts_rank" in sql + assert "memory_type" in sql + params = cursor.execute.call_args[0][1] + assert "hello" in " ".join(str(p) for p in params) + + +def test_search_by_fulltext_applies_threshold(graph_db): + db, cursor = graph_db + cursor.fetchall.return_value = [("n1", 0.3)] + + results = db.search_by_fulltext( + query_words=["hello"], top_k=10, user_name="alice", threshold=0.5 + ) + + assert results == [] + + +def test_search_by_fulltext_quotes_words_for_tsquery(graph_db): + db, cursor = graph_db + cursor.fetchall.return_value = [] + + db.search_by_fulltext(query_words=["it's", "foo"], top_k=10, user_name="alice") + + # single quotes inside the term must be escaped to avoid breaking to_tsquery + # execute params layout: (user_name, [scope, status, ...filters], tsquery_string, + # tsquery_string, top_k) — the two tsquery_string slots feed the SELECT ts_rank() + # and the WHERE @@ to_tsquery(); pick the first one via [-3] + tsquery_param = str(cursor.execute.call_args[0][1][-3]) + assert "it''s" in tsquery_param + assert "'foo'" in tsquery_param + assert "|" in tsquery_param + + +def test_search_by_fulltext_no_results(graph_db): + db, cursor = graph_db + cursor.fetchall.return_value = [] + + assert db.search_by_fulltext(query_words=["nothing"], top_k=10, user_name="alice") == [] + + +def test_search_by_fulltext_handles_search_filter(graph_db): + db, cursor = graph_db + cursor.fetchall.return_value = [("n1", 0.9)] + + results = db.search_by_fulltext( + query_words=["hello"], + top_k=10, + user_name="alice", + search_filter={"importance": "2"}, + ) + + assert len(results) == 1 + sql = cursor.execute.call_args[0][0] + assert "importance" in sql + + +def test_search_by_fulltext_accepts_kwargs_from_recall_path(graph_db): + """recall.py passes cube_name / knowledgebase_ids which must not blow up.""" + db, cursor = graph_db + cursor.fetchall.return_value = [] + + results = db.search_by_fulltext( + query_words=["hello"], + top_k=10, + status="activated", + scope="LongTermMemory", + search_filter=None, + filter=None, + user_name="alice", + cube_name="cube-a", + knowledgebase_ids=["kb1"], + tsquery_config="jiebaqry", + ) + + assert results == []