Soumil RathiandClaude Opus 4.6
3ae22be8a2
- Remove unused imports: get_update_memory_messages, get_fact_retrieval_messages
- Fix pytz.timezone("US/Pacific") → timezone.utc in redis vector store
- Remove unused entity_types variable in test_entity_extraction
- Regenerate pnpm-lock.yaml after neo4j-driver removal
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
352 lines
13 KiB
Python
352 lines
13 KiB
Python
import json
|
|
import logging
|
|
from datetime import datetime, timezone
|
|
from functools import reduce
|
|
|
|
import numpy as np
|
|
import redis
|
|
from redis.commands.search.query import Query
|
|
from redisvl.index import SearchIndex
|
|
from redisvl.query import TextQuery, VectorQuery
|
|
from redisvl.query.filter import Tag
|
|
|
|
from mem0.memory.utils import extract_json
|
|
from mem0.vector_stores.base import VectorStoreBase
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# TODO: Improve as these are not the best fields for the Redis's perspective. Might do away with them.
|
|
DEFAULT_FIELDS = [
|
|
{"name": "memory_id", "type": "tag"},
|
|
{"name": "hash", "type": "tag"},
|
|
{"name": "agent_id", "type": "tag"},
|
|
{"name": "run_id", "type": "tag"},
|
|
{"name": "user_id", "type": "tag"},
|
|
{"name": "memory", "type": "text"},
|
|
{"name": "metadata", "type": "text"},
|
|
# TODO: Although it is numeric but also accepts string
|
|
{"name": "created_at", "type": "numeric"},
|
|
{"name": "updated_at", "type": "numeric"},
|
|
{
|
|
"name": "embedding",
|
|
"type": "vector",
|
|
"attrs": {"distance_metric": "cosine", "algorithm": "flat", "datatype": "float32"},
|
|
},
|
|
]
|
|
|
|
excluded_keys = {"user_id", "agent_id", "run_id", "hash", "data", "created_at", "updated_at"}
|
|
|
|
|
|
class MemoryResult:
|
|
def __init__(self, id: str, payload: dict, score: float = None):
|
|
self.id = id
|
|
self.payload = payload
|
|
self.score = score
|
|
|
|
|
|
class RedisDB(VectorStoreBase):
|
|
def __init__(
|
|
self,
|
|
redis_url: str,
|
|
collection_name: str,
|
|
embedding_model_dims: int,
|
|
):
|
|
"""
|
|
Initialize the Redis vector store.
|
|
|
|
Args:
|
|
redis_url (str): Redis URL.
|
|
collection_name (str): Collection name.
|
|
embedding_model_dims (int): Embedding model dimensions.
|
|
"""
|
|
self.embedding_model_dims = embedding_model_dims
|
|
index_schema = {
|
|
"name": collection_name,
|
|
"prefix": f"mem0:{collection_name}",
|
|
}
|
|
|
|
fields = DEFAULT_FIELDS.copy()
|
|
fields[-1]["attrs"]["dims"] = embedding_model_dims
|
|
|
|
self.schema = {"index": index_schema, "fields": fields}
|
|
|
|
self.client = redis.Redis.from_url(redis_url)
|
|
self.index = SearchIndex.from_dict(self.schema)
|
|
self.index.set_client(self.client)
|
|
self.index.create(overwrite=True)
|
|
|
|
def create_col(self, name=None, vector_size=None, distance=None):
|
|
"""
|
|
Create a new collection (index) in Redis.
|
|
|
|
Args:
|
|
name (str, optional): Name for the collection. Defaults to None, which uses the current collection_name.
|
|
vector_size (int, optional): Size of the vector embeddings. Defaults to None, which uses the current embedding_model_dims.
|
|
distance (str, optional): Distance metric to use. Defaults to None, which uses 'cosine'.
|
|
|
|
Returns:
|
|
The created index object.
|
|
"""
|
|
# Use provided parameters or fall back to instance attributes
|
|
collection_name = name or self.schema["index"]["name"]
|
|
embedding_dims = vector_size or self.embedding_model_dims
|
|
distance_metric = distance or "cosine"
|
|
|
|
# Create a new schema with the specified parameters
|
|
index_schema = {
|
|
"name": collection_name,
|
|
"prefix": f"mem0:{collection_name}",
|
|
}
|
|
|
|
# Copy the default fields and update the vector field with the specified dimensions
|
|
fields = DEFAULT_FIELDS.copy()
|
|
fields[-1]["attrs"]["dims"] = embedding_dims
|
|
fields[-1]["attrs"]["distance_metric"] = distance_metric
|
|
|
|
# Create the schema
|
|
schema = {"index": index_schema, "fields": fields}
|
|
|
|
# Create the index
|
|
index = SearchIndex.from_dict(schema)
|
|
index.set_client(self.client)
|
|
index.create(overwrite=True)
|
|
|
|
# Update instance attributes if creating a new collection
|
|
if name:
|
|
self.schema = schema
|
|
self.index = index
|
|
|
|
return index
|
|
|
|
def insert(self, vectors: list, payloads: list = None, ids: list = None):
|
|
data = []
|
|
for vector, payload, id in zip(vectors, payloads, ids):
|
|
# Start with required fields
|
|
entry = {
|
|
"memory_id": id,
|
|
"hash": payload["hash"],
|
|
"memory": payload["data"],
|
|
"created_at": int(datetime.fromisoformat(payload["created_at"]).timestamp()),
|
|
"embedding": np.array(vector, dtype=np.float32).tobytes(),
|
|
}
|
|
|
|
# Conditionally add optional fields
|
|
for field in ["agent_id", "run_id", "user_id"]:
|
|
if field in payload:
|
|
entry[field] = payload[field]
|
|
|
|
# Add metadata excluding specific keys
|
|
entry["metadata"] = json.dumps({k: v for k, v in payload.items() if k not in excluded_keys})
|
|
|
|
data.append(entry)
|
|
self.index.load(data, id_field="memory_id")
|
|
|
|
def search(self, query: str, vectors: list, top_k: int = 5, filters: dict = None):
|
|
conditions = [Tag(key) == value for key, value in filters.items() if value is not None]
|
|
filter = reduce(lambda x, y: x & y, conditions)
|
|
|
|
v = VectorQuery(
|
|
vector=np.array(vectors, dtype=np.float32).tobytes(),
|
|
vector_field_name="embedding",
|
|
return_fields=["memory_id", "hash", "agent_id", "run_id", "user_id", "memory", "metadata", "created_at"],
|
|
filter_expression=filter,
|
|
num_results=top_k,
|
|
)
|
|
|
|
results = self.index.query(v)
|
|
|
|
return [
|
|
MemoryResult(
|
|
id=result["memory_id"],
|
|
score=float(result["vector_distance"]),
|
|
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 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}")
|
|
|
|
def update(self, vector_id=None, vector=None, payload=None):
|
|
data = {
|
|
"memory_id": vector_id,
|
|
"hash": payload["hash"],
|
|
"memory": payload["data"],
|
|
"created_at": int(datetime.fromisoformat(payload["created_at"]).timestamp()),
|
|
"updated_at": int(datetime.fromisoformat(payload["updated_at"]).timestamp()),
|
|
}
|
|
|
|
# Only update embedding if vector is provided
|
|
if vector is not None:
|
|
data["embedding"] = np.array(vector, dtype=np.float32).tobytes()
|
|
|
|
for field in ["agent_id", "run_id", "user_id"]:
|
|
if field in payload:
|
|
data[field] = payload[field]
|
|
|
|
data["metadata"] = json.dumps({k: v for k, v in payload.items() if k not in excluded_keys})
|
|
self.index.load(data=[data], keys=[f"{self.schema['index']['prefix']}:{vector_id}"], id_field="memory_id")
|
|
|
|
def get(self, vector_id):
|
|
result = self.index.fetch(vector_id)
|
|
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()},
|
|
}
|
|
|
|
return MemoryResult(id=result["memory_id"], payload=payload)
|
|
|
|
def list_cols(self):
|
|
return self.index.listall()
|
|
|
|
def delete_col(self):
|
|
self.index.delete()
|
|
|
|
def col_info(self, name):
|
|
return self.index.info()
|
|
|
|
def reset(self):
|
|
"""
|
|
Reset the index by deleting and recreating it.
|
|
"""
|
|
collection_name = self.schema["index"]["name"]
|
|
logger.warning(f"Resetting index {collection_name}...")
|
|
self.delete_col()
|
|
|
|
self.index = SearchIndex.from_dict(self.schema)
|
|
self.index.set_client(self.client)
|
|
self.index.create(overwrite=True)
|
|
|
|
# or use
|
|
# self.create_col(collection_name, self.embedding_model_dims)
|
|
|
|
# Recreate the index with the same parameters
|
|
self.create_col(collection_name, self.embedding_model_dims)
|
|
|
|
def list(self, filters: dict = None, top_k: int = None) -> list:
|
|
"""
|
|
List all recent created memories from the vector store.
|
|
"""
|
|
conditions = [Tag(key) == value for key, value in filters.items() if value is not None]
|
|
filter = reduce(lambda x, y: x & y, conditions)
|
|
query = Query(str(filter)).sort_by("created_at", asc=False)
|
|
if top_k is not None:
|
|
query = Query(str(filter)).sort_by("created_at", asc=False).paging(0, top_k)
|
|
|
|
results = self.index.search(query)
|
|
return [
|
|
[
|
|
MemoryResult(
|
|
id=result["memory_id"],
|
|
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 result.__dict__.get("updated_at")
|
|
else {}
|
|
),
|
|
**{
|
|
field: result[field]
|
|
for field in ["agent_id", "run_id", "user_id"]
|
|
if field in result.__dict__
|
|
},
|
|
**{k: v for k, v in json.loads(extract_json(result["metadata"])).items()},
|
|
},
|
|
)
|
|
for result in results.docs
|
|
]
|
|
]
|