Sfoglia il codice sorgente

[Mem0] Update platform client, improve deduction logic and update client docs (#1510)

Deshraj Yadav 1 anno fa
parent
commit
c7b9498693
3 ha cambiato i file con 411 aggiunte e 367 eliminazioni
  1. 238 222
      docs/platform/quickstart.mdx
  2. 172 144
      mem0/client/main.py
  3. 1 1
      pyproject.toml

+ 238 - 222
docs/platform/quickstart.mdx

@@ -27,308 +27,315 @@ client = MemoryClient(api_key="your-api-key")
 
 ## 4. Memory Operations
 
-We provide a simple yet customizable interface for performing CRUD operations on memory. Here is how you can create and get memories:
+Mem0 provides a simple and customizable interface for performing CRUD operations on memory.
 
 
 ### 4.1 Create Memories
 
-For users (long-term memory):
+You can create long-term and short-term memories for your users, AI Agents, etc. Here are some examples:
 
-```python
-# create long-term memory for users
-client.add("Remember my name is Deshraj Yadav.", user_id="deshraj")
-client.add("I like to eat pizza and go out on weekends.", user_id="deshraj")
-client.add("Oh I am actually allergic to cheese to cannot eat pizza anymore.", user_id="deshraj")
-```
+#### Long-term memory for a user
 
-Output:
-```python
-{'message': 'Memory added successfully!'}
+```python Code
+from mem0 import MemoryClient
+client = MemoryClient(api_key="your-api-key")
+
+messages = [
+    {"role": "user", "content": "Hi, I'm Alex. I'm a vegetarian and I'm allergic to nuts."},
+    {"role": "assistant", "content": "Hello Alex! I've noted that you're a vegetarian and have a nut allergy. I'll keep this in mind for any food-related recommendations or discussions."}
+]
+client.add(messages, user_id="alex")
 ```
 
-You can see all the memory operations happening on the platform itself.
+#### Short-term memory for a user session
 
-![Mem0 Platform Activity](/images/platform/activity.png)
+```python Code
+from mem0 import MemoryClient
+client = MemoryClient(api_key="your-api-key")
 
+messages = [
+    {"role": "user", "content": "I'm planning a trip to Japan next month."},
+    {"role": "assistant", "content": "That's exciting, Alex! A trip to Japan next month sounds wonderful. Would you like some recommendations for vegetarian-friendly restaurants in Japan?"},
+    {"role": "user", "content": "Yes, please! Especially in Tokyo."},
+    {"role": "assistant", "content": "Great! I'll remember that you're interested in vegetarian restaurants in Tokyo for your upcoming trip. I'll prepare a list for you in our next interaction."}
+]
+client.add(messages, user_id="alex123", session_id="trip-planning-2024")
+```
 
-You can also add memories for a particular session or for an AI agent that you are building:
+#### Long-term memory for agents
 
-- For user sessions (short-term memory):
+```python Code
+from mem0 import MemoryClient
+client = MemoryClient(api_key="your-api-key")
 
-```python
-client.add("Deshraj is building Gmail AI agent", user_id="deshraj", session_id="session-1")
+messages = [
+    {"role": "system", "content": "You are a personalized travel assistant. Remember user preferences and provide tailored recommendations."},
+    {"role": "assistant", "content": "Understood. I'll maintain personalized travel preferences for each user and provide customized recommendations based on their dietary restrictions, interests, and past interactions."}
+]
+client.add(messages, agent_id="travel-assistant")
 ```
 
-- For agents (long-term memory):
+You can monitor memory operations on the platform:
 
-```python
-client.add("Return short responses when responding to emails", agent_id="gmail-agent")
-```
+![Mem0 Platform Activity](/images/platform/activity.png)
 
-### 4.2 Retrieve Memories
 
+### 4.2 Search Relevant Memories
 
-<CodeGroup>
+You can also get related memories for a given natural language question using our search method.
 
+<CodeGroup>
 ```python Code
-client.get_all(user_id="deshraj")
+query = "What do you know about me?"
+client.search(query, user_id="alex")
 ```
 
-```python Output
+```json Output
 [
-  {
-    'id': 'dbce6e06-6adf-40b8-9187-3d30bd13b741',
-    'agent': None,
-    'consumer': {
-      'id': 8,
-      'user_id': 'deshraj',
-      'metadata': None,
-      'created_at': '2024-07-17T16:47:23.899900-07:00',
-      'updated_at': '2024-07-17T16:47:23.899918-07:00'
-    },
-    'app': None,
-    'run': None,
-    'hash': '57288ac8a87c4ac8d3ac7f2075d264ca',
-    'input': 'Remember my name is Deshraj Yadav.',
-    'text': 'My name is Deshraj Yadav.',
-    'metadata': None,
-    'created_at': '2024-07-17T16:47:25.670180-07:00',
-    'updated_at': '2024-07-17T16:47:25.670197-07:00'
-  },
-  {
-    'id': 'f6dec5d1-b5db-45f5-a2fb-3979a0f27d30',
-    'agent': None,
-    'consumer': {
-      'id': 8,
-      'user_id': 'deshraj',
-      # ... other consumer fields ...
-    },
-    # ... other fields ...
-    'text': 'I am allergic to cheese so I cannot eat pizza anymore.',
-    # ... remaining fields ...
-  },
-  # ... additional memory entries ...
+    {
+        "id": "7f165f7e-b411-4afe-b7e5-35789b72c4a5",
+        "memory": "Name: Alex. Vegetarian. Allergic to nuts.",
+        "input": [
+            {
+                "role": "user",
+                "content": "Hi, I'm Alex. I'm a vegetarian and I'm allergic to nuts."
+            },
+            {
+                "role": "assistant",
+                "content": "Hello Alex! I've noted that you're a vegetarian and have a nut allergy. I'll keep this in mind for any food-related recommendations or discussions."
+            }
+        ],
+        "user_id": "alex",
+        "hash": "9ee7e1455e84d1dab700ed8749aed75a",
+        "metadata": null,
+        "created_at": "2024-07-20T01:30:36.275141-07:00",
+        "updated_at": "2024-07-20T01:30:36.275172-07:00"
+    }
 ]
 ```
 </CodeGroup>
 
-Similarly, you can get all memories for an agent:
+Similarly, you can search for agent memories by passing the `agent_id` argument. 
 
-```python
-agent_memories = client.get_all(agent_id="gmail-agent")
+```python Code
+client.search("What are the learnings from previous runs?", agent_id="travel-assistant")
 ```
 
-Get specific memory:
+### 4.3 Get All Memories
+
+Fetch all memories for a user, agent, or session using the get_all() method.
+
+#### Get all memories of an AI Agent
 
 <CodeGroup>
 
 ```python Code
-memory = client.get(memory_id="dbce6e06-6adf-40b8-9187-3d30bd13b741")
+client.get_all(agent_id="travel-assistant")
 ```
 
-```python Output
-{
-    'id': 'dbce6e06-6adf-40b8-9187-3d30bd13b741',
-    'agent': None,
-    'consumer': {
-      'id': 8,
-      'user_id': 'deshraj',
-      'metadata': None,
-      'created_at': '2024-07-17T16:47:23.899900-07:00',
-      'updated_at': '2024-07-17T16:47:23.899918-07:00'
-    },
-    'app': None,
-    'run': None,
-    'hash': '57288ac8a87c4ac8d3ac7f2075d264ca',
-    'input': 'Remember my name is Deshraj Yadav.',
-    'text': 'My name is Deshraj Yadav.',
-    'metadata': None,
-    'created_at': '2024-07-17T16:47:25.670180-07:00',
-    'updated_at': '2024-07-17T16:47:25.670197-07:00'
-}
+```json Output
+[
+    {
+        "id": "48677fae-16c4-4d57-9c27-f984f290722a",
+        "memory": "Will remember personalized travel preferences for each user.",
+        "input": [
+            {
+                "role": "system",
+                "content": "You are a personalized travel assistant. Remember user preferences and provide tailored recommendations."
+            },
+            {
+                "role": "assistant",
+                "content": "Understood. I'll maintain personalized travel preferences for each user and provide customized recommendations based on their dietary restrictions, interests, and past interactions."
+            }
+        ],
+        "agent_id": "travel-assistant",
+        "hash": "644d5f40af13d0525305d8dfadc00fd1",
+        "metadata": null,
+        "created_at": "2024-07-20T01:25:22.122299-07:00",
+        "updated_at": "2024-07-20T01:25:22.122327-07:00"
+    }
+]
 ```
 </CodeGroup>
 
-### 4.3 Update Memory
-
-You can also update specific memory by using the following method:
+#### Get all memories of user
 
 <CodeGroup>
-
 ```python Code
-client.update(memory_id, data="Updated name is Deshraj Kumar")
+user_memories = client.get_all(user_id="alex")
 ```
 
-```python Output
-{
-    'id': 'dbce6e06-6adf-40b8-9187-3d30bd13b741',
-    'agent': None,
-    'consumer': {
-      'id': 8,
-      'user_id': 'deshraj',
-      'metadata': None,
-      'created_at': '2024-07-17T16:47:23.899900-07:00',
-      'updated_at': '2024-07-17T16:47:23.899918-07:00'
-    },
-    'app': None,
-    'run': None,
-    'hash': '57288ac8a87c4ac8d3ac7f2075d264ca',
-    'input': 'Updated name is Deshraj Kumar.',
-    'text': 'Name is Deshraj Kumar.',
-    'metadata': None,
-    'created_at': '2024-07-17T16:47:25.670180-07:00',
-    'updated_at': '2024-07-17T16:47:25.670197-07:00'
-}
+```json Output
+[
+    {
+        "id": "7f165f7e-b411-4afe-b7e5-35789b72c4a5",
+        "memory": "Name: Alex. Vegetarian. Allergic to nuts.",
+        "input": [
+            {
+                "role": "user",
+                "content": "Hi, I'm Alex. I'm a vegetarian and I'm allergic to nuts."
+            },
+            {
+                "role": "assistant",
+                "content": "Hello Alex! I've noted that you're a vegetarian and have a nut allergy. I'll keep this in mind for any food-related recommendations or discussions."
+            }
+        ],
+        "user_id": "alex",
+        "hash": "9ee7e1455e84d1dab700ed8749aed75a",
+        "metadata": null,
+        "created_at": "2024-07-20T01:30:36.275141-07:00",
+        "updated_at": "2024-07-20T01:30:36.275172-07:00"
+    }
+]
 ```
 </CodeGroup>
 
-
-### 4.4 Memory History
-
-Get history of how a memory has changed over time
+#### Get short-term memories for a session
 
 <CodeGroup>
 
 ```python Code
-history = client.history(memory_id)
+short_term_memories = client.get_all(user_id="alex123", session_id="trip-planning-2024")
 ```
 
 ```python Output
 [
     {
-        'id': '51193804-2ee6-4f81-b4e7-497e98b70858',
-        'memory': {
-            'id': 'dbce6e06-6adf-40b8-9187-3d30bd13b741',
-            'agent': None,
-            'consumer': {
-                'id': 8,
-                'user_id': 'deshraj',
-                'metadata': None,
-                'created_at': '2024-07-17T16:47:23.899900-07:00',
-                'updated_at': '2024-07-17T16:47:23.899918-07:00'
+        "id": "582bbe6d-506b-48c6-a4c6-5df3b1e63428",
+        "memory": "Planning a trip to Japan next month. Interested in vegetarian restaurants in Tokyo.",
+        "input": [
+            {
+                "role": "user",
+                "content": "I'm planning a trip to Japan next month."
             },
-            'app': None,
-            'run': None,
-            'hash': '57288ac8a87c4ac8d3ac7f2075d264ca',
-            'input': 'Remember my name is Deshraj Yadav.',
-            'text': 'My name is Deshraj Yadav.',
-            'metadata': None,
-            'created_at': '2024-07-17T16:47:25.670180-07:00',
-            'updated_at': '2024-07-17T16:47:25.670197-07:00'
-        },
-        'hash': '57288ac8a87c4ac8d3ac7f2075d264ca',
-        'event': 'ADD',
-        'input': 'Remember my name is Deshraj Yadav.',
-        'previous_text': None,
-        'text': 'My name is Deshraj Yadav.',
-        'metadata': None,
-        'created_at': '2024-07-17T16:47:25.686899-07:00',
-        'updated_at': '2024-07-17T16:47:25.670197-07:00',
-        'change_description': 'Memory ADD event'
+            {
+                "role": "assistant",
+                "content": "That's exciting, Alex! A trip to Japan next month sounds wonderful. Would you like some recommendations for vegetarian-friendly restaurants in Japan?"
+            },
+            {
+                "role": "user",
+                "content": "Yes, please! Especially in Tokyo."
+            },
+            {
+                "role": "assistant",
+                "content": "Great! I'll remember that you're interested in vegetarian restaurants in Tokyo for your upcoming trip. I'll prepare a list for you in our next interaction."
+            }
+        ],
+        "user_id": "alex123",
+        "hash": "d2088c936e259f2f5d2d75543d31401c",
+        "metadata": null,
+        "created_at": "2024-07-20T01:34:09.748379-07:00",
+        "updated_at": "2024-07-20T01:34:09.748391-07:00"
     }
 ]
 ```
 </CodeGroup>
 
-### 4.5 Search for relevant memories
+
+#### Get specific memory
 
 <CodeGroup>
 
 ```python Code
-client.search("What does Deshraj like to eat?", user_id="deshraj", limit=3)
+memory = client.get(memory_id="582bbe6d-506b-48c6-a4c6-5df3b1e63428")
 ```
 
 ```python Output
-[
-  {
-    "id": "dbce6e06-6adf-40b8-9187-3d30bd13b741",
-    "agent": null,
-    "consumer": {
-      "id": 8,
-      "user_id": "deshraj",
-      "metadata": null,
-      "created_at": "...",
-      "updated_at": "..."
-    },
-    "app": null,
-    "run": null,
-    "hash": "57288ac8a87c4ac8d3ac7f2075d264ca",
-    "input": "Remember my name is Deshraj Yadav.",
-    "text": "My name is Deshraj Yadav.",
-    "metadata": null,
-    "created_at": "2024-07-17T16:47:25.670180-07:00",
-    "updated_at": "..."
-  },
-  {
-    "id": "091dbed6-74b4-4e15-b765-81be2abe0d6b",
-    "agent": null,
-    "consumer": {
-      "id": 8,
-      "user_id": "deshraj",
-      "metadata": null,
-      "created_at": "...",
-      "updated_at": "..."
-    },
-    "app": null,
-    "run": null,
-    "hash": "622a5a24d5ac54136414a22ec12f9520",
-    "input": "Oh I am actually allergic to cheese to cannot eat pizza anymore.",
-    "text": "I like to eat pizza and go out on weekends.",
+{
+    "id": "582bbe6d-506b-48c6-a4c6-5df3b1e63428",
+    "memory": "Planning a trip to Japan next month. Interested in vegetarian restaurants in Tokyo.",
+    "input": [
+        {
+            "role": "user",
+            "content": "I'm planning a trip to Japan next month."
+        },
+        {
+            "role": "assistant",
+            "content": "That's exciting, Alex! A trip to Japan next month sounds wonderful. Would you like some recommendations for vegetarian-friendly restaurants in Japan?"
+        },
+        {
+            "role": "user",
+            "content": "Yes, please! Especially in Tokyo."
+        },
+        {
+            "role": "assistant",
+            "content": "Great! I'll remember that you're interested in vegetarian restaurants in Tokyo for your upcoming trip. I'll prepare a list for you in our next interaction."
+        }
+    ],
+    "user_id": "alex123",
+    "hash": "d2088c936e259f2f5d2d75543d31401c",
     "metadata": null,
-    "created_at": "2024-07-17T16:49:24.276695-07:00",
-    "updated_at": "..."
-  },
+    "created_at": "2024-07-20T01:34:09.748379-07:00",
+    "updated_at": "2024-07-20T01:34:09.748391-07:00"
+}
+```
+</CodeGroup>
+
+### 4.4 Memory History
+
+Get history of how a memory has changed over time
+
+<CodeGroup>
+
+```python Code
+# Add some message to create history
+messages = [{"role": "user", "content": "I recently tried chicken and I loved it. I'm thinking of trying more non-vegetarian dishes.."}]
+client.add(messages, user_id="alex")
+
+# Add second message to update history
+messages.append({'role': 'user', 'content': 'I turned vegetarian now.'})
+client.add(messages, user_id="alex")
+
+# Get history of how memory changed over time
+memory_id = "<memory-id-here>"
+history = client.history(memory_id)
+print(history)
+```
+
+```json Output
+[
   {
-    "id": "5fb8f85d-3383-4bad-9d46-f171272478a4",
-    "agent": null,
-    "consumer": {
-      "id": 8,
-      "user_id": "deshraj",
-      "metadata": null,
-      "created_at": "...",
-      "updated_at": "..."
-    },
-    "app": null,
-    "run": {
-      "id": 1,
-      "run_id": "session-1",
-      "name": "",
-      "metadata": null,
-      "created_at": "...",
-      "updated_at": "..."
-    },
-    "hash": "179ced9649ac2b85350ece4946b1ee9b",
-    "input": "Deshraj is building Gmail AI agent",
-    "text": "Deshraj is building Gmail AI agent",
-    "metadata": null,
-    "created_at": "2024-07-17T16:52:41.278920-07:00",
-    "updated_at": "..."
+    'id': '5a718c7e-7ed2-4c72-8c55-d7896dc61ca3',
+    'memory_id': 'c6c5da66-d851-4653-8a09-ec7b2c1c9cc9',
+    'input': [
+      {
+        'role': 'user',
+        'content': "I recently tried chicken and I loved it. I'm thinking of trying more non-vegetarian dishes.."
+      }
+    ],
+    'old_memory': None,
+    'new_memory': 'Recently tried chicken and loved it. Interested in trying more non-vegetarian dishes.',
+    'event': 'ADD',
+    'metadata': None,
+    'created_at': '2024-07-20T02:13:35.842720-07:00',
+    'updated_at': '2024-07-20T02:13:35.818065-07:00'
   },
   {
-    "id": "f6dec5d1-b5db-45f5-a2fb-3979a0f27d30",
-    "agent": null,
-    "consumer": {
-      "id": 8,
-      "user_id": "deshraj",
-      "metadata": null,
-      "created_at": "...",
-      "updated_at": "..."
-    },
-    "app": null,
-    "run": null,
-    "hash": "19248f0766044b5973fc0ef1bf3955ef",
-    "input": "Oh I am actually allergic to cheese to cannot eat pizza anymore.",
-    "text": "I am allergic to cheese so I cannot eat pizza anymore.",
-    "metadata": null,
-    "created_at": "2024-07-17T16:49:38.622084-07:00",
-    "updated_at": "..."
+    'id': 'e52e2c73-4b97-4154-9844-aa44c6ee57f5',
+    'memory_id': 'c6c5da66-d851-4653-8a09-ec7b2c1c9cc9',
+    'input': [
+      {
+        'role': 'user',
+        'content': "I recently tried chicken and I loved it. I'm thinking of trying more non-vegetarian dishes.."
+      },
+      {
+        'role': 'user',
+        'content': 'I turned vegetarian now.'
+      }
+    ],
+    'old_memory': 'Recently tried chicken and loved it. Interested in trying more non-vegetarian dishes.',
+    'new_memory': '',
+    'event': 'DELETE',
+    'metadata': None,
+    'created_at': '2024-07-20T02:13:38.070492-07:00',
+    'updated_at': '2024-07-20T02:13:38.060219-07:00'
   }
 ]
 ```
 </CodeGroup>
 
 
-### 4.6 Delete Memory
+### 4.5 Delete Memory
 
 Delete specific memory:
 
@@ -356,3 +363,12 @@ client.delete_all(user_id="alex")
 {'message': 'Memories deleted successfully!'}
 ```
 </CodeGroup>
+
+Fun fact: You can also delete the memory using the `add()` method by passing natural language command. 
+
+```python Code
+client.add("Delete all of my food preferences", user_id="alex")
+
+If you have any questions, please feel free to reach out to us using one of the following methods:
+
+<Snippet file="get-help.mdx" />

+ 172 - 144
mem0/client/main.py

@@ -1,8 +1,10 @@
-import httpx
-import os
 import logging
-import warnings
-from typing import Optional, Dict, Any
+import os
+from functools import wraps
+from typing import Any, Dict, List, Optional, Union
+
+import httpx
+
 from mem0.memory.setup import setup_config
 from mem0.memory.telemetry import capture_client_event
 
@@ -12,232 +14,258 @@ logger = logging.getLogger(__name__)
 setup_config()
 
 
+class APIError(Exception):
+    """Exception raised for errors in the API."""
+
+    pass
+
+
+def api_error_handler(func):
+    """Decorator to handle API errors consistently."""
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except httpx.HTTPStatusError as e:
+            logger.error(f"HTTP error occurred: {e}")
+            raise APIError(f"API request failed: {e.response.text}")
+        except httpx.RequestError as e:
+            logger.error(f"Request error occurred: {e}")
+            raise APIError(f"Request failed: {str(e)}")
+
+    return wrapper
+
+
 class MemoryClient:
+    """Client for interacting with the Mem0 API.
+
+    This class provides methods to create, retrieve, search, and delete memories
+    using the Mem0 API.
+
+    Attributes:
+        api_key (str): The API key for authenticating with the Mem0 API.
+        host (str): The base URL for the Mem0 API.
+        client (httpx.Client): The HTTP client used for making API requests.
+    """
+
     def __init__(self, api_key: Optional[str] = None, host: Optional[str] = None):
-        """
-        Initialize the Mem0 client.
+        """Initialize the MemoryClient.
 
         Args:
-            api_key (Optional[str]): API Key from Mem0 Platform. Defaults to environment variable 'MEM0_API_KEY' if not provided.
-            host (Optional[str]): API host URL. Defaults to 'https://api.mem0.ai/v1'.
+            api_key: The API key for authenticating with the Mem0 API. If not provided,
+                     it will attempt to use the MEM0_API_KEY environment variable.
+            host: The base URL for the Mem0 API. Defaults to "https://api.mem0.ai/v1".
+
+        Raises:
+            ValueError: If no API key is provided or found in the environment.
         """
         self.api_key = api_key or os.getenv("MEM0_API_KEY")
         self.host = host or "https://api.mem0.ai/v1"
+
+        if not self.api_key:
+            raise ValueError("API Key not provided. Please provide an API Key.")
+
         self.client = httpx.Client(
             base_url=self.host,
             headers={"Authorization": f"Token {self.api_key}"},
+            timeout=60,
         )
         self._validate_api_key()
         capture_client_event("client.init", self)
 
     def _validate_api_key(self):
-        if not self.api_key:
-            warnings.warn("API Key not provided. Please provide an API Key.")
-        response = self.client.get("/memories/", params={"user_id": "test"})
-        if response.status_code != 200:
+        """Validate the API key by making a test request."""
+        try:
+            response = self.client.get("/memories/", params={"user_id": "test"})
+            response.raise_for_status()
+        except httpx.HTTPStatusError:
             raise ValueError(
                 "Invalid API Key. Please get a valid API Key from https://app.mem0.ai"
             )
 
+    @api_error_handler
     def add(
-        self,
-        data: str,
-        user_id: Optional[str] = None,
-        agent_id: Optional[str] = None,
-        session_id: Optional[str] = None,
-        metadata: Optional[Dict[str, Any]] = None,
-        filters: Optional[Dict[str, Any]] = None,
+        self, messages: Union[str, List[Dict[str, str]]], **kwargs
     ) -> Dict[str, Any]:
-        """
-        Create a new memory.
+        """Add a new memory.
 
         Args:
-            data (str): The data to be stored in the memory.
-            user_id (Optional[str]): User ID to save the memory specific to a user. Defaults to None.
-            agent_id (Optional[str]): Agent ID for agent-specific memory. Defaults to None.
-            session_id (Optional[str]): Run ID to save memory for a specific session. Defaults to None.
-            metadata (Optional[Dict[str, Any]]): Metadata to be saved with the memory. Defaults to None.
-            filters (Optional[Dict[str, Any]]): Filters to apply to the memory. Defaults to None.
+            messages: Either a string message or a list of message dictionaries.
+            **kwargs: Additional parameters such as user_id, agent_id, session_id, metadata, filters.
 
         Returns:
-            Dict[str, Any]: The response from the server.
+            A dictionary containing the API response.
+
+        Raises:
+            APIError: If the API request fails.
         """
+        payload = self._prepare_payload(messages, kwargs)
+        response = self.client.post("/memories/", json=payload)
+        response.raise_for_status()
         capture_client_event("client.add", self)
-        payload = {"text": data}
-        if metadata:
-            payload["metadata"] = metadata
-        if filters:
-            payload["filters"] = filters
-        if user_id:
-            payload["user_id"] = user_id
-        if agent_id:
-            payload["agent_id"] = agent_id
-        if session_id:
-            payload["run_id"] = session_id
-
-        response = self.client.post("/memories/", json=payload, timeout=60)
-        if response.status_code != 200:
-            logger.error(response.json())
-            raise ValueError(f"Failed to add memory. Response: {response.json()}")
         return response.json()
 
+    @api_error_handler
     def get(self, memory_id: str) -> Dict[str, Any]:
-        """
-        Get a memory by ID.
+        """Retrieve a specific memory by ID.
 
         Args:
-            memory_id (str): Memory ID.
+            memory_id: The ID of the memory to retrieve.
 
         Returns:
-            Dict[str, Any]: The memory data.
+            A dictionary containing the memory data.
+
+        Raises:
+            APIError: If the API request fails.
         """
-        capture_client_event("client.get", self)
         response = self.client.get(f"/memories/{memory_id}/")
+        response.raise_for_status()
+        capture_client_event("client.get", self)
         return response.json()
 
-    def get_all(
-        self,
-        user_id: Optional[str] = None,
-        agent_id: Optional[str] = None,
-        session_id: Optional[str] = None,
-        limit: int = 100,
-    ) -> Dict[str, Any]:
-        """
-        Get all memories.
+    @api_error_handler
+    def get_all(self, **kwargs) -> Dict[str, Any]:
+        """Retrieve all memories, with optional filtering.
 
         Args:
-            user_id (Optional[str]): User ID to filter memories. Defaults to None.
-            agent_id (Optional[str]): Agent ID to filter memories. Defaults to None.
-            session_id (Optional[str]): Run ID to filter memories. Defaults to None.
-            limit (int): Number of memories to return. Defaults to 100.
+            **kwargs: Optional parameters for filtering (user_id, agent_id, session_id, limit).
 
         Returns:
-            Dict[str, Any]: The list of memories.
+            A dictionary containing the list of memories.
+
+        Raises:
+            APIError: If the API request fails.
         """
-        params = {
-            "user_id": user_id,
-            "agent_id": agent_id,
-            "run_id": session_id,
-            "limit": limit,
-        }
-        response = self.client.get(
-            "/memories/", params={k: v for k, v in params.items() if v is not None}
-        )
+        params = self._prepare_params(kwargs)
+        response = self.client.get("/memories/", params=params)
+        response.raise_for_status()
         capture_client_event(
-            "client.get_all", self, {"filters": len(params), "limit": limit}
+            "client.get_all",
+            self,
+            {"filters": len(params), "limit": kwargs.get("limit", 100)},
         )
         return response.json()
 
-    def search(
-        self,
-        query: str,
-        user_id: Optional[str] = None,
-        agent_id: Optional[str] = None,
-        session_id: Optional[str] = None,
-        limit: int = 100,
-        filters: Optional[Dict[str, Any]] = None,
-    ) -> Dict[str, Any]:
-        """
-        Search memories.
+    @api_error_handler
+    def search(self, query: str, **kwargs) -> Dict[str, Any]:
+        """Search memories based on a query.
 
         Args:
-            query (str): Query to search for in the memories.
-            user_id (Optional[str]): User ID to filter memories. Defaults to None.
-            agent_id (Optional[str]): Agent ID to filter memories. Defaults to None.
-            session_id (Optional[str]): Run ID to filter memories. Defaults to None.
-            limit (int): Number of memories to return. Defaults to 100.
-            filters (Optional[Dict[str, Any]]): Filters to apply to the search. Defaults to None.
+            query: The search query string.
+            **kwargs: Additional parameters such as user_id, agent_id, session_id, limit, filters.
 
         Returns:
-            Dict[str, Any]: The search results.
-        """
-        payload = {
-            "text": query,
-            "limit": limit,
-            "filters": filters,
-            "user_id": user_id,
-            "agent_id": agent_id,
-            "run_id": session_id,
-        }
-        response = self.client.post("/memories/search/", json=payload)
-        capture_client_event("client.search", self, {"limit": limit})
-        return response.json()
+            A dictionary containing the search results.
 
-    def update(self, memory_id: str, data: str) -> Dict[str, Any]:
+        Raises:
+            APIError: If the API request fails.
         """
-        Update a memory by ID.
-
-        Args:
-            memory_id (str): Memory ID.
-            data (str): Data to update in the memory.
-
-        Returns:
-            Dict[str, Any]: The response from the server.
-        """
-        capture_client_event("client.update", self)
-        response = self.client.put(f"/memories/{memory_id}/", json={"text": data})
+        payload = {"query": query}
+        payload.update({k: v for k, v in kwargs.items() if v is not None})
+        response = self.client.post("/memories/search/", json=payload)
+        response.raise_for_status()
+        capture_client_event("client.search", self, {"limit": kwargs.get("limit", 100)})
         return response.json()
 
+    @api_error_handler
     def delete(self, memory_id: str) -> Dict[str, Any]:
-        """
-        Delete a memory by ID.
+        """Delete a specific memory by ID.
 
         Args:
-            memory_id (str): Memory ID.
+            memory_id: The ID of the memory to delete.
 
         Returns:
-            Dict[str, Any]: The response from the server.
+            A dictionary containing the API response.
+
+        Raises:
+            APIError: If the API request fails.
         """
-        capture_client_event("client.delete", self)
         response = self.client.delete(f"/memories/{memory_id}/")
+        response.raise_for_status()
+        capture_client_event("client.delete", self)
         return response.json()
 
-    def delete_all(
-        self,
-        user_id: Optional[str] = None,
-        agent_id: Optional[str] = None,
-        session_id: Optional[str] = None,
-    ) -> Dict[str, Any]:
-        """
-        Delete all memories.
+    @api_error_handler
+    def delete_all(self, **kwargs) -> Dict[str, Any]:
+        """Delete all memories, with optional filtering.
 
         Args:
-            user_id (Optional[str]): User ID to filter memories. Defaults to None.
-            agent_id (Optional[str]): Agent ID to filter memories. Defaults to None.
-            session_id (Optional[str]): Run ID to filter memories. Defaults to None.
+            **kwargs: Optional parameters for filtering (user_id, agent_id, session_id).
 
         Returns:
-            Dict[str, Any]: The response from the server.
+            A dictionary containing the API response.
+
+        Raises:
+            APIError: If the API request fails.
         """
-        params = {"user_id": user_id, "agent_id": agent_id, "run_id": session_id}
-        response = self.client.delete(
-            "/memories/", params={k: v for k, v in params.items() if v is not None}
-        )
+        params = self._prepare_params(kwargs)
+        response = self.client.delete("/memories/", params=params)
+        response.raise_for_status()
         capture_client_event("client.delete_all", self, {"params": len(params)})
         return response.json()
 
+    @api_error_handler
     def history(self, memory_id: str) -> Dict[str, Any]:
-        """
-        Get history of a memory by ID.
+        """Retrieve the history of a specific memory.
 
         Args:
-            memory_id (str): Memory ID.
+            memory_id: The ID of the memory to retrieve history for.
 
         Returns:
-            Dict[str, Any]: The memory history.
+            A dictionary containing the memory history.
+
+        Raises:
+            APIError: If the API request fails.
         """
         response = self.client.get(f"/memories/{memory_id}/history/")
+        response.raise_for_status()
         capture_client_event("client.history", self)
         return response.json()
 
     def reset(self):
-        """
-        Reset the client. (Not implemented yet)
+        """Reset the client. (Not implemented)
+
+        Raises:
+            NotImplementedError: This method is not implemented yet.
         """
         raise NotImplementedError("Reset is not implemented yet")
 
     def chat(self):
-        """
-        Start a chat with the Mem0 AI. (Not implemented yet)
+        """Start a chat with the Mem0 AI. (Not implemented)
+
+        Raises:
+            NotImplementedError: This method is not implemented yet.
         """
         raise NotImplementedError("Chat is not implemented yet")
+
+    def _prepare_payload(
+        self, messages: Union[str, List[Dict[str, str]], None], kwargs: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Prepare the payload for API requests.
+
+        Args:
+            messages: The messages to include in the payload.
+            kwargs: Additional keyword arguments to include in the payload.
+
+        Returns:
+            A dictionary containing the prepared payload.
+        """
+        payload = {}
+        if isinstance(messages, str):
+            payload["messages"] = [{"role": "user", "content": messages}]
+        elif isinstance(messages, list):
+            payload["messages"] = messages
+        payload.update({k: v for k, v in kwargs.items() if v is not None})
+        return payload
+
+    def _prepare_params(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
+        """Prepare query parameters for API requests.
+
+        Args:
+            kwargs: Keyword arguments to include in the parameters.
+
+        Returns:
+            A dictionary containing the prepared parameters.
+        """
+        return {k: v for k, v in kwargs.items() if v is not None}

+ 1 - 1
pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "mem0ai"
-version = "0.0.8"
+version = "0.0.9"
 description = "Long-term memory for AI Agents"
 authors = ["Deshraj Yadav <deshraj@mem0.ai>", "Taranjeet Singh <taranjeet@mem0.ai>"]
 exclude = [