GML 2193 Add short-term and long-term agent memory in the KB graph (replace last-4 SQLite history)

common/config.py

@@ -368,6 +368,40 @@ def get_mcp_servers(graphname=None):
     return resolve_mcp_servers(global_list, graph_list)
 
 
+def get_memory_config(graphname=None):
+    """Return validated global + per-graph agent-memory settings.
+
+    Secrets stay in the environment; non-secret tuning lives in the existing
+    server config. Per-graph values override global values.
+    """
+    value = dict(server_config.get("memory_config") or {})
+    value.update(_load_graph_config(graphname).get("memory_config") or {})
+    database_url = os.getenv("GRAPHRAG_MEMORY_DATABASE_URL") or value.get(
+        "database_url"
+    )
+    env_enabled = os.getenv("GRAPHRAG_MEMORY_ENABLED")
+    enabled = (
+        env_enabled.lower() in {"1", "true", "yes"}
+        if env_enabled is not None
+        else bool(value.get("enabled", database_url))
+    )
+    context = int(value.get("stm_context_tokens", 6000))
+    if context <= 0:
+        raise ValueError("stm_context_tokens must be greater than 0")
+    return {
+        **value,
+        "enabled": enabled and bool(database_url),
+        "database_url": database_url,
+        "stm_context_tokens": context,
+        "ltm_top_k": max(1, int(value.get("ltm_top_k", 8))),
+        "ltm_score_threshold": float(value.get("ltm_score_threshold", 0.2)),
+        "fact_extract_async": bool(value.get("fact_extract_async", True)),
+        "max_facts_per_scope": max(1, int(value.get("max_facts_per_scope", 1000))),
+        "retry_limit": max(1, int(value.get("retry_limit", 8))),
+        "embedding_dimension": int(value.get("embedding_dimension", embedding_dimension)),
+    }
+
+
 PATH_PREFIX = os.getenv("PATH_PREFIX", "")
 PRODUCTION = os.getenv("PRODUCTION", "false").lower() == "true"
 

common/db/schema_utils.py

@@ -74,6 +74,7 @@
     "Content",
     "Community",
     "Image",
+    "MemoryEmbedding",
 })
 
 
@@ -221,6 +222,11 @@ def is_structural_type(name: str) -> bool:
     return is_reserved_word(name)
 
 
+def is_hidden_memory_type(name: str) -> bool:
+    """Internal memory index types never belong in agent-visible schema."""
+    return bool(name) and name.casefold() == "memoryembedding"
+
+
 # -----------------------------------------------------------------------------
 # Canonical proposal dataclass
 # -----------------------------------------------------------------------------
@@ -1251,6 +1257,7 @@ def render_schema_rep(conn, exclude_structural: bool = False) -> AllowedSchema:
         all_verts = conn.getVertexTypes() or []
     except Exception:
         all_verts = []
+    all_verts = [v for v in all_verts if not is_hidden_memory_type(v)]
     domain_verts = (
         [v for v in all_verts if not is_structural_type(v)]
         if exclude_structural else list(all_verts)
@@ -1403,6 +1410,7 @@ async def render_schema_rep_async(
         all_verts = await conn.getVertexTypes() or []
     except Exception:
         all_verts = []
+    all_verts = [v for v in all_verts if not is_hidden_memory_type(v)]
     domain_verts = (
         [v for v in all_verts if not is_structural_type(v)]
         if exclude_structural else list(all_verts)

common/embeddings/tigergraph_embedding_store.py

@@ -15,6 +15,7 @@
 import logging
 import traceback
 import json
+from pathlib import Path
 from time import sleep, time
 from collections import defaultdict
 from typing import Iterable, List, Optional, Tuple
@@ -34,6 +35,22 @@
 logger = logging.getLogger(__name__)
 
 
+def _extract_gsql_job(source: str, job_name: str) -> str:
+    marker = f"CREATE SCHEMA_CHANGE JOB {job_name}"
+    start = source.find(marker)
+    if start < 0:
+        raise RuntimeError(f"Missing schema-change job {job_name}")
+    depth = 0
+    for index, char in enumerate(source[start:], start):
+        if char == "{":
+            depth += 1
+        elif char == "}":
+            depth -= 1
+            if depth == 0:
+                return source[start : index + 1]
+    raise RuntimeError(f"Unclosed schema-change job {job_name}")
+
+
 class TigerGraphEmbeddingStore(EmbeddingStore):
     def __init__(
         self,
@@ -46,6 +63,7 @@ def __init__(
         self.support_ai_instance = support_ai_instance
         self.default_vector_attribute = default_vector_attribute
         self.vector_attr_cache = {}
+        self._memory_index_ready = False
 
         if isinstance(conn.apiToken, tuple):
             token = conn.apiToken[0]
@@ -152,9 +170,130 @@ def install_vector_queries(self):
         else:
             logger.info(f"All vector queries already installed, skipping.")
 
+    def ensure_memory_index(self, dimension: int = 1536) -> None:
+        """Create the internal vector-only memory type and its queries once."""
+        if self._memory_index_ready:
+            return
+        graphname = self.conn.graphname
+        if not graphname or graphname == "MyGraph":
+            raise RuntimeError("Memory embeddings require a bound graph")
+        schema = self.conn.gsql(f"USE GRAPH {graphname}\n ls")
+        has_vertex = "- VERTEX MemoryEmbedding" in schema
+        # Vector attributes are listed under "Vector Embeddings:", not
+        # next to the VERTEX line, so do not slice the vertex block.
+        self.refreshvector_attr_cache()
+        has_vector = self.has_vector_attribute("MemoryEmbedding", "embedding")
+        results = []
+        # Vertex must exist before the ALTER VECTOR job can be created.
+        if not has_vertex or not has_vector:
+            body = Path(
+                "common/gsql/supportai/SupportAI_Schema_Memory_Vector.gsql"
+            ).read_text().replace("dimension=1536", f"dimension={dimension}")
+            if not has_vertex:
+                results.append(self._run_memory_schema_job(
+                    graphname,
+                    "add_graphrag_memory_vertex",
+                    _extract_gsql_job(body, "add_graphrag_memory_vertex"),
+                ))
+            if not has_vector:
+                results.append(self._run_memory_schema_job(
+                    graphname,
+                    "add_graphrag_memory_vector",
+                    _extract_gsql_job(body, "add_graphrag_memory_vector"),
+                ))
+        if results:
+            self.vector_attr_cache = {}
+            schema = self.conn.gsql(f"USE GRAPH {graphname}\n ls")
+            self.refreshvector_attr_cache()
+            if (
+                "- VERTEX MemoryEmbedding" not in schema
+                or not self.has_vector_attribute("MemoryEmbedding", "embedding")
+            ):
+                raise RuntimeError(
+                    "Memory vector schema bootstrap failed: " + "\n".join(results)
+                )
+        installed = [
+            endpoint.split("/")[-1]
+            for endpoint in self.conn.getEndpoints(dynamic=True)
+            if f"/{graphname}/" in endpoint
+        ]
+        missing = []
+        for name in ("Memory_Search_Embeddings", "Memory_Delete_Embedding"):
+            if name not in installed:
+                missing.append(name)
+                body = Path(f"common/gsql/memory/{name}.gsql").read_text()
+                self.conn.gsql(f"USE GRAPH {graphname}\nBEGIN\n{body}\nEND")
+        if missing:
+            self.conn.gsql(
+                f"USE GRAPH {graphname}\nINSTALL QUERY " + ", ".join(missing)
+            )
+        self._memory_index_ready = True
+
+    def _run_memory_schema_job(
+        self, graphname: str, job_name: str, body: str
+    ) -> str:
+        create_result = self.conn.gsql(f"USE GRAPH {graphname}\n{body}")
+        run_result = self.conn.gsql(
+            f"USE GRAPH {graphname}\nRUN SCHEMA_CHANGE JOB {job_name}"
+        )
+        try:
+            self.conn.gsql(f"USE GRAPH {graphname}\nDROP JOB {job_name}")
+        except Exception:
+            logger.warning("Could not drop memory schema-change job %s", job_name)
+        return f"{create_result}\n{run_result}"
+
+    def upsert_memory_embedding(self, index_id: str, vector: List[float]) -> None:
+        """Persist a mem0-provided vector without embedding the text again."""
+        self.ensure_memory_index(len(vector))
+        result = self.conn.upsertData(
+            json.dumps(
+                {
+                    "vertices": {
+                        "MemoryEmbedding": {
+                            index_id: {"embedding": {"value": vector}}
+                        }
+                    }
+                }
+            )
+        )
+        if not result:
+            raise RuntimeError(f"Failed to upsert memory embedding {index_id}")
+
+    def search_memory_embeddings(
+        self,
+        scope_pattern: str,
+        query_vector: List[float],
+        top_k: int,
+    ) -> List[Tuple[str, float]]:
+        self.ensure_memory_index(len(query_vector))
+        result = self.conn.runInstalledQuery(
+            "Memory_Search_Embeddings",
+            {
+                "scope_pattern": scope_pattern,
+                "query_vector": query_vector,
+                "top_k": top_k,
+            },
+        )
+        hits = []
+        for block in result or []:
+            for item in block.get("hits", []):
+                value = item.get("value", item)
+                if value.get("index_id"):
+                    hits.append(
+                        (str(value["index_id"]), float(value.get("score", 0)))
+                    )
+        return sorted(hits, key=lambda pair: pair[1], reverse=True)
+
+    def delete_memory_embedding(self, index_id: str) -> None:
+        self.ensure_memory_index()
+        self.conn.runInstalledQuery(
+            "Memory_Delete_Embedding", {"index_id": index_id}
+        )
+
     def set_graphname(self, graphname):
         self.conn.graphname = graphname
         self.vector_attr_cache = {}
+        self._memory_index_ready = False
         if self.conn.apiToken or self.conn.jwtToken:
             self.conn.getToken()
         # Re-verify GDS presence on every graphname switch. Cheap when
@@ -188,6 +327,7 @@ def set_connection(self, conn):
              )
 
         self.vector_attr_cache = {}
+        self._memory_index_ready = False
         self.install_vector_queries()
 
     def refreshvector_attr_cache(self):

common/gsql/memory/Memory_Delete_Embedding.gsql

@@ -0,0 +1,9 @@
+/*
+ * Copyright (c) 2024-2026 TigerGraph, Inc.
+ */
+
+CREATE OR REPLACE DISTRIBUTED QUERY Memory_Delete_Embedding(STRING index_id) {
+  targets = SELECT s FROM MemoryEmbedding:s WHERE s.id == index_id
+            POST-ACCUM DELETE(s);
+  PRINT index_id;
+}

common/gsql/memory/Memory_Search_Embeddings.gsql

@@ -0,0 +1,34 @@
+/*
+ * Copyright (c) 2024-2026 TigerGraph, Inc.
+ * Returns index IDs only. Fact text and ownership remain in PostgreSQL.
+ */
+
+CREATE OR REPLACE DISTRIBUTED QUERY Memory_Search_Embeddings(
+  STRING scope_pattern,
+  LIST<FLOAT> query_vector,
+  INT top_k = 8
+) {
+  TYPEDEF TUPLE<STRING index_id, FLOAT score> MemoryHit;
+  HeapAccum<MemoryHit>(top_k, score DESC) @@hits;
+
+  candidates = SELECT m
+               FROM MemoryEmbedding:m
+               WHERE m.id LIKE scope_pattern
+                 AND m.embedding.size() > 0;
+
+  nearest = vectorSearch(
+    {MemoryEmbedding.embedding},
+    query_vector,
+    top_k,
+    {candidate_set: candidates}
+  );
+
+  ranked = SELECT m
+           FROM nearest:m
+           POST-ACCUM @@hits += MemoryHit(
+             m.id,
+             1 - gds.vector.distance(query_vector, m.embedding, "COSINE")
+           );
+
+  PRINT @@hits AS hits;
+}

common/gsql/supportai/SupportAI_Schema_Memory_Vector.gsql

@@ -0,0 +1,14 @@
+/*
+ * Copyright (c) 2024-2026 TigerGraph, Inc.
+ * Internal memory index: canonical payloads remain in PostgreSQL.
+ * Vertex and vector attribute are separate jobs — TigerGraph cannot
+ * ADD a vertex and ALTER it in the same schema-change job.
+ */
+
+CREATE SCHEMA_CHANGE JOB add_graphrag_memory_vertex {
+  ADD VERTEX MemoryEmbedding(PRIMARY_ID id STRING) WITH primary_id_as_attribute="true";
+}
+
+CREATE SCHEMA_CHANGE JOB add_graphrag_memory_vector {
+  ALTER VERTEX MemoryEmbedding ADD VECTOR ATTRIBUTE embedding(dimension=1536, metric="cosine");
+}

common/memory/__init__.py

@@ -0,0 +1,11 @@
+"""PostgreSQL-backed short- and long-term agent memory."""
+
+__all__ = ["MemoryService"]
+
+
+def __getattr__(name):
+    if name == "MemoryService":
+        from common.memory.service import MemoryService
+
+        return MemoryService
+    raise AttributeError(name)

common/memory/compaction.py

@@ -0,0 +1,53 @@
+"""Token-based STM compaction helpers."""
+
+from __future__ import annotations
+
+try:
+    import tiktoken
+except ImportError:  # pragma: no cover - production requirements include it
+    tiktoken = None
+
+
+def token_count(text: str) -> int:
+    if not text:
+        return 0
+    try:
+        if tiktoken is None:
+            raise RuntimeError("tokenizer unavailable")
+        return len(tiktoken.get_encoding("cl100k_base").encode(text))
+    except Exception:
+        return max(1, len(text) // 4)
+
+
+def message_text(messages) -> str:
+    return "\n".join(
+        f"{item.role}: {item.content}"
+        if hasattr(item, "role")
+        else f"{item.get('role')}: {item.get('content')}"
+        for item in messages
+    )
+
+
+KEEP_NEWEST_TURNS = 2
+
+
+def select_compaction_range(summary: str, messages, *, budget_tokens: int):
+    """Summarize oldest turns once STM exceeds the agent budget."""
+    total = token_count(summary) + token_count(message_text(messages))
+    if total <= budget_tokens:
+        return []
+    keep_count = KEEP_NEWEST_TURNS * 2
+    eligible = messages[:-keep_count] if len(messages) > keep_count else []
+    selected = []
+    remaining = total
+    for item in eligible:
+        selected.append(item)
+        rendered = (
+            f"{item.role}: {item.content}"
+            if hasattr(item, "role")
+            else f"{item.get('role')}: {item.get('content')}"
+        )
+        remaining -= token_count(rendered)
+        if remaining <= budget_tokens:
+            break
+    return selected