Add custom prompt

Amnah199 · Amnah199 · commit df96c5612daf · 2025-11-24T20:09:05.000+05:00
diff --git a/haystack_experimental/components/memory_agents/agent.py b/haystack_experimental/components/memory_agents/agent.py
@@ -254,5 +254,6 @@ def run(  # noqa: PLR0915
             result["last_message"] = msgs[-1]
 
         # Add the new conversation as memories to the memory store
-        self.memory_store.add_memories(result["messages"])
+        user_messages = [message for message in result["messages"] if message.role == "user"]
+        self.memory_store.add_memories(user_messages)
         return result
diff --git a/haystack_experimental/memory/example.py b/haystack_experimental/memory/example.py
@@ -2,12 +2,25 @@
 from haystack.dataclasses import ChatMessage
 
 from haystack_experimental.components.memory_agents.agent import Agent
-from haystack_experimental.memory.src.mem0.memory_store import Mem0MemoryStore
+from haystack_experimental.memory.src.mem0.memory_store import Mem0MemoryStore, Mem0Scope
 
-memory_store = Mem0MemoryStore(user_id="haystack_mem0")
+memory_store = Mem0MemoryStore(scope=Mem0Scope(user_id="haystack_mem0"))
 
-chat_generator = OpenAIChatGenerator()
-agent = Agent(chat_generator=chat_generator, memory_store=memory_store)
 
-answer = agent.run(messages=[ChatMessage.from_user(" suggest me some music and a drink with it to relax.")])
-print(answer)
+messages = [
+    ChatMessage.from_user("I like to listen to Russian pop music"),
+    ChatMessage.from_user("I liked cold spanish latte with oat milk"),
+    ChatMessage.from_user("I live in Florence Italy and I love mountains"),
+    ChatMessage.from_user("""I am a software engineer and I like building application in python.
+                                  Most of my projects are related to NLP and LLM agents.
+                                  I find it easier to use Haystack framework to build my projects."""),
+    ChatMessage.from_user("""I work in a startup and I am the CEO of the company.
+                                  I have a team of 10 people and we are building a
+                                  platform for small businesses to manage their customers and sales."""),
+]
+
+memory_store.add_memories(messages)
+
+result = memory_store.search_memories()
+
+print(result)
diff --git a/haystack_experimental/memory/src/mem0/memory_store.py b/haystack_experimental/memory/src/mem0/memory_store.py
@@ -8,42 +8,65 @@
 from haystack.dataclasses.chat_message import ChatMessage
 from haystack.lazy_imports import LazyImport
 
+from .utils import Mem0Scope
+
 with LazyImport(message="Run 'pip install mem0ai'") as mem0_import:
-    from mem0 import MemoryClient
+    from mem0 import Memory, MemoryClient
 
 
 class Mem0MemoryStore:
     """
     A memory store implementation using Mem0 as the backend.
 
-    :param api_key: The Mem0 API key (if not provided, uses MEM0_API_KEY environment variable)
-    :param user_id: The user ID for the memory store.
-    :param memory_config: Configuration dictionary for Mem0 client
     """
 
-    def __init__(self, user_id: str, api_key: Optional[str] = None, memory_config: Optional[dict[str, Any]] = None):
+    def __init__(
+        self,
+        scope: Mem0Scope,
+        api_key: Optional[str] = None,
+        memory_config: Optional[dict[str, Any]] = None,
+        search_criteria: Optional[dict[str, Any]] = None,
+    ):
+        """
+        Initialize the Mem0 memory store.
+
+        :param scope: The scope for the memory store. This defines the identifiers to retrieve or update memories.
+        :param api_key: The Mem0 API key (if not provided, uses MEM0_API_KEY environment variable)
+        :param memory_config: Configuration dictionary for Mem0 client
+        :param search_criteria: Set the search configuration for the memory store.
+            This can include query, filters, and top_k.
+        """
+
         mem0_import.check()
         self.api_key = api_key or os.getenv("MEM0_API_KEY")
         if not self.api_key:
             raise ValueError("Mem0 API key must be provided either as parameter or MEM0_API_KEY environment variable.")
 
-        self.user_id = user_id
+        self.scope = scope
         self.memory_config = memory_config
 
         # If a memory config is provided, use it to initialize the Mem0 client
         if self.memory_config:
-            self.client = MemoryClient.from_config(self.memory_config)
+            self.client = Memory.from_config(self.memory_config)
         else:
-            self.client = MemoryClient(api_key=self.api_key)
+            self.client = MemoryClient(
+                api_key=self.api_key,
+            )
 
         # User can set the search criteria using the set_search_criteria method
-        self.search_criteria = {}
+        self.search_criteria = search_criteria
+        if not self.search_criteria:
+            self.search_criteria = {
+                "query": None,
+                "filters": None,
+                "top_k": 10,
+            }
 
     def to_dict(self) -> dict[str, Any]:
         """Serialize the store configuration to a dictionary."""
         return default_to_dict(
             self,
-            user_id=self.user_id,
+            scope=self.scope,
             api_key=self.api_key,
             memory_config=self.memory_config,
             search_criteria=self.search_criteria,
@@ -54,7 +77,7 @@ def from_dict(cls, data: dict[str, Any]) -> "Mem0MemoryStore":
         """Deserialize the store from a dictionary."""
         return default_from_dict(cls, data)
 
-    def add_memories(self, messages: list[ChatMessage]) -> list[str]:
+    def add_memories(self, messages: list[ChatMessage], infer: bool = True) -> list[str]:
         """
         Add ChatMessage memories to Mem0.
 
@@ -69,53 +92,52 @@ def add_memories(self, messages: list[ChatMessage]) -> list[str]:
             mem0_message = [{"role": "user", "content": message.text}]
 
             try:
-                # Mem0 primarily uses user_id as the main identifier
-                # org_id and session_id are stored in metadata for filtering
-                result = self.client.add(messages=mem0_message, user_id=self.user_id, metadata=message.meta)
+                result = self.client.add(
+                    messages=mem0_message, metadata=message.meta, infer=infer, **self.scope.get_scope()
+                )
                 # Mem0 returns different response formats, handle both
                 memory_id = result.get("id") or result.get("memory_id") or str(result)
                 added_ids.append(memory_id)
             except Exception as e:
                 raise RuntimeError(f"Failed to add memory message: {e}") from e
 
-        return added_ids
-
-    def set_search_criteria(
-        self, query: Optional[str] = None, filters: Optional[dict[str, Any]] = None, top_k: Optional[int] = None
-    ):
-        """
-        Set the memory configuration for the memory store.
-        """
-        self.search_criteria = {"query": query, "filters": filters, "top_k": top_k}
+        return list(added_ids)
 
     def search_memories(
-        self, query: Optional[str] = None, filters: Optional[dict[str, Any]] = None, top_k: int = 8
+        self,
+        query: Optional[str] = None,
+        filters: Optional[dict[str, Any]] = None,
+        top_k: int = 5,
+        search_criteria: Optional[dict[str, Any]] = None,
     ) -> list[ChatMessage]:
         """
         Search for memories in Mem0.
 
         :param query: Text query to search for. If not provided, all memories will be returned.
         :param filters: Additional filters to apply on search. For more details on mem0 filters, see https://mem0.ai/docs/search/
         :param top_k: Maximum number of results to return
+        :param search_criteria: Search criteria to search memories from the store.
+            This can include query, filters, and top_k.
         :returns: List of ChatMessage memories matching the criteria
         """
         # Prepare filters for Mem0
+        search_criteria = search_criteria or self.search_criteria
 
-        search_query = query or self.search_criteria.get("query", None)
-        search_filters = filters or self.search_criteria.get("filters", {})
-        search_top_k = top_k or self.search_criteria.get("top_k", 10)
+        search_query = query or search_criteria.get("query", None)
+        search_filters = filters or search_criteria.get("filters", {})
+        search_top_k = top_k or search_criteria.get("top_k", 10)
 
         if search_filters:
-            mem0_filters = {"AND": [{"user_id": self.user_id}, search_filters]}
+            mem0_filters = search_filters
         else:
-            mem0_filters = {"user_id": self.user_id}
+            mem0_filters = self.scope.get_scope()
 
         try:
             if not search_query:
-                memories = self.client.get_all(filters=mem0_filters, top_k=search_top_k)
+                memories = self.client.get_all(filters=mem0_filters, **self.scope.get_scope())
             else:
                 memories = self.client.search(
-                    query=search_query, limit=search_top_k, filters=mem0_filters, user_id=self.user_id
+                    query=search_query, limit=search_top_k, filters=mem0_filters, **self.scope.get_scope()
                 )
             messages = [
                 ChatMessage.from_user(text=memory["memory"], meta=memory["metadata"]) for memory in memories["results"]
@@ -135,7 +157,7 @@ def delete_all_memories(self, user_id: Optional[str] = None):
         :param user_id: User identifier for scoping the deletion
         """
         try:
-            self.client.delete_all(user_id=user_id or self.user_id)
+            self.client.delete_all(**self.scope.get_scope())
         except Exception as e:
             raise RuntimeError(f"Failed to delete memories for user {user_id}: {e}") from e
 
@@ -149,15 +171,3 @@ def delete_memory(self, memory_id: str):
             self.client.delete(memory_id=memory_id)
         except Exception as e:
             raise RuntimeError(f"Failed to delete memory {memory_id}: {e}") from e
-
-    def get_memory(self, memory_id: str) -> ChatMessage:
-        """
-        Get memory from Mem0.
-
-        :param memory_id: The ID of the memory to get.
-        :returns: The memory.
-        """
-        try:
-            return self.client.get(memory_id=memory_id)
-        except Exception as e:
-            raise RuntimeError(f"Failed to get memory {memory_id}: {e}") from e
diff --git a/haystack_experimental/memory/src/mem0/utils.py b/haystack_experimental/memory/src/mem0/utils.py
@@ -0,0 +1,130 @@
+from typing import Any, Optional
+
+from haystack import default_from_dict
+
+
+class Mem0Scope:
+    """
+    This object is used to scope the memory store.
+    """
+
+    def __init__(self, user_id: Optional[str] = None, run_id: Optional[str] = None, session_id: Optional[str] = None):
+        self.user_id = user_id
+        self.run_id = run_id
+        self.session_id = session_id
+
+        self._check_one_id_is_set()
+
+    def _check_one_id_is_set(self):
+        # we need to check that at least one of the ids is set
+        if not self.user_id and not self.run_id and not self.session_id:
+            raise ValueError("At least one of user_id, run_id, or session_id must be set")
+
+    def _to_dict(self) -> dict[str, Any]:
+        return {
+            "user_id": self.user_id,
+            "run_id": self.run_id,
+            "session_id": self.session_id,
+        }
+
+    @classmethod
+    def _from_dict(cls, data: dict[str, Any]) -> "Mem0Scope":
+        return default_from_dict(cls, data)
+
+    def _get_scope(self) -> dict[str, Any]:
+        return {key: value for key, value in self.to_dict().items() if value is not None}
+
+
+SEMANTIC_MEMORY_EXTRACTION_PROMPT = """
+You are extracting semantic memories from a conversation chain. Semantic memories are general
+preferences or personality traits or facts about the user that occurred during the conversation.
+
+
+WHAT TO EXTRACT:
+- General preferences or personality traits
+- Facts about the user
+
+
+WHAT NOT TO EXTRACT:
+- Vague statements without concrete facts
+- Specific events or actions that occurred during the conversation
+- Contextual details from the conversation flow
+- Problem-solution pairs that emerged during the conversation
+- Questions without answers
+- Speculative or hypothetical statements
+
+
+OUTPUT FORMAT:
+You MUST return a JSON object with ONLY a "facts" key containing an array of strings.
+Each fact should be a string that contains fact about the user.
+
+{
+    "facts": [
+        "User likes to listen to Russian pop music"
+    ]
+}
+
+IMPORTANT:
+- Return JSON with ONLY the "facts" key
+- Each fact must be a STRING, not an object
+- If no semantic memories can be extracted, return {"facts": []}
+
+EXAMPLES:
+
+Example 1:
+Conversation:
+USER: I like to listen to Russian pop music
+ASSISTANT: That's interesting! What other types of music do you like?
+USER: I also like to listen to jazz music
+ASSISTANT: Jazz is great! Do you have a favorite jazz artist?
+USER: Yes, I like to listen to John Coltrane
+
+Extracted semantic memory:
+{
+    "facts": [
+        "User likes to listen to Russian pop music"
+        "User likes to listen to jazz music"
+        "User likes to listen to John Coltrane"
+    ]
+}
+
+Example 2:
+Conversation:
+USER: I live in Florence Italy and I love mountains.
+ASSISTANT: That sounds exciting! Which countries are you planning to visit?
+USER: I'm thinking of France, Italy, and Spain. I've already booked flights to Paris for next month.
+ASSISTANT: Great itinerary! Paris is beautiful. Have you been to Europe before?
+USER: Yes, I visited London last year for a conference. It was my first time in Europe.
+
+Extracted semantic memory:
+{
+    "facts": [
+        "User lives in Florence Italy and loves mountains"
+    ]
+}
+
+Example 3 (Empty Output):
+Conversation:
+USER: Hello, how are you?
+ASSISTANT: I'm doing well, thank you! How can I help you today?
+USER: Just checking in, nothing specific.
+
+Extracted semantic memory:
+{
+    "facts": []
+}
+
+
+"""
+
+semantic_memory_config = {
+    "llm": {
+        "provider": "openai",
+        "config": {
+            "model": "gpt-4.1-nano-2025-04-14",
+            "temperature": 0.2,
+            "max_tokens": 2000,
+        },
+    },
+    "custom_fact_extraction_prompt": SEMANTIC_MEMORY_EXTRACTION_PROMPT,
+}
