diff --git a/common/config.py b/common/config.py index d6144ca5..6e57a020 100644 --- a/common/config.py +++ b/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" diff --git a/common/db/schema_utils.py b/common/db/schema_utils.py index 748171ef..382f1b89 100644 --- a/common/db/schema_utils.py +++ b/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) diff --git a/common/embeddings/tigergraph_embedding_store.py b/common/embeddings/tigergraph_embedding_store.py index 47dbed8e..e98826a0 100644 --- a/common/embeddings/tigergraph_embedding_store.py +++ b/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): diff --git a/common/gsql/memory/Memory_Delete_Embedding.gsql b/common/gsql/memory/Memory_Delete_Embedding.gsql new file mode 100644 index 00000000..cdd181ef --- /dev/null +++ b/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; +} diff --git a/common/gsql/memory/Memory_Search_Embeddings.gsql b/common/gsql/memory/Memory_Search_Embeddings.gsql new file mode 100644 index 00000000..3b85ef03 --- /dev/null +++ b/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 query_vector, + INT top_k = 8 +) { + TYPEDEF TUPLE MemoryHit; + HeapAccum(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; +} diff --git a/common/gsql/supportai/SupportAI_Schema_Memory_Vector.gsql b/common/gsql/supportai/SupportAI_Schema_Memory_Vector.gsql new file mode 100644 index 00000000..dbbc44ef --- /dev/null +++ b/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"); +} diff --git a/common/memory/__init__.py b/common/memory/__init__.py new file mode 100644 index 00000000..8c8b4404 --- /dev/null +++ b/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) diff --git a/common/memory/compaction.py b/common/memory/compaction.py new file mode 100644 index 00000000..8683e97e --- /dev/null +++ b/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 diff --git a/common/memory/mem0_adapter.py b/common/memory/mem0_adapter.py new file mode 100644 index 00000000..229a1abc --- /dev/null +++ b/common/memory/mem0_adapter.py @@ -0,0 +1,222 @@ +# Copyright (c) 2024-2026 TigerGraph, Inc. +"""Adapters that keep mem0 payloads in PostgreSQL and vectors in TigerGraph.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Any +from urllib.parse import unquote + +from common.memory.store import PostgresMemoryStore, scope_hash + +os.environ.setdefault("MEM0_TELEMETRY", "false") + +try: + from mem0.vector_stores.base import VectorStoreBase +except ImportError: # Allows the application to start when memory is disabled. + class VectorStoreBase: # type: ignore[no-redef] + pass + + +@dataclass +class MemoryVectorResult: + id: str + score: float + payload: dict + + +def split_user_id(value: str) -> tuple[str, str]: + try: + graphname, username = value.split("\0", 1) + except ValueError as exc: + raise ValueError("mem0 user_id must include graph and username") from exc + if not graphname or not username: + raise ValueError("mem0 user_id contains an empty scope component") + return username, graphname + + +class PostgresHistoryAdapter: + """Drop-in replacement for mem0's hardcoded SQLiteManager.""" + + def __init__(self, store: PostgresMemoryStore, message_limit: int = 10): + self.store = store + self.message_limit = message_limit + + def save_messages(self, messages: list[dict], session_scope: str) -> None: + # The chat path synchronously persists the same messages before mem0 + # runs. Keeping another rolling-message table would duplicate STM. + return None + + def get_last_messages(self, session_scope: str, limit: int = 10) -> list[dict]: + scope = {} + for part in session_scope.split("&"): + key, separator, value = part.partition("=") + if separator: + scope[key] = unquote(value) + user_id, run_id = scope.get("user_id"), scope.get("run_id") + if not user_id or not run_id: + return [] + username, graphname = split_user_id(user_id) + messages = self.store.get_messages(username, graphname, run_id) + return [ + { + "role": item.get("role"), + "content": item.get("content"), + "created_at": str(item.get("create_ts") or ""), + } + for item in messages[-min(limit, self.message_limit) :] + ] + + def add_history( + self, + memory_id, + old_memory, + new_memory, + event, + *, + created_at=None, + updated_at=None, + is_deleted=0, + actor_id=None, + role=None, + ) -> None: + self.store.add_history(locals()) + + def batch_add_history(self, records: list[dict]) -> None: + for record in records: + self.store.add_history(record) + + def get_history(self, memory_id: str) -> list[dict]: + return self.store.get_history(memory_id) + + def reset(self) -> None: + # Global reset is deliberately unsupported in a multi-tenant database. + raise RuntimeError("Global mem0 history reset is disabled") + + def close(self) -> None: + return None + + +class SplitMemoryVectorStore(VectorStoreBase): + """mem0 VectorStoreBase with PostgreSQL payloads and TigerGraph vectors.""" + + def __init__( + self, + store: PostgresMemoryStore, + embedding_store_resolver: Any, + collection_name: str = "graphrag_memory", + ): + self.store = store + self.embedding_store_resolver = embedding_store_resolver + self.collection_name = collection_name + + @staticmethod + def _scope(payload_or_filters: dict) -> tuple[str, str, str | None]: + user_id = payload_or_filters.get("user_id") + if not user_id: + raise ValueError("All memory operations require user_id") + username, graphname = split_user_id(str(user_id)) + run_id = payload_or_filters.get("run_id") + return username, graphname, str(run_id) if run_id else None + + def create_col(self, name, vector_size, distance): + return None + + def insert(self, vectors, payloads=None, ids=None): + payloads, ids = payloads or [], ids or [] + for vector, payload, fact_id in zip(vectors, payloads, ids): + username, graphname, run_id = self._scope(payload) + row = self.store.put_fact( + str(fact_id), + username, + graphname, + run_id, + str(payload.get("data") or ""), + dict(payload), + ) + self.embedding_store_resolver(graphname).upsert_memory_embedding( + row.index_id, vector + ) + self.store.mark_fact_indexed(row.fact_id) + + def search(self, query, vectors, top_k=5, filters=None): + username, graphname, run_id = self._scope(filters or {}) + hits = self.embedding_store_resolver(graphname).search_memory_embeddings( + f"{scope_hash(username, graphname)}:%", + vectors, + top_k=max(top_k * 2, top_k), + ) + score_by_id = {hit_id: score for hit_id, score in hits} + rows = self.store.facts_by_index_ids(username, graphname, score_by_id) + results = [] + for row in rows: + if run_id and str(row.payload.get("run_id") or "") != run_id: + continue + results.append( + MemoryVectorResult( + id=row.fact_id, + score=float(score_by_id[row.index_id]), + payload=dict(row.payload), + ) + ) + if len(results) >= top_k: + break + return results + + def delete(self, vector_id): + row = self.store.get_fact(str(vector_id)) + if row: + self.store.delete_fact(row.fact_id) + self.embedding_store_resolver( + row.graphname + ).delete_memory_embedding(row.index_id) + self.store.mark_fact_unindexed(row.fact_id) + + def update(self, vector_id, vector=None, payload=None): + existing = self.store.get_fact(str(vector_id)) + if not existing: + return + merged = dict(existing.payload) + if payload: + merged.update(payload) + row = self.store.put_fact( + existing.fact_id, + existing.username, + existing.graphname, + existing.conversation_id, + str(merged.get("data") or existing.text), + merged, + ) + if vector is not None: + self.embedding_store_resolver( + row.graphname + ).upsert_memory_embedding(row.index_id, vector) + self.store.mark_fact_indexed(row.fact_id) + + def get(self, vector_id): + row = self.store.get_fact(str(vector_id)) + if not row: + return None + return MemoryVectorResult(row.fact_id, 1.0, dict(row.payload)) + + def list(self, filters=None, top_k=None): + username, graphname, run_id = self._scope(filters or {}) + rows = self.store.list_facts(username, graphname, run_id) + if top_k is not None: + rows = rows[:top_k] + return [[ + MemoryVectorResult(row.fact_id, 1.0, dict(row.payload)) for row in rows + ]] + + def list_cols(self): + return ["memory_facts"] + + def delete_col(self): + raise RuntimeError("Global memory collection deletion is disabled") + + def col_info(self): + return {"name": "memory_facts"} + + def reset(self): + raise RuntimeError("Global memory reset is disabled") diff --git a/common/memory/models.py b/common/memory/models.py new file mode 100644 index 00000000..f4c5996d --- /dev/null +++ b/common/memory/models.py @@ -0,0 +1,157 @@ +# Copyright (c) 2024-2026 TigerGraph, Inc. +"""SQLAlchemy models for PostgreSQL-backed conversation memory.""" + +from __future__ import annotations + +import uuid +from datetime import datetime, timezone +from typing import Optional + +from sqlalchemy import ( + JSON, + Boolean, + DateTime, + Float, + Index, + Integer, + String, + Text, + UniqueConstraint, +) +from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column + + +def utcnow() -> datetime: + return datetime.now(timezone.utc) + + +class Base(DeclarativeBase): + pass + + +class Conversation(Base): + __tablename__ = "memory_conversations" + __table_args__ = ( + UniqueConstraint( + "username", "graphname", "conversation_id", name="uq_memory_conversation_scope" + ), + Index("ix_memory_conversation_owner", "username", "graphname", "updated_at"), + ) + + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + conversation_id: Mapped[str] = mapped_column(String(36), nullable=False) + username: Mapped[str] = mapped_column(String(256), nullable=False) + graphname: Mapped[str] = mapped_column(String(256), nullable=False) + name: Mapped[str] = mapped_column(String(512), default="", nullable=False) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), default=utcnow, onupdate=utcnow + ) + deleted_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True)) + + +class Message(Base): + __tablename__ = "memory_messages" + __table_args__ = ( + UniqueConstraint("message_id", name="uq_memory_message_id"), + Index("ix_memory_message_thread", "conversation_id", "sequence"), + ) + + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + conversation_id: Mapped[str] = mapped_column(String(36), nullable=False) + message_id: Mapped[str] = mapped_column(String(36), nullable=False) + parent_id: Mapped[Optional[str]] = mapped_column(String(36)) + sequence: Mapped[int] = mapped_column(Integer, nullable=False) + model: Mapped[Optional[str]] = mapped_column(String(256)) + content: Mapped[str] = mapped_column(Text, default="", nullable=False) + role: Mapped[str] = mapped_column(String(32), nullable=False) + response_time: Mapped[Optional[float]] = mapped_column(Float) + feedback: Mapped[Optional[int]] = mapped_column(Integer) + comment: Mapped[Optional[str]] = mapped_column(Text) + answered_question: Mapped[Optional[bool]] = mapped_column(Boolean) + response_type: Mapped[Optional[str]] = mapped_column(String(64)) + query_sources: Mapped[Optional[dict]] = mapped_column(JSON) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), default=utcnow, onupdate=utcnow + ) + deleted_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True)) + + +class ConversationSummary(Base): + __tablename__ = "memory_conversation_summaries" + + conversation_id: Mapped[str] = mapped_column(String(36), primary_key=True) + summary: Mapped[str] = mapped_column(Text, default="", nullable=False) + last_summarized_sequence: Mapped[int] = mapped_column(Integer, default=0, nullable=False) + source_tokens: Mapped[int] = mapped_column(Integer, default=0, nullable=False) + summary_tokens: Mapped[int] = mapped_column(Integer, default=0, nullable=False) + model: Mapped[Optional[str]] = mapped_column(String(256)) + version: Mapped[int] = mapped_column(Integer, default=1, nullable=False) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), default=utcnow, onupdate=utcnow + ) + + +class MemoryFact(Base): + __tablename__ = "memory_facts" + __table_args__ = ( + Index("ix_memory_fact_scope", "username", "graphname", "deleted_at"), + Index("ix_memory_fact_conversation", "conversation_id"), + ) + + fact_id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid.uuid4()) + ) + index_id: Mapped[str] = mapped_column(String(160), unique=True, nullable=False) + username: Mapped[str] = mapped_column(String(256), nullable=False) + graphname: Mapped[str] = mapped_column(String(256), nullable=False) + conversation_id: Mapped[Optional[str]] = mapped_column(String(36)) + text: Mapped[str] = mapped_column(Text, nullable=False) + payload: Mapped[dict] = mapped_column(JSON, default=dict, nullable=False) + indexed: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), default=utcnow, onupdate=utcnow + ) + expires_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True)) + deleted_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True)) + + +class MemoryHistory(Base): + __tablename__ = "memory_fact_history" + __table_args__ = (Index("ix_memory_history_fact", "memory_id", "created_at"),) + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid.uuid4()) + ) + memory_id: Mapped[str] = mapped_column(String(36), nullable=False) + old_memory: Mapped[Optional[str]] = mapped_column(Text) + new_memory: Mapped[Optional[str]] = mapped_column(Text) + event: Mapped[str] = mapped_column(String(32), nullable=False) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow) + updated_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True)) + is_deleted: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False) + actor_id: Mapped[Optional[str]] = mapped_column(String(256)) + role: Mapped[Optional[str]] = mapped_column(String(32)) + + +class MemoryOutbox(Base): + __tablename__ = "memory_outbox" + __table_args__ = ( + UniqueConstraint("operation", "fact_id", name="uq_memory_outbox_operation"), + Index("ix_memory_outbox_pending", "completed_at", "next_attempt_at"), + ) + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid.uuid4()) + ) + operation: Mapped[str] = mapped_column(String(16), nullable=False) + fact_id: Mapped[str] = mapped_column(String(36), nullable=False) + attempts: Mapped[int] = mapped_column(Integer, default=0, nullable=False) + last_error: Mapped[Optional[str]] = mapped_column(Text) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow) + next_attempt_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow) + completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True)) + + diff --git a/common/memory/runtime.py b/common/memory/runtime.py new file mode 100644 index 00000000..d96c3a48 --- /dev/null +++ b/common/memory/runtime.py @@ -0,0 +1,193 @@ +# Copyright (c) 2024-2026 TigerGraph, Inc. +"""Pinned in-process mem0 runtime using GraphRAG storage adapters.""" + +from __future__ import annotations + +import importlib.metadata + +from common.config import get_chat_config, get_embedding_service, get_llm_service +from common.memory.mem0_adapter import PostgresHistoryAdapter +from common.memory.store import PostgresMemoryStore, memory_user_id + +MEM0_VERSION = "2.1.0" + + +class GraphRAGMem0Embedder: + """Reuse GraphRAG's configured embedding client inside mem0.""" + + def __init__(self, config=None): + self.config = config + self.embedding_service = get_embedding_service() + + def embed(self, text: str, memory_action=None): + return self.embedding_service.embed_query(text) + + def embed_batch(self, texts, memory_action=None): + if not texts: + return [] + if isinstance(texts, str): + return [self.embed(texts, memory_action)] + return self.embedding_service.embed_documents(list(texts)) + + +class GraphRAGMem0LLM: + """mem0 LLMBase stand-in that calls GraphRAG's langchain chat model. + + mem0's bundled LangchainLLM imports the uninstalled ``langchain`` meta + package. This adapter uses the already-constructed GraphRAG model + (AFC-disabled for Gemini) and speaks mem0's generate_response contract. + """ + + def __init__(self, config=None): + self.config = config + self.langchain_model = getattr(config, "model", None) + if self.langchain_model is None: + raise ValueError("GraphRAG mem0 LLM requires the GraphRAG chat model") + + def generate_response( + self, + messages, + response_format=None, + tools=None, + tool_choice="auto", + **kwargs, + ): + langchain_messages = [] + for message in messages: + role = message["role"] + content = message["content"] + if role == "system": + langchain_messages.append(("system", content)) + elif role == "user": + langchain_messages.append(("human", content)) + elif role == "assistant": + langchain_messages.append(("ai", content)) + if not langchain_messages: + raise ValueError("No valid messages found in the messages list") + + model = self.langchain_model + if tools: + model = model.bind_tools(tools=tools, tool_choice=tool_choice) + response = model.invoke(langchain_messages, **kwargs) + if not tools: + return response.content + return { + "content": response.content, + "tool_calls": [ + {"name": call["name"], "arguments": call["args"]} + for call in getattr(response, "tool_calls", []) or [] + ], + } + + +class GraphRAGVectorConfig: + def __init__( + self, + *, + store, + embedding_store_resolver, + collection_name="graphrag_memory", + ): + self.store = store + self.embedding_store_resolver = embedding_store_resolver + self.collection_name = collection_name + + def model_dump(self): + return { + "store": self.store, + "embedding_store_resolver": self.embedding_store_resolver, + "collection_name": self.collection_name, + } + + +class Mem0Runtime: + """Uses mem0 public operations; only persistence is adapted.""" + + def __init__( + self, + graphname: str, + store: PostgresMemoryStore, + embedding_store_resolver, + ): + installed = importlib.metadata.version("mem0ai") + if installed != MEM0_VERSION: + raise RuntimeError( + f"Unsupported mem0ai {installed}; expected pinned {MEM0_VERSION}" + ) + from mem0.configs.base import MemoryConfig + from mem0.configs.llms.base import BaseLlmConfig + from mem0.memory.main import Memory + import mem0.memory.main as mem0_main + from mem0.utils.factory import EmbedderFactory, LlmFactory, VectorStoreFactory + + chat = get_chat_config(graphname) + # mem0's native gemini client enables AFC. Reuse GraphRAG's langchain + # wrapper so extract/search share the same AFC-disabled model. + # mem0's own "langchain" provider imports the uninstalled langchain + # meta package, so the factory slot is pointed at GraphRAGMem0LLM. + langchain_model = get_llm_service(chat).llm + LlmFactory.provider_to_class["langchain"] = ( + "common.memory.runtime.GraphRAGMem0LLM", + BaseLlmConfig, + ) + + config = MemoryConfig( + **{ + "version": "v1.1", + "history_db_path": ":memory:", + "vector_store": { + "provider": "qdrant", + "config": { + "collection_name": "graphrag_memory_bootstrap", + "path": "/tmp/graphrag-memory-bootstrap", + }, + }, + "llm": { + "provider": "langchain", + "config": {"model": langchain_model}, + }, + "embedder": {"provider": "openai", "config": {}}, + "custom_instructions": ( + "Store only durable user facts, preferences, decisions, and goals. " + "Never store credentials, secrets, API keys, transient greetings, " + "copied graph evidence, or unsupported claims made by the assistant." + ), + } + ) + VectorStoreFactory.provider_to_class["graphrag"] = ( + "common.memory.mem0_adapter.SplitMemoryVectorStore" + ) + EmbedderFactory.provider_to_class["graphrag"] = ( + "common.memory.runtime.GraphRAGMem0Embedder" + ) + config.vector_store.provider = "graphrag" + config.vector_store.config = GraphRAGVectorConfig( + store=store, + embedding_store_resolver=embedding_store_resolver, + ) + config.embedder.provider = "graphrag" + config.embedder.config = {} + # mem0 telemetry creates a second vector store; GraphRAG already has + # application-level metrics and must not duplicate the memory index. + mem0_main.MEM0_TELEMETRY = False + self.memory = Memory(config) + old_db = self.memory.db + self.memory.db = PostgresHistoryAdapter(store) + old_db.close() + self.graphname = graphname + + def add(self, messages: list[dict], username: str, conversation_id: str): + return self.memory.add( + messages, + user_id=memory_user_id(username, self.graphname), + run_id=conversation_id, + infer=True, + ) + + def search(self, query: str, username: str, *, top_k: int, threshold: float): + return self.memory.search( + query, + top_k=top_k, + threshold=threshold, + filters={"user_id": memory_user_id(username, self.graphname)}, + ) diff --git a/common/memory/service.py b/common/memory/service.py new file mode 100644 index 00000000..188e8b87 --- /dev/null +++ b/common/memory/service.py @@ -0,0 +1,241 @@ +# Copyright (c) 2024-2026 TigerGraph, Inc. +"""Single application service for transcript, STM compaction, and mem0 LTM.""" + +from __future__ import annotations + +import asyncio +import logging +import re +import threading +from collections import defaultdict + +from common.config import ( + get_chat_config, + get_embedding_service, + get_embedding_store, + get_llm_service, + get_memory_config, +) +from common.memory.compaction import message_text, select_compaction_range, token_count +from common.memory.runtime import Mem0Runtime +from common.memory.store import PostgresMemoryStore + +logger = logging.getLogger(__name__) + +_SECRET_PATTERN = re.compile( + r"(?:api[_ -]?key|password|secret|bearer\s+[A-Za-z0-9._-]+|" + r"-----BEGIN [A-Z ]*PRIVATE KEY-----)", + re.IGNORECASE, +) + +_stores: dict[str, PostgresMemoryStore] = {} +_runtimes: dict[tuple[str, str], Mem0Runtime] = {} +_global_lock = threading.Lock() +_compaction_locks: defaultdict[str, asyncio.Lock] = defaultdict(asyncio.Lock) + + +class MemoryService: + def __init__(self, graphname: str): + self.graphname = graphname + self.config = get_memory_config(graphname) + if not self.config["enabled"]: + raise RuntimeError("Agent memory is not enabled") + database_url = self.config["database_url"] + with _global_lock: + store = _stores.get(database_url) + if store is None: + store = PostgresMemoryStore(database_url) + _stores[database_url] = store + self.store = store + + @classmethod + def enabled(cls, graphname: str) -> bool: + return bool(get_memory_config(graphname)["enabled"]) + + def _runtime(self) -> Mem0Runtime: + key = (self.config["database_url"], self.graphname) + with _global_lock: + runtime = _runtimes.get(key) + if runtime is None: + runtime = Mem0Runtime( + self.graphname, self.store, get_embedding_store + ) + _runtimes[key] = runtime + return runtime + + def list_conversations(self, username: str) -> list[dict]: + return self.store.list_conversations(username, self.graphname) + + def get_messages(self, username: str, conversation_id: str) -> list[dict]: + return self.store.get_messages(username, self.graphname, conversation_id) + + def save_message(self, username: str, message: dict) -> dict: + return self.store.save_message(username, self.graphname, message) + + def feedback_messages(self, username: str, include_all: bool) -> list[dict]: + return self.store.feedback_messages( + username, self.graphname, include_all=include_all + ) + + def delete_conversation(self, username: str, conversation_id: str) -> None: + self.store.delete_conversation(username, self.graphname, conversation_id) + self.flush_outbox() + + def history_pairs(self, username: str, conversation_id: str) -> list[dict]: + messages = self.get_messages(username, conversation_id) + by_parent = { + item.get("parent_id"): item + for item in messages + if item.get("role") == "system" and item.get("parent_id") + } + return [ + { + "query": item.get("content", ""), + "response": by_parent[item["message_id"]].get("content", ""), + "create_ts": by_parent[item["message_id"]].get("create_ts"), + "update_ts": by_parent[item["message_id"]].get("update_ts"), + } + for item in messages + if item.get("role") == "user" and item.get("message_id") in by_parent + ] + + async def compact_if_needed(self, username: str, conversation_id: str) -> bool: + lock = _compaction_locks[ + f"{self.graphname}\0{username}\0{conversation_id}" + ] + async with lock: + summary, messages = self.store.unsummarized_messages( + username, self.graphname, conversation_id + ) + old_summary = summary.summary if summary else "" + selected = select_compaction_range( + old_summary, + messages, + budget_tokens=self.config["stm_context_tokens"], + ) + if not selected: + return False + selected_text = message_text(selected) + prompt = ( + "Update the rolling conversation summary. Preserve user goals, " + "decisions, unresolved questions, referenced entities, and constraints. " + "Do not invent facts. Return only the updated summary.\n\n" + f"Existing summary:\n{old_summary or '(none)'}\n\n" + f"New conversation range:\n{selected_text}" + ) + llm = get_llm_service(get_chat_config(self.graphname)) + response = await asyncio.to_thread(llm.llm.invoke, prompt) + summary_text = llm._message_text(response).strip() + self.store.update_summary( + username, + self.graphname, + conversation_id, + summary_text=summary_text, + through_sequence=selected[-1].sequence, + source_tokens=token_count(selected_text), + summary_tokens=token_count(summary_text), + model=str(get_chat_config(self.graphname).get("llm_model", "")), + ) + return True + + def short_term_context(self, username: str, conversation_id: str) -> str: + summary, messages = self.store.unsummarized_messages( + username, self.graphname, conversation_id + ) + # The in-flight user turn is already the agent question. Feeding it + # back as STM makes the planner treat the question as "provided + # context" and skip graph retrieval. + if messages and getattr(messages[-1], "role", None) == "user": + messages = messages[:-1] + sections = [] + if summary and summary.summary: + sections.append("Conversation summary:\n" + summary.summary) + budget = self.config["stm_context_tokens"] - sum( + token_count(section) for section in sections + ) + selected = [] + for item in reversed(messages): + rendered = f"{item.role}: {item.content}" + tokens = token_count(rendered) + if selected and tokens > budget: + break + if tokens <= budget: + selected.append(rendered) + budget -= tokens + if selected: + sections.append("Recent exact turns:\n" + "\n".join(reversed(selected))) + return "\n\n".join(sections) + + def recall(self, username: str, question: str) -> list[dict]: + if self.store.expire_facts(username, self.graphname): + self.flush_outbox() + # LTM is user+graph, not conversation. Skip mem0 entirely when this + # scope has never stored a fact — first chat in a new graph must not + # pay Mem0Runtime + TigerGraph index startup for an empty search. + if not self.store.has_facts(username, self.graphname): + return [] + result = self._runtime().search( + question, + username, + top_k=self.config["ltm_top_k"], + threshold=self.config["ltm_score_threshold"], + ) + return result.get("results", []) + + def memory_context( + self, username: str, conversation_id: str, question: str + ) -> str: + sections = [] + stm = self.short_term_context(username, conversation_id) + if stm: + sections.append(stm) + try: + facts = self.recall(username, question) + except Exception as exc: + logger.warning("Long-term memory recall unavailable: %s", exc) + facts = [] + if facts: + sections.append( + "Relevant long-term user memory (not graph evidence):\n" + + "\n".join(f"- {item.get('memory', '')}" for item in facts) + ) + return "\n\n".join(sections) + + def extract(self, username: str, conversation_id: str, messages: list[dict]): + safe_messages = [ + message + for message in messages + if not _SECRET_PATTERN.search(str(message.get("content") or "")) + ] + if not safe_messages: + return {"results": []} + result = self._runtime().add(safe_messages, username, conversation_id) + self.store.apply_retention( + username, self.graphname, self.config["ltm_retention_days"] + ) + self.store.enforce_fact_cap( + username, self.graphname, self.config["max_facts_per_scope"] + ) + self.flush_outbox() + return result + + def flush_outbox(self) -> None: + for item in self.store.pending_outbox(): + fact = self.store.get_fact(item.fact_id, include_deleted=True) + try: + if item.operation == "delete": + if fact: + get_embedding_store( + fact.graphname + ).delete_memory_embedding(fact.index_id) + self.store.complete_outbox(item.id, indexed=False) + elif fact: + if not fact.indexed: + vector = get_embedding_service().embed_query(fact.text) + get_embedding_store( + fact.graphname + ).upsert_memory_embedding(fact.index_id, vector) + self.store.complete_outbox(item.id, indexed=True) + except Exception as exc: + logger.warning("Memory index outbox retry failed: %s", exc) + self.store.fail_outbox(item.id, str(exc)) diff --git a/common/memory/store.py b/common/memory/store.py new file mode 100644 index 00000000..efa329e3 --- /dev/null +++ b/common/memory/store.py @@ -0,0 +1,612 @@ +# Copyright (c) 2024-2026 TigerGraph, Inc. +"""Canonical relational store for chat transcripts and mem0 payloads.""" + +from __future__ import annotations + +import hashlib +from contextlib import contextmanager +from datetime import datetime, timedelta +from time import sleep +from typing import Iterable, Iterator + +from sqlalchemy import create_engine, func, select, update +from sqlalchemy.engine import Engine +from sqlalchemy.exc import OperationalError +from sqlalchemy.orm import Session, sessionmaker + +from common.memory.models import ( + Base, + Conversation, + ConversationSummary, + MemoryFact, + MemoryHistory, + MemoryOutbox, + Message, + utcnow, +) + + +def memory_user_id(username: str, graphname: str) -> str: + """mem0 user ID; graph is part of the tenant boundary.""" + return f"{graphname}\0{username}" + + +def scope_hash(username: str, graphname: str) -> str: + return hashlib.sha256(memory_user_id(username, graphname).encode()).hexdigest()[:32] + + +def fact_index_id(username: str, graphname: str, fact_id: str) -> str: + return f"{scope_hash(username, graphname)}:{fact_id}" + + +class PostgresMemoryStore: + def __init__(self, database_url: str, *, create_schema: bool = True): + connect_args = {"check_same_thread": False} if database_url.startswith("sqlite") else {} + self.engine: Engine = create_engine( + database_url, pool_pre_ping=True, future=True, connect_args=connect_args + ) + self._session_factory = sessionmaker(self.engine, expire_on_commit=False) + if create_schema: + for attempt in range(6): + try: + Base.metadata.create_all(self.engine) + break + except OperationalError: + if attempt == 5: + raise + sleep(2 ** attempt) + + @contextmanager + def session(self) -> Iterator[Session]: + session = self._session_factory() + try: + yield session + session.commit() + except Exception: + session.rollback() + raise + finally: + session.close() + + def healthcheck(self) -> bool: + with self.session() as session: + session.execute(select(1)) + return True + + @staticmethod + def conversation_json(row: Conversation) -> dict: + return { + "id": row.id, + "user_id": row.username, + "graphname": row.graphname, + "conversation_id": row.conversation_id, + "name": row.name, + "create_ts": row.created_at, + "update_ts": row.updated_at, + "delete_ts": row.deleted_at, + } + + @staticmethod + def message_json(row: Message) -> dict: + return { + "id": row.id, + "conversation_id": row.conversation_id, + "message_id": row.message_id, + "parent_id": row.parent_id, + "model": row.model, + "content": row.content, + "role": row.role, + "response_time": row.response_time, + "feedback": row.feedback, + "comment": row.comment, + "answered_question": row.answered_question, + "response_type": row.response_type, + "query_sources": row.query_sources, + "create_ts": row.created_at, + "update_ts": row.updated_at, + "delete_ts": row.deleted_at, + } + + def _owned_conversation( + self, session: Session, username: str, graphname: str, conversation_id: str + ) -> Conversation | None: + return session.scalar( + select(Conversation).where( + Conversation.username == username, + Conversation.graphname == graphname, + Conversation.conversation_id == conversation_id, + Conversation.deleted_at.is_(None), + ) + ) + + def list_conversations(self, username: str, graphname: str) -> list[dict]: + with self.session() as session: + rows = session.scalars( + select(Conversation) + .where( + Conversation.username == username, + Conversation.graphname == graphname, + Conversation.deleted_at.is_(None), + ) + .order_by(Conversation.updated_at.desc()) + ).all() + return [self.conversation_json(row) for row in rows] + + def get_messages( + self, username: str, graphname: str, conversation_id: str + ) -> list[dict]: + with self.session() as session: + if self._owned_conversation(session, username, graphname, conversation_id) is None: + return [] + rows = session.scalars( + select(Message) + .where( + Message.conversation_id == conversation_id, + Message.deleted_at.is_(None), + ) + .order_by(Message.sequence) + ).all() + return [self.message_json(row) for row in rows] + + def save_message( + self, + username: str, + graphname: str, + message: dict, + *, + conversation_name: str | None = None, + created_at: datetime | None = None, + updated_at: datetime | None = None, + ) -> dict: + conversation_id = str(message["conversation_id"]) + message_id = str(message["message_id"]) + with self.session() as session: + conversation = self._owned_conversation( + session, username, graphname, conversation_id + ) + if conversation is None: + if session.scalar( + select(Conversation).where( + Conversation.conversation_id == conversation_id + ) + ): + raise PermissionError("Conversation belongs to another scope") + content = str(message.get("content") or "") + conversation = Conversation( + username=username, + graphname=graphname, + conversation_id=conversation_id, + name=conversation_name + or (content[:80] if message.get("role") == "user" else "New conversation"), + created_at=created_at or utcnow(), + updated_at=updated_at or created_at or utcnow(), + ) + session.add(conversation) + session.flush() + + existing = session.scalar(select(Message).where(Message.message_id == message_id)) + if existing: + if existing.conversation_id != conversation_id: + raise PermissionError("Message belongs to another conversation") + existing.feedback = message.get("feedback") + existing.comment = message.get("comment") + existing.updated_at = updated_at or utcnow() + return self.message_json(existing) + + sequence = int( + session.scalar( + select(func.coalesce(func.max(Message.sequence), 0)).where( + Message.conversation_id == conversation_id + ) + ) + ) + 1 + row = Message( + conversation_id=conversation_id, + message_id=message_id, + parent_id=message.get("parent_id"), + sequence=sequence, + model=message.get("model"), + content=str(message.get("content") or ""), + role=str(message.get("role") or "user"), + response_time=message.get("response_time"), + feedback=message.get("feedback"), + comment=message.get("comment"), + answered_question=message.get("answered_question"), + response_type=message.get("response_type"), + query_sources=message.get("query_sources"), + created_at=created_at or utcnow(), + updated_at=updated_at or created_at or utcnow(), + ) + session.add(row) + conversation.updated_at = row.updated_at + session.flush() + return self.message_json(row) + + def feedback_messages( + self, username: str, graphname: str, *, include_all: bool = False + ) -> list[dict]: + with self.session() as session: + stmt = ( + select(Message) + .join( + Conversation, + Conversation.conversation_id == Message.conversation_id, + ) + .where( + Conversation.graphname == graphname, + Conversation.deleted_at.is_(None), + Message.deleted_at.is_(None), + ) + ) + if not include_all: + stmt = stmt.where(Conversation.username == username) + return [ + self.message_json(row) + for row in session.scalars(stmt.order_by(Message.created_at.desc())).all() + ] + + def delete_conversation( + self, username: str, graphname: str, conversation_id: str + ) -> list[str]: + with self.session() as session: + conversation = self._owned_conversation( + session, username, graphname, conversation_id + ) + if conversation is None: + return [] + now = utcnow() + conversation.deleted_at = now + session.execute( + update(Message) + .where(Message.conversation_id == conversation_id) + .values(deleted_at=now) + ) + facts = list( + session.scalars( + select(MemoryFact).where( + MemoryFact.username == username, + MemoryFact.graphname == graphname, + MemoryFact.conversation_id == conversation_id, + MemoryFact.deleted_at.is_(None), + ) + ).all() + ) + for fact in facts: + fact.deleted_at = now + self._enqueue(session, "delete", fact.fact_id) + return [fact.fact_id for fact in facts] + + def get_summary( + self, username: str, graphname: str, conversation_id: str + ) -> ConversationSummary | None: + with self.session() as session: + if self._owned_conversation(session, username, graphname, conversation_id) is None: + return None + return session.get(ConversationSummary, conversation_id) + + def unsummarized_messages( + self, username: str, graphname: str, conversation_id: str + ) -> tuple[ConversationSummary | None, list[Message]]: + with self.session() as session: + if self._owned_conversation(session, username, graphname, conversation_id) is None: + return None, [] + summary = session.get(ConversationSummary, conversation_id) + after = summary.last_summarized_sequence if summary else 0 + messages = list( + session.scalars( + select(Message) + .where( + Message.conversation_id == conversation_id, + Message.sequence > after, + Message.deleted_at.is_(None), + ) + .order_by(Message.sequence) + ).all() + ) + return summary, messages + + def update_summary( + self, + username: str, + graphname: str, + conversation_id: str, + *, + summary_text: str, + through_sequence: int, + source_tokens: int, + summary_tokens: int, + model: str, + ) -> None: + with self.session() as session: + conversation = self._owned_conversation( + session, username, graphname, conversation_id + ) + if conversation is None: + raise PermissionError("Conversation does not belong to caller") + # Row lock serializes compaction for this conversation on PostgreSQL. + session.execute( + select(Conversation) + .where(Conversation.id == conversation.id) + .with_for_update() + ) + row = session.get(ConversationSummary, conversation_id) + if row is None: + row = ConversationSummary( + conversation_id=conversation_id, + summary="", + last_summarized_sequence=0, + source_tokens=0, + summary_tokens=0, + version=0, + ) + session.add(row) + if through_sequence <= row.last_summarized_sequence: + return + row.summary = summary_text + row.last_summarized_sequence = through_sequence + row.source_tokens = source_tokens + row.summary_tokens = summary_tokens + row.model = model + row.version = (row.version or 0) + 1 + + def put_fact( + self, + fact_id: str, + username: str, + graphname: str, + conversation_id: str | None, + text: str, + payload: dict, + ) -> MemoryFact: + with self.session() as session: + row = session.get(MemoryFact, fact_id) + if row is None: + row = MemoryFact( + fact_id=fact_id, + index_id=fact_index_id(username, graphname, fact_id), + username=username, + graphname=graphname, + conversation_id=conversation_id, + text=text, + payload=payload, + ) + session.add(row) + elif row.username != username or row.graphname != graphname: + raise PermissionError("Memory fact scope cannot change") + else: + row.text = text + row.payload = payload + row.deleted_at = None + self._enqueue(session, "upsert", fact_id) + session.flush() + return row + + def get_fact(self, fact_id: str, *, include_deleted: bool = False) -> MemoryFact | None: + with self.session() as session: + stmt = select(MemoryFact).where(MemoryFact.fact_id == fact_id) + if not include_deleted: + stmt = stmt.where(MemoryFact.deleted_at.is_(None)) + return session.scalar(stmt) + + def has_facts(self, username: str, graphname: str) -> bool: + with self.session() as session: + return session.scalar( + select(MemoryFact.fact_id).where( + MemoryFact.username == username, + MemoryFact.graphname == graphname, + MemoryFact.deleted_at.is_(None), + ).limit(1) + ) is not None + + def list_facts( + self, username: str, graphname: str, conversation_id: str | None = None + ) -> list[MemoryFact]: + with self.session() as session: + stmt = select(MemoryFact).where( + MemoryFact.username == username, + MemoryFact.graphname == graphname, + MemoryFact.deleted_at.is_(None), + ) + if conversation_id is not None: + stmt = stmt.where(MemoryFact.conversation_id == conversation_id) + return list(session.scalars(stmt.order_by(MemoryFact.created_at.desc())).all()) + + def facts_by_index_ids( + self, username: str, graphname: str, index_ids: Iterable[str] + ) -> list[MemoryFact]: + ordered = list(index_ids) + if not ordered: + return [] + with self.session() as session: + rows = session.scalars( + select(MemoryFact).where( + MemoryFact.username == username, + MemoryFact.graphname == graphname, + MemoryFact.index_id.in_(ordered), + MemoryFact.deleted_at.is_(None), + ) + ).all() + by_id = {row.index_id: row for row in rows} + return [by_id[index] for index in ordered if index in by_id] + + def delete_fact(self, fact_id: str) -> None: + with self.session() as session: + row = session.get(MemoryFact, fact_id) + if row and row.deleted_at is None: + row.deleted_at = utcnow() + self._enqueue(session, "delete", fact_id) + + def enforce_fact_cap(self, username: str, graphname: str, maximum: int) -> int: + with self.session() as session: + rows = list( + session.scalars( + select(MemoryFact) + .where( + MemoryFact.username == username, + MemoryFact.graphname == graphname, + MemoryFact.deleted_at.is_(None), + ) + .order_by(MemoryFact.updated_at.desc()) + ).all() + ) + evicted = rows[maximum:] + for row in evicted: + row.deleted_at = utcnow() + self._enqueue(session, "delete", row.fact_id) + return len(evicted) + + def apply_retention( + self, username: str, graphname: str, retention_days: int + ) -> None: + expires_at = utcnow() + timedelta(days=retention_days) + with self.session() as session: + rows = session.scalars( + select(MemoryFact).where( + MemoryFact.username == username, + MemoryFact.graphname == graphname, + MemoryFact.deleted_at.is_(None), + MemoryFact.expires_at.is_(None), + ) + ).all() + for row in rows: + row.expires_at = expires_at + + def expire_facts(self, username: str, graphname: str) -> int: + with self.session() as session: + rows = session.scalars( + select(MemoryFact).where( + MemoryFact.username == username, + MemoryFact.graphname == graphname, + MemoryFact.deleted_at.is_(None), + MemoryFact.expires_at.is_not(None), + MemoryFact.expires_at <= utcnow(), + ) + ).all() + for row in rows: + row.deleted_at = utcnow() + self._enqueue(session, "delete", row.fact_id) + return len(rows) + + def add_history(self, record: dict) -> None: + with self.session() as session: + session.add( + MemoryHistory( + memory_id=str(record.get("memory_id")), + old_memory=record.get("old_memory"), + new_memory=record.get("new_memory"), + event=str(record.get("event") or "ADD"), + created_at=_parse_datetime(record.get("created_at")) or utcnow(), + updated_at=_parse_datetime(record.get("updated_at")), + is_deleted=bool(record.get("is_deleted", False)), + actor_id=record.get("actor_id"), + role=record.get("role"), + ) + ) + + def get_history(self, memory_id: str) -> list[dict]: + with self.session() as session: + rows = session.scalars( + select(MemoryHistory) + .where(MemoryHistory.memory_id == memory_id) + .order_by(MemoryHistory.created_at) + ).all() + return [ + { + "id": row.id, + "memory_id": row.memory_id, + "old_memory": row.old_memory, + "new_memory": row.new_memory, + "event": row.event, + "created_at": row.created_at.isoformat(), + "updated_at": row.updated_at.isoformat() if row.updated_at else None, + "is_deleted": row.is_deleted, + "actor_id": row.actor_id, + "role": row.role, + } + for row in rows + ] + + @staticmethod + def _enqueue(session: Session, operation: str, fact_id: str) -> None: + item = session.scalar( + select(MemoryOutbox).where( + MemoryOutbox.operation == operation, MemoryOutbox.fact_id == fact_id + ) + ) + if item: + item.completed_at = None + item.next_attempt_at = utcnow() + item.last_error = None + else: + session.add(MemoryOutbox(operation=operation, fact_id=fact_id)) + + def pending_outbox(self, limit: int = 100) -> list[MemoryOutbox]: + with self.session() as session: + return list( + session.scalars( + select(MemoryOutbox) + .where( + MemoryOutbox.completed_at.is_(None), + MemoryOutbox.next_attempt_at <= utcnow(), + ) + .order_by(MemoryOutbox.created_at) + .limit(limit) + ).all() + ) + + def complete_outbox(self, outbox_id: str, *, indexed: bool) -> None: + with self.session() as session: + item = session.get(MemoryOutbox, outbox_id) + if not item: + return + item.completed_at = utcnow() + fact = session.get(MemoryFact, item.fact_id) + if fact: + fact.indexed = indexed + + def mark_fact_indexed(self, fact_id: str) -> None: + with self.session() as session: + fact = session.get(MemoryFact, fact_id) + if fact: + fact.indexed = True + items = session.scalars( + select(MemoryOutbox).where( + MemoryOutbox.fact_id == fact_id, + MemoryOutbox.operation == "upsert", + MemoryOutbox.completed_at.is_(None), + ) + ).all() + for item in items: + item.completed_at = utcnow() + + def mark_fact_unindexed(self, fact_id: str) -> None: + with self.session() as session: + fact = session.get(MemoryFact, fact_id) + if fact: + fact.indexed = False + items = session.scalars( + select(MemoryOutbox).where( + MemoryOutbox.fact_id == fact_id, + MemoryOutbox.operation == "delete", + MemoryOutbox.completed_at.is_(None), + ) + ).all() + for item in items: + item.completed_at = utcnow() + + def fail_outbox(self, outbox_id: str, error: str) -> None: + with self.session() as session: + item = session.get(MemoryOutbox, outbox_id) + if item: + item.attempts += 1 + item.last_error = error[:4000] + delay = min(3600, 2 ** min(item.attempts, 10)) + item.next_attempt_at = utcnow() + timedelta(seconds=delay) + +def _parse_datetime(value) -> datetime | None: + if value is None or isinstance(value, datetime): + return value + try: + return datetime.fromisoformat(str(value).replace("Z", "+00:00")) + except ValueError: + return None diff --git a/common/requirements.txt b/common/requirements.txt index d20bec7a..dd33f670 100644 --- a/common/requirements.txt +++ b/common/requirements.txt @@ -202,3 +202,5 @@ rpds-py>=0.30.0 PyJWT>=2.13.0 annotated-doc>=0.0.4 typing-inspection>=0.4.2 +mem0ai==2.1.0 +psycopg[binary]>=3.2.8,<4.0.0 diff --git a/docker-compose.yml b/docker-compose.yml index 1034eada..c12bf186 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -10,7 +10,7 @@ services: - 8000:8000 depends_on: - graphrag-ecc - - chat-history + - postgres # - tigergraph environment: SERVER_CONFIG: "/code/configs/server_config.json" @@ -35,20 +35,20 @@ services: volumes: - ./configs/:/code/configs - chat-history: - image: tigergraph/chat-history:latest - platform: linux/amd64 - container_name: chat-history - build: - context: chat-history/ - dockerfile: Dockerfile - ports: - - 8002:8002 + postgres: + image: postgres:16-alpine + container_name: graphrag-postgres environment: - CONFIG_FILES: "/configs/server_config.json" - LOGLEVEL: "INFO" + POSTGRES_DB: graphrag + POSTGRES_USER: graphrag + POSTGRES_PASSWORD: graphrag + healthcheck: + test: ["CMD-SHELL", "pg_isready -U graphrag -d graphrag"] + interval: 5s + timeout: 5s + retries: 12 volumes: - - ./configs/:/configs + - graphrag_postgres_data:/var/lib/postgresql/data graphrag-ui: image: tigergraph/graphrag-ui:latest @@ -84,5 +84,6 @@ services: # volumes: # - tigergraph_data:/home/tigergraph/tigergraph/data # -# volumes: -# tigergraph_data: +volumes: + graphrag_postgres_data: +# tigergraph_data: diff --git a/docs/tutorials/configs/server_config.json b/docs/tutorials/configs/server_config.json index 2ee25f4e..a39ead56 100644 --- a/docs/tutorials/configs/server_config.json +++ b/docs/tutorials/configs/server_config.json @@ -10,7 +10,6 @@ "graphrag_config": { "reuse_embedding": true, "ecc": "http://graphrag-ecc:8001", - "chat_history_api": "http://chat-history:8002", "chunker_config": {} }, "llm_config": { @@ -38,14 +37,16 @@ } } }, - "chat_config": { - "apiPort": "8002", - "dbPath": "chats.db", - "dbLogPath": "db.log", - "logPath": "requestLogs.jsonl", - "conversationAccessRoles": [ - "superuser", - "globaldesigner" - ] + "memory_config": { + "enabled": true, + "database_url": "postgresql+psycopg://graphrag:graphrag@postgres:5432/graphrag", + "stm_context_tokens": 6000, + "ltm_top_k": 5, + "ltm_score_threshold": 0.7, + "ltm_retention_days": 365, + "max_facts_per_scope": 1000, + "fact_extract_async": true, + "outbox_batch_size": 100, + "outbox_max_retries": 8 } } diff --git a/graphrag-k8s.yml b/graphrag-k8s.yml index 2ee262ff..2fe64b4d 100644 --- a/graphrag-k8s.yml +++ b/graphrag-k8s.yml @@ -32,6 +32,13 @@ spec: value: "INFO" - name: USE_CYPHER value: "true" + - name: GRAPHRAG_MEMORY_ENABLED + value: "true" + - name: GRAPHRAG_MEMORY_DATABASE_URL + valueFrom: + secretKeyRef: + name: graphrag-postgres + key: database-url volumeMounts: - name: configs mountPath: /code/configs @@ -105,54 +112,82 @@ spec: - port: 8001 targetPort: 8001 +--- +apiVersion: v1 +kind: Secret +metadata: + name: graphrag-postgres + namespace: graphrag +type: Opaque +stringData: + password: change-me + database-url: postgresql+psycopg://graphrag:change-me@postgres:5432/graphrag + +--- +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: graphrag-postgres + namespace: graphrag +spec: + accessModes: + - ReadWriteOnce + resources: + requests: + storage: 10Gi + --- apiVersion: apps/v1 kind: Deployment metadata: - name: chat-history + name: postgres namespace: graphrag spec: replicas: 1 selector: matchLabels: - app: chat-history + app: postgres template: metadata: labels: - app: chat-history + app: postgres spec: containers: - - name: chat-history - image: tigergraph/chat-history:latest + - name: postgres + image: postgres:16-alpine imagePullPolicy: IfNotPresent ports: - - containerPort: 8002 + - containerPort: 5432 env: - - name: CONFIG_FILES - value: "/configs/server_config.json" - - name: LOGLEVEL - value: "INFO" + - name: POSTGRES_DB + value: graphrag + - name: POSTGRES_USER + value: graphrag + - name: POSTGRES_PASSWORD + valueFrom: + secretKeyRef: + name: graphrag-postgres + key: password volumeMounts: - - name: configs - mountPath: /configs + - name: data + mountPath: /var/lib/postgresql/data volumes: - - name: configs - hostPath: - path: /path/to/graphrag/configs - type: DirectoryOrCreate + - name: data + persistentVolumeClaim: + claimName: graphrag-postgres --- apiVersion: v1 kind: Service metadata: - name: chat-history + name: postgres namespace: graphrag spec: selector: - app: chat-history + app: postgres ports: - - port: 8002 - targetPort: 8002 + - port: 5432 + targetPort: 5432 --- apiVersion: apps/v1 diff --git a/graphrag-ui/src/components/Interact.tsx b/graphrag-ui/src/components/Interact.tsx index f5fff486..9b4c29ee 100644 --- a/graphrag-ui/src/components/Interact.tsx +++ b/graphrag-ui/src/components/Interact.tsx @@ -40,9 +40,10 @@ export const Interactions: FC = ({ const sendFeedback = async (action: Feedback, message: Message) => { const creds = sessionStorage.getItem("auth"); + const graph = sessionStorage.getItem("selectedGraph") || ""; setFeedback(action); message.feedback = action; - await fetch(`${GRAPHRAG_URL}/ui/feedback`, { + await fetch(`${GRAPHRAG_URL}/ui/${encodeURIComponent(graph)}/feedback`, { method: "POST", body: JSON.stringify(message), headers: { diff --git a/graphrag-ui/src/components/SideMenu.tsx b/graphrag-ui/src/components/SideMenu.tsx index e79f8908..12a54d93 100644 --- a/graphrag-ui/src/components/SideMenu.tsx +++ b/graphrag-ui/src/components/SideMenu.tsx @@ -58,9 +58,10 @@ import { useConfirm } from "@/hooks/useConfirm"; import { useAlert } from "@/hooks/useAlert"; import { useNavigate } from "react-router-dom"; -// TODO make dynamic -const WS_HISTORY_URL = "/ui/user"; -const WS_CONVO_URL = "/ui/conversation"; +const historyUrl = () => + `/ui/${encodeURIComponent(sessionStorage.getItem("selectedGraph") || "")}/user`; +const conversationUrl = () => + `/ui/${encodeURIComponent(sessionStorage.getItem("selectedGraph") || "")}/conversation`; // How many conversations to load at a time. Only the visible ones have their // messages fetched, so a long history can't flood the browser with requests. const PAGE_SIZE = 10; @@ -132,7 +133,7 @@ const SideMenu = ({ method: "GET", headers: { Authorization: creds, "Content-Type": "application/json" }, }; - const response = await fetch(`${WS_HISTORY_URL}/${username}`, settings); + const response = await fetch(`${historyUrl()}/${username}`, settings); if (!response.ok) return []; const data = await safeJson(response); if (!Array.isArray(data) || data.length === 0) return []; @@ -155,7 +156,7 @@ const SideMenu = ({ const results = await Promise.all( items.map(async (item: any) => { try { - const r = await fetch(`${WS_CONVO_URL}/${item.conversation_id}`, settings); + const r = await fetch(`${conversationUrl()}/${item.conversation_id}`, settings); if (!r.ok) return null; const content = await safeJson(r); let lastUpdateTime = item.update_ts || item.create_ts; @@ -234,7 +235,7 @@ const SideMenu = ({ const chunk = older.slice(i, i + BATCH); await Promise.all( chunk.map((c: any) => - fetch(`${WS_CONVO_URL}/${c.conversation_id}`, settings).catch(() => {}) + fetch(`${conversationUrl()}/${c.conversation_id}`, settings).catch(() => {}) ) ); } @@ -269,7 +270,7 @@ const SideMenu = ({ for (let i = 0; i < ids.length; i += BATCH) { const results = await Promise.all( ids.slice(i, i + BATCH).map((id) => - fetch(`${WS_CONVO_URL}/${id}`, settings) + fetch(`${conversationUrl()}/${id}`, settings) .then((r) => (r.ok ? id : null)) .catch(() => null) ) @@ -420,7 +421,7 @@ const SideMenu = ({ } } - const response = await fetch(`${WS_CONVO_URL}/${id}`, settings); + const response = await fetch(`${conversationUrl()}/${id}`, settings); if (!response.ok) { return; } diff --git a/graphrag/Dockerfile b/graphrag/Dockerfile index ffffe13a..ea0c8083 100644 --- a/graphrag/Dockerfile +++ b/graphrag/Dockerfile @@ -4,7 +4,8 @@ WORKDIR /code COPY common/requirements.txt requirements.txt RUN apt-get update && apt-get upgrade -y -RUN pip install -r requirements.txt +RUN pip install -r requirements.txt \ + && python -m spacy download en_core_web_sm COPY graphrag/app /code COPY common /code/common diff --git a/graphrag/app/agent/agent_graph.py b/graphrag/app/agent/agent_graph.py index 597fbf15..1c570e1c 100644 --- a/graphrag/app/agent/agent_graph.py +++ b/graphrag/app/agent/agent_graph.py @@ -220,7 +220,7 @@ def contextualize_question(self, question: str, conversation) -> str: return question try: history_lines = [] - for turn in conversation[-4:]: + for turn in conversation: if isinstance(turn, dict): q = turn.get("query", "") a = turn.get("response", "") diff --git a/graphrag/app/agent/agent_router.py b/graphrag/app/agent/agent_router.py index 52efd772..5793ee0a 100644 --- a/graphrag/app/agent/agent_router.py +++ b/graphrag/app/agent/agent_router.py @@ -18,6 +18,7 @@ from pydantic import BaseModel, Field from common.logs.logwriter import LogWriter from common.logs.log import req_id_cv +from common.db.schema_utils import is_hidden_memory_type from pyTigerGraph.pyTigerGraph import TigerGraphConnection import logging @@ -41,7 +42,11 @@ def route_question(self, question: str, conversation: list[dict[str, str]] = Non str: The datasource to use for the question. """ LogWriter.info(f"request_id={req_id_cv.get()} ENTRY route_question with {question}") - v_types = self.db_conn.getVertexTypes() + v_types = [ + value + for value in self.db_conn.getVertexTypes() + if not is_hidden_memory_type(value) + ] e_types = self.db_conn.getEdgeTypes() router_parser = PydanticOutputParser[RouterResponse](pydantic_object=RouterResponse) diff --git a/graphrag/app/agent/method_selector.py b/graphrag/app/agent/method_selector.py index fece7dd9..a8bf13b0 100644 --- a/graphrag/app/agent/method_selector.py +++ b/graphrag/app/agent/method_selector.py @@ -33,6 +33,7 @@ from common.logs.log import req_id_cv from common.logs.logwriter import LogWriter +from common.db.schema_utils import is_hidden_memory_type logger = logging.getLogger(__name__) @@ -250,7 +251,11 @@ def choose( # Stage B — LLM. Schema types are passed in to anchor the prompt. try: - v_types = self.db_conn.getVertexTypes() + v_types = [ + value + for value in self.db_conn.getVertexTypes() + if not is_hidden_memory_type(value) + ] e_types = self.db_conn.getEdgeTypes() except Exception as e: # noqa: BLE001 - schema lookup is best-effort logger.warning( diff --git a/graphrag/app/routers/ui.py b/graphrag/app/routers/ui.py index 455e9f92..e4ec4af6 100644 --- a/graphrag/app/routers/ui.py +++ b/graphrag/app/routers/ui.py @@ -67,6 +67,7 @@ from common.utils.user_messages import NO_ANSWER from common.metrics.prometheus_metrics import metrics as pmetrics from common.metrics.tg_proxy import TigerGraphConnectionProxy +from common.memory import MemoryService from common.utils.graph_locks import acquire_graph_lock, release_graph_lock, acquire_rebuild_lock, release_rebuild_lock, get_rebuilding_graph, get_current_operation # Cache the last successful ECC status response per graph so the UI @@ -474,22 +475,13 @@ def _parse_auth_header(authorization: str | None) -> HTTPBasicCredentials: ) -def _chat_history_auth_header(creds: HTTPBasicCredentials) -> str: - """Build the Basic-auth header used when proxying to chat-history. - - Chat-history identifies the caller by the Basic-auth username only - (it ignores the password). For sentinel logins - (``__graphrag_token__`` / ``__GSQL__secret``) we substitute the - TG-resolved username so conversations get stored / fetched under - the user's real identity instead of the sentinel string. - """ +def _resolved_username(creds: HTTPBasicCredentials) -> str: + """Return the authenticated TigerGraph principal used for memory scope.""" try: _, _, resolved = _get_user_role_details(creds.username, creds.password) except Exception: resolved = creds.username - username = resolved or creds.username - encoded = base64.b64encode(f"{username}:{creds.password}".encode()).decode() - return f"Basic {encoded}" + return resolved or creds.username def _ecc_auth_header(creds: HTTPBasicCredentials) -> str: @@ -790,26 +782,17 @@ def schema_reserved_names( } -@router.post(f"{route_prefix}/feedback") -def add_feedback( +@router.post(f"{route_prefix}/{{graphname}}/feedback") +async def add_feedback( + graphname: ValidGraphName, message: Message, creds: Annotated[tuple[list[str], HTTPBasicCredentials], Depends(ui_basic_auth)], ): creds = creds[1] - try: - res = httpx.post( - f"{graphrag_config['chat_history_api']}/conversation", - json=message.model_dump(), - headers={"Authorization": _chat_history_auth_header(creds)}, - ) - res.raise_for_status() - except Exception as e: - exc = traceback.format_exc() - logger.debug_pii( - f"/ui/feedback request_id={req_id_cv.get()} Exception Trace:\n{exc}" - ) - raise e - + service = MemoryService(graphname) + await asyncio.to_thread( + service.save_message, _resolved_username(creds), message.model_dump() + ) return {"message": "feedback saved", "message_id": message.message_id} @@ -2691,27 +2674,19 @@ async def serve_image_from_vertex( raise HTTPException(status_code=500, detail=f"Error serving image: {str(e)}") -@router.get(route_prefix + "/user/{user_id}") +@router.get(route_prefix + "/{graphname}/user/{user_id}") async def get_user_conversations( + graphname: ValidGraphName, user_id: str, creds: Annotated[tuple[list[str], HTTPBasicCredentials], Depends(ui_basic_auth)], ): creds = creds[1] - try: - async with httpx.AsyncClient() as client: - res = await client.get( - f"{graphrag_config['chat_history_api']}/user/{user_id}", - headers={"Authorization": _chat_history_auth_header(creds)}, - ) - res.raise_for_status() - except Exception as e: - exc = traceback.format_exc() - logger.debug_pii( - f"/ui/user/{user_id} request_id={req_id_cv.get()} Exception Trace:\n{exc}" - ) - raise e - - return res.json() + username = _resolved_username(creds) + if user_id != username: + raise HTTPException(status_code=403, detail="Cannot read another user's history") + return await asyncio.to_thread( + MemoryService(graphname).list_conversations, username + ) @router.get(route_prefix + "/roles") @@ -2724,77 +2699,46 @@ async def get_user_roles( return {"roles": roles, "graph_roles": graph_roles} -@router.get(route_prefix + "/conversation/{conversation_id}") +@router.get(route_prefix + "/{graphname}/conversation/{conversation_id}") async def get_conversation_contents( + graphname: ValidGraphName, conversation_id: str, creds: Annotated[tuple[list[str], HTTPBasicCredentials], Depends(ui_basic_auth)], ): creds = creds[1] - try: - async with httpx.AsyncClient() as client: - res = await client.get( - f"{graphrag_config['chat_history_api']}/conversation/{conversation_id}", - headers={"Authorization": _chat_history_auth_header(creds)}, - ) - res.raise_for_status() - except Exception as e: - exc = traceback.format_exc() - logger.debug_pii( - f"/conversation/{conversation_id} request_id={req_id_cv.get()} Exception Trace:\n{exc}" - ) - raise e + return await asyncio.to_thread( + MemoryService(graphname).get_messages, + _resolved_username(creds), + conversation_id, + ) - return res.json() -@router.get(route_prefix + "/get_feedback") +@router.get(route_prefix + "/{graphname}/get_feedback") async def get_conversation_feedback( + graphname: ValidGraphName, creds: Annotated[tuple[list[str], HTTPBasicCredentials], Depends(ui_basic_auth)], ): - creds = creds[1] - try: - async with httpx.AsyncClient() as client: - res = await client.get( - f"{graphrag_config['chat_history_api']}/get_feedback", - headers={"Authorization": _chat_history_auth_header(creds)}, - ) - res.raise_for_status() - except httpx.HTTPStatusError as e: - logger.error(f"HTTP error occurred: {e}") - raise HTTPException(status_code=e.response.status_code, detail="Failed to fetch feedback") - except Exception as e: - exc = traceback.format_exc() - logger.debug_pii( - f"/get_feedback request_id={req_id_cv.get()} Exception Trace:\n{exc}" - ) - raise HTTPException(status_code=500, detail="Internal server error") - - return res.json() + roles, credentials = creds + return await asyncio.to_thread( + MemoryService(graphname).feedback_messages, + _resolved_username(credentials), + _is_superadmin(roles), + ) -@router.delete(route_prefix + "/conversation/{conversation_id}") +@router.delete(route_prefix + "/{graphname}/conversation/{conversation_id}") async def delete_conversation( + graphname: ValidGraphName, conversation_id: str, creds: Annotated[tuple[list[str], HTTPBasicCredentials], Depends(ui_basic_auth)], ): """Delete a conversation and all its messages.""" creds = creds[1] - try: - async with httpx.AsyncClient() as client: - res = await client.delete( - f"{graphrag_config['chat_history_api']}/conversation/{conversation_id}", - headers={"Authorization": _chat_history_auth_header(creds)}, - ) - res.raise_for_status() - except httpx.HTTPStatusError as e: - logger.error(f"HTTP error occurred: {e}") - raise HTTPException(status_code=e.response.status_code, detail="Failed to delete conversation") - except Exception as e: - exc = traceback.format_exc() - logger.debug_pii( - f"/conversation/{conversation_id} DELETE request_id={req_id_cv.get()} Exception Trace:\n{exc}" - ) - raise HTTPException(status_code=500, detail="Internal server error") - + await asyncio.to_thread( + MemoryService(graphname).delete_conversation, + _resolved_username(creds), + conversation_id, + ) return {"message": "Conversation deleted successfully"} @@ -2885,8 +2829,7 @@ async def run_agent( async with asyncio.TaskGroup() as tg: # run agent a_resp = tg.create_task( - # TODO: make num mesages in history configureable - a_question_for_agent(data, conversation_history[-4:]) + a_question_for_agent(data, conversation_history) ) # sample Q and emit events if ws: @@ -2935,48 +2878,23 @@ async def run_agent( async def load_conversation_history( - conversation_id: str, usr_creds: HTTPBasicCredentials + conversation_id: str, + usr_creds: HTTPBasicCredentials, + graphname: str, + username: str | None = None, ) -> list[dict[str, str]]: - """ - Load conversation history from the chat history service. - Returns a list of dicts with 'query', 'response', 'create_ts', and 'update_ts' keys. - """ if not conversation_id or conversation_id == "new": return [] - - ch = graphrag_config.get("chat_history_api") - if ch is None: - LogWriter.info("chat-history not enabled, returning empty history") - return [] - - headers = {"Authorization": _chat_history_auth_header(usr_creds)} try: - async with httpx.AsyncClient() as client: - res = await client.get( - f"{ch}/conversation/{conversation_id}", - headers=headers, - ) - res.raise_for_status() - conversation_data = res.json() - # Convert conversation messages to the format expected by the agent - history = [] - for msg in conversation_data: - if msg.get("role") == "user": - # Find the corresponding system response - for response_msg in conversation_data: - if (response_msg.get("role") == "system" and - response_msg.get("parent_id") == msg.get("message_id")): - history.append({ - "query": msg.get("content", ""), - "response": response_msg.get("content", ""), - "create_ts": response_msg.get("create_ts"), - "update_ts": response_msg.get("update_ts"), - }) - break - - LogWriter.info(f"Loaded {len(history)} conversation history entries for conversation {conversation_id}") - return history - + history = await asyncio.to_thread( + MemoryService(graphname).history_pairs, + username or _resolved_username(usr_creds), + conversation_id, + ) + LogWriter.info( + f"Loaded {len(history)} conversation history entries for conversation {conversation_id}" + ) + return history except Exception as e: exc = traceback.format_exc() logger.debug_pii(f"Error loading conversation history for {conversation_id}\nException Trace:\n{exc}") @@ -2985,23 +2903,53 @@ async def load_conversation_history( async def write_message_to_history( - message: Message, usr_creds: HTTPBasicCredentials + message: Message, + usr_creds: HTTPBasicCredentials, + graphname: str, + username: str | None = None, ): - ch = graphrag_config.get("chat_history_api") - if ch is not None: - headers = {"Authorization": _chat_history_auth_header(usr_creds)} - try: - async with httpx.AsyncClient() as client: - res = await client.post( - f"{ch}/conversation", headers=headers, json=message.model_dump() - ) - res.raise_for_status() - except Exception: # catch all exceptions to log them, but don't raise - exc = traceback.format_exc() - logger.debug_pii(f"Error writing chat history\nException Trace:\n{exc}") + await asyncio.to_thread( + MemoryService(graphname).save_message, + username or _resolved_username(usr_creds), + message.model_dump(), + ) + + +def _with_memory_context(conversation_history, memory_context): + history = list(conversation_history or []) + if memory_context: + history.append( + { + "query": "Memory context (not knowledge-graph evidence)", + "response": memory_context, + } + ) + return history - else: - LogWriter.info(f"chat-history not enabled. chat-history url: {ch}") + +async def _maintain_memory( + service: MemoryService, + username: str, + conversation_id: str, + user_text: str, + assistant_text: str, +) -> None: + try: + await asyncio.to_thread( + service.extract, + username, + conversation_id, + [ + {"role": "user", "content": user_text}, + {"role": "assistant", "content": assistant_text}, + ], + ) + except Exception as exc: + logger.warning("mem0 extraction deferred/failed: %s", exc) + try: + await service.compact_if_needed(username, conversation_id) + except Exception as exc: + logger.warning("STM compaction failed: %s", exc) # Recognized agentic orchestrator styles. A value routed to the agentic engine # that is not one of these (e.g. a classic retriever name like "hybrid") would @@ -3084,8 +3032,16 @@ async def graph_query( ).decode() _, conn = ws_basic_auth(auth_header, graphname) try: + username = _resolved_username(creds) + memory_service = MemoryService(graphname) # Load conversation history if conversation_id is provided - conversation_history = await load_conversation_history(conversation_id, creds) if conversation_id else [] + conversation_history = ( + await load_conversation_history( + conversation_id, creds, graphname, username + ) + if conversation_id + else [] + ) # Use provided conversation ID or generate new one if not conversation_id or conversation_id == "new": @@ -3111,9 +3067,16 @@ async def graph_query( role=Role.USER, ) # save message - await write_message_to_history(message, creds) + await write_message_to_history(message, creds, graphname, username) prev_id = message.message_id + memory_context = await asyncio.to_thread( + memory_service.memory_context, username, convo_id, data + ) + conversation_history = _with_memory_context( + conversation_history, memory_context + ) + # generate response and keep track of response time start = time.monotonic() resp = await run_agent( @@ -3135,7 +3098,18 @@ async def graph_query( response_type=resp.response_type, query_sources=resp.query_sources, ) - await write_message_to_history(message, creds) + await write_message_to_history(message, creds, graphname, username) + maintenance = _maintain_memory( + memory_service, + username, + convo_id, + data, + resp.natural_language_response, + ) + if memory_service.config["fact_extract_async"]: + asyncio.create_task(maintenance) + else: + await maintenance await asyncio.to_thread(_save_trace_log, message.message_id, convo_id, data, resp, elapsed, creds.username) prev_id = message.message_id @@ -3240,9 +3214,12 @@ async def chat( f"WebSocket conversation_id received: {conversation_id or 'empty'} " f"(graph={graphname}, mode={mode or 'config'}, selection={rag_pattern})" ) - + + memory_service = MemoryService(graphname) # Load conversation history if not a new conversation - conversation_history = await load_conversation_history(conversation_id, usr_creds) + conversation_history = await load_conversation_history( + conversation_id, usr_creds, graphname, ws_username + ) # Use provided conversation ID or generate new one if conversation_id == "new" or not conversation_id: @@ -3277,13 +3254,22 @@ async def chat( role=Role.USER, ) # save message - await write_message_to_history(message, usr_creds) + await write_message_to_history( + message, usr_creds, graphname, ws_username + ) prev_id = message.message_id + memory_context = await asyncio.to_thread( + memory_service.memory_context, ws_username, convo_id, data + ) + agent_history = _with_memory_context( + conversation_history, memory_context + ) + # generate response and keep track of response time start = time.monotonic() resp = await run_agent( - agent, data, conversation_history, graphname, websocket, + agent, data, agent_history, graphname, websocket, is_superadmin=ws_is_superadmin, ) elapsed = time.monotonic() - start @@ -3301,7 +3287,20 @@ async def chat( response_type=resp.response_type, query_sources=resp.query_sources, ) - await write_message_to_history(message, usr_creds) + await write_message_to_history( + message, usr_creds, graphname, ws_username + ) + maintenance = _maintain_memory( + memory_service, + ws_username, + convo_id, + data, + resp.natural_language_response, + ) + if memory_service.config["fact_extract_async"]: + asyncio.create_task(maintenance) + else: + await maintenance await asyncio.to_thread(_save_trace_log, message.message_id, convo_id, data, resp, elapsed, ws_username) prev_id = message.message_id diff --git a/graphrag/app/tools/find_existing_query.py b/graphrag/app/tools/find_existing_query.py index c9c25d6d..abc69a27 100644 --- a/graphrag/app/tools/find_existing_query.py +++ b/graphrag/app/tools/find_existing_query.py @@ -75,7 +75,9 @@ def _get_installed_queries(self) -> List[str]: graphname = self.conn.graphname installed_queries = [ q.split("/")[-1] for q in endpoints - if f"/{graphname}/" in q and q.endswith(("GET", "POST")) + if f"/{graphname}/" in q + and q.endswith(("GET", "POST")) + and not q.split("/")[-1].casefold().startswith("memory_") ] return installed_queries except Exception as e: diff --git a/graphrag/app/tools/map_question_to_schema.py b/graphrag/app/tools/map_question_to_schema.py index 35d742dd..e7a00c04 100644 --- a/graphrag/app/tools/map_question_to_schema.py +++ b/graphrag/app/tools/map_question_to_schema.py @@ -27,7 +27,7 @@ from common.logs.log import req_id_cv from common.logs.logwriter import LogWriter from common.db.connections import get_schema_ver -from common.db.schema_utils import read_type_metadata +from common.db.schema_utils import is_hidden_memory_type, read_type_metadata logger = logging.getLogger(__name__) @@ -103,7 +103,11 @@ def _run(self, query: str, conversation: List[Dict[str, str]]) -> str: schema_ver = get_schema_ver(self.conn) if schema_ver is None or self.schema_ver != schema_ver: self.schema_ver = schema_ver if schema_ver is not None else -1 - self.vertices = self.conn.getVertexTypes() + self.vertices = [ + vertex + for vertex in self.conn.getVertexTypes() + if not is_hidden_memory_type(vertex) + ] self.edges = self.conn.getEdgeTypes() for vertex in self.vertices: diff --git a/graphrag/app/tools/validation_utils.py b/graphrag/app/tools/validation_utils.py index 7a5a9817..65748231 100644 --- a/graphrag/app/tools/validation_utils.py +++ b/graphrag/app/tools/validation_utils.py @@ -7,6 +7,7 @@ import logging from common.logs.log import req_id_cv from common.logs.logwriter import LogWriter +from common.db.schema_utils import is_hidden_memory_type logger = logging.getLogger(__name__) @@ -33,6 +34,10 @@ def validate_schema(conn, v_types, e_types, v_attrs, e_attrs): if v_types: vertices = conn.getVertexTypes() for v in v_types: + if is_hidden_memory_type(v): + raise MapQuestionToSchemaException( + "Internal memory schema is not available to agents." + ) logger.debug( f"request_id={req_id_cv.get()} validate_schema() validating vertex_type={v}" ) @@ -83,6 +88,10 @@ def validate_function_call(conn, generated_call: str, valid_functions: list) -> # handle installed queries LogWriter.info(f"request_id={req_id_cv.get()} ENTRY validate_function_call()") generated_call = generated_call.strip().strip("\n").strip("\t") + if "memoryembedding" in generated_call.casefold() or "memory_" in generated_call.casefold(): + raise InvalidFunctionCallException( + "Internal memory schema and queries are not available to agents." + ) # LogWriter.info(f"generated_call: {generated_call}") # LogWriter.info(f"valid_headers: {valid_headers}") endpoints = conn.getEndpoints(dynamic=True) # installed queries in database diff --git a/graphrag/tests/test_memory_store.py b/graphrag/tests/test_memory_store.py new file mode 100644 index 00000000..f9391377 --- /dev/null +++ b/graphrag/tests/test_memory_store.py @@ -0,0 +1,238 @@ +from pathlib import Path +import importlib.metadata +import inspect + +import pytest + +from common.memory.mem0_adapter import ( + PostgresHistoryAdapter, + SplitMemoryVectorStore, +) +from common.memory.compaction import select_compaction_range +from common.memory.store import PostgresMemoryStore, memory_user_id + + +@pytest.fixture +def store(tmp_path: Path): + return PostgresMemoryStore(f"sqlite:///{tmp_path / 'memory.db'}") + + +def _message(conversation_id, message_id, role="user", content="hello", parent_id=None): + return { + "conversation_id": conversation_id, + "message_id": message_id, + "parent_id": parent_id, + "role": role, + "content": content, + } + + +def test_has_facts_is_user_and_graph_scoped(store): + assert store.has_facts("alice", "g1") is False + store.put_fact( + "fact-1", + username="alice", + graphname="g1", + conversation_id="c1", + text="prefers concise answers", + payload={}, + ) + assert store.has_facts("alice", "g1") is True + assert store.has_facts("bob", "g1") is False + assert store.has_facts("alice", "g2") is False + + +def test_conversations_are_user_and_graph_scoped(store): + store.save_message("alice", "g1", _message("c1", "m1")) + + assert len(store.list_conversations("alice", "g1")) == 1 + assert store.list_conversations("bob", "g1") == [] + assert store.list_conversations("alice", "g2") == [] + assert store.get_messages("bob", "g1", "c1") == [] + + +def test_message_id_cannot_cross_conversations(store): + store.save_message("alice", "g1", _message("c1", "m1")) + with pytest.raises(PermissionError): + store.save_message("alice", "g1", _message("c2", "m1")) + + +def test_mem0_history_reads_canonical_transcript_without_duplicate_table(store): + store.save_message("alice", "g1", _message("c1", "m1", content="question")) + store.save_message( + "alice", + "g1", + _message("c1", "m2", role="system", content="answer", parent_id="m1"), + ) + adapter = PostgresHistoryAdapter(store) + scope = f"run_id=c1&user_id={memory_user_id('alice', 'g1')}" + + adapter.save_messages([{"role": "user", "content": "duplicate"}], scope) + messages = adapter.get_last_messages(scope) + + assert [message["content"] for message in messages] == ["question", "answer"] + + +class _FakeEmbeddingStore: + def __init__(self): + self.vectors = {} + + def upsert_memory_embedding(self, index_id, vector): + self.vectors[index_id] = vector + + def search_memory_embeddings(self, scope_pattern, query_vector, top_k): + prefix = scope_pattern.removesuffix("%") + return [ + (index_id, 0.9) + for index_id in self.vectors + if index_id.startswith(prefix) + ][:top_k] + + def delete_memory_embedding(self, index_id): + self.vectors.pop(index_id, None) + + +def test_split_store_keeps_payload_in_sql_and_vector_in_tigergraph(store): + indexes = {} + + def resolver(graphname): + return indexes.setdefault(graphname, _FakeEmbeddingStore()) + + adapter = SplitMemoryVectorStore(store, resolver) + payload = { + "user_id": memory_user_id("alice", "g1"), + "run_id": "c1", + "data": "Alice prefers concise answers", + } + adapter.insert([[0.1, 0.2]], [payload], ["fact-1"]) + + fact = store.get_fact("fact-1") + assert fact.text == "Alice prefers concise answers" + assert list(indexes["g1"].vectors.values()) == [[0.1, 0.2]] + assert adapter.search( + "concise", [0.1, 0.2], filters={"user_id": payload["user_id"]} + )[0].id == "fact-1" + assert adapter.search( + "concise", + [0.1, 0.2], + filters={"user_id": memory_user_id("bob", "g1")}, + ) == [] + + +def test_summary_advances_only_forward(store): + store.save_message("alice", "g1", _message("c1", "m1")) + store.update_summary( + "alice", + "g1", + "c1", + summary_text="first", + through_sequence=1, + source_tokens=10, + summary_tokens=2, + model="test", + ) + store.update_summary( + "alice", + "g1", + "c1", + summary_text="stale", + through_sequence=1, + source_tokens=10, + summary_tokens=2, + model="test", + ) + + assert store.get_summary("alice", "g1", "c1").summary == "first" + + +def test_compaction_is_token_triggered_not_message_count(): + many_short_messages = [ + {"role": "user", "content": "ok"}, + {"role": "system", "content": "yes"}, + ] * 20 + assert select_compaction_range( + "", + many_short_messages, + budget_tokens=1000, + ) == [] + + +def test_compaction_preserves_newest_turns_under_token_pressure(): + messages = [ + { + "role": "user" if index % 2 == 0 else "system", + "content": f"{index}-" + "x" * 400, + } + for index in range(12) + ] + selected = select_compaction_range( + "", + messages, + budget_tokens=500, + ) + + assert selected + assert len(selected) <= len(messages) - 4 + assert all(item not in selected for item in messages[-4:]) + + +def test_pinned_mem0_vector_store_contract(): + pytest.importorskip("mem0") + assert importlib.metadata.version("mem0ai") == "2.1.0" + assert not inspect.isabstract(SplitMemoryVectorStore) + + +def test_graphrag_mem0_llm_uses_bound_langchain_model(): + from common.memory.runtime import GraphRAGMem0LLM + + class _FakeModel: + def __init__(self): + self.bound = None + self.invoked = None + + def bind_tools(self, tools=None, tool_choice="auto"): + self.bound = (tools, tool_choice) + return self + + def invoke(self, messages, **kwargs): + self.invoked = messages + return type("Resp", (), { + "content": "ok", + "tool_calls": [{"name": "add_memory", "args": {"data": "x"}}], + })() + + model = _FakeModel() + llm = GraphRAGMem0LLM(type("Cfg", (), {"model": model})()) + result = llm.generate_response( + [{"role": "user", "content": "hi"}], + tools=[{"function": {"name": "add_memory"}}], + ) + + assert model.bound is not None + assert model.invoked == [("human", "hi")] + assert result["tool_calls"][0]["name"] == "add_memory" + + +def test_fact_cap_is_scoped_and_enqueues_vector_deletion(store): + for index in range(3): + store.put_fact( + f"fact-{index}", + username="alice", + graphname="g1", + conversation_id="c1", + text=f"preference {index}", + payload={}, + ) + store.put_fact( + "other-user", + username="bob", + graphname="g1", + conversation_id="c2", + text="other", + payload={}, + ) + + assert store.enforce_fact_cap("alice", "g1", 2) == 1 + assert len(store.list_facts("alice", "g1")) == 2 + assert len(store.list_facts("bob", "g1")) == 1 + assert any(item.operation == "delete" for item in store.pending_outbox()) diff --git a/scripts/drop_legacy_memory_schema.py b/scripts/drop_legacy_memory_schema.py new file mode 100644 index 00000000..7b0df1a6 --- /dev/null +++ b/scripts/drop_legacy_memory_schema.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +"""One-off: drop superseded graph-backed memory types from TigerGraph.""" + +import sys + +import requests +from requests.auth import HTTPBasicAuth + +HOST = "http://tigergraph:14240" +AUTH = HTTPBasicAuth("tigergraph", "tigergraph") +GRAPH = "multihop" +JOB = "drop_legacy_memory" + + +def gsql(command: str, timeout: int = 60) -> str: + print(f"GSQL ({timeout}s): {command.replace(chr(10), ' | ')}", flush=True) + response = requests.post( + f"{HOST}/gsql/v1/statements", + auth=AUTH, + data=command.encode(), + headers={"Content-Type": "text/plain"}, + timeout=timeout, + ) + response.raise_for_status() + text = response.text + print(text, flush=True) + return text + + +def main() -> None: + schema = gsql(f"USE GRAPH {GRAPH}\nls") + needed = ("MemorySession", "MemoryTurn", "MemoryFact", "MEMORY_HAS_TURN", "MEMORY_PRODUCED_FACT") + if not any(name in schema for name in needed): + print("Nothing to drop") + return + + if f"CREATE SCHEMA_CHANGE JOB {JOB}" in schema: + print("Using existing drop job", flush=True) + result = gsql(f"USE GRAPH {GRAPH}\nRUN SCHEMA_CHANGE JOB {JOB}", timeout=300) + else: + result = gsql( + f"USE GRAPH {GRAPH}\n" + f"CREATE SCHEMA_CHANGE JOB {JOB} {{\n" + " DROP EDGE MEMORY_HAS_TURN;\n" + " DROP EDGE MEMORY_PRODUCED_FACT;\n" + " DROP VERTEX MemoryFact;\n" + " DROP VERTEX MemoryTurn;\n" + " DROP VERTEX MemorySession;\n" + "}\n" + f"RUN SCHEMA_CHANGE JOB {JOB}", + timeout=300, + ) + leftover = gsql(f"USE GRAPH {GRAPH}\nls") + remaining = [name for name in needed if name in leftover] + if remaining: + sys.exit(f"Still present: {remaining}\nLast job output:\n{result}") + gsql( + f"USE GRAPH {GRAPH}\n" + f"DROP JOB {JOB}\n" + "DROP JOB repair_graphrag_memory_schema\n" + "DROP JOB add_graphrag_memory_vector" + ) + print("Legacy memory types removed") + + +if __name__ == "__main__": + main()