From 4cc91d750564c656d82379e5b760d28de0702d66 Mon Sep 17 00:00:00 2001
From: Dev Khant <devkhant24@gmail.com>
Date: Thu, 7 Nov 2024 10:46:08 +0530
Subject: [PATCH] Add support for Org/Proj ID (#2014)

---
 docs/api-reference/overview.mdx | 12 +++++---
 mem0/client/main.py             | 53 +++++++++++++++++++++++++++++++--
 2 files changed, 58 insertions(+), 7 deletions(-)

diff --git a/docs/api-reference/overview.mdx b/docs/api-reference/overview.mdx
index c68708bb..9d10851b 100644
--- a/docs/api-reference/overview.mdx
+++ b/docs/api-reference/overview.mdx
@@ -32,9 +32,10 @@ Example with the mem0 Python package:
 ```python
 from mem0 import MemoryClient
 
+# Recommended: Using organization and project IDs
 client = MemoryClient(
-  organization='YOUR_ORG_NAME',
-  project='YOUR_PROJECT_NAME',
+    org_id='YOUR_ORG_ID',
+    project_id='YOUR_PROJECT_ID',
 )
 ```
 
@@ -43,12 +44,15 @@ Example with the mem0 Node.js package:
 ```javascript
 import { MemoryClient } from "mem0ai";
 
+// Recommended: Using organization and project IDs
 const client = new MemoryClient({
-    organization: "YOUR_ORG_NAME",
-    project: "YOUR_PROJECT_NAME"
+    orgId: "YOUR_ORG_ID",
+    projectId: "YOUR_PROJECT_ID"
 });
 ```
 
+> **Note**: The use of `organization` and `project` parameters is deprecated and will be removed in version `0.1.40`. Please use `org_id` and `project_id` (Python) or `orgId` and `projectId` (Node.js) instead.
+
 ## Getting Started
 
 To begin using the Mem0 API, you'll need to:
diff --git a/mem0/client/main.py b/mem0/client/main.py
index 3996b05a..b1268e35 100644
--- a/mem0/client/main.py
+++ b/mem0/client/main.py
@@ -2,6 +2,7 @@ import logging
 import os
 from functools import wraps
 from typing import Any, Dict, List, Optional, Union
+import warnings
 
 import httpx
 
@@ -55,6 +56,8 @@ class MemoryClient:
         host: Optional[str] = None,
         organization: Optional[str] = None,
         project: Optional[str] = None,
+        org_id: Optional[str] = None,
+        project_id: Optional[str] = None,
     ):
         """Initialize the MemoryClient.
 
@@ -62,8 +65,10 @@ class MemoryClient:
             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".
-            org_name: The name of the organization. Optional.
-            project_name: The name of the project. Optional.
+            organization: (Deprecated) The name of the organization. Use org_id instead.
+            project: (Deprecated) The name of the project. Use project_id instead.
+            org_id: The ID of the organization.
+            project_id: The ID of the project.
 
         Raises:
             ValueError: If no API key is provided or found in the environment.
@@ -72,11 +77,21 @@ class MemoryClient:
         self.host = host or "https://api.mem0.ai"
         self.organization = organization
         self.project = project
+        self.org_id = org_id
+        self.project_id = project_id
         self.user_id = get_user_id()
 
         if not self.api_key:
             raise ValueError("Mem0 API Key not provided. Please provide an API Key.")
 
+        if organization or project:
+            warnings.warn(
+                "Using 'organization' and 'project' parameters is deprecated and will be removed in version 0.1.40. "
+                "Please use 'org_id' and 'project_id' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         self.client = httpx.Client(
             base_url=self.host,
             headers={"Authorization": f"Token {self.api_key}", "Mem0-User-ID": self.user_id},
@@ -337,7 +352,30 @@ class MemoryClient:
 
         Returns:
             A dictionary containing the prepared parameters.
+
+        Raises:
+            ValueError: If both org_id/project_id and org_name/project_name are provided.
         """
+        has_new = bool(self.org_id or self.project_id)
+        has_old = bool(self.organization or self.project)
+        
+        if has_new and has_old:
+            raise ValueError(
+                "Please use either org_id/project_id or org_name/project_name, not both. "
+                "Note that org_name/project_name are deprecated."
+            )
+
+        # Add org_id and project_id if available
+        if self.org_id:
+            kwargs["org_id"] = self.org_id
+        if self.project_id:
+            kwargs["project_id"] = self.project_id
+
+        # Add deprecated org_name and project_name for backward compatibility
+        if self.organization:
+            kwargs["org_name"] = self.organization
+        if self.project:
+            kwargs["project_name"] = self.project
 
         return {k: v for k, v in kwargs.items() if v is not None}
 
@@ -351,8 +389,17 @@ class AsyncMemoryClient:
         host: Optional[str] = None,
         organization: Optional[str] = None,
         project: Optional[str] = None,
+        org_id: Optional[str] = None,
+        project_id: Optional[str] = None,
     ):
-        self.sync_client = MemoryClient(api_key, host, organization, project)
+        self.sync_client = MemoryClient(
+            api_key, 
+            host, 
+            organization, 
+            project,
+            org_id,
+            project_id
+        )
         self.async_client = httpx.AsyncClient(
             base_url=self.sync_client.host,
             headers=self.sync_client.client.headers,