pratik/knowledge-base
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Added support for run id agent id

Browse source
This commit is contained in:
Pratik Narola 2025-09-02 17:28:41 +00:00
parent c864fe8895
commit aa7742c5ad
1 changed files with 29 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
backend/mem0_manager.py
View file

@ -32,10 +32,10 @@ class Mem0Manager:
"embedder": { "embedder": {
"provider": "ollama", "provider": "ollama",
"config": { "config": {
"model": "hf.co/Qwen/Qwen3-Embedding-0.6B-GGUF:Q8_0", "model": "hf.co/Qwen/Qwen3-Embedding-8B-GGUF:Q8_0",
# "api_key": settings.embedder_api_key, # "api_key": settings.embedder_api_key,
"ollama_base_url": "http://host.docker.internal:11434", "ollama_base_url": "http://host.docker.internal:11434",
"embedding_dims": 1024 "embedding_dims": 4096
} }
}, },
"vector_store": { "vector_store": {
@ -46,7 +46,7 @@ class Mem0Manager:
"password": settings.postgres_password, "password": settings.postgres_password,
"host": settings.postgres_host, "host": settings.postgres_host,
"port": settings.postgres_port, "port": settings.postgres_port,
"embedding_model_dims": 1024 "embedding_model_dims": 4096
} }
}, },
"graph_store": { "graph_store": {
@ -69,7 +69,6 @@ class Mem0Manager:
# Pure passthrough methods - no custom logic # Pure passthrough methods - no custom logic
@timed("add_memories") @timed("add_memories")
async def add_memories( async def add_memories(
self, self,
@ -91,12 +90,9 @@ class Mem0Manager:
# Direct Mem0 add with metadata support # Direct Mem0 add with metadata support
combined_metadata = metadata or {} combined_metadata = metadata or {}
if agent_id: result = self.memory.add(formatted_messages, user_id=user_id,
combined_metadata["agent_id"] = agent_id agent_id=agent_id, run_id=run_id,
if run_id: metadata=combined_metadata if combined_metadata else None)
combined_metadata["run_id"] = run_id
result = self.memory.add(formatted_messages, user_id=user_id, metadata=combined_metadata if combined_metadata else None)
return { return {
"added_memories": result if isinstance(result, list) else [result], "added_memories": result if isinstance(result, list) else [result],
@ -115,19 +111,19 @@ class Mem0Manager:
limit: int = 5, limit: int = 5,
threshold: Optional[float] = None, threshold: Optional[float] = None,
filters: Optional[Dict[str, Any]] = None, filters: Optional[Dict[str, Any]] = None,
keyword_search: bool = False, # keyword_search: bool = False,
rerank: bool = False, # rerank: bool = False,
filter_memories: bool = False, # filter_memories: bool = False,
agent_id: Optional[str] = None, agent_id: Optional[str] = None,
run_id: Optional[str] = None run_id: Optional[str] = None
) -> Dict[str, Any]: ) -> Dict[str, Any]:
"""Search memories - native Mem0 pattern (5 lines vs 70).""" """Search memories - native Mem0 pattern"""
try: try:
# Minimal empty query protection for API compatibility # Minimal empty query protection for API compatibility
if not query or query.strip() == "": if not query or query.strip() == "":
return {"memories": [], "total_count": 0, "query": query, "note": "Empty query provided, no results returned. Use a specific query to search memories."} return {"memories": [], "total_count": 0, "query": query, "note": "Empty query provided, no results returned. Use a specific query to search memories."}
# Direct Mem0 search - trust native handling # Direct Mem0 search - trust native handling
result = self.memory.search(query=query, user_id=user_id, limit=limit) result = self.memory.search(query=query, user_id=user_id, agent_id=agent_id, run_id=run_id, limit=limit, threshold=threshold, filters=filters)
return {"memories": result.get("results", []), "total_count": len(result.get("results", [])), "query": query} return {"memories": result.get("results", []), "total_count": len(result.get("results", [])), "query": query}
except Exception as e: except Exception as e:
logger.error(f"Error searching memories: {e}") logger.error(f"Error searching memories: {e}")
@ -144,7 +140,7 @@ class Mem0Manager:
"""Get all memories for a user - native Mem0 pattern.""" """Get all memories for a user - native Mem0 pattern."""
try: try:
# Direct Mem0 get_all call - trust native parameter handling # Direct Mem0 get_all call - trust native parameter handling
result = self.memory.get_all(user_id=user_id, limit=limit) result = self.memory.get_all(user_id=user_id, limit=limit, agent_id=agent_id, run_id=run_id, filters=filters)
return result.get("results", []) return result.get("results", [])
except Exception as e: except Exception as e:
logger.error(f"Error getting user memories: {e}") logger.error(f"Error getting user memories: {e}")
@ -177,7 +173,7 @@ class Mem0Manager:
logger.error(f"Error deleting memory: {e}") logger.error(f"Error deleting memory: {e}")
raise e raise e
async def delete_user_memories(self, user_id: str) -> Dict[str, Any]: async def delete_user_memories(self, user_id: Optional[str]) -> Dict[str, Any]:
"""Delete all user memories - pure Mem0 passthrough.""" """Delete all user memories - pure Mem0 passthrough."""
try: try:
self.memory.delete_all(user_id=user_id) self.memory.delete_all(user_id=user_id)
@ -200,12 +196,14 @@ class Mem0Manager:
raise e raise e
async def get_graph_relationships(self, user_id: str) -> Dict[str, Any]: async def get_graph_relationships(self, user_id: Optional[str], agent_id: Optional[str], run_id: Optional[str]) -> Dict[str, Any]:
"""Get graph relationships - using correct Mem0 get_all() method.""" """Get graph relationships - using correct Mem0 get_all() method."""
try: try:
# Use get_all() to retrieve memories with graph relationships # Use get_all() to retrieve memories with graph relationships
result = self.memory.get_all( result = self.memory.get_all(
user_id=user_id, user_id=user_id,
agent_id=agent_id,
run_id=run_id,
limit=50 limit=50
) )
@ -231,6 +229,8 @@ class Mem0Manager:
"relationships": relationships, "relationships": relationships,
"entities": entities, "entities": entities,
"user_id": user_id, "user_id": user_id,
"agent_id": agent_id,
"run_id": run_id,
"total_memories": len(result.get("results", [])), "total_memories": len(result.get("results", [])),
"total_relationships": len(relationships) "total_relationships": len(relationships)
} }
@ -242,6 +242,8 @@ class Mem0Manager:
"relationships": [], "relationships": [],
"entities": [], "entities": [],
"user_id": user_id, "user_id": user_id,
"agent_id": agent_id,
"run_id": run_id,
"total_memories": 0, "total_memories": 0,
"total_relationships": 0, "total_relationships": 0,
"error": str(e) "error": str(e)
@ -251,14 +253,16 @@ class Mem0Manager:
async def chat_with_memory( async def chat_with_memory(
self, self,
message: str, message: str,
user_id: str = "default", user_id: Optional[str] = None,
context: Optional[List[Dict[str, str]]] = None, agent_id: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None run_id: Optional[str] = None,
# context: Optional[List[Dict[str, str]]] = None,
# metadata: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]: ) -> Dict[str, Any]:
"""Chat with memory - native Mem0 pattern (15 lines vs 95).""" """Chat with memory - native Mem0 pattern (15 lines vs 95)."""
try: try:
# Retrieve relevant memories using direct Mem0 search # Retrieve relevant memories using direct Mem0 search
search_result = self.memory.search(query=message, user_id=user_id, limit=3) search_result = self.memory.search(query=message, user_id=user_id, agent_id=agent_id, run_id=run_id, limit=10, threshold=0.3)
relevant_memories = search_result.get("results", []) relevant_memories = search_result.get("results", [])
memories_str = "\n".join(f"- {entry['memory']}" for entry in relevant_memories) memories_str = "\n".join(f"- {entry['memory']}" for entry in relevant_memories)