From 73e6dfdcee67173c5bb083b79087950a16e2f002 Mon Sep 17 00:00:00 2001 From: Even Date: Mon, 20 Jul 2026 17:50:49 +0800 Subject: [PATCH 1/2] feat: add OceanBase/seekdb vector and graph storage providers Add optional OceanBase / seekdb backends reusing the existing BaseVecDB and BaseGraphDB contracts, without changing any default behavior. - vec_dbs/oceanbase.py: OceanBaseVecDB on top of pyseekdb's Collection API, serving General Memory; require a positive vector_dimension in config. - graph_dbs/oceanbase.py: OceanBaseGraphDB ported from the postgres backend (nodes + edges + JSON + VECTOR) over the MySQL-compatible protocol, with a thread-safe connection pool, atomic multi-step deletes, and identifier whitelisting (table_prefix / search_filter keys). - Register "oceanbase" / "seekdb" aliases in the vec/graph factories and config factories; add GraphDBError; declare the optional "ob-mem" extra. - Add contract tests for both providers. --- pyproject.toml | 6 + src/memos/api/config.py | 24 + src/memos/api/handlers/config_builders.py | 2 + src/memos/configs/graph_db.py | 62 ++ src/memos/configs/vec_db.py | 20 + src/memos/exceptions.py | 3 + src/memos/graph_dbs/factory.py | 3 + src/memos/graph_dbs/oceanbase.py | 1198 +++++++++++++++++++++ src/memos/vec_dbs/factory.py | 3 + src/memos/vec_dbs/oceanbase.py | 311 ++++++ tests/graph_dbs/test_oceanbase.py | 146 +++ tests/vec_dbs/test_oceanbase.py | 188 ++++ 12 files changed, 1966 insertions(+) create mode 100644 src/memos/graph_dbs/oceanbase.py create mode 100644 src/memos/vec_dbs/oceanbase.py create mode 100644 tests/graph_dbs/test_oceanbase.py create mode 100644 tests/vec_dbs/test_oceanbase.py diff --git a/pyproject.toml b/pyproject.toml index 00d240ba8..0f8811db3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -112,6 +112,11 @@ tavily = [ "tavily-python (>=0.5.0,<1.0.0)", ] +# OceanBase / seekdb unified vector + graph provider +ob-mem = [ + "pyseekdb (>=1.4.0,<2.0.0)", # Unified client for seekdb / OceanBase (Collection + raw SQL) +] + # All optional dependencies # Allow users to install with `pip install MemoryOS[all]` all = [ @@ -140,6 +145,7 @@ all = [ "rake-nltk (>=1.0.6,<1.1.0)", "alibabacloud-oss-v2 (>=1.2.2,<1.2.3)", "tavily-python (>=0.5.0,<1.0.0)", + "pyseekdb (>=1.4.0,<2.0.0)", # Uncategorized dependencies ] diff --git a/src/memos/api/config.py b/src/memos/api/config.py index 4d2029969..1b517f875 100644 --- a/src/memos/api/config.py +++ b/src/memos/api/config.py @@ -896,6 +896,30 @@ def get_postgres_config(user_id: str | None = None) -> dict[str, Any]: "maxconn": int(os.getenv("POSTGRES_MAX_CONN", "20")), } + @staticmethod + def get_oceanbase_config(user_id: str | None = None) -> dict[str, Any]: + """Get OceanBase / seekdb configuration for MemOS graph storage. + + Single MySQL-compatible database with logical tenant isolation by + ``user_name``; nodes/edges live in ``{table_prefix}_nodes`` / ``_edges``. + """ + user_name = os.getenv("MEMOS_USER_NAME", "default") + if user_id: + user_name = f"memos_{user_id.replace('-', '')}" + + return { + "host": os.getenv("OCEANBASE_HOST", "localhost"), + "port": int(os.getenv("OCEANBASE_PORT", "2881")), + "user": os.getenv("OCEANBASE_USER", "root"), + "password": os.getenv("OCEANBASE_PASSWORD", ""), + "db_name": os.getenv("OCEANBASE_DB", "memos"), + "table_prefix": os.getenv("OCEANBASE_TABLE_PREFIX", "memos_graph"), + "user_name": user_name, + "use_multi_db": False, + "embedding_dimension": int(os.getenv("EMBEDDING_DIMENSION", "768")), + "maxconn": int(os.getenv("OCEANBASE_MAX_CONN", "20")), + } + @staticmethod def get_mysql_config() -> dict[str, Any]: """Get MySQL configuration.""" diff --git a/src/memos/api/handlers/config_builders.py b/src/memos/api/handlers/config_builders.py index 0a083e284..eb7bb615f 100644 --- a/src/memos/api/handlers/config_builders.py +++ b/src/memos/api/handlers/config_builders.py @@ -41,6 +41,8 @@ def build_graph_db_config(user_id: str = "default") -> dict[str, Any]: "neo4j": APIConfig.get_neo4j_config(user_id=user_id), "polardb": APIConfig.get_polardb_config(user_id=user_id), "postgres": APIConfig.get_postgres_config(user_id=user_id), + "oceanbase": APIConfig.get_oceanbase_config(user_id=user_id), + "seekdb": APIConfig.get_oceanbase_config(user_id=user_id), } # Support both GRAPH_DB_BACKEND and legacy NEO4J_BACKEND env vars diff --git a/src/memos/configs/graph_db.py b/src/memos/configs/graph_db.py index 98de09812..d41d7b72d 100644 --- a/src/memos/configs/graph_db.py +++ b/src/memos/configs/graph_db.py @@ -1,3 +1,5 @@ +import re + from typing import Any, ClassVar from pydantic import BaseModel, Field, field_validator, model_validator @@ -241,6 +243,64 @@ def validate_config(self): return self +class OceanBaseGraphDBConfig(BaseConfig): + """ + OceanBase / seekdb graph configuration for MemOS (via pyseekdb). + + Stores graph nodes, edges, JSON properties and node embeddings in a single + MySQL-compatible database. Logical tenant isolation is done via ``user_name``. + + Example: + --- + host = "127.0.0.1" + port = 2881 + user = "root" + password = "secret" + db_name = "memos" + user_name = "alice" + use_multi_db = False + """ + + host: str = Field(..., description="Database host") + port: int = Field(default=2881, description="Database port") + user: str = Field(default="root", description="Database user") + password: str = Field(default="", description="Database password") + db_name: str = Field(..., description="Database name") + table_prefix: str = Field(default="memos_graph", description="Prefix for the node/edge tables") + user_name: str | None = Field( + default=None, + description="Logical user/tenant ID for data isolation", + ) + use_multi_db: bool = Field( + default=False, + description="If False: use single database with logical isolation by user_name", + ) + embedding_dimension: int = Field(default=768, description="Dimension of vector embedding") + maxconn: int = Field( + default=20, + description="Maximum number of connections in the connection pool", + ) + + @field_validator("table_prefix") + @classmethod + def validate_table_prefix(cls, value: str) -> str: + """Reject prefixes that could break out of the identifier when built into DDL.""" + if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", value): + raise ValueError( + "`table_prefix` must match [A-Za-z_][A-Za-z0-9_]* (letters, digits, underscore)" + ) + return value + + @model_validator(mode="after") + def validate_config(self): + """Validate config.""" + if not self.db_name: + raise ValueError("`db_name` must be provided") + if not self.use_multi_db and not self.user_name: + raise ValueError("In single-database mode, `user_name` must be provided") + return self + + class GraphDBConfigFactory(BaseModel): backend: str = Field(..., description="Backend for graph database") config: dict[str, Any] = Field(..., description="Configuration for the graph database backend") @@ -250,6 +310,8 @@ class GraphDBConfigFactory(BaseModel): "neo4j-community": Neo4jCommunityGraphDBConfig, "polardb": PolarDBGraphDBConfig, "postgres": PostgresGraphDBConfig, + "oceanbase": OceanBaseGraphDBConfig, + "seekdb": OceanBaseGraphDBConfig, } @field_validator("backend") diff --git a/src/memos/configs/vec_db.py b/src/memos/configs/vec_db.py index 9fdb83a35..c159ffa98 100644 --- a/src/memos/configs/vec_db.py +++ b/src/memos/configs/vec_db.py @@ -54,6 +54,24 @@ class MilvusVecDBConfig(BaseVecDBConfig): password: str = Field(default="", description="Password for Milvus connection") +class OceanBaseVecDBConfig(BaseVecDBConfig): + """Configuration for OceanBase / seekdb vector database (via pyseekdb).""" + + host: str = Field(..., description="Host for the seekdb / OceanBase server") + port: int = Field(default=2881, description="Port for the seekdb / OceanBase server") + user: str = Field(default="root", description="Username for the connection") + password: str = Field(default="", description="Password for the connection") + database: str = Field(default="memos", description="Database name") + + @field_validator("vector_dimension") + @classmethod + def validate_vector_dimension(cls, value: int | None) -> int: + """OceanBase requires a concrete vector dimension to build the HNSW index.""" + if value is None or isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError("`vector_dimension` must be a positive integer for OceanBase") + return value + + class VectorDBConfigFactory(BaseConfig): """Factory class for creating vector database configurations.""" @@ -63,6 +81,8 @@ class VectorDBConfigFactory(BaseConfig): backend_to_class: ClassVar[dict[str, Any]] = { "qdrant": QdrantVecDBConfig, "milvus": MilvusVecDBConfig, + "oceanbase": OceanBaseVecDBConfig, + "seekdb": OceanBaseVecDBConfig, } @field_validator("backend") diff --git a/src/memos/exceptions.py b/src/memos/exceptions.py index 28b83c3df..a3a7698a5 100644 --- a/src/memos/exceptions.py +++ b/src/memos/exceptions.py @@ -21,6 +21,9 @@ class MemCubeError(MemOSError): ... class VectorDBError(MemOSError): ... +class GraphDBError(MemOSError): ... + + class LLMError(MemOSError): ... diff --git a/src/memos/graph_dbs/factory.py b/src/memos/graph_dbs/factory.py index 93b5971ec..c86316813 100644 --- a/src/memos/graph_dbs/factory.py +++ b/src/memos/graph_dbs/factory.py @@ -4,6 +4,7 @@ from memos.graph_dbs.base import BaseGraphDB from memos.graph_dbs.neo4j import Neo4jGraphDB from memos.graph_dbs.neo4j_community import Neo4jCommunityGraphDB +from memos.graph_dbs.oceanbase import OceanBaseGraphDB from memos.graph_dbs.polardb import PolarDBGraphDB from memos.graph_dbs.postgres import PostgresGraphDB @@ -16,6 +17,8 @@ class GraphStoreFactory(BaseGraphDB): "neo4j-community": Neo4jCommunityGraphDB, "polardb": PolarDBGraphDB, "postgres": PostgresGraphDB, + "oceanbase": OceanBaseGraphDB, + "seekdb": OceanBaseGraphDB, } @classmethod diff --git a/src/memos/graph_dbs/oceanbase.py b/src/memos/graph_dbs/oceanbase.py new file mode 100644 index 000000000..4eef0a562 --- /dev/null +++ b/src/memos/graph_dbs/oceanbase.py @@ -0,0 +1,1198 @@ +"""OceanBase / seekdb graph backend for MemOS. + +Ported from the PostgreSQL + pgvector backend (``graph_dbs/postgres.py``): graph +data is stored as relational tables (nodes + edges) with JSON properties and a +``VECTOR`` embedding column, all in a single MySQL-compatible database. + +Access goes through the ``pymysql`` driver over seekdb / OceanBase's +MySQL-compatible protocol. All parameterized DML/DQL uses +``cursor.execute(sql, params)`` (MySQL ``%s`` paramstyle) so user input is bound, +never string-concatenated. + +Tables (``{prefix}`` defaults to ``memos_graph``): +- ``{prefix}_nodes``: memory nodes with JSON properties and vector embeddings +- ``{prefix}_edges``: relationships between memory nodes +""" + +import json +import queue +import re +import threading + +from contextlib import contextmanager, suppress +from datetime import datetime +from typing import Any, Literal + +from memos.configs.graph_db import OceanBaseGraphDBConfig +from memos.dependency import require_python_package +from memos.exceptions import GraphDBError +from memos.graph_dbs.base import BaseGraphDB +from memos.log import get_logger + + +logger = get_logger(__name__) + +_SAFE_FIELD_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +class _PyMySQLConnectionPool: + """Minimal thread-safe PyMySQL connection pool. + + PyMySQL's DB-API ``threadsafety`` is 1: a connection must not be shared + across threads. Each operation therefore borrows a dedicated connection and + returns it afterwards. Connections are created lazily up to ``maxconn``; + ``acquire()`` blocks once the pool is exhausted until a connection is freed. + """ + + def __init__(self, connect_fn, maxconn: int): + self._connect_fn = connect_fn + self._maxconn = max(1, int(maxconn)) + self._idle: queue.LifoQueue = queue.LifoQueue(maxsize=self._maxconn) + self._lock = threading.Lock() + self._created = 0 + self._closed = False + + def acquire(self): + if self._closed: + raise GraphDBError("Connection pool is closed") + with self._lock: + if self._created < self._maxconn and self._idle.empty(): + self._created += 1 + try: + return self._connect_fn() + except Exception: + self._created -= 1 + raise + return self._idle.get() + + def release(self, conn) -> None: + if conn is None: + return + if self._closed: + with suppress(Exception): + conn.close() + return + try: + self._idle.put_nowait(conn) + except queue.Full: + with suppress(Exception): + conn.close() + + def close(self) -> None: + self._closed = True + while True: + try: + conn = self._idle.get_nowait() + except queue.Empty: + break + with suppress(Exception): + conn.close() + + +def _prepare_node_metadata(metadata: dict[str, Any]) -> dict[str, Any]: + """Ensure metadata has proper datetime fields and normalized types.""" + now = datetime.utcnow().isoformat() + metadata.setdefault("created_at", now) + metadata.setdefault("updated_at", now) + + embedding = metadata.get("embedding") + if embedding and isinstance(embedding, list): + metadata["embedding"] = [float(x) for x in embedding] + + return metadata + + +def _normalize_dt(value: Any) -> str: + """Normalize an ISO datetime string to a MySQL-compatible ``DATETIME`` literal.""" + if value is None: + value = datetime.utcnow().isoformat() + text = str(value).replace("T", " ") + # Drop timezone suffix / fractional beyond seconds tolerated by DATETIME. + if "+" in text: + text = text.split("+", 1)[0] + return text.strip() + + +class OceanBaseGraphDB(BaseGraphDB): + """OceanBase / seekdb implementation of a graph memory store (via pymysql). + + seekdb / OceanBase expose a MySQL-compatible protocol; pyseekdb's Collection + client does not run arbitrary SQL, so relational graph operations use the + ``pymysql`` driver directly (already a core MemOS dependency). + """ + + @require_python_package( + import_name="pymysql", + install_command="pip install pymysql", + install_link="https://pypi.org/project/pymysql/", + ) + def __init__(self, config: OceanBaseGraphDBConfig): + """Open the MySQL-protocol connection and ensure the schema exists.""" + import pymysql + + self._pymysql = pymysql + self.config = config + self.user_name = config.user_name + self.dim = config.embedding_dimension + self.nodes_table = f"{config.table_prefix}_nodes" + self.edges_table = f"{config.table_prefix}_edges" + self._closed = False + self._conn_kwargs = { + "host": config.host, + "port": config.port, + "user": config.user, + "password": config.password, + "database": config.db_name, + "autocommit": True, + "charset": "utf8mb4", + } + self._pool = _PyMySQLConnectionPool( + connect_fn=lambda: self._pymysql.connect(**self._conn_kwargs), + maxconn=config.maxconn, + ) + + logger.info( + "Connecting to OceanBase/seekdb: %s:%s/%s", + config.host, + config.port, + config.db_name, + ) + self._init_schema() + + # ========================================================================= + # Connection / execution helpers + # ========================================================================= + + def _ensure_open(self) -> None: + if self._closed: + raise GraphDBError("OceanBaseGraphDB connection is closed") + + def _healthy(self, conn): + """Ping the borrowed connection, replacing it if the socket is dead.""" + try: + conn.ping(reconnect=True) + return conn + except Exception: + with suppress(Exception): + conn.close() + return self._pymysql.connect(**self._conn_kwargs) + + @contextmanager + def _borrow(self): + """Borrow a single connection for one auto-committed operation.""" + self._ensure_open() + conn = self._healthy(self._pool.acquire()) + try: + yield conn + finally: + self._pool.release(conn) + + @contextmanager + def _transaction(self): + """Borrow a connection and run several statements as one atomic unit.""" + self._ensure_open() + conn = self._healthy(self._pool.acquire()) + cur = None + try: + cur = conn.cursor() + conn.begin() + try: + yield cur + conn.commit() + except Exception: + with suppress(Exception): + conn.rollback() + raise + finally: + if cur is not None: + with suppress(Exception): + cur.close() + self._pool.release(conn) + + def _query(self, sql: str, params: tuple | list | None = None) -> list[tuple]: + with self._borrow() as conn: + cur = conn.cursor() + try: + cur.execute(sql, tuple(params) if params else ()) + return list(cur.fetchall()) + finally: + with suppress(Exception): + cur.close() + + def _query_one(self, sql: str, params: tuple | list | None = None) -> tuple | None: + with self._borrow() as conn: + cur = conn.cursor() + try: + cur.execute(sql, tuple(params) if params else ()) + return cur.fetchone() + finally: + with suppress(Exception): + cur.close() + + def _execute(self, sql: str, params: tuple | list | None = None) -> int: + with self._borrow() as conn: + cur = conn.cursor() + try: + cur.execute(sql, tuple(params) if params else ()) + rowcount = cur.rowcount if cur.rowcount is not None else 0 + with suppress(Exception): + conn.commit() + return rowcount + finally: + with suppress(Exception): + cur.close() + + @staticmethod + def _vec_literal(vector: list[float]) -> str: + """Render a vector as an OceanBase vector literal string.""" + return "[" + ",".join(str(float(x)) for x in vector) + "]" + + @staticmethod + def _placeholders(values: list[Any]) -> str: + return ", ".join(["%s"] * len(values)) + + def _init_schema(self) -> None: + """Create node/edge tables and indexes if they don't exist.""" + self._execute( + f""" + CREATE TABLE IF NOT EXISTS {self.nodes_table} ( + id VARCHAR(255) PRIMARY KEY, + memory LONGTEXT, + properties JSON, + embedding VECTOR({self.dim}), + user_name VARCHAR(255), + created_at DATETIME DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME DEFAULT CURRENT_TIMESTAMP + ) + """ + ) + self._execute( + f""" + CREATE TABLE IF NOT EXISTS {self.edges_table} ( + id BIGINT AUTO_INCREMENT PRIMARY KEY, + source_id VARCHAR(255) NOT NULL, + target_id VARCHAR(255) NOT NULL, + edge_type VARCHAR(255) NOT NULL, + created_at DATETIME DEFAULT CURRENT_TIMESTAMP, + UNIQUE KEY uq_edge (source_id, target_id, edge_type) + ) + """ + ) + for ddl in ( + f"CREATE INDEX idx_{self.config.table_prefix}_nodes_user " + f"ON {self.nodes_table}(user_name)", + f"CREATE INDEX idx_{self.config.table_prefix}_edges_source " + f"ON {self.edges_table}(source_id)", + f"CREATE INDEX idx_{self.config.table_prefix}_edges_target " + f"ON {self.edges_table}(target_id)", + f"CREATE VECTOR INDEX idx_{self.config.table_prefix}_nodes_embedding " + f"ON {self.nodes_table}(embedding) WITH (distance=cosine, type=hnsw)", + ): + with suppress(Exception): + self._execute(ddl) + logger.info( + "OceanBase graph schema initialized (%s, %s)", self.nodes_table, self.edges_table + ) + + # ========================================================================= + # Node management + # ========================================================================= + + def remove_oldest_memory( + self, memory_type: str, keep_latest: int, user_name: str | None = None + ) -> None: + """Remove all memories of a type except the latest ``keep_latest`` entries.""" + user_name = user_name or self.user_name + keep_latest = int(keep_latest) + + with self._transaction() as cur: + cur.execute( + f""" + SELECT id FROM {self.nodes_table} + WHERE user_name = %s + AND JSON_UNQUOTE(JSON_EXTRACT(properties, '$.memory_type')) = %s + ORDER BY updated_at DESC + LIMIT %s, 18446744073709551615 + """, + (user_name, memory_type, keep_latest), + ) + ids_to_delete = [row[0] for row in cur.fetchall()] + if not ids_to_delete: + return + + ph = self._placeholders(ids_to_delete) + cur.execute( + f"DELETE FROM {self.edges_table} WHERE source_id IN ({ph}) OR target_id IN ({ph})", + (*ids_to_delete, *ids_to_delete), + ) + cur.execute( + f"DELETE FROM {self.nodes_table} WHERE id IN ({ph})", + tuple(ids_to_delete), + ) + logger.info( + "Removed %s oldest %s memories for user %s", + len(ids_to_delete), + memory_type, + user_name, + ) + + def add_node( + self, id: str, memory: str, metadata: dict[str, Any], user_name: str | None = None + ) -> None: + """Add (or upsert) a memory node.""" + user_name = user_name or self.user_name + metadata = _prepare_node_metadata(metadata.copy()) + + embedding = metadata.pop("embedding", None) + created_at = _normalize_dt(metadata.pop("created_at", None)) + updated_at = _normalize_dt(metadata.pop("updated_at", None)) + + if metadata.get("sources"): + metadata["sources"] = [ + json.dumps(s) if not isinstance(s, str) else s for s in metadata["sources"] + ] + + if embedding: + self._execute( + f""" + INSERT INTO {self.nodes_table} + (id, memory, properties, embedding, user_name, created_at, updated_at) + VALUES (%s, %s, %s, %s, %s, %s, %s) + ON DUPLICATE KEY UPDATE + memory = VALUES(memory), + properties = VALUES(properties), + embedding = VALUES(embedding), + updated_at = VALUES(updated_at) + """, + ( + id, + memory, + json.dumps(metadata), + self._vec_literal(embedding), + user_name, + created_at, + updated_at, + ), + ) + else: + self._execute( + f""" + INSERT INTO {self.nodes_table} + (id, memory, properties, user_name, created_at, updated_at) + VALUES (%s, %s, %s, %s, %s, %s) + ON DUPLICATE KEY UPDATE + memory = VALUES(memory), + properties = VALUES(properties), + updated_at = VALUES(updated_at) + """, + (id, memory, json.dumps(metadata), user_name, created_at, updated_at), + ) + + def add_nodes_batch(self, nodes: list[dict[str, Any]], user_name: str | None = None) -> None: + """Batch add memory nodes.""" + for node in nodes: + self.add_node( + id=node["id"], + memory=node["memory"], + metadata=node.get("metadata", {}), + user_name=user_name, + ) + + def update_node(self, id: str, fields: dict[str, Any], user_name: str | None = None) -> None: + """Update node fields (merging into JSON properties).""" + user_name = user_name or self.user_name + if not fields: + return + + current = self.get_node(id, user_name=user_name) + if not current: + return + + props = current.get("metadata", {}).copy() + fields = fields.copy() + embedding = fields.pop("embedding", None) + memory = fields.pop("memory", current.get("memory", "")) + props.update(fields) + props.pop("embedding", None) + props["updated_at"] = datetime.utcnow().isoformat() + props.pop("created_at", None) + + if embedding: + self._execute( + f""" + UPDATE {self.nodes_table} + SET memory = %s, properties = %s, embedding = %s, updated_at = NOW() + WHERE id = %s AND user_name = %s + """, + (memory, json.dumps(props), self._vec_literal(embedding), id, user_name), + ) + else: + self._execute( + f""" + UPDATE {self.nodes_table} + SET memory = %s, properties = %s, updated_at = NOW() + WHERE id = %s AND user_name = %s + """, + (memory, json.dumps(props), id, user_name), + ) + + def delete_node(self, id: str, user_name: str | None = None) -> None: + """Delete a node and its incident edges (atomically).""" + user_name = user_name or self.user_name + with self._transaction() as cur: + cur.execute( + f"DELETE FROM {self.edges_table} WHERE source_id = %s OR target_id = %s", + (id, id), + ) + cur.execute( + f"DELETE FROM {self.nodes_table} WHERE id = %s AND user_name = %s", + (id, user_name), + ) + + def get_node(self, id: str, include_embedding: bool = False, **kwargs) -> dict[str, Any] | None: + """Get a single node by ID.""" + user_name = kwargs.get("user_name") or self.user_name + cols = "id, memory, properties, created_at, updated_at" + if include_embedding: + cols += ", embedding" + row = self._query_one( + f"SELECT {cols} FROM {self.nodes_table} WHERE id = %s AND user_name = %s", + (id, user_name), + ) + if not row: + return None + return self._parse_row(row, include_embedding) + + def get_nodes( + self, ids: list, include_embedding: bool = False, **kwargs + ) -> list[dict[str, Any]]: + """Get multiple nodes by IDs.""" + if not ids: + return [] + user_name = kwargs.get("user_name") or self.user_name + cols = "id, memory, properties, created_at, updated_at" + if include_embedding: + cols += ", embedding" + ph = self._placeholders(ids) + rows = self._query( + f"SELECT {cols} FROM {self.nodes_table} WHERE id IN ({ph}) AND user_name = %s", + (*ids, user_name), + ) + return [self._parse_row(row, include_embedding) for row in rows] + + def _parse_row(self, row, include_embedding: bool = False) -> dict[str, Any]: + """Parse a database row into a node dict.""" + raw_props = row[2] + props = raw_props if isinstance(raw_props, dict) else json.loads(raw_props or "{}") + props["created_at"] = self._iso(row[3]) + props["updated_at"] = self._iso(row[4]) + result = { + "id": row[0], + "memory": row[1] or "", + "metadata": props, + } + if include_embedding and len(row) > 5: + result["metadata"]["embedding"] = self._parse_embedding(row[5]) + return result + + @staticmethod + def _iso(value: Any) -> str | None: + if value is None: + return None + if isinstance(value, datetime): + return value.isoformat() + return str(value) + + @staticmethod + def _parse_embedding(value: Any) -> Any: + if value is None or isinstance(value, list): + return value + text = str(value).strip().strip("[]") + if not text: + return [] + try: + return [float(x) for x in text.split(",")] + except ValueError: + return value + + # ========================================================================= + # Metadata filter building (MySQL/OceanBase JSON dialect) + # ========================================================================= + + @staticmethod + def _is_safe_field_name(field: str) -> bool: + return bool(_SAFE_FIELD_RE.match(field)) + + def _field_expr(self, key: str) -> tuple[str, str]: + """Build (text_expr, json_expr) SQL fragments for a filter key.""" + direct_columns = {"id", "memory", "user_name", "created_at", "updated_at"} + if key in direct_columns: + return key, key + + if key.startswith("info."): + sub_key = key[5:] + if not self._is_safe_field_name(sub_key): + raise ValueError(f"Invalid filter field: {key}") + path = f"$.info.{sub_key}" + else: + if not self._is_safe_field_name(key): + raise ValueError(f"Invalid filter field: {key}") + path = f"$.{key}" + + text_expr = f"JSON_UNQUOTE(JSON_EXTRACT(properties, '{path}'))" + json_expr = f"JSON_EXTRACT(properties, '{path}')" + return text_expr, json_expr + + def _build_single_filter_condition( + self, condition_dict: dict[str, Any], params: list[Any] + ) -> str | None: + """Build SQL for a single filter condition dict.""" + if not condition_dict: + return None + + array_fields = {"tags", "sources", "file_ids"} + timestamp_fields = {"created_at", "updated_at"} + parts: list[str] = [] + + for key, value in condition_dict.items(): + text_expr, json_expr = self._field_expr(key) + raw_key = key[5:] if key.startswith("info.") else key + + if isinstance(value, dict): + for op, op_value in value.items(): + if op in ("gt", "lt", "gte", "lte"): + op_map = {"gt": ">", "lt": "<", "gte": ">=", "lte": "<="} + sql_op = op_map[op] + if raw_key in timestamp_fields or raw_key.endswith("_at"): + parts.append( + f"CAST({text_expr} AS DATETIME) {sql_op} CAST(%s AS DATETIME)" + ) + params.append(op_value) + else: + parts.append( + f"CAST(NULLIF({text_expr}, '') AS DECIMAL(38,10)) {sql_op} %s" + ) + params.append(op_value) + elif op == "contains": + if raw_key in array_fields: + parts.append(f"JSON_CONTAINS({json_expr}, %s)") + params.append(json.dumps(op_value)) + else: + parts.append(f"{text_expr} LIKE %s") + params.append(f"%{op_value}%") + elif op == "in": + if not isinstance(op_value, list): + raise ValueError( + f"in operator expects list for '{key}', got {type(op_value).__name__}" + ) + if raw_key in array_fields: + parts.append(f"JSON_OVERLAPS({json_expr}, %s)") + params.append(json.dumps([str(v) for v in op_value])) + else: + ph = self._placeholders(op_value) + parts.append(f"{text_expr} IN ({ph})") + params.extend([str(v) for v in op_value]) + elif op == "like": + parts.append(f"{text_expr} LIKE %s") + params.append(f"%{op_value}%") + else: + raise ValueError(f"Unsupported filter operator: {op}") + elif raw_key in array_fields: + if isinstance(value, list): + parts.append(f"JSON_CONTAINS({json_expr}, %s)") + params.append(json.dumps(value)) + else: + parts.append(f"JSON_CONTAINS({json_expr}, %s)") + params.append(json.dumps([value])) + else: + parts.append(f"{text_expr} = %s") + params.append(str(value)) + + if not parts: + return None + return " AND ".join(parts) + + def _build_filter_where_clause(self, filter_dict: dict[str, Any], params: list[Any]) -> str: + """Build a SQL WHERE fragment from a filter dict.""" + if not filter_dict: + return "" + + if "and" in filter_dict: + and_conditions = filter_dict.get("and") + if not isinstance(and_conditions, list): + raise ValueError("Invalid filter format: 'and' must be a list") + parts: list[str] = [] + for cond in and_conditions: + if isinstance(cond, dict): + cond_sql = self._build_single_filter_condition(cond, params) + if cond_sql: + parts.append(f"({cond_sql})") + return " AND ".join(parts) + + if "or" in filter_dict: + or_conditions = filter_dict.get("or") + if not isinstance(or_conditions, list): + raise ValueError("Invalid filter format: 'or' must be a list") + parts = [] + for cond in or_conditions: + if isinstance(cond, dict): + cond_sql = self._build_single_filter_condition(cond, params) + if cond_sql: + parts.append(f"({cond_sql})") + return f"({' OR '.join(parts)})" if parts else "" + + cond_sql = self._build_single_filter_condition(filter_dict, params) + return cond_sql or "" + + def delete_node_by_prams( + self, + writable_cube_ids: list[str] | None = None, + memory_ids: list[str] | None = None, + file_ids: list[str] | None = None, + filter: dict | None = None, + ) -> int: + """Delete nodes by memory_ids, file_ids, or filter (and clean up edges).""" + where_conditions: list[str] = [] + params: list[Any] = [] + + if memory_ids: + ph = self._placeholders(memory_ids) + where_conditions.append(f"id IN ({ph})") + params.extend(memory_ids) + + if file_ids: + file_conditions: list[str] = [] + for file_id in file_ids: + file_conditions.append("JSON_CONTAINS(JSON_EXTRACT(properties, '$.file_ids'), %s)") + params.append(json.dumps([file_id])) + if file_conditions: + where_conditions.append(f"({' OR '.join(file_conditions)})") + + if filter: + filter_where = self._build_filter_where_clause(filter, params) + if filter_where: + where_conditions.append(f"({filter_where})") + + if not where_conditions: + logger.warning("[delete_node_by_prams] No memory_ids, file_ids, or filter provided") + return 0 + + if writable_cube_ids: + ph = self._placeholders(writable_cube_ids) + where_conditions.append(f"user_name IN ({ph})") + params.extend(writable_cube_ids) + + where_clause = " AND ".join(where_conditions) + + with self._transaction() as cur: + cur.execute(f"SELECT id FROM {self.nodes_table} WHERE {where_clause}", params) + ids = [row[0] for row in cur.fetchall()] + if not ids: + return 0 + + ph = self._placeholders(ids) + cur.execute( + f"DELETE FROM {self.edges_table} WHERE source_id IN ({ph}) OR target_id IN ({ph})", + (*ids, *ids), + ) + cur.execute( + f"DELETE FROM {self.nodes_table} WHERE id IN ({ph})", + tuple(ids), + ) + deleted = cur.rowcount if cur.rowcount is not None else 0 + logger.info("[delete_node_by_prams] Deleted %s nodes", deleted) + return deleted + + # ========================================================================= + # Edge management + # ========================================================================= + + def add_edge( + self, source_id: str, target_id: str, type: str, user_name: str | None = None + ) -> None: + """Create an edge between nodes (idempotent).""" + self._execute( + f""" + INSERT INTO {self.edges_table} (source_id, target_id, edge_type) + VALUES (%s, %s, %s) + ON DUPLICATE KEY UPDATE edge_type = VALUES(edge_type) + """, + (source_id, target_id, type), + ) + + def delete_edge( + self, source_id: str, target_id: str, type: str, user_name: str | None = None + ) -> None: + """Delete an edge.""" + self._execute( + f""" + DELETE FROM {self.edges_table} + WHERE source_id = %s AND target_id = %s AND edge_type = %s + """, + (source_id, target_id, type), + ) + + def edge_exists(self, source_id: str, target_id: str, type: str) -> bool: + """Check if an edge exists.""" + row = self._query_one( + f""" + SELECT 1 FROM {self.edges_table} + WHERE source_id = %s AND target_id = %s AND edge_type = %s + LIMIT 1 + """, + (source_id, target_id, type), + ) + return row is not None + + # ========================================================================= + # Graph queries + # ========================================================================= + + def get_neighbors( + self, id: str, type: str, direction: Literal["in", "out", "both"] = "out" + ) -> list[str]: + """Get neighboring node IDs by relationship type and direction.""" + if direction == "out": + rows = self._query( + f"SELECT target_id FROM {self.edges_table} WHERE source_id = %s AND edge_type = %s", + (id, type), + ) + elif direction == "in": + rows = self._query( + f"SELECT source_id FROM {self.edges_table} WHERE target_id = %s AND edge_type = %s", + (id, type), + ) + else: + rows = self._query( + f""" + SELECT target_id FROM {self.edges_table} WHERE source_id = %s AND edge_type = %s + UNION + SELECT source_id FROM {self.edges_table} WHERE target_id = %s AND edge_type = %s + """, + (id, type, id, type), + ) + return [row[0] for row in rows] + + def get_path(self, source_id: str, target_id: str, max_depth: int = 3) -> list[str]: + """Get a shortest path from source to target via a recursive CTE. + + Mirrors the PostgreSQL backend, which accumulates the path in a native + array. OceanBase/MySQL have no array type, so a JSON array is used + instead: ``JSON_ARRAY_APPEND`` grows the path per element (no delimiter, + so node ids may contain any character) and ``JSON_CONTAINS`` guards + against revisiting a node (cycle guard). + """ + if source_id == target_id: + return [source_id] + + row = self._query_one( + f""" + WITH RECURSIVE paths AS ( + SELECT + target_id, + JSON_ARRAY(source_id) AS nodes, + 1 AS depth + FROM {self.edges_table} + WHERE source_id = %s + UNION ALL + SELECT + e.target_id, + JSON_ARRAY_APPEND(p.nodes, '$', e.source_id), + p.depth + 1 + FROM {self.edges_table} e + JOIN paths p ON e.source_id = p.target_id + WHERE p.depth < %s + AND NOT JSON_CONTAINS(p.nodes, JSON_QUOTE(e.source_id)) + ) + SELECT JSON_ARRAY_APPEND(nodes, '$', target_id) AS full_path + FROM paths + WHERE target_id = %s + ORDER BY depth + LIMIT 1 + """, + (source_id, max_depth, target_id), + ) + if not row or not row[0]: + return [] + full_path = row[0] + return full_path if isinstance(full_path, list) else json.loads(full_path) + + def get_subgraph(self, center_id: str, depth: int = 2) -> list[str]: + """Get the node IDs of the subgraph around a center node via a recursive CTE. + + seekdb/OceanBase do not support ``UNION`` (DISTINCT) inside a recursive + CTE, so ``UNION ALL`` is used together with a JSON-array visited-set to + guard against cycles (mirroring ``get_path``); the outer + ``SELECT DISTINCT`` collapses the duplicate rows that ``UNION ALL`` keeps. + """ + rows = self._query( + f""" + WITH RECURSIVE subgraph AS ( + SELECT + CAST(%s AS CHAR(255)) AS node_id, + 0 AS level, + JSON_ARRAY(%s) AS visited + UNION ALL + SELECT + CASE WHEN e.source_id = s.node_id THEN e.target_id ELSE e.source_id END, + s.level + 1, + JSON_ARRAY_APPEND( + s.visited, '$', + CASE WHEN e.source_id = s.node_id THEN e.target_id ELSE e.source_id END + ) + FROM subgraph s + JOIN {self.edges_table} e + ON (e.source_id = s.node_id OR e.target_id = s.node_id) + WHERE s.level < %s + AND NOT JSON_CONTAINS( + s.visited, + JSON_QUOTE( + CASE WHEN e.source_id = s.node_id THEN e.target_id ELSE e.source_id END + ) + ) + ) + SELECT DISTINCT node_id FROM subgraph + """, + (center_id, center_id, depth), + ) + return [row[0] for row in rows] + + def get_context_chain(self, id: str, type: str = "FOLLOWS") -> list[str]: + """Get the ordered context chain following a relationship type.""" + return self.get_neighbors(id, type, "out") + + # ========================================================================= + # Search operations + # ========================================================================= + + def search_by_embedding( + self, + vector: list[float], + top_k: int = 5, + 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, + **kwargs, + ) -> list[dict]: + """Search nodes by vector similarity, applying tenant filters before ranking.""" + user_name = user_name or self.user_name + + conditions = ["embedding IS NOT NULL"] + params: list[Any] = [] + + if user_name: + conditions.append("user_name = %s") + params.append(user_name) + + if scope: + conditions.append("JSON_UNQUOTE(JSON_EXTRACT(properties, '$.memory_type')) = %s") + params.append(scope) + + if status: + conditions.append("JSON_UNQUOTE(JSON_EXTRACT(properties, '$.status')) = %s") + params.append(status) + else: + conditions.append( + "(JSON_UNQUOTE(JSON_EXTRACT(properties, '$.status')) = 'activated' " + "OR JSON_EXTRACT(properties, '$.status') IS NULL)" + ) + + if search_filter: + for k, v in search_filter.items(): + if not self._is_safe_field_name(k): + raise ValueError(f"Invalid search_filter field: {k}") + conditions.append(f"JSON_UNQUOTE(JSON_EXTRACT(properties, '$.{k}')) = %s") + params.append(str(v)) + + where_clause = " AND ".join(conditions) + vec = self._vec_literal(vector) + + rows = self._query( + f""" + SELECT id, 1 - cosine_distance(embedding, %s) AS score + FROM {self.nodes_table} + WHERE {where_clause} + ORDER BY cosine_distance(embedding, %s) + LIMIT %s + """, + (vec, *params, vec, top_k), + ) + + results = [] + for row in rows: + score = float(row[1]) + if threshold is None or score >= threshold: + results.append({"id": row[0], "score": score}) + return results + + def get_by_metadata( + self, + filters: list[dict[str, Any]], + status: str | None = None, + user_name: str | None = None, + filter: dict | None = None, + knowledgebase_ids: list[str] | None = None, + user_name_flag: bool = True, + ) -> list[str]: + """Get node IDs matching metadata filters.""" + user_name = user_name or self.user_name + + conditions: list[str] = [] + params: list[Any] = [] + + if user_name_flag and user_name: + conditions.append("user_name = %s") + params.append(user_name) + + if status: + conditions.append("JSON_UNQUOTE(JSON_EXTRACT(properties, '$.status')) = %s") + params.append(status) + + for f in filters: + field = f["field"] + op = f.get("op", "=") + value = f["value"] + if not self._is_safe_field_name(field): + raise ValueError(f"Invalid filter field: {field}") + text_expr = f"JSON_UNQUOTE(JSON_EXTRACT(properties, '$.{field}'))" + json_expr = f"JSON_EXTRACT(properties, '$.{field}')" + + if op == "=": + conditions.append(f"{text_expr} = %s") + params.append(str(value)) + elif op == "in": + ph = self._placeholders(value) + conditions.append(f"{text_expr} IN ({ph})") + params.extend([str(v) for v in value]) + elif op in (">", ">=", "<", "<="): + conditions.append(f"CAST({text_expr} AS DECIMAL(38,10)) {op} %s") + params.append(value) + elif op == "contains": + conditions.append(f"JSON_CONTAINS({json_expr}, %s)") + params.append(json.dumps([value])) + + where_clause = " AND ".join(conditions) if conditions else "TRUE" + + rows = self._query(f"SELECT id FROM {self.nodes_table} WHERE {where_clause}", params) + return [row[0] for row in rows] + + def get_all_memory_items( + self, + scope: str, + include_embedding: bool = False, + status: str | None = None, + filter: dict | None = None, + knowledgebase_ids: list[str] | None = None, + **kwargs, + ) -> list[dict]: + """Get all memory items of a specific memory_type.""" + user_name = kwargs.get("user_name") or self.user_name + + conditions = [ + "JSON_UNQUOTE(JSON_EXTRACT(properties, '$.memory_type')) = %s", + "user_name = %s", + ] + params: list[Any] = [scope, user_name] + + if status: + conditions.append("JSON_UNQUOTE(JSON_EXTRACT(properties, '$.status')) = %s") + params.append(status) + + where_clause = " AND ".join(conditions) + cols = "id, memory, properties, created_at, updated_at" + if include_embedding: + cols += ", embedding" + rows = self._query(f"SELECT {cols} FROM {self.nodes_table} WHERE {where_clause}", params) + return [self._parse_row(row, include_embedding) for row in rows] + + def get_structure_optimization_candidates( + self, scope: str, include_embedding: bool = False + ) -> list[dict]: + """Find isolated nodes (no incident edges) for a given scope.""" + user_name = self.user_name + cols = "m.id, m.memory, m.properties, m.created_at, m.updated_at" + if include_embedding: + cols += ", m.embedding" + rows = self._query( + f""" + SELECT {cols} + FROM {self.nodes_table} m + LEFT JOIN {self.edges_table} e1 ON m.id = e1.source_id + LEFT JOIN {self.edges_table} e2 ON m.id = e2.target_id + WHERE JSON_UNQUOTE(JSON_EXTRACT(m.properties, '$.memory_type')) = %s + AND m.user_name = %s + AND JSON_UNQUOTE(JSON_EXTRACT(m.properties, '$.status')) = 'activated' + AND e1.id IS NULL + AND e2.id IS NULL + """, + (scope, user_name), + ) + return [self._parse_row(row, include_embedding) for row in rows] + + # ========================================================================= + # Maintenance + # ========================================================================= + + def deduplicate_nodes(self) -> None: + """Not implemented - handled at application level.""" + + def get_grouped_counts( + 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 the given JSON property fields.""" + user_name = user_name or self.user_name + if not group_fields: + raise ValueError("group_fields cannot be empty") + for field in group_fields: + if not self._is_safe_field_name(field): + raise ValueError(f"Invalid group field: {field}") + + select_fields = ", ".join( + [ + f"JSON_UNQUOTE(JSON_EXTRACT(properties, '$.{field}')) AS {field}" + for field in group_fields + ] + ) + group_by = ", ".join( + [f"JSON_UNQUOTE(JSON_EXTRACT(properties, '$.{field}'))" for field in group_fields] + ) + + conditions = ["user_name = %s"] + query_params: list[Any] = [user_name] + + if where_clause: + where_clause = where_clause.strip() + if where_clause.upper().startswith("WHERE"): + where_clause = where_clause[5:].strip() + if where_clause: + conditions.append(where_clause) + if params: + query_params.extend(params.values()) + + where_sql = " AND ".join(conditions) + + rows = self._query( + f""" + SELECT {select_fields}, COUNT(*) AS count + FROM {self.nodes_table} + WHERE {where_sql} + GROUP BY {group_by} + """, + query_params, + ) + results = [] + for row in rows: + result = {} + for i, field in enumerate(group_fields): + result[field] = row[i] + result["count"] = row[len(group_fields)] + results.append(result) + return results + + def detect_conflicts(self) -> list[tuple[str, str]]: + """Not implemented.""" + return [] + + def merge_nodes(self, id1: str, id2: str) -> str: + """Not implemented.""" + raise NotImplementedError + + def clear(self, user_name: str | None = None) -> None: + """Clear all graph data for the given tenant (never the whole database).""" + user_name = user_name or self.user_name + with self._transaction() as cur: + cur.execute( + f"SELECT id FROM {self.nodes_table} WHERE user_name = %s", + (user_name,), + ) + ids = [row[0] for row in cur.fetchall()] + if ids: + ph = self._placeholders(ids) + cur.execute( + f"DELETE FROM {self.edges_table} " + f"WHERE source_id IN ({ph}) OR target_id IN ({ph})", + (*ids, *ids), + ) + cur.execute( + f"DELETE FROM {self.nodes_table} WHERE user_name = %s", + (user_name,), + ) + logger.info("Cleared all graph data for user %s", user_name) + + def drop_database(self, user_name: str | None = None) -> None: + """Scoped clear: remove the current tenant's graph data only. + + Redefined (vs. dropping a physical database) because the OceanBase / + seekdb database is shared across MemOS modules. + """ + self.clear(user_name=user_name) + + def export_graph(self, include_embedding: bool = False, **kwargs) -> dict[str, Any]: + """Export all graph data for the given tenant.""" + user_name = kwargs.get("user_name") or self.user_name + cols = "id, memory, properties, created_at, updated_at" + if include_embedding: + cols += ", embedding" + node_rows = self._query( + f""" + SELECT {cols} FROM {self.nodes_table} + WHERE user_name = %s + ORDER BY created_at DESC + """, + (user_name,), + ) + nodes = [self._parse_row(row, include_embedding) for row in node_rows] + + node_ids = [n["id"] for n in nodes] + if node_ids: + ph = self._placeholders(node_ids) + edge_rows = self._query( + f""" + SELECT source_id, target_id, edge_type + FROM {self.edges_table} + WHERE source_id IN ({ph}) OR target_id IN ({ph}) + """, + (*node_ids, *node_ids), + ) + edges = [{"source": row[0], "target": row[1], "type": row[2]} for row in edge_rows] + else: + edges = [] + + return { + "nodes": nodes, + "edges": edges, + "total_nodes": len(nodes), + "total_edges": len(edges), + } + + def import_graph(self, data: dict[str, Any], user_name: str | None = None) -> None: + """Import graph data.""" + user_name = user_name or self.user_name + for node in data.get("nodes", []): + self.add_node( + id=node["id"], + memory=node.get("memory", ""), + metadata=node.get("metadata", {}), + user_name=user_name, + ) + for edge in data.get("edges", []): + self.add_edge( + source_id=edge["source"], + target_id=edge["target"], + type=edge["type"], + ) + + def close(self): + """Close the connection pool; further queries are rejected afterwards.""" + if self._closed: + return + self._closed = True + self._pool.close() diff --git a/src/memos/vec_dbs/factory.py b/src/memos/vec_dbs/factory.py index f2950b4ea..7fbdeccd5 100644 --- a/src/memos/vec_dbs/factory.py +++ b/src/memos/vec_dbs/factory.py @@ -3,6 +3,7 @@ from memos.configs.vec_db import VectorDBConfigFactory from memos.vec_dbs.base import BaseVecDB from memos.vec_dbs.milvus import MilvusVecDB +from memos.vec_dbs.oceanbase import OceanBaseVecDB from memos.vec_dbs.qdrant import QdrantVecDB @@ -12,6 +13,8 @@ class VecDBFactory(BaseVecDB): backend_to_class: ClassVar[dict[str, Any]] = { "qdrant": QdrantVecDB, "milvus": MilvusVecDB, + "oceanbase": OceanBaseVecDB, + "seekdb": OceanBaseVecDB, } @classmethod diff --git a/src/memos/vec_dbs/oceanbase.py b/src/memos/vec_dbs/oceanbase.py new file mode 100644 index 000000000..f3884328e --- /dev/null +++ b/src/memos/vec_dbs/oceanbase.py @@ -0,0 +1,311 @@ +"""OceanBase / seekdb vector database backend for MemOS. + +Implemented on top of pyseekdb's Collection / vector API (chromadb-style). The +same implementation connects to seekdb Server and OceanBase Server; only the +endpoint differs. + +The MemOS ``VecDBItem`` is mapped onto a pyseekdb Collection record as: +- ``vector`` -> ``embeddings`` +- ``memory`` -> ``documents`` (best-effort, taken from payload) +- ``payload`` -> ``metadatas`` (stored directly; pyseekdb's ``metadata`` is a JSON + column, so the full nested payload is persisted and read back verbatim). +""" + +from typing import Any + +from memos.configs.vec_db import OceanBaseVecDBConfig +from memos.dependency import require_python_package +from memos.log import get_logger +from memos.vec_dbs.base import BaseVecDB +from memos.vec_dbs.item import VecDBItem + + +logger = get_logger(__name__) + +# MemOS distance metric -> pyseekdb distance metric. +_DISTANCE_MAP = { + "cosine": "cosine", + "euclidean": "l2", + "dot": "inner_product", +} + + +class OceanBaseVecDB(BaseVecDB): + """OceanBase / seekdb vector database implementation (via pyseekdb).""" + + @require_python_package( + import_name="pyseekdb", + install_command="pip install MemoryOS[ob-mem]", + install_link="https://github.com/oceanbase/pyseekdb", + ) + def __init__(self, config: OceanBaseVecDBConfig): + """Initialize the client and ensure the default collection exists.""" + import pyseekdb + + self.config = config + self._distance = _DISTANCE_MAP.get(config.distance_metric or "cosine", "cosine") + + self.client = pyseekdb.Client( + host=config.host, + port=config.port, + database=config.database, + user=config.user, + password=config.password, + ) + self.collection = None + self.create_collection() + + # ------------------------------------------------------------------ + # Collection management + # ------------------------------------------------------------------ + + def create_collection(self) -> None: + """Create the default collection if it does not exist.""" + from pyseekdb import HNSWConfiguration + + name = self.config.collection_name + if self.client.has_collection(name): + logger.warning("Collection '%s' already exists. Skipping creation.", name) + self.collection = self.client.get_collection(name, embedding_function=None) + return + + configuration = HNSWConfiguration( + dimension=self.config.vector_dimension, + distance=self._distance, + ) + self.collection = self.client.create_collection( + name=name, + configuration=configuration, + embedding_function=None, + ) + logger.info( + "Collection '%s' created with %s dimensions (distance=%s).", + name, + self.config.vector_dimension, + self._distance, + ) + + def list_collections(self) -> list[str]: + """List all collections.""" + return [c.name for c in self.client.list_collections()] + + def delete_collection(self, name: str) -> None: + """Delete a collection.""" + self.client.delete_collection(name) + + def collection_exists(self, name: str) -> bool: + """Check if a collection exists.""" + return bool(self.client.has_collection(name)) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + @staticmethod + def _document_of(payload: dict[str, Any] | None) -> str: + """Best-effort document text for full-text features.""" + if payload and isinstance(payload.get("memory"), str): + return payload["memory"] + return "" + + def _distance_to_score(self, distance: float | None) -> float | None: + """Normalize a pyseekdb distance to a 'higher is more similar' score. + + For ``cosine`` the score is the cosine similarity (pyseekdb returns + ``1 - cosine_similarity``), which is already a meaningful absolute value. + + For ``l2`` and ``inner_product`` the raw distance is unbounded (l2) or + unbounded/signed (inner_product), which makes absolute-threshold + comparisons unreliable. We defensively squash them into a bounded + similarity in ``(0, 1]`` / ``(0, 1)`` while preserving the exact same + monotonic ordering as the raw distance, so relative ranking is unchanged. + """ + if distance is None: + return None + d = float(distance) + if self._distance == "cosine": + return 1.0 - d + if self._distance == "inner_product": + # Higher raw distance -> more similar. Squash to (0, 1) monotonically. + return 0.5 * (1.0 + d / (1.0 + abs(d))) + # l2 / euclidean: non-negative distance, smaller is closer. Map to (0, 1]. + return 1.0 / (1.0 + d) + + # ------------------------------------------------------------------ + # Search / read + # ------------------------------------------------------------------ + + def search( + self, query_vector: list[float], top_k: int, filter: dict[str, Any] | None = None + ) -> list[VecDBItem]: + """Search for similar items in the collection.""" + response = self.collection.query( + query_embeddings=query_vector, + n_results=top_k, + where=filter or None, + include=["documents", "metadatas", "embeddings"], + ) + ids = (response.get("ids") or [[]])[0] + metadatas = (response.get("metadatas") or [[]])[0] + embeddings = (response.get("embeddings") or [[]])[0] + distances = (response.get("distances") or [[]])[0] + + results: list[VecDBItem] = [] + for idx, item_id in enumerate(ids): + metadata = metadatas[idx] if idx < len(metadatas) else None + vector = embeddings[idx] if idx < len(embeddings) else None + distance = distances[idx] if idx < len(distances) else None + results.append( + VecDBItem( + id=item_id, + vector=list(vector) if vector is not None else None, + payload=metadata or None, + score=self._distance_to_score(distance), + ) + ) + logger.info("OceanBase vector search completed with %s results.", len(results)) + return results + + def _items_from_get(self, response: dict[str, Any]) -> list[VecDBItem]: + """Convert a pyseekdb ``get`` response (flat lists) into VecDBItems.""" + ids = response.get("ids") or [] + metadatas = response.get("metadatas") or [] + embeddings = response.get("embeddings") or [] + items: list[VecDBItem] = [] + for idx, item_id in enumerate(ids): + metadata = metadatas[idx] if idx < len(metadatas) else None + vector = embeddings[idx] if idx < len(embeddings) else None + items.append( + VecDBItem( + id=item_id, + vector=list(vector) if vector is not None else None, + payload=metadata or None, + ) + ) + return items + + def get_by_id(self, id: str) -> VecDBItem | None: + """Get a single item by ID.""" + response = self.collection.get(ids=[id], include=["metadatas", "embeddings"]) + items = self._items_from_get(response) + return items[0] if items else None + + def get_by_ids(self, ids: list[str]) -> list[VecDBItem]: + """Get multiple items by their IDs.""" + if not ids: + return [] + response = self.collection.get(ids=ids, include=["metadatas", "embeddings"]) + return self._items_from_get(response) + + def get_by_filter(self, filter: dict[str, Any], scroll_limit: int = 100) -> list[VecDBItem]: + """ + Retrieve all items matching the given filter, paginating through results. + + Args: + filter: Payload filters to match against stored items + scroll_limit: Maximum number of items to retrieve per scroll request + """ + items: list[VecDBItem] = [] + offset = 0 + while True: + response = self.collection.get( + where=filter or None, + limit=scroll_limit, + offset=offset, + include=["metadatas", "embeddings"], + ) + batch = self._items_from_get(response) + items.extend(batch) + if len(batch) < scroll_limit: + break + offset += scroll_limit + return items + + def get_all(self, scroll_limit: int = 100) -> list[VecDBItem]: + """Retrieve all items in the collection.""" + return self.get_by_filter({}, scroll_limit=scroll_limit) + + def count(self, filter: dict[str, Any] | None = None) -> int: + """Count items, optionally with a filter.""" + if not filter: + return int(self.collection.count()) + return len(self.get_by_filter(filter)) + + # ------------------------------------------------------------------ + # Write + # ------------------------------------------------------------------ + + @staticmethod + def _normalize_item(item: VecDBItem | dict[str, Any]) -> VecDBItem: + if isinstance(item, dict): + return VecDBItem.from_dict(item.copy()) + return item + + def _columns_from( + self, data: list[VecDBItem | dict[str, Any]] + ) -> tuple[list[str], list[list[float] | None], list[str], list[dict[str, Any] | None]]: + """Split items into the parallel column lists expected by pyseekdb. + + ``payload`` is mapped onto ``metadatas`` and ``payload['memory']`` onto + ``documents`` (best-effort). ``payload`` may be ``None`` for a given item; + pyseekdb persists that as a SQL ``NULL`` metadata value. + """ + ids: list[str] = [] + embeddings: list[list[float] | None] = [] + documents: list[str] = [] + metadatas: list[dict[str, Any] | None] = [] + for raw in data: + item = self._normalize_item(raw) + ids.append(item.id) + embeddings.append(item.vector) + documents.append(self._document_of(item.payload)) + metadatas.append(item.payload) + return ids, embeddings, documents, metadatas + + def add(self, data: list[VecDBItem | dict[str, Any]]) -> None: + """Add data to the collection.""" + ids, embeddings, documents, metadatas = self._columns_from(data) + if not ids: + return + self.collection.add( + ids=ids, + embeddings=embeddings, + documents=documents, + metadatas=metadatas, + ) + + def update(self, id: str, data: VecDBItem | dict[str, Any]) -> None: + """Update an item in the collection.""" + item = self._normalize_item(data) + metadata = item.payload + if item.vector: + self.collection.update( + ids=id, + embeddings=item.vector, + documents=self._document_of(item.payload), + metadatas=metadata, + ) + else: + self.collection.update(ids=id, metadatas=metadata) + + def upsert(self, data: list[VecDBItem | dict[str, Any]]) -> None: + """Add or update data in the collection.""" + ids, embeddings, documents, metadatas = self._columns_from(data) + if not ids: + return + self.collection.upsert( + ids=ids, + embeddings=embeddings, + documents=documents, + metadatas=metadatas, + ) + + def delete(self, ids: list[str]) -> None: + """Delete items from the collection by IDs.""" + if not ids: + return + self.collection.delete(ids=ids) + + def ensure_payload_indexes(self, fields: list[str]) -> None: + """No-op: pyseekdb Collections index metadata automatically.""" + logger.debug("ensure_payload_indexes is a no-op for OceanBase Collection API: %s", fields) diff --git a/tests/graph_dbs/test_oceanbase.py b/tests/graph_dbs/test_oceanbase.py new file mode 100644 index 000000000..c4a58a194 --- /dev/null +++ b/tests/graph_dbs/test_oceanbase.py @@ -0,0 +1,146 @@ +from unittest.mock import MagicMock, patch + +import pytest + +from memos.configs.graph_db import GraphDBConfigFactory +from memos.graph_dbs.factory import GraphStoreFactory + + +@pytest.fixture +def cursor(): + return MagicMock(name="cursor") + + +@pytest.fixture +def graph_db(cursor): + conn = MagicMock(name="conn") + conn.cursor.return_value = cursor + config = GraphDBConfigFactory.model_validate( + { + "backend": "oceanbase", + "config": { + "host": "127.0.0.1", + "port": 2881, + "user": "root", + "password": "", + "db_name": "memos", + "user_name": "alice", + "embedding_dimension": 3, + }, + } + ) + with patch("pymysql.connect", return_value=conn): + db = GraphStoreFactory.from_config(config) + return db + + +def _executed_sql(cursor): + return [str(call.args[0]) for call in cursor.execute.call_args_list] + + +def test_backend_registered(): + for backend in ("oceanbase", "seekdb"): + assert backend in GraphStoreFactory.backend_to_class + + +def test_schema_initialized(graph_db, cursor): + ddl = " ".join(_executed_sql(cursor)) + assert "memos_graph_nodes" in ddl + assert "memos_graph_edges" in ddl + + +def test_add_node(graph_db, cursor): + graph_db.add_node("n1", "hello", {"status": "activated"}) + sqls = _executed_sql(cursor) + assert any("INSERT INTO memos_graph_nodes" in s for s in sqls) + + +def test_add_node_with_embedding(graph_db, cursor): + graph_db.add_node("n1", "hello", {"status": "activated", "embedding": [0.1, 0.2, 0.3]}) + insert_sql = next(s for s in _executed_sql(cursor) if "INSERT INTO memos_graph_nodes" in s) + assert "embedding" in insert_sql + + +def test_get_node(graph_db, cursor): + cursor.fetchone.return_value = ( + "n1", + "hello", + '{"status": "activated"}', + "2024-01-01 00:00:00", + "2024-01-02 00:00:00", + ) + node = graph_db.get_node("n1") + assert node["id"] == "n1" + assert node["memory"] == "hello" + assert node["metadata"]["status"] == "activated" + assert node["metadata"]["created_at"] == "2024-01-01 00:00:00" + + +def test_get_node_missing(graph_db, cursor): + cursor.fetchone.return_value = None + assert graph_db.get_node("nope") is None + + +def test_edge_exists_true(graph_db, cursor): + cursor.fetchone.return_value = (1,) + assert graph_db.edge_exists("a", "b", "FOLLOWS") is True + + +def test_edge_exists_false(graph_db, cursor): + cursor.fetchone.return_value = None + assert graph_db.edge_exists("a", "b", "FOLLOWS") is False + + +def test_get_neighbors(graph_db, cursor): + cursor.fetchall.return_value = [("b",), ("c",)] + assert graph_db.get_neighbors("a", "FOLLOWS", "out") == ["b", "c"] + + +def test_search_by_embedding_threshold(graph_db, cursor): + cursor.fetchall.return_value = [("n1", 0.9), ("n2", 0.4)] + results = graph_db.search_by_embedding([0.1, 0.2, 0.3], top_k=5, threshold=0.5) + assert results == [{"id": "n1", "score": 0.9}] + + +def test_get_path_cte(graph_db, cursor): + # Recursive CTE returns a single row whose JSON array holds the full path. + cursor.fetchone.return_value = ('["a", "b", "c"]',) + assert graph_db.get_path("a", "c", max_depth=3) == ["a", "b", "c"] + + +def test_get_path_not_found(graph_db, cursor): + cursor.fetchone.return_value = None + assert graph_db.get_path("a", "z", max_depth=3) == [] + + +def test_get_path_same_node(graph_db): + assert graph_db.get_path("a", "a") == ["a"] + + +def test_get_subgraph_cte(graph_db, cursor): + # Recursive CTE returns one node id per row. + cursor.fetchall.return_value = [("a",), ("b",), ("c",)] + result = set(graph_db.get_subgraph("a", depth=1)) + assert result == {"a", "b", "c"} + + +def test_get_by_metadata(graph_db, cursor): + cursor.fetchall.return_value = [("n1",), ("n2",)] + ids = graph_db.get_by_metadata([{"field": "topic", "value": "psychology"}]) + assert ids == ["n1", "n2"] + + +def test_clear_is_tenant_scoped(graph_db, cursor): + cursor.fetchall.return_value = [("n1",)] + graph_db.clear() + sqls = _executed_sql(cursor) + delete_nodes = [s for s in sqls if "DELETE FROM memos_graph_nodes" in s] + assert delete_nodes + assert all("user_name = %s" in s for s in delete_nodes) + + +def test_drop_database_never_drops_physical_db(graph_db, cursor): + cursor.fetchall.return_value = [("n1",)] + graph_db.drop_database() + all_sql = " ".join(_executed_sql(cursor)).upper() + assert "DROP DATABASE" not in all_sql diff --git a/tests/vec_dbs/test_oceanbase.py b/tests/vec_dbs/test_oceanbase.py new file mode 100644 index 000000000..aa13d1d45 --- /dev/null +++ b/tests/vec_dbs/test_oceanbase.py @@ -0,0 +1,188 @@ +import sys +import types +import uuid + +from unittest.mock import MagicMock + +import pytest + +from memos.configs.vec_db import VectorDBConfigFactory +from memos.vec_dbs.factory import VecDBFactory +from memos.vec_dbs.item import VecDBItem + + +@pytest.fixture +def fake_pyseekdb(): + """Inject a fake ``pyseekdb`` module so the provider can be constructed.""" + module = types.ModuleType("pyseekdb") + module.Client = MagicMock(name="Client") + module.HNSWConfiguration = MagicMock(name="HNSWConfiguration") + sys.modules["pyseekdb"] = module + try: + yield module + finally: + sys.modules.pop("pyseekdb", None) + + +@pytest.fixture +def config(): + return VectorDBConfigFactory.model_validate( + { + "backend": "oceanbase", + "config": { + "collection_name": "test_collection", + "vector_dimension": 3, + "distance_metric": "cosine", + "host": "127.0.0.1", + "port": 2881, + "user": "root", + "password": "", + "database": "memos", + }, + } + ) + + +@pytest.fixture +def collection(): + return MagicMock(name="collection") + + +@pytest.fixture +def vec_db(config, fake_pyseekdb, collection): + client = fake_pyseekdb.Client.return_value + client.has_collection.return_value = False + client.create_collection.return_value = collection + client.get_collection.return_value = collection + return VecDBFactory.from_config(config) + + +def test_backend_registered(): + for backend in ("oceanbase", "seekdb"): + assert backend in VecDBFactory.backend_to_class + + +def test_create_collection(vec_db, collection): + vec_db.client.create_collection.assert_called_once() + assert vec_db.collection is collection + assert vec_db.config.collection_name == "test_collection" + + +def test_add(vec_db, collection): + item_id = str(uuid.uuid4()) + vec_db.add( + [ + { + "id": item_id, + "vector": [0.1, 0.2, 0.3], + "payload": {"memory": "hi", "status": "activated"}, + } + ] + ) + collection.add.assert_called_once() + kwargs = collection.add.call_args.kwargs + assert kwargs["ids"] == [item_id] + assert kwargs["embeddings"] == [[0.1, 0.2, 0.3]] + assert kwargs["documents"] == ["hi"] + # payload is stored directly in the JSON metadata column (nested included) + meta = kwargs["metadatas"][0] + assert meta == {"memory": "hi", "status": "activated"} + + +def test_search_scores_and_payload(vec_db, collection): + item_id = str(uuid.uuid4()) + collection.query.return_value = { + "ids": [[item_id]], + "metadatas": [[{"memory": "hello", "metadata": {"status": "activated"}}]], + "embeddings": [[[0.1, 0.2, 0.3]]], + "distances": [[0.1]], + } + results = vec_db.search([0.1, 0.2, 0.3], top_k=1) + assert len(results) == 1 + assert isinstance(results[0], VecDBItem) + # cosine distance 0.1 -> similarity 0.9 (higher is more similar) + assert results[0].score == pytest.approx(0.9) + # nested payload is returned verbatim + assert results[0].payload == {"memory": "hello", "metadata": {"status": "activated"}} + + +def test_get_by_id(vec_db, collection): + item_id = str(uuid.uuid4()) + collection.get.return_value = { + "ids": [item_id], + "metadatas": [{"memory": "x"}], + "embeddings": [[0.1, 0.2, 0.3]], + } + result = vec_db.get_by_id(item_id) + assert isinstance(result, VecDBItem) + assert result.vector == [0.1, 0.2, 0.3] + assert result.payload == {"memory": "x"} + + +def test_get_by_id_missing(vec_db, collection): + collection.get.return_value = {"ids": [], "metadatas": [], "embeddings": []} + assert vec_db.get_by_id(str(uuid.uuid4())) is None + + +def test_get_all_paginates(vec_db, collection): + first = str(uuid.uuid4()) + second = str(uuid.uuid4()) + collection.get.side_effect = [ + {"ids": [first], "metadatas": [{"memory": "a"}], "embeddings": [[0.1, 0.2, 0.3]]}, + {"ids": [second], "metadatas": [{"memory": "b"}], "embeddings": [[0.4, 0.5, 0.6]]}, + {"ids": [], "metadatas": [], "embeddings": []}, + ] + results = vec_db.get_all(scroll_limit=1) + assert [r.id for r in results] == [first, second] + + +def test_count_no_filter(vec_db, collection): + collection.count.return_value = 7 + assert vec_db.count() == 7 + + +def test_count_with_filter(vec_db, collection): + collection.get.return_value = { + "ids": [str(uuid.uuid4())], + "metadatas": [{"status": "activated"}], + "embeddings": [[0.1, 0.2, 0.3]], + } + assert vec_db.count({"status": "activated"}) == 1 + + +def test_update_vector(vec_db, collection): + item_id = str(uuid.uuid4()) + vec_db.update(item_id, {"id": item_id, "vector": [0.4, 0.5, 0.6], "payload": {"memory": "y"}}) + collection.update.assert_called_once() + kwargs = collection.update.call_args.kwargs + assert kwargs["embeddings"] == [0.4, 0.5, 0.6] + + +def test_update_metadata_only(vec_db, collection): + item_id = str(uuid.uuid4()) + vec_db.update(item_id, {"id": item_id, "payload": {"status": "archived"}}) + collection.update.assert_called_once() + kwargs = collection.update.call_args.kwargs + assert "embeddings" not in kwargs + assert kwargs["metadatas"]["status"] == "archived" + + +def test_delete(vec_db, collection): + ids = [str(uuid.uuid4()), str(uuid.uuid4())] + vec_db.delete(ids) + collection.delete.assert_called_once_with(ids=ids) + + +def test_delete_collection(vec_db): + vec_db.delete_collection("some_collection") + vec_db.client.delete_collection.assert_called_once_with("some_collection") + + +def test_euclidean_score_ordering(config, fake_pyseekdb, collection): + config.config.distance_metric = "euclidean" + client = fake_pyseekdb.Client.return_value + client.has_collection.return_value = False + client.create_collection.return_value = collection + db = VecDBFactory.from_config(config) + # For l2, smaller distance must yield a higher score. + assert db._distance_to_score(0.2) > db._distance_to_score(0.5) From c7111cec11a348ef2ffee67ae92e2f65390634e9 Mon Sep 17 00:00:00 2001 From: Even Date: Mon, 20 Jul 2026 19:40:36 +0800 Subject: [PATCH 2/2] chore: update pyseekdb version constraints and enhance OceanBase connection handling - Updated pyseekdb version constraints in pyproject.toml to restrict to <1.5.0. - Increased default embedding dimension in APIConfig from 768 to 1024. - Improved connection handling in OceanBaseGraphDB and OceanBaseVecDB to ensure better resource management and error handling. - Added validation for table prefix length in OceanBaseGraphDB to prevent identifier overflow. - Enhanced logging for empty password configurations in OceanBaseGraphDB. - Updated tests to reflect changes in search behavior and connection management. --- pyproject.toml | 4 +- src/memos/api/config.py | 2 +- src/memos/configs/graph_db.py | 17 +++++++ src/memos/graph_dbs/oceanbase.py | 75 ++++++++++++++++++++----------- src/memos/vec_dbs/oceanbase.py | 12 +++-- tests/graph_dbs/test_oceanbase.py | 5 ++- 6 files changed, 81 insertions(+), 34 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0f8811db3..c9c320f05 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -114,7 +114,7 @@ tavily = [ # OceanBase / seekdb unified vector + graph provider ob-mem = [ - "pyseekdb (>=1.4.0,<2.0.0)", # Unified client for seekdb / OceanBase (Collection + raw SQL) + "pyseekdb (>=1.4.0,<1.5.0)", # Unified client for seekdb / OceanBase (Collection + raw SQL) ] # All optional dependencies @@ -145,7 +145,7 @@ all = [ "rake-nltk (>=1.0.6,<1.1.0)", "alibabacloud-oss-v2 (>=1.2.2,<1.2.3)", "tavily-python (>=0.5.0,<1.0.0)", - "pyseekdb (>=1.4.0,<2.0.0)", + "pyseekdb (>=1.4.0,<1.5.0)", # Uncategorized dependencies ] diff --git a/src/memos/api/config.py b/src/memos/api/config.py index 1b517f875..a173cc44a 100644 --- a/src/memos/api/config.py +++ b/src/memos/api/config.py @@ -916,7 +916,7 @@ def get_oceanbase_config(user_id: str | None = None) -> dict[str, Any]: "table_prefix": os.getenv("OCEANBASE_TABLE_PREFIX", "memos_graph"), "user_name": user_name, "use_multi_db": False, - "embedding_dimension": int(os.getenv("EMBEDDING_DIMENSION", "768")), + "embedding_dimension": int(os.getenv("EMBEDDING_DIMENSION", "1024")), "maxconn": int(os.getenv("OCEANBASE_MAX_CONN", "20")), } diff --git a/src/memos/configs/graph_db.py b/src/memos/configs/graph_db.py index d41d7b72d..24aa9dd9d 100644 --- a/src/memos/configs/graph_db.py +++ b/src/memos/configs/graph_db.py @@ -6,6 +6,10 @@ from memos.configs.base import BaseConfig from memos.configs.vec_db import VectorDBConfigFactory +from memos.log import get_logger + + +logger = get_logger(__name__) class BaseGraphDBConfig(BaseConfig): @@ -285,10 +289,18 @@ class OceanBaseGraphDBConfig(BaseConfig): @classmethod def validate_table_prefix(cls, value: str) -> str: """Reject prefixes that could break out of the identifier when built into DDL.""" + # Longest derived identifier is idx_{prefix}_nodes_embedding (prefix + 20 chars); + # MySQL / OceanBase cap identifiers at 64 characters. + max_prefix_len = 44 if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", value): raise ValueError( "`table_prefix` must match [A-Za-z_][A-Za-z0-9_]* (letters, digits, underscore)" ) + if len(value) > max_prefix_len: + raise ValueError( + f"`table_prefix` must be at most {max_prefix_len} characters to keep derived " + "index names within MySQL/OceanBase's 64-character identifier limit" + ) return value @model_validator(mode="after") @@ -298,6 +310,11 @@ def validate_config(self): raise ValueError("`db_name` must be provided") if not self.use_multi_db and not self.user_name: raise ValueError("In single-database mode, `user_name` must be provided") + if not self.password: + logger.warning( + "OceanBase graph DB is configured with an empty password; " + "make sure this is intentional (e.g. a local dev instance)." + ) return self diff --git a/src/memos/graph_dbs/oceanbase.py b/src/memos/graph_dbs/oceanbase.py index 4eef0a562..647353c5d 100644 --- a/src/memos/graph_dbs/oceanbase.py +++ b/src/memos/graph_dbs/oceanbase.py @@ -55,16 +55,30 @@ def __init__(self, connect_fn, maxconn: int): def acquire(self): if self._closed: raise GraphDBError("Connection pool is closed") + # Decide whether to create a new connection while holding the lock, but run + # the (potentially slow) connect outside the lock so it never blocks peers. with self._lock: - if self._created < self._maxconn and self._idle.empty(): + should_create = self._created < self._maxconn and self._idle.empty() + if should_create: self._created += 1 - try: - return self._connect_fn() - except Exception: + if should_create: + try: + return self._connect_fn() + except Exception: + with self._lock: self._created -= 1 - raise + raise return self._idle.get() + def discard(self, conn) -> None: + """Permanently drop a (likely broken) connection and free its pool slot.""" + if conn is not None: + with suppress(Exception): + conn.close() + with self._lock: + if self._created > 0: + self._created -= 1 + def release(self, conn) -> None: if conn is None: return @@ -106,10 +120,13 @@ def _normalize_dt(value: Any) -> str: """Normalize an ISO datetime string to a MySQL-compatible ``DATETIME`` literal.""" if value is None: value = datetime.utcnow().isoformat() - text = str(value).replace("T", " ") - # Drop timezone suffix / fractional beyond seconds tolerated by DATETIME. - if "+" in text: - text = text.split("+", 1)[0] + text = str(value).replace("T", " ").strip() + # Strip the timezone suffix (+HH:MM / -HH:MM / Z) without touching the date's + # own hyphens: only the time part after the first space can carry a tz. + if " " in text: + date_part, time_part = text.split(" ", 1) + time_part = re.split(r"[Z+\-]", time_part, maxsplit=1)[0].strip() + text = f"{date_part} {time_part}" return text.strip() @@ -167,21 +184,26 @@ def _ensure_open(self) -> None: if self._closed: raise GraphDBError("OceanBaseGraphDB connection is closed") - def _healthy(self, conn): - """Ping the borrowed connection, replacing it if the socket is dead.""" + def _acquire_live(self): + """Get a live connection from the pool. + + If the borrowed connection's socket is dead, discard it through the pool + (so ``_created`` stays accurate) and acquire a fresh one, avoiding the + counter drift that a raw out-of-band reconnect would cause. + """ + conn = self._pool.acquire() try: conn.ping(reconnect=True) return conn except Exception: - with suppress(Exception): - conn.close() - return self._pymysql.connect(**self._conn_kwargs) + self._pool.discard(conn) + return self._pool.acquire() @contextmanager def _borrow(self): """Borrow a single connection for one auto-committed operation.""" self._ensure_open() - conn = self._healthy(self._pool.acquire()) + conn = self._acquire_live() try: yield conn finally: @@ -191,7 +213,7 @@ def _borrow(self): def _transaction(self): """Borrow a connection and run several statements as one atomic unit.""" self._ensure_open() - conn = self._healthy(self._pool.acquire()) + conn = self._acquire_live() cur = None try: cur = conn.cursor() @@ -1047,10 +1069,20 @@ def get_grouped_counts( params: dict[str, Any] | None = None, user_name: str | None = None, ) -> list[dict[str, Any]]: - """Count nodes grouped by the given JSON property fields.""" + """Count nodes grouped by the given JSON property fields. + + ``where_clause`` is intentionally not supported: splicing a raw SQL + string would be an injection vector, and no current caller uses it. + Express additional filters through the structured filter helpers instead. + """ user_name = user_name or self.user_name if not group_fields: raise ValueError("group_fields cannot be empty") + if where_clause: + raise ValueError( + "get_grouped_counts() does not support a raw `where_clause`; " + "use structured filters instead" + ) for field in group_fields: if not self._is_safe_field_name(field): raise ValueError(f"Invalid group field: {field}") @@ -1068,15 +1100,6 @@ def get_grouped_counts( conditions = ["user_name = %s"] query_params: list[Any] = [user_name] - if where_clause: - where_clause = where_clause.strip() - if where_clause.upper().startswith("WHERE"): - where_clause = where_clause[5:].strip() - if where_clause: - conditions.append(where_clause) - if params: - query_params.extend(params.values()) - where_sql = " AND ".join(conditions) rows = self._query( diff --git a/src/memos/vec_dbs/oceanbase.py b/src/memos/vec_dbs/oceanbase.py index f3884328e..3a4cfe8dc 100644 --- a/src/memos/vec_dbs/oceanbase.py +++ b/src/memos/vec_dbs/oceanbase.py @@ -43,7 +43,13 @@ def __init__(self, config: OceanBaseVecDBConfig): import pyseekdb self.config = config - self._distance = _DISTANCE_MAP.get(config.distance_metric or "cosine", "cosine") + metric = config.distance_metric or "cosine" + if metric not in _DISTANCE_MAP: + raise ValueError( + f"Unsupported distance_metric '{metric}'. " + f"Valid options are: {list(_DISTANCE_MAP.keys())}" + ) + self._distance = _DISTANCE_MAP[metric] self.client = pyseekdb.Client( host=config.host, @@ -143,7 +149,7 @@ def search( query_embeddings=query_vector, n_results=top_k, where=filter or None, - include=["documents", "metadatas", "embeddings"], + include=["metadatas", "embeddings", "distances"], ) ids = (response.get("ids") or [[]])[0] metadatas = (response.get("metadatas") or [[]])[0] @@ -278,7 +284,7 @@ def update(self, id: str, data: VecDBItem | dict[str, Any]) -> None: """Update an item in the collection.""" item = self._normalize_item(data) metadata = item.payload - if item.vector: + if item.vector is not None: self.collection.update( ids=id, embeddings=item.vector, diff --git a/tests/graph_dbs/test_oceanbase.py b/tests/graph_dbs/test_oceanbase.py index c4a58a194..d8d3bbcc6 100644 --- a/tests/graph_dbs/test_oceanbase.py +++ b/tests/graph_dbs/test_oceanbase.py @@ -97,9 +97,10 @@ def test_get_neighbors(graph_db, cursor): def test_search_by_embedding_threshold(graph_db, cursor): - cursor.fetchall.return_value = [("n1", 0.9), ("n2", 0.4)] + # score == threshold (0.5) must be included (implementation uses score >= threshold). + cursor.fetchall.return_value = [("n1", 0.9), ("n3", 0.5), ("n2", 0.4)] results = graph_db.search_by_embedding([0.1, 0.2, 0.3], top_k=5, threshold=0.5) - assert results == [{"id": "n1", "score": 0.9}] + assert results == [{"id": "n1", "score": 0.9}, {"id": "n3", "score": 0.5}] def test_get_path_cte(graph_db, cursor):