Code formatting (#1986)

This commit is contained in:
Dev Khant
2024-10-29 11:32:07 +05:30
committed by GitHub
parent dca74a1ec0
commit 605558da9d
13 changed files with 119 additions and 149 deletions

View File

@@ -117,7 +117,9 @@ class MemoryClient:
payload = self._prepare_payload(messages, kwargs)
response = self.client.post("/v1/memories/", json=payload)
response.raise_for_status()
capture_client_event("client.add", self)
if "metadata" in kwargs:
del kwargs["metadata"]
capture_client_event("client.add", self, {"keys": list(kwargs.keys())})
return response.json()
@api_error_handler
@@ -135,7 +137,7 @@ class MemoryClient:
"""
response = self.client.get(f"/v1/memories/{memory_id}/")
response.raise_for_status()
capture_client_event("client.get", self)
capture_client_event("client.get", self, {"memory_id": memory_id})
return response.json()
@api_error_handler
@@ -159,10 +161,12 @@ class MemoryClient:
elif version == "v2":
response = self.client.post(f"/{version}/memories/", json=params)
response.raise_for_status()
if "metadata" in kwargs:
del kwargs["metadata"]
capture_client_event(
"client.get_all",
self,
{"filters": len(params), "limit": kwargs.get("limit", 100)},
{"api_version": version, "keys": list(kwargs.keys())},
)
return response.json()
@@ -186,7 +190,9 @@ class MemoryClient:
payload.update({k: v for k, v in kwargs.items() if v is not None})
response = self.client.post(f"/{version}/memories/search/", json=payload)
response.raise_for_status()
capture_client_event("client.search", self, {"limit": kwargs.get("limit", 100)})
if "metadata" in kwargs:
del kwargs["metadata"]
capture_client_event("client.search", self, {"api_version": version, "keys": list(kwargs.keys())})
return response.json()
@api_error_handler
@@ -199,7 +205,7 @@ class MemoryClient:
Returns:
Dict[str, Any]: The response from the server.
"""
capture_client_event("client.update", self)
capture_client_event("client.update", self, {"memory_id": memory_id})
response = self.client.put(f"/v1/memories/{memory_id}/", json={"text": data})
response.raise_for_status()
return response.json()
@@ -219,7 +225,7 @@ class MemoryClient:
"""
response = self.client.delete(f"/v1/memories/{memory_id}/")
response.raise_for_status()
capture_client_event("client.delete", self)
capture_client_event("client.delete", self, {"memory_id": memory_id})
return response.json()
@api_error_handler
@@ -239,7 +245,7 @@ class MemoryClient:
params = self._prepare_params(kwargs)
response = self.client.delete("/v1/memories/", params=params)
response.raise_for_status()
capture_client_event("client.delete_all", self, {"params": len(params)})
capture_client_event("client.delete_all", self, {"keys": list(kwargs.keys())})
return response.json()
@api_error_handler
@@ -257,7 +263,7 @@ class MemoryClient:
"""
response = self.client.get(f"/v1/memories/{memory_id}/history/")
response.raise_for_status()
capture_client_event("client.history", self)
capture_client_event("client.history", self, {"memory_id": memory_id})
return response.json()
@api_error_handler
@@ -390,14 +396,16 @@ class AsyncMemoryClient:
payload = self.sync_client._prepare_payload(messages, kwargs)
response = await self.async_client.post("/v1/memories/", json=payload)
response.raise_for_status()
capture_client_event("async_client.add", self.sync_client)
if "metadata" in kwargs:
del kwargs["metadata"]
capture_client_event("async_client.add", self.sync_client, {"keys": list(kwargs.keys())})
return response.json()
@api_error_handler
async def get(self, memory_id: str) -> Dict[str, Any]:
response = await self.async_client.get(f"/v1/memories/{memory_id}/")
response.raise_for_status()
capture_client_event("async_client.get", self.sync_client)
capture_client_event("async_client.get", self.sync_client, {"memory_id": memory_id})
return response.json()
@api_error_handler
@@ -408,8 +416,10 @@ class AsyncMemoryClient:
elif version == "v2":
response = await self.async_client.post(f"/{version}/memories/", json=params)
response.raise_for_status()
if "metadata" in kwargs:
del kwargs["metadata"]
capture_client_event(
"async_client.get_all", self.sync_client, {"filters": len(params), "limit": kwargs.get("limit", 100)}
"async_client.get_all", self.sync_client, {"api_version": version, "keys": list(kwargs.keys())}
)
return response.json()
@@ -419,21 +429,25 @@ class AsyncMemoryClient:
payload.update(self.sync_client._prepare_params(kwargs))
response = await self.async_client.post(f"/{version}/memories/search/", json=payload)
response.raise_for_status()
capture_client_event("async_client.search", self.sync_client, {"limit": kwargs.get("limit", 100)})
if "metadata" in kwargs:
del kwargs["metadata"]
capture_client_event(
"async_client.search", self.sync_client, {"api_version": version, "keys": list(kwargs.keys())}
)
return response.json()
@api_error_handler
async def update(self, memory_id: str, data: str) -> Dict[str, Any]:
response = await self.async_client.put(f"/v1/memories/{memory_id}/", json={"text": data})
response.raise_for_status()
capture_client_event("async_client.update", self.sync_client)
capture_client_event("async_client.update", self.sync_client, {"memory_id": memory_id})
return response.json()
@api_error_handler
async def delete(self, memory_id: str) -> Dict[str, Any]:
response = await self.async_client.delete(f"/v1/memories/{memory_id}/")
response.raise_for_status()
capture_client_event("async_client.delete", self.sync_client)
capture_client_event("async_client.delete", self.sync_client, {"memory_id": memory_id})
return response.json()
@api_error_handler
@@ -441,14 +455,14 @@ class AsyncMemoryClient:
params = self.sync_client._prepare_params(kwargs)
response = await self.async_client.delete("/v1/memories/", params=params)
response.raise_for_status()
capture_client_event("async_client.delete_all", self.sync_client, {"params": len(params)})
capture_client_event("async_client.delete_all", self.sync_client, {"keys": list(kwargs.keys())})
return response.json()
@api_error_handler
async def history(self, memory_id: str) -> List[Dict[str, Any]]:
response = await self.async_client.get(f"/v1/memories/{memory_id}/history/")
response.raise_for_status()
capture_client_event("async_client.history", self.sync_client)
capture_client_event("async_client.history", self.sync_client, {"memory_id": memory_id})
return response.json()
@api_error_handler

View File

@@ -1,5 +1,3 @@
import subprocess
import sys
from typing import Any, ClassVar, Dict, Optional
from pydantic import BaseModel, Field, model_validator

View File

@@ -37,11 +37,11 @@ class Memory(MemoryBase):
self.llm = LlmFactory.create(self.config.llm.provider, self.config.llm.config)
self.db = SQLiteManager(self.config.history_db_path)
self.collection_name = self.config.vector_store.config.collection_name
self.version = self.config.version
self.api_version = self.config.version
self.enable_graph = False
if self.version == "v1.1" and self.config.graph_store.config:
if self.api_version == "v1.1" and self.config.graph_store.config:
from mem0.memory.graph_memory import MemoryGraph
self.graph = MemoryGraph(self.config)
@@ -119,7 +119,7 @@ class Memory(MemoryBase):
vector_store_result = future1.result()
graph_result = future2.result()
if self.version == "v1.1":
if self.api_version == "v1.1":
return {
"results": vector_store_result,
"relations": graph_result,
@@ -226,13 +226,13 @@ class Memory(MemoryBase):
except Exception as e:
logging.error(f"Error in new_memories_with_actions: {e}")
capture_event("mem0.add", self)
capture_event("mem0.add", self, {"version": self.api_version, "keys": list(filters.keys())})
return returned_memories
def _add_to_graph(self, messages, filters):
added_entities = []
if self.version == "v1.1" and self.enable_graph:
if self.api_version == "v1.1" and self.enable_graph:
if filters["user_id"]:
self.graph.user_id = filters["user_id"]
elif filters["agent_id"]:
@@ -305,13 +305,13 @@ class Memory(MemoryBase):
if run_id:
filters["run_id"] = run_id
capture_event("mem0.get_all", self, {"filters": len(filters), "limit": limit})
capture_event("mem0.get_all", self, {"limit": limit, "keys": list(filters.keys())})
with concurrent.futures.ThreadPoolExecutor() as executor:
future_memories = executor.submit(self._get_all_from_vector_store, filters, limit)
future_graph_entities = (
executor.submit(self.graph.get_all, filters, limit)
if self.version == "v1.1" and self.enable_graph
if self.api_version == "v1.1" and self.enable_graph
else None
)
@@ -322,7 +322,7 @@ class Memory(MemoryBase):
all_memories = future_memories.result()
graph_entities = future_graph_entities.result() if future_graph_entities else None
if self.version == "v1.1":
if self.api_version == "v1.1":
if self.enable_graph:
return {"results": all_memories, "relations": graph_entities}
else:
@@ -398,14 +398,14 @@ class Memory(MemoryBase):
capture_event(
"mem0.search",
self,
{"filters": len(filters), "limit": limit, "version": self.version},
{"limit": limit, "version": self.api_version, "keys": list(filters.keys())},
)
with concurrent.futures.ThreadPoolExecutor() as executor:
future_memories = executor.submit(self._search_vector_store, query, filters, limit)
future_graph_entities = (
executor.submit(self.graph.search, query, filters, limit)
if self.version == "v1.1" and self.enable_graph
if self.api_version == "v1.1" and self.enable_graph
else None
)
@@ -416,7 +416,7 @@ class Memory(MemoryBase):
original_memories = future_memories.result()
graph_entities = future_graph_entities.result() if future_graph_entities else None
if self.version == "v1.1":
if self.api_version == "v1.1":
if self.enable_graph:
return {"results": original_memories, "relations": graph_entities}
else:
@@ -518,14 +518,14 @@ class Memory(MemoryBase):
"At least one filter is required to delete all memories. If you want to delete all memories, use the `reset()` method."
)
capture_event("mem0.delete_all", self, {"filters": len(filters)})
capture_event("mem0.delete_all", self, {"keys": list(filters.keys())})
memories = self.vector_store.list(filters=filters)[0]
for memory in memories:
self._delete_memory(memory.id)
logger.info(f"Deleted {len(memories)} memories")
if self.version == "v1.1" and self.enable_graph:
if self.api_version == "v1.1" and self.enable_graph:
self.graph.delete_all(filters)
return {"message": "Memories deleted successfully!"}
@@ -561,6 +561,7 @@ class Memory(MemoryBase):
payloads=[metadata],
)
self.db.add_history(memory_id, None, data, "ADD", created_at=metadata["created_at"])
capture_event("mem0._create_memory", self, {"memory_id": memory_id})
return memory_id
def _update_memory(self, memory_id, data, existing_embeddings, metadata=None):
@@ -603,6 +604,7 @@ class Memory(MemoryBase):
created_at=new_metadata["created_at"],
updated_at=new_metadata["updated_at"],
)
capture_event("mem0._update_memory", self, {"memory_id": memory_id})
return memory_id
def _delete_memory(self, memory_id):
@@ -611,6 +613,7 @@ class Memory(MemoryBase):
prev_value = existing_memory.payload["data"]
self.vector_store.delete(vector_id=memory_id)
self.db.add_history(memory_id, prev_value, None, "DELETE", is_deleted=1)
capture_event("mem0._delete_memory", self, {"memory_id": memory_id})
return memory_id
def reset(self):

View File

@@ -67,7 +67,7 @@ def capture_event(event_name, memory_instance, additional_data=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__}",
"function": f"{memory_instance.__class__.__module__}.{memory_instance.__class__.__name__}.{memory_instance.version}",
"function": f"{memory_instance.__class__.__module__}.{memory_instance.__class__.__name__}.{memory_instance.api_version}",
}
if additional_data:
event_data.update(additional_data)