diff --git a/docs/components/llms/models/aws_bedrock.mdx b/docs/components/llms/models/aws_bedrock.mdx index 5c1b2501..7578524a 100644 --- a/docs/components/llms/models/aws_bedrock.mdx +++ b/docs/components/llms/models/aws_bedrock.mdx @@ -24,7 +24,7 @@ config = { "config": { "model": "arn:aws:bedrock:us-east-1:123456789012:model/your-model-name", "temperature": 0.2, - "max_tokens": 1500, + "max_tokens": 2000, } } } diff --git a/docs/components/llms/models/deepseek.mdx b/docs/components/llms/models/deepseek.mdx index 65626c11..77f226ab 100644 --- a/docs/components/llms/models/deepseek.mdx +++ b/docs/components/llms/models/deepseek.mdx @@ -19,7 +19,7 @@ config = { "config": { "model": "deepseek-chat", # default model "temperature": 0.2, - "max_tokens": 1500, + "max_tokens": 2000, "top_p": 1.0 } } diff --git a/docs/components/llms/models/gemini.mdx b/docs/components/llms/models/gemini.mdx index f020a2dd..a40c2650 100644 --- a/docs/components/llms/models/gemini.mdx +++ b/docs/components/llms/models/gemini.mdx @@ -19,7 +19,7 @@ config = { "config": { "model": "gemini-1.5-flash-latest", "temperature": 0.2, - "max_tokens": 1500, + "max_tokens": 2000, } } } diff --git a/docs/components/llms/models/google_AI.mdx b/docs/components/llms/models/google_AI.mdx index 1546e332..b1c22b1e 100644 --- a/docs/components/llms/models/google_AI.mdx +++ b/docs/components/llms/models/google_AI.mdx @@ -19,7 +19,7 @@ config = { "config": { "model": "gemini/gemini-pro", "temperature": 0.2, - "max_tokens": 1500, + "max_tokens": 2000, } } } diff --git a/docs/components/llms/models/groq.mdx b/docs/components/llms/models/groq.mdx index d27aa17b..b0b7531d 100644 --- a/docs/components/llms/models/groq.mdx +++ b/docs/components/llms/models/groq.mdx @@ -17,7 +17,7 @@ config = { "config": { "model": "mixtral-8x7b-32768", "temperature": 0.1, - "max_tokens": 1000, + "max_tokens": 2000, } } } diff --git a/docs/components/llms/models/litellm.mdx b/docs/components/llms/models/litellm.mdx index 4a6f47cd..00b5afb0 100644 --- a/docs/components/llms/models/litellm.mdx +++ b/docs/components/llms/models/litellm.mdx @@ -14,7 +14,7 @@ config = { "config": { "model": "gpt-4o-mini", "temperature": 0.2, - "max_tokens": 1500, + "max_tokens": 2000, } } } diff --git a/docs/components/llms/models/openai.mdx b/docs/components/llms/models/openai.mdx index 2002baa4..da3419c2 100644 --- a/docs/components/llms/models/openai.mdx +++ b/docs/components/llms/models/openai.mdx @@ -18,7 +18,7 @@ config = { "config": { "model": "gpt-4o", "temperature": 0.2, - "max_tokens": 1500, + "max_tokens": 2000, } } } diff --git a/docs/components/llms/models/together.mdx b/docs/components/llms/models/together.mdx index 2cac8d79..dd685a29 100644 --- a/docs/components/llms/models/together.mdx +++ b/docs/components/llms/models/together.mdx @@ -15,7 +15,7 @@ config = { "config": { "model": "mistralai/Mixtral-8x7B-Instruct-v0.1", "temperature": 0.2, - "max_tokens": 1500, + "max_tokens": 2000, } } } diff --git a/docs/components/llms/models/xAI.mdx b/docs/components/llms/models/xAI.mdx index c9e06125..d7a365f5 100644 --- a/docs/components/llms/models/xAI.mdx +++ b/docs/components/llms/models/xAI.mdx @@ -21,7 +21,7 @@ config = { "config": { "model": "grok-2-latest", "temperature": 0.1, - "max_tokens": 1000, + "max_tokens": 2000, } } } diff --git a/docs/examples/mem0-with-ollama.mdx b/docs/examples/mem0-with-ollama.mdx index de87ece6..664e00b6 100644 --- a/docs/examples/mem0-with-ollama.mdx +++ b/docs/examples/mem0-with-ollama.mdx @@ -37,7 +37,7 @@ config = { "config": { "model": "llama3.1:latest", "temperature": 0, - "max_tokens": 8000, + "max_tokens": 2000, "ollama_base_url": "http://localhost:11434", # Ensure this URL is correct }, }, diff --git a/docs/features/custom-prompts.mdx b/docs/features/custom-prompts.mdx index a7cbd427..b33de779 100644 --- a/docs/features/custom-prompts.mdx +++ b/docs/features/custom-prompts.mdx @@ -53,7 +53,7 @@ config = { "config": { "model": "gpt-4o", "temperature": 0.2, - "max_tokens": 1500, + "max_tokens": 2000, } }, "custom_prompt": custom_prompt, diff --git a/docs/integrations/llama-index.mdx b/docs/integrations/llama-index.mdx index c514dd66..ac0551f0 100644 --- a/docs/integrations/llama-index.mdx +++ b/docs/integrations/llama-index.mdx @@ -80,7 +80,7 @@ config = { "config": { "model": "gpt-4o", "temperature": 0.2, - "max_tokens": 1500, + "max_tokens": 2000, }, }, "embedder": { diff --git a/docs/open-source/graph_memory/overview.mdx b/docs/open-source/graph_memory/overview.mdx index 80e2b238..849b0530 100644 --- a/docs/open-source/graph_memory/overview.mdx +++ b/docs/open-source/graph_memory/overview.mdx @@ -81,7 +81,7 @@ config = { "config": { "model": "gpt-4o", "temperature": 0.2, - "max_tokens": 1500, + "max_tokens": 2000, } }, "graph_store": { diff --git a/mem0/configs/llms/base.py b/mem0/configs/llms/base.py index 36e32257..a681db33 100644 --- a/mem0/configs/llms/base.py +++ b/mem0/configs/llms/base.py @@ -16,7 +16,7 @@ class BaseLlmConfig(ABC): model: Optional[str] = None, temperature: float = 0.1, api_key: Optional[str] = None, - max_tokens: int = 3000, + max_tokens: int = 2000, top_p: float = 0.1, top_k: int = 1, # Openrouter specific @@ -48,7 +48,7 @@ class BaseLlmConfig(ABC): :type temperature: float, optional :param api_key: OpenAI API key to be use, defaults to None :type api_key: Optional[str], optional - :param max_tokens: Controls how many tokens are generated, defaults to 3000 + :param max_tokens: Controls how many tokens are generated, defaults to 2000 :type max_tokens: int, optional :param top_p: Controls the diversity of words. Higher values (closer to 1) make word selection more diverse, defaults to 1 diff --git a/mem0/embeddings/azure_openai.py b/mem0/embeddings/azure_openai.py index 8b0e89f0..5801932c 100644 --- a/mem0/embeddings/azure_openai.py +++ b/mem0/embeddings/azure_openai.py @@ -26,7 +26,7 @@ class AzureOpenAIEmbedding(EmbeddingBase): default_headers=default_headers, ) - def embed(self, text, memory_action:Optional[Literal["add", "search", "update"]] = None): + def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None): """ Get the embedding for the given text using OpenAI. diff --git a/mem0/embeddings/base.py b/mem0/embeddings/base.py index bb924c46..ed328128 100644 --- a/mem0/embeddings/base.py +++ b/mem0/embeddings/base.py @@ -18,7 +18,7 @@ class EmbeddingBase(ABC): self.config = config @abstractmethod - def embed(self, text, memory_action:Optional[Literal["add", "search", "update"]]): + def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]]): """ Get the embedding for the given text. diff --git a/mem0/embeddings/gemini.py b/mem0/embeddings/gemini.py index 7af86061..d549db78 100644 --- a/mem0/embeddings/gemini.py +++ b/mem0/embeddings/gemini.py @@ -18,7 +18,7 @@ class GoogleGenAIEmbedding(EmbeddingBase): genai.configure(api_key=api_key) - def embed(self, text, memory_action:Optional[Literal["add", "search", "update"]] = None): + def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None): """ Get the embedding for the given text using Google Generative AI. Args: diff --git a/mem0/embeddings/huggingface.py b/mem0/embeddings/huggingface.py index 88b4579c..913d5882 100644 --- a/mem0/embeddings/huggingface.py +++ b/mem0/embeddings/huggingface.py @@ -16,7 +16,7 @@ class HuggingFaceEmbedding(EmbeddingBase): self.config.embedding_dims = self.config.embedding_dims or self.model.get_sentence_embedding_dimension() - def embed(self, text, memory_action:Optional[Literal["add", "search", "update"]] = None): + def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None): """ Get the embedding for the given text using Hugging Face. diff --git a/mem0/embeddings/ollama.py b/mem0/embeddings/ollama.py index 2f824928..6f6d1181 100644 --- a/mem0/embeddings/ollama.py +++ b/mem0/embeddings/ollama.py @@ -39,7 +39,7 @@ class OllamaEmbedding(EmbeddingBase): if not any(model.get("name") == self.config.model for model in local_models): self.client.pull(self.config.model) - def embed(self, text, memory_action:Optional[Literal["add", "search", "update"]] = None): + def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None): """ Get the embedding for the given text using Ollama. diff --git a/mem0/embeddings/openai.py b/mem0/embeddings/openai.py index def289fa..a774544c 100644 --- a/mem0/embeddings/openai.py +++ b/mem0/embeddings/openai.py @@ -18,7 +18,7 @@ class OpenAIEmbedding(EmbeddingBase): base_url = self.config.openai_base_url or os.getenv("OPENAI_API_BASE") self.client = OpenAI(api_key=api_key, base_url=base_url) - def embed(self, text, memory_action:Optional[Literal["add", "search", "update"]] = None): + def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None): """ Get the embedding for the given text using OpenAI. diff --git a/mem0/embeddings/together.py b/mem0/embeddings/together.py index 86be9050..b3eca0b7 100644 --- a/mem0/embeddings/together.py +++ b/mem0/embeddings/together.py @@ -17,7 +17,7 @@ class TogetherEmbedding(EmbeddingBase): self.config.embedding_dims = self.config.embedding_dims or 768 self.client = Together(api_key=api_key) - def embed(self, text, memory_action:Optional[Literal["add", "search", "update"]] = None): + def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None): """ Get the embedding for the given text using OpenAI. diff --git a/mem0/embeddings/vertexai.py b/mem0/embeddings/vertexai.py index e309deca..380b7ea5 100644 --- a/mem0/embeddings/vertexai.py +++ b/mem0/embeddings/vertexai.py @@ -13,13 +13,13 @@ class VertexAIEmbedding(EmbeddingBase): self.config.model = self.config.model or "text-embedding-004" self.config.embedding_dims = self.config.embedding_dims or 256 - + self.embedding_types = { "add": self.config.memory_add_embedding_type or "RETRIEVAL_DOCUMENT", "update": self.config.memory_update_embedding_type or "RETRIEVAL_DOCUMENT", - "search": self.config.memory_search_embedding_type or "RETRIEVAL_QUERY" + "search": self.config.memory_search_embedding_type or "RETRIEVAL_QUERY", } - + credentials_path = self.config.vertex_credentials_json if credentials_path: @@ -31,7 +31,7 @@ class VertexAIEmbedding(EmbeddingBase): self.model = TextEmbeddingModel.from_pretrained(self.config.model) - def embed(self, text, memory_action:Optional[Literal["add", "search", "update"]] = None): + def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None): """ Get the embedding for the given text using Vertex AI. @@ -45,9 +45,9 @@ class VertexAIEmbedding(EmbeddingBase): if memory_action is not None: if memory_action not in self.embedding_types: raise ValueError(f"Invalid memory action: {memory_action}") - + embedding_type = self.embedding_types[memory_action] - + text_input = TextEmbeddingInput(text=text, task_type=embedding_type) embeddings = self.model.get_embeddings(texts=[text_input], output_dimensionality=self.config.embedding_dims) diff --git a/mem0/llms/groq.py b/mem0/llms/groq.py index e970a1ee..38a1c8a0 100644 --- a/mem0/llms/groq.py +++ b/mem0/llms/groq.py @@ -84,4 +84,4 @@ class GroqLLM(LLMBase): params["tool_choice"] = tool_choice response = self.client.chat.completions.create(**params) - return self._parse_response(response, tools) + return self._parse_response(response, tools) \ No newline at end of file diff --git a/mem0/llms/openai.py b/mem0/llms/openai.py index 38daf650..a9c302f8 100644 --- a/mem0/llms/openai.py +++ b/mem0/llms/openai.py @@ -63,7 +63,6 @@ class OpenAILLM(LLMBase): response_format=None, tools: Optional[List[Dict]] = None, tool_choice: str = "auto", - max_tokens: int = 100, ): """ Generate a response based on the given messages using OpenAI. @@ -81,7 +80,7 @@ class OpenAILLM(LLMBase): "model": self.config.model, "messages": messages, "temperature": self.config.temperature, - "max_tokens": max_tokens, + "max_tokens": self.config.max_tokens, "top_p": self.config.top_p, } diff --git a/tests/llms/test_azure_openai.py b/tests/llms/test_azure_openai.py index 079ef349..77cb88e9 100644 --- a/tests/llms/test_azure_openai.py +++ b/tests/llms/test_azure_openai.py @@ -128,4 +128,4 @@ def test_generate_with_http_proxies(default_headers): api_version=None, default_headers=default_headers, ) - mock_http_client.assert_called_once_with(proxies="http://testproxy.mem0.net:8000") + mock_http_client.assert_called_once_with(proxies="http://testproxy.mem0.net:8000") \ No newline at end of file