fix: OSS Python SDK hygiene batch, 9 small bug fixes (#6770)

This commit is contained in:
Kartik
2026-08-03 20:36:28 +05:30
committed by GitHub
parent 5f77d86caf
commit 8ca9a0f2c0
17 changed files with 201 additions and 86 deletions
+3 -3
View File
@@ -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()
+1 -1
View File
@@ -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"
+1 -13
View File
@@ -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):
+5 -3
View File
@@ -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:
+6 -3
View File
@@ -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
return fallback_docs[:final_top_k] if final_top_k else fallback_docs
+9 -5
View File
@@ -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
return fallback_docs[:final_top_k] if final_top_k else fallback_docs
+11 -5
View File
@@ -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
return fallback_docs[:final_top_k] if final_top_k else fallback_docs
+6 -3
View File
@@ -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
return fallback_docs[:final_top_k] if final_top_k else fallback_docs
+4 -2
View File
@@ -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
-1
View File
@@ -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"))
@@ -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__)
@@ -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)
+25 -1
View File
@@ -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) == ""
+32
View File
@@ -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
@@ -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)
+1 -44
View File
@@ -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)]
+9
View File
@@ -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"}