feat(oss): port v3 pipeline with hybrid search, entity extraction, and additive scoring (#4805)
Co-authored-by: Soumil Rathi <soumilrathi@gmail.com> Co-authored-by: Saket Aryan <saketaryan2002@gmail.com> Co-authored-by: chaithanyak42 <chaithanya.kumar42a@gmail.com> Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
soumil-rathi
co-authored by
Soumil Rathi
Saket Aryan
chaithanyak42
kartik-mem0
Soumil Rathi
Saket Aryan
chaithanyak42
kartik-mem0
parent
57f944e18a
commit
a488e19044
+42
-24
@@ -29,6 +29,9 @@ warnings.filterwarnings("default", category=DeprecationWarning)
|
||||
# Setup user config
|
||||
setup_config()
|
||||
|
||||
# Entity parameters that must be passed via filters, not top-level
|
||||
ENTITY_PARAMS = frozenset({"user_id", "agent_id", "app_id", "run_id"})
|
||||
|
||||
|
||||
class MemoryClient:
|
||||
"""Client for interacting with the Mem0 API.
|
||||
@@ -162,12 +165,9 @@ class MemoryClient:
|
||||
elif not isinstance(messages, list):
|
||||
raise ValueError(f"messages must be str, dict, or list[dict], got {type(messages).__name__}")
|
||||
|
||||
# Force v1.1 format for all add operations
|
||||
kwargs["output_format"] = "v1.1"
|
||||
|
||||
kwargs = self._prepare_params(kwargs)
|
||||
payload = self._prepare_payload(messages, kwargs)
|
||||
response = self.client.post("/v1/memories/", json=payload)
|
||||
response = self.client.post("/v3/memories/", json=payload)
|
||||
response.raise_for_status()
|
||||
if "metadata" in kwargs:
|
||||
del kwargs["metadata"]
|
||||
@@ -204,8 +204,7 @@ class MemoryClient:
|
||||
|
||||
Args:
|
||||
options: Typed options for the get_all operation (GetAllMemoryOptions).
|
||||
**kwargs: Optional parameters for filtering (user_id, agent_id,
|
||||
app_id, top_k, page, page_size).
|
||||
**kwargs: Optional parameters for filtering (filters, page, page_size).
|
||||
|
||||
Returns:
|
||||
A dictionary containing memories in v1.1 format: {"results": [...]}
|
||||
@@ -218,6 +217,14 @@ class MemoryClient:
|
||||
NetworkError: If network connectivity issues occur.
|
||||
MemoryNotFoundError: If the memory doesn't exist (for updates/deletes).
|
||||
"""
|
||||
# Reject top-level entity params - must use filters instead
|
||||
invalid_keys = ENTITY_PARAMS & set(kwargs.keys())
|
||||
if invalid_keys:
|
||||
raise ValueError(
|
||||
f"Top-level entity parameters {invalid_keys} are not supported in get_all(). "
|
||||
f"Use filters={{'user_id': '...'}} instead."
|
||||
)
|
||||
|
||||
kwargs = {**(options.model_dump(exclude_unset=True) if options else {}), **kwargs}
|
||||
params = self._prepare_params(kwargs)
|
||||
|
||||
@@ -267,11 +274,19 @@ class MemoryClient:
|
||||
NetworkError: If network connectivity issues occur.
|
||||
MemoryNotFoundError: If the memory doesn't exist (for updates/deletes).
|
||||
"""
|
||||
# Reject top-level entity params - must use filters instead
|
||||
invalid_keys = ENTITY_PARAMS & set(kwargs.keys())
|
||||
if invalid_keys:
|
||||
raise ValueError(
|
||||
f"Top-level entity parameters {invalid_keys} are not supported in search(). "
|
||||
f"Use filters={{'user_id': '...'}} instead."
|
||||
)
|
||||
|
||||
kwargs = {**(options.model_dump(exclude_unset=True) if options else {}), **kwargs}
|
||||
params = self._prepare_params(kwargs)
|
||||
payload = {"query": query, **params}
|
||||
|
||||
response = self.client.post("/v2/memories/search/", json=payload)
|
||||
response = self.client.post("/v3/memories/search/", json=payload)
|
||||
response.raise_for_status()
|
||||
if "metadata" in kwargs:
|
||||
del kwargs["metadata"]
|
||||
@@ -283,12 +298,7 @@ class MemoryClient:
|
||||
"sync_type": "sync",
|
||||
},
|
||||
)
|
||||
result = response.json()
|
||||
|
||||
# Ensure v1.1 format (wrap raw list if needed)
|
||||
if isinstance(result, list):
|
||||
return {"results": result}
|
||||
return result
|
||||
return response.json()
|
||||
|
||||
@api_error_handler
|
||||
def update(
|
||||
@@ -1083,12 +1093,9 @@ class AsyncMemoryClient:
|
||||
elif not isinstance(messages, list):
|
||||
raise ValueError(f"messages must be str, dict, or list[dict], got {type(messages).__name__}")
|
||||
|
||||
# Force v1.1 format for all add operations
|
||||
kwargs["output_format"] = "v1.1"
|
||||
|
||||
kwargs = self._prepare_params(kwargs)
|
||||
payload = self._prepare_payload(messages, kwargs)
|
||||
response = await self.async_client.post("/v1/memories/", json=payload)
|
||||
response = await self.async_client.post("/v3/memories/", json=payload)
|
||||
response.raise_for_status()
|
||||
if "metadata" in kwargs:
|
||||
del kwargs["metadata"]
|
||||
@@ -1122,6 +1129,14 @@ class AsyncMemoryClient:
|
||||
NetworkError: If network connectivity issues occur.
|
||||
MemoryNotFoundError: If the memory doesn't exist (for updates/deletes).
|
||||
"""
|
||||
# Reject top-level entity params - must use filters instead
|
||||
invalid_keys = ENTITY_PARAMS & set(kwargs.keys())
|
||||
if invalid_keys:
|
||||
raise ValueError(
|
||||
f"Top-level entity parameters {invalid_keys} are not supported in get_all(). "
|
||||
f"Use filters={{'user_id': '...'}} instead."
|
||||
)
|
||||
|
||||
kwargs = {**(options.model_dump(exclude_unset=True) if options else {}), **kwargs}
|
||||
params = self._prepare_params(kwargs)
|
||||
|
||||
@@ -1171,11 +1186,19 @@ class AsyncMemoryClient:
|
||||
NetworkError: If network connectivity issues occur.
|
||||
MemoryNotFoundError: If the memory doesn't exist (for updates/deletes).
|
||||
"""
|
||||
# Reject top-level entity params - must use filters instead
|
||||
invalid_keys = ENTITY_PARAMS & set(kwargs.keys())
|
||||
if invalid_keys:
|
||||
raise ValueError(
|
||||
f"Top-level entity parameters {invalid_keys} are not supported in search(). "
|
||||
f"Use filters={{'user_id': '...'}} instead."
|
||||
)
|
||||
|
||||
kwargs = {**(options.model_dump(exclude_unset=True) if options else {}), **kwargs}
|
||||
params = self._prepare_params(kwargs)
|
||||
payload = {"query": query, **params}
|
||||
|
||||
response = await self.async_client.post("/v2/memories/search/", json=payload)
|
||||
response = await self.async_client.post("/v3/memories/search/", json=payload)
|
||||
response.raise_for_status()
|
||||
if "metadata" in kwargs:
|
||||
del kwargs["metadata"]
|
||||
@@ -1187,12 +1210,7 @@ class AsyncMemoryClient:
|
||||
"sync_type": "async",
|
||||
},
|
||||
)
|
||||
result = response.json()
|
||||
|
||||
# Ensure v1.1 format (wrap raw list if needed)
|
||||
if isinstance(result, list):
|
||||
return {"results": result}
|
||||
return result
|
||||
return response.json()
|
||||
|
||||
@api_error_handler
|
||||
async def update(
|
||||
|
||||
+20
-13
@@ -2,6 +2,9 @@
|
||||
|
||||
These models provide IDE autocompletion, runtime validation, and type safety.
|
||||
Methods accept both typed options and **kwargs for backward compatibility.
|
||||
|
||||
Identity fields (user_id, agent_id, app_id, run_id) must be passed inside
|
||||
the ``filters`` dict — the v3 API does not accept them at the top level.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
@@ -9,18 +12,16 @@ from typing import Any, Dict, List, Optional, Union
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class EntityOptions(BaseModel):
|
||||
"""Identity options for add/delete operations (top-level entity IDs)."""
|
||||
class AddMemoryOptions(BaseModel):
|
||||
"""Options for the add() method.
|
||||
|
||||
user_id: Optional[str] = Field(default=None, description="The user ID to associate with the memory")
|
||||
agent_id: Optional[str] = Field(default=None, description="The agent ID to associate with the memory")
|
||||
app_id: Optional[str] = Field(default=None, description="The app ID to associate with the memory")
|
||||
run_id: Optional[str] = Field(default=None, description="The run ID to associate with the memory")
|
||||
|
||||
|
||||
class AddMemoryOptions(EntityOptions):
|
||||
"""Options for the add() method."""
|
||||
Identity fields (user_id, agent_id, app_id, run_id) must be passed inside
|
||||
the ``filters`` dict — the v3 API does not accept them at the top level.
|
||||
"""
|
||||
|
||||
filters: Optional[Dict[str, Any]] = Field(
|
||||
default=None, description="Filters containing entity IDs (e.g. {'user_id': '...'})"
|
||||
)
|
||||
metadata: Optional[Dict[str, Any]] = Field(default=None, description="Additional metadata for the memory")
|
||||
infer: Optional[bool] = Field(default=None, description="Whether to infer memories from the input")
|
||||
custom_categories: Optional[List[Dict[str, Any]]] = Field(
|
||||
@@ -72,10 +73,16 @@ class GetAllMemoryOptions(BaseModel):
|
||||
categories: Optional[List[str]] = Field(default=None, description="Categories to filter by")
|
||||
|
||||
|
||||
class DeleteAllMemoryOptions(EntityOptions):
|
||||
"""Options for the delete_all() method."""
|
||||
class DeleteAllMemoryOptions(BaseModel):
|
||||
"""Options for the delete_all() method.
|
||||
|
||||
pass
|
||||
Identity fields (user_id, agent_id, app_id, run_id) must be passed inside
|
||||
the ``filters`` dict — the API does not accept them at the top level.
|
||||
"""
|
||||
|
||||
filters: Optional[Dict[str, Any]] = Field(
|
||||
default=None, description="Filters containing entity IDs (e.g. {'user_id': '...'})"
|
||||
)
|
||||
|
||||
|
||||
class UpdateMemoryOptions(BaseModel):
|
||||
|
||||
@@ -5,7 +5,6 @@ from pydantic import BaseModel, Field
|
||||
|
||||
from mem0.configs.rerankers.config import RerankerConfig
|
||||
from mem0.embeddings.configs import EmbedderConfig
|
||||
from mem0.graphs.configs import GraphStoreConfig
|
||||
from mem0.llms.configs import LlmConfig
|
||||
from mem0.vector_stores.configs import VectorStoreConfig
|
||||
|
||||
@@ -44,10 +43,6 @@ class MemoryConfig(BaseModel):
|
||||
description="Path to the history database",
|
||||
default=os.path.join(mem0_dir, "history.db"),
|
||||
)
|
||||
graph_store: GraphStoreConfig = Field(
|
||||
description="Configuration for the graph",
|
||||
default_factory=GraphStoreConfig,
|
||||
)
|
||||
reranker: Optional[RerankerConfig] = Field(
|
||||
description="Configuration for the reranker",
|
||||
default=None,
|
||||
@@ -60,10 +55,6 @@ class MemoryConfig(BaseModel):
|
||||
description="Custom instructions for fact extraction",
|
||||
default=None,
|
||||
)
|
||||
custom_update_memory_prompt: Optional[str] = Field(
|
||||
description="Custom prompt for the update memory",
|
||||
default=None,
|
||||
)
|
||||
|
||||
|
||||
class AzureConfig(BaseModel):
|
||||
|
||||
+604
-1
@@ -1,4 +1,5 @@
|
||||
from datetime import datetime
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
|
||||
MEMORY_ANSWER_PROMPT = """
|
||||
You are an expert at answering questions based on the provided memories. Your task is to provide accurate and concise answers to the questions by leveraging the information given in the memories.
|
||||
@@ -457,3 +458,605 @@ def get_update_memory_messages(retrieved_old_memory_dict, response_content, cust
|
||||
|
||||
Do not return anything except the JSON format.
|
||||
"""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# V3 Additive Extraction Prompt (ADD-only with memory linking)
|
||||
# Ported from platform/backend/shared/core/config/prompts.py
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
ADDITIVE_EXTRACTION_PROMPT = """
|
||||
|
||||
# ROLE
|
||||
|
||||
You are a Memory Extractor — a precise, evidence-bound processor responsible for extracting rich, contextual memories from conversations. Your sole operation is ADD: identify every piece of memorable information and produce self-contained, contextually rich factual statements.
|
||||
|
||||
You extract from BOTH user and assistant messages. User messages reveal personal facts, preferences, plans, and experiences. Assistant messages contain recommendations, plans, suggestions, and actionable information the user may later reference.
|
||||
|
||||
Accuracy and completeness are critical. Every piece of memorable information must be captured — a missed extraction means lost context that degrades future personalization. When a conversation covers multiple topics, extract each one separately. Do not let a dominant topic cause you to miss secondary information.
|
||||
|
||||
# INPUTS
|
||||
|
||||
## New Messages
|
||||
|
||||
The current conversation turn(s) with "role" (user/assistant) and "content".
|
||||
|
||||
Both roles contain extractable information:
|
||||
- **User messages**: Personal facts, preferences, plans, experiences, things done / never done before, opinions, requests, implicit preferences revealed through questions
|
||||
- **Assistant messages**: Specific recommendations given, plans or schedules created, information researched, solutions provided, agreements reached
|
||||
|
||||
Attribute correctly: use "User" for user-stated facts. For assistant-generated content, frame in terms of the user's context (e.g., "User was recommended X" or "User's plan includes X as discussed in conversation").
|
||||
|
||||
Do NOT extract:
|
||||
- Vague assistant characterizations ("you seem passionate", "that sounds stressful") unless the user explicitly confirms them
|
||||
- Generic assistant acknowledgments ("Sure!", "Great question!")
|
||||
- Assistant meta-commentary about its own capabilities
|
||||
|
||||
|
||||
## Summary
|
||||
|
||||
A narrative summary of the user's profile from prior conversations. May be empty for new users. Use it to enrich extractions — it holds established context like names, locations, and relationships.
|
||||
|
||||
|
||||
## Recently Extracted Memories
|
||||
|
||||
Memories already captured from recent messages in this session (up to 20). This is your primary deduplication reference — do not re-extract information already captured here.
|
||||
|
||||
|
||||
## Existing Memories
|
||||
|
||||
Memories currently in the system relevant to this conversation. Formatted as:
|
||||
[{"id": "uuid-string", "text": "..."}, ...]
|
||||
|
||||
Use these ONLY for deduplication and linking — do NOT extract new memories from Existing Memories. Your extractions must come exclusively from New Messages. If new information in New Messages is semantically equivalent to an Existing Memory with no meaningful new context, skip it.
|
||||
|
||||
When a new memory is related to an Existing Memory — same topic, overlapping entities, updated/shifted preference, follow-up event, or continuation of a narrative — include the Existing Memory's ID in the new memory's "linked_memory_ids" array. Your ADD output IDs remain sequential ("0", "1", ...) but linked_memory_ids uses the UUIDs from this list.
|
||||
|
||||
|
||||
IMPORTANT: An existing memory about an entity (e.g., "User has a dog named Max") does NOT mean all information about that entity has been captured. New events, activities, experiences, or details about a known entity MUST still be extracted as separate memories and linked back. Only skip extraction when the specific fact or event itself is already captured — not merely because the entity appears in an existing memory. "User has a dog named Max" and "User went on a camping trip with Max where they hiked and swam" are two distinct memories, not duplicates.
|
||||
|
||||
|
||||
## Last k Messages
|
||||
|
||||
Recent messages (up to 20) preceding New Messages. Use to resolve references and pronouns in New Messages.
|
||||
|
||||
|
||||
## Observation Date
|
||||
|
||||
When the conversation actually took place (e.g., "2023-05-24"). This is your ONLY temporal anchor for resolving time references.
|
||||
|
||||
Resolve ALL relative references against Observation Date:
|
||||
- "yesterday" → day before Observation Date
|
||||
- "last week" → week preceding Observation Date
|
||||
- "next month" → month following Observation Date
|
||||
- "recently" → shortly before Observation Date
|
||||
- "just finished", "today" → on or near Observation Date
|
||||
|
||||
CRITICAL: "User went to Paris last week" is useless 6 months later. "User went to Paris the week of May 15, 2023" is meaningful forever. Always ground relative references to specific dates.
|
||||
|
||||
|
||||
## Current Date
|
||||
|
||||
Today's system date. May be years after Observation Date. Do NOT use this to resolve temporal references in messages — only Observation Date grounds user and assistant statements.
|
||||
|
||||
|
||||
## Optional Inputs
|
||||
|
||||
- **includes**: Topics to focus on
|
||||
- **excludes**: Topics to skip
|
||||
- **custom_instructions**: User-defined rules (highest priority)
|
||||
- **feedback_str**: Adjust extraction based on this feedback
|
||||
|
||||
|
||||
# GUIDELINES
|
||||
|
||||
## What to Extract
|
||||
|
||||
Extract ALL memorable information from both user and assistant messages. Think broadly:
|
||||
|
||||
**From user messages:**
|
||||
- Personal details, preferences, plans, relationships, professional context
|
||||
- Health/wellness, opinions, hobbies, emotional states
|
||||
- Entity attributes (breed, model, color, make, size)
|
||||
- Implicit preferences revealed through requests
|
||||
- **Shared content and reference material** — when a user shares documents, case studies, articles, data, specifications, stat blocks, code, or any structured information, extract the key factual data FROM that content. The user shared it because they want it remembered.
|
||||
- Firsts and milestones — 'first call-out', 'just started', 'recently joined', etc.
|
||||
- Specific foods, meals, and who was present (e.g. 'dinner with mom — salads, sandwiches, homemade desserts').
|
||||
- Inspiration and motivation — what inspired someone to start something, who encouraged them.
|
||||
|
||||
**From assistant messages (ONLY when genuinely new):**
|
||||
- Specific recommendations given (books, restaurants, products, services)
|
||||
- Plans or schedules created for the user
|
||||
- Information researched or provided (facts, instructions, solutions)
|
||||
- Agreements reached during conversation
|
||||
- **Personal facts, experiences, and details shared by named speakers** — in multi-speaker conversations, the "assistant" role may represent a real person sharing their own life (e.g., "Maria: I just got a new cat named Bailey"). Extract their personal information with the same rigor as user-stated facts, attributed to the speaker by name.
|
||||
|
||||
Do NOT extract from assistant messages that merely restate, summarize, or confirm what the user already said. The user's own words are the primary source — if the user said it and the assistant echoed it, extract only once from the user's version. Note: a single assistant message may contain BOTH an echo AND new personal facts — skip the echo portion but still extract the new facts.
|
||||
|
||||
Do NOT extract: greetings, filler, vague acknowledgments, or content too generic to be useful.
|
||||
|
||||
**When in doubt, extract.** A slightly redundant memory is far less costly than a missing one. The deduplication system downstream will handle true duplicates — your job is to ensure nothing meaningful is lost.
|
||||
|
||||
### Casual Topics Are Still Extractable
|
||||
|
||||
Conversations about pets, hobbies, childhood memories, funny anecdotes, and personal preferences are NOT "chitchat" to be skipped. In a personal memory system, these casual revelations are often the MOST valuable — someone's pet's name, a childhood activity with a parent, a funny incident, a new hobby. Only skip messages that are PURELY phatic ("Hi!", "Sounds good!", "Thanks!") with zero informational content.
|
||||
|
||||
### Extract Incidental Facts, Not Just Requests
|
||||
|
||||
When a user asks a question or makes a request, their message often contains INCIDENTAL PERSONAL FACTS stated as context. These facts are just as extractable as the request itself:
|
||||
|
||||
- "I've harvested cherry tomatoes from my garden — any companion plant suggestions?" → Extract BOTH "User grows cherry tomatoes in their garden"
|
||||
- "I just started 'The Nightingale' by Kristin Hannah — can you recommend similar books?" → Extract BOTH "User started reading 'The Nightingale' by Kristin Hannah on [date]"
|
||||
- "As an aspiring stand-up comedian, can you suggest Netflix comedy specials?" → Extract BOTH the career aspiration
|
||||
- "My daughter Sara loves painting — where can I find kids' art classes?" → Extract "User has a daughter named Sara who loves painting"
|
||||
|
||||
Do NOT let the request overshadow the facts. A question about companion plants is transient; the fact that the user grows cherry tomatoes is a persistent personal detail worth remembering.
|
||||
|
||||
**IMPORTANT — Extract ALL dimensions of a conversation.** A single session may contain career facts, entertainment preferences, scheduled plans, and personal opinions. Extract each dimension as a separate memory. Do not let one dominant topic cause you to miss secondary information.
|
||||
|
||||
### Shared Photos and Images
|
||||
|
||||
When a message contains a photo description (e.g., "[Shared photo: ...]" or describes sharing/showing an image), extract factual information from BOTH the surrounding conversation text AND the photo description. The photo description provides visual context that may contain important details:
|
||||
|
||||
- A photo of a group at a park → extract the activity (e.g., "had a picnic at the park")
|
||||
- A photo showing a specific object, place, or person → extract what is depicted
|
||||
- A photo with visible text (signs, posters, book covers) → extract the text content
|
||||
|
||||
## Memory Quality Standards
|
||||
|
||||
### Contextually Rich, Not Atomic
|
||||
Capture the full picture — fact AND surrounding context — in a single unified memory, not scattered fragments.
|
||||
|
||||
Bad: "User has a dog" | Good: "User has a dog named Poppy and their morning walks together are the highlight of their day"
|
||||
|
||||
This applies especially to **transitions and changes**. When the user describes changing, switching, replacing, stopping, or trying something new in place of something else, the memory MUST capture the transition — what the new state is AND what it replaces or changes from. The relationship between old and new is critical context. Without it, the system has an isolated new fact with no understanding of what changed.
|
||||
|
||||
Bad: "User prefers oat milk lattes"
|
||||
Good: "User switched from almond milk to oat milk lattes after developing an almond sensitivity"
|
||||
|
||||
Bad: "User is taking online Spanish classes on Wednesdays"
|
||||
Good: "User switched from in-person French classes to online Spanish classes on Wednesdays after relocating"
|
||||
|
||||
When the change is explicitly temporary or a trial, capture that too — "for a month", "trying out", "testing" — these signal the old arrangement may resume.
|
||||
|
||||
### Clean Factual Statements
|
||||
Preserve the FULL meaning including emotional reactions, motivations, and subjective experiences. Remove filler words and conversation mechanics (greetings, "like", "you know"), but KEEP:
|
||||
- Emotional states: "scared but reassured", "happy and thankful", "liberated and empowered"
|
||||
- Motivations and reasons: "motivated by her own journey and the support she received"
|
||||
- Subjective descriptions: "resilient", "therapeutic", "nerve-wracking"
|
||||
|
||||
### Self-Contained
|
||||
Every memory must be understandable on its own. Replace all pronouns with specific names or "User."
|
||||
|
||||
### Concise but Complete (15-80 words, up to 100 for detail-rich content)
|
||||
1-2 sentences per memory (up to 3 for content with multiple proper nouns, specific quantities, or enumerated items). When a topic has too many details, split into multiple focused memories rather than compressing details away. NEVER sacrifice a proper noun, title, date, or specific detail to meet a word count — completeness beats brevity.
|
||||
|
||||
### Temporally Grounded
|
||||
Preserve exact dates, durations, and temporal relationships. Convert relative → absolute using Observation Date (NOT Current Date). NEVER convert absolute → vague. "18 days" stays "18 days", not "some time."
|
||||
|
||||
### Numerically Precise
|
||||
Preserve exact quantities as stated. "416 pages" stays "416 pages", not "about 400 pages."
|
||||
|
||||
### Preserve Specific Details — Never Generalize Concrete Information
|
||||
|
||||
When information contains specific details — whether quantities, identifiers, descriptions, visual details, quoted text, named objects, proper nouns, or any concrete information — those specifics MUST survive extraction. Replacing a specific detail with a vague category is a critical error.
|
||||
|
||||
#### Proper Nouns and Titles Should be Preserved
|
||||
|
||||
Book titles, movie titles, game names, song titles, restaurant names, neighborhood names, brand names, character names, and named places are the HIGHEST-VALUE details in a memory. Users search by name — a memory without the name is unfindable. ALWAYS preserve exact proper nouns:
|
||||
|
||||
- "watched 'Eternal Sunshine of the Spotless Mind'" → KEEP the full title
|
||||
- "went to Woodhaven for a road trip" → KEEP "Woodhaven"
|
||||
- "tried the new restaurant Osteria Francescana" → KEEP "Osteria Francescana", NOT "a new restaurant"
|
||||
- "reading 'A Court of Thorns and Roses'" → KEEP the title in quotes, NOT "a fantasy book"
|
||||
- "his favorite character is Aragorn from Lord of the Rings" → KEEP "Aragorn" and "Lord of the Rings"
|
||||
|
||||
#### Qualifiers and Specific Attributes Are Essential
|
||||
|
||||
Never generalize specific qualifiers. The qualifier is almost always the detail that matters most for recall:
|
||||
|
||||
- "promoted to assistant manager" → KEEP "assistant manager", NOT "manager"
|
||||
- "ordered grilled salmon and roasted vegetables" → KEEP "grilled salmon and roasted vegetables", NOT "healthy meal"
|
||||
- "started doing aerial yoga" → KEEP "aerial yoga", NOT "yoga" or "a workout class"
|
||||
- "painted a forest scene in watercolors" → KEEP "a forest scene in watercolors", NOT "started painting"
|
||||
- "drove a Ferrari 488 GTB" → KEEP "Ferrari 488 GTB", NOT "sports car"
|
||||
- "scored 3 goals in the semifinal" → KEEP "3 goals in the semifinal", NOT "scored several goals"
|
||||
- "walks her dogs multiple times a day" → KEEP "multiple times a day", NOT "regularly" or "daily"
|
||||
|
||||
If the input is specific, the memory must be equally specific. The concrete details are precisely what distinguishes a useful memory from a useless one. NEVER replace a specific noun, number, title, or description with a vague category or paraphrase — this destroys the information the user actually shared.
|
||||
|
||||
### Meaning-Preserving
|
||||
Capture the EXACT meaning of what was said. Read carefully:
|
||||
- "Didn't get to bed until 2 AM" = went TO BED at 2 AM (late bedtime), NOT "slept until 2 AM" (late wakeup)
|
||||
- "Can't stop eating chocolate" = eats a lot of chocolate, NOT has stopped eating chocolate
|
||||
- "I used to love hiking" = no longer loves hiking, NOT currently loves hiking
|
||||
|
||||
Misinterpreting the user's words is worse than not extracting at all.
|
||||
|
||||
|
||||
## Integrity Rules
|
||||
|
||||
- **No Fabrication**: Every detail must trace to the inputs. If you can't point to where it came from, don't include it.
|
||||
- **No Implicit Attribute Inference**: Don't infer gender, age, ethnicity, etc. from names or context. Only record explicitly stated attributes.
|
||||
- **Correct Attribution**: Distinguish user-stated facts from assistant-provided information. Frame assistant content appropriately.
|
||||
- **No Echo Extraction**: When an assistant message restates, summarizes, or confirms information the user already provided in the same conversation, do NOT extract it again from the assistant's message. Only extract from assistant messages when they contribute genuinely NEW information not already present in the user's messages — specific recommendations, newly created plans or schedules, researched facts, or solutions the assistant provided that the user did not state themselves. If the user says "I want daily check-ins at 7:30 AM" and the assistant responds "I've set up daily check-ins at 7:30 AM", that is already captured from the user's message — do not extract a second memory from the assistant's echo.
|
||||
- **No Within-Response Duplication**: Each piece of information must appear exactly ONCE in your output, regardless of how many messages mention it. Before finalizing your output, review your extractions and remove any that are semantically equivalent to another extraction in the same response. Two memories about the same fact phrased differently are redundant — keep the richer one and drop the other.
|
||||
- **No Meta-Extraction**: Extract the CONTENT of what was shared, not a description of the user's action. When a user shares a document, data, or reference material, extract the actual facts FROM that material.
|
||||
- WRONG: "User asked for the introductory paragraph to be shortened" / "User shared a case summary for optimization"
|
||||
- RIGHT: "The Bajimaya v Reward Homes case involved construction starting in 2014, contract signed in 2015, with completion due by October 2015" / "The tribunal found Reward Homes breached its contract through poor workmanship, waterproofing defects, and non-compliance with the Building Code of Australia"
|
||||
- WRONG: "Assistant created a D&D adventure with enemies"
|
||||
- RIGHT: "The Lost Temple of the Djinn adventure includes 4 Mummies (AC 11, 45 HP), 2 Construct Guardians (AC 17, 110 HP), and 6 Skeletal Warriors (AC 12, 22 HP)"
|
||||
- **No Detail Contamination from Context**: When extracting from New Messages, do NOT import or merge details from Existing Memories or Recent Memories into the new extraction UNLESS the new message explicitly references those details. If the New Message says "I had a great meal" and an Existing Memory says "User's favorite restaurant is Olive Garden," do NOT produce "User had a great meal at Olive Garden" — the new message never mentioned the restaurant. Each extraction must be faithful to its source message only.
|
||||
|
||||
|
||||
## Memory Linking
|
||||
|
||||
When extracting a new memory, check if it relates to any Existing Memory. Add related Existing Memory IDs to "linked_memory_ids". Link when:
|
||||
|
||||
- **Same entity/topic**: New fact about a person, place, or thing already mentioned
|
||||
- **Updated preference**: A changed or evolved opinion on something previously captured
|
||||
- **Continuation**: Follow-up event or next step in a previously captured narrative
|
||||
- **Contradiction**: New information that conflicts with an existing memory
|
||||
|
||||
Do NOT link memories that merely share a vague theme. Links should be specific and meaningful — the linked memories should be about the same specific entity, event, or topic. If no existing memories are related, omit linked_memory_ids or pass an empty array.
|
||||
|
||||
|
||||
# EXAMPLES
|
||||
|
||||
|
||||
## Example 1: Multi-Topic Extraction
|
||||
|
||||
Summary: ""
|
||||
Recently Extracted: []
|
||||
Existing Memories: []
|
||||
New Messages:
|
||||
[{"role": "user", "content": "Hey! I'm Marcus. I just got promoted to Senior Engineer at Shopify last week - been grinding for two years for this. My wife Elena and I celebrated with dinner at Osteria Francescana, it's our go-to spot for special occasions. We're also expecting our first baby in March!"},
|
||||
{"role": "assistant", "content": "Congratulations on everything, Marcus! What exciting times."}]
|
||||
Observation Date: 2025-08-19
|
||||
|
||||
Output:
|
||||
{"memory": [
|
||||
{"id": "0", "text": "User's name is Marcus and was promoted to Senior Engineer at Shopify around August 12, 2025 after working toward it for two years"},
|
||||
{"id": "1", "text": "Marcus has a wife named Elena and they celebrate special occasions at Osteria Francescana, their go-to restaurant"},
|
||||
{"id": "2", "text": "Marcus and his wife Elena are expecting their first baby in March 2026"}
|
||||
]}
|
||||
|
||||
Three distinct topics — career, relationship/dining, family milestone — each get their own memory with full context.
|
||||
|
||||
|
||||
## Example 2: Extracting from Assistant Recommendations
|
||||
|
||||
Summary: "User is an aspiring stand-up comedian interested in improving their craft."
|
||||
Recently Extracted: []
|
||||
Existing Memories: []
|
||||
New Messages:
|
||||
[{"role": "user", "content": "Can you recommend some sports documentaries on Netflix with strong storytelling? I love \"The Last Dance\" by Michael Jordan."},
|
||||
{"role": "assistant", "content": "Great taste! Here are some Netflix documentaries known for their storytelling: 1) \"Formula 1: Drive to Survive\" (behind the scenes of Formula 1 racing) 2) \"Athlete A\" (investigative look at USA Gymnastics) 3) \"The Battered Bastards of Baseball\" (independent baseball story). All focus on powerful, narrative-driven sports stories."}]
|
||||
Observation Date: 2023-06-01
|
||||
|
||||
Output:
|
||||
{"memory": [
|
||||
{"id": "0", "text": "User enjoys watching sports documentaries on Netflix with strong storytelling, such as 'The Last Dance' featuring Michael Jordan"},
|
||||
{"id": "1", "text": "User was recommended the following sports documentaries on Netflix for storytelling: 'Formula 1: Drive to Survive', 'Athlete A', and 'The Battered Bastards of Baseball'"}
|
||||
]}
|
||||
|
||||
The user's viewing preference (Netflix stand-up comedy) is extracted alongside the assistant's specific recommendations. Both are valuable for future personalization.
|
||||
|
||||
|
||||
## Example 3: Nothing to Extract
|
||||
|
||||
Summary: "User is a product manager named David."
|
||||
Existing Memories: [{"id": "0", "text": "David is a product manager at a fintech startup"}]
|
||||
New Messages:
|
||||
[{"role": "user", "content": "Hey, good morning!"},
|
||||
{"role": "assistant", "content": "Good morning, David! How can I help you today?"}]
|
||||
Observation Date: 2025-08-19
|
||||
|
||||
Output: {"memory": []}
|
||||
|
||||
## Example 5: Deduplication — Skip Already Captured
|
||||
|
||||
Recently Extracted: ["Marcus was promoted to Senior Engineer at Shopify around August 12, 2025"]
|
||||
Existing Memories: [{"id": "0", "text": "Marcus was promoted to Senior Engineer at Shopify around August 12, 2025"}]
|
||||
New Messages:
|
||||
[{"role": "user", "content": "Still can't believe I got the senior engineer promotion at Shopify!"}]
|
||||
Observation Date: 2025-08-19
|
||||
|
||||
Output: {"memory": []}
|
||||
|
||||
|
||||
## Example 6: Extract ALL Dimensions — Don't Miss Secondary Info
|
||||
|
||||
Summary: "User is an aspiring actor."
|
||||
Recently Extracted: []
|
||||
Existing Memories: []
|
||||
New Messages:
|
||||
[{"role": "user", "content": "As an aspiring actor, I'm looking for advice on improving my craft. Can you recommend some films on Netflix with strong acting performances like Daniel Day-Lewis in 'There Will Be Blood'? I also want to find online resources for acting techniques."},
|
||||
{"role": "assistant", "content": "For Netflix films with great acting, check out 'Marriage Story' and 'The Irishman'. For acting techniques, I'd recommend 'An Actor Prepares' by Stanislavski and the MasterClass by Helen Mirren."}]
|
||||
Observation Date: 2023-06-01
|
||||
|
||||
Output:
|
||||
{"memory": [
|
||||
{"id": "0", "text": "User is an aspiring actor seeking to improve their craft through studying films with strong performances and acting technique resources"},
|
||||
{"id": "1", "text": "User enjoys watching films on Netflix with outstanding acting, especially performances like Daniel Day-Lewis in 'There Will Be Blood'"},
|
||||
{"id": "2", "text": "User was recommended 'Marriage Story' and 'The Irishman' for performance study, 'An Actor Prepares' by Stanislavski, and Helen Mirren's MasterClass for acting techniques"}
|
||||
]}
|
||||
|
||||
Three dimensions: (1) career aspiration, (2) entertainment viewing preference, (3) specific recommendations. Each extracted separately.
|
||||
|
||||
|
||||
## Example 7: Vague Temporal References with Historical Observation Date
|
||||
|
||||
Recently Extracted: ["User started reading 'The Hitchhiker's Guide to the Galaxy' on January 16, 2022"]
|
||||
Existing Memories: [{"id": "0", "text": "User started reading 'The Hitchhiker's Guide to the Galaxy' on January 16, 2022"}]
|
||||
New Messages:
|
||||
[{"role": "user", "content": "I've actually listened to Ready Player One as an audiobook recently and enjoyed the pop culture references."}]
|
||||
Observation Date: 2022-01-16
|
||||
Current Date: 2026-02-18
|
||||
|
||||
Output:
|
||||
{"memory": [{"id": "0", "text": "User listened to the Ready Player One audiobook around early January 2022 and enjoyed the pop culture references"}]}
|
||||
|
||||
"Recently" is grounded to the Observation Date (January 2022), NOT Current Date (February 2026). The Hitchhiker's Guide memory already exists — not re-extracted.
|
||||
|
||||
|
||||
## Example 8: Document / Reference Material — Extract Content, Not Actions
|
||||
|
||||
Summary: ""
|
||||
Recently Extracted: []
|
||||
Existing Memories: []
|
||||
New Messages:
|
||||
[{"role": "user", "content": "I want you to remember this case. If you understand, just say acknowledged. Bajimaya v Reward Homes Pty Ltd [2021] NSWCATAP 297 — The construction began in 2014, contract signed in 2015 with completion due by October 2015. The plaintiff received keys in December 2016 and found defects including incomplete works, poor workmanship, and non-compliance with the building code. The tribunal found the builder breached contract."},
|
||||
{"role": "assistant", "content": "Acknowledged."}]
|
||||
Observation Date: 2024-03-10
|
||||
|
||||
Output:
|
||||
{"memory": [
|
||||
{"id": "0", "text": "Bajimaya v Reward Homes Pty Ltd [2021] NSWCATAP 297: construction of the home began in 2014, contract signed in 2015, with completion due by October 2015. Keys were delivered in December 2016."},
|
||||
{"id": "1", "text": "In Bajimaya v Reward Homes, the plaintiff found defects including incomplete works, poor workmanship, and non-compliance with the Building Code of Australia after receiving the home in December 2016."},
|
||||
{"id": "2", "text": "The tribunal found Reward Homes Pty Ltd breached its contract with Mr. Bajimaya by failing to complete work in a proper and workmanlike manner and failing to comply with plans, specifications, and the Building Code."}
|
||||
]}
|
||||
|
||||
The user shared reference material to be remembered. Extract the actual factual content — dates, parties, findings — NOT "User shared a case summary" or "User asked to remember a case."
|
||||
|
||||
|
||||
## Example 9: Structured Data with Counts and Specifics
|
||||
|
||||
Summary: ""
|
||||
Recently Extracted: []
|
||||
Existing Memories: []
|
||||
New Messages:
|
||||
[{"role": "user", "content": "Here are the enemy stat blocks for our D&D campaign: Mummies (4): AC 11, HP 45, Speed 20 ft, with Curse of the Pharaohs (DC 15 Wisdom) and Mummy Rot (DC 15 Constitution). Construct Guardians (2): AC 17, HP 110, Speed 30 ft, with Immutable Form, Magic Resistance, and Siege Monster. Skeletal Warriors (6): AC 12, HP 22, Speed 30 ft, with Undead Fortitude."},
|
||||
{"role": "assistant", "content": "Got it! I've noted all the stat blocks. Ready when you want to start the encounter."}]
|
||||
Observation Date: 2024-01-15
|
||||
|
||||
Output:
|
||||
{"memory": [
|
||||
{"id": "0", "text": "User's D&D campaign encounter includes 4 Mummies (AC 11, 45 HP, Speed 20 ft) with Curse of the Pharaohs (DC 15 Wisdom save) and Mummy Rot (DC 15 Constitution save)"},
|
||||
{"id": "1", "text": "User's D&D campaign encounter includes 2 Construct Guardians (AC 17, 110 HP, Speed 30 ft) with Immutable Form, Magic Resistance, and Siege Monster traits"},
|
||||
{"id": "2", "text": "User's D&D campaign encounter includes 6 Skeletal Warriors (AC 12, 22 HP, Speed 30 ft) with the Undead Fortitude trait"}
|
||||
]}
|
||||
|
||||
Every count (4 Mummies, 2 Construct Guardians, 6 Skeletal Warriors) and every specific value (AC, HP, DCs, trait names) is preserved. Dropping the counts or stat values would destroy the most queryable information.
|
||||
|
||||
|
||||
## Example 10: Memory Linking — Connecting Related Memories
|
||||
|
||||
Summary: ""
|
||||
Recently Extracted: []
|
||||
Existing Memories: [{"id": "a1b2c3d4-5678-9abc-def0-111111111111", "text": "User has a dog named Poppy, a golden retriever"}, {"id": "b2c3d4e5-6789-abcd-ef01-222222222222", "text": "User works as a Senior Engineer at Shopify"}]
|
||||
New Messages:
|
||||
[{"role": "user", "content": "Poppy had her vet checkup yesterday — she's healthy but needs to lose a few pounds. Also, I'm switching teams at work next month to the payments platform."}]
|
||||
Observation Date: 2025-03-15
|
||||
|
||||
Output:
|
||||
{"memory": [
|
||||
{"id": "0", "text": "User's dog Poppy had a vet checkup around March 14, 2025, is healthy but needs to lose weight", "linked_memory_ids": ["a1b2c3d4-5678-9abc-def0-111111111111"]},
|
||||
{"id": "1", "text": "User is switching teams at Shopify to the payments platform in April 2025", "linked_memory_ids": ["b2c3d4e5-6789-abcd-ef01-222222222222"]}
|
||||
]}
|
||||
|
||||
Both new memories link to related existing memories — the vet checkup links to the existing Poppy memory, and the team switch links to the existing Shopify memory. This enables the system to build a graph of related memories.
|
||||
|
||||
|
||||
## Example 11: Long Multi-Topic Conversation — Don't Stop After First Topic
|
||||
|
||||
Summary: ""
|
||||
Recently Extracted: []
|
||||
Existing Memories: []
|
||||
New Messages:
|
||||
[{"role": "user", "content": "I adopted a puppy named Max last weekend! He's a beagle mix."},
|
||||
{"role": "assistant", "content": "Congratulations! How's he settling in?"},
|
||||
{"role": "user", "content": "Great! Oh, and I also started pottery classes on Tuesdays. Made a mug with my daughter's face on it."},
|
||||
{"role": "assistant", "content": "Fun! Sounds like a lot going on."},
|
||||
{"role": "user", "content": "Yeah — my sister just moved to Portland too. I'm happy but honestly a bit overwhelmed. My boss gave me a promotion to team lead last week as well."}]
|
||||
Observation Date: 2025-03-10
|
||||
|
||||
Output:
|
||||
{"memory": [
|
||||
{"id": "0", "text": "User adopted a beagle mix puppy named Max around March 1-2, 2025"},
|
||||
{"id": "1", "text": "User started taking pottery classes on Tuesdays"},
|
||||
{"id": "2", "text": "User made a ceramic mug with their daughter's face on it in pottery class"},
|
||||
{"id": "3", "text": "User's sister recently moved to Portland"},
|
||||
{"id": "4", "text": "User was promoted to team lead around March 3, 2025, and feels happy but overwhelmed about all the recent changes"}
|
||||
]}
|
||||
|
||||
FIVE topics across 5 messages — each one extracted separately. Do not stop after the first topic (the puppy). The pottery mug detail, the sister's move, and the emotional reaction to the promotion are all distinct, extractable facts.
|
||||
|
||||
|
||||
## Example 12: Multi-Speaker Conversation — Extract From ALL Speakers
|
||||
|
||||
Summary: "John has a dog named Max."
|
||||
Recently Extracted: []
|
||||
Existing Memories: [{"id": "a1b2c3d4-0000-0000-0000-111111111111", "text": "John has a dog named Max"}]
|
||||
New Messages:
|
||||
[{"role": "user", "content": "John: Max and I had a blast on our camping trip last summer. We hiked, swam, and made great memories. It was a really peaceful experience."},
|
||||
{"role": "assistant", "content": "Maria: That sounds amazing! I actually just got a new cat named Bailey last week — she's been such a joy already. Camping with pets is so soul-nourishing."},
|
||||
{"role": "user", "content": "John: Congrats on Bailey! Here's a picture of my family too — that was from a trip we took for my daughter Sara's birthday last fall."}]
|
||||
Observation Date: 2023-08-11
|
||||
|
||||
Output:
|
||||
{"memory": [
|
||||
{"id": "0", "text": "John and his dog Max went on a camping trip in the summer of 2023 where they hiked, swam, and found it a peaceful experience", "linked_memory_ids": ["a1b2c3d4-0000-0000-0000-111111111111"]},
|
||||
{"id": "1", "text": "Maria got a new cat named Bailey around early August 2023 and describes her as a joy"},
|
||||
{"id": "2", "text": "John has a daughter named Sara and the family took a trip for her birthday in fall 2022"}
|
||||
]}
|
||||
|
||||
Three key lessons: (1) The existing memory "John has a dog named Max" does NOT mean all Max-related information is captured — the camping trip is a new event with specific activities (hiking, swimming) and must be extracted and linked. (2) Maria is a named speaker in the "assistant" role but shares a genuine personal fact (new cat Bailey) — this MUST be extracted with the same rigor as user facts. Her echo ("that sounds amazing", "camping is soul-nourishing") is correctly skipped, but her personal fact is not. (3) Sara's name and the birthday trip are separate factual details that each deserve their own extraction.
|
||||
|
||||
|
||||
# CRITICAL: Exhaustive Extraction Checklist
|
||||
|
||||
Before producing output, mentally scan the ENTIRE conversation — every single message — and verify:
|
||||
1. Have you extracted at least one memory from every distinct topic or subject change in the conversation?
|
||||
2. Have you extracted facts from messages in the MIDDLE and END of the conversation, not just the beginning?
|
||||
3. For conversations with 10+ messages, you should typically extract 5-15 memories. If you have fewer than 3, re-read the conversation — you are almost certainly missing information.
|
||||
4. Re-read each user message individually: does EVERY specific fact, preference, experience, or event mentioned in that message have a corresponding extraction? If a single message mentions two distinct facts (e.g., an allergy AND a hobby), both must be captured.
|
||||
|
||||
A common failure mode is "first topic dominance" — the extractor captures the first major topic thoroughly, then treats subsequent topics as filler. This is WRONG. Every topic mentioned deserves extraction if it contains memorable facts. If a chunk has 8 messages covering 4 different topics, you MUST produce memories for all 4 topics — not just the first or most prominent one.
|
||||
|
||||
|
||||
# OUTPUT FORMAT
|
||||
|
||||
Return ONLY valid JSON parsable by json.loads(). No text, reasoning, explanations, or wrappers.
|
||||
|
||||
## Structure
|
||||
|
||||
{
|
||||
"memory": [
|
||||
{"id": "0", "text": "First extracted memory", "attributed_to": "user", "linked_memory_ids": ["uuid-of-related-existing-memory"]},
|
||||
{"id": "1", "text": "Second extracted memory", "attributed_to": "assistant"}
|
||||
]
|
||||
}
|
||||
|
||||
## Fields
|
||||
|
||||
- **id** (string, required): Sequential integers as strings starting at "0".
|
||||
- **text** (string, required): A contextually rich, self-contained factual statement (15-80 words).
|
||||
- **attributed_to** (string, required): Who this memory is about. Use "user" for facts stated by or about the user (preferences, plans, personal facts). Use "assistant" for information provided by the assistant (recommendations, confirmations, plans created, information researched).
|
||||
- **linked_memory_ids** (array of strings, optional): IDs of Existing Memories that this new memory relates to. Use the exact IDs from the Existing Memories list. Omit or pass [] if no existing memories are related.
|
||||
|
||||
## Rules
|
||||
|
||||
- Extract every piece of memorable information as a separate memory object.
|
||||
- If nothing is worth extracting, return: {"memory": []}
|
||||
- No duplicate IDs. Use double quotes. No trailing commas.
|
||||
|
||||
"""
|
||||
|
||||
|
||||
AGENT_CONTEXT_SUFFIX = """
|
||||
|
||||
## Entity Context
|
||||
|
||||
The primary entity is an AI agent. Frame memories from the agent's perspective:
|
||||
- For user-stated facts, frame as agent knowledge: "Agent was informed that [fact]" or "Agent learned that [fact]"
|
||||
- For agent actions, use direct statements: "Agent recommended [X]" or "Agent specializes in [domain]"
|
||||
- For agent configuration or instructions, capture directly: "Agent is configured to [behavior]"
|
||||
|
||||
The attributed_to field should still reflect the original source: "user" for facts the user stated, "assistant" for things the agent said or did.
|
||||
"""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# V3 Prompt Builder — constructs the user-side prompt for additive extraction
|
||||
# Ported from platform/backend/shared/core/utils/prompt_builder.py
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
PAST_MESSAGE_TRUNCATION_LIMIT = 300
|
||||
|
||||
|
||||
def _truncate_content(text, limit=PAST_MESSAGE_TRUNCATION_LIMIT):
|
||||
"""Truncate text to limit characters, appending '...' when shortened."""
|
||||
if len(text) <= limit:
|
||||
return text
|
||||
return text[:limit] + "..."
|
||||
|
||||
|
||||
def _format_summary(summary):
|
||||
"""Extract summary text from a string or dict with a 'summary' key."""
|
||||
if isinstance(summary, dict):
|
||||
return summary.get("summary", "")
|
||||
return summary or ""
|
||||
|
||||
|
||||
def _format_conversation_history(messages):
|
||||
"""Format message dicts into 'role: content' lines with truncation."""
|
||||
if not messages:
|
||||
return ""
|
||||
result = ""
|
||||
for msg in messages:
|
||||
role = msg.get("role", "")
|
||||
content = msg.get("message") or msg.get("content", "")
|
||||
if role and content:
|
||||
result += f"{role}: {_truncate_content(content)}\n"
|
||||
return result
|
||||
|
||||
|
||||
def _serialize_memories(memories):
|
||||
"""JSON-serialize a list of memory objects, defaulting to '[]'."""
|
||||
return json.dumps(memories or [], ensure_ascii=False)
|
||||
|
||||
|
||||
def _format_new_messages(new_messages):
|
||||
"""Pass through if already a string, otherwise JSON-serialize."""
|
||||
if isinstance(new_messages, str):
|
||||
return new_messages
|
||||
return json.dumps(new_messages or [], ensure_ascii=False)
|
||||
|
||||
|
||||
def _resolve_dates(current_date=None, observation_date=None):
|
||||
"""Resolve current and observation dates, defaulting to today."""
|
||||
if current_date is None:
|
||||
current_date = datetime.now(timezone.utc).date().isoformat()
|
||||
if observation_date is None:
|
||||
observation_date = current_date
|
||||
return current_date, observation_date
|
||||
|
||||
|
||||
def generate_additive_extraction_prompt(
|
||||
summary=None,
|
||||
recently_extracted_memories=None,
|
||||
existing_memories=None,
|
||||
new_messages=None,
|
||||
*,
|
||||
last_k_messages=None,
|
||||
current_date=None,
|
||||
timestamp=None,
|
||||
custom_instructions=None,
|
||||
use_input_language=False,
|
||||
):
|
||||
"""Build the user prompt for additive (ADD-only) extraction with linking.
|
||||
|
||||
Pairs with ADDITIVE_EXTRACTION_PROMPT system prompt.
|
||||
The LLM will produce only ADD operations, with optional linked_memory_ids.
|
||||
"""
|
||||
current_date, observation_date = _resolve_dates(current_date, timestamp)
|
||||
|
||||
sections = []
|
||||
sections.append(f"## Summary\n{_format_summary(summary)}")
|
||||
sections.append(f"## Last k Messages\n{_format_conversation_history(last_k_messages)}")
|
||||
sections.append(f"## Recently Extracted Memories\n{_serialize_memories(recently_extracted_memories)}")
|
||||
sections.append(f"## Existing Memories\n{_serialize_memories(existing_memories)}")
|
||||
sections.append(f"## New Messages\n{_format_new_messages(new_messages)}")
|
||||
sections.append(f"## Observation Date\n{observation_date}")
|
||||
sections.append(f"## Current Date\n{current_date}")
|
||||
|
||||
if custom_instructions:
|
||||
sections.append(f"## Custom Instructions\n{custom_instructions}")
|
||||
|
||||
if use_input_language:
|
||||
sections.append(
|
||||
"## Language Requirement\n"
|
||||
"CRITICAL: Respond in the SAME LANGUAGE and SCRIPT as the input messages.\n"
|
||||
"1. Match the language of the user's messages exactly — if they write in Korean, extract in Korean; Japanese in Japanese; etc.\n"
|
||||
"2. Preserve the exact script/alphabet of the input.\n"
|
||||
"3. Do NOT translate or transliterate into English unless the input is already in English.\n"
|
||||
"4. Maintain all quality standards (contextual richness, temporal grounding, etc.) regardless of language.\n"
|
||||
"5. Technical terms, proper nouns, and brand names should be preserved in their original form as used in the input.\n"
|
||||
"6. If the input mixes languages (e.g., Hinglish), preserve both the mixed language style AND the script.\n"
|
||||
"7. For Japanese: explicitly resolve omitted subjects using conversation context.\n"
|
||||
"8. For CJK languages: maintain appropriate formality level from the source text."
|
||||
)
|
||||
|
||||
sections.append("# Output:")
|
||||
return "\n\n".join(sections)
|
||||
|
||||
@@ -53,3 +53,20 @@ class AzureOpenAIEmbedding(EmbeddingBase):
|
||||
"""
|
||||
text = text.replace("\n", " ")
|
||||
return self.client.embeddings.create(input=[text], model=self.config.model).data[0].embedding
|
||||
|
||||
def embed_batch(self, texts, memory_action="add"):
|
||||
"""Embed multiple texts in a single Azure OpenAI API call.
|
||||
|
||||
Automatically chunks into batches of 100 to stay within API limits.
|
||||
"""
|
||||
MAX_BATCH = 100
|
||||
texts = [text.replace("\n", " ") for text in texts]
|
||||
all_embeddings = []
|
||||
for i in range(0, len(texts), MAX_BATCH):
|
||||
chunk = texts[i : i + MAX_BATCH]
|
||||
response = self.client.embeddings.create(
|
||||
input=chunk,
|
||||
model=self.config.model,
|
||||
)
|
||||
all_embeddings.extend(item.embedding for item in sorted(response.data, key=lambda x: x.index))
|
||||
return all_embeddings
|
||||
|
||||
@@ -29,3 +29,19 @@ class EmbeddingBase(ABC):
|
||||
list: The embedding vector.
|
||||
"""
|
||||
pass
|
||||
|
||||
def embed_batch(self, texts, memory_action="add"):
|
||||
"""Embed multiple texts. Override in subclasses for native batch support.
|
||||
|
||||
Default implementation calls embed() sequentially for each text.
|
||||
Subclasses with native batch APIs (e.g., OpenAI) should override
|
||||
this for better performance.
|
||||
|
||||
Args:
|
||||
texts: List of text strings to embed.
|
||||
memory_action: The action context ("add", "search", "update").
|
||||
|
||||
Returns:
|
||||
List of embedding vectors (list of floats), one per input text.
|
||||
"""
|
||||
return [self.embed(text, memory_action) for text in texts]
|
||||
|
||||
@@ -53,3 +53,24 @@ class OpenAIEmbedding(EmbeddingBase):
|
||||
if self._pass_dimensions_to_api:
|
||||
kwargs["dimensions"] = self.config.embedding_dims
|
||||
return self.client.embeddings.create(**kwargs).data[0].embedding
|
||||
|
||||
def embed_batch(self, texts, memory_action="add"):
|
||||
"""Embed multiple texts in a single OpenAI API call.
|
||||
|
||||
Automatically chunks into batches of 100 to stay within API limits.
|
||||
"""
|
||||
MAX_BATCH = 100
|
||||
texts = [text.replace("\n", " ") for text in texts]
|
||||
all_embeddings = []
|
||||
for i in range(0, len(texts), MAX_BATCH):
|
||||
chunk = texts[i : i + MAX_BATCH]
|
||||
kwargs = {
|
||||
"input": chunk,
|
||||
"model": self.config.model,
|
||||
"encoding_format": "float",
|
||||
}
|
||||
if self._pass_dimensions_to_api:
|
||||
kwargs["dimensions"] = self.config.embedding_dims
|
||||
response = self.client.embeddings.create(**kwargs)
|
||||
all_embeddings.extend(item.embedding for item in sorted(response.data, key=lambda x: x.index))
|
||||
return all_embeddings
|
||||
|
||||
@@ -321,24 +321,6 @@ class VectorStoreError(MemoryError):
|
||||
super().__init__(message, error_code, details, suggestion, debug_info)
|
||||
|
||||
|
||||
class GraphStoreError(MemoryError):
|
||||
"""Raised when graph store operations fail.
|
||||
|
||||
This exception is raised when graph store operations fail,
|
||||
such as relationship creation, entity management, or graph queries.
|
||||
|
||||
Example:
|
||||
raise GraphStoreError(
|
||||
message="Graph store operation failed",
|
||||
error_code="GRAPH_001",
|
||||
details={"operation": "create_relationship", "entity": "user_123"},
|
||||
suggestion="Please check your graph store configuration and connection"
|
||||
)
|
||||
"""
|
||||
def __init__(self, message: str, error_code: str = "GRAPH_001", details: dict = None,
|
||||
suggestion: str = "Please check your graph store configuration and connection",
|
||||
debug_info: dict = None):
|
||||
super().__init__(message, error_code, details, suggestion, debug_info)
|
||||
|
||||
|
||||
class EmbeddingError(MemoryError):
|
||||
|
||||
@@ -1,136 +0,0 @@
|
||||
from typing import Optional, Union
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator, model_validator
|
||||
|
||||
from mem0.llms.configs import LlmConfig
|
||||
|
||||
|
||||
class Neo4jConfig(BaseModel):
|
||||
url: Optional[str] = Field(None, description="Host address for the graph database")
|
||||
username: Optional[str] = Field(None, description="Username for the graph database")
|
||||
password: Optional[str] = Field(None, description="Password for the graph database")
|
||||
database: Optional[str] = Field(None, description="Database for the graph database")
|
||||
base_label: Optional[bool] = Field(None, description="Whether to use base node label __Entity__ for all entities")
|
||||
|
||||
@model_validator(mode="before")
|
||||
def check_host_port_or_path(cls, values):
|
||||
url, username, password = (
|
||||
values.get("url"),
|
||||
values.get("username"),
|
||||
values.get("password"),
|
||||
)
|
||||
if not url or not username or not password:
|
||||
raise ValueError("Please provide 'url', 'username' and 'password'.")
|
||||
return values
|
||||
|
||||
|
||||
class MemgraphConfig(BaseModel):
|
||||
url: Optional[str] = Field(None, description="Host address for the graph database")
|
||||
username: Optional[str] = Field(None, description="Username for the graph database")
|
||||
password: Optional[str] = Field(None, description="Password for the graph database")
|
||||
|
||||
@model_validator(mode="before")
|
||||
def check_host_port_or_path(cls, values):
|
||||
url, username, password = (
|
||||
values.get("url"),
|
||||
values.get("username"),
|
||||
values.get("password"),
|
||||
)
|
||||
if not url or not username or not password:
|
||||
raise ValueError("Please provide 'url', 'username' and 'password'.")
|
||||
return values
|
||||
|
||||
|
||||
class NeptuneConfig(BaseModel):
|
||||
app_id: Optional[str] = Field("Mem0", description="APP_ID for the connection")
|
||||
endpoint: Optional[str] = (
|
||||
Field(
|
||||
None,
|
||||
description="Endpoint to connect to a Neptune-DB Cluster as 'neptune-db://<host>' or Neptune Analytics Server as 'neptune-graph://<graphid>'",
|
||||
),
|
||||
)
|
||||
base_label: Optional[bool] = Field(None, description="Whether to use base node label __Entity__ for all entities")
|
||||
collection_name: Optional[str] = Field(None, description="vector_store collection name to store vectors when using Neptune-DB Clusters")
|
||||
|
||||
@model_validator(mode="before")
|
||||
def check_host_port_or_path(cls, values):
|
||||
endpoint = values.get("endpoint")
|
||||
if not endpoint:
|
||||
raise ValueError("Please provide 'endpoint' with the format as 'neptune-db://<endpoint>' or 'neptune-graph://<graphid>'.")
|
||||
if endpoint.startswith("neptune-db://"):
|
||||
# This is a Neptune DB Graph
|
||||
return values
|
||||
elif endpoint.startswith("neptune-graph://"):
|
||||
# This is a Neptune Analytics Graph
|
||||
graph_identifier = endpoint.replace("neptune-graph://", "")
|
||||
if not graph_identifier.startswith("g-"):
|
||||
raise ValueError("Provide a valid 'graph_identifier'.")
|
||||
values["graph_identifier"] = graph_identifier
|
||||
return values
|
||||
else:
|
||||
raise ValueError(
|
||||
"You must provide an endpoint to create a NeptuneServer as either neptune-db://<endpoint> or neptune-graph://<graphid>"
|
||||
)
|
||||
|
||||
|
||||
class KuzuConfig(BaseModel):
|
||||
db: Optional[str] = Field(":memory:", description="Path to a Kuzu database file")
|
||||
|
||||
|
||||
class ApacheAgeConfig(BaseModel):
|
||||
host: Optional[str] = Field("localhost", description="PostgreSQL server hostname")
|
||||
port: Optional[int] = Field(5432, description="PostgreSQL server port")
|
||||
database: Optional[str] = Field(None, description="PostgreSQL database name")
|
||||
username: Optional[str] = Field(None, description="PostgreSQL username")
|
||||
password: Optional[str] = Field(None, description="PostgreSQL password")
|
||||
graph_name: Optional[str] = Field("mem0_graph", description="Name of the Apache AGE graph")
|
||||
|
||||
@model_validator(mode="before")
|
||||
def check_required_fields(cls, values):
|
||||
database, username, password = (
|
||||
values.get("database"),
|
||||
values.get("username"),
|
||||
values.get("password"),
|
||||
)
|
||||
if not database or not username or not password:
|
||||
raise ValueError("Please provide 'database', 'username' and 'password'.")
|
||||
return values
|
||||
|
||||
|
||||
class GraphStoreConfig(BaseModel):
|
||||
provider: str = Field(
|
||||
description="Provider of the data store (e.g., 'neo4j', 'memgraph', 'neptune', 'kuzu', 'apache_age')",
|
||||
default="neo4j",
|
||||
)
|
||||
config: Union[Neo4jConfig, MemgraphConfig, NeptuneConfig, KuzuConfig, ApacheAgeConfig] = Field(
|
||||
description="Configuration for the specific data store", default=None
|
||||
)
|
||||
llm: Optional[LlmConfig] = Field(description="LLM configuration for querying the graph store", default=None)
|
||||
custom_prompt: Optional[str] = Field(
|
||||
description="Custom prompt to fetch entities from the given text", default=None
|
||||
)
|
||||
threshold: float = Field(
|
||||
description="Threshold for embedding similarity when matching nodes during graph ingestion. "
|
||||
"Range: 0.0 to 1.0. Higher values require closer matches. "
|
||||
"Use lower values (e.g., 0.5-0.7) for distinct entities with similar embeddings. "
|
||||
"Use higher values (e.g., 0.9+) when you want stricter matching.",
|
||||
default=0.7,
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
)
|
||||
|
||||
@field_validator("config")
|
||||
def validate_config(cls, v, values):
|
||||
provider = values.data.get("provider")
|
||||
if provider == "neo4j":
|
||||
return Neo4jConfig(**v.model_dump())
|
||||
elif provider == "memgraph":
|
||||
return MemgraphConfig(**v.model_dump())
|
||||
elif provider == "neptune" or provider == "neptunedb":
|
||||
return NeptuneConfig(**v.model_dump())
|
||||
elif provider == "kuzu":
|
||||
return KuzuConfig(**v.model_dump())
|
||||
elif provider == "apache_age":
|
||||
return ApacheAgeConfig(**v.model_dump())
|
||||
else:
|
||||
raise ValueError(f"Unsupported graph store provider: {provider}")
|
||||
@@ -1,515 +0,0 @@
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from mem0.memory.utils import format_entities, remove_spaces_from_entities
|
||||
|
||||
try:
|
||||
from rank_bm25 import BM25Okapi
|
||||
except ImportError:
|
||||
raise ImportError("rank_bm25 is not installed. Please install it using pip install rank-bm25")
|
||||
|
||||
from mem0.graphs.tools import (
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
DELETE_MEMORY_TOOL_GRAPH,
|
||||
EXTRACT_ENTITIES_STRUCT_TOOL,
|
||||
EXTRACT_ENTITIES_TOOL,
|
||||
RELATIONS_STRUCT_TOOL,
|
||||
RELATIONS_TOOL,
|
||||
)
|
||||
from mem0.graphs.utils import EXTRACT_RELATIONS_PROMPT, get_delete_messages
|
||||
from mem0.utils.factory import EmbedderFactory, LlmFactory, VectorStoreFactory
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class NeptuneBase(ABC):
|
||||
"""
|
||||
Abstract base class for neptune (neptune analytics and neptune db) calls using OpenCypher
|
||||
to store/retrieve data
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _create_embedding_model(config):
|
||||
"""
|
||||
:return: the Embedder model used for memory store
|
||||
"""
|
||||
return EmbedderFactory.create(
|
||||
config.embedder.provider,
|
||||
config.embedder.config,
|
||||
{"enable_embeddings": True},
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _create_llm(config, llm_provider):
|
||||
"""
|
||||
:return: the llm model used for memory store
|
||||
"""
|
||||
return LlmFactory.create(llm_provider, config.llm.config)
|
||||
|
||||
@staticmethod
|
||||
def _create_vector_store(vector_store_provider, config):
|
||||
"""
|
||||
:param vector_store_provider: name of vector store
|
||||
:param config: the vector_store configuration
|
||||
:return:
|
||||
"""
|
||||
return VectorStoreFactory.create(vector_store_provider, config.vector_store.config)
|
||||
|
||||
def add(self, data, filters):
|
||||
"""
|
||||
Adds data to the graph.
|
||||
|
||||
Args:
|
||||
data (str): The data to add to the graph.
|
||||
filters (dict): A dictionary containing filters to be applied during the addition.
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
to_be_added = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
to_be_deleted = self._get_delete_entities_from_search_output(search_output, data, filters)
|
||||
|
||||
deleted_entities = self._delete_entities(to_be_deleted, filters["user_id"])
|
||||
added_entities = self._add_entities(to_be_added, filters["user_id"], entity_type_map)
|
||||
|
||||
return {"deleted_entities": deleted_entities, "added_entities": added_entities}
|
||||
|
||||
def _retrieve_nodes_from_data(self, data, filters):
|
||||
"""
|
||||
Extract all entities mentioned in the query.
|
||||
"""
|
||||
_tools = [EXTRACT_ENTITIES_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [EXTRACT_ENTITIES_STRUCT_TOOL]
|
||||
search_results = self.llm.generate_response(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": f"You are a smart assistant who understands entities and their types in a given text. If user message contains self reference such as 'I', 'me', 'my' etc. then use {filters['user_id']} as the source entity. Extract all the entities from the text. ***DO NOT*** answer the question itself if the given text is a question.",
|
||||
},
|
||||
{"role": "user", "content": data},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entity_type_map = {}
|
||||
|
||||
try:
|
||||
for tool_call in search_results["tool_calls"]:
|
||||
if tool_call["name"] != "extract_entities":
|
||||
continue
|
||||
for item in tool_call.get("arguments", {}).get("entities", []):
|
||||
entity_type_map[item["entity"]] = item["entity_type"]
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Error in search tool: {e}, llm_provider={self.llm_provider}, search_results={search_results}"
|
||||
)
|
||||
|
||||
entity_type_map = {k.lower().replace(" ", "_"): v.lower().replace(" ", "_") for k, v in entity_type_map.items()}
|
||||
return entity_type_map
|
||||
|
||||
def _establish_nodes_relations_from_data(self, data, filters, entity_type_map):
|
||||
"""
|
||||
Establish relations among the extracted nodes.
|
||||
"""
|
||||
if self.config.graph_store.custom_prompt:
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": EXTRACT_RELATIONS_PROMPT.replace("USER_ID", filters["user_id"]).replace(
|
||||
"CUSTOM_PROMPT", f"4. {self.config.graph_store.custom_prompt}"
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": data},
|
||||
]
|
||||
else:
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": EXTRACT_RELATIONS_PROMPT.replace("USER_ID", filters["user_id"]),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"List of entities: {list(entity_type_map.keys())}. \n\nText: {data}",
|
||||
},
|
||||
]
|
||||
|
||||
_tools = [RELATIONS_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [RELATIONS_STRUCT_TOOL]
|
||||
|
||||
extracted_entities = self.llm.generate_response(
|
||||
messages=messages,
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entities = []
|
||||
if extracted_entities["tool_calls"]:
|
||||
entities = extracted_entities["tool_calls"][0].get("arguments", {}).get("entities", [])
|
||||
|
||||
entities = self._remove_spaces_from_entities(entities)
|
||||
logger.debug(f"Extracted entities: {entities}")
|
||||
return entities
|
||||
|
||||
def _remove_spaces_from_entities(self, entity_list):
|
||||
return remove_spaces_from_entities(entity_list, sanitize_relationship=False)
|
||||
|
||||
def _get_delete_entities_from_search_output(self, search_output, data, filters):
|
||||
"""
|
||||
Get the entities to be deleted from the search output.
|
||||
"""
|
||||
|
||||
search_output_string = format_entities(search_output)
|
||||
system_prompt, user_prompt = get_delete_messages(search_output_string, data, filters["user_id"])
|
||||
|
||||
_tools = [DELETE_MEMORY_TOOL_GRAPH]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
]
|
||||
|
||||
memory_updates = self.llm.generate_response(
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
to_be_deleted = []
|
||||
for item in memory_updates["tool_calls"]:
|
||||
if item["name"] == "delete_graph_memory":
|
||||
to_be_deleted.append(item["arguments"])
|
||||
# in case if it is not in the correct format
|
||||
to_be_deleted = self._remove_spaces_from_entities(to_be_deleted)
|
||||
logger.debug(f"Deleted relationships: {to_be_deleted}")
|
||||
return to_be_deleted
|
||||
|
||||
def _delete_entities(self, to_be_deleted, user_id):
|
||||
"""
|
||||
Delete the entities from the graph.
|
||||
"""
|
||||
|
||||
results = []
|
||||
for item in to_be_deleted:
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
# Delete the specific relationship between nodes
|
||||
cypher, params = self._delete_entities_cypher(source, destination, relationship, user_id)
|
||||
result = self.graph.query(cypher, params=params)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
@abstractmethod
|
||||
def _delete_entities_cypher(self, source, destination, relationship, user_id):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for deleting entities in the graph DB
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
def _add_entities(self, to_be_added, user_id, entity_type_map):
|
||||
"""
|
||||
Add the new entities to the graph. Merge the nodes if they already exist.
|
||||
"""
|
||||
|
||||
results = []
|
||||
for item in to_be_added:
|
||||
# entities
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
# types
|
||||
source_type = entity_type_map.get(source, "__User__")
|
||||
destination_type = entity_type_map.get(destination, "__User__")
|
||||
|
||||
# embeddings
|
||||
source_embedding = self.embedding_model.embed(source)
|
||||
dest_embedding = self.embedding_model.embed(destination)
|
||||
|
||||
# search for the nodes with the closest embeddings
|
||||
source_node_search_result = self._search_source_node(source_embedding, user_id, threshold=self.threshold)
|
||||
destination_node_search_result = self._search_destination_node(dest_embedding, user_id, threshold=self.threshold)
|
||||
|
||||
cypher, params = self._add_entities_cypher(
|
||||
source_node_search_result,
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination_node_search_result,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id,
|
||||
)
|
||||
result = self.graph.query(cypher, params=params)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
def _add_entities_cypher(
|
||||
self,
|
||||
source_node_list,
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination_node_list,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
"""
|
||||
if not destination_node_list and source_node_list:
|
||||
return self._add_entities_by_source_cypher(
|
||||
source_node_list,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id)
|
||||
elif destination_node_list and not source_node_list:
|
||||
return self._add_entities_by_destination_cypher(
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination_node_list,
|
||||
relationship,
|
||||
user_id)
|
||||
elif source_node_list and destination_node_list:
|
||||
return self._add_relationship_entities_cypher(
|
||||
source_node_list,
|
||||
destination_node_list,
|
||||
relationship,
|
||||
user_id)
|
||||
# else source_node_list and destination_node_list are empty
|
||||
return self._add_new_entities_cypher(
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id)
|
||||
|
||||
@abstractmethod
|
||||
def _add_entities_by_source_cypher(
|
||||
self,
|
||||
source_node_list,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def _add_entities_by_destination_cypher(
|
||||
self,
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination_node_list,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def _add_relationship_entities_cypher(
|
||||
self,
|
||||
source_node_list,
|
||||
destination_node_list,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def _add_new_entities_cypher(
|
||||
self,
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
pass
|
||||
|
||||
def search(self, query, filters, top_k=100):
|
||||
"""
|
||||
Search for memories and related graph data.
|
||||
|
||||
Args:
|
||||
query (str): Query to search for.
|
||||
filters (dict): A dictionary containing filters to be applied during the search.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing:
|
||||
- "contexts": List of search results from the base data store.
|
||||
- "entities": List of related graph data based on the query.
|
||||
"""
|
||||
|
||||
entity_type_map = self._retrieve_nodes_from_data(query, filters)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
|
||||
if not search_output:
|
||||
return []
|
||||
|
||||
search_outputs_sequence = [
|
||||
[item["source"], item["relationship"], item["destination"]] for item in search_output
|
||||
]
|
||||
bm25 = BM25Okapi(search_outputs_sequence)
|
||||
|
||||
tokenized_query = query.split(" ")
|
||||
reranked_results = bm25.get_top_n(tokenized_query, search_outputs_sequence, n=5)
|
||||
|
||||
search_results = []
|
||||
for item in reranked_results:
|
||||
search_results.append({"source": item[0], "relationship": item[1], "destination": item[2]})
|
||||
|
||||
return search_results
|
||||
|
||||
def _search_source_node(self, source_embedding, user_id, threshold=0.9):
|
||||
cypher, params = self._search_source_node_cypher(source_embedding, user_id, threshold)
|
||||
result = self.graph.query(cypher, params=params)
|
||||
return result
|
||||
|
||||
@abstractmethod
|
||||
def _search_source_node_cypher(self, source_embedding, user_id, threshold):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for source nodes
|
||||
"""
|
||||
pass
|
||||
|
||||
def _search_destination_node(self, destination_embedding, user_id, threshold=0.9):
|
||||
cypher, params = self._search_destination_node_cypher(destination_embedding, user_id, threshold)
|
||||
result = self.graph.query(cypher, params=params)
|
||||
return result
|
||||
|
||||
@abstractmethod
|
||||
def _search_destination_node_cypher(self, destination_embedding, user_id, threshold):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for destination nodes
|
||||
"""
|
||||
pass
|
||||
|
||||
def delete(self, data, filters):
|
||||
"""
|
||||
Delete graph entities associated with the given memory text.
|
||||
|
||||
Extracts entities and relationships from the memory text using the same
|
||||
pipeline as add(), then deletes the matching relationships in the graph.
|
||||
|
||||
Args:
|
||||
data (str): The memory text whose graph entities should be removed.
|
||||
filters (dict): Scope filters (user_id, agent_id, run_id).
|
||||
"""
|
||||
try:
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
if not entity_type_map:
|
||||
logger.debug("No entities found in memory text, skipping graph cleanup")
|
||||
return
|
||||
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
if to_be_deleted:
|
||||
self._delete_entities(to_be_deleted, filters["user_id"])
|
||||
except Exception as e:
|
||||
logger.error(f"Error during graph cleanup for memory delete: {e}")
|
||||
|
||||
def delete_all(self, filters):
|
||||
cypher, params = self._delete_all_cypher(filters)
|
||||
self.graph.query(cypher, params=params)
|
||||
|
||||
@abstractmethod
|
||||
def _delete_all_cypher(self, filters):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to delete all edges/nodes in the memory store
|
||||
"""
|
||||
pass
|
||||
|
||||
def get_all(self, filters, top_k=100):
|
||||
"""
|
||||
Retrieves all nodes and relationships from the graph database based on filtering criteria.
|
||||
|
||||
Args:
|
||||
filters (dict): A dictionary containing filters to be applied during the retrieval.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
Returns:
|
||||
list: A list of dictionaries, each containing:
|
||||
- 'contexts': The base data store response for each memory.
|
||||
- 'entities': A list of strings representing the nodes and relationships
|
||||
"""
|
||||
|
||||
# return all nodes and relationships
|
||||
query, params = self._get_all_cypher(filters, top_k)
|
||||
results = self.graph.query(query, params=params)
|
||||
|
||||
final_results = []
|
||||
for result in results:
|
||||
final_results.append(
|
||||
{
|
||||
"source": result["source"],
|
||||
"relationship": result["relationship"],
|
||||
"target": result["target"],
|
||||
}
|
||||
)
|
||||
|
||||
logger.debug(f"Retrieved {len(final_results)} relationships")
|
||||
|
||||
return final_results
|
||||
|
||||
@abstractmethod
|
||||
def _get_all_cypher(self, filters, top_k):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to get all edges/nodes in the memory store
|
||||
"""
|
||||
pass
|
||||
|
||||
def _search_graph_db(self, node_list, filters, top_k=100):
|
||||
"""
|
||||
Search similar nodes among and their respective incoming and outgoing relations.
|
||||
"""
|
||||
result_relations = []
|
||||
|
||||
for node in node_list:
|
||||
n_embedding = self.embedding_model.embed(node)
|
||||
cypher_query, params = self._search_graph_db_cypher(n_embedding, filters, top_k)
|
||||
ans = self.graph.query(cypher_query, params=params)
|
||||
result_relations.extend(ans)
|
||||
|
||||
return result_relations
|
||||
|
||||
@abstractmethod
|
||||
def _search_graph_db_cypher(self, n_embedding, filters, top_k):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for similar nodes in the memory store
|
||||
"""
|
||||
pass
|
||||
|
||||
# Reset is not defined in base.py
|
||||
def reset(self):
|
||||
"""
|
||||
Reset the graph by clearing all nodes and relationships.
|
||||
|
||||
link: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/neptune-graph/client/reset_graph.html
|
||||
"""
|
||||
|
||||
logger.warning("Clearing graph...")
|
||||
graph_id = self.graph.graph_identifier
|
||||
self.graph.client.reset_graph(
|
||||
graphIdentifier=graph_id,
|
||||
skipSnapshot=True,
|
||||
)
|
||||
waiter = self.graph.client.get_waiter("graph_available")
|
||||
waiter.wait(graphIdentifier=graph_id, WaiterConfig={"Delay": 10, "MaxAttempts": 60})
|
||||
@@ -1,511 +0,0 @@
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from .base import NeptuneBase
|
||||
|
||||
try:
|
||||
from langchain_aws import NeptuneGraph
|
||||
except ImportError:
|
||||
raise ImportError("langchain_aws is not installed. Please install it using 'make install_all'.")
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class MemoryGraph(NeptuneBase):
|
||||
def __init__(self, config):
|
||||
"""
|
||||
Initialize the Neptune DB memory store.
|
||||
"""
|
||||
|
||||
self.config = config
|
||||
|
||||
self.graph = None
|
||||
endpoint = self.config.graph_store.config.endpoint
|
||||
if endpoint and endpoint.startswith("neptune-db://"):
|
||||
host = endpoint.replace("neptune-db://", "")
|
||||
port = 8182
|
||||
self.graph = NeptuneGraph(host, port)
|
||||
|
||||
if not self.graph:
|
||||
raise ValueError("Unable to create a Neptune-DB client: missing 'endpoint' in config")
|
||||
|
||||
self.node_label = ":`__Entity__`" if self.config.graph_store.config.base_label else ""
|
||||
|
||||
self.embedding_model = NeptuneBase._create_embedding_model(self.config)
|
||||
|
||||
# Default to openai if no specific provider is configured
|
||||
self.llm_provider = "openai"
|
||||
if self.config.graph_store.llm:
|
||||
self.llm_provider = self.config.graph_store.llm.provider
|
||||
elif self.config.llm.provider:
|
||||
self.llm_provider = self.config.llm.provider
|
||||
|
||||
# fetch the vector store as a provider
|
||||
self.vector_store_provider = self.config.vector_store.provider
|
||||
if self.config.graph_store.config.collection_name:
|
||||
vector_store_collection_name = self.config.graph_store.config.collection_name
|
||||
else:
|
||||
vector_store_config = self.config.vector_store.config
|
||||
if vector_store_config.collection_name:
|
||||
vector_store_collection_name = vector_store_config.collection_name + "_neptune_vector_store"
|
||||
else:
|
||||
vector_store_collection_name = "mem0_neptune_vector_store"
|
||||
self.config.vector_store.config.collection_name = vector_store_collection_name
|
||||
self.vector_store = NeptuneBase._create_vector_store(self.vector_store_provider, self.config)
|
||||
|
||||
self.llm = NeptuneBase._create_llm(self.config, self.llm_provider)
|
||||
self.user_id = None
|
||||
# Use threshold from graph_store config, default to 0.7 for backward compatibility
|
||||
self.threshold = self.config.graph_store.threshold if hasattr(self.config.graph_store, 'threshold') else 0.7
|
||||
self.vector_store_limit=5
|
||||
|
||||
def _delete_entities_cypher(self, source, destination, relationship, user_id):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for deleting entities in the graph DB
|
||||
|
||||
:param source: source node
|
||||
:param destination: destination node
|
||||
:param relationship: relationship label
|
||||
:param user_id: user_id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{name: $source_name, user_id: $user_id}})
|
||||
-[r:{relationship}]->
|
||||
(m {self.node_label} {{name: $dest_name, user_id: $user_id}})
|
||||
DELETE r
|
||||
RETURN
|
||||
n.name AS source,
|
||||
m.name AS target,
|
||||
type(r) AS relationship
|
||||
"""
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(f"_delete_entities\n query={cypher}")
|
||||
return cypher, params
|
||||
|
||||
def _add_entities_by_source_cypher(
|
||||
self,
|
||||
source_node_list,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
|
||||
:param source_node_list: list of source nodes
|
||||
:param destination: destination name
|
||||
:param dest_embedding: destination embedding
|
||||
:param destination_type: destination node label
|
||||
:param relationship: relationship label
|
||||
:param user_id: user id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
destination_id = str(uuid.uuid4())
|
||||
destination_payload = {
|
||||
"name": destination,
|
||||
"type": destination_type,
|
||||
"user_id": user_id,
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
self.vector_store.insert(
|
||||
vectors=[dest_embedding],
|
||||
payloads=[destination_payload],
|
||||
ids=[destination_id],
|
||||
)
|
||||
|
||||
destination_label = self.node_label if self.node_label else f":`{destination_type}`"
|
||||
destination_extra_set = f", destination:`{destination_type}`" if self.node_label else ""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source {{user_id: $user_id}})
|
||||
WHERE id(source) = $source_id
|
||||
SET source.mentions = coalesce(source.mentions, 0) + 1
|
||||
WITH source
|
||||
MERGE (destination {destination_label} {{`~id`: $destination_id, name: $destination_name, user_id: $user_id}})
|
||||
ON CREATE SET
|
||||
destination.created = timestamp(),
|
||||
destination.updated = timestamp(),
|
||||
destination.mentions = 1
|
||||
{destination_extra_set}
|
||||
ON MATCH SET
|
||||
destination.mentions = coalesce(destination.mentions, 0) + 1,
|
||||
destination.updated = timestamp()
|
||||
WITH source, destination
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = timestamp(),
|
||||
r.updated = timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1,
|
||||
r.updated = timestamp()
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target, id(destination) AS destination_id
|
||||
"""
|
||||
|
||||
params = {
|
||||
"source_id": source_node_list[0]["id(source_candidate)"],
|
||||
"destination_id": destination_id,
|
||||
"destination_name": destination,
|
||||
"dest_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
|
||||
logger.debug(
|
||||
f"_add_entities:\n source_node_search_result={source_node_list[0]}\n query={cypher}"
|
||||
)
|
||||
return cypher, params
|
||||
|
||||
def _add_entities_by_destination_cypher(
|
||||
self,
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination_node_list,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
|
||||
:param source: source node name
|
||||
:param source_embedding: source node embedding
|
||||
:param source_type: source node label
|
||||
:param destination_node_list: list of dest nodes
|
||||
:param relationship: relationship label
|
||||
:param user_id: user id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
source_id = str(uuid.uuid4())
|
||||
source_payload = {
|
||||
"name": source,
|
||||
"type": source_type,
|
||||
"user_id": user_id,
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
self.vector_store.insert(
|
||||
vectors=[source_embedding],
|
||||
payloads=[source_payload],
|
||||
ids=[source_id],
|
||||
)
|
||||
|
||||
source_label = self.node_label if self.node_label else f":`{source_type}`"
|
||||
source_extra_set = f", source:`{source_type}`" if self.node_label else ""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (destination {{user_id: $user_id}})
|
||||
WHERE id(destination) = $destination_id
|
||||
SET
|
||||
destination.mentions = coalesce(destination.mentions, 0) + 1,
|
||||
destination.updated = timestamp()
|
||||
WITH destination
|
||||
MERGE (source {source_label} {{`~id`: $source_id, name: $source_name, user_id: $user_id}})
|
||||
ON CREATE SET
|
||||
source.created = timestamp(),
|
||||
source.updated = timestamp(),
|
||||
source.mentions = 1
|
||||
{source_extra_set}
|
||||
ON MATCH SET
|
||||
source.mentions = coalesce(source.mentions, 0) + 1,
|
||||
source.updated = timestamp()
|
||||
WITH source, destination
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = timestamp(),
|
||||
r.updated = timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1,
|
||||
r.updated = timestamp()
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"destination_id": destination_node_list[0]["id(destination_candidate)"],
|
||||
"source_id": source_id,
|
||||
"source_name": source,
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(
|
||||
f"_add_entities:\n destination_node_search_result={destination_node_list[0]}\n query={cypher}"
|
||||
)
|
||||
return cypher, params
|
||||
|
||||
def _add_relationship_entities_cypher(
|
||||
self,
|
||||
source_node_list,
|
||||
destination_node_list,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
|
||||
:param source_node_list: list of source node ids
|
||||
:param destination_node_list: list of dest node ids
|
||||
:param relationship: relationship label
|
||||
:param user_id: user id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source {{user_id: $user_id}})
|
||||
WHERE id(source) = $source_id
|
||||
SET
|
||||
source.mentions = coalesce(source.mentions, 0) + 1,
|
||||
source.updated = timestamp()
|
||||
WITH source
|
||||
MATCH (destination {{user_id: $user_id}})
|
||||
WHERE id(destination) = $destination_id
|
||||
SET
|
||||
destination.mentions = coalesce(destination.mentions) + 1,
|
||||
destination.updated = timestamp()
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created_at = timestamp(),
|
||||
r.updated_at = timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET r.mentions = coalesce(r.mentions, 0) + 1
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
params = {
|
||||
"source_id": source_node_list[0]["id(source_candidate)"],
|
||||
"destination_id": destination_node_list[0]["id(destination_candidate)"],
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(
|
||||
f"_add_entities:\n destination_node_search_result={destination_node_list[0]}\n source_node_search_result={source_node_list[0]}\n query={cypher}"
|
||||
)
|
||||
return cypher, params
|
||||
|
||||
def _add_new_entities_cypher(
|
||||
self,
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
|
||||
:param source: source node name
|
||||
:param source_embedding: source node embedding
|
||||
:param source_type: source node label
|
||||
:param destination: destination name
|
||||
:param dest_embedding: destination embedding
|
||||
:param destination_type: destination node label
|
||||
:param relationship: relationship label
|
||||
:param user_id: user id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
source_id = str(uuid.uuid4())
|
||||
source_payload = {
|
||||
"name": source,
|
||||
"type": source_type,
|
||||
"user_id": user_id,
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
destination_id = str(uuid.uuid4())
|
||||
destination_payload = {
|
||||
"name": destination,
|
||||
"type": destination_type,
|
||||
"user_id": user_id,
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
self.vector_store.insert(
|
||||
vectors=[source_embedding, dest_embedding],
|
||||
payloads=[source_payload, destination_payload],
|
||||
ids=[source_id, destination_id],
|
||||
)
|
||||
|
||||
source_label = self.node_label if self.node_label else f":`{source_type}`"
|
||||
source_extra_set = f", source:`{source_type}`" if self.node_label else ""
|
||||
destination_label = self.node_label if self.node_label else f":`{destination_type}`"
|
||||
destination_extra_set = f", destination:`{destination_type}`" if self.node_label else ""
|
||||
|
||||
cypher = f"""
|
||||
MERGE (n {source_label} {{name: $source_name, user_id: $user_id, `~id`: $source_id}})
|
||||
ON CREATE SET n.created = timestamp(),
|
||||
n.mentions = 1
|
||||
{source_extra_set}
|
||||
ON MATCH SET n.mentions = coalesce(n.mentions, 0) + 1
|
||||
WITH n
|
||||
MERGE (m {destination_label} {{name: $dest_name, user_id: $user_id, `~id`: $dest_id}})
|
||||
ON CREATE SET m.created = timestamp(),
|
||||
m.mentions = 1
|
||||
{destination_extra_set}
|
||||
ON MATCH SET m.mentions = coalesce(m.mentions, 0) + 1
|
||||
WITH n, m
|
||||
MERGE (n)-[rel:{relationship}]->(m)
|
||||
ON CREATE SET rel.created = timestamp(), rel.mentions = 1
|
||||
ON MATCH SET rel.mentions = coalesce(rel.mentions, 0) + 1
|
||||
RETURN n.name AS source, type(rel) AS relationship, m.name AS target
|
||||
"""
|
||||
params = {
|
||||
"source_id": source_id,
|
||||
"dest_id": destination_id,
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"source_embedding": source_embedding,
|
||||
"dest_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(
|
||||
f"_add_new_entities_cypher:\n query={cypher}"
|
||||
)
|
||||
return cypher, params
|
||||
|
||||
def _search_source_node_cypher(self, source_embedding, user_id, threshold):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for source nodes
|
||||
|
||||
:param source_embedding: source vector
|
||||
:param user_id: user_id to use
|
||||
:param threshold: the threshold for similarity
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
source_nodes = self.vector_store.search(
|
||||
query="",
|
||||
vectors=source_embedding,
|
||||
top_k=self.vector_store_limit,
|
||||
filters={"user_id": user_id},
|
||||
)
|
||||
|
||||
ids = [n.id for n in filter(lambda s: s.score > threshold, source_nodes)]
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source_candidate {self.node_label})
|
||||
WHERE source_candidate.user_id = $user_id AND id(source_candidate) IN $ids
|
||||
RETURN id(source_candidate)
|
||||
"""
|
||||
|
||||
params = {
|
||||
"ids": ids,
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
"threshold": threshold,
|
||||
}
|
||||
logger.debug(f"_search_source_node\n query={cypher}")
|
||||
return cypher, params
|
||||
|
||||
def _search_destination_node_cypher(self, destination_embedding, user_id, threshold):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for destination nodes
|
||||
|
||||
:param source_embedding: source vector
|
||||
:param user_id: user_id to use
|
||||
:param threshold: the threshold for similarity
|
||||
:return: str, dict
|
||||
"""
|
||||
destination_nodes = self.vector_store.search(
|
||||
query="",
|
||||
vectors=destination_embedding,
|
||||
top_k=self.vector_store_limit,
|
||||
filters={"user_id": user_id},
|
||||
)
|
||||
|
||||
ids = [n.id for n in filter(lambda d: d.score > threshold, destination_nodes)]
|
||||
|
||||
cypher = f"""
|
||||
MATCH (destination_candidate {self.node_label})
|
||||
WHERE destination_candidate.user_id = $user_id AND id(destination_candidate) IN $ids
|
||||
RETURN id(destination_candidate)
|
||||
"""
|
||||
|
||||
params = {
|
||||
"ids": ids,
|
||||
"destination_embedding": destination_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
|
||||
logger.debug(f"_search_destination_node\n query={cypher}")
|
||||
return cypher, params
|
||||
|
||||
def _delete_all_cypher(self, filters):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to delete all edges/nodes in the memory store
|
||||
|
||||
:param filters: search filters
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
# remove the vector store index
|
||||
self.vector_store.reset()
|
||||
|
||||
# create a query that: deletes the nodes of the graph_store
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{user_id: $user_id}})
|
||||
DETACH DELETE n
|
||||
"""
|
||||
params = {"user_id": filters["user_id"]}
|
||||
|
||||
logger.debug(f"delete_all query={cypher}")
|
||||
return cypher, params
|
||||
|
||||
def _get_all_cypher(self, filters, top_k):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to get all edges/nodes in the memory store
|
||||
|
||||
:param filters: search filters
|
||||
:param top_k: return limit
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{user_id: $user_id}})-[r]->(m {self.node_label} {{user_id: $user_id}})
|
||||
RETURN n.name AS source, type(r) AS relationship, m.name AS target
|
||||
LIMIT $limit
|
||||
"""
|
||||
params = {"user_id": filters["user_id"], "limit": top_k}
|
||||
return cypher, params
|
||||
|
||||
def _search_graph_db_cypher(self, n_embedding, filters, top_k):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for similar nodes in the memory store
|
||||
|
||||
:param n_embedding: node vector
|
||||
:param filters: search filters
|
||||
:param top_k: return limit
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
# search vector store for applicable nodes using cosine similarity
|
||||
search_nodes = self.vector_store.search(
|
||||
query="",
|
||||
vectors=n_embedding,
|
||||
top_k=self.vector_store_limit,
|
||||
filters=filters,
|
||||
)
|
||||
|
||||
ids = [n.id for n in search_nodes]
|
||||
|
||||
cypher_query = f"""
|
||||
MATCH (n {self.node_label})-[r]->(m)
|
||||
WHERE n.user_id = $user_id AND id(n) IN $n_ids
|
||||
RETURN n.name AS source, id(n) AS source_id, type(r) AS relationship, id(r) AS relation_id, m.name AS destination, id(m) AS destination_id
|
||||
UNION
|
||||
MATCH (m)-[r]->(n {self.node_label})
|
||||
RETURN m.name AS source, id(m) AS source_id, type(r) AS relationship, id(r) AS relation_id, n.name AS destination, id(n) AS destination_id
|
||||
LIMIT $limit
|
||||
"""
|
||||
params = {
|
||||
"n_ids": ids,
|
||||
"user_id": filters["user_id"],
|
||||
"limit": top_k,
|
||||
}
|
||||
logger.debug(f"_search_graph_db\n query={cypher_query}")
|
||||
|
||||
return cypher_query, params
|
||||
@@ -1,475 +0,0 @@
|
||||
import logging
|
||||
|
||||
from .base import NeptuneBase
|
||||
|
||||
try:
|
||||
from botocore.config import Config
|
||||
from langchain_aws import NeptuneAnalyticsGraph
|
||||
except ImportError:
|
||||
raise ImportError("langchain_aws is not installed. Please install it using 'make install_all'.")
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MemoryGraph(NeptuneBase):
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
|
||||
self.graph = None
|
||||
endpoint = self.config.graph_store.config.endpoint
|
||||
app_id = self.config.graph_store.config.app_id
|
||||
if endpoint and endpoint.startswith("neptune-graph://"):
|
||||
graph_identifier = endpoint.replace("neptune-graph://", "")
|
||||
self.graph = NeptuneAnalyticsGraph(graph_identifier = graph_identifier,
|
||||
config = Config(user_agent_appid=app_id))
|
||||
|
||||
if not self.graph:
|
||||
raise ValueError("Unable to create a Neptune client: missing 'endpoint' in config")
|
||||
|
||||
self.node_label = ":`__Entity__`" if self.config.graph_store.config.base_label else ""
|
||||
|
||||
self.embedding_model = NeptuneBase._create_embedding_model(self.config)
|
||||
|
||||
# Default to openai if no specific provider is configured
|
||||
self.llm_provider = "openai"
|
||||
if self.config.llm.provider:
|
||||
self.llm_provider = self.config.llm.provider
|
||||
if self.config.graph_store.llm:
|
||||
self.llm_provider = self.config.graph_store.llm.provider
|
||||
|
||||
self.llm = NeptuneBase._create_llm(self.config, self.llm_provider)
|
||||
self.user_id = None
|
||||
# Use threshold from graph_store config, default to 0.7 for backward compatibility
|
||||
self.threshold = self.config.graph_store.threshold if hasattr(self.config.graph_store, 'threshold') else 0.7
|
||||
|
||||
def _delete_entities_cypher(self, source, destination, relationship, user_id):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for deleting entities in the graph DB
|
||||
|
||||
:param source: source node
|
||||
:param destination: destination node
|
||||
:param relationship: relationship label
|
||||
:param user_id: user_id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{name: $source_name, user_id: $user_id}})
|
||||
-[r:{relationship}]->
|
||||
(m {self.node_label} {{name: $dest_name, user_id: $user_id}})
|
||||
DELETE r
|
||||
RETURN
|
||||
n.name AS source,
|
||||
m.name AS target,
|
||||
type(r) AS relationship
|
||||
"""
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(f"_delete_entities\n query={cypher}")
|
||||
return cypher, params
|
||||
|
||||
def _add_entities_by_source_cypher(
|
||||
self,
|
||||
source_node_list,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
|
||||
:param source_node_list: list of source nodes
|
||||
:param destination: destination name
|
||||
:param dest_embedding: destination embedding
|
||||
:param destination_type: destination node label
|
||||
:param relationship: relationship label
|
||||
:param user_id: user id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
destination_label = self.node_label if self.node_label else f":`{destination_type}`"
|
||||
destination_extra_set = f", destination:`{destination_type}`" if self.node_label else ""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source {{user_id: $user_id}})
|
||||
WHERE id(source) = $source_id
|
||||
SET source.mentions = coalesce(source.mentions, 0) + 1
|
||||
WITH source
|
||||
MERGE (destination {destination_label} {{name: $destination_name, user_id: $user_id}})
|
||||
ON CREATE SET
|
||||
destination.created = timestamp(),
|
||||
destination.updated = timestamp(),
|
||||
destination.mentions = 1
|
||||
{destination_extra_set}
|
||||
ON MATCH SET
|
||||
destination.mentions = coalesce(destination.mentions, 0) + 1,
|
||||
destination.updated = timestamp()
|
||||
WITH source, destination, $dest_embedding as dest_embedding
|
||||
CALL neptune.algo.vectors.upsert(destination, dest_embedding)
|
||||
WITH source, destination
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = timestamp(),
|
||||
r.updated = timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1,
|
||||
r.updated = timestamp()
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"source_id": source_node_list[0]["id(source_candidate)"],
|
||||
"destination_name": destination,
|
||||
"dest_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(
|
||||
f"_add_entities:\n source_node_search_result={source_node_list[0]}\n query={cypher}"
|
||||
)
|
||||
return cypher, params
|
||||
|
||||
def _add_entities_by_destination_cypher(
|
||||
self,
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination_node_list,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
|
||||
:param source: source node name
|
||||
:param source_embedding: source node embedding
|
||||
:param source_type: source node label
|
||||
:param destination_node_list: list of dest nodes
|
||||
:param relationship: relationship label
|
||||
:param user_id: user id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
source_label = self.node_label if self.node_label else f":`{source_type}`"
|
||||
source_extra_set = f", source:`{source_type}`" if self.node_label else ""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (destination {{user_id: $user_id}})
|
||||
WHERE id(destination) = $destination_id
|
||||
SET
|
||||
destination.mentions = coalesce(destination.mentions, 0) + 1,
|
||||
destination.updated = timestamp()
|
||||
WITH destination
|
||||
MERGE (source {source_label} {{name: $source_name, user_id: $user_id}})
|
||||
ON CREATE SET
|
||||
source.created = timestamp(),
|
||||
source.updated = timestamp(),
|
||||
source.mentions = 1
|
||||
{source_extra_set}
|
||||
ON MATCH SET
|
||||
source.mentions = coalesce(source.mentions, 0) + 1,
|
||||
source.updated = timestamp()
|
||||
WITH source, destination, $source_embedding as source_embedding
|
||||
CALL neptune.algo.vectors.upsert(source, source_embedding)
|
||||
WITH source, destination
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = timestamp(),
|
||||
r.updated = timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1,
|
||||
r.updated = timestamp()
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"destination_id": destination_node_list[0]["id(destination_candidate)"],
|
||||
"source_name": source,
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(
|
||||
f"_add_entities:\n destination_node_search_result={destination_node_list[0]}\n query={cypher}"
|
||||
)
|
||||
return cypher, params
|
||||
|
||||
def _add_relationship_entities_cypher(
|
||||
self,
|
||||
source_node_list,
|
||||
destination_node_list,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
|
||||
:param source_node_list: list of source node ids
|
||||
:param destination_node_list: list of dest node ids
|
||||
:param relationship: relationship label
|
||||
:param user_id: user id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source {{user_id: $user_id}})
|
||||
WHERE id(source) = $source_id
|
||||
SET
|
||||
source.mentions = coalesce(source.mentions, 0) + 1,
|
||||
source.updated = timestamp()
|
||||
WITH source
|
||||
MATCH (destination {{user_id: $user_id}})
|
||||
WHERE id(destination) = $destination_id
|
||||
SET
|
||||
destination.mentions = coalesce(destination.mentions) + 1,
|
||||
destination.updated = timestamp()
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created_at = timestamp(),
|
||||
r.updated_at = timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET r.mentions = coalesce(r.mentions, 0) + 1
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
params = {
|
||||
"source_id": source_node_list[0]["id(source_candidate)"],
|
||||
"destination_id": destination_node_list[0]["id(destination_candidate)"],
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(
|
||||
f"_add_entities:\n destination_node_search_result={destination_node_list[0]}\n source_node_search_result={source_node_list[0]}\n query={cypher}"
|
||||
)
|
||||
return cypher, params
|
||||
|
||||
def _add_new_entities_cypher(
|
||||
self,
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
|
||||
:param source: source node name
|
||||
:param source_embedding: source node embedding
|
||||
:param source_type: source node label
|
||||
:param destination: destination name
|
||||
:param dest_embedding: destination embedding
|
||||
:param destination_type: destination node label
|
||||
:param relationship: relationship label
|
||||
:param user_id: user id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
source_label = self.node_label if self.node_label else f":`{source_type}`"
|
||||
source_extra_set = f", source:`{source_type}`" if self.node_label else ""
|
||||
destination_label = self.node_label if self.node_label else f":`{destination_type}`"
|
||||
destination_extra_set = f", destination:`{destination_type}`" if self.node_label else ""
|
||||
|
||||
cypher = f"""
|
||||
MERGE (n {source_label} {{name: $source_name, user_id: $user_id}})
|
||||
ON CREATE SET n.created = timestamp(),
|
||||
n.updated = timestamp(),
|
||||
n.mentions = 1
|
||||
{source_extra_set}
|
||||
ON MATCH SET
|
||||
n.mentions = coalesce(n.mentions, 0) + 1,
|
||||
n.updated = timestamp()
|
||||
WITH n, $source_embedding as source_embedding
|
||||
CALL neptune.algo.vectors.upsert(n, source_embedding)
|
||||
WITH n
|
||||
MERGE (m {destination_label} {{name: $dest_name, user_id: $user_id}})
|
||||
ON CREATE SET
|
||||
m.created = timestamp(),
|
||||
m.updated = timestamp(),
|
||||
m.mentions = 1
|
||||
{destination_extra_set}
|
||||
ON MATCH SET
|
||||
m.updated = timestamp(),
|
||||
m.mentions = coalesce(m.mentions, 0) + 1
|
||||
WITH n, m, $dest_embedding as dest_embedding
|
||||
CALL neptune.algo.vectors.upsert(m, dest_embedding)
|
||||
WITH n, m
|
||||
MERGE (n)-[rel:{relationship}]->(m)
|
||||
ON CREATE SET
|
||||
rel.created = timestamp(),
|
||||
rel.updated = timestamp(),
|
||||
rel.mentions = 1
|
||||
ON MATCH SET
|
||||
rel.updated = timestamp(),
|
||||
rel.mentions = coalesce(rel.mentions, 0) + 1
|
||||
RETURN n.name AS source, type(rel) AS relationship, m.name AS target
|
||||
"""
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"source_embedding": source_embedding,
|
||||
"dest_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(
|
||||
f"_add_new_entities_cypher:\n query={cypher}"
|
||||
)
|
||||
return cypher, params
|
||||
|
||||
def _search_source_node_cypher(self, source_embedding, user_id, threshold):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for source nodes
|
||||
|
||||
:param source_embedding: source vector
|
||||
:param user_id: user_id to use
|
||||
:param threshold: the threshold for similarity
|
||||
:return: str, dict
|
||||
"""
|
||||
cypher = f"""
|
||||
MATCH (source_candidate {self.node_label})
|
||||
WHERE source_candidate.user_id = $user_id
|
||||
|
||||
WITH source_candidate, $source_embedding as v_embedding
|
||||
CALL neptune.algo.vectors.distanceByEmbedding(
|
||||
v_embedding,
|
||||
source_candidate,
|
||||
{{metric:"CosineSimilarity"}}
|
||||
) YIELD distance
|
||||
WITH source_candidate, distance AS cosine_similarity
|
||||
WHERE cosine_similarity >= $threshold
|
||||
|
||||
WITH source_candidate, cosine_similarity
|
||||
ORDER BY cosine_similarity DESC
|
||||
LIMIT 1
|
||||
|
||||
RETURN id(source_candidate), cosine_similarity
|
||||
"""
|
||||
|
||||
params = {
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
"threshold": threshold,
|
||||
}
|
||||
logger.debug(f"_search_source_node\n query={cypher}")
|
||||
return cypher, params
|
||||
|
||||
def _search_destination_node_cypher(self, destination_embedding, user_id, threshold):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for destination nodes
|
||||
|
||||
:param source_embedding: source vector
|
||||
:param user_id: user_id to use
|
||||
:param threshold: the threshold for similarity
|
||||
:return: str, dict
|
||||
"""
|
||||
cypher = f"""
|
||||
MATCH (destination_candidate {self.node_label})
|
||||
WHERE destination_candidate.user_id = $user_id
|
||||
|
||||
WITH destination_candidate, $destination_embedding as v_embedding
|
||||
CALL neptune.algo.vectors.distanceByEmbedding(
|
||||
v_embedding,
|
||||
destination_candidate,
|
||||
{{metric:"CosineSimilarity"}}
|
||||
) YIELD distance
|
||||
WITH destination_candidate, distance AS cosine_similarity
|
||||
WHERE cosine_similarity >= $threshold
|
||||
|
||||
WITH destination_candidate, cosine_similarity
|
||||
ORDER BY cosine_similarity DESC
|
||||
LIMIT 1
|
||||
|
||||
RETURN id(destination_candidate), cosine_similarity
|
||||
"""
|
||||
params = {
|
||||
"destination_embedding": destination_embedding,
|
||||
"user_id": user_id,
|
||||
"threshold": threshold,
|
||||
}
|
||||
|
||||
logger.debug(f"_search_destination_node\n query={cypher}")
|
||||
return cypher, params
|
||||
|
||||
def _delete_all_cypher(self, filters):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to delete all edges/nodes in the memory store
|
||||
|
||||
:param filters: search filters
|
||||
:return: str, dict
|
||||
"""
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{user_id: $user_id}})
|
||||
DETACH DELETE n
|
||||
"""
|
||||
params = {"user_id": filters["user_id"]}
|
||||
|
||||
logger.debug(f"delete_all query={cypher}")
|
||||
return cypher, params
|
||||
|
||||
def _get_all_cypher(self, filters, top_k):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to get all edges/nodes in the memory store
|
||||
|
||||
:param filters: search filters
|
||||
:param top_k: return limit
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{user_id: $user_id}})-[r]->(m {self.node_label} {{user_id: $user_id}})
|
||||
RETURN n.name AS source, type(r) AS relationship, m.name AS target
|
||||
LIMIT $limit
|
||||
"""
|
||||
params = {"user_id": filters["user_id"], "limit": top_k}
|
||||
return cypher, params
|
||||
|
||||
def _search_graph_db_cypher(self, n_embedding, filters, top_k):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for similar nodes in the memory store
|
||||
|
||||
:param n_embedding: node vector
|
||||
:param filters: search filters
|
||||
:param top_k: return limit
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
cypher_query = f"""
|
||||
MATCH (n {self.node_label})
|
||||
WHERE n.user_id = $user_id
|
||||
WITH n, $n_embedding as n_embedding
|
||||
CALL neptune.algo.vectors.distanceByEmbedding(
|
||||
n_embedding,
|
||||
n,
|
||||
{{metric:"CosineSimilarity"}}
|
||||
) YIELD distance
|
||||
WITH n, distance as similarity
|
||||
WHERE similarity >= $threshold
|
||||
CALL {{
|
||||
WITH n
|
||||
MATCH (n)-[r]->(m)
|
||||
RETURN n.name AS source, id(n) AS source_id, type(r) AS relationship, id(r) AS relation_id, m.name AS destination, id(m) AS destination_id
|
||||
UNION ALL
|
||||
WITH n
|
||||
MATCH (m)-[r]->(n)
|
||||
RETURN m.name AS source, id(m) AS source_id, type(r) AS relationship, id(r) AS relation_id, n.name AS destination, id(n) AS destination_id
|
||||
}}
|
||||
WITH distinct source, source_id, relationship, relation_id, destination, destination_id, similarity
|
||||
RETURN source, source_id, relationship, relation_id, destination, destination_id, similarity
|
||||
ORDER BY similarity DESC
|
||||
LIMIT $limit
|
||||
"""
|
||||
params = {
|
||||
"n_embedding": n_embedding,
|
||||
"threshold": self.threshold,
|
||||
"user_id": filters["user_id"],
|
||||
"limit": top_k,
|
||||
}
|
||||
logger.debug(f"_search_graph_db\n query={cypher_query}")
|
||||
|
||||
return cypher_query, params
|
||||
@@ -1,371 +0,0 @@
|
||||
UPDATE_MEMORY_TOOL_GRAPH = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "update_graph_memory",
|
||||
"description": "Update the relationship key of an existing graph memory based on new information. This function should be called when there's a need to modify an existing relationship in the knowledge graph. The update should only be performed if the new information is more recent, more accurate, or provides additional context compared to the existing information. The source and destination nodes of the relationship must remain the same as in the existing graph memory; only the relationship itself can be updated.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the source node in the relationship to be updated. This should match an existing node in the graph.",
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the destination node in the relationship to be updated. This should match an existing node in the graph.",
|
||||
},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "The new or updated relationship between the source and destination nodes. This should be a concise, clear description of how the two nodes are connected.",
|
||||
},
|
||||
},
|
||||
"required": ["source", "destination", "relationship"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
ADD_MEMORY_TOOL_GRAPH = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "add_graph_memory",
|
||||
"description": "Add a new graph memory to the knowledge graph. This function creates a new relationship between two nodes, potentially creating new nodes if they don't exist.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the source node in the new relationship. This can be an existing node or a new node to be created.",
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the destination node in the new relationship. This can be an existing node or a new node to be created.",
|
||||
},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "The type of relationship between the source and destination nodes. This should be a concise, clear description of how the two nodes are connected.",
|
||||
},
|
||||
"source_type": {
|
||||
"type": "string",
|
||||
"description": "The type or category of the source node. This helps in classifying and organizing nodes in the graph.",
|
||||
},
|
||||
"destination_type": {
|
||||
"type": "string",
|
||||
"description": "The type or category of the destination node. This helps in classifying and organizing nodes in the graph.",
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"source",
|
||||
"destination",
|
||||
"relationship",
|
||||
"source_type",
|
||||
"destination_type",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
NOOP_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "noop",
|
||||
"description": "No operation should be performed to the graph entities. This function is called when the system determines that no changes or additions are necessary based on the current input or context. It serves as a placeholder action when no other actions are required, ensuring that the system can explicitly acknowledge situations where no modifications to the graph are needed.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": [],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
RELATIONS_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "establish_relationships",
|
||||
"description": "Establish relationships among the entities based on the provided text.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"entities": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {"type": "string", "description": "The source entity of the relationship."},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "The relationship between the source and destination entities.",
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "The destination entity of the relationship.",
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"source",
|
||||
"relationship",
|
||||
"destination",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
}
|
||||
},
|
||||
"required": ["entities"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
EXTRACT_ENTITIES_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "extract_entities",
|
||||
"description": "Extract entities and their types from the text.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"entities": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"entity": {"type": "string", "description": "The name or identifier of the entity."},
|
||||
"entity_type": {"type": "string", "description": "The type or category of the entity."},
|
||||
},
|
||||
"required": ["entity", "entity_type"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"description": "An array of entities with their types.",
|
||||
}
|
||||
},
|
||||
"required": ["entities"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
UPDATE_MEMORY_STRUCT_TOOL_GRAPH = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "update_graph_memory",
|
||||
"description": "Update the relationship key of an existing graph memory based on new information. This function should be called when there's a need to modify an existing relationship in the knowledge graph. The update should only be performed if the new information is more recent, more accurate, or provides additional context compared to the existing information. The source and destination nodes of the relationship must remain the same as in the existing graph memory; only the relationship itself can be updated.",
|
||||
"strict": True,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the source node in the relationship to be updated. This should match an existing node in the graph.",
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the destination node in the relationship to be updated. This should match an existing node in the graph.",
|
||||
},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "The new or updated relationship between the source and destination nodes. This should be a concise, clear description of how the two nodes are connected.",
|
||||
},
|
||||
},
|
||||
"required": ["source", "destination", "relationship"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
ADD_MEMORY_STRUCT_TOOL_GRAPH = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "add_graph_memory",
|
||||
"description": "Add a new graph memory to the knowledge graph. This function creates a new relationship between two nodes, potentially creating new nodes if they don't exist.",
|
||||
"strict": True,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the source node in the new relationship. This can be an existing node or a new node to be created.",
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the destination node in the new relationship. This can be an existing node or a new node to be created.",
|
||||
},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "The type of relationship between the source and destination nodes. This should be a concise, clear description of how the two nodes are connected.",
|
||||
},
|
||||
"source_type": {
|
||||
"type": "string",
|
||||
"description": "The type or category of the source node. This helps in classifying and organizing nodes in the graph.",
|
||||
},
|
||||
"destination_type": {
|
||||
"type": "string",
|
||||
"description": "The type or category of the destination node. This helps in classifying and organizing nodes in the graph.",
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"source",
|
||||
"destination",
|
||||
"relationship",
|
||||
"source_type",
|
||||
"destination_type",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
NOOP_STRUCT_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "noop",
|
||||
"description": "No operation should be performed to the graph entities. This function is called when the system determines that no changes or additions are necessary based on the current input or context. It serves as a placeholder action when no other actions are required, ensuring that the system can explicitly acknowledge situations where no modifications to the graph are needed.",
|
||||
"strict": True,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": [],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
RELATIONS_STRUCT_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "establish_relations",
|
||||
"description": "Establish relationships among the entities based on the provided text.",
|
||||
"strict": True,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"entities": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "The source entity of the relationship.",
|
||||
},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "The relationship between the source and destination entities.",
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "The destination entity of the relationship.",
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"source",
|
||||
"relationship",
|
||||
"destination",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
}
|
||||
},
|
||||
"required": ["entities"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
EXTRACT_ENTITIES_STRUCT_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "extract_entities",
|
||||
"description": "Extract entities and their types from the text.",
|
||||
"strict": True,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"entities": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"entity": {"type": "string", "description": "The name or identifier of the entity."},
|
||||
"entity_type": {"type": "string", "description": "The type or category of the entity."},
|
||||
},
|
||||
"required": ["entity", "entity_type"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"description": "An array of entities with their types.",
|
||||
}
|
||||
},
|
||||
"required": ["entities"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "delete_graph_memory",
|
||||
"description": "Delete the relationship between two nodes. This function deletes the existing relationship.",
|
||||
"strict": True,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the source node in the relationship.",
|
||||
},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "The existing relationship between the source and destination nodes that needs to be deleted.",
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the destination node in the relationship.",
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"source",
|
||||
"relationship",
|
||||
"destination",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
DELETE_MEMORY_TOOL_GRAPH = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "delete_graph_memory",
|
||||
"description": "Delete the relationship between two nodes. This function deletes the existing relationship.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the source node in the relationship.",
|
||||
},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "The existing relationship between the source and destination nodes that needs to be deleted.",
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the destination node in the relationship.",
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"source",
|
||||
"relationship",
|
||||
"destination",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -1,97 +0,0 @@
|
||||
UPDATE_GRAPH_PROMPT = """
|
||||
You are an AI expert specializing in graph memory management and optimization. Your task is to analyze existing graph memories alongside new information, and update the relationships in the memory list to ensure the most accurate, current, and coherent representation of knowledge.
|
||||
|
||||
Input:
|
||||
1. Existing Graph Memories: A list of current graph memories, each containing source, target, and relationship information.
|
||||
2. New Graph Memory: Fresh information to be integrated into the existing graph structure.
|
||||
|
||||
Guidelines:
|
||||
1. Identification: Use the source and target as primary identifiers when matching existing memories with new information.
|
||||
2. Conflict Resolution:
|
||||
- If new information contradicts an existing memory:
|
||||
a) For matching source and target but differing content, update the relationship of the existing memory.
|
||||
b) If the new memory provides more recent or accurate information, update the existing memory accordingly.
|
||||
3. Comprehensive Review: Thoroughly examine each existing graph memory against the new information, updating relationships as necessary. Multiple updates may be required.
|
||||
4. Consistency: Maintain a uniform and clear style across all memories. Each entry should be concise yet comprehensive.
|
||||
5. Semantic Coherence: Ensure that updates maintain or improve the overall semantic structure of the graph.
|
||||
6. Temporal Awareness: If timestamps are available, consider the recency of information when making updates.
|
||||
7. Relationship Refinement: Look for opportunities to refine relationship descriptions for greater precision or clarity.
|
||||
8. Redundancy Elimination: Identify and merge any redundant or highly similar relationships that may result from the update.
|
||||
|
||||
Memory Format:
|
||||
source -- RELATIONSHIP -- destination
|
||||
|
||||
Task Details:
|
||||
======= Existing Graph Memories:=======
|
||||
{existing_memories}
|
||||
|
||||
======= New Graph Memory:=======
|
||||
{new_memories}
|
||||
|
||||
Output:
|
||||
Provide a list of update instructions, each specifying the source, target, and the new relationship to be set. Only include memories that require updates.
|
||||
"""
|
||||
|
||||
EXTRACT_RELATIONS_PROMPT = """
|
||||
|
||||
You are an advanced algorithm designed to extract structured information from text to construct knowledge graphs. Your goal is to capture comprehensive and accurate information. Follow these key principles:
|
||||
|
||||
1. Extract only explicitly stated information from the text.
|
||||
2. Establish relationships among the entities provided.
|
||||
3. Use "USER_ID" as the source entity for any self-references (e.g., "I," "me," "my," etc.) in user messages.
|
||||
CUSTOM_PROMPT
|
||||
|
||||
Relationships:
|
||||
- Use consistent, general, and timeless relationship types.
|
||||
- Example: Prefer "professor" over "became_professor."
|
||||
- Relationships should only be established among the entities explicitly mentioned in the user message.
|
||||
|
||||
Entity Consistency:
|
||||
- Ensure that relationships are coherent and logically align with the context of the message.
|
||||
- Maintain consistent naming for entities across the extracted data.
|
||||
|
||||
Strive to construct a coherent and easily understandable knowledge graph by establishing all the relationships among the entities and adherence to the user’s context.
|
||||
|
||||
Adhere strictly to these guidelines to ensure high-quality knowledge graph extraction."""
|
||||
|
||||
DELETE_RELATIONS_SYSTEM_PROMPT = """
|
||||
You are a graph memory manager specializing in identifying, managing, and optimizing relationships within graph-based memories. Your primary task is to analyze a list of existing relationships and determine which ones should be deleted based on the new information provided.
|
||||
Input:
|
||||
1. Existing Graph Memories: A list of current graph memories, each containing source, relationship, and destination information.
|
||||
2. New Text: The new information to be integrated into the existing graph structure.
|
||||
3. Use "USER_ID" as node for any self-references (e.g., "I," "me," "my," etc.) in user messages.
|
||||
|
||||
Guidelines:
|
||||
1. Identification: Use the new information to evaluate existing relationships in the memory graph.
|
||||
2. Deletion Criteria: Delete a relationship only if it meets at least one of these conditions:
|
||||
- Outdated or Inaccurate: The new information is more recent or accurate.
|
||||
- Contradictory: The new information conflicts with or negates the existing information.
|
||||
3. DO NOT DELETE if their is a possibility of same type of relationship but different destination nodes.
|
||||
4. Comprehensive Analysis:
|
||||
- Thoroughly examine each existing relationship against the new information and delete as necessary.
|
||||
- Multiple deletions may be required based on the new information.
|
||||
5. Semantic Integrity:
|
||||
- Ensure that deletions maintain or improve the overall semantic structure of the graph.
|
||||
- Avoid deleting relationships that are NOT contradictory/outdated to the new information.
|
||||
6. Temporal Awareness: Prioritize recency when timestamps are available.
|
||||
7. Necessity Principle: Only DELETE relationships that must be deleted and are contradictory/outdated to the new information to maintain an accurate and coherent memory graph.
|
||||
|
||||
Note: DO NOT DELETE if their is a possibility of same type of relationship but different destination nodes.
|
||||
|
||||
For example:
|
||||
Existing Memory: alice -- loves_to_eat -- pizza
|
||||
New Information: Alice also loves to eat burger.
|
||||
|
||||
Do not delete in the above example because there is a possibility that Alice loves to eat both pizza and burger.
|
||||
|
||||
Memory Format:
|
||||
source -- relationship -- destination
|
||||
|
||||
Provide a list of deletion instructions, each specifying the relationship to be deleted.
|
||||
"""
|
||||
|
||||
|
||||
def get_delete_messages(existing_memories_string, data, user_id):
|
||||
return DELETE_RELATIONS_SYSTEM_PROMPT.replace(
|
||||
"USER_ID", user_id
|
||||
), f"Here are the existing memories: {existing_memories_string} \n\n New Information: {data}"
|
||||
@@ -1,595 +0,0 @@
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
|
||||
from mem0.memory.utils import format_entities, sanitize_relationship_for_cypher
|
||||
|
||||
try:
|
||||
import age
|
||||
except ImportError:
|
||||
raise ImportError("apache-age-python is not installed. Please install it using pip install apache-age-python")
|
||||
|
||||
try:
|
||||
from rank_bm25 import BM25Okapi
|
||||
except ImportError:
|
||||
raise ImportError("rank_bm25 is not installed. Please install it using pip install rank-bm25")
|
||||
|
||||
from mem0.graphs.tools import (
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
DELETE_MEMORY_TOOL_GRAPH,
|
||||
EXTRACT_ENTITIES_STRUCT_TOOL,
|
||||
EXTRACT_ENTITIES_TOOL,
|
||||
RELATIONS_STRUCT_TOOL,
|
||||
RELATIONS_TOOL,
|
||||
)
|
||||
from mem0.graphs.utils import EXTRACT_RELATIONS_PROMPT, get_delete_messages
|
||||
from mem0.utils.factory import EmbedderFactory, LlmFactory
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _cosine_similarity(vec1, vec2):
|
||||
"""Compute cosine similarity between two vectors without numpy."""
|
||||
dot = sum(a * b for a, b in zip(vec1, vec2))
|
||||
norm1 = sum(a * a for a in vec1) ** 0.5
|
||||
norm2 = sum(b * b for b in vec2) ** 0.5
|
||||
if norm1 == 0 or norm2 == 0:
|
||||
return 0.0
|
||||
return dot / (norm1 * norm2)
|
||||
|
||||
|
||||
def _get_similar_nodes(nodes, query_embedding, filters, threshold):
|
||||
"""Find nodes above the similarity threshold from a fetched node list.
|
||||
|
||||
Shared by ``_find_similar_node`` and ``_search_graph_db`` to avoid
|
||||
duplicating the client-side cosine similarity logic.
|
||||
"""
|
||||
matches = []
|
||||
for node in nodes:
|
||||
props = node if isinstance(node, dict) else {}
|
||||
stored_emb = props.get("embedding")
|
||||
if not stored_emb:
|
||||
continue
|
||||
if isinstance(stored_emb, str):
|
||||
stored_emb = json.loads(stored_emb)
|
||||
|
||||
if filters.get("agent_id") and props.get("agent_id") != filters["agent_id"]:
|
||||
continue
|
||||
if filters.get("run_id") and props.get("run_id") != filters["run_id"]:
|
||||
continue
|
||||
|
||||
sim = _cosine_similarity(query_embedding, stored_emb)
|
||||
if sim >= threshold:
|
||||
matches.append({"name": props.get("name"), "similarity": sim, "props": props})
|
||||
|
||||
matches.sort(key=lambda x: x["similarity"], reverse=True)
|
||||
return matches
|
||||
|
||||
|
||||
class MemoryGraph:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
|
||||
graph_cfg = self.config.graph_store.config
|
||||
self.graph_name = graph_cfg.graph_name
|
||||
|
||||
# Connect using the Apache AGE Python driver (psycopg2-based)
|
||||
self.ag = age.connect(
|
||||
graph=self.graph_name,
|
||||
host=graph_cfg.host,
|
||||
port=graph_cfg.port,
|
||||
dbname=graph_cfg.database,
|
||||
user=graph_cfg.username,
|
||||
password=graph_cfg.password,
|
||||
)
|
||||
|
||||
self.embedding_model = EmbedderFactory.create(
|
||||
self.config.embedder.provider, self.config.embedder.config, self.config.vector_store.config
|
||||
)
|
||||
|
||||
# Default to openai if no specific provider is configured
|
||||
self.llm_provider = "openai"
|
||||
if self.config.llm and self.config.llm.provider:
|
||||
self.llm_provider = self.config.llm.provider
|
||||
if self.config.graph_store and self.config.graph_store.llm and self.config.graph_store.llm.provider:
|
||||
self.llm_provider = self.config.graph_store.llm.provider
|
||||
|
||||
# Get LLM config with proper null checks
|
||||
llm_config = None
|
||||
if self.config.graph_store and self.config.graph_store.llm and hasattr(self.config.graph_store.llm, "config"):
|
||||
llm_config = self.config.graph_store.llm.config
|
||||
elif hasattr(self.config.llm, "config"):
|
||||
llm_config = self.config.llm.config
|
||||
self.llm = LlmFactory.create(self.llm_provider, llm_config)
|
||||
self.user_id = None
|
||||
self.threshold = self.config.graph_store.threshold if hasattr(self.config.graph_store, "threshold") else 0.7
|
||||
|
||||
# -- helpers ---------------------------------------------------------------
|
||||
|
||||
def _exec_cypher(self, cypher_stmt, cols=None, params=None):
|
||||
"""Execute a Cypher query via the AGE driver and return results.
|
||||
|
||||
Uses ``ag.execCypher`` which delegates to psycopg2's safe parameter
|
||||
substitution (``%s`` placeholders). The *cols* argument specifies the
|
||||
column names in the ``AS (…)`` clause — when ``None`` the driver
|
||||
defaults to a single ``v agtype`` column.
|
||||
|
||||
When *cols* are provided, returns a list of dicts keyed by column name.
|
||||
When *cols* is ``None``, returns vertex/edge property dicts or raw values.
|
||||
"""
|
||||
cursor = self.ag.execCypher(cypher_stmt, cols=cols, params=params)
|
||||
rows = cursor.fetchall()
|
||||
if not rows:
|
||||
return []
|
||||
|
||||
col_names = [desc[0] for desc in cursor.description] if cursor.description else None
|
||||
results = []
|
||||
for row in rows:
|
||||
if col_names and len(col_names) > 1:
|
||||
record = {}
|
||||
for i, col_name in enumerate(col_names):
|
||||
val = row[i]
|
||||
if hasattr(val, "properties"):
|
||||
record[col_name] = val.properties
|
||||
else:
|
||||
record[col_name] = val
|
||||
results.append(record)
|
||||
else:
|
||||
val = row[0] if len(row) == 1 else row
|
||||
if hasattr(val, "properties"):
|
||||
results.append(val.properties)
|
||||
else:
|
||||
results.append(val)
|
||||
return results
|
||||
|
||||
def _fetch_user_nodes_with_embeddings(self, user_id):
|
||||
"""Fetch all nodes with embeddings for a given user_id."""
|
||||
return self._exec_cypher(
|
||||
"MATCH (n {user_id: %s}) WHERE n.embedding IS NOT NULL RETURN n",
|
||||
params=(user_id,),
|
||||
)
|
||||
|
||||
def _find_similar_node(self, embedding, filters, threshold=0.9):
|
||||
"""Find the most similar existing node by cosine similarity.
|
||||
|
||||
Apache AGE does not have a built-in vector index, so we fetch all
|
||||
node embeddings matching the filters and compute cosine similarity
|
||||
on the client side. This is adequate for moderate graph sizes; for
|
||||
very large graphs consider pairing AGE with pgvector.
|
||||
"""
|
||||
nodes = self._fetch_user_nodes_with_embeddings(filters["user_id"])
|
||||
matches = _get_similar_nodes(nodes, embedding, filters, threshold)
|
||||
return matches[0]["props"] if matches else None
|
||||
|
||||
def _merge_node(self, user_id, name, embedding, agent_id=None, run_id=None):
|
||||
"""Create a node if it doesn't exist, or update mentions if it does.
|
||||
|
||||
Apache AGE does not support ``ON CREATE SET`` / ``ON MATCH SET``, so
|
||||
we use ``MERGE … SET`` which always applies the SET clause. Embeddings
|
||||
and optional filter properties are set in a single query.
|
||||
"""
|
||||
set_parts = [
|
||||
"n.embedding = %s",
|
||||
"n.mentions = coalesce(n.mentions, 0) + 1",
|
||||
"n.created = coalesce(n.created, %s)",
|
||||
]
|
||||
params = [user_id, name, json.dumps(embedding), int(time.time() * 1000)]
|
||||
|
||||
if agent_id:
|
||||
set_parts.append("n.agent_id = %s")
|
||||
params.append(agent_id)
|
||||
if run_id:
|
||||
set_parts.append("n.run_id = %s")
|
||||
params.append(run_id)
|
||||
|
||||
set_clause = ", ".join(set_parts)
|
||||
self._exec_cypher(
|
||||
f"MERGE (n {{user_id: %s, name: %s}}) SET {set_clause}",
|
||||
params=tuple(params),
|
||||
)
|
||||
|
||||
def close(self):
|
||||
"""Close the underlying database connection."""
|
||||
if self.ag:
|
||||
self.ag.close()
|
||||
|
||||
# -- public API ------------------------------------------------------------
|
||||
|
||||
def add(self, data, filters):
|
||||
"""
|
||||
Adds data to the graph.
|
||||
|
||||
Args:
|
||||
data (str): The data to add to the graph.
|
||||
filters (dict): A dictionary containing filters to be applied during the addition.
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
to_be_added = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
to_be_deleted = self._get_delete_entities_from_search_output(search_output, data, filters)
|
||||
|
||||
deleted_entities = self._delete_entities(to_be_deleted, filters)
|
||||
added_entities = self._add_entities(to_be_added, filters, entity_type_map)
|
||||
|
||||
return {"deleted_entities": deleted_entities, "added_entities": added_entities}
|
||||
|
||||
def search(self, query, filters, top_k=100):
|
||||
"""
|
||||
Search for memories and related graph data.
|
||||
|
||||
Args:
|
||||
query (str): Query to search for.
|
||||
filters (dict): A dictionary containing filters to be applied during the search.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
|
||||
Returns:
|
||||
list: A list of dicts with keys "source", "relationship", "destination".
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(query, filters)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
|
||||
if not search_output:
|
||||
return []
|
||||
|
||||
search_outputs_sequence = [
|
||||
[item["source"], item["relationship"], item["destination"]] for item in search_output
|
||||
]
|
||||
bm25 = BM25Okapi(search_outputs_sequence)
|
||||
|
||||
tokenized_query = query.split(" ")
|
||||
reranked_results = bm25.get_top_n(tokenized_query, search_outputs_sequence, n=5)
|
||||
|
||||
search_results = []
|
||||
for item in reranked_results:
|
||||
search_results.append({"source": item[0], "relationship": item[1], "destination": item[2]})
|
||||
|
||||
logger.info(f"Returned {len(search_results)} search results")
|
||||
return search_results
|
||||
|
||||
def delete(self, data, filters):
|
||||
"""
|
||||
Delete graph entities associated with the given memory text.
|
||||
|
||||
Extracts entities and relationships from the memory text using the same
|
||||
pipeline as add(), then deletes the matching relationships in the graph.
|
||||
|
||||
Args:
|
||||
data (str): The memory text whose graph entities should be removed.
|
||||
filters (dict): Scope filters (user_id, agent_id, run_id).
|
||||
"""
|
||||
try:
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
if not entity_type_map:
|
||||
logger.debug("No entities found in memory text, skipping graph cleanup")
|
||||
return
|
||||
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
if to_be_deleted:
|
||||
self._delete_entities(to_be_deleted, filters)
|
||||
except Exception as e:
|
||||
logger.error(f"Error during graph cleanup for memory delete: {e}")
|
||||
|
||||
def delete_all(self, filters):
|
||||
"""Delete all nodes and relationships for a user or specific agent."""
|
||||
where_parts = ["n.user_id = %s"]
|
||||
params = [filters["user_id"]]
|
||||
if filters.get("agent_id"):
|
||||
where_parts.append("n.agent_id = %s")
|
||||
params.append(filters["agent_id"])
|
||||
if filters.get("run_id"):
|
||||
where_parts.append("n.run_id = %s")
|
||||
params.append(filters["run_id"])
|
||||
where_clause = " AND ".join(where_parts)
|
||||
|
||||
self._exec_cypher(
|
||||
f"MATCH (n) WHERE {where_clause} DETACH DELETE n",
|
||||
params=tuple(params),
|
||||
)
|
||||
self.ag.commit()
|
||||
|
||||
def get_all(self, filters, top_k=100):
|
||||
"""
|
||||
Retrieves all nodes and relationships from the graph database based on optional filtering criteria.
|
||||
|
||||
Args:
|
||||
filters (dict): A dictionary containing filters to be applied during the retrieval.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
Returns:
|
||||
list: A list of dictionaries, each containing:
|
||||
- 'source': The source node name.
|
||||
- 'relationship': The relationship type.
|
||||
- 'target': The target node name.
|
||||
"""
|
||||
where_parts = ["n.user_id = %s", "m.user_id = %s"]
|
||||
params = [filters["user_id"], filters["user_id"]]
|
||||
if filters.get("agent_id"):
|
||||
where_parts.extend(["n.agent_id = %s", "m.agent_id = %s"])
|
||||
params.extend([filters["agent_id"], filters["agent_id"]])
|
||||
if filters.get("run_id"):
|
||||
where_parts.extend(["n.run_id = %s", "m.run_id = %s"])
|
||||
params.extend([filters["run_id"], filters["run_id"]])
|
||||
where_clause = " AND ".join(where_parts)
|
||||
params.append(top_k)
|
||||
|
||||
results = self._exec_cypher(
|
||||
f"MATCH (n)-[r]->(m) WHERE {where_clause} "
|
||||
f"RETURN n.name, type(r), m.name LIMIT %s",
|
||||
cols=["source", "relationship", "target"],
|
||||
params=tuple(params),
|
||||
)
|
||||
|
||||
final_results = []
|
||||
for result in results:
|
||||
final_results.append(
|
||||
{
|
||||
"source": result["source"],
|
||||
"relationship": result["relationship"],
|
||||
"target": result["target"],
|
||||
}
|
||||
)
|
||||
|
||||
logger.info(f"Retrieved {len(final_results)} relationships")
|
||||
return final_results
|
||||
|
||||
# -- LLM-driven extraction -------------------------------------------------
|
||||
|
||||
def _retrieve_nodes_from_data(self, data, filters):
|
||||
"""Extracts all the entities mentioned in the query."""
|
||||
_tools = [EXTRACT_ENTITIES_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [EXTRACT_ENTITIES_STRUCT_TOOL]
|
||||
search_results = self.llm.generate_response(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": f"You are a smart assistant who understands entities and their types in a given text. If user message contains self reference such as 'I', 'me', 'my' etc. then use {filters['user_id']} as the source entity. Extract all the entities from the text. ***DO NOT*** answer the question itself if the given text is a question.",
|
||||
},
|
||||
{"role": "user", "content": data},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entity_type_map = {}
|
||||
|
||||
try:
|
||||
for tool_call in search_results["tool_calls"]:
|
||||
if tool_call["name"] != "extract_entities":
|
||||
continue
|
||||
for item in tool_call.get("arguments", {}).get("entities", []):
|
||||
if "entity" in item and "entity_type" in item:
|
||||
entity_type_map[item["entity"]] = item["entity_type"]
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Error in search tool: {e}, llm_provider={self.llm_provider}, search_results={search_results}"
|
||||
)
|
||||
|
||||
entity_type_map = {k.lower().replace(" ", "_"): v.lower().replace(" ", "_") for k, v in entity_type_map.items()}
|
||||
logger.debug(f"Entity type map: {entity_type_map}\n search_results={search_results}")
|
||||
return entity_type_map
|
||||
|
||||
def _establish_nodes_relations_from_data(self, data, filters, entity_type_map):
|
||||
"""Establish relations among the extracted nodes."""
|
||||
|
||||
user_identity = f"user_id: {filters['user_id']}"
|
||||
if filters.get("agent_id"):
|
||||
user_identity += f", agent_id: {filters['agent_id']}"
|
||||
if filters.get("run_id"):
|
||||
user_identity += f", run_id: {filters['run_id']}"
|
||||
|
||||
if self.config.graph_store.custom_prompt:
|
||||
system_content = EXTRACT_RELATIONS_PROMPT.replace("USER_ID", user_identity)
|
||||
system_content = system_content.replace("CUSTOM_PROMPT", f"4. {self.config.graph_store.custom_prompt}")
|
||||
messages = [
|
||||
{"role": "system", "content": system_content},
|
||||
{"role": "user", "content": data},
|
||||
]
|
||||
else:
|
||||
system_content = EXTRACT_RELATIONS_PROMPT.replace("USER_ID", user_identity)
|
||||
messages = [
|
||||
{"role": "system", "content": system_content},
|
||||
{"role": "user", "content": f"List of entities: {list(entity_type_map.keys())}. \n\nText: {data}"},
|
||||
]
|
||||
|
||||
_tools = [RELATIONS_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [RELATIONS_STRUCT_TOOL]
|
||||
|
||||
extracted_entities = self.llm.generate_response(
|
||||
messages=messages,
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entities = []
|
||||
if extracted_entities and extracted_entities.get("tool_calls"):
|
||||
entities = extracted_entities["tool_calls"][0].get("arguments", {}).get("entities", [])
|
||||
|
||||
entities = self._remove_spaces_from_entities(entities)
|
||||
logger.debug(f"Extracted entities: {entities}")
|
||||
return entities
|
||||
|
||||
# -- graph DB operations ---------------------------------------------------
|
||||
|
||||
def _search_graph_db(self, node_list, filters, top_k=100):
|
||||
"""Search similar nodes and their respective incoming and outgoing relations."""
|
||||
result_relations = []
|
||||
|
||||
for node in node_list:
|
||||
n_embedding = self.embedding_model.embed(node)
|
||||
|
||||
nodes = self._fetch_user_nodes_with_embeddings(filters["user_id"])
|
||||
similar_nodes = _get_similar_nodes(nodes, n_embedding, filters, self.threshold)
|
||||
|
||||
# Build WHERE clause for relationship target filtering
|
||||
rel_where_parts = ["m.user_id = %s"]
|
||||
rel_params_suffix = [filters["user_id"]]
|
||||
if filters.get("agent_id"):
|
||||
rel_where_parts.append("m.agent_id = %s")
|
||||
rel_params_suffix.append(filters["agent_id"])
|
||||
if filters.get("run_id"):
|
||||
rel_where_parts.append("m.run_id = %s")
|
||||
rel_params_suffix.append(filters["run_id"])
|
||||
rel_where = " AND ".join(rel_where_parts)
|
||||
|
||||
# For each similar node, fetch its relationships
|
||||
for sn in similar_nodes[:top_k]:
|
||||
node_name = sn["name"]
|
||||
similarity = sn["similarity"]
|
||||
|
||||
out_params = (filters["user_id"], node_name) + tuple(rel_params_suffix)
|
||||
out_results = self._exec_cypher(
|
||||
f"MATCH (n {{user_id: %s, name: %s}})-[r]->(m) "
|
||||
f"WHERE {rel_where} "
|
||||
f"RETURN n.name, type(r), m.name",
|
||||
cols=["source", "relationship", "destination"],
|
||||
params=out_params,
|
||||
)
|
||||
|
||||
in_params = (filters["user_id"], node_name) + tuple(rel_params_suffix)
|
||||
in_results = self._exec_cypher(
|
||||
f"MATCH (n {{user_id: %s, name: %s}})<-[r]-(m) "
|
||||
f"WHERE {rel_where} "
|
||||
f"RETURN m.name, type(r), n.name",
|
||||
cols=["source", "relationship", "destination"],
|
||||
params=in_params,
|
||||
)
|
||||
|
||||
for rel in out_results + in_results:
|
||||
rel["similarity"] = similarity
|
||||
result_relations.append(rel)
|
||||
|
||||
return result_relations
|
||||
|
||||
def _get_delete_entities_from_search_output(self, search_output, data, filters):
|
||||
"""Get the entities to be deleted from the search output."""
|
||||
search_output_string = format_entities(search_output)
|
||||
|
||||
user_identity = f"user_id: {filters['user_id']}"
|
||||
if filters.get("agent_id"):
|
||||
user_identity += f", agent_id: {filters['agent_id']}"
|
||||
if filters.get("run_id"):
|
||||
user_identity += f", run_id: {filters['run_id']}"
|
||||
|
||||
system_prompt, user_prompt = get_delete_messages(search_output_string, data, user_identity)
|
||||
|
||||
_tools = [DELETE_MEMORY_TOOL_GRAPH]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [DELETE_MEMORY_STRUCT_TOOL_GRAPH]
|
||||
|
||||
memory_updates = self.llm.generate_response(
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
to_be_deleted = []
|
||||
for item in memory_updates.get("tool_calls", []):
|
||||
if item.get("name") == "delete_graph_memory":
|
||||
to_be_deleted.append(item.get("arguments"))
|
||||
to_be_deleted = self._remove_spaces_from_entities(to_be_deleted)
|
||||
logger.debug(f"Deleted relationships: {to_be_deleted}")
|
||||
return to_be_deleted
|
||||
|
||||
def _delete_entities(self, to_be_deleted, filters):
|
||||
"""Delete the entities from the graph."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id")
|
||||
run_id = filters.get("run_id")
|
||||
results = []
|
||||
|
||||
try:
|
||||
for item in to_be_deleted:
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
where_parts = [
|
||||
"n.user_id = %s", "n.name = %s",
|
||||
"m.user_id = %s", "m.name = %s",
|
||||
]
|
||||
params = [user_id, source, user_id, destination]
|
||||
if agent_id:
|
||||
where_parts.extend(["n.agent_id = %s", "m.agent_id = %s"])
|
||||
params.extend([agent_id, agent_id])
|
||||
if run_id:
|
||||
where_parts.extend(["n.run_id = %s", "m.run_id = %s"])
|
||||
params.extend([run_id, run_id])
|
||||
where_clause = " AND ".join(where_parts)
|
||||
|
||||
result = self._exec_cypher(
|
||||
f"MATCH (n)-[r:{relationship}]->(m) "
|
||||
f"WHERE {where_clause} "
|
||||
f"DELETE r "
|
||||
f"RETURN n.name, type(r), m.name",
|
||||
cols=["source", "relationship", "target"],
|
||||
params=tuple(params),
|
||||
)
|
||||
results.append(result)
|
||||
|
||||
self.ag.commit()
|
||||
except Exception:
|
||||
self.ag.rollback()
|
||||
raise
|
||||
|
||||
return results
|
||||
|
||||
def _add_entities(self, to_be_added, filters, entity_type_map):
|
||||
"""Add new entities to the graph. Merge nodes if they already exist.
|
||||
|
||||
Apache AGE does not support ``ON CREATE SET`` / ``ON MATCH SET``, so we
|
||||
use ``MERGE … SET`` which always applies. The ``coalesce`` pattern
|
||||
ensures ``created`` is only set on the first merge.
|
||||
"""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id")
|
||||
run_id = filters.get("run_id")
|
||||
results = []
|
||||
|
||||
try:
|
||||
for item in to_be_added:
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
source_embedding = self.embedding_model.embed(source)
|
||||
dest_embedding = self.embedding_model.embed(destination)
|
||||
|
||||
source_match = self._find_similar_node(source_embedding, filters, threshold=self.threshold)
|
||||
dest_match = self._find_similar_node(dest_embedding, filters, threshold=self.threshold)
|
||||
|
||||
effective_source = source_match["name"] if source_match else source
|
||||
effective_dest = dest_match["name"] if dest_match else destination
|
||||
|
||||
# Merge source and destination nodes
|
||||
self._merge_node(user_id, effective_source, source_embedding, agent_id, run_id)
|
||||
self._merge_node(user_id, effective_dest, dest_embedding, agent_id, run_id)
|
||||
|
||||
# Merge relationship
|
||||
result = self._exec_cypher(
|
||||
f"MATCH (s {{user_id: %s, name: %s}}), (d {{user_id: %s, name: %s}}) "
|
||||
f"MERGE (s)-[r:{relationship}]->(d) "
|
||||
f"RETURN s.name, type(r), d.name",
|
||||
cols=["source", "relationship", "target"],
|
||||
params=(user_id, effective_source, user_id, effective_dest),
|
||||
)
|
||||
results.append(result)
|
||||
|
||||
self.ag.commit()
|
||||
except Exception:
|
||||
self.ag.rollback()
|
||||
raise
|
||||
|
||||
return results
|
||||
|
||||
def _remove_spaces_from_entities(self, entity_list):
|
||||
for item in entity_list:
|
||||
item["source"] = item["source"].lower().replace(" ", "_")
|
||||
item["relationship"] = sanitize_relationship_for_cypher(item["relationship"].lower().replace(" ", "_"))
|
||||
item["destination"] = item["destination"].lower().replace(" ", "_")
|
||||
return entity_list
|
||||
|
||||
def reset(self):
|
||||
"""Reset the graph by clearing all nodes and relationships."""
|
||||
logger.warning("Clearing graph...")
|
||||
self._exec_cypher("MATCH (n) DETACH DELETE n")
|
||||
self.ag.commit()
|
||||
@@ -1,744 +0,0 @@
|
||||
import logging
|
||||
|
||||
from mem0.memory.utils import format_entities, remove_spaces_from_entities
|
||||
|
||||
try:
|
||||
from langchain_neo4j import Neo4jGraph
|
||||
except ImportError:
|
||||
raise ImportError("langchain_neo4j is not installed. Please install it using pip install langchain-neo4j")
|
||||
|
||||
try:
|
||||
from rank_bm25 import BM25Okapi
|
||||
except ImportError:
|
||||
raise ImportError("rank_bm25 is not installed. Please install it using pip install rank-bm25")
|
||||
|
||||
from mem0.graphs.tools import (
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
DELETE_MEMORY_TOOL_GRAPH,
|
||||
EXTRACT_ENTITIES_STRUCT_TOOL,
|
||||
EXTRACT_ENTITIES_TOOL,
|
||||
RELATIONS_STRUCT_TOOL,
|
||||
RELATIONS_TOOL,
|
||||
)
|
||||
from mem0.graphs.utils import EXTRACT_RELATIONS_PROMPT, get_delete_messages
|
||||
from mem0.utils.factory import EmbedderFactory, LlmFactory
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MemoryGraph:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
self.graph = Neo4jGraph(
|
||||
url=self.config.graph_store.config.url,
|
||||
username=self.config.graph_store.config.username,
|
||||
password=self.config.graph_store.config.password,
|
||||
database=self.config.graph_store.config.database,
|
||||
refresh_schema=False,
|
||||
driver_config={"notifications_min_severity": "OFF"},
|
||||
)
|
||||
self.embedding_model = EmbedderFactory.create(
|
||||
self.config.embedder.provider, self.config.embedder.config, self.config.vector_store.config
|
||||
)
|
||||
self.node_label = ":`__Entity__`" if self.config.graph_store.config.base_label else ""
|
||||
|
||||
if self.config.graph_store.config.base_label:
|
||||
# Safely add user_id index
|
||||
try:
|
||||
self.graph.query(f"CREATE INDEX entity_single IF NOT EXISTS FOR (n {self.node_label}) ON (n.user_id)")
|
||||
except Exception:
|
||||
pass
|
||||
try: # Safely try to add composite index (Enterprise only)
|
||||
self.graph.query(
|
||||
f"CREATE INDEX entity_composite IF NOT EXISTS FOR (n {self.node_label}) ON (n.name, n.user_id)"
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Default to openai if no specific provider is configured
|
||||
self.llm_provider = "openai"
|
||||
if self.config.llm and self.config.llm.provider:
|
||||
self.llm_provider = self.config.llm.provider
|
||||
if self.config.graph_store and self.config.graph_store.llm and self.config.graph_store.llm.provider:
|
||||
self.llm_provider = self.config.graph_store.llm.provider
|
||||
|
||||
# Get LLM config with proper null checks
|
||||
llm_config = None
|
||||
if self.config.graph_store and self.config.graph_store.llm and hasattr(self.config.graph_store.llm, "config"):
|
||||
llm_config = self.config.graph_store.llm.config
|
||||
elif hasattr(self.config.llm, "config"):
|
||||
llm_config = self.config.llm.config
|
||||
self.llm = LlmFactory.create(self.llm_provider, llm_config)
|
||||
self.user_id = None
|
||||
# Use threshold from graph_store config, default to 0.7 for backward compatibility
|
||||
self.threshold = self.config.graph_store.threshold if hasattr(self.config.graph_store, 'threshold') else 0.7
|
||||
|
||||
def add(self, data, filters):
|
||||
"""
|
||||
Adds data to the graph.
|
||||
|
||||
Args:
|
||||
data (str): The data to add to the graph.
|
||||
filters (dict): A dictionary containing filters to be applied during the addition.
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
to_be_added = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
to_be_deleted = self._get_delete_entities_from_search_output(search_output, data, filters)
|
||||
|
||||
# TODO: Batch queries with APOC plugin
|
||||
# TODO: Add more filter support
|
||||
deleted_entities = self._delete_entities(to_be_deleted, filters)
|
||||
added_entities = self._add_entities(to_be_added, filters, entity_type_map)
|
||||
|
||||
return {"deleted_entities": deleted_entities, "added_entities": added_entities}
|
||||
|
||||
def search(self, query, filters, top_k=100):
|
||||
"""
|
||||
Search for memories and related graph data.
|
||||
|
||||
Args:
|
||||
query (str): Query to search for.
|
||||
filters (dict): A dictionary containing filters to be applied during the search.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing:
|
||||
- "contexts": List of search results from the base data store.
|
||||
- "entities": List of related graph data based on the query.
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(query, filters)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
|
||||
if not search_output:
|
||||
return []
|
||||
|
||||
search_outputs_sequence = [
|
||||
[item["source"], item["relationship"], item["destination"]] for item in search_output
|
||||
]
|
||||
bm25 = BM25Okapi(search_outputs_sequence)
|
||||
|
||||
tokenized_query = query.split(" ")
|
||||
reranked_results = bm25.get_top_n(tokenized_query, search_outputs_sequence, n=5)
|
||||
|
||||
search_results = []
|
||||
for item in reranked_results:
|
||||
search_results.append({"source": item[0], "relationship": item[1], "destination": item[2]})
|
||||
|
||||
logger.info(f"Returned {len(search_results)} search results")
|
||||
|
||||
return search_results
|
||||
|
||||
def delete(self, data, filters):
|
||||
"""
|
||||
Delete graph entities associated with the given memory text.
|
||||
|
||||
Extracts entities and relationships from the memory text using the same
|
||||
pipeline as add(), then soft-deletes the matching relationships in the graph.
|
||||
|
||||
Args:
|
||||
data (str): The memory text whose graph entities should be removed.
|
||||
filters (dict): Scope filters (user_id, agent_id, run_id).
|
||||
"""
|
||||
try:
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
if not entity_type_map:
|
||||
logger.debug("No entities found in memory text, skipping graph cleanup")
|
||||
return
|
||||
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
if to_be_deleted:
|
||||
self._delete_entities(to_be_deleted, filters)
|
||||
except Exception as e:
|
||||
logger.error(f"Error during graph cleanup for memory delete: {e}")
|
||||
|
||||
def delete_all(self, filters):
|
||||
# Build node properties for filtering
|
||||
node_props = ["user_id: $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
node_props.append("agent_id: $agent_id")
|
||||
if filters.get("run_id"):
|
||||
node_props.append("run_id: $run_id")
|
||||
node_props_str = ", ".join(node_props)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{{node_props_str}}})
|
||||
DETACH DELETE n
|
||||
"""
|
||||
params = {"user_id": filters["user_id"]}
|
||||
if filters.get("agent_id"):
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
params["run_id"] = filters["run_id"]
|
||||
self.graph.query(cypher, params=params)
|
||||
|
||||
def get_all(self, filters, top_k=100):
|
||||
"""
|
||||
Retrieves all nodes and relationships from the graph database based on optional filtering criteria.
|
||||
Args:
|
||||
filters (dict): A dictionary containing filters to be applied during the retrieval.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
Returns:
|
||||
list: A list of dictionaries, each containing:
|
||||
- 'contexts': The base data store response for each memory.
|
||||
- 'entities': A list of strings representing the nodes and relationships
|
||||
"""
|
||||
params = {"user_id": filters["user_id"], "limit": top_k}
|
||||
|
||||
# Build node properties based on filters
|
||||
node_props = ["user_id: $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
node_props.append("agent_id: $agent_id")
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
node_props.append("run_id: $run_id")
|
||||
params["run_id"] = filters["run_id"]
|
||||
node_props_str = ", ".join(node_props)
|
||||
|
||||
query = f"""
|
||||
MATCH (n {self.node_label} {{{node_props_str}}})-[r]->(m {self.node_label} {{{node_props_str}}})
|
||||
WHERE r.valid IS NULL OR r.valid = true
|
||||
RETURN n.name AS source, type(r) AS relationship, m.name AS target
|
||||
LIMIT $limit
|
||||
"""
|
||||
results = self.graph.query(query, params=params)
|
||||
|
||||
final_results = []
|
||||
for result in results:
|
||||
final_results.append(
|
||||
{
|
||||
"source": result["source"],
|
||||
"relationship": result["relationship"],
|
||||
"target": result["target"],
|
||||
}
|
||||
)
|
||||
|
||||
logger.info(f"Retrieved {len(final_results)} relationships")
|
||||
|
||||
return final_results
|
||||
|
||||
def _retrieve_nodes_from_data(self, data, filters):
|
||||
"""Extracts all the entities mentioned in the query."""
|
||||
_tools = [EXTRACT_ENTITIES_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [EXTRACT_ENTITIES_STRUCT_TOOL]
|
||||
search_results = self.llm.generate_response(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": f"You are a smart assistant who understands entities and their types in a given text. If user message contains self reference such as 'I', 'me', 'my' etc. then use {filters['user_id']} as the source entity. Extract all the entities from the text. ***DO NOT*** answer the question itself if the given text is a question.",
|
||||
},
|
||||
{"role": "user", "content": data},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entity_type_map = {}
|
||||
|
||||
try:
|
||||
for tool_call in search_results["tool_calls"]:
|
||||
if tool_call["name"] != "extract_entities":
|
||||
continue
|
||||
for item in tool_call.get("arguments", {}).get("entities", []):
|
||||
entity_type_map[item["entity"]] = item["entity_type"]
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Error in search tool: {e}, llm_provider={self.llm_provider}, search_results={search_results}"
|
||||
)
|
||||
|
||||
entity_type_map = {k.lower().replace(" ", "_"): v.lower().replace(" ", "_") for k, v in entity_type_map.items()}
|
||||
logger.debug(f"Entity type map: {entity_type_map}\n search_results={search_results}")
|
||||
return entity_type_map
|
||||
|
||||
def _establish_nodes_relations_from_data(self, data, filters, entity_type_map):
|
||||
"""Establish relations among the extracted nodes."""
|
||||
|
||||
# Compose user identification string for prompt
|
||||
user_identity = f"user_id: {filters['user_id']}"
|
||||
if filters.get("agent_id"):
|
||||
user_identity += f", agent_id: {filters['agent_id']}"
|
||||
if filters.get("run_id"):
|
||||
user_identity += f", run_id: {filters['run_id']}"
|
||||
|
||||
if self.config.graph_store.custom_prompt:
|
||||
system_content = EXTRACT_RELATIONS_PROMPT.replace("USER_ID", user_identity)
|
||||
# Add the custom prompt line if configured
|
||||
system_content = system_content.replace("CUSTOM_PROMPT", f"4. {self.config.graph_store.custom_prompt}")
|
||||
messages = [
|
||||
{"role": "system", "content": system_content},
|
||||
{"role": "user", "content": data},
|
||||
]
|
||||
else:
|
||||
system_content = EXTRACT_RELATIONS_PROMPT.replace("USER_ID", user_identity)
|
||||
messages = [
|
||||
{"role": "system", "content": system_content},
|
||||
{"role": "user", "content": f"List of entities: {list(entity_type_map.keys())}. \n\nText: {data}"},
|
||||
]
|
||||
|
||||
_tools = [RELATIONS_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [RELATIONS_STRUCT_TOOL]
|
||||
|
||||
extracted_entities = self.llm.generate_response(
|
||||
messages=messages,
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entities = []
|
||||
if extracted_entities.get("tool_calls"):
|
||||
entities = extracted_entities["tool_calls"][0].get("arguments", {}).get("entities", [])
|
||||
|
||||
entities = self._remove_spaces_from_entities(entities)
|
||||
logger.debug(f"Extracted entities: {entities}")
|
||||
return entities
|
||||
|
||||
def _search_graph_db(self, node_list, filters, top_k=100):
|
||||
"""Search similar nodes among and their respective incoming and outgoing relations."""
|
||||
result_relations = []
|
||||
|
||||
# Build node properties for filtering
|
||||
node_props = ["user_id: $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
node_props.append("agent_id: $agent_id")
|
||||
if filters.get("run_id"):
|
||||
node_props.append("run_id: $run_id")
|
||||
node_props_str = ", ".join(node_props)
|
||||
|
||||
for node in node_list:
|
||||
n_embedding = self.embedding_model.embed(node)
|
||||
|
||||
cypher_query = f"""
|
||||
MATCH (n {self.node_label} {{{node_props_str}}})
|
||||
WHERE n.embedding IS NOT NULL
|
||||
WITH n, round(2 * vector.similarity.cosine(n.embedding, $n_embedding) - 1, 4) AS similarity // denormalize for backward compatibility
|
||||
WHERE similarity >= $threshold
|
||||
CALL {{
|
||||
WITH n
|
||||
MATCH (n)-[r]->(m {self.node_label} {{{node_props_str}}})
|
||||
WHERE r.valid IS NULL OR r.valid = true
|
||||
RETURN n.name AS source, elementId(n) AS source_id, type(r) AS relationship, elementId(r) AS relation_id, m.name AS destination, elementId(m) AS destination_id
|
||||
UNION
|
||||
WITH n
|
||||
MATCH (n)<-[r]-(m {self.node_label} {{{node_props_str}}})
|
||||
WHERE r.valid IS NULL OR r.valid = true
|
||||
RETURN m.name AS source, elementId(m) AS source_id, type(r) AS relationship, elementId(r) AS relation_id, n.name AS destination, elementId(n) AS destination_id
|
||||
}}
|
||||
WITH distinct source, source_id, relationship, relation_id, destination, destination_id, similarity
|
||||
RETURN source, source_id, relationship, relation_id, destination, destination_id, similarity
|
||||
ORDER BY similarity DESC
|
||||
LIMIT $limit
|
||||
"""
|
||||
|
||||
params = {
|
||||
"n_embedding": n_embedding,
|
||||
"threshold": self.threshold,
|
||||
"user_id": filters["user_id"],
|
||||
"limit": top_k,
|
||||
}
|
||||
if filters.get("agent_id"):
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
params["run_id"] = filters["run_id"]
|
||||
|
||||
ans = self.graph.query(cypher_query, params=params)
|
||||
result_relations.extend(ans)
|
||||
|
||||
return result_relations
|
||||
|
||||
def _get_delete_entities_from_search_output(self, search_output, data, filters):
|
||||
"""Get the entities to be deleted from the search output."""
|
||||
search_output_string = format_entities(search_output)
|
||||
|
||||
# Compose user identification string for prompt
|
||||
user_identity = f"user_id: {filters['user_id']}"
|
||||
if filters.get("agent_id"):
|
||||
user_identity += f", agent_id: {filters['agent_id']}"
|
||||
if filters.get("run_id"):
|
||||
user_identity += f", run_id: {filters['run_id']}"
|
||||
|
||||
system_prompt, user_prompt = get_delete_messages(search_output_string, data, user_identity)
|
||||
|
||||
_tools = [DELETE_MEMORY_TOOL_GRAPH]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
]
|
||||
|
||||
memory_updates = self.llm.generate_response(
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
to_be_deleted = []
|
||||
for item in memory_updates.get("tool_calls", []):
|
||||
if item.get("name") == "delete_graph_memory":
|
||||
to_be_deleted.append(item.get("arguments"))
|
||||
# Clean entities formatting
|
||||
to_be_deleted = self._remove_spaces_from_entities(to_be_deleted)
|
||||
logger.debug(f"Deleted relationships: {to_be_deleted}")
|
||||
return to_be_deleted
|
||||
|
||||
def _delete_entities(self, to_be_deleted, filters):
|
||||
"""Delete the entities from the graph."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id", None)
|
||||
run_id = filters.get("run_id", None)
|
||||
results = []
|
||||
|
||||
for item in to_be_deleted:
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
# Build the agent filter for the query
|
||||
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"user_id": user_id,
|
||||
}
|
||||
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
params["run_id"] = run_id
|
||||
|
||||
# Build node properties for filtering
|
||||
source_props = ["name: $source_name", "user_id: $user_id"]
|
||||
dest_props = ["name: $dest_name", "user_id: $user_id"]
|
||||
if agent_id:
|
||||
source_props.append("agent_id: $agent_id")
|
||||
dest_props.append("agent_id: $agent_id")
|
||||
if run_id:
|
||||
source_props.append("run_id: $run_id")
|
||||
dest_props.append("run_id: $run_id")
|
||||
source_props_str = ", ".join(source_props)
|
||||
dest_props_str = ", ".join(dest_props)
|
||||
|
||||
# Soft-delete: mark relationship as invalid instead of removing it,
|
||||
# enabling temporal reasoning over historical graph state.
|
||||
# See: https://github.com/mem0ai/mem0/issues/4187
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{{source_props_str}}})
|
||||
-[r:{relationship}]->
|
||||
(m {self.node_label} {{{dest_props_str}}})
|
||||
WHERE r.valid IS NULL OR r.valid = true
|
||||
SET r.valid = false, r.invalidated_at = datetime()
|
||||
RETURN
|
||||
n.name AS source,
|
||||
m.name AS target,
|
||||
type(r) AS relationship
|
||||
"""
|
||||
|
||||
result = self.graph.query(cypher, params=params)
|
||||
results.append(result)
|
||||
|
||||
return results
|
||||
|
||||
def _add_entities(self, to_be_added, filters, entity_type_map):
|
||||
"""Add the new entities to the graph. Merge the nodes if they already exist."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id", None)
|
||||
run_id = filters.get("run_id", None)
|
||||
results = []
|
||||
for item in to_be_added:
|
||||
# entities
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
# types
|
||||
source_type = entity_type_map.get(source, "__User__")
|
||||
source_label = self.node_label if self.node_label else f":`{source_type}`"
|
||||
source_extra_set = f", source:`{source_type}`" if self.node_label else ""
|
||||
destination_type = entity_type_map.get(destination, "__User__")
|
||||
destination_label = self.node_label if self.node_label else f":`{destination_type}`"
|
||||
destination_extra_set = f", destination:`{destination_type}`" if self.node_label else ""
|
||||
|
||||
# embeddings
|
||||
source_embedding = self.embedding_model.embed(source)
|
||||
dest_embedding = self.embedding_model.embed(destination)
|
||||
|
||||
# search for the nodes with the closest embeddings
|
||||
source_node_search_result = self._search_source_node(source_embedding, filters, threshold=self.threshold)
|
||||
destination_node_search_result = self._search_destination_node(dest_embedding, filters, threshold=self.threshold)
|
||||
|
||||
# TODO: Create a cypher query and common params for all the cases
|
||||
if not destination_node_search_result and source_node_search_result:
|
||||
# Build destination MERGE properties
|
||||
merge_props = ["name: $destination_name", "user_id: $user_id"]
|
||||
if agent_id:
|
||||
merge_props.append("agent_id: $agent_id")
|
||||
if run_id:
|
||||
merge_props.append("run_id: $run_id")
|
||||
merge_props_str = ", ".join(merge_props)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source)
|
||||
WHERE elementId(source) = $source_id
|
||||
SET source.mentions = coalesce(source.mentions, 0) + 1
|
||||
WITH source
|
||||
MERGE (destination {destination_label} {{{merge_props_str}}})
|
||||
ON CREATE SET
|
||||
destination.created = timestamp(),
|
||||
destination.mentions = 1
|
||||
{destination_extra_set}
|
||||
ON MATCH SET
|
||||
destination.mentions = coalesce(destination.mentions, 0) + 1
|
||||
WITH source, destination
|
||||
CALL db.create.setNodeVectorProperty(destination, 'embedding', $destination_embedding)
|
||||
WITH source, destination
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created_at = timestamp(),
|
||||
r.updated_at = timestamp(),
|
||||
r.mentions = 1,
|
||||
r.valid = true
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1,
|
||||
r.valid = true,
|
||||
r.updated_at = timestamp(),
|
||||
r.invalidated_at = null
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"source_id": source_node_search_result[0]["elementId(source_candidate)"],
|
||||
"destination_name": destination,
|
||||
"destination_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
params["run_id"] = run_id
|
||||
|
||||
elif destination_node_search_result and not source_node_search_result:
|
||||
# Build source MERGE properties
|
||||
merge_props = ["name: $source_name", "user_id: $user_id"]
|
||||
if agent_id:
|
||||
merge_props.append("agent_id: $agent_id")
|
||||
if run_id:
|
||||
merge_props.append("run_id: $run_id")
|
||||
merge_props_str = ", ".join(merge_props)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (destination)
|
||||
WHERE elementId(destination) = $destination_id
|
||||
SET destination.mentions = coalesce(destination.mentions, 0) + 1
|
||||
WITH destination
|
||||
MERGE (source {source_label} {{{merge_props_str}}})
|
||||
ON CREATE SET
|
||||
source.created = timestamp(),
|
||||
source.mentions = 1
|
||||
{source_extra_set}
|
||||
ON MATCH SET
|
||||
source.mentions = coalesce(source.mentions, 0) + 1
|
||||
WITH source, destination
|
||||
CALL db.create.setNodeVectorProperty(source, 'embedding', $source_embedding)
|
||||
WITH source, destination
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created_at = timestamp(),
|
||||
r.updated_at = timestamp(),
|
||||
r.mentions = 1,
|
||||
r.valid = true
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1,
|
||||
r.valid = true,
|
||||
r.updated_at = timestamp(),
|
||||
r.invalidated_at = null
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"destination_id": destination_node_search_result[0]["elementId(destination_candidate)"],
|
||||
"source_name": source,
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
params["run_id"] = run_id
|
||||
|
||||
elif source_node_search_result and destination_node_search_result:
|
||||
cypher = f"""
|
||||
MATCH (source)
|
||||
WHERE elementId(source) = $source_id
|
||||
SET source.mentions = coalesce(source.mentions, 0) + 1
|
||||
WITH source
|
||||
MATCH (destination)
|
||||
WHERE elementId(destination) = $destination_id
|
||||
SET destination.mentions = coalesce(destination.mentions, 0) + 1
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created_at = timestamp(),
|
||||
r.updated_at = timestamp(),
|
||||
r.mentions = 1,
|
||||
r.valid = true
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1,
|
||||
r.valid = true,
|
||||
r.updated_at = timestamp(),
|
||||
r.invalidated_at = null
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"source_id": source_node_search_result[0]["elementId(source_candidate)"],
|
||||
"destination_id": destination_node_search_result[0]["elementId(destination_candidate)"],
|
||||
"user_id": user_id,
|
||||
}
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
params["run_id"] = run_id
|
||||
|
||||
else:
|
||||
# Build dynamic MERGE props for both source and destination
|
||||
source_props = ["name: $source_name", "user_id: $user_id"]
|
||||
dest_props = ["name: $dest_name", "user_id: $user_id"]
|
||||
if agent_id:
|
||||
source_props.append("agent_id: $agent_id")
|
||||
dest_props.append("agent_id: $agent_id")
|
||||
if run_id:
|
||||
source_props.append("run_id: $run_id")
|
||||
dest_props.append("run_id: $run_id")
|
||||
source_props_str = ", ".join(source_props)
|
||||
dest_props_str = ", ".join(dest_props)
|
||||
|
||||
cypher = f"""
|
||||
MERGE (source {source_label} {{{source_props_str}}})
|
||||
ON CREATE SET source.created = timestamp(),
|
||||
source.mentions = 1
|
||||
{source_extra_set}
|
||||
ON MATCH SET source.mentions = coalesce(source.mentions, 0) + 1
|
||||
WITH source
|
||||
CALL db.create.setNodeVectorProperty(source, 'embedding', $source_embedding)
|
||||
WITH source
|
||||
MERGE (destination {destination_label} {{{dest_props_str}}})
|
||||
ON CREATE SET destination.created = timestamp(),
|
||||
destination.mentions = 1
|
||||
{destination_extra_set}
|
||||
ON MATCH SET destination.mentions = coalesce(destination.mentions, 0) + 1
|
||||
WITH source, destination
|
||||
CALL db.create.setNodeVectorProperty(destination, 'embedding', $dest_embedding)
|
||||
WITH source, destination
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created_at = timestamp(),
|
||||
r.updated_at = timestamp(),
|
||||
r.mentions = 1,
|
||||
r.valid = true
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1,
|
||||
r.valid = true,
|
||||
r.updated_at = timestamp(),
|
||||
r.invalidated_at = null
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"source_embedding": source_embedding,
|
||||
"dest_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
params["run_id"] = run_id
|
||||
result = self.graph.query(cypher, params=params)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
def _remove_spaces_from_entities(self, entity_list):
|
||||
return remove_spaces_from_entities(entity_list, sanitize_relationship=True)
|
||||
|
||||
def _search_source_node(self, source_embedding, filters, threshold=0.9):
|
||||
# Build WHERE conditions
|
||||
where_conditions = ["source_candidate.embedding IS NOT NULL", "source_candidate.user_id = $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
where_conditions.append("source_candidate.agent_id = $agent_id")
|
||||
if filters.get("run_id"):
|
||||
where_conditions.append("source_candidate.run_id = $run_id")
|
||||
where_clause = " AND ".join(where_conditions)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source_candidate {self.node_label})
|
||||
WHERE {where_clause}
|
||||
|
||||
WITH source_candidate,
|
||||
round(2 * vector.similarity.cosine(source_candidate.embedding, $source_embedding) - 1, 4) AS source_similarity // denormalize for backward compatibility
|
||||
WHERE source_similarity >= $threshold
|
||||
|
||||
WITH source_candidate, source_similarity
|
||||
ORDER BY source_similarity DESC
|
||||
LIMIT 1
|
||||
|
||||
RETURN elementId(source_candidate)
|
||||
"""
|
||||
|
||||
params = {
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": filters["user_id"],
|
||||
"threshold": threshold,
|
||||
}
|
||||
if filters.get("agent_id"):
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
params["run_id"] = filters["run_id"]
|
||||
|
||||
result = self.graph.query(cypher, params=params)
|
||||
return result
|
||||
|
||||
def _search_destination_node(self, destination_embedding, filters, threshold=0.9):
|
||||
# Build WHERE conditions
|
||||
where_conditions = ["destination_candidate.embedding IS NOT NULL", "destination_candidate.user_id = $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
where_conditions.append("destination_candidate.agent_id = $agent_id")
|
||||
if filters.get("run_id"):
|
||||
where_conditions.append("destination_candidate.run_id = $run_id")
|
||||
where_clause = " AND ".join(where_conditions)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (destination_candidate {self.node_label})
|
||||
WHERE {where_clause}
|
||||
|
||||
WITH destination_candidate,
|
||||
round(2 * vector.similarity.cosine(destination_candidate.embedding, $destination_embedding) - 1, 4) AS destination_similarity // denormalize for backward compatibility
|
||||
|
||||
WHERE destination_similarity >= $threshold
|
||||
|
||||
WITH destination_candidate, destination_similarity
|
||||
ORDER BY destination_similarity DESC
|
||||
LIMIT 1
|
||||
|
||||
RETURN elementId(destination_candidate)
|
||||
"""
|
||||
|
||||
params = {
|
||||
"destination_embedding": destination_embedding,
|
||||
"user_id": filters["user_id"],
|
||||
"threshold": threshold,
|
||||
}
|
||||
if filters.get("agent_id"):
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
params["run_id"] = filters["run_id"]
|
||||
|
||||
result = self.graph.query(cypher, params=params)
|
||||
return result
|
||||
|
||||
# Reset is not defined in base.py
|
||||
def reset(self):
|
||||
"""Reset the graph by clearing all nodes and relationships."""
|
||||
logger.warning("Clearing graph...")
|
||||
cypher_query = """
|
||||
MATCH (n) DETACH DELETE n
|
||||
"""
|
||||
return self.graph.query(cypher_query)
|
||||
@@ -1,732 +0,0 @@
|
||||
import logging
|
||||
|
||||
from mem0.memory.utils import format_entities, remove_spaces_from_entities
|
||||
|
||||
try:
|
||||
import kuzu
|
||||
except ImportError:
|
||||
raise ImportError("kuzu is not installed. Please install it using pip install kuzu")
|
||||
|
||||
try:
|
||||
from rank_bm25 import BM25Okapi
|
||||
except ImportError:
|
||||
raise ImportError("rank_bm25 is not installed. Please install it using pip install rank-bm25")
|
||||
|
||||
from mem0.graphs.tools import (
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
DELETE_MEMORY_TOOL_GRAPH,
|
||||
EXTRACT_ENTITIES_STRUCT_TOOL,
|
||||
EXTRACT_ENTITIES_TOOL,
|
||||
RELATIONS_STRUCT_TOOL,
|
||||
RELATIONS_TOOL,
|
||||
)
|
||||
from mem0.graphs.utils import EXTRACT_RELATIONS_PROMPT, get_delete_messages
|
||||
from mem0.utils.factory import EmbedderFactory, LlmFactory
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MemoryGraph:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
|
||||
self.embedding_model = EmbedderFactory.create(
|
||||
self.config.embedder.provider,
|
||||
self.config.embedder.config,
|
||||
self.config.vector_store.config,
|
||||
)
|
||||
self.embedding_dims = self.embedding_model.config.embedding_dims
|
||||
|
||||
if self.embedding_dims is None or self.embedding_dims <= 0:
|
||||
raise ValueError(f"embedding_dims must be a positive integer. Given: {self.embedding_dims}")
|
||||
|
||||
self.db = kuzu.Database(self.config.graph_store.config.db)
|
||||
self.graph = kuzu.Connection(self.db)
|
||||
|
||||
self.node_label = ":Entity"
|
||||
self.rel_label = ":CONNECTED_TO"
|
||||
self.kuzu_create_schema()
|
||||
|
||||
# Default to openai if no specific provider is configured
|
||||
self.llm_provider = "openai"
|
||||
if self.config.llm and self.config.llm.provider:
|
||||
self.llm_provider = self.config.llm.provider
|
||||
if self.config.graph_store and self.config.graph_store.llm and self.config.graph_store.llm.provider:
|
||||
self.llm_provider = self.config.graph_store.llm.provider
|
||||
# Get LLM config with proper null checks
|
||||
llm_config = None
|
||||
if self.config.graph_store and self.config.graph_store.llm and hasattr(self.config.graph_store.llm, "config"):
|
||||
llm_config = self.config.graph_store.llm.config
|
||||
elif hasattr(self.config.llm, "config"):
|
||||
llm_config = self.config.llm.config
|
||||
self.llm = LlmFactory.create(self.llm_provider, llm_config)
|
||||
|
||||
self.user_id = None
|
||||
# Use threshold from graph_store config, default to 0.7 for backward compatibility
|
||||
self.threshold = self.config.graph_store.threshold if hasattr(self.config.graph_store, 'threshold') else 0.7
|
||||
|
||||
def kuzu_create_schema(self):
|
||||
self.kuzu_execute(
|
||||
"""
|
||||
CREATE NODE TABLE IF NOT EXISTS Entity(
|
||||
id SERIAL PRIMARY KEY,
|
||||
user_id STRING,
|
||||
agent_id STRING,
|
||||
run_id STRING,
|
||||
name STRING,
|
||||
mentions INT64,
|
||||
created TIMESTAMP,
|
||||
embedding FLOAT[]);
|
||||
"""
|
||||
)
|
||||
self.kuzu_execute(
|
||||
"""
|
||||
CREATE REL TABLE IF NOT EXISTS CONNECTED_TO(
|
||||
FROM Entity TO Entity,
|
||||
name STRING,
|
||||
mentions INT64,
|
||||
created TIMESTAMP,
|
||||
updated TIMESTAMP
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
def kuzu_execute(self, query, parameters=None):
|
||||
results = self.graph.execute(query, parameters)
|
||||
return list(results.rows_as_dict())
|
||||
|
||||
def add(self, data, filters):
|
||||
"""
|
||||
Adds data to the graph.
|
||||
|
||||
Args:
|
||||
data (str): The data to add to the graph.
|
||||
filters (dict): A dictionary containing filters to be applied during the addition.
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
to_be_added = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
to_be_deleted = self._get_delete_entities_from_search_output(search_output, data, filters)
|
||||
|
||||
deleted_entities = self._delete_entities(to_be_deleted, filters)
|
||||
added_entities = self._add_entities(to_be_added, filters, entity_type_map)
|
||||
|
||||
return {"deleted_entities": deleted_entities, "added_entities": added_entities}
|
||||
|
||||
def search(self, query, filters, top_k=5):
|
||||
"""
|
||||
Search for memories and related graph data.
|
||||
|
||||
Args:
|
||||
query (str): Query to search for.
|
||||
filters (dict): A dictionary containing filters to be applied during the search.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing:
|
||||
- "contexts": List of search results from the base data store.
|
||||
- "entities": List of related graph data based on the query.
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(query, filters)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
|
||||
if not search_output:
|
||||
return []
|
||||
|
||||
search_outputs_sequence = [
|
||||
[item["source"], item["relationship"], item["destination"]] for item in search_output
|
||||
]
|
||||
bm25 = BM25Okapi(search_outputs_sequence)
|
||||
|
||||
tokenized_query = query.split(" ")
|
||||
reranked_results = bm25.get_top_n(tokenized_query, search_outputs_sequence, n=top_k)
|
||||
|
||||
search_results = []
|
||||
for item in reranked_results:
|
||||
search_results.append({"source": item[0], "relationship": item[1], "destination": item[2]})
|
||||
|
||||
logger.info(f"Returned {len(search_results)} search results")
|
||||
|
||||
return search_results
|
||||
|
||||
def delete(self, data, filters):
|
||||
"""
|
||||
Delete graph entities associated with the given memory text.
|
||||
|
||||
Extracts entities and relationships from the memory text using the same
|
||||
pipeline as add(), then deletes the matching relationships in the graph.
|
||||
|
||||
Args:
|
||||
data (str): The memory text whose graph entities should be removed.
|
||||
filters (dict): Scope filters (user_id, agent_id, run_id).
|
||||
"""
|
||||
try:
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
if not entity_type_map:
|
||||
logger.debug("No entities found in memory text, skipping graph cleanup")
|
||||
return
|
||||
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
if to_be_deleted:
|
||||
self._delete_entities(to_be_deleted, filters)
|
||||
except Exception as e:
|
||||
logger.error(f"Error during graph cleanup for memory delete: {e}")
|
||||
|
||||
def delete_all(self, filters):
|
||||
# Build node properties for filtering
|
||||
node_props = ["user_id: $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
node_props.append("agent_id: $agent_id")
|
||||
if filters.get("run_id"):
|
||||
node_props.append("run_id: $run_id")
|
||||
node_props_str = ", ".join(node_props)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{{node_props_str}}})
|
||||
DETACH DELETE n
|
||||
"""
|
||||
params = {"user_id": filters["user_id"]}
|
||||
if filters.get("agent_id"):
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
params["run_id"] = filters["run_id"]
|
||||
self.kuzu_execute(cypher, parameters=params)
|
||||
|
||||
def get_all(self, filters, top_k=100):
|
||||
"""
|
||||
Retrieves all nodes and relationships from the graph database based on optional filtering criteria.
|
||||
Args:
|
||||
filters (dict): A dictionary containing filters to be applied during the retrieval.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
Returns:
|
||||
list: A list of dictionaries, each containing:
|
||||
- 'contexts': The base data store response for each memory.
|
||||
- 'entities': A list of strings representing the nodes and relationships
|
||||
"""
|
||||
|
||||
params = {
|
||||
"user_id": filters["user_id"],
|
||||
"limit": top_k,
|
||||
}
|
||||
# Build node properties based on filters
|
||||
node_props = ["user_id: $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
node_props.append("agent_id: $agent_id")
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
node_props.append("run_id: $run_id")
|
||||
params["run_id"] = filters["run_id"]
|
||||
node_props_str = ", ".join(node_props)
|
||||
|
||||
query = f"""
|
||||
MATCH (n {self.node_label} {{{node_props_str}}})-[r]->(m {self.node_label} {{{node_props_str}}})
|
||||
RETURN
|
||||
n.name AS source,
|
||||
r.name AS relationship,
|
||||
m.name AS target
|
||||
LIMIT $limit
|
||||
"""
|
||||
results = self.kuzu_execute(query, parameters=params)
|
||||
|
||||
final_results = []
|
||||
for result in results:
|
||||
final_results.append(
|
||||
{
|
||||
"source": result["source"],
|
||||
"relationship": result["relationship"],
|
||||
"target": result["target"],
|
||||
}
|
||||
)
|
||||
|
||||
logger.info(f"Retrieved {len(final_results)} relationships")
|
||||
|
||||
return final_results
|
||||
|
||||
def _retrieve_nodes_from_data(self, data, filters):
|
||||
"""Extracts all the entities mentioned in the query."""
|
||||
_tools = [EXTRACT_ENTITIES_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [EXTRACT_ENTITIES_STRUCT_TOOL]
|
||||
search_results = self.llm.generate_response(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": f"You are a smart assistant who understands entities and their types in a given text. If user message contains self reference such as 'I', 'me', 'my' etc. then use {filters['user_id']} as the source entity. Extract all the entities from the text. ***DO NOT*** answer the question itself if the given text is a question.",
|
||||
},
|
||||
{"role": "user", "content": data},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entity_type_map = {}
|
||||
|
||||
try:
|
||||
for tool_call in search_results["tool_calls"]:
|
||||
if tool_call["name"] != "extract_entities":
|
||||
continue
|
||||
for item in tool_call.get("arguments", {}).get("entities", []):
|
||||
entity_type_map[item["entity"]] = item["entity_type"]
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Error in search tool: {e}, llm_provider={self.llm_provider}, search_results={search_results}"
|
||||
)
|
||||
|
||||
entity_type_map = {k.lower().replace(" ", "_"): v.lower().replace(" ", "_") for k, v in entity_type_map.items()}
|
||||
logger.debug(f"Entity type map: {entity_type_map}\n search_results={search_results}")
|
||||
return entity_type_map
|
||||
|
||||
def _establish_nodes_relations_from_data(self, data, filters, entity_type_map):
|
||||
"""Establish relations among the extracted nodes."""
|
||||
|
||||
# Compose user identification string for prompt
|
||||
user_identity = f"user_id: {filters['user_id']}"
|
||||
if filters.get("agent_id"):
|
||||
user_identity += f", agent_id: {filters['agent_id']}"
|
||||
if filters.get("run_id"):
|
||||
user_identity += f", run_id: {filters['run_id']}"
|
||||
|
||||
if self.config.graph_store.custom_prompt:
|
||||
system_content = EXTRACT_RELATIONS_PROMPT.replace("USER_ID", user_identity)
|
||||
# Add the custom prompt line if configured
|
||||
system_content = system_content.replace("CUSTOM_PROMPT", f"4. {self.config.graph_store.custom_prompt}")
|
||||
messages = [
|
||||
{"role": "system", "content": system_content},
|
||||
{"role": "user", "content": data},
|
||||
]
|
||||
else:
|
||||
system_content = EXTRACT_RELATIONS_PROMPT.replace("USER_ID", user_identity)
|
||||
messages = [
|
||||
{"role": "system", "content": system_content},
|
||||
{"role": "user", "content": f"List of entities: {list(entity_type_map.keys())}. \n\nText: {data}"},
|
||||
]
|
||||
|
||||
_tools = [RELATIONS_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [RELATIONS_STRUCT_TOOL]
|
||||
|
||||
extracted_entities = self.llm.generate_response(
|
||||
messages=messages,
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entities = []
|
||||
if extracted_entities.get("tool_calls"):
|
||||
entities = extracted_entities["tool_calls"][0].get("arguments", {}).get("entities", [])
|
||||
|
||||
entities = self._remove_spaces_from_entities(entities)
|
||||
logger.debug(f"Extracted entities: {entities}")
|
||||
return entities
|
||||
|
||||
def _search_graph_db(self, node_list, filters, top_k=100, threshold=None):
|
||||
"""Search similar nodes among and their respective incoming and outgoing relations."""
|
||||
result_relations = []
|
||||
|
||||
params = {
|
||||
"threshold": threshold if threshold else self.threshold,
|
||||
"user_id": filters["user_id"],
|
||||
"limit": top_k,
|
||||
}
|
||||
# Build node properties for filtering
|
||||
node_props = ["user_id: $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
node_props.append("agent_id: $agent_id")
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
node_props.append("run_id: $run_id")
|
||||
params["run_id"] = filters["run_id"]
|
||||
node_props_str = ", ".join(node_props)
|
||||
|
||||
for node in node_list:
|
||||
n_embedding = self.embedding_model.embed(node)
|
||||
params["n_embedding"] = n_embedding
|
||||
|
||||
results = []
|
||||
for match_fragment in [
|
||||
f"(n)-[r]->(m {self.node_label} {{{node_props_str}}}) WITH n as src, r, m as dst, similarity",
|
||||
f"(m {self.node_label} {{{node_props_str}}})-[r]->(n) WITH m as src, r, n as dst, similarity"
|
||||
]:
|
||||
results.extend(self.kuzu_execute(
|
||||
f"""
|
||||
MATCH (n {self.node_label} {{{node_props_str}}})
|
||||
WHERE n.embedding IS NOT NULL
|
||||
WITH n, array_cosine_similarity(n.embedding, CAST($n_embedding,'FLOAT[{self.embedding_dims}]')) AS similarity
|
||||
WHERE similarity >= CAST($threshold, 'DOUBLE')
|
||||
MATCH {match_fragment}
|
||||
RETURN
|
||||
src.name AS source,
|
||||
id(src) AS source_id,
|
||||
r.name AS relationship,
|
||||
id(r) AS relation_id,
|
||||
dst.name AS destination,
|
||||
id(dst) AS destination_id,
|
||||
similarity
|
||||
LIMIT $limit
|
||||
""",
|
||||
parameters=params))
|
||||
|
||||
# Kuzu does not support sort/limit over unions. Do it manually for now.
|
||||
result_relations.extend(sorted(results, key=lambda x: x["similarity"], reverse=True)[:top_k])
|
||||
|
||||
return result_relations
|
||||
|
||||
def _get_delete_entities_from_search_output(self, search_output, data, filters):
|
||||
"""Get the entities to be deleted from the search output."""
|
||||
search_output_string = format_entities(search_output)
|
||||
|
||||
# Compose user identification string for prompt
|
||||
user_identity = f"user_id: {filters['user_id']}"
|
||||
if filters.get("agent_id"):
|
||||
user_identity += f", agent_id: {filters['agent_id']}"
|
||||
if filters.get("run_id"):
|
||||
user_identity += f", run_id: {filters['run_id']}"
|
||||
|
||||
system_prompt, user_prompt = get_delete_messages(search_output_string, data, user_identity)
|
||||
|
||||
_tools = [DELETE_MEMORY_TOOL_GRAPH]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
]
|
||||
|
||||
memory_updates = self.llm.generate_response(
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
to_be_deleted = []
|
||||
for item in memory_updates.get("tool_calls", []):
|
||||
if item.get("name") == "delete_graph_memory":
|
||||
to_be_deleted.append(item.get("arguments"))
|
||||
# Clean entities formatting
|
||||
to_be_deleted = self._remove_spaces_from_entities(to_be_deleted)
|
||||
logger.debug(f"Deleted relationships: {to_be_deleted}")
|
||||
return to_be_deleted
|
||||
|
||||
def _delete_entities(self, to_be_deleted, filters):
|
||||
"""Delete the entities from the graph."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id", None)
|
||||
run_id = filters.get("run_id", None)
|
||||
results = []
|
||||
|
||||
for item in to_be_deleted:
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"user_id": user_id,
|
||||
"relationship_name": relationship,
|
||||
}
|
||||
# Build node properties for filtering
|
||||
source_props = ["name: $source_name", "user_id: $user_id"]
|
||||
dest_props = ["name: $dest_name", "user_id: $user_id"]
|
||||
if agent_id:
|
||||
source_props.append("agent_id: $agent_id")
|
||||
dest_props.append("agent_id: $agent_id")
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
source_props.append("run_id: $run_id")
|
||||
dest_props.append("run_id: $run_id")
|
||||
params["run_id"] = run_id
|
||||
source_props_str = ", ".join(source_props)
|
||||
dest_props_str = ", ".join(dest_props)
|
||||
|
||||
# Delete the specific relationship between nodes
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{{source_props_str}}})
|
||||
-[r {self.rel_label} {{name: $relationship_name}}]->
|
||||
(m {self.node_label} {{{dest_props_str}}})
|
||||
DELETE r
|
||||
RETURN
|
||||
n.name AS source,
|
||||
r.name AS relationship,
|
||||
m.name AS target
|
||||
"""
|
||||
|
||||
result = self.kuzu_execute(cypher, parameters=params)
|
||||
results.append(result)
|
||||
|
||||
return results
|
||||
|
||||
def _add_entities(self, to_be_added, filters, entity_type_map):
|
||||
"""Add the new entities to the graph. Merge the nodes if they already exist."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id", None)
|
||||
run_id = filters.get("run_id", None)
|
||||
results = []
|
||||
for item in to_be_added:
|
||||
# entities
|
||||
source = item["source"]
|
||||
source_label = self.node_label
|
||||
|
||||
destination = item["destination"]
|
||||
destination_label = self.node_label
|
||||
|
||||
relationship = item["relationship"]
|
||||
relationship_label = self.rel_label
|
||||
|
||||
# embeddings
|
||||
source_embedding = self.embedding_model.embed(source)
|
||||
dest_embedding = self.embedding_model.embed(destination)
|
||||
|
||||
# search for the nodes with the closest embeddings
|
||||
source_node_search_result = self._search_source_node(source_embedding, filters, threshold=self.threshold)
|
||||
destination_node_search_result = self._search_destination_node(dest_embedding, filters, threshold=self.threshold)
|
||||
|
||||
if not destination_node_search_result and source_node_search_result:
|
||||
params = {
|
||||
"table_id": source_node_search_result[0]["id"]["table"],
|
||||
"offset_id": source_node_search_result[0]["id"]["offset"],
|
||||
"destination_name": destination,
|
||||
"destination_embedding": dest_embedding,
|
||||
"relationship_name": relationship,
|
||||
"user_id": user_id,
|
||||
}
|
||||
# Build source MERGE properties
|
||||
merge_props = ["name: $destination_name", "user_id: $user_id"]
|
||||
if agent_id:
|
||||
merge_props.append("agent_id: $agent_id")
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
merge_props.append("run_id: $run_id")
|
||||
params["run_id"] = run_id
|
||||
merge_props_str = ", ".join(merge_props)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source)
|
||||
WHERE id(source) = internal_id($table_id, $offset_id)
|
||||
SET source.mentions = coalesce(source.mentions, 0) + 1
|
||||
WITH source
|
||||
MERGE (destination {destination_label} {{{merge_props_str}}})
|
||||
ON CREATE SET
|
||||
destination.created = current_timestamp(),
|
||||
destination.mentions = 1,
|
||||
destination.embedding = CAST($destination_embedding,'FLOAT[{self.embedding_dims}]')
|
||||
ON MATCH SET
|
||||
destination.mentions = coalesce(destination.mentions, 0) + 1,
|
||||
destination.embedding = CAST($destination_embedding,'FLOAT[{self.embedding_dims}]')
|
||||
WITH source, destination
|
||||
MERGE (source)-[r {relationship_label} {{name: $relationship_name}}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = current_timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1
|
||||
RETURN
|
||||
source.name AS source,
|
||||
r.name AS relationship,
|
||||
destination.name AS target
|
||||
"""
|
||||
elif destination_node_search_result and not source_node_search_result:
|
||||
params = {
|
||||
"table_id": destination_node_search_result[0]["id"]["table"],
|
||||
"offset_id": destination_node_search_result[0]["id"]["offset"],
|
||||
"source_name": source,
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
"relationship_name": relationship,
|
||||
}
|
||||
# Build source MERGE properties
|
||||
merge_props = ["name: $source_name", "user_id: $user_id"]
|
||||
if agent_id:
|
||||
merge_props.append("agent_id: $agent_id")
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
merge_props.append("run_id: $run_id")
|
||||
params["run_id"] = run_id
|
||||
merge_props_str = ", ".join(merge_props)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (destination)
|
||||
WHERE id(destination) = internal_id($table_id, $offset_id)
|
||||
SET destination.mentions = coalesce(destination.mentions, 0) + 1
|
||||
WITH destination
|
||||
MERGE (source {source_label} {{{merge_props_str}}})
|
||||
ON CREATE SET
|
||||
source.created = current_timestamp(),
|
||||
source.mentions = 1,
|
||||
source.embedding = CAST($source_embedding,'FLOAT[{self.embedding_dims}]')
|
||||
ON MATCH SET
|
||||
source.mentions = coalesce(source.mentions, 0) + 1,
|
||||
source.embedding = CAST($source_embedding,'FLOAT[{self.embedding_dims}]')
|
||||
WITH source, destination
|
||||
MERGE (source)-[r {relationship_label} {{name: $relationship_name}}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = current_timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1
|
||||
RETURN
|
||||
source.name AS source,
|
||||
r.name AS relationship,
|
||||
destination.name AS target
|
||||
"""
|
||||
elif source_node_search_result and destination_node_search_result:
|
||||
cypher = f"""
|
||||
MATCH (source)
|
||||
WHERE id(source) = internal_id($src_table, $src_offset)
|
||||
SET source.mentions = coalesce(source.mentions, 0) + 1
|
||||
WITH source
|
||||
MATCH (destination)
|
||||
WHERE id(destination) = internal_id($dst_table, $dst_offset)
|
||||
SET destination.mentions = coalesce(destination.mentions, 0) + 1
|
||||
MERGE (source)-[r {relationship_label} {{name: $relationship_name}}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = current_timestamp(),
|
||||
r.updated = current_timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET r.mentions = coalesce(r.mentions, 0) + 1
|
||||
RETURN
|
||||
source.name AS source,
|
||||
r.name AS relationship,
|
||||
destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"src_table": source_node_search_result[0]["id"]["table"],
|
||||
"src_offset": source_node_search_result[0]["id"]["offset"],
|
||||
"dst_table": destination_node_search_result[0]["id"]["table"],
|
||||
"dst_offset": destination_node_search_result[0]["id"]["offset"],
|
||||
"relationship_name": relationship,
|
||||
}
|
||||
else:
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"relationship_name": relationship,
|
||||
"source_embedding": source_embedding,
|
||||
"dest_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
# Build dynamic MERGE props for both source and destination
|
||||
source_props = ["name: $source_name", "user_id: $user_id"]
|
||||
dest_props = ["name: $dest_name", "user_id: $user_id"]
|
||||
if agent_id:
|
||||
source_props.append("agent_id: $agent_id")
|
||||
dest_props.append("agent_id: $agent_id")
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
source_props.append("run_id: $run_id")
|
||||
dest_props.append("run_id: $run_id")
|
||||
params["run_id"] = run_id
|
||||
source_props_str = ", ".join(source_props)
|
||||
dest_props_str = ", ".join(dest_props)
|
||||
|
||||
cypher = f"""
|
||||
MERGE (source {source_label} {{{source_props_str}}})
|
||||
ON CREATE SET
|
||||
source.created = current_timestamp(),
|
||||
source.mentions = 1,
|
||||
source.embedding = CAST($source_embedding,'FLOAT[{self.embedding_dims}]')
|
||||
ON MATCH SET
|
||||
source.mentions = coalesce(source.mentions, 0) + 1,
|
||||
source.embedding = CAST($source_embedding,'FLOAT[{self.embedding_dims}]')
|
||||
WITH source
|
||||
MERGE (destination {destination_label} {{{dest_props_str}}})
|
||||
ON CREATE SET
|
||||
destination.created = current_timestamp(),
|
||||
destination.mentions = 1,
|
||||
destination.embedding = CAST($dest_embedding,'FLOAT[{self.embedding_dims}]')
|
||||
ON MATCH SET
|
||||
destination.mentions = coalesce(destination.mentions, 0) + 1,
|
||||
destination.embedding = CAST($dest_embedding,'FLOAT[{self.embedding_dims}]')
|
||||
WITH source, destination
|
||||
MERGE (source)-[rel {relationship_label} {{name: $relationship_name}}]->(destination)
|
||||
ON CREATE SET
|
||||
rel.created = current_timestamp(),
|
||||
rel.mentions = 1
|
||||
ON MATCH SET
|
||||
rel.mentions = coalesce(rel.mentions, 0) + 1
|
||||
RETURN
|
||||
source.name AS source,
|
||||
rel.name AS relationship,
|
||||
destination.name AS target
|
||||
"""
|
||||
|
||||
result = self.kuzu_execute(cypher, parameters=params)
|
||||
results.append(result)
|
||||
|
||||
return results
|
||||
|
||||
def _remove_spaces_from_entities(self, entity_list):
|
||||
return remove_spaces_from_entities(entity_list, sanitize_relationship=False)
|
||||
|
||||
def _search_source_node(self, source_embedding, filters, threshold=0.9):
|
||||
params = {
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": filters["user_id"],
|
||||
"threshold": threshold,
|
||||
}
|
||||
where_conditions = ["source_candidate.embedding IS NOT NULL", "source_candidate.user_id = $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
where_conditions.append("source_candidate.agent_id = $agent_id")
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
where_conditions.append("source_candidate.run_id = $run_id")
|
||||
params["run_id"] = filters["run_id"]
|
||||
where_clause = " AND ".join(where_conditions)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source_candidate {self.node_label})
|
||||
WHERE {where_clause}
|
||||
|
||||
WITH source_candidate,
|
||||
array_cosine_similarity(source_candidate.embedding, CAST($source_embedding,'FLOAT[{self.embedding_dims}]')) AS source_similarity
|
||||
|
||||
WHERE source_similarity >= $threshold
|
||||
|
||||
WITH source_candidate, source_similarity
|
||||
ORDER BY source_similarity DESC
|
||||
LIMIT 2
|
||||
|
||||
RETURN id(source_candidate) as id, source_similarity
|
||||
"""
|
||||
|
||||
return self.kuzu_execute(cypher, parameters=params)
|
||||
|
||||
def _search_destination_node(self, destination_embedding, filters, threshold=0.9):
|
||||
params = {
|
||||
"destination_embedding": destination_embedding,
|
||||
"user_id": filters["user_id"],
|
||||
"threshold": threshold,
|
||||
}
|
||||
where_conditions = ["destination_candidate.embedding IS NOT NULL", "destination_candidate.user_id = $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
where_conditions.append("destination_candidate.agent_id = $agent_id")
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
where_conditions.append("destination_candidate.run_id = $run_id")
|
||||
params["run_id"] = filters["run_id"]
|
||||
where_clause = " AND ".join(where_conditions)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (destination_candidate {self.node_label})
|
||||
WHERE {where_clause}
|
||||
|
||||
WITH destination_candidate,
|
||||
array_cosine_similarity(destination_candidate.embedding, CAST($destination_embedding,'FLOAT[{self.embedding_dims}]')) AS destination_similarity
|
||||
|
||||
WHERE destination_similarity >= $threshold
|
||||
|
||||
WITH destination_candidate, destination_similarity
|
||||
ORDER BY destination_similarity DESC
|
||||
LIMIT 2
|
||||
|
||||
RETURN id(destination_candidate) as id, destination_similarity
|
||||
"""
|
||||
|
||||
return self.kuzu_execute(cypher, parameters=params)
|
||||
|
||||
# Reset is not defined in base.py
|
||||
def reset(self):
|
||||
"""Reset the graph by clearing all nodes and relationships."""
|
||||
logger.warning("Clearing graph...")
|
||||
cypher_query = """
|
||||
MATCH (n) DETACH DELETE n
|
||||
"""
|
||||
return self.kuzu_execute(cypher_query)
|
||||
+1090
-832
File diff suppressed because it is too large
Load Diff
@@ -1,708 +0,0 @@
|
||||
import logging
|
||||
|
||||
from mem0.memory.utils import format_entities, remove_spaces_from_entities
|
||||
|
||||
try:
|
||||
from langchain_memgraph.graphs.memgraph import Memgraph
|
||||
except ImportError:
|
||||
raise ImportError("langchain_memgraph is not installed. Please install it using pip install langchain-memgraph")
|
||||
|
||||
try:
|
||||
from rank_bm25 import BM25Okapi
|
||||
except ImportError:
|
||||
raise ImportError("rank_bm25 is not installed. Please install it using pip install rank-bm25")
|
||||
|
||||
from mem0.graphs.tools import (
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
DELETE_MEMORY_TOOL_GRAPH,
|
||||
EXTRACT_ENTITIES_STRUCT_TOOL,
|
||||
EXTRACT_ENTITIES_TOOL,
|
||||
RELATIONS_STRUCT_TOOL,
|
||||
RELATIONS_TOOL,
|
||||
)
|
||||
from mem0.graphs.utils import EXTRACT_RELATIONS_PROMPT, get_delete_messages
|
||||
from mem0.utils.factory import EmbedderFactory, LlmFactory
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MemoryGraph:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
self.graph = Memgraph(
|
||||
self.config.graph_store.config.url,
|
||||
self.config.graph_store.config.username,
|
||||
self.config.graph_store.config.password,
|
||||
)
|
||||
self.embedding_model = EmbedderFactory.create(
|
||||
self.config.embedder.provider,
|
||||
self.config.embedder.config,
|
||||
{"enable_embeddings": True},
|
||||
)
|
||||
|
||||
# Default to openai if no specific provider is configured
|
||||
self.llm_provider = "openai"
|
||||
if self.config.llm and self.config.llm.provider:
|
||||
self.llm_provider = self.config.llm.provider
|
||||
if self.config.graph_store and self.config.graph_store.llm and self.config.graph_store.llm.provider:
|
||||
self.llm_provider = self.config.graph_store.llm.provider
|
||||
|
||||
# Get LLM config with proper null checks
|
||||
llm_config = None
|
||||
if self.config.graph_store and self.config.graph_store.llm and hasattr(self.config.graph_store.llm, "config"):
|
||||
llm_config = self.config.graph_store.llm.config
|
||||
elif hasattr(self.config.llm, "config"):
|
||||
llm_config = self.config.llm.config
|
||||
self.llm = LlmFactory.create(self.llm_provider, llm_config)
|
||||
self.user_id = None
|
||||
# Use threshold from graph_store config, default to 0.7 for backward compatibility
|
||||
self.threshold = self.config.graph_store.threshold if hasattr(self.config.graph_store, 'threshold') else 0.7
|
||||
|
||||
# Setup Memgraph:
|
||||
# 1. Create vector index (created Entity label on all nodes)
|
||||
# 2. Create label property index for performance optimizations
|
||||
embedding_dims = self.config.embedder.config["embedding_dims"]
|
||||
index_info = self._fetch_existing_indexes()
|
||||
|
||||
# Create vector index if not exists
|
||||
if not self._vector_index_exists(index_info, "memzero"):
|
||||
self.graph.query(
|
||||
f"CREATE VECTOR INDEX memzero ON :Entity(embedding) WITH CONFIG {{'dimension': {embedding_dims}, 'capacity': 1000, 'metric': 'cos'}};"
|
||||
)
|
||||
|
||||
# Create label+property index if not exists
|
||||
if not self._label_property_index_exists(index_info, "Entity", "user_id"):
|
||||
self.graph.query("CREATE INDEX ON :Entity(user_id);")
|
||||
|
||||
# Create label index if not exists
|
||||
if not self._label_index_exists(index_info, "Entity"):
|
||||
self.graph.query("CREATE INDEX ON :Entity;")
|
||||
|
||||
def add(self, data, filters):
|
||||
"""
|
||||
Adds data to the graph.
|
||||
|
||||
Args:
|
||||
data (str): The data to add to the graph.
|
||||
filters (dict): A dictionary containing filters to be applied during the addition.
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
to_be_added = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
to_be_deleted = self._get_delete_entities_from_search_output(search_output, data, filters)
|
||||
|
||||
# TODO: Batch queries with APOC plugin
|
||||
# TODO: Add more filter support
|
||||
deleted_entities = self._delete_entities(to_be_deleted, filters)
|
||||
added_entities = self._add_entities(to_be_added, filters, entity_type_map)
|
||||
|
||||
return {"deleted_entities": deleted_entities, "added_entities": added_entities}
|
||||
|
||||
def search(self, query, filters, top_k=100):
|
||||
"""
|
||||
Search for memories and related graph data.
|
||||
|
||||
Args:
|
||||
query (str): Query to search for.
|
||||
filters (dict): A dictionary containing filters to be applied during the search.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing:
|
||||
- "contexts": List of search results from the base data store.
|
||||
- "entities": List of related graph data based on the query.
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(query, filters)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
|
||||
if not search_output:
|
||||
return []
|
||||
|
||||
search_outputs_sequence = [
|
||||
[item["source"], item["relationship"], item["destination"]] for item in search_output
|
||||
]
|
||||
bm25 = BM25Okapi(search_outputs_sequence)
|
||||
|
||||
tokenized_query = query.split(" ")
|
||||
reranked_results = bm25.get_top_n(tokenized_query, search_outputs_sequence, n=5)
|
||||
|
||||
search_results = []
|
||||
for item in reranked_results:
|
||||
search_results.append({"source": item[0], "relationship": item[1], "destination": item[2]})
|
||||
|
||||
logger.info(f"Returned {len(search_results)} search results")
|
||||
|
||||
return search_results
|
||||
|
||||
def delete(self, data, filters):
|
||||
"""
|
||||
Delete graph entities associated with the given memory text.
|
||||
|
||||
Extracts entities and relationships from the memory text using the same
|
||||
pipeline as add(), then deletes the matching relationships in the graph.
|
||||
|
||||
Args:
|
||||
data (str): The memory text whose graph entities should be removed.
|
||||
filters (dict): Scope filters (user_id, agent_id).
|
||||
"""
|
||||
try:
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
if not entity_type_map:
|
||||
logger.debug("No entities found in memory text, skipping graph cleanup")
|
||||
return
|
||||
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
if to_be_deleted:
|
||||
self._delete_entities(to_be_deleted, filters)
|
||||
except Exception as e:
|
||||
logger.error(f"Error during graph cleanup for memory delete: {e}")
|
||||
|
||||
def delete_all(self, filters):
|
||||
"""Delete all nodes and relationships for a user or specific agent."""
|
||||
if filters.get("agent_id"):
|
||||
cypher = """
|
||||
MATCH (n:Entity {user_id: $user_id, agent_id: $agent_id})
|
||||
DETACH DELETE n
|
||||
"""
|
||||
params = {"user_id": filters["user_id"], "agent_id": filters["agent_id"]}
|
||||
else:
|
||||
cypher = """
|
||||
MATCH (n:Entity {user_id: $user_id})
|
||||
DETACH DELETE n
|
||||
"""
|
||||
params = {"user_id": filters["user_id"]}
|
||||
self.graph.query(cypher, params=params)
|
||||
|
||||
def get_all(self, filters, top_k=100):
|
||||
"""
|
||||
Retrieves all nodes and relationships from the graph database based on optional filtering criteria.
|
||||
|
||||
Args:
|
||||
filters (dict): A dictionary containing filters to be applied during the retrieval.
|
||||
Supports 'user_id' (required) and 'agent_id' (optional).
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
Returns:
|
||||
list: A list of dictionaries, each containing:
|
||||
- 'source': The source node name.
|
||||
- 'relationship': The relationship type.
|
||||
- 'target': The target node name.
|
||||
"""
|
||||
# Build query based on whether agent_id is provided
|
||||
if filters.get("agent_id"):
|
||||
query = """
|
||||
MATCH (n:Entity {user_id: $user_id, agent_id: $agent_id})-[r]->(m:Entity {user_id: $user_id, agent_id: $agent_id})
|
||||
RETURN n.name AS source, type(r) AS relationship, m.name AS target
|
||||
LIMIT $limit
|
||||
"""
|
||||
params = {"user_id": filters["user_id"], "agent_id": filters["agent_id"], "limit": top_k}
|
||||
else:
|
||||
query = """
|
||||
MATCH (n:Entity {user_id: $user_id})-[r]->(m:Entity {user_id: $user_id})
|
||||
RETURN n.name AS source, type(r) AS relationship, m.name AS target
|
||||
LIMIT $limit
|
||||
"""
|
||||
params = {"user_id": filters["user_id"], "limit": top_k}
|
||||
|
||||
results = self.graph.query(query, params=params)
|
||||
|
||||
final_results = []
|
||||
for result in results:
|
||||
final_results.append(
|
||||
{
|
||||
"source": result["source"],
|
||||
"relationship": result["relationship"],
|
||||
"target": result["target"],
|
||||
}
|
||||
)
|
||||
|
||||
logger.info(f"Retrieved {len(final_results)} relationships")
|
||||
|
||||
return final_results
|
||||
|
||||
def _retrieve_nodes_from_data(self, data, filters):
|
||||
"""Extracts all the entities mentioned in the query."""
|
||||
_tools = [EXTRACT_ENTITIES_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [EXTRACT_ENTITIES_STRUCT_TOOL]
|
||||
search_results = self.llm.generate_response(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": f"You are a smart assistant who understands entities and their types in a given text. If user message contains self reference such as 'I', 'me', 'my' etc. then use {filters['user_id']} as the source entity. Extract all the entities from the text. ***DO NOT*** answer the question itself if the given text is a question.",
|
||||
},
|
||||
{"role": "user", "content": data},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entity_type_map = {}
|
||||
|
||||
try:
|
||||
for tool_call in search_results["tool_calls"]:
|
||||
if tool_call["name"] != "extract_entities":
|
||||
continue
|
||||
for item in tool_call.get("arguments", {}).get("entities", []):
|
||||
if "entity" in item and "entity_type" in item:
|
||||
entity_type_map[item["entity"]] = item["entity_type"]
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Error in search tool: {e}, llm_provider={self.llm_provider}, search_results={search_results}"
|
||||
)
|
||||
|
||||
entity_type_map = {k.lower().replace(" ", "_"): v.lower().replace(" ", "_") for k, v in entity_type_map.items()}
|
||||
logger.debug(f"Entity type map: {entity_type_map}\n search_results={search_results}")
|
||||
return entity_type_map
|
||||
|
||||
def _establish_nodes_relations_from_data(self, data, filters, entity_type_map):
|
||||
"""Eshtablish relations among the extracted nodes."""
|
||||
if self.config.graph_store.custom_prompt:
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": EXTRACT_RELATIONS_PROMPT.replace("USER_ID", filters["user_id"]).replace(
|
||||
"CUSTOM_PROMPT", f"4. {self.config.graph_store.custom_prompt}"
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": data},
|
||||
]
|
||||
else:
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": EXTRACT_RELATIONS_PROMPT.replace("USER_ID", filters["user_id"]),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"List of entities: {list(entity_type_map.keys())}. \n\nText: {data}",
|
||||
},
|
||||
]
|
||||
|
||||
_tools = [RELATIONS_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [RELATIONS_STRUCT_TOOL]
|
||||
|
||||
extracted_entities = self.llm.generate_response(
|
||||
messages=messages,
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entities = []
|
||||
if extracted_entities and extracted_entities.get("tool_calls"):
|
||||
entities = extracted_entities["tool_calls"][0].get("arguments", {}).get("entities", [])
|
||||
|
||||
entities = self._remove_spaces_from_entities(entities)
|
||||
logger.debug(f"Extracted entities: {entities}")
|
||||
return entities
|
||||
|
||||
def _search_graph_db(self, node_list, filters, top_k=100):
|
||||
"""Search similar nodes among and their respective incoming and outgoing relations."""
|
||||
result_relations = []
|
||||
|
||||
for node in node_list:
|
||||
n_embedding = self.embedding_model.embed(node)
|
||||
|
||||
# Build query based on whether agent_id is provided
|
||||
if filters.get("agent_id"):
|
||||
cypher_query = """
|
||||
CALL vector_search.search("memzero", $limit, $n_embedding)
|
||||
YIELD distance, node, similarity
|
||||
WITH node AS n, similarity
|
||||
WHERE n:Entity AND n.user_id = $user_id AND n.agent_id = $agent_id AND n.embedding IS NOT NULL AND similarity >= $threshold
|
||||
MATCH (n)-[r]->(m:Entity)
|
||||
RETURN n.name AS source, id(n) AS source_id, type(r) AS relationship, id(r) AS relation_id, m.name AS destination, id(m) AS destination_id, similarity
|
||||
UNION
|
||||
CALL vector_search.search("memzero", $limit, $n_embedding)
|
||||
YIELD distance, node, similarity
|
||||
WITH node AS n, similarity
|
||||
WHERE n:Entity AND n.user_id = $user_id AND n.agent_id = $agent_id AND n.embedding IS NOT NULL AND similarity >= $threshold
|
||||
MATCH (m:Entity)-[r]->(n)
|
||||
RETURN m.name AS source, id(m) AS source_id, type(r) AS relationship, id(r) AS relation_id, n.name AS destination, id(n) AS destination_id, similarity
|
||||
ORDER BY similarity DESC
|
||||
LIMIT $limit;
|
||||
"""
|
||||
params = {
|
||||
"n_embedding": n_embedding,
|
||||
"threshold": self.threshold,
|
||||
"user_id": filters["user_id"],
|
||||
"agent_id": filters["agent_id"],
|
||||
"limit": top_k,
|
||||
}
|
||||
else:
|
||||
cypher_query = """
|
||||
CALL vector_search.search("memzero", $limit, $n_embedding)
|
||||
YIELD distance, node, similarity
|
||||
WITH node AS n, similarity
|
||||
WHERE n:Entity AND n.user_id = $user_id AND n.embedding IS NOT NULL AND similarity >= $threshold
|
||||
MATCH (n)-[r]->(m:Entity)
|
||||
RETURN n.name AS source, id(n) AS source_id, type(r) AS relationship, id(r) AS relation_id, m.name AS destination, id(m) AS destination_id, similarity
|
||||
UNION
|
||||
CALL vector_search.search("memzero", $limit, $n_embedding)
|
||||
YIELD distance, node, similarity
|
||||
WITH node AS n, similarity
|
||||
WHERE n:Entity AND n.user_id = $user_id AND n.embedding IS NOT NULL AND similarity >= $threshold
|
||||
MATCH (m:Entity)-[r]->(n)
|
||||
RETURN m.name AS source, id(m) AS source_id, type(r) AS relationship, id(r) AS relation_id, n.name AS destination, id(n) AS destination_id, similarity
|
||||
ORDER BY similarity DESC
|
||||
LIMIT $limit;
|
||||
"""
|
||||
params = {
|
||||
"n_embedding": n_embedding,
|
||||
"threshold": self.threshold,
|
||||
"user_id": filters["user_id"],
|
||||
"limit": top_k,
|
||||
}
|
||||
|
||||
ans = self.graph.query(cypher_query, params=params)
|
||||
result_relations.extend(ans)
|
||||
|
||||
return result_relations
|
||||
|
||||
def _get_delete_entities_from_search_output(self, search_output, data, filters):
|
||||
"""Get the entities to be deleted from the search output."""
|
||||
search_output_string = format_entities(search_output)
|
||||
system_prompt, user_prompt = get_delete_messages(search_output_string, data, filters["user_id"])
|
||||
|
||||
_tools = [DELETE_MEMORY_TOOL_GRAPH]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
]
|
||||
|
||||
memory_updates = self.llm.generate_response(
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
to_be_deleted = []
|
||||
for item in memory_updates["tool_calls"]:
|
||||
if item["name"] == "delete_graph_memory":
|
||||
to_be_deleted.append(item["arguments"])
|
||||
# in case if it is not in the correct format
|
||||
to_be_deleted = self._remove_spaces_from_entities(to_be_deleted)
|
||||
logger.debug(f"Deleted relationships: {to_be_deleted}")
|
||||
return to_be_deleted
|
||||
|
||||
def _delete_entities(self, to_be_deleted, filters):
|
||||
"""Delete the entities from the graph."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id", None)
|
||||
results = []
|
||||
|
||||
for item in to_be_deleted:
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
# Build the agent filter for the query
|
||||
agent_filter = ""
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"user_id": user_id,
|
||||
}
|
||||
|
||||
if agent_id:
|
||||
agent_filter = "AND n.agent_id = $agent_id AND m.agent_id = $agent_id"
|
||||
params["agent_id"] = agent_id
|
||||
|
||||
# Delete the specific relationship between nodes
|
||||
cypher = f"""
|
||||
MATCH (n:Entity {{name: $source_name, user_id: $user_id}})
|
||||
-[r:{relationship}]->
|
||||
(m:Entity {{name: $dest_name, user_id: $user_id}})
|
||||
WHERE 1=1 {agent_filter}
|
||||
DELETE r
|
||||
RETURN
|
||||
n.name AS source,
|
||||
m.name AS target,
|
||||
type(r) AS relationship
|
||||
"""
|
||||
|
||||
result = self.graph.query(cypher, params=params)
|
||||
results.append(result)
|
||||
|
||||
return results
|
||||
|
||||
# added Entity label to all nodes for vector search to work
|
||||
def _add_entities(self, to_be_added, filters, entity_type_map):
|
||||
"""Add the new entities to the graph. Merge the nodes if they already exist."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id", None)
|
||||
results = []
|
||||
|
||||
for item in to_be_added:
|
||||
# entities
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
# types
|
||||
source_type = entity_type_map.get(source, "__User__")
|
||||
destination_type = entity_type_map.get(destination, "__User__")
|
||||
|
||||
# embeddings
|
||||
source_embedding = self.embedding_model.embed(source)
|
||||
dest_embedding = self.embedding_model.embed(destination)
|
||||
|
||||
# search for the nodes with the closest embeddings
|
||||
source_node_search_result = self._search_source_node(source_embedding, filters, threshold=self.threshold)
|
||||
destination_node_search_result = self._search_destination_node(dest_embedding, filters, threshold=self.threshold)
|
||||
|
||||
# Prepare agent_id for node creation
|
||||
agent_id_clause = ""
|
||||
if agent_id:
|
||||
agent_id_clause = ", agent_id: $agent_id"
|
||||
|
||||
# TODO: Create a cypher query and common params for all the cases
|
||||
if not destination_node_search_result and source_node_search_result:
|
||||
cypher = f"""
|
||||
MATCH (source:Entity)
|
||||
WHERE id(source) = $source_id
|
||||
MERGE (destination:{destination_type}:Entity {{name: $destination_name, user_id: $user_id{agent_id_clause}}})
|
||||
ON CREATE SET
|
||||
destination.created = timestamp(),
|
||||
destination.embedding = $destination_embedding,
|
||||
destination:Entity
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = timestamp()
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"source_id": source_node_search_result[0]["id(source_candidate)"],
|
||||
"destination_name": destination,
|
||||
"destination_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
|
||||
elif destination_node_search_result and not source_node_search_result:
|
||||
cypher = f"""
|
||||
MATCH (destination:Entity)
|
||||
WHERE id(destination) = $destination_id
|
||||
MERGE (source:{source_type}:Entity {{name: $source_name, user_id: $user_id{agent_id_clause}}})
|
||||
ON CREATE SET
|
||||
source.created = timestamp(),
|
||||
source.embedding = $source_embedding,
|
||||
source:Entity
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = timestamp()
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"destination_id": destination_node_search_result[0]["id(destination_candidate)"],
|
||||
"source_name": source,
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
|
||||
elif source_node_search_result and destination_node_search_result:
|
||||
cypher = f"""
|
||||
MATCH (source:Entity)
|
||||
WHERE id(source) = $source_id
|
||||
MATCH (destination:Entity)
|
||||
WHERE id(destination) = $destination_id
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created_at = timestamp(),
|
||||
r.updated_at = timestamp()
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
params = {
|
||||
"source_id": source_node_search_result[0]["id(source_candidate)"],
|
||||
"destination_id": destination_node_search_result[0]["id(destination_candidate)"],
|
||||
"user_id": user_id,
|
||||
}
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
|
||||
else:
|
||||
cypher = f"""
|
||||
MERGE (n:{source_type}:Entity {{name: $source_name, user_id: $user_id{agent_id_clause}}})
|
||||
ON CREATE SET n.created = timestamp(), n.embedding = $source_embedding, n:Entity
|
||||
ON MATCH SET n.embedding = $source_embedding
|
||||
MERGE (m:{destination_type}:Entity {{name: $dest_name, user_id: $user_id{agent_id_clause}}})
|
||||
ON CREATE SET m.created = timestamp(), m.embedding = $dest_embedding, m:Entity
|
||||
ON MATCH SET m.embedding = $dest_embedding
|
||||
MERGE (n)-[rel:{relationship}]->(m)
|
||||
ON CREATE SET rel.created = timestamp()
|
||||
RETURN n.name AS source, type(rel) AS relationship, m.name AS target
|
||||
"""
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"source_embedding": source_embedding,
|
||||
"dest_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
|
||||
result = self.graph.query(cypher, params=params)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
def _remove_spaces_from_entities(self, entity_list):
|
||||
return remove_spaces_from_entities(entity_list, sanitize_relationship=True)
|
||||
|
||||
def _search_source_node(self, source_embedding, filters, threshold=0.9):
|
||||
"""Search for source nodes with similar embeddings."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id", None)
|
||||
|
||||
if agent_id:
|
||||
cypher = """
|
||||
CALL vector_search.search("memzero", 1, $source_embedding)
|
||||
YIELD distance, node, similarity
|
||||
WITH node AS source_candidate, similarity
|
||||
WHERE source_candidate.user_id = $user_id
|
||||
AND source_candidate.agent_id = $agent_id
|
||||
AND similarity >= $threshold
|
||||
RETURN id(source_candidate);
|
||||
"""
|
||||
params = {
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
"agent_id": agent_id,
|
||||
"threshold": threshold,
|
||||
}
|
||||
else:
|
||||
cypher = """
|
||||
CALL vector_search.search("memzero", 1, $source_embedding)
|
||||
YIELD distance, node, similarity
|
||||
WITH node AS source_candidate, similarity
|
||||
WHERE source_candidate.user_id = $user_id
|
||||
AND similarity >= $threshold
|
||||
RETURN id(source_candidate);
|
||||
"""
|
||||
params = {
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
"threshold": threshold,
|
||||
}
|
||||
|
||||
result = self.graph.query(cypher, params=params)
|
||||
return result
|
||||
|
||||
def _search_destination_node(self, destination_embedding, filters, threshold=0.9):
|
||||
"""Search for destination nodes with similar embeddings."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id", None)
|
||||
|
||||
if agent_id:
|
||||
cypher = """
|
||||
CALL vector_search.search("memzero", 1, $destination_embedding)
|
||||
YIELD distance, node, similarity
|
||||
WITH node AS destination_candidate, similarity
|
||||
WHERE node.user_id = $user_id
|
||||
AND node.agent_id = $agent_id
|
||||
AND similarity >= $threshold
|
||||
RETURN id(destination_candidate);
|
||||
"""
|
||||
params = {
|
||||
"destination_embedding": destination_embedding,
|
||||
"user_id": user_id,
|
||||
"agent_id": agent_id,
|
||||
"threshold": threshold,
|
||||
}
|
||||
else:
|
||||
cypher = """
|
||||
CALL vector_search.search("memzero", 1, $destination_embedding)
|
||||
YIELD distance, node, similarity
|
||||
WITH node AS destination_candidate, similarity
|
||||
WHERE node.user_id = $user_id
|
||||
AND similarity >= $threshold
|
||||
RETURN id(destination_candidate);
|
||||
"""
|
||||
params = {
|
||||
"destination_embedding": destination_embedding,
|
||||
"user_id": user_id,
|
||||
"threshold": threshold,
|
||||
}
|
||||
|
||||
result = self.graph.query(cypher, params=params)
|
||||
return result
|
||||
|
||||
|
||||
def _vector_index_exists(self, index_info, index_name):
|
||||
"""
|
||||
Check if a vector index exists, compatible with both Memgraph versions.
|
||||
|
||||
Args:
|
||||
index_info (dict): Index information from _fetch_existing_indexes
|
||||
index_name (str): Name of the index to check
|
||||
|
||||
Returns:
|
||||
bool: True if index exists, False otherwise
|
||||
"""
|
||||
vector_indexes = index_info.get("vector_index_exists", [])
|
||||
|
||||
# Check for index by name regardless of version-specific format differences
|
||||
return any(
|
||||
idx.get("index_name") == index_name or
|
||||
idx.get("index name") == index_name or
|
||||
idx.get("name") == index_name
|
||||
for idx in vector_indexes
|
||||
)
|
||||
|
||||
def _label_property_index_exists(self, index_info, label, property_name):
|
||||
"""
|
||||
Check if a label+property index exists, compatible with both versions.
|
||||
|
||||
Args:
|
||||
index_info (dict): Index information from _fetch_existing_indexes
|
||||
label (str): Label name
|
||||
property_name (str): Property name
|
||||
|
||||
Returns:
|
||||
bool: True if index exists, False otherwise
|
||||
"""
|
||||
indexes = index_info.get("index_exists", [])
|
||||
|
||||
return any(
|
||||
(idx.get("index type") == "label+property" or idx.get("index_type") == "label+property") and
|
||||
(idx.get("label") == label) and
|
||||
(idx.get("property") == property_name or property_name in str(idx.get("properties", "")))
|
||||
for idx in indexes
|
||||
)
|
||||
|
||||
def _label_index_exists(self, index_info, label):
|
||||
"""
|
||||
Check if a label index exists, compatible with both versions.
|
||||
|
||||
Args:
|
||||
index_info (dict): Index information from _fetch_existing_indexes
|
||||
label (str): Label name
|
||||
|
||||
Returns:
|
||||
bool: True if index exists, False otherwise
|
||||
"""
|
||||
indexes = index_info.get("index_exists", [])
|
||||
|
||||
return any(
|
||||
(idx.get("index type") == "label" or idx.get("index_type") == "label") and
|
||||
(idx.get("label") == label)
|
||||
for idx in indexes
|
||||
)
|
||||
|
||||
def _fetch_existing_indexes(self):
|
||||
"""
|
||||
Retrieves information about existing indexes and vector indexes in the Memgraph database.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing lists of existing indexes and vector indexes.
|
||||
"""
|
||||
try:
|
||||
index_exists = list(self.graph.query("SHOW INDEX INFO;"))
|
||||
vector_index_exists = list(self.graph.query("SHOW VECTOR INDEX INFO;"))
|
||||
return {"index_exists": index_exists, "vector_index_exists": vector_index_exists}
|
||||
except Exception as e:
|
||||
logger.warning(f"Error fetching indexes: {e}. Returning empty index info.")
|
||||
return {"index_exists": [], "vector_index_exists": []}
|
||||
+132
-3
@@ -2,6 +2,7 @@ import logging
|
||||
import sqlite3
|
||||
import threading
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -14,6 +15,7 @@ class SQLiteManager:
|
||||
self._lock = threading.Lock()
|
||||
self._migrate_history_table()
|
||||
self._create_history_table()
|
||||
self._create_messages_table()
|
||||
|
||||
def _migrate_history_table(self) -> None:
|
||||
"""
|
||||
@@ -123,6 +125,28 @@ class SQLiteManager:
|
||||
logger.error(f"Failed to create history table: {e}")
|
||||
raise
|
||||
|
||||
def _create_messages_table(self) -> None:
|
||||
with self._lock:
|
||||
try:
|
||||
self.connection.execute("BEGIN")
|
||||
self.connection.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS messages (
|
||||
id TEXT PRIMARY KEY,
|
||||
session_scope TEXT,
|
||||
role TEXT,
|
||||
content TEXT,
|
||||
name TEXT,
|
||||
created_at DATETIME
|
||||
)
|
||||
"""
|
||||
)
|
||||
self.connection.execute("COMMIT")
|
||||
except Exception as e:
|
||||
self.connection.execute("ROLLBACK")
|
||||
logger.error(f"Failed to create messages table: {e}")
|
||||
raise
|
||||
|
||||
def add_history(
|
||||
self,
|
||||
memory_id: str,
|
||||
@@ -166,6 +190,40 @@ class SQLiteManager:
|
||||
logger.error(f"Failed to add history record: {e}")
|
||||
raise
|
||||
|
||||
def batch_add_history(self, records: List[Dict[str, Any]]) -> None:
|
||||
with self._lock:
|
||||
try:
|
||||
self.connection.execute("BEGIN")
|
||||
self.connection.executemany(
|
||||
"""
|
||||
INSERT INTO history (
|
||||
id, memory_id, old_memory, new_memory, event,
|
||||
created_at, updated_at, is_deleted, actor_id, role
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
[
|
||||
(
|
||||
str(uuid.uuid4()),
|
||||
record.get("memory_id"),
|
||||
record.get("old_memory"),
|
||||
record.get("new_memory"),
|
||||
record.get("event"),
|
||||
record.get("created_at"),
|
||||
record.get("updated_at"),
|
||||
record.get("is_deleted", 0),
|
||||
record.get("actor_id"),
|
||||
record.get("role"),
|
||||
)
|
||||
for record in records
|
||||
],
|
||||
)
|
||||
self.connection.execute("COMMIT")
|
||||
except Exception as e:
|
||||
self.connection.execute("ROLLBACK")
|
||||
logger.error(f"Failed to batch add history records: {e}")
|
||||
raise
|
||||
|
||||
def get_history(self, memory_id: str) -> List[Dict[str, Any]]:
|
||||
with self._lock:
|
||||
cur = self.connection.execute(
|
||||
@@ -196,18 +254,89 @@ class SQLiteManager:
|
||||
for r in rows
|
||||
]
|
||||
|
||||
def save_messages(self, messages: List[Dict[str, Any]], session_scope: str) -> None:
|
||||
if not messages:
|
||||
return
|
||||
with self._lock:
|
||||
try:
|
||||
self.connection.execute("BEGIN")
|
||||
now = datetime.now(timezone.utc).isoformat()
|
||||
for message in messages:
|
||||
self.connection.execute(
|
||||
"""
|
||||
INSERT INTO messages (id, session_scope, role, content, name, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
str(uuid.uuid4()),
|
||||
session_scope,
|
||||
message.get("role"),
|
||||
message.get("content"),
|
||||
message.get("name"),
|
||||
now,
|
||||
),
|
||||
)
|
||||
# Evict old messages beyond the most recent 10 for this scope.
|
||||
# Wrapped in a derived table to force SQLite to materialize the
|
||||
# ORDER BY before the outer NOT IN evaluates it.
|
||||
self.connection.execute(
|
||||
"""
|
||||
DELETE FROM messages WHERE session_scope = ? AND id NOT IN (
|
||||
SELECT id FROM (
|
||||
SELECT id FROM messages WHERE session_scope = ? ORDER BY created_at DESC LIMIT 10
|
||||
)
|
||||
)
|
||||
""",
|
||||
(session_scope, session_scope),
|
||||
)
|
||||
self.connection.execute("COMMIT")
|
||||
except Exception as e:
|
||||
self.connection.execute("ROLLBACK")
|
||||
logger.error(f"Failed to save messages: {e}")
|
||||
raise
|
||||
|
||||
def get_last_messages(self, session_scope: str, limit: int = 10) -> List[Dict[str, Any]]:
|
||||
with self._lock:
|
||||
# Subquery picks the latest N rows (DESC + LIMIT), outer query
|
||||
# re-sorts them chronologically (ASC) for the caller.
|
||||
cur = self.connection.execute(
|
||||
"""
|
||||
SELECT role, content, name, created_at FROM (
|
||||
SELECT role, content, name, created_at
|
||||
FROM messages
|
||||
WHERE session_scope = ?
|
||||
ORDER BY created_at DESC
|
||||
LIMIT ?
|
||||
) ORDER BY created_at ASC
|
||||
""",
|
||||
(session_scope, limit),
|
||||
)
|
||||
rows = cur.fetchall()
|
||||
|
||||
return [
|
||||
{
|
||||
"role": r[0],
|
||||
"content": r[1],
|
||||
"name": r[2],
|
||||
"created_at": r[3],
|
||||
}
|
||||
for r in rows
|
||||
]
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Drop and recreate the history table."""
|
||||
"""Drop and recreate the history and messages tables."""
|
||||
with self._lock:
|
||||
try:
|
||||
self.connection.execute("BEGIN")
|
||||
self.connection.execute("DROP TABLE IF EXISTS history")
|
||||
self.connection.execute("DROP TABLE IF EXISTS messages")
|
||||
self.connection.execute("COMMIT")
|
||||
self._create_history_table()
|
||||
except Exception as e:
|
||||
self.connection.execute("ROLLBACK")
|
||||
logger.error(f"Failed to reset history table: {e}")
|
||||
logger.error(f"Failed to reset tables: {e}")
|
||||
raise
|
||||
self._create_history_table()
|
||||
self._create_messages_table()
|
||||
|
||||
def close(self) -> None:
|
||||
if self.connection:
|
||||
|
||||
@@ -169,9 +169,6 @@ def capture_event(event_name, memory_instance, additional_data=None):
|
||||
"collection": memory_instance.collection_name,
|
||||
"vector_size": memory_instance.embedding_model.config.embedding_dims,
|
||||
"history_store": "sqlite",
|
||||
"graph_store": f"{memory_instance.graph.__class__.__module__}.{memory_instance.graph.__class__.__name__}"
|
||||
if memory_instance.config.graph_store.config
|
||||
else None,
|
||||
"vector_store": f"{memory_instance.vector_store.__class__.__module__}.{memory_instance.vector_store.__class__.__name__}",
|
||||
"llm": f"{memory_instance.llm.__class__.__module__}.{memory_instance.llm.__class__.__name__}",
|
||||
"embedding_model": f"{memory_instance.embedding_model.__class__.__module__}.{memory_instance.embedding_model.__class__.__name__}",
|
||||
|
||||
@@ -0,0 +1,357 @@
|
||||
"""
|
||||
Entity extraction from text using spaCy NLP.
|
||||
|
||||
Extracts four types of entities from text:
|
||||
- **Proper nouns**: Capitalized multi-word sequences (person names, places, brands)
|
||||
- **Quoted text**: Text in single or double quotes (titles, specific terms)
|
||||
- **Noun compounds**: Multi-word noun phrases with specific modifiers (e.g., "machine learning")
|
||||
- **Noun fallback**: Single nouns from circumstantial compound patterns
|
||||
|
||||
Public API:
|
||||
extract_entities(text: str) -> List[Tuple[str, str]]
|
||||
|
||||
Internal:
|
||||
_extract_entities_from_doc(doc) -> List[Tuple[str, str]]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import re
|
||||
from typing import List, Tuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Words that are too generic to be useful as entity heads
|
||||
_GENERIC_HEADS = {
|
||||
"thing", "stuff", "way", "time", "experience", "situation", "case",
|
||||
"fact", "matter", "issue", "idea", "thought", "feeling", "place",
|
||||
"area", "part", "kind", "type", "sort", "lot", "bit", "day", "year",
|
||||
"week", "month", "moment", "instance", "example", "technique",
|
||||
"method", "approach", "process", "step", "tool", "result", "outcome",
|
||||
"goal", "task", "item", "topic", "scale", "size", "level", "degree",
|
||||
"amount", "number", "style", "look", "color", "colour", "shape",
|
||||
"form", "piece", "section", "side", "end", "edge", "surface", "point",
|
||||
}
|
||||
|
||||
# Modifiers that describe circumstance, not content
|
||||
_CIRCUMSTANTIAL_MODS = {
|
||||
"solo", "individual", "team", "group", "joint", "collaborative",
|
||||
"first", "last", "next", "previous", "final", "initial", "main", "side",
|
||||
}
|
||||
|
||||
# Adjectives too vague to make a compound entity specific
|
||||
_NON_SPECIFIC_ADJ = {
|
||||
"many", "few", "several", "some", "any", "all", "most", "more",
|
||||
"less", "much", "little", "enough", "various", "numerous", "multiple",
|
||||
"countless", "great", "good", "bad", "nice", "terrible", "awful",
|
||||
"awesome", "amazing", "wonderful", "horrible", "excellent", "poor",
|
||||
"best", "worst", "fine", "okay", "new", "old", "recent", "past",
|
||||
"future", "current", "previous", "next", "last", "first", "latest",
|
||||
"early", "late", "former", "modern", "ancient", "big", "small",
|
||||
"large", "tiny", "huge", "enormous", "long", "short", "tall", "high",
|
||||
"low", "wide", "narrow", "thick", "thin", "deep", "shallow",
|
||||
"similar", "different", "same", "other", "another", "such", "certain",
|
||||
"important", "main", "major", "minor", "key", "primary", "real",
|
||||
"actual", "true", "whole", "entire", "full", "complete", "total",
|
||||
"basic", "simple", "interesting", "boring", "exciting", "special",
|
||||
"particular", "general", "common", "unique", "rare", "typical",
|
||||
"usual", "normal", "regular", "possible", "likely", "potential",
|
||||
"available", "necessary", "only", "solo", "individual", "team",
|
||||
"group", "joint", "collaborative", "final", "initial", "side",
|
||||
}
|
||||
|
||||
# Generic tail words to strip from compound entities
|
||||
_GENERIC_ENDINGS = {
|
||||
"work", "works", "job", "jobs", "task", "tasks", "stuff", "things",
|
||||
"thing", "info", "information", "details", "data", "content",
|
||||
"material", "materials", "activities", "activity", "efforts", "effort",
|
||||
"options", "option", "choices", "choice", "results", "result",
|
||||
"output", "outputs", "products", "product", "items", "item",
|
||||
}
|
||||
|
||||
# Capitalized single words that are too generic to be proper nouns
|
||||
_GENERIC_CAPS = {
|
||||
"works", "items", "things", "stuff", "resources", "options", "tips",
|
||||
"ideas", "steps", "ways", "methods", "tools", "features", "benefits",
|
||||
"examples", "details", "notes", "instructions", "guidelines",
|
||||
"recommendations", "suggestions", "overview", "summary", "conclusion",
|
||||
"introduction", "pros", "cons", "advantages", "disadvantages",
|
||||
}
|
||||
|
||||
# Markdown/formatting markers to skip during extraction
|
||||
_FORMATTING_MARKERS = {"*", "-", "+", "\u2022", "\u2013", "\u2014", "#", "##", "###", "**", "__"}
|
||||
|
||||
|
||||
def _is_sentence_start(tokens: list, idx: int) -> bool:
|
||||
"""Check if a token is at the start of a sentence or after formatting."""
|
||||
if idx == 0:
|
||||
return True
|
||||
tok = tokens[idx]
|
||||
if tok.is_sent_start:
|
||||
return True
|
||||
prev = tokens[idx - 1].text
|
||||
return prev in ".!?:" or prev in _FORMATTING_MARKERS or "\n" in prev
|
||||
|
||||
|
||||
def _strip_generic_ending(toks: list) -> list:
|
||||
"""Remove generic trailing words from compound token sequences."""
|
||||
if len(toks) <= 1:
|
||||
return toks
|
||||
last = toks[-1].lemma_.lower() if hasattr(toks[-1], "lemma_") else toks[-1].lower()
|
||||
return toks[:-1] if last in _GENERIC_ENDINGS and len(toks) > 2 else toks
|
||||
|
||||
|
||||
def _lemmatize_compound(toks: list) -> str:
|
||||
"""Join compound tokens, lemmatizing nouns."""
|
||||
return " ".join(t.lemma_ if t.pos_ == "NOUN" else t.text for t in toks)
|
||||
|
||||
|
||||
def _has_artifacts(txt: str) -> bool:
|
||||
"""Check for formatting artifacts that indicate non-entity text."""
|
||||
return any(
|
||||
[
|
||||
"**" in txt or "__" in txt or ":*" in txt,
|
||||
re.search(r"\s\*\s|\s\*$|^\*\s", txt),
|
||||
" " in txt or "\n" in txt or "\t" in txt,
|
||||
len(txt) > 100,
|
||||
txt.startswith(("\u2022", "-", "+", "\u2013", "\u2014")),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def extract_entities(text: str) -> List[Tuple[str, str]]:
|
||||
"""Extract named entities, quoted text, and noun compounds from text.
|
||||
|
||||
This is the public API that accepts a string. It loads the spaCy model
|
||||
internally and delegates to _extract_entities_from_doc().
|
||||
|
||||
Args:
|
||||
text: Input text to extract entities from.
|
||||
|
||||
Returns:
|
||||
Deduplicated list of (entity_type, entity_text) tuples.
|
||||
Entity types: PROPER, QUOTED, COMPOUND, NOUN.
|
||||
Returns empty list if spaCy is unavailable.
|
||||
"""
|
||||
from mem0.utils.spacy_models import get_nlp_full
|
||||
|
||||
nlp = get_nlp_full()
|
||||
if nlp is None:
|
||||
return []
|
||||
|
||||
doc = nlp(text)
|
||||
return _extract_entities_from_doc(doc)
|
||||
|
||||
|
||||
def extract_entities_batch(texts: List[str], batch_size: int = 32) -> List[List[Tuple[str, str]]]:
|
||||
"""Extract entities from multiple texts using spaCy's nlp.pipe() for batched NER.
|
||||
|
||||
Uses spaCy's efficient batch processing pipeline instead of calling
|
||||
nlp() individually per text. Significantly faster for multiple texts.
|
||||
|
||||
Args:
|
||||
texts: List of input texts to extract entities from.
|
||||
batch_size: Number of texts to process in each spaCy batch.
|
||||
|
||||
Returns:
|
||||
List of entity lists, one per input text. Each entity list contains
|
||||
(entity_type, entity_text) tuples. Returns list of empty lists if
|
||||
spaCy is unavailable.
|
||||
"""
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
from mem0.utils.spacy_models import get_nlp_full
|
||||
|
||||
nlp = get_nlp_full()
|
||||
if nlp is None:
|
||||
return [[] for _ in texts]
|
||||
|
||||
results = []
|
||||
for doc in nlp.pipe(texts, batch_size=batch_size):
|
||||
results.append(_extract_entities_from_doc(doc))
|
||||
return results
|
||||
|
||||
|
||||
def _extract_entities_from_doc(doc) -> List[Tuple[str, str]]:
|
||||
"""Extract entities from a spaCy Doc object.
|
||||
|
||||
Ported from platform's shared.core.utils.entity_extraction.extract_entities().
|
||||
"""
|
||||
entities: List[Tuple[str, str]] = []
|
||||
text = doc.text
|
||||
tokens = list(doc)
|
||||
|
||||
# === PROPER NOUN SEQUENCES ===
|
||||
i = 0
|
||||
while i < len(tokens):
|
||||
tok = tokens[i]
|
||||
if tok.text in _FORMATTING_MARKERS:
|
||||
i += 1
|
||||
continue
|
||||
is_cap = tok.text and tok.text[0].isupper()
|
||||
is_label = i + 1 < len(tokens) and tokens[i + 1].text == ":"
|
||||
|
||||
if is_cap and not is_label and tok.pos_ in {"PROPN", "NOUN", "ADJ"}:
|
||||
seq = [(tok, i)]
|
||||
j = i + 1
|
||||
while j < len(tokens):
|
||||
t = tokens[j]
|
||||
if (t.text and t.text[0].isupper()) or t.text.lower() in {
|
||||
"'s", "of", "the", "in", "and", "for", "at", "is",
|
||||
}:
|
||||
seq.append((t, j))
|
||||
j += 1
|
||||
else:
|
||||
break
|
||||
# Strip trailing function words
|
||||
while seq and seq[-1][0].text.lower() in {"of", "the", "in", "and", "for", "at", "is", "'s"}:
|
||||
seq.pop()
|
||||
if seq:
|
||||
has_mid_cap = any(
|
||||
not _is_sentence_start(tokens, idx)
|
||||
for (t, idx) in seq
|
||||
if t.text[0].isupper() and t.text.lower() not in {"'s", "of", "the", "in", "and", "for", "at", "is"}
|
||||
)
|
||||
if has_mid_cap:
|
||||
phrase = "".join(t.text_with_ws for (t, idx) in seq).strip()
|
||||
if len(phrase) > 2:
|
||||
entities.append(("PROPER", phrase))
|
||||
i = j
|
||||
else:
|
||||
i += 1
|
||||
|
||||
# === QUOTED TEXT ===
|
||||
for m in re.finditer(r'"([^"]+)"', text):
|
||||
if len(m.group(1).strip()) > 2:
|
||||
entities.append(("QUOTED", m.group(1).strip()))
|
||||
for m in re.finditer(r"(?:^|[\s\(\[{,;])'([^']+)'(?=[\s\.,;:!?\)\]]|$)", text):
|
||||
if len(m.group(1).strip()) > 2:
|
||||
entities.append(("QUOTED", m.group(1).strip()))
|
||||
|
||||
# === NOUN-NOUN COMPOUNDS ===
|
||||
for chunk in doc.noun_chunks:
|
||||
chunk_tokens = list(chunk)
|
||||
split_indices: list = []
|
||||
poss_splits: list = []
|
||||
for idx, tok in enumerate(chunk_tokens):
|
||||
if tok.dep_ == "case" and tok.text in {"'s", "\u2019s", "'"}:
|
||||
split_indices.append(idx)
|
||||
poss_splits.append(idx)
|
||||
elif tok.pos_ == "PUNCT" and tok.text in {"'", '"', "\u2018", "\u2019", "\u201c", "\u201d"}:
|
||||
split_indices.append(idx)
|
||||
|
||||
if split_indices:
|
||||
groups: list = []
|
||||
prev = 0
|
||||
for split_idx in split_indices:
|
||||
if split_idx > prev:
|
||||
groups.append(chunk_tokens[prev:split_idx])
|
||||
if split_idx in poss_splits:
|
||||
next_split = next((s for s in split_indices if s > split_idx), None)
|
||||
owned = chunk_tokens[split_idx + 1: next_split if next_split else len(chunk_tokens)]
|
||||
if owned:
|
||||
first_content = next((t for t in owned if t.pos_ not in {"PUNCT", "PART"}), None)
|
||||
if not (first_content and first_content.text and first_content.text[0].isupper()):
|
||||
prev = next_split if next_split else len(chunk_tokens)
|
||||
continue
|
||||
prev = split_idx + 1
|
||||
if prev < len(chunk_tokens):
|
||||
groups.append(chunk_tokens[prev:])
|
||||
else:
|
||||
groups = [chunk_tokens]
|
||||
|
||||
for group in groups:
|
||||
if not group:
|
||||
continue
|
||||
head = next((t for t in reversed(group) if t.pos_ in {"NOUN", "PROPN"}), None)
|
||||
if not head:
|
||||
continue
|
||||
head_generic = head.lemma_.lower() in _GENERIC_HEADS
|
||||
content = [
|
||||
t
|
||||
for t in group
|
||||
if t.pos_ not in {"DET", "PRON", "PUNCT", "PART", "ADP", "SCONJ", "NUM"} and (t.pos_ == "ADJ" or not t.is_stop)
|
||||
]
|
||||
if not content:
|
||||
continue
|
||||
|
||||
compound_toks = [t for t in content if t.dep_ == "compound"]
|
||||
adj_toks = [t for t in content if t.pos_ == "ADJ" or t.dep_ == "amod"]
|
||||
has_spec_adj = any(t.lemma_.lower() not in _NON_SPECIFIC_ADJ for t in adj_toks)
|
||||
if head_generic and not has_spec_adj and not compound_toks:
|
||||
continue
|
||||
|
||||
if compound_toks:
|
||||
is_circ = any(t.lemma_.lower() in _CIRCUMSTANTIAL_MODS for t in compound_toks)
|
||||
if is_circ:
|
||||
val = head.lemma_ if head.pos_ == "NOUN" else head.text
|
||||
if len(val) > 2:
|
||||
entities.append(("NOUN", val))
|
||||
else:
|
||||
filtered = _strip_generic_ending(
|
||||
[t for t in content if not (t.pos_ == "ADJ" and t.lemma_.lower() in _NON_SPECIFIC_ADJ)]
|
||||
)
|
||||
if filtered:
|
||||
phrase = _lemmatize_compound(filtered)
|
||||
if len(phrase) > 3 and " " in phrase:
|
||||
entities.append(("COMPOUND", phrase))
|
||||
elif len(content) > 1 and has_spec_adj:
|
||||
filtered = _strip_generic_ending(
|
||||
[t for t in content if not ((t.pos_ == "ADJ" or t.dep_ == "amod") and t.lemma_.lower() in _NON_SPECIFIC_ADJ)]
|
||||
)
|
||||
if filtered:
|
||||
phrase = _lemmatize_compound(filtered)
|
||||
if len(phrase) > 3 and " " in phrase:
|
||||
entities.append(("COMPOUND", phrase))
|
||||
|
||||
# === FALLBACK: Mis-tagged VERB heads ===
|
||||
processed = {e[1].lower() for e in entities if e[0] == "COMPOUND"}
|
||||
generic_verb_heads = _GENERIC_HEADS | {"find", "buy", "purchase", "sale", "deal", "trip", "visit"}
|
||||
|
||||
def collect_compounds(head):
|
||||
return [t for t in doc if t.head == head and t.dep_ == "compound"]
|
||||
|
||||
for tok in doc:
|
||||
if tok.pos_ == "VERB" and tok.dep_ in {"pobj", "dobj", "nsubj"}:
|
||||
comps = sorted(collect_compounds(tok), key=lambda t: t.i)
|
||||
if comps:
|
||||
phrase_toks = comps if tok.lemma_.lower() in generic_verb_heads else comps + [tok]
|
||||
phrase = " ".join(t.text for t in phrase_toks)
|
||||
if phrase.lower() not in processed and len(phrase) > 3 and " " in phrase:
|
||||
entities.append(("COMPOUND", phrase))
|
||||
processed.add(phrase.lower())
|
||||
|
||||
# === DEDUPLICATION & CLEANUP ===
|
||||
seen: set = set()
|
||||
deduped = []
|
||||
for t, e in entities:
|
||||
k = e.lower().strip()
|
||||
if k not in seen and len(k) > 2:
|
||||
seen.add(k)
|
||||
deduped.append((t, e))
|
||||
|
||||
cleaned: List[Tuple[str, str]] = []
|
||||
for etype, etext in deduped:
|
||||
txt = re.sub(r"^\*+\s*|\s*\*+$", "", etext.strip())
|
||||
txt = re.sub(r"\s*:+$", "", txt)
|
||||
txt = re.sub(r"^\d+\s*\.\s*", "", txt)
|
||||
if not txt or len(txt) <= 2 or _has_artifacts(txt):
|
||||
continue
|
||||
if etype == "PROPER" and " " not in txt and txt.lower() in _GENERIC_CAPS:
|
||||
continue
|
||||
cleaned.append((etype, txt))
|
||||
|
||||
# Keep best type per entity (PROPER > COMPOUND > QUOTED > NOUN)
|
||||
type_pri = {"PROPER": 0, "COMPOUND": 1, "QUOTED": 2, "NOUN": 3, "VERB": 4}
|
||||
best: dict = {}
|
||||
for t, e in cleaned:
|
||||
k = e.lower()
|
||||
if k not in best or type_pri.get(t, 99) < type_pri.get(best[k][0], 99):
|
||||
best[k] = (t, e)
|
||||
deduped = list(best.values())
|
||||
|
||||
# Remove entities that are substrings of longer entities
|
||||
all_lower = [e[1].lower() for e in deduped]
|
||||
return [(t, e) for t, e in deduped if not any(e.lower() != o and e.lower() in o for o in all_lower)]
|
||||
@@ -209,30 +209,6 @@ class VectorStoreFactory:
|
||||
return instance
|
||||
|
||||
|
||||
class GraphStoreFactory:
|
||||
"""
|
||||
Factory for creating MemoryGraph instances for different graph store providers.
|
||||
Usage: GraphStoreFactory.create(provider_name, config)
|
||||
"""
|
||||
|
||||
provider_to_class = {
|
||||
"memgraph": "mem0.memory.memgraph_memory.MemoryGraph",
|
||||
"neptune": "mem0.graphs.neptune.neptunegraph.MemoryGraph",
|
||||
"neptunedb": "mem0.graphs.neptune.neptunedb.MemoryGraph",
|
||||
"kuzu": "mem0.memory.kuzu_memory.MemoryGraph",
|
||||
"apache_age": "mem0.memory.apache_age_memory.MemoryGraph",
|
||||
"default": "mem0.memory.graph_memory.MemoryGraph",
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def create(cls, provider_name, config):
|
||||
class_type = cls.provider_to_class.get(provider_name, cls.provider_to_class["default"])
|
||||
try:
|
||||
GraphClass = load_class(class_type)
|
||||
except (ImportError, AttributeError) as e:
|
||||
raise ImportError(f"Could not import MemoryGraph for provider '{provider_name}': {e}")
|
||||
return GraphClass(config)
|
||||
|
||||
|
||||
class RerankerFactory:
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
"""
|
||||
BM25 lemmatization for consistent keyword matching.
|
||||
|
||||
Uses spaCy's lemmatizer for better handling of:
|
||||
- Verb forms: attending/attends/attended -> attend
|
||||
- Comparatives/superlatives: older/oldest -> old
|
||||
- Plurals: memories -> memory
|
||||
- Avoids over-stemming: organization != organize
|
||||
|
||||
Also includes original -ing forms alongside lemmas to handle cases
|
||||
where spaCy's context-dependent lemmatization produces inconsistent
|
||||
results (e.g., "meeting" as noun vs verb -> different lemmas).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def lemmatize_for_bm25(text: str) -> str:
|
||||
"""Lemmatize text for BM25 matching.
|
||||
|
||||
Returns space-joined lemmas for full-text search. Falls back to
|
||||
the original text if spaCy is unavailable.
|
||||
"""
|
||||
from mem0.utils.spacy_models import get_nlp_lemma
|
||||
|
||||
nlp = get_nlp_lemma()
|
||||
if nlp is None:
|
||||
return text
|
||||
|
||||
doc = nlp(text.lower())
|
||||
tokens = []
|
||||
|
||||
for token in doc:
|
||||
if token.is_punct or token.is_stop:
|
||||
continue
|
||||
|
||||
lemma = token.lemma_
|
||||
if lemma.isalnum():
|
||||
tokens.append(lemma)
|
||||
|
||||
# Also add original if it ends in -ing and differs from lemma.
|
||||
# This handles noun/verb ambiguity (meeting/meet, attending/attend).
|
||||
if token.text.endswith("ing") and token.text != lemma and token.text.isalnum():
|
||||
tokens.append(token.text)
|
||||
|
||||
return " ".join(tokens)
|
||||
@@ -0,0 +1,126 @@
|
||||
"""
|
||||
Scoring utilities for hybrid retrieval.
|
||||
|
||||
Provides:
|
||||
- **BM25 normalization**: Sigmoid normalization of raw BM25 scores to [0, 1].
|
||||
- **BM25 parameter selection**: Query-length-adaptive sigmoid parameters.
|
||||
- **Additive scoring**: Combined scoring with semantic + BM25 + entity boost.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
def get_bm25_params(query: str, *, lemmatized: Optional[str] = None) -> tuple:
|
||||
"""Get BM25 sigmoid parameters based on query length.
|
||||
|
||||
Longer queries tend to have higher raw BM25 scores, so we adjust
|
||||
the sigmoid midpoint and steepness accordingly.
|
||||
|
||||
Returns:
|
||||
(midpoint, steepness) for sigmoid normalization.
|
||||
"""
|
||||
if lemmatized is None:
|
||||
from mem0.utils.lemmatization import lemmatize_for_bm25
|
||||
|
||||
lemmatized = lemmatize_for_bm25(query)
|
||||
num_terms = len(lemmatized.split()) if lemmatized else 1
|
||||
|
||||
if num_terms <= 3:
|
||||
return 5.0, 0.7
|
||||
elif num_terms <= 6:
|
||||
return 7.0, 0.6
|
||||
elif num_terms <= 9:
|
||||
return 9.0, 0.5
|
||||
elif num_terms <= 15:
|
||||
return 10.0, 0.5
|
||||
else:
|
||||
return 12.0, 0.5
|
||||
|
||||
|
||||
def normalize_bm25(raw_score: float, midpoint: float, steepness: float) -> float:
|
||||
"""Normalize BM25 score to [0, 1] using logistic sigmoid.
|
||||
|
||||
Args:
|
||||
raw_score: Raw BM25 score (unbounded, typically 0-20+).
|
||||
midpoint: Score at which sigmoid outputs 0.5.
|
||||
steepness: Controls how quickly sigmoid transitions.
|
||||
|
||||
Returns:
|
||||
Normalized score in range [0, 1].
|
||||
"""
|
||||
return 1.0 / (1.0 + math.exp(-steepness * (raw_score - midpoint)))
|
||||
|
||||
|
||||
ENTITY_BOOST_WEIGHT = 0.5
|
||||
|
||||
|
||||
def score_and_rank(
|
||||
semantic_results: List[Dict[str, Any]],
|
||||
bm25_scores: Dict[str, float],
|
||||
entity_boosts: Dict[str, float],
|
||||
threshold: float,
|
||||
top_k: int,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Score candidates additively and return top-k results.
|
||||
|
||||
For each candidate:
|
||||
semantic_score is taken from the result's score field.
|
||||
combined = (semantic + bm25 + entity_boost) / max_possible
|
||||
|
||||
Threshold gates the semantic score BEFORE combining -- candidates
|
||||
below the threshold are excluded even if BM25/entity would boost them.
|
||||
|
||||
The divisor adapts based on which signals are active:
|
||||
- Semantic only: max_possible = 1.0
|
||||
- Semantic + BM25: max_possible = 2.0
|
||||
- Semantic + BM25 + entity: max_possible = 2.5
|
||||
- Semantic + entity (no BM25): max_possible = 1.5
|
||||
|
||||
Returns:
|
||||
List of scored result dicts sorted by combined score descending.
|
||||
"""
|
||||
has_bm25 = bool(bm25_scores)
|
||||
has_entity = bool(entity_boosts)
|
||||
|
||||
max_possible = 1.0
|
||||
if has_bm25:
|
||||
max_possible += 1.0
|
||||
if has_entity:
|
||||
max_possible += ENTITY_BOOST_WEIGHT
|
||||
|
||||
scored: List[Dict[str, Any]] = []
|
||||
|
||||
for result in semantic_results:
|
||||
mem_id = result.get("id")
|
||||
if mem_id is None:
|
||||
continue
|
||||
|
||||
semantic_score = result.get("score", 0.0)
|
||||
if semantic_score < threshold:
|
||||
continue
|
||||
|
||||
mem_id_str = str(mem_id)
|
||||
bm25_score = bm25_scores.get(mem_id_str, 0.0)
|
||||
entity_boost = entity_boosts.get(mem_id_str, 0.0)
|
||||
|
||||
raw_combined = semantic_score + bm25_score + entity_boost
|
||||
combined = min(raw_combined / max_possible, 1.0)
|
||||
|
||||
scored.append(
|
||||
{
|
||||
"id": mem_id_str,
|
||||
"score": combined,
|
||||
"score_breakdown": {
|
||||
"semantic": semantic_score,
|
||||
"bm25": bm25_score,
|
||||
"entity_boost": entity_boost,
|
||||
},
|
||||
"payload": result.get("payload"),
|
||||
}
|
||||
)
|
||||
|
||||
scored.sort(key=lambda x: x["score"], reverse=True)
|
||||
return scored[:top_k]
|
||||
@@ -0,0 +1,91 @@
|
||||
"""
|
||||
Shared spaCy model loader.
|
||||
|
||||
Consolidates spaCy model loading into a single module so that
|
||||
entity_extraction and lemmatization share one instance instead of
|
||||
each loading their own copy from disk.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import threading
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_nlp_full = None
|
||||
_nlp_lemma = None
|
||||
_load_failed_full = False
|
||||
_load_failed_lemma = False
|
||||
_lock = threading.Lock()
|
||||
|
||||
|
||||
def _ensure_model_available():
|
||||
"""Download en_core_web_sm if spaCy is installed but model is missing."""
|
||||
try:
|
||||
import spacy
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"spaCy is not installed. Install it with: pip install mem0ai[nlp]"
|
||||
)
|
||||
|
||||
if not spacy.util.is_package("en_core_web_sm"):
|
||||
logger.info("Downloading spaCy model en_core_web_sm...")
|
||||
try:
|
||||
from spacy.cli import download
|
||||
|
||||
download("en_core_web_sm")
|
||||
logger.info("spaCy model en_core_web_sm downloaded successfully")
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Failed to download spaCy model en_core_web_sm: {e}. "
|
||||
"Please install manually: python -m spacy download en_core_web_sm"
|
||||
) from e
|
||||
|
||||
|
||||
def get_nlp_full():
|
||||
"""Return spaCy model with all pipelines (NER, tagger, etc.) for entity extraction."""
|
||||
global _nlp_full, _load_failed_full
|
||||
if _load_failed_full:
|
||||
return None
|
||||
if _nlp_full is not None:
|
||||
return _nlp_full
|
||||
with _lock:
|
||||
if _nlp_full is not None:
|
||||
return _nlp_full
|
||||
if _load_failed_full:
|
||||
return None
|
||||
try:
|
||||
_ensure_model_available()
|
||||
import spacy
|
||||
|
||||
_nlp_full = spacy.load("en_core_web_sm")
|
||||
logger.info("spaCy full model loaded")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load spaCy full model: {e}")
|
||||
_load_failed_full = True
|
||||
return None
|
||||
return _nlp_full
|
||||
|
||||
|
||||
def get_nlp_lemma():
|
||||
"""Return spaCy model with only lemmatizer for BM25 text processing."""
|
||||
global _nlp_lemma, _load_failed_lemma
|
||||
if _load_failed_lemma:
|
||||
return None
|
||||
if _nlp_lemma is not None:
|
||||
return _nlp_lemma
|
||||
with _lock:
|
||||
if _nlp_lemma is not None:
|
||||
return _nlp_lemma
|
||||
if _load_failed_lemma:
|
||||
return None
|
||||
try:
|
||||
_ensure_model_available()
|
||||
import spacy
|
||||
|
||||
_nlp_lemma = spacy.load("en_core_web_sm", disable=["ner", "parser"])
|
||||
logger.info("spaCy lemma model loaded")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load spaCy lemma model: {e}")
|
||||
_load_failed_lemma = True
|
||||
return None
|
||||
return _nlp_lemma
|
||||
@@ -246,6 +246,34 @@ class AzureAISearch(VectorStoreBase):
|
||||
results.append(OutputData(id=result["id"], score=result["@search.score"], payload=payload))
|
||||
return results
|
||||
|
||||
def keyword_search(self, query, top_k=5, filters=None):
|
||||
"""Search for memories using keyword/BM25 text matching (no vector queries).
|
||||
|
||||
Args:
|
||||
query (str): The text query to search for.
|
||||
top_k (int): Maximum number of results to return. Defaults to 5.
|
||||
filters (Dict, optional): Filters to apply to the search.
|
||||
|
||||
Returns:
|
||||
List[OutputData]: Search results with id, score, and payload.
|
||||
"""
|
||||
filter_expression = None
|
||||
if filters:
|
||||
filter_expression = self._build_filter_expression(filters)
|
||||
|
||||
search_results = self.search_client.search(
|
||||
search_text=query,
|
||||
filter=filter_expression,
|
||||
top=top_k,
|
||||
search_fields=["payload"],
|
||||
)
|
||||
|
||||
results = []
|
||||
for result in search_results:
|
||||
payload = json.loads(extract_json(result["payload"]))
|
||||
results.append(OutputData(id=result["id"], score=result["@search.score"], payload=payload))
|
||||
return results
|
||||
|
||||
def delete(self, vector_id):
|
||||
"""
|
||||
Delete a vector by ID.
|
||||
|
||||
@@ -178,17 +178,32 @@ class AzureMySQL(VectorStoreBase):
|
||||
dims = vector_size or self.embedding_model_dims
|
||||
|
||||
with self._get_cursor(commit=True) as cur:
|
||||
# Create table with vector column
|
||||
# Create table with vector column and a generated column for fulltext keyword search
|
||||
cur.execute(f"""
|
||||
CREATE TABLE IF NOT EXISTS `{table_name}` (
|
||||
id VARCHAR(255) PRIMARY KEY,
|
||||
vector JSON,
|
||||
payload JSON,
|
||||
text_lemmatized VARCHAR(1000) GENERATED ALWAYS AS
|
||||
(CAST(payload->>'$.text_lemmatized' AS CHAR(1000))) STORED,
|
||||
INDEX idx_payload_keys ((CAST(payload AS CHAR(255)) ARRAY))
|
||||
)
|
||||
""")
|
||||
logger.info(f"Created collection '{table_name}' with vector dimension {dims}")
|
||||
|
||||
# Add FULLTEXT index on text_lemmatized for keyword_search()
|
||||
try:
|
||||
cur.execute(f"""
|
||||
CREATE FULLTEXT INDEX ft_text_lemmatized
|
||||
ON `{table_name}` (text_lemmatized)
|
||||
""")
|
||||
logger.info(f"Created FULLTEXT index on '{table_name}.text_lemmatized'")
|
||||
except Exception as e:
|
||||
logger.debug(
|
||||
f"Could not create FULLTEXT index on '{table_name}.text_lemmatized': {e}. "
|
||||
"It may already exist or FULLTEXT may not be supported."
|
||||
)
|
||||
|
||||
def insert(self, vectors: List[List[float]], payloads: Optional[List[Dict]] = None, ids: Optional[List[str]] = None):
|
||||
"""
|
||||
Insert vectors into the collection.
|
||||
@@ -300,6 +315,61 @@ class AzureMySQL(VectorStoreBase):
|
||||
for r in scored_results
|
||||
]
|
||||
|
||||
def keyword_search(self, query, top_k=5, filters=None):
|
||||
"""
|
||||
Search for memories using MySQL FULLTEXT search via MATCH() AGAINST().
|
||||
|
||||
This method attempts to use a FULLTEXT index on the text_lemmatized column.
|
||||
If the column or index does not exist, it returns None gracefully.
|
||||
|
||||
Args:
|
||||
query (str): The text query for keyword-based search.
|
||||
top_k (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (dict, optional): Filters to apply to the search. Defaults to None.
|
||||
|
||||
Returns:
|
||||
list: Search results in the same format as search(), or None if FULLTEXT
|
||||
search is not supported on this collection.
|
||||
"""
|
||||
try:
|
||||
filter_conditions = []
|
||||
filter_params = []
|
||||
|
||||
if filters:
|
||||
for k, v in filters.items():
|
||||
filter_conditions.append("JSON_EXTRACT(payload, %s) = %s")
|
||||
filter_params.extend([f"$.{k}", json.dumps(v)])
|
||||
|
||||
filter_clause = ""
|
||||
if filter_conditions:
|
||||
filter_clause = " AND " + " AND ".join(filter_conditions)
|
||||
|
||||
with self._get_cursor() as cur:
|
||||
query_sql = f"""
|
||||
SELECT id, payload,
|
||||
MATCH(text_lemmatized) AGAINST(%s IN NATURAL LANGUAGE MODE) AS score
|
||||
FROM `{self.collection_name}`
|
||||
WHERE MATCH(text_lemmatized) AGAINST(%s IN NATURAL LANGUAGE MODE)
|
||||
{filter_clause}
|
||||
ORDER BY score DESC
|
||||
LIMIT %s
|
||||
"""
|
||||
params = [query, query] + filter_params + [top_k]
|
||||
cur.execute(query_sql, params)
|
||||
results = cur.fetchall()
|
||||
|
||||
return [
|
||||
OutputData(
|
||||
id=r['id'],
|
||||
score=float(r['score']),
|
||||
payload=json.loads(r['payload']) if isinstance(r['payload'], str) else r['payload'],
|
||||
)
|
||||
for r in results
|
||||
]
|
||||
except Exception as e:
|
||||
logger.debug(f"Keyword search not available for collection {self.collection_name}: {e}")
|
||||
return None
|
||||
|
||||
def delete(self, vector_id: str):
|
||||
"""
|
||||
Delete a vector by ID.
|
||||
|
||||
@@ -27,6 +27,7 @@ try:
|
||||
VectorIndex,
|
||||
)
|
||||
from pymochow.model.table import (
|
||||
BM25SearchRequest,
|
||||
FloatVector,
|
||||
Partition,
|
||||
Row,
|
||||
@@ -227,6 +228,48 @@ class BaiduDB(VectorStoreBase):
|
||||
|
||||
return output
|
||||
|
||||
def keyword_search(self, query, top_k=5, filters=None):
|
||||
"""
|
||||
Perform keyword-based search using Baidu Mochow's BM25 search.
|
||||
|
||||
Args:
|
||||
query (str): The text query to search for.
|
||||
top_k (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (Dict, optional): Filters to apply to the search.
|
||||
|
||||
Returns:
|
||||
list: Search results, or None if the table lacks an inverted index.
|
||||
"""
|
||||
try:
|
||||
search_filter = None
|
||||
if filters:
|
||||
search_filter = self._create_filter(filters)
|
||||
|
||||
request = BM25SearchRequest(
|
||||
index_name="data_bm25_idx",
|
||||
search_text=query,
|
||||
limit=top_k,
|
||||
filter=search_filter,
|
||||
)
|
||||
|
||||
projections = ["id", "metadata"]
|
||||
res = self._table.bm25_search(request=request, projections=projections)
|
||||
|
||||
output = []
|
||||
for row in res.rows:
|
||||
row_data = row.get("row", {})
|
||||
output_data = OutputData(
|
||||
id=row_data.get("id"),
|
||||
score=row.get("score", 0.0),
|
||||
payload=row_data.get("metadata", {}),
|
||||
)
|
||||
output.append(output_data)
|
||||
|
||||
return output
|
||||
except Exception as e:
|
||||
logger.error(f"Error during keyword search for query '{query}': {e}")
|
||||
return None
|
||||
|
||||
def delete(self, vector_id):
|
||||
"""
|
||||
Delete a vector by ID.
|
||||
|
||||
@@ -56,3 +56,37 @@ class VectorStoreBase(ABC):
|
||||
def reset(self):
|
||||
"""Reset by delete the collection and recreate it."""
|
||||
pass
|
||||
|
||||
def keyword_search(self, query: str, top_k: int = 5, filters: dict = None):
|
||||
"""Keyword/BM25 full-text search. Returns None if not supported by this store.
|
||||
|
||||
Override in subclasses that support native keyword/BM25 search.
|
||||
Returns results in the same format as search() -- list of objects with
|
||||
id, score, and payload attributes.
|
||||
|
||||
Args:
|
||||
query: The search query text (should be lemmatized for best results).
|
||||
top_k: Maximum number of results to return.
|
||||
filters: Optional metadata filters (same format as search filters).
|
||||
|
||||
Returns:
|
||||
List of search results with id, score, payload, or None if not supported.
|
||||
"""
|
||||
return None
|
||||
|
||||
def search_batch(self, queries: list, vectors_list: list, top_k: int = 1, filters: dict = None):
|
||||
"""Batch search for multiple queries at once.
|
||||
|
||||
Default implementation calls search() sequentially. Override in subclasses
|
||||
with native batch support (e.g., Qdrant query_batch_points).
|
||||
|
||||
Args:
|
||||
queries: List of query texts.
|
||||
vectors_list: List of query vectors (one per query).
|
||||
top_k: Maximum results per query.
|
||||
filters: Optional metadata filters applied to all queries.
|
||||
|
||||
Returns:
|
||||
List of result lists, one per query.
|
||||
"""
|
||||
return [self.search(q, v, top_k=top_k, filters=filters) for q, v in zip(queries, vectors_list)]
|
||||
|
||||
@@ -515,6 +515,55 @@ class Databricks(VectorStoreBase):
|
||||
logger.error(f"Search failed: {e}")
|
||||
raise
|
||||
|
||||
def keyword_search(self, query, top_k=5, filters=None):
|
||||
"""
|
||||
Search for memories using full-text keyword search.
|
||||
|
||||
Only supported for DELTA_SYNC index type. Returns None for DIRECT_ACCESS indexes.
|
||||
|
||||
Args:
|
||||
query (str): Search query text.
|
||||
top_k (int): Maximum number of results. Defaults to 5.
|
||||
filters (dict, optional): Filters to apply.
|
||||
|
||||
Returns:
|
||||
List[MemoryResult] or None: Search results, or None if index type is DIRECT_ACCESS.
|
||||
"""
|
||||
if self.index_type == VectorIndexType.DIRECT_ACCESS:
|
||||
logger.warning("keyword_search is not supported for DIRECT_ACCESS index type.")
|
||||
return None
|
||||
|
||||
try:
|
||||
filters_json = json.dumps(filters) if filters else None
|
||||
|
||||
sdk_results = self.client.vector_search_indexes.query_index(
|
||||
index_name=self.fully_qualified_index_name,
|
||||
columns=self.column_names,
|
||||
query_text=query,
|
||||
num_results=top_k,
|
||||
query_type="FULL_TEXT",
|
||||
filters_json=filters_json,
|
||||
)
|
||||
|
||||
result_data = sdk_results.result if hasattr(sdk_results, "result") else sdk_results
|
||||
data_array = result_data.data_array if getattr(result_data, "data_array", None) else []
|
||||
|
||||
memory_results = []
|
||||
for row in data_array:
|
||||
row_dict = dict(zip(self.column_names, row)) if isinstance(row, (list, tuple)) else row
|
||||
score = row_dict.get("score") or (
|
||||
row[-1] if isinstance(row, (list, tuple)) and len(row) > len(self.column_names) else None
|
||||
)
|
||||
payload = {k: row_dict.get(k) for k in self.column_names}
|
||||
payload["data"] = payload.get("memory", "")
|
||||
memory_id = row_dict.get("memory_id") or row_dict.get("id")
|
||||
memory_results.append(MemoryResult(id=memory_id, score=score, payload=payload))
|
||||
return memory_results
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Keyword search failed: {e}")
|
||||
raise
|
||||
|
||||
def delete(self, vector_id):
|
||||
"""
|
||||
Delete a vector by ID from the Delta table.
|
||||
|
||||
@@ -158,6 +158,49 @@ class ElasticsearchDB(VectorStoreBase):
|
||||
|
||||
return results
|
||||
|
||||
def keyword_search(self, query, top_k=5, filters=None):
|
||||
"""Search for memories using BM25 keyword matching.
|
||||
|
||||
Args:
|
||||
query (str): The text query to search for.
|
||||
top_k (int): Maximum number of results to return. Defaults to 5.
|
||||
filters (Dict, optional): Filters to apply to the search.
|
||||
|
||||
Returns:
|
||||
List[OutputData]: Search results with id, score, and payload.
|
||||
"""
|
||||
# Build a multi_match query across text fields in metadata
|
||||
should_clauses = [
|
||||
{"match": {"metadata.data": query}},
|
||||
{"match": {"metadata.text_lemmatized": query}},
|
||||
]
|
||||
|
||||
bool_query = {
|
||||
"should": should_clauses,
|
||||
"minimum_should_match": 1,
|
||||
}
|
||||
|
||||
if filters:
|
||||
filter_conditions = []
|
||||
for key, value in filters.items():
|
||||
filter_conditions.append({"term": {f"metadata.{key}": value}})
|
||||
bool_query["filter"] = filter_conditions
|
||||
|
||||
search_query = {
|
||||
"size": top_k,
|
||||
"query": {"bool": bool_query},
|
||||
}
|
||||
|
||||
response = self.client.search(index=self.collection_name, body=search_query)
|
||||
|
||||
results = []
|
||||
for hit in response["hits"]["hits"]:
|
||||
results.append(
|
||||
OutputData(id=hit["_id"], score=hit["_score"], payload=hit.get("_source", {}).get("metadata", {}))
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
def delete(self, vector_id: str) -> None:
|
||||
"""Delete a vector by ID."""
|
||||
self.client.delete(index=self.collection_name, id=vector_id)
|
||||
|
||||
@@ -11,7 +11,7 @@ try:
|
||||
except ImportError:
|
||||
raise ImportError("The 'pymilvus' library is required. Please install it using 'pip install pymilvus'.")
|
||||
|
||||
from pymilvus import CollectionSchema, DataType, FieldSchema, MilvusClient
|
||||
from pymilvus import CollectionSchema, DataType, FieldSchema, Function, FunctionType, MilvusClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -73,14 +73,34 @@ class MilvusDB(VectorStoreBase):
|
||||
FieldSchema(name="id", dtype=DataType.VARCHAR, is_primary=True, max_length=512),
|
||||
FieldSchema(name="vectors", dtype=DataType.FLOAT_VECTOR, dim=vector_size),
|
||||
FieldSchema(name="metadata", dtype=DataType.JSON),
|
||||
# Text field for BM25 full-text search (auto-tokenized by Milvus analyzer)
|
||||
FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535, enable_analyzer=True),
|
||||
# Sparse vector field populated automatically by the BM25 function below
|
||||
FieldSchema(name="sparse", dtype=DataType.SPARSE_FLOAT_VECTOR),
|
||||
]
|
||||
|
||||
schema = CollectionSchema(fields, enable_dynamic_field=True)
|
||||
|
||||
index = self.client.prepare_index_params(
|
||||
# Add BM25 function so Milvus auto-generates sparse vectors from the text field
|
||||
bm25_function = Function(
|
||||
name="bm25",
|
||||
input_field_names=["text"],
|
||||
output_field_names=["sparse"],
|
||||
function_type=FunctionType.BM25,
|
||||
)
|
||||
schema.add_function(bm25_function)
|
||||
|
||||
index_params = self.client.prepare_index_params()
|
||||
index_params.add_index(
|
||||
field_name="vectors", metric_type=metric_type, index_type="AUTOINDEX", index_name="vector_index"
|
||||
)
|
||||
self.client.create_collection(collection_name=collection_name, schema=schema, index_params=index)
|
||||
index_params.add_index(
|
||||
field_name="sparse",
|
||||
index_type="SPARSE_INVERTED_INDEX",
|
||||
metric_type="BM25",
|
||||
index_name="sparse_index",
|
||||
)
|
||||
self.client.create_collection(collection_name=collection_name, schema=schema, index_params=index_params)
|
||||
|
||||
def insert(self, ids, vectors, payloads, **kwargs: Optional[dict[str, any]]):
|
||||
"""Insert vectors into a collection.
|
||||
@@ -92,7 +112,13 @@ class MilvusDB(VectorStoreBase):
|
||||
"""
|
||||
# Batch insert all records at once for better performance and consistency
|
||||
data = [
|
||||
{"id": idx, "vectors": embedding, "metadata": metadata}
|
||||
{
|
||||
"id": idx,
|
||||
"vectors": embedding,
|
||||
"metadata": metadata,
|
||||
# Populate the text field for BM25 sparse search; prefer lemmatized text, fall back to raw data
|
||||
"text": (metadata.get("text_lemmatized") or metadata.get("data", ""))[:65535] if metadata else "",
|
||||
}
|
||||
for idx, embedding, metadata in zip(ids, vectors, payloads)
|
||||
]
|
||||
self.client.insert(collection_name=self.collection_name, data=data, **kwargs)
|
||||
@@ -163,6 +189,39 @@ class MilvusDB(VectorStoreBase):
|
||||
result = self._parse_output(data=hits[0])
|
||||
return result
|
||||
|
||||
def keyword_search(self, query, top_k=5, filters=None):
|
||||
"""
|
||||
Search for memories using BM25-based full-text search via Milvus sparse vector support.
|
||||
|
||||
Milvus 2.5+ supports native BM25 via full-text search with a SPARSE_FLOAT_VECTOR field.
|
||||
This method attempts to use that capability. If the collection does not have a sparse
|
||||
field configured, it returns None gracefully.
|
||||
|
||||
Args:
|
||||
query (str): The text query for keyword-based search.
|
||||
top_k (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (dict, optional): Filters to apply to the search. Defaults to None.
|
||||
|
||||
Returns:
|
||||
list: Search results in the same format as search(), or None if sparse search
|
||||
is not supported on this collection.
|
||||
"""
|
||||
try:
|
||||
query_filter = self._create_filter(filters) if filters else None
|
||||
hits = self.client.search(
|
||||
collection_name=self.collection_name,
|
||||
data=[query],
|
||||
anns_field="sparse",
|
||||
limit=top_k,
|
||||
filter=query_filter,
|
||||
output_fields=["*"],
|
||||
)
|
||||
result = self._parse_output(data=hits[0])
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.debug(f"Keyword search not available for collection {self.collection_name}: {e}")
|
||||
return None
|
||||
|
||||
def delete(self, vector_id):
|
||||
"""
|
||||
Delete a vector by ID.
|
||||
@@ -192,7 +251,10 @@ class MilvusDB(VectorStoreBase):
|
||||
if payload is None:
|
||||
payload = existing[0].get("metadata")
|
||||
|
||||
schema = {"id": vector_id, "vectors": vector, "metadata": payload}
|
||||
text = ""
|
||||
if payload:
|
||||
text = (payload.get("text_lemmatized") or payload.get("data", ""))[:65535]
|
||||
schema = {"id": vector_id, "vectors": vector, "metadata": payload, "text": text}
|
||||
self.client.upsert(collection_name=self.collection_name, data=schema)
|
||||
|
||||
def get(self, vector_id):
|
||||
|
||||
@@ -85,6 +85,41 @@ class MongoDB(VectorStoreBase):
|
||||
logger.info(
|
||||
f"Search index '{self.index_name}' created successfully for collection '{self.collection_name}'."
|
||||
)
|
||||
|
||||
# Create Atlas Search text index for keyword_search()
|
||||
text_index_name = f"{self.collection_name}_text_search_index"
|
||||
try:
|
||||
found_text_indexes = list(collection.list_search_indexes(name=text_index_name))
|
||||
if not found_text_indexes:
|
||||
text_search_index_model = SearchIndexModel(
|
||||
name=text_index_name,
|
||||
definition={
|
||||
"mappings": {
|
||||
"dynamic": False,
|
||||
"fields": {
|
||||
"payload": {
|
||||
"type": "document",
|
||||
"fields": {
|
||||
"data": {"type": "string"},
|
||||
"text_lemmatized": {"type": "string"},
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
},
|
||||
)
|
||||
collection.create_search_index(text_search_index_model)
|
||||
logger.info(
|
||||
f"Text search index '{text_index_name}' created successfully for collection '{self.collection_name}'."
|
||||
)
|
||||
else:
|
||||
logger.info(f"Text search index '{text_index_name}' already exists in collection '{self.collection_name}'.")
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Could not create text search index '{text_index_name}': {e}. "
|
||||
"Atlas Search may not be available. keyword_search() will not work."
|
||||
)
|
||||
|
||||
return collection
|
||||
except PyMongoError as e:
|
||||
logger.error(f"Error creating collection and search index: {e}")
|
||||
@@ -168,6 +203,58 @@ class MongoDB(VectorStoreBase):
|
||||
output = [OutputData(id=str(doc["_id"]), score=doc.get("score"), payload=doc.get("payload")) for doc in results]
|
||||
return output
|
||||
|
||||
def keyword_search(self, query, top_k=5, filters=None):
|
||||
"""
|
||||
Perform keyword-based search using MongoDB Atlas Search.
|
||||
|
||||
Args:
|
||||
query (str): The text query to search for.
|
||||
top_k (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (Dict, optional): Filters to apply to the search.
|
||||
|
||||
Returns:
|
||||
List[OutputData]: Search results, or None if Atlas Search index is not available.
|
||||
"""
|
||||
try:
|
||||
collection = self.client[self.db_name][self.collection_name]
|
||||
search_index_name = f"{self.collection_name}_text_search_index"
|
||||
|
||||
pipeline = [
|
||||
{
|
||||
"$search": {
|
||||
"index": search_index_name,
|
||||
"text": {
|
||||
"query": query,
|
||||
"path": ["payload.data", "payload.text_lemmatized"],
|
||||
},
|
||||
}
|
||||
},
|
||||
{"$set": {"score": {"$meta": "searchScore"}}},
|
||||
{"$project": {"embedding": 0}},
|
||||
]
|
||||
|
||||
# Add filter stage if filters are provided
|
||||
if filters:
|
||||
filter_conditions = []
|
||||
for key, value in filters.items():
|
||||
filter_conditions.append({"payload." + key: value})
|
||||
if filter_conditions:
|
||||
pipeline.insert(1, {"$match": {"$and": filter_conditions}})
|
||||
|
||||
pipeline.append({"$limit": top_k})
|
||||
|
||||
results = list(collection.aggregate(pipeline))
|
||||
logger.info(f"Keyword search completed. Found {len(results)} documents.")
|
||||
|
||||
output = [
|
||||
OutputData(id=str(doc["_id"]), score=doc.get("score"), payload=doc.get("payload"))
|
||||
for doc in results
|
||||
]
|
||||
return output
|
||||
except Exception as e:
|
||||
logger.error(f"Error during keyword search for query '{query}': {e}")
|
||||
return None
|
||||
|
||||
def delete(self, vector_id: str) -> None:
|
||||
"""
|
||||
Delete a vector by ID.
|
||||
|
||||
@@ -203,6 +203,57 @@ class OpenSearchDB(VectorStoreBase):
|
||||
logger.error(f"Error during search: {e}", exc_info=True)
|
||||
return []
|
||||
|
||||
def keyword_search(self, query, top_k=5, filters=None):
|
||||
"""Search for memories using BM25 keyword matching.
|
||||
|
||||
Args:
|
||||
query (str): The text query to search for.
|
||||
top_k (int): Maximum number of results to return. Defaults to 5.
|
||||
filters (Dict, optional): Filters to apply to the search.
|
||||
|
||||
Returns:
|
||||
List[OutputData]: Search results with id, score, and payload.
|
||||
"""
|
||||
# Build a multi_match query across text fields in payload
|
||||
should_clauses = [
|
||||
{"match": {"payload.data": query}},
|
||||
{"match": {"payload.text_lemmatized": query}},
|
||||
]
|
||||
|
||||
bool_query = {
|
||||
"should": should_clauses,
|
||||
"minimum_should_match": 1,
|
||||
}
|
||||
|
||||
# Apply filters consistently with the existing search() method
|
||||
filter_clauses = []
|
||||
if filters:
|
||||
for key in ["user_id", "run_id", "agent_id"]:
|
||||
value = filters.get(key)
|
||||
if value:
|
||||
filter_clauses.append({"term": {f"payload.{key}.keyword": value}})
|
||||
|
||||
if filter_clauses:
|
||||
bool_query["filter"] = filter_clauses
|
||||
|
||||
query_body = {
|
||||
"size": top_k,
|
||||
"query": {"bool": bool_query},
|
||||
}
|
||||
|
||||
try:
|
||||
response = self.client.search(index=self.collection_name, body=query_body)
|
||||
|
||||
hits = response["hits"]["hits"]
|
||||
results = [
|
||||
OutputData(id=hit["_source"].get("id"), score=hit["_score"], payload=hit["_source"].get("payload", {}))
|
||||
for hit in hits[:top_k]
|
||||
]
|
||||
return results
|
||||
except Exception as e:
|
||||
logger.error(f"Error during keyword search: {e}")
|
||||
return []
|
||||
|
||||
def delete(self, vector_id: str) -> None:
|
||||
"""Delete a vector by custom ID."""
|
||||
# First, find the document by custom ID
|
||||
|
||||
@@ -179,6 +179,13 @@ class PGVector(VectorStoreBase):
|
||||
USING hnsw (vector vector_cosine_ops)
|
||||
"""
|
||||
)
|
||||
cur.execute(
|
||||
f"""
|
||||
CREATE INDEX IF NOT EXISTS {self.collection_name}_text_lemmatized_idx
|
||||
ON {self.collection_name}
|
||||
USING gin(to_tsvector('simple', payload->>'text_lemmatized'));
|
||||
"""
|
||||
)
|
||||
|
||||
def insert(self, vectors: list[list[float]], payloads=None, ids=None) -> None:
|
||||
logger.info(f"Inserting {len(vectors)} vectors into collection {self.collection_name}")
|
||||
@@ -243,6 +250,50 @@ class PGVector(VectorStoreBase):
|
||||
results = cur.fetchall()
|
||||
return [OutputData(id=str(r[0]), score=float(r[1]), payload=r[2]) for r in results]
|
||||
|
||||
def keyword_search(self, query, top_k=5, filters=None):
|
||||
"""
|
||||
Search using PostgreSQL full-text search on lemmatized text.
|
||||
|
||||
Args:
|
||||
query (str): The search query text.
|
||||
top_k (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (dict, optional): Filters to apply to the search. Defaults to None.
|
||||
|
||||
Returns:
|
||||
List[OutputData]: Search results ranked by text relevance.
|
||||
"""
|
||||
filter_conditions = []
|
||||
filter_params = []
|
||||
|
||||
if filters:
|
||||
for k, v in filters.items():
|
||||
filter_conditions.append("payload->>%s = %s")
|
||||
filter_params.extend([k, str(v)])
|
||||
|
||||
filter_clause = ""
|
||||
if filter_conditions:
|
||||
filter_clause = "AND " + " AND ".join(filter_conditions)
|
||||
|
||||
try:
|
||||
with self._get_cursor() as cur:
|
||||
cur.execute(
|
||||
f"""
|
||||
SELECT id, ts_rank_cd(to_tsvector('simple', payload->>'text_lemmatized'), plainto_tsquery('simple', %s)) AS score, payload
|
||||
FROM {self.collection_name}
|
||||
WHERE to_tsvector('simple', payload->>'text_lemmatized') @@ plainto_tsquery('simple', %s)
|
||||
{filter_clause}
|
||||
ORDER BY score DESC
|
||||
LIMIT %s
|
||||
""",
|
||||
(query, query, *filter_params, top_k),
|
||||
)
|
||||
|
||||
results = cur.fetchall()
|
||||
return [OutputData(id=str(r[0]), score=float(r[1]), payload=r[2]) for r in results]
|
||||
except Exception as e:
|
||||
logger.debug(f"Keyword search failed: {e}")
|
||||
return None
|
||||
|
||||
def delete(self, vector_id: str) -> None:
|
||||
"""
|
||||
Delete a vector by ID.
|
||||
|
||||
@@ -241,6 +241,42 @@ class PineconeDB(VectorStoreBase):
|
||||
results = self._parse_output(response.matches)
|
||||
return results
|
||||
|
||||
def keyword_search(self, query, top_k=5, filters=None):
|
||||
"""
|
||||
Search using BM25 sparse vectors for keyword-based retrieval.
|
||||
|
||||
Args:
|
||||
query (str): The search query text.
|
||||
top_k (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (dict, optional): Filters to apply to the search. Defaults to None.
|
||||
|
||||
Returns:
|
||||
List[OutputData]: Search results, or None if hybrid search is not configured.
|
||||
"""
|
||||
if not self.hybrid_search or self.sparse_encoder is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
filter_dict = self._create_filter(filters) if filters else None
|
||||
|
||||
sparse_vector = self.sparse_encoder.encode_queries(query)
|
||||
|
||||
query_params = {
|
||||
"sparse_vector": sparse_vector,
|
||||
"top_k": top_k,
|
||||
"include_metadata": True,
|
||||
"include_values": False,
|
||||
}
|
||||
|
||||
if filter_dict:
|
||||
query_params["filter"] = filter_dict
|
||||
|
||||
response = self.index.query(**query_params, namespace=self.namespace)
|
||||
|
||||
return self._parse_output(response.matches)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def delete(self, vector_id: Union[str, int]):
|
||||
"""
|
||||
Delete a vector by ID.
|
||||
|
||||
+121
-15
@@ -2,7 +2,7 @@ import logging
|
||||
import re
|
||||
from typing import Optional
|
||||
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client import QdrantClient, models
|
||||
from qdrant_client.models import (
|
||||
DatetimeRange,
|
||||
Distance,
|
||||
@@ -16,6 +16,8 @@ from qdrant_client.models import (
|
||||
PointStruct,
|
||||
PointVectors,
|
||||
Range,
|
||||
SparseVector,
|
||||
SparseVectorParams,
|
||||
VectorParams,
|
||||
)
|
||||
|
||||
@@ -64,7 +66,7 @@ class Qdrant(VectorStoreBase):
|
||||
if host and port:
|
||||
params["host"] = host
|
||||
params["port"] = port
|
||||
|
||||
|
||||
if not params:
|
||||
params["path"] = path
|
||||
self.is_local = True
|
||||
@@ -76,11 +78,44 @@ class Qdrant(VectorStoreBase):
|
||||
self.collection_name = collection_name
|
||||
self.embedding_model_dims = embedding_model_dims
|
||||
self.on_disk = on_disk
|
||||
self._bm25_encoder = None
|
||||
self.create_col(embedding_model_dims, on_disk)
|
||||
|
||||
def _get_bm25_encoder(self):
|
||||
"""Lazy-load the BM25 sparse text encoder (fastembed)."""
|
||||
if self._bm25_encoder is None:
|
||||
try:
|
||||
from fastembed import SparseTextEmbedding
|
||||
self._bm25_encoder = SparseTextEmbedding(model_name="Qdrant/bm25")
|
||||
logger.info("BM25 encoder loaded (fastembed Qdrant/bm25)")
|
||||
except ImportError:
|
||||
logger.warning("fastembed not installed — BM25 keyword search disabled. Install with: pip install fastembed")
|
||||
self._bm25_encoder = False # sentinel: tried and failed
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load BM25 encoder: {e}")
|
||||
self._bm25_encoder = False
|
||||
return self._bm25_encoder if self._bm25_encoder is not False else None
|
||||
|
||||
def _encode_bm25(self, text: str) -> SparseVector | None:
|
||||
"""Encode text into a BM25 sparse vector."""
|
||||
encoder = self._get_bm25_encoder()
|
||||
if encoder is None:
|
||||
return None
|
||||
try:
|
||||
results = list(encoder.embed([text]))
|
||||
if results:
|
||||
sparse = results[0]
|
||||
return SparseVector(
|
||||
indices=sparse.indices.tolist(),
|
||||
values=sparse.values.tolist(),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(f"BM25 encoding failed: {e}")
|
||||
return None
|
||||
|
||||
def create_col(self, vector_size: int, on_disk: bool, distance: Distance = Distance.COSINE):
|
||||
"""
|
||||
Create a new collection.
|
||||
Create a new collection with dense vectors and BM25 sparse vectors.
|
||||
|
||||
Args:
|
||||
vector_size (int): Size of the vectors to be stored.
|
||||
@@ -98,6 +133,11 @@ class Qdrant(VectorStoreBase):
|
||||
self.client.create_collection(
|
||||
collection_name=self.collection_name,
|
||||
vectors_config=VectorParams(size=vector_size, distance=distance, on_disk=on_disk),
|
||||
sparse_vectors_config={
|
||||
"bm25": SparseVectorParams(
|
||||
modifier=models.Modifier.IDF,
|
||||
),
|
||||
},
|
||||
)
|
||||
self._create_filter_indexes()
|
||||
|
||||
@@ -107,9 +147,9 @@ class Qdrant(VectorStoreBase):
|
||||
if self.is_local:
|
||||
logger.debug("Skipping payload index creation for local Qdrant (not supported)")
|
||||
return
|
||||
|
||||
|
||||
common_fields = ["user_id", "agent_id", "run_id", "actor_id"]
|
||||
|
||||
|
||||
for field in common_fields:
|
||||
try:
|
||||
self.client.create_payload_index(
|
||||
@@ -123,7 +163,8 @@ class Qdrant(VectorStoreBase):
|
||||
|
||||
def insert(self, vectors: list, payloads: list = None, ids: list = None):
|
||||
"""
|
||||
Insert vectors into a collection.
|
||||
Insert vectors into a collection, including BM25 sparse vectors
|
||||
computed from the text_lemmatized payload field.
|
||||
|
||||
Args:
|
||||
vectors (list): List of vectors to insert.
|
||||
@@ -131,14 +172,21 @@ class Qdrant(VectorStoreBase):
|
||||
ids (list, optional): List of IDs corresponding to vectors. Defaults to None.
|
||||
"""
|
||||
logger.info(f"Inserting {len(vectors)} vectors into collection {self.collection_name}")
|
||||
points = [
|
||||
PointStruct(
|
||||
id=idx if ids is None else ids[idx],
|
||||
vector=vector,
|
||||
payload=payloads[idx] if payloads else {},
|
||||
)
|
||||
for idx, vector in enumerate(vectors)
|
||||
]
|
||||
points = []
|
||||
for idx, vector in enumerate(vectors):
|
||||
payload = payloads[idx] if payloads else {}
|
||||
point_id = idx if ids is None else ids[idx]
|
||||
|
||||
# Build named vectors: dense + optional BM25 sparse
|
||||
named_vectors = {"": vector}
|
||||
text_for_bm25 = payload.get("text_lemmatized") or payload.get("data", "")
|
||||
if text_for_bm25:
|
||||
sparse = self._encode_bm25(text_for_bm25)
|
||||
if sparse is not None:
|
||||
named_vectors["bm25"] = sparse
|
||||
|
||||
points.append(PointStruct(id=point_id, vector=named_vectors, payload=payload))
|
||||
|
||||
self.client.upsert(collection_name=self.collection_name, points=points)
|
||||
|
||||
# ISO 8601 datetime pattern for detecting datetime strings in range filters
|
||||
@@ -330,6 +378,53 @@ class Qdrant(VectorStoreBase):
|
||||
)
|
||||
return hits.points
|
||||
|
||||
def search_batch(self, queries: list, vectors_list: list, top_k: int = 1, filters: dict = None):
|
||||
"""Batch search using Qdrant's query_batch_points for efficiency."""
|
||||
query_filter = self._create_filter(filters) if filters else None
|
||||
requests = [
|
||||
models.QueryRequest(query=vec, filter=query_filter, limit=top_k)
|
||||
for vec in vectors_list
|
||||
]
|
||||
try:
|
||||
results = self.client.query_batch_points(
|
||||
collection_name=self.collection_name,
|
||||
requests=requests,
|
||||
)
|
||||
return [r.points for r in results]
|
||||
except Exception as e:
|
||||
logger.warning(f"Batch search failed, falling back to sequential: {e}")
|
||||
return [self.search(q, v, top_k=top_k, filters=filters) for q, v in zip(queries, vectors_list)]
|
||||
|
||||
def keyword_search(self, query, top_k=5, filters=None):
|
||||
"""
|
||||
Search using BM25 sparse vectors for keyword-based retrieval.
|
||||
|
||||
Args:
|
||||
query (str): The search query text.
|
||||
top_k (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (dict, optional): Filters to apply to the search. Defaults to None.
|
||||
|
||||
Returns:
|
||||
list: Search results, or None if BM25 is not available.
|
||||
"""
|
||||
sparse_query = self._encode_bm25(query)
|
||||
if sparse_query is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
query_filter = self._create_filter(filters) if filters else None
|
||||
hits = self.client.query_points(
|
||||
collection_name=self.collection_name,
|
||||
query=sparse_query,
|
||||
using="bm25",
|
||||
query_filter=query_filter,
|
||||
limit=top_k,
|
||||
)
|
||||
return hits.points
|
||||
except Exception as e:
|
||||
logger.debug(f"BM25 keyword search failed: {e}")
|
||||
return None
|
||||
|
||||
def delete(self, vector_id: int):
|
||||
"""
|
||||
Delete a vector by ID.
|
||||
@@ -354,9 +449,20 @@ class Qdrant(VectorStoreBase):
|
||||
payload (dict, optional): Updated payload. Defaults to None.
|
||||
"""
|
||||
if vector is not None and payload is not None:
|
||||
point = PointStruct(id=vector_id, vector=vector, payload=payload)
|
||||
# Full update: attach BM25 sparse vector alongside dense vector
|
||||
named_vectors = {"": vector}
|
||||
text_for_bm25 = payload.get("text_lemmatized") or payload.get("data", "")
|
||||
if text_for_bm25:
|
||||
sparse = self._encode_bm25(text_for_bm25)
|
||||
if sparse is not None:
|
||||
named_vectors["bm25"] = sparse
|
||||
point = PointStruct(id=vector_id, vector=named_vectors, payload=payload)
|
||||
self.client.upsert(collection_name=self.collection_name, points=[point])
|
||||
else:
|
||||
# Partial update: use Qdrant's dedicated endpoints.
|
||||
# Note: BM25 sparse vector cannot be refreshed via set_payload alone;
|
||||
# payload-only updates will leave any existing BM25 vector stale. In
|
||||
# practice v3 re-embeds on memory text change, so this is acceptable.
|
||||
if payload is not None:
|
||||
self.client.set_payload(
|
||||
collection_name=self.collection_name,
|
||||
|
||||
@@ -7,7 +7,7 @@ import numpy as np
|
||||
import redis
|
||||
from redis.commands.search.query import Query
|
||||
from redisvl.index import SearchIndex
|
||||
from redisvl.query import VectorQuery
|
||||
from redisvl.query import TextQuery, VectorQuery
|
||||
from redisvl.query.filter import Tag
|
||||
|
||||
from mem0.memory.utils import extract_json
|
||||
@@ -181,6 +181,60 @@ class RedisDB(VectorStoreBase):
|
||||
for result in results
|
||||
]
|
||||
|
||||
def keyword_search(self, query, top_k=5, filters=None):
|
||||
"""
|
||||
Search for memories using BM25 keyword search on the memory field.
|
||||
|
||||
Args:
|
||||
query (str): Search query text.
|
||||
top_k (int): Maximum number of results. Defaults to 5.
|
||||
filters (dict, optional): Filters to apply (user_id, agent_id, run_id).
|
||||
|
||||
Returns:
|
||||
List[MemoryResult]: Search results.
|
||||
"""
|
||||
filter_expression = None
|
||||
if filters:
|
||||
conditions = [Tag(key) == value for key, value in filters.items() if value is not None]
|
||||
if conditions:
|
||||
filter_expression = reduce(lambda x, y: x & y, conditions)
|
||||
|
||||
t = TextQuery(
|
||||
text=query,
|
||||
text_field_name="memory",
|
||||
return_fields=["memory_id", "hash", "agent_id", "run_id", "user_id", "memory", "metadata", "created_at"],
|
||||
filter_expression=filter_expression,
|
||||
num_results=top_k,
|
||||
)
|
||||
|
||||
results = self.index.query(t)
|
||||
|
||||
return [
|
||||
MemoryResult(
|
||||
id=result["memory_id"],
|
||||
score=result.get("text_score", 1.0),
|
||||
payload={
|
||||
"hash": result["hash"],
|
||||
"data": result["memory"],
|
||||
"created_at": datetime.fromtimestamp(
|
||||
int(result["created_at"]), tz=timezone.utc
|
||||
).isoformat(timespec="microseconds"),
|
||||
**(
|
||||
{
|
||||
"updated_at": datetime.fromtimestamp(
|
||||
int(result["updated_at"]), tz=timezone.utc
|
||||
).isoformat(timespec="microseconds")
|
||||
}
|
||||
if "updated_at" in result
|
||||
else {}
|
||||
),
|
||||
**{field: result[field] for field in ["agent_id", "run_id", "user_id"] if field in result},
|
||||
**{k: v for k, v in json.loads(extract_json(result["metadata"])).items()},
|
||||
},
|
||||
)
|
||||
for result in results
|
||||
]
|
||||
|
||||
def delete(self, vector_id):
|
||||
self.index.drop_keys(f"{self.schema['index']['prefix']}:{vector_id}")
|
||||
|
||||
|
||||
@@ -149,6 +149,45 @@ class UpstashVector(VectorStoreBase):
|
||||
for res in response
|
||||
]
|
||||
|
||||
def keyword_search(self, query, top_k=5, filters=None):
|
||||
"""
|
||||
Perform keyword-based search using Upstash's BM25 sparse search.
|
||||
|
||||
Args:
|
||||
query (str): The text query to search for.
|
||||
top_k (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (Dict, optional): Filters to apply to the search.
|
||||
|
||||
Returns:
|
||||
List[OutputData]: Search results, or None if sparse/BM25 search is not supported.
|
||||
"""
|
||||
try:
|
||||
filters_str = (
|
||||
" AND ".join([f"{k} = {self._stringify(v)}" for k, v in filters.items()])
|
||||
if filters
|
||||
else None
|
||||
)
|
||||
|
||||
response = self.client.query(
|
||||
data=query,
|
||||
top_k=top_k,
|
||||
filter=filters_str or "",
|
||||
include_metadata=True,
|
||||
namespace=self.collection_name,
|
||||
)
|
||||
|
||||
return [
|
||||
OutputData(
|
||||
id=res.id,
|
||||
score=res.score,
|
||||
payload=res.metadata,
|
||||
)
|
||||
for res in response
|
||||
]
|
||||
except Exception as e:
|
||||
logger.error(f"Error during keyword search for query '{query}': {e}")
|
||||
return None
|
||||
|
||||
def delete(self, vector_id: int):
|
||||
"""
|
||||
Delete a vector by ID.
|
||||
|
||||
@@ -276,6 +276,11 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
logger.error("Stack trace: %s", traceback.format_exc())
|
||||
raise
|
||||
|
||||
def keyword_search(self, query, top_k=5, filters=None):
|
||||
# Vertex AI hybrid search requires sparse embeddings configuration.
|
||||
# Not yet supported - requires HybridQuery with sparse encoder setup.
|
||||
return None
|
||||
|
||||
def delete(self, vector_id: Optional[str] = None, ids: Optional[List[str]] = None) -> bool:
|
||||
"""
|
||||
Delete vectors from the Matching Engine index.
|
||||
|
||||
@@ -223,6 +223,55 @@ class Weaviate(VectorStoreBase):
|
||||
)
|
||||
return results
|
||||
|
||||
def keyword_search(self, query, top_k=5, filters=None):
|
||||
"""
|
||||
Search for memories using BM25 keyword search.
|
||||
|
||||
Args:
|
||||
query (str): Search query text.
|
||||
top_k (int): Maximum number of results. Defaults to 5.
|
||||
filters (dict, optional): Filters to apply (user_id, agent_id, run_id).
|
||||
|
||||
Returns:
|
||||
List[OutputData]: Search results.
|
||||
"""
|
||||
collection = self.client.collections.get(str(self.collection_name))
|
||||
filter_conditions = []
|
||||
if filters:
|
||||
for key, value in filters.items():
|
||||
if value and key in ["user_id", "agent_id", "run_id"]:
|
||||
filter_conditions.append(Filter.by_property(key).equal(value))
|
||||
combined_filter = Filter.all_of(filter_conditions) if filter_conditions else None
|
||||
response = collection.query.bm25(
|
||||
query=query,
|
||||
query_properties=["data"],
|
||||
limit=top_k,
|
||||
filters=combined_filter,
|
||||
return_properties=["hash", "created_at", "updated_at", "user_id", "agent_id", "run_id", "data", "category"],
|
||||
return_metadata=MetadataQuery(score=True),
|
||||
)
|
||||
results = []
|
||||
for obj in response.objects:
|
||||
payload = obj.properties.copy()
|
||||
|
||||
for id_field in ["run_id", "agent_id", "user_id"]:
|
||||
if id_field in payload and payload[id_field] is None:
|
||||
del payload[id_field]
|
||||
|
||||
payload["id"] = str(obj.uuid).split("'")[0]
|
||||
if obj.metadata.score is not None:
|
||||
score = obj.metadata.score
|
||||
else:
|
||||
score = 1.0
|
||||
results.append(
|
||||
OutputData(
|
||||
id=str(obj.uuid),
|
||||
score=score,
|
||||
payload=payload,
|
||||
)
|
||||
)
|
||||
return results
|
||||
|
||||
def delete(self, vector_id):
|
||||
"""
|
||||
Delete a vector by ID.
|
||||
|
||||
Reference in New Issue
Block a user