feat: add ockam support

Author: SanjoDeundiak

mem0/configs/embeddings/base.py

@@ -37,6 +37,8 @@ def __init__(
         aws_access_key_id: Optional[str] = None,
         aws_secret_access_key: Optional[str] = None,
         aws_region: Optional[str] = "us-west-2",
+        # Ockam Model specific
+        ockam_model: Optional = None,
     ):
         """
         Initializes a configuration class instance for the Embeddings.
@@ -101,3 +103,6 @@ def __init__(
         self.aws_access_key_id = aws_access_key_id
         self.aws_secret_access_key = aws_secret_access_key
         self.aws_region = aws_region
+
+        # Ockam Model specific
+        self.ockam_model = ockam_model

mem0/configs/llms/base.py

@@ -50,6 +50,8 @@ def __init__(
         aws_access_key_id: Optional[str] = None,
         aws_secret_access_key: Optional[str] = None,
         aws_region: Optional[str] = "us-west-2",
+        # Ockam Model specific
+        ockam_model: Optional = None,
     ):
         """
         Initializes a configuration class instance for the LLM.
@@ -150,3 +152,6 @@ def __init__(
         self.aws_access_key_id = aws_access_key_id
         self.aws_secret_access_key = aws_secret_access_key
         self.aws_region = aws_region
+
+        # Ockam Model specific
+        self.ockam_model = ockam_model

mem0/configs/vector_stores/inmemory.py

@@ -0,0 +1,19 @@
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class InMemoryConfig(BaseModel):
+    collection_name: str = Field("mem0", description="Default name for the collection")
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_extra_fields(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+        allowed_fields = set(cls.model_fields.keys())
+        input_fields = set(values.keys())
+        extra_fields = input_fields - allowed_fields
+        if extra_fields:
+            raise ValueError(
+                f"Extra fields not allowed: {', '.join(extra_fields)}. Please input only the following fields: {', '.join(allowed_fields)}"
+            )
+        return values

mem0/embeddings/base.py

@@ -18,7 +18,7 @@ def __init__(self, config: Optional[BaseEmbedderConfig] = None):
             self.config = config
 
     @abstractmethod
-    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]]):
+    async def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]]):
         """
         Get the embedding for the given text.
 

mem0/embeddings/configs.py

@@ -24,6 +24,7 @@ def validate_config(cls, v, values):
             "lmstudio",
             "langchain",
             "aws_bedrock",
+            "ockam",
         ]:
             return v
         else:

mem0/embeddings/ockam.py

@@ -0,0 +1,33 @@
+from typing import Literal, Optional
+
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.base import EmbeddingBase
+
+try:
+    import litellm
+except ImportError:
+    raise ImportError("The 'litellm' library is required. Please install it using 'pip install litellm'.")
+
+
+class OckamEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        if not self.config.ockam_model:
+            raise ValueError("'ockam_model' is required for 'OckamEmbedding'.")
+
+    async def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Get the embedding for the given text using Ollama.
+
+        Args:
+            text (str): The text to embed.
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
+        Returns:
+            list: The embedding vector.
+        """
+        ockam_model = self.config.ockam_model
+        router = ockam_model.router()
+        kwargs = ockam_model.kwargs
+        response = await router.aembedding(model=ockam_model.name, input=text, **kwargs)
+        return response.data[0]["embedding"]

mem0/llms/base.py

@@ -17,7 +17,7 @@ def __init__(self, config: Optional[BaseLlmConfig] = None):
             self.config = config
 
     @abstractmethod
-    def generate_response(self, messages, tools: Optional[List[Dict]] = None, tool_choice: str = "auto"):
+    async def generate_response(self, messages, tools: Optional[List[Dict]] = None, tool_choice: str = "auto"):
         """
         Generate a response based on the given messages.
 

mem0/llms/configs.py

@@ -28,6 +28,7 @@ def validate_config(cls, v, values):
             "lmstudio",
             "vllm",
             "langchain",
+            "ockam",
         ):
             return v
         else:

mem0/llms/ockam.py

@@ -0,0 +1,93 @@
+import json
+from typing import Dict, List, Optional
+
+try:
+    import litellm
+except ImportError:
+    raise ImportError("The 'litellm' library is required. Please install it using 'pip install litellm'.")
+
+from mem0.configs.llms.base import BaseLlmConfig
+from mem0.llms.base import LLMBase
+
+
+class OckamLLM(LLMBase):
+    def __init__(self, config: Optional[BaseLlmConfig] = None):
+        super().__init__(config)
+
+        if not self.config.ockam_model:
+            raise ValueError("'ockam_model' is required for 'OckamLLM'.")
+
+    def _parse_response(self, response, tools):
+        """
+        Process the response based on whether tools are used or not.
+
+        Args:
+            response: The raw response from API.
+            tools: The list of tools provided in the request.
+
+        Returns:
+            str or dict: The processed response.
+        """
+        if tools:
+            processed_response = {
+                "content": response.choices[0].message.content,
+                "tool_calls": [],
+            }
+
+            if response.choices[0].message.tool_calls:
+                for tool_call in response.choices[0].message.tool_calls:
+                    processed_response["tool_calls"].append(
+                        {
+                            "name": tool_call.function.name,
+                            "arguments": json.loads(tool_call.function.arguments),
+                        }
+                    )
+
+            return processed_response
+        else:
+            return response.choices[0].message.content
+
+    async def generate_response(
+        self,
+        messages: List[Dict[str, str]],
+        response_format=None,
+        tools: Optional[List[Dict]] = None,
+        tool_choice: str = "auto",
+    ):
+        """
+        Generate a response based on the given messages using Litellm.
+
+        Args:
+            messages (list): List of message dicts containing 'role' and 'content'.
+            response_format (str or object, optional): Format of the response. Defaults to "text".
+            tools (list, optional): List of tools that the model can call. Defaults to None.
+            tool_choice (str, optional): Tool choice method. Defaults to "auto".
+
+        Returns:
+            str: The generated response.
+        """
+
+        # FIXME
+        # if not litellm.supports_function_calling(self.config.model):
+        #     raise ValueError(f"Model '{self.config.model}' in litellm does not support function calling.")
+
+        ockam_model = self.config.ockam_model
+
+        params = {
+            "model": ockam_model.name,
+            "messages": messages,
+            "temperature": self.config.temperature,
+            "max_tokens": self.config.max_tokens,
+            "top_p": self.config.top_p,
+        }
+        if response_format:
+            params["response_format"] = response_format
+        if tools:  # TODO: Remove tools if no issues found with new memory addition logic
+            params["tools"] = tools
+            params["tool_choice"] = tool_choice
+
+        router = ockam_model.router()
+        kwargs = {**ockam_model.kwargs, **params}
+
+        response = await router.acompletion(**kwargs)
+        return self._parse_response(response, tools)