klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Formatting (#2750)

Browse Source
This commit is contained in:
Dev Khant
2025-05-22 01:17:29 +05:30
committed by GitHub
parent dff91154a7
commit d85fcda037
71 changed files with 1391 additions and 1823 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
mem0/vector_stores/azure_ai_search.py
View File

@@ -377,4 +377,3 @@ class AzureAISearch(VectorStoreBase):
except Exception as e:
logger.error(f"Error resetting index {self.index_name}: {e}")
raise

2
mem0/vector_stores/base.py
View File

@@ -51,7 +51,7 @@ class VectorStoreBase(ABC):
def list(self, filters=None, limit=None):
"""List all memories."""
pass
@abstractmethod
def reset(self):
"""Reset by delete the collection and recreate it."""

2
mem0/vector_stores/chroma.py
View File

@@ -221,7 +221,7 @@ class ChromaDB(VectorStoreBase):
"""
results = self.collection.get(where=filters, limit=limit)
return [self._parse_output(results)]
def reset(self):
"""Reset the index by deleting and recreating it."""
logger.warning(f"Resetting index {self.collection_name}...")

9
mem0/vector_stores/elasticsearch.py
View File

@@ -58,7 +58,12 @@ class ElasticsearchDB(VectorStoreBase):
"mappings": {
"properties": {
"text": {"type": "text"},
"vector": {"type": "dense_vector", "dims": self.embedding_model_dims, "index": True, "similarity": "cosine"},
"vector": {
"type": "dense_vector",
"dims": self.embedding_model_dims,
"index": True,
"similarity": "cosine",
},
"metadata": {"type": "object", "properties": {"user_id": {"type": "keyword"}}},
}
},
@@ -222,7 +227,7 @@ class ElasticsearchDB(VectorStoreBase):
)
return [results]
def reset(self):
"""Reset the index by deleting and recreating it."""
logger.warning(f"Resetting index {self.collection_name}...")

2
mem0/vector_stores/faiss.py
View File

@@ -465,7 +465,7 @@ class FAISS(VectorStoreBase):
break
return [results]
def reset(self):
"""Reset the index by deleting and recreating it."""
logger.warning(f"Resetting index {self.collection_name}...")

6
mem0/vector_stores/langchain.py
View File

@@ -14,6 +14,7 @@ from mem0.vector_stores.base import VectorStoreBase
logger = logging.getLogger(__name__)
class OutputData(BaseModel):
id: Optional[str] # memory id
score: Optional[float] # distance
@@ -162,10 +163,7 @@ class Langchain(VectorStoreBase):
if filters and "user_id" in filters:
where_clause = {"user_id": filters["user_id"]}
result = self.client._collection.get(
where=where_clause,
limit=limit
)
result = self.client._collection.get(where=where_clause, limit=limit)
# Convert the result to the expected format
if result and isinstance(result, dict):

2
mem0/vector_stores/milvus.py
View File

@@ -237,7 +237,7 @@ class MilvusDB(VectorStoreBase):
obj = OutputData(id=data.get("id"), score=None, payload=data.get("metadata"))
memories.append(obj)
return [memories]
def reset(self):
"""Reset the index by deleting and recreating it."""
logger.warning(f"Resetting index {self.collection_name}...")

100
mem0/vector_stores/opensearch.py
View File

@@ -1,6 +1,6 @@
import logging
from typing import Any, Dict, List, Optional
import time
from typing import Any, Dict, List, Optional
try:
from opensearchpy import OpenSearch, RequestsHttpConnection
@@ -34,7 +34,7 @@ class OpenSearchDB(VectorStoreBase):
use_ssl=config.use_ssl,
verify_certs=config.verify_certs,
connection_class=RequestsHttpConnection,
pool_maxsize=20
pool_maxsize=20,
)
self.collection_name = config.collection_name
@@ -69,9 +69,7 @@ class OpenSearchDB(VectorStoreBase):
def create_col(self, name: str, vector_size: int) -> None:
"""Create a new collection (index in OpenSearch)."""
index_settings = {
"settings": {
"index.knn": True
},
"settings": {"index.knn": True},
"mappings": {
"properties": {
"vector_field": {
@@ -82,7 +80,7 @@ class OpenSearchDB(VectorStoreBase):
"payload": {"type": "object"},
"id": {"type": "keyword"},
}
}
},
}
if not self.client.indices.exists(index=name):
@@ -102,9 +100,7 @@ class OpenSearchDB(VectorStoreBase):
except Exception:
retry_count += 1
if retry_count == max_retries:
raise TimeoutError(
f"Index {name} creation timed out after {max_retries} seconds"
)
raise TimeoutError(f"Index {name} creation timed out after {max_retries} seconds")
time.sleep(0.5)
def insert(
@@ -145,10 +141,7 @@ class OpenSearchDB(VectorStoreBase):
}
# Start building the full query
query_body = {
"size": limit * 2,
"query": None
}
query_body = {"size": limit * 2, "query": None}
# Prepare filter conditions if applicable
filter_clauses = []
@@ -156,18 +149,11 @@ class OpenSearchDB(VectorStoreBase):
for key in ["user_id", "run_id", "agent_id"]:
value = filters.get(key)
if value:
filter_clauses.append({
"term": {f"payload.{key}.keyword": value}
})
filter_clauses.append({"term": {f"payload.{key}.keyword": value}})
# Combine knn with filters if needed
if filter_clauses:
query_body["query"] = {
"bool": {
"must": knn_query,
"filter": filter_clauses
}
}
query_body["query"] = {"bool": {"must": knn_query, "filter": filter_clauses}}
else:
query_body["query"] = knn_query
@@ -176,11 +162,7 @@ class OpenSearchDB(VectorStoreBase):
hits = response["hits"]["hits"]
results = [
OutputData(
id=hit["_source"].get("id"),
score=hit["_score"],
payload=hit["_source"].get("payload", {})
)
OutputData(id=hit["_source"].get("id"), score=hit["_score"], payload=hit["_source"].get("payload", {}))
for hit in hits
]
return results
@@ -188,13 +170,7 @@ class OpenSearchDB(VectorStoreBase):
def delete(self, vector_id: str) -> None:
"""Delete a vector by custom ID."""
# First, find the document by custom ID
search_query = {
"query": {
"term": {
"id": vector_id
}
}
}
search_query = {"query": {"term": {"id": vector_id}}}
response = self.client.search(index=self.collection_name, body=search_query)
hits = response.get("hits", {}).get("hits", [])
@@ -207,18 +183,11 @@ class OpenSearchDB(VectorStoreBase):
# Delete using the actual document ID
self.client.delete(index=self.collection_name, id=opensearch_id)
def update(self, vector_id: str, vector: Optional[List[float]] = None, payload: Optional[Dict] = None) -> None:
"""Update a vector and its payload using the custom 'id' field."""
# First, find the document by custom ID
search_query = {
"query": {
"term": {
"id": vector_id
}
}
}
search_query = {"query": {"term": {"id": vector_id}}}
response = self.client.search(index=self.collection_name, body=search_query)
hits = response.get("hits", {}).get("hits", [])
@@ -241,7 +210,6 @@ class OpenSearchDB(VectorStoreBase):
except Exception:
pass
def get(self, vector_id: str) -> Optional[OutputData]:
"""Retrieve a vector by ID."""
try:
@@ -251,13 +219,7 @@ class OpenSearchDB(VectorStoreBase):
self.create_col(self.collection_name, self.embedding_model_dims)
return None
search_query = {
"query": {
"term": {
"id": vector_id
}
}
}
search_query = {"query": {"term": {"id": vector_id}}}
response = self.client.search(index=self.collection_name, body=search_query)
hits = response["hits"]["hits"]
@@ -265,11 +227,7 @@ class OpenSearchDB(VectorStoreBase):
if not hits:
return None
return OutputData(
id=hits[0]["_source"].get("id"),
score=1.0,
payload=hits[0]["_source"].get("payload", {})
)
return OutputData(id=hits[0]["_source"].get("id"), score=1.0, payload=hits[0]["_source"].get("payload", {}))
except Exception as e:
logger.error(f"Error retrieving vector {vector_id}: {str(e)}")
return None
@@ -287,30 +245,19 @@ class OpenSearchDB(VectorStoreBase):
return self.client.indices.get(index=name)
def list(self, filters: Optional[Dict] = None, limit: Optional[int] = None) -> List[OutputData]:
try:
"""List all memories with optional filters."""
query: Dict = {
"query": {
"match_all": {}
}
}
query: Dict = {"query": {"match_all": {}}}
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}
})
filter_clauses.append({"term": {f"payload.{key}.keyword": value}})
if filter_clauses:
query["query"] = {
"bool": {
"filter": filter_clauses
}
}
query["query"] = {"bool": {"filter": filter_clauses}}
if limit:
query["size"] = limit
@@ -318,18 +265,15 @@ class OpenSearchDB(VectorStoreBase):
response = self.client.search(index=self.collection_name, body=query)
hits = response["hits"]["hits"]
return [[
OutputData(
id=hit["_source"].get("id"),
score=1.0,
payload=hit["_source"].get("payload", {})
)
for hit in hits
]]
return [
[
OutputData(id=hit["_source"].get("id"), score=1.0, payload=hit["_source"].get("payload", {}))
for hit in hits
]
]
except Exception:
return []
def reset(self):
"""Reset the index by deleting and recreating it."""
logger.warning(f"Resetting index {self.collection_name}...")

2
mem0/vector_stores/pgvector.py
View File

@@ -286,7 +286,7 @@ class PGVector(VectorStoreBase):
self.cur.close()
if hasattr(self, "conn"):
self.conn.close()
def reset(self):
"""Reset the index by deleting and recreating it."""
logger.warning(f"Resetting index {self.collection_name}...")

2
mem0/vector_stores/qdrant.py
View File

@@ -232,7 +232,7 @@ class Qdrant(VectorStoreBase):
with_vectors=False,
)
return result
def reset(self):
"""Reset the index by deleting and recreating it."""
logger.warning(f"Resetting index {self.collection_name}...")

11
mem0/vector_stores/redis.py
View File

@@ -88,7 +88,7 @@ class RedisDB(VectorStoreBase):
The created index object.
"""
# Use provided parameters or fall back to instance attributes
collection_name = name or self.schema['index']['name']
collection_name = name or self.schema["index"]["name"]
embedding_dims = vector_size or self.embedding_model_dims
distance_metric = distance or "cosine"
@@ -237,17 +237,16 @@ class RedisDB(VectorStoreBase):
"""
Reset the index by deleting and recreating it.
"""
collection_name = self.schema['index']['name']
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)
# 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)

2
mem0/vector_stores/supabase.py
View File

@@ -229,7 +229,7 @@ class Supabase(VectorStoreBase):
records = self.collection.fetch(ids=ids)
return [[OutputData(id=str(record[0]), score=None, payload=record[2]) for record in records]]
def reset(self):
"""Reset the index by deleting and recreating it."""
logger.warning(f"Resetting index {self.collection_name}...")

3
mem0/vector_stores/upstash_vector.py
View File

@@ -285,10 +285,9 @@ class UpstashVector(VectorStoreBase):
- Per-namespace vector and pending vector counts
"""
return self.client.info()
def reset(self):
"""
Reset the Upstash Vector index.
"""
self.delete_col()

2
mem0/vector_stores/weaviate.py
View File

@@ -308,7 +308,7 @@ class Weaviate(VectorStoreBase):
payload["id"] = str(obj.uuid).split("'")[0]
results.append(OutputData(id=str(obj.uuid).split("'")[0], score=1.0, payload=payload))
return [results]
def reset(self):
"""Reset the index by deleting and recreating it."""
logger.warning(f"Resetting index {self.collection_name}...")