From 8ca9a0f2c006706f6f0b683e1b8d6c715d6eda86 Mon Sep 17 00:00:00 2001 From: Kartik Date: Mon, 3 Aug 2026 20:36:28 +0530 Subject: [PATCH] fix: OSS Python SDK hygiene batch, 9 small bug fixes (#6770) --- mem0/embeddings/fastembed.py | 6 +- mem0/embeddings/huggingface.py | 2 +- mem0/embeddings/ollama.py | 14 +--- mem0/memory/utils.py | 8 +- mem0/reranker/cohere_reranker.py | 9 ++- mem0/reranker/huggingface_reranker.py | 14 ++-- .../reranker/sentence_transformer_reranker.py | 16 ++-- mem0/reranker/zero_entropy_reranker.py | 9 ++- mem0/utils/factory.py | 6 +- mem0/vector_stores/mongodb.py | 1 - mem0/vector_stores/vertex_ai_vector_search.py | 1 - .../embeddings/test_huggingface_embeddings.py | 14 +++- tests/memory/test_memory_utils.py | 26 ++++++- tests/rerankers/conftest.py | 32 ++++++++ .../test_reranker_fallback_no_mutation.py | 75 +++++++++++++++++++ .../rerankers/test_reranker_fallback_topk.py | 45 +---------- tests/utils/test_factory.py | 9 +++ 17 files changed, 201 insertions(+), 86 deletions(-) create mode 100644 tests/rerankers/test_reranker_fallback_no_mutation.py diff --git a/mem0/embeddings/fastembed.py b/mem0/embeddings/fastembed.py index fadaedb8c..9434e5089 100644 --- a/mem0/embeddings/fastembed.py +++ b/mem0/embeddings/fastembed.py @@ -1,7 +1,7 @@ -from typing import Optional, Literal +from typing import Literal, Optional -from mem0.embeddings.base import EmbeddingBase from mem0.configs.embeddings.base import BaseEmbedderConfig +from mem0.embeddings.base import EmbeddingBase try: from fastembed import TextEmbedding @@ -29,4 +29,4 @@ class FastEmbedEmbedding(EmbeddingBase): """ text = text.replace("\n", " ") embeddings = list(self.dense_model.embed(text)) - return embeddings[0] + return embeddings[0].tolist() diff --git a/mem0/embeddings/huggingface.py b/mem0/embeddings/huggingface.py index 230fb7e20..7e175f36e 100644 --- a/mem0/embeddings/huggingface.py +++ b/mem0/embeddings/huggingface.py @@ -17,7 +17,7 @@ class HuggingFaceEmbedding(EmbeddingBase): super().__init__(config) if self.config.huggingface_base_url: - self.client = OpenAI(base_url=self.config.huggingface_base_url) + self.client = OpenAI(base_url=self.config.huggingface_base_url, api_key=self.config.api_key) self.config.model = self.config.model or "tei" else: self.config.model = self.config.model or "multi-qa-MiniLM-L6-cos-v1" diff --git a/mem0/embeddings/ollama.py b/mem0/embeddings/ollama.py index c5a74e36e..55c164f2c 100644 --- a/mem0/embeddings/ollama.py +++ b/mem0/embeddings/ollama.py @@ -1,5 +1,3 @@ -import subprocess -import sys from typing import Literal, Optional from mem0.configs.embeddings.base import BaseEmbedderConfig @@ -8,17 +6,7 @@ from mem0.embeddings.base import EmbeddingBase try: from ollama import Client except ImportError: - user_input = input("The 'ollama' library is required. Install it now? [y/N]: ") - if user_input.lower() == "y": - try: - subprocess.check_call([sys.executable, "-m", "pip", "install", "ollama"]) - from ollama import Client - except subprocess.CalledProcessError: - print("Failed to install 'ollama'. Please install it manually using 'pip install ollama'.") - sys.exit(1) - else: - print("The required 'ollama' library is not installed.") - sys.exit(1) + raise ImportError("The 'ollama' library is required. Please install it using 'pip install ollama'.") class OllamaEmbedding(EmbeddingBase): diff --git a/mem0/memory/utils.py b/mem0/memory/utils.py index dd7b1e4cc..b22d60a27 100644 --- a/mem0/memory/utils.py +++ b/mem0/memory/utils.py @@ -121,6 +121,8 @@ def remove_code_blocks(content: str) -> str: - If a code block is detected, it returns only the inner content, stripping out the markers. - If no code block markers are found, the original content is returned as-is. """ + if content is None: + return "" pattern = r"^```[a-zA-Z0-9]*\n([\s\S]*?)\n```$" match = re.match(pattern, content.strip()) match_res=match.group(1).strip() if match else content.strip() @@ -213,8 +215,8 @@ def parse_vision_messages(messages, llm=None, vision_details="auto"): try: description = get_image_description(image_url, llm, vision_details) returned_messages.append({"role": role, "content": description}) - except Exception: - raise Exception(f"Error while downloading {image_url}.") + except Exception as e: + raise Exception(f"Error while downloading {image_url}.") from e else: # Regular text content returned_messages.append(msg) @@ -227,7 +229,7 @@ def process_telemetry_filters(filters): Process the telemetry filters """ if filters is None: - return {} + return [], {} encoded_ids = {} if "user_id" in filters: diff --git a/mem0/reranker/cohere_reranker.py b/mem0/reranker/cohere_reranker.py index 281fabcc6..6063ad0a7 100644 --- a/mem0/reranker/cohere_reranker.py +++ b/mem0/reranker/cohere_reranker.py @@ -1,6 +1,6 @@ import logging import os -from typing import List, Dict, Any +from typing import Any, Dict, List from mem0.reranker.base import BaseReranker @@ -84,7 +84,10 @@ class CohereReranker(BaseReranker): except Exception as e: # Fallback to original order if reranking fails logger.warning("Cohere reranking failed, falling back to original order: %s", e) + fallback_docs = [] for doc in documents: - doc['rerank_score'] = 0.0 + fallback_doc = doc.copy() + fallback_doc['rerank_score'] = 0.0 + fallback_docs.append(fallback_doc) final_top_k = top_k or self.config.top_k - return documents[:final_top_k] if final_top_k else documents \ No newline at end of file + return fallback_docs[:final_top_k] if final_top_k else fallback_docs \ No newline at end of file diff --git a/mem0/reranker/huggingface_reranker.py b/mem0/reranker/huggingface_reranker.py index be87d3dba..6af1bee2e 100644 --- a/mem0/reranker/huggingface_reranker.py +++ b/mem0/reranker/huggingface_reranker.py @@ -1,14 +1,15 @@ import logging -from typing import List, Dict, Any, Union +from typing import Any, Dict, List, Union + import numpy as np -from mem0.reranker.base import BaseReranker from mem0.configs.rerankers.base import BaseRerankerConfig from mem0.configs.rerankers.huggingface import HuggingFaceRerankerConfig +from mem0.reranker.base import BaseReranker try: - from transformers import AutoTokenizer, AutoModelForSequenceClassification import torch + from transformers import AutoModelForSequenceClassification, AutoTokenizer TRANSFORMERS_AVAILABLE = True except ImportError: TRANSFORMERS_AVAILABLE = False @@ -163,7 +164,10 @@ class HuggingFaceReranker(BaseReranker): except Exception as e: # Fallback to original order if reranking fails logger.warning("HuggingFace reranking failed, falling back to original order: %s", e) + fallback_docs = [] for doc in documents: - doc['rerank_score'] = 0.0 + fallback_doc = doc.copy() + fallback_doc['rerank_score'] = 0.0 + fallback_docs.append(fallback_doc) final_top_k = top_k or self.config.top_k - return documents[:final_top_k] if final_top_k else documents \ No newline at end of file + return fallback_docs[:final_top_k] if final_top_k else fallback_docs \ No newline at end of file diff --git a/mem0/reranker/sentence_transformer_reranker.py b/mem0/reranker/sentence_transformer_reranker.py index d891294f1..2b92f3e80 100644 --- a/mem0/reranker/sentence_transformer_reranker.py +++ b/mem0/reranker/sentence_transformer_reranker.py @@ -1,10 +1,13 @@ import logging -from typing import List, Dict, Any, Union +from typing import Any, Dict, List, Union + import numpy as np -from mem0.reranker.base import BaseReranker from mem0.configs.rerankers.base import BaseRerankerConfig -from mem0.configs.rerankers.sentence_transformer import SentenceTransformerRerankerConfig +from mem0.configs.rerankers.sentence_transformer import ( + SentenceTransformerRerankerConfig, +) +from mem0.reranker.base import BaseReranker try: from sentence_transformers import CrossEncoder @@ -108,7 +111,10 @@ class SentenceTransformerReranker(BaseReranker): except Exception as e: # Fallback to original order if reranking fails logger.warning("SentenceTransformer reranking failed, falling back to original order: %s", e) + fallback_docs = [] for doc in documents: - doc['rerank_score'] = 0.0 + fallback_doc = doc.copy() + fallback_doc['rerank_score'] = 0.0 + fallback_docs.append(fallback_doc) final_top_k = top_k or self.config.top_k - return documents[:final_top_k] if final_top_k else documents \ No newline at end of file + return fallback_docs[:final_top_k] if final_top_k else fallback_docs \ No newline at end of file diff --git a/mem0/reranker/zero_entropy_reranker.py b/mem0/reranker/zero_entropy_reranker.py index dcf71bfaf..68228c85d 100644 --- a/mem0/reranker/zero_entropy_reranker.py +++ b/mem0/reranker/zero_entropy_reranker.py @@ -1,6 +1,6 @@ import logging import os -from typing import List, Dict, Any +from typing import Any, Dict, List from mem0.reranker.base import BaseReranker @@ -95,7 +95,10 @@ class ZeroEntropyReranker(BaseReranker): except Exception as e: # Fallback to original order if reranking fails logger.warning("Zero Entropy reranking failed, falling back to original order: %s", e) + fallback_docs = [] for doc in documents: - doc['rerank_score'] = 0.0 + fallback_doc = doc.copy() + fallback_doc['rerank_score'] = 0.0 + fallback_docs.append(fallback_doc) final_top_k = top_k or self.config.top_k - return documents[:final_top_k] if final_top_k else documents \ No newline at end of file + return fallback_docs[:final_top_k] if final_top_k else fallback_docs \ No newline at end of file diff --git a/mem0/utils/factory.py b/mem0/utils/factory.py index 866a4da73..30a1079bb 100644 --- a/mem0/utils/factory.py +++ b/mem0/utils/factory.py @@ -19,7 +19,9 @@ from mem0.configs.rerankers.base import BaseRerankerConfig from mem0.configs.rerankers.cohere import CohereRerankerConfig from mem0.configs.rerankers.huggingface import HuggingFaceRerankerConfig from mem0.configs.rerankers.llm import LLMRerankerConfig -from mem0.configs.rerankers.sentence_transformer import SentenceTransformerRerankerConfig +from mem0.configs.rerankers.sentence_transformer import ( + SentenceTransformerRerankerConfig, +) from mem0.configs.rerankers.zero_entropy import ZeroEntropyRerankerConfig from mem0.embeddings.mock import MockEmbeddings @@ -86,7 +88,7 @@ class LlmFactory: config = config_class(**kwargs) elif isinstance(config, dict): # Merge dict config with kwargs - config.update(kwargs) + config = {**config, **kwargs} config = config_class(**config) elif isinstance(config, BaseLlmConfig): # Convert base config to provider-specific config if needed diff --git a/mem0/vector_stores/mongodb.py b/mem0/vector_stores/mongodb.py index a7738c6f2..53c0b7de6 100644 --- a/mem0/vector_stores/mongodb.py +++ b/mem0/vector_stores/mongodb.py @@ -15,7 +15,6 @@ except ImportError: from mem0.vector_stores.base import VectorStoreBase logger = logging.getLogger(__name__) -logging.basicConfig(level=logging.INFO) _DRIVER_METADATA = DriverInfo(name="Mem0", version=version("mem0ai")) diff --git a/mem0/vector_stores/vertex_ai_vector_search.py b/mem0/vector_stores/vertex_ai_vector_search.py index 8ba1ceea7..0954bf90c 100644 --- a/mem0/vector_stores/vertex_ai_vector_search.py +++ b/mem0/vector_stores/vertex_ai_vector_search.py @@ -22,7 +22,6 @@ from mem0.configs.vector_stores.vertex_ai_vector_search import ( from mem0.vector_stores.base import VectorStoreBase # Configure logging -logging.basicConfig(level=logging.DEBUG) logger = logging.getLogger(__name__) diff --git a/tests/embeddings/test_huggingface_embeddings.py b/tests/embeddings/test_huggingface_embeddings.py index 91f3ffe40..ed2ee61e8 100644 --- a/tests/embeddings/test_huggingface_embeddings.py +++ b/tests/embeddings/test_huggingface_embeddings.py @@ -94,7 +94,7 @@ def test_embed_with_huggingface_base_url(): embedder = HuggingFaceEmbedding(config) result = embedder.embed("Hello from custom endpoint") - mock_openai.assert_called_once_with(base_url="http://localhost:8080") + mock_openai.assert_called_once_with(base_url="http://localhost:8080", api_key=None) mock_client.embeddings.create.assert_called_once_with( input="Hello from custom endpoint", model="my-custom-model", @@ -103,6 +103,18 @@ def test_embed_with_huggingface_base_url(): assert result == [0.1, 0.2, 0.3] +def test_embed_with_huggingface_base_url_forwards_api_key(): + config = BaseEmbedderConfig( + huggingface_base_url="http://localhost:8080", + model="my-custom-model", + api_key="tei-token", + ) + with patch("mem0.embeddings.huggingface.OpenAI") as mock_openai: + HuggingFaceEmbedding(config) + + mock_openai.assert_called_once_with(base_url="http://localhost:8080", api_key="tei-token") + + def test_embed_batch_sentence_transformer(mock_sentence_transformer): config = BaseEmbedderConfig() embedder = HuggingFaceEmbedding(config) diff --git a/tests/memory/test_memory_utils.py b/tests/memory/test_memory_utils.py index 49c55c5cb..3bab79429 100644 --- a/tests/memory/test_memory_utils.py +++ b/tests/memory/test_memory_utils.py @@ -1,9 +1,12 @@ -import pytest from unittest.mock import Mock +import pytest + from mem0.memory.utils import ( parse_messages, parse_vision_messages, + process_telemetry_filters, + remove_code_blocks, remove_spaces_from_entities, sanitize_relationship_for_cypher, ) @@ -114,6 +117,17 @@ class TestParseVisionMessages: parse_vision_messages(messages, llm=mock_llm) mock_llm.generate_response.assert_not_called() + def test_download_failure_preserves_original_exception(self): + mock_llm = Mock() + mock_llm.generate_response.side_effect = ValueError("network down") + messages = [ + {"role": "user", "content": {"type": "image_url", "image_url": {"url": "https://example.com/img.png"}}} + ] + with pytest.raises(Exception) as exc_info: + parse_vision_messages(messages, llm=mock_llm) + assert isinstance(exc_info.value.__cause__, ValueError) + assert "network down" in str(exc_info.value.__cause__) + class TestRemoveSpacesFromEntities: """ @@ -168,3 +182,13 @@ class TestRemoveSpacesFromEntities: f = remove_spaces_from_entities([dict(base)], sanitize_relationship=False)[0]["relationship"] assert t == sanitize_relationship_for_cypher("a/b") assert f == "a/b" + + +class TestProcessTelemetryFilters: + def test_none_filters_returns_empty_list_and_dict(self): + assert process_telemetry_filters(None) == ([], {}) + + +class TestRemoveCodeBlocks: + def test_none_content_returns_empty_string(self): + assert remove_code_blocks(None) == "" diff --git a/tests/rerankers/conftest.py b/tests/rerankers/conftest.py index 258d959cb..edae4f581 100644 --- a/tests/rerankers/conftest.py +++ b/tests/rerankers/conftest.py @@ -1,3 +1,5 @@ +import sys +from types import ModuleType from unittest.mock import MagicMock, patch import pytest @@ -9,3 +11,33 @@ def mock_llm(): mock_llm_instance = MagicMock() mock_factory.create.return_value = mock_llm_instance yield mock_factory, mock_llm_instance + + +@pytest.fixture +def mock_cohere(monkeypatch): + """Provide a fake ``cohere`` module so CohereReranker imports/constructs.""" + fake_cohere = ModuleType("cohere") + fake_client = MagicMock() + fake_cohere.Client = MagicMock(return_value=fake_client) + monkeypatch.setitem(sys.modules, "cohere", fake_cohere) + + import mem0.reranker.cohere_reranker as cohere_reranker + + monkeypatch.setattr(cohere_reranker, "cohere", fake_cohere, raising=False) + monkeypatch.setattr(cohere_reranker, "COHERE_AVAILABLE", True, raising=False) + return cohere_reranker, fake_client + + +@pytest.fixture +def mock_zero_entropy(monkeypatch): + """Provide a fake ``zeroentropy`` module so ZeroEntropyReranker imports.""" + fake_module = ModuleType("zeroentropy") + fake_client = MagicMock() + fake_module.ZeroEntropy = MagicMock(return_value=fake_client) + monkeypatch.setitem(sys.modules, "zeroentropy", fake_module) + + import mem0.reranker.zero_entropy_reranker as zero_entropy_reranker + + monkeypatch.setattr(zero_entropy_reranker, "ZeroEntropy", fake_module.ZeroEntropy, raising=False) + monkeypatch.setattr(zero_entropy_reranker, "ZERO_ENTROPY_AVAILABLE", True, raising=False) + return zero_entropy_reranker, fake_client diff --git a/tests/rerankers/test_reranker_fallback_no_mutation.py b/tests/rerankers/test_reranker_fallback_no_mutation.py new file mode 100644 index 000000000..3a33ded60 --- /dev/null +++ b/tests/rerankers/test_reranker_fallback_no_mutation.py @@ -0,0 +1,75 @@ +"""Regression tests pinning that the reranker fallback path never mutates caller-owned documents.""" + +from unittest.mock import MagicMock, patch + +from mem0.configs.rerankers.cohere import CohereRerankerConfig +from mem0.configs.rerankers.huggingface import HuggingFaceRerankerConfig +from mem0.configs.rerankers.sentence_transformer import ( + SentenceTransformerRerankerConfig, +) +from mem0.configs.rerankers.zero_entropy import ZeroEntropyRerankerConfig +from mem0.reranker.huggingface_reranker import HuggingFaceReranker +from mem0.reranker.sentence_transformer_reranker import SentenceTransformerReranker + + +def _docs(n): + return [{"memory": f"doc{i}"} for i in range(n)] + + +class TestCohereFallbackNoMutation: + def test_fallback_does_not_mutate_original_documents(self, mock_cohere): + module, fake_client = mock_cohere + fake_client.rerank.side_effect = RuntimeError("API error") + + reranker = module.CohereReranker(CohereRerankerConfig(api_key="test-key")) + documents = _docs(3) + result = reranker.rerank("query", documents) + + assert all("rerank_score" not in doc for doc in documents) + assert all(doc["rerank_score"] == 0.0 for doc in result) + + +class TestZeroEntropyFallbackNoMutation: + def test_fallback_does_not_mutate_original_documents(self, mock_zero_entropy): + module, fake_client = mock_zero_entropy + fake_client.models.rerank.side_effect = RuntimeError("API error") + + reranker = module.ZeroEntropyReranker(ZeroEntropyRerankerConfig(api_key="test-key")) + documents = _docs(3) + result = reranker.rerank("query", documents) + + assert all("rerank_score" not in doc for doc in documents) + assert all(doc["rerank_score"] == 0.0 for doc in result) + + +class TestHuggingFaceFallbackNoMutation: + def test_fallback_does_not_mutate_original_documents(self): + with ( + patch("mem0.reranker.huggingface_reranker.AutoTokenizer") as mock_tokenizer_cls, + patch("mem0.reranker.huggingface_reranker.AutoModelForSequenceClassification") as mock_model_cls, + ): + mock_tokenizer = MagicMock(side_effect=RuntimeError("tokenizer error")) + mock_tokenizer_cls.from_pretrained.return_value = mock_tokenizer + mock_model_cls.from_pretrained.return_value = MagicMock() + + reranker = HuggingFaceReranker(HuggingFaceRerankerConfig()) + documents = _docs(3) + result = reranker.rerank("query", documents) + + assert all("rerank_score" not in doc for doc in documents) + assert all(doc["rerank_score"] == 0.0 for doc in result) + + +class TestSentenceTransformerFallbackNoMutation: + def test_fallback_does_not_mutate_original_documents(self): + with patch("mem0.reranker.sentence_transformer_reranker.CrossEncoder") as mock_cross_encoder_cls: + mock_model = MagicMock() + mock_model.predict.side_effect = RuntimeError("predict error") + mock_cross_encoder_cls.return_value = mock_model + + reranker = SentenceTransformerReranker(SentenceTransformerRerankerConfig()) + documents = _docs(3) + result = reranker.rerank("query", documents) + + assert all("rerank_score" not in doc for doc in documents) + assert all(doc["rerank_score"] == 0.0 for doc in result) diff --git a/tests/rerankers/test_reranker_fallback_topk.py b/tests/rerankers/test_reranker_fallback_topk.py index a2767e8ae..aa6410a32 100644 --- a/tests/rerankers/test_reranker_fallback_topk.py +++ b/tests/rerankers/test_reranker_fallback_topk.py @@ -1,52 +1,9 @@ -"""Regression tests for the reranker fallback path honoring ``config.top_k``. - -When the underlying rerank call fails, the reranker falls back to returning the -documents in their original order. That fallback must still respect the -configured ``top_k`` limit, exactly like the success path does. The HuggingFace -and SentenceTransformer rerankers already behave this way; these tests pin the -same contract for the Cohere and ZeroEntropy rerankers. -""" - -import sys -from types import ModuleType -from unittest.mock import MagicMock - -import pytest +"""Regression tests pinning that the reranker fallback path still honors ``config.top_k``.""" from mem0.configs.rerankers.cohere import CohereRerankerConfig from mem0.configs.rerankers.zero_entropy import ZeroEntropyRerankerConfig -@pytest.fixture -def mock_cohere(monkeypatch): - """Provide a fake ``cohere`` module so CohereReranker imports/constructs.""" - fake_cohere = ModuleType("cohere") - fake_client = MagicMock() - fake_cohere.Client = MagicMock(return_value=fake_client) - monkeypatch.setitem(sys.modules, "cohere", fake_cohere) - - import mem0.reranker.cohere_reranker as cohere_reranker - - monkeypatch.setattr(cohere_reranker, "cohere", fake_cohere, raising=False) - monkeypatch.setattr(cohere_reranker, "COHERE_AVAILABLE", True, raising=False) - return cohere_reranker, fake_client - - -@pytest.fixture -def mock_zero_entropy(monkeypatch): - """Provide a fake ``zeroentropy`` module so ZeroEntropyReranker imports.""" - fake_module = ModuleType("zeroentropy") - fake_client = MagicMock() - fake_module.ZeroEntropy = MagicMock(return_value=fake_client) - monkeypatch.setitem(sys.modules, "zeroentropy", fake_module) - - import mem0.reranker.zero_entropy_reranker as zero_entropy_reranker - - monkeypatch.setattr(zero_entropy_reranker, "ZeroEntropy", fake_module.ZeroEntropy, raising=False) - monkeypatch.setattr(zero_entropy_reranker, "ZERO_ENTROPY_AVAILABLE", True, raising=False) - return zero_entropy_reranker, fake_client - - def _docs(n): return [{"memory": f"doc{i}"} for i in range(n)] diff --git a/tests/utils/test_factory.py b/tests/utils/test_factory.py index 0f07bd2d5..cbb63890b 100644 --- a/tests/utils/test_factory.py +++ b/tests/utils/test_factory.py @@ -51,3 +51,12 @@ def test_base_to_provider_without_reasoning_fields_still_builds(): assert isinstance(built, AnthropicConfig) assert built.model == "claude-3-5-sonnet-20240620" + + +def test_dict_config_not_mutated_by_kwargs(): + config = {"model": "gpt-4o-mini"} + + with patch("mem0.utils.factory.load_class", return_value=Mock(return_value=Mock())): + LlmFactory.create("openai", config, api_key="secret") + + assert config == {"model": "gpt-4o-mini"}