Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/semantic-search.md
Original file line number Diff line number Diff line change
Expand Up @@ -394,6 +394,7 @@ All settings use the `BASIC_MEMORY_` environment prefix:
| `reranker_model` | `BASIC_MEMORY_RERANKER_MODEL` | `jinaai/jina-reranker-v1-tiny-en` | Model identifier. LiteLLM requires explicit `provider/model` routing. |
| `reranker_candidates` | `BASIC_MEMORY_RERANKER_CANDIDATES` | `20` | Number of leading retrieval results rescored on every page. Larger values can improve recall but increase latency and provider usage. |
| `reranker_max_document_chars` | `BASIC_MEMORY_RERANKER_MAX_DOCUMENT_CHARS` | `0` | Maximum characters sent per candidate. `0` sends the full matched text; a positive cap bounds latency and request size. |
| `reranker_timeout` | `BASIC_MEMORY_RERANKER_TIMEOUT` | `30.0` | Maximum seconds for each LiteLLM rerank request. FastEmbed runs locally and ignores this setting. |
| `reranker_api_base` | `BASIC_MEMORY_RERANKER_API_BASE` | Unset | Optional custom endpoint for the LiteLLM provider. |
| `reranker_api_key` | `BASIC_MEMORY_RERANKER_API_KEY` | Unset | Optional credential passed directly to LiteLLM. When unset, LiteLLM resolves provider credentials from its normal environment variables. |

Expand Down
6 changes: 6 additions & 0 deletions src/basic_memory/config_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -450,6 +450,12 @@ def __init__(self, **data: Any) -> None: ...
"most-relevant matched chunk leads the text, so a modest cap keeps most of the signal.",
ge=0,
)
reranker_timeout: float = Field(
default=30.0,
description="Maximum seconds allowed for each LiteLLM rerank request. "
"FastEmbed runs locally and ignores this setting.",
gt=0,
)
reranker_api_base: str | None = Field(
default=None,
description="Optional custom API base URL for the litellm reranker provider "
Expand Down
14 changes: 13 additions & 1 deletion src/basic_memory/mcp/tools/chatgpt_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
from basic_memory.mcp.client_info import is_openai_mcp_client
from basic_memory.mcp.server import mcp
from basic_memory.mcp.tools.read_note import read_note
from basic_memory.mcp.tools.search import search_notes
from basic_memory.mcp.tools.search import _SERVICE_UNAVAILABLE_HEADING, search_notes
from basic_memory.schemas.search import SearchResponse, SearchResult


Expand Down Expand Up @@ -182,6 +182,18 @@ async def search(
)

if isinstance(results, str):
# Trigger: search_notes translated an API 503 into its retryable outage response.
# Why: OpenAI clients need to distinguish a temporary provider failure from an
# internal adapter error before deciding whether to retry.
# Outcome: preserve the retry signal in the Actions-compatible error payload.
if results.startswith(_SERVICE_UNAVAILABLE_HEADING):
return _text_content(
{
"results": [],
"error": "Search temporarily unavailable",
"error_message": "Search temporarily unavailable, retry shortly",
}
)
logger.warning(f"Search failed with error: {results[:100]}...")
search_results = {
"results": [],
Expand Down
3 changes: 2 additions & 1 deletion src/basic_memory/repository/fastembed_rerank_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
from loguru import logger
from requests import exceptions as requests_exceptions

from basic_memory.config_models import DEFAULT_FASTEMBED_RERANK_MODEL
from basic_memory.repository.rerank_provider import validate_rerank_scores
from basic_memory.repository.semantic_errors import (
RerankProviderContractError,
Expand Down Expand Up @@ -68,7 +69,7 @@ class FastEmbedRerankProvider:

def __init__(
self,
model_name: str = "Xenova/ms-marco-MiniLM-L-6-v2",
model_name: str = DEFAULT_FASTEMBED_RERANK_MODEL,
*,
cache_dir: str | None = None,
threads: int | None = None,
Expand Down
16 changes: 10 additions & 6 deletions src/basic_memory/repository/rerank_provider_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,12 +20,12 @@
from basic_memory.repository.rerank_provider import RerankProvider

# Key on the fields that change the loaded provider's identity: provider, model,
# (for the litellm path) the endpoint/key routing, and the resolved cache dir. The
# cache dir matters because the fastembed provider is constructed with it — omitting
# it (as an earlier version did) lets two configs with different cache dirs share one
# singleton pointing at the wrong directory, the #741/#872 class of bug the embedding
# factory guards against. CPU-derived thread counts stay out (they drift per call).
type RerankCacheKey = tuple[str, str, str | None, str | None, str]
# (for the litellm path) the endpoint/key routing and timeout, and the resolved cache
# dir. The cache dir matters because the fastembed provider is constructed with it —
# omitting it (as an earlier version did) lets two configs with different cache dirs
# share one singleton pointing at the wrong directory, the #741/#872 class of bug the
# embedding factory guards against. CPU-derived thread counts stay out (they drift per call).
type RerankCacheKey = tuple[str, str, str | None, str | None, float | None, str]

_RERANK_PROVIDER_CACHE: dict[RerankCacheKey, RerankProvider] = {}
_RERANK_PROVIDER_CACHE_LOCK = Lock()
Expand All @@ -35,14 +35,17 @@ def _rerank_cache_key(app_config: BasicMemoryConfig) -> RerankCacheKey:
provider_name = app_config.reranker_provider.strip().lower()
api_base_digest = None
api_key_digest = None
timeout = None
if provider_name == "litellm":
api_base_digest = _sensitive_value_digest(app_config.reranker_api_base)
api_key_digest = _sensitive_value_digest(app_config.reranker_api_key)
timeout = app_config.reranker_timeout
return (
provider_name,
app_config.reranker_model,
api_base_digest,
api_key_digest,
timeout,
_resolve_cache_dir(app_config),
)

Expand Down Expand Up @@ -91,6 +94,7 @@ def create_rerank_provider(app_config: BasicMemoryConfig) -> RerankProvider | No
model_name=app_config.reranker_model,
api_key=app_config.reranker_api_key,
api_base=app_config.reranker_api_base,
timeout=app_config.reranker_timeout,
)
else:
raise ValueError(f"Unsupported reranker provider: {provider_name}")
Expand Down
25 changes: 24 additions & 1 deletion tests/mcp/tools/test_chatgpt_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,29 @@ async def fake_search_notes_fn(*args, **kwargs):
assert "error_details" in content


@pytest.mark.asyncio
async def test_search_retryable_outage_returns_explicit_retry_message(
monkeypatch, client, test_project, context_state
):
"""The search_notes 503 shape remains retryable through the ChatGPT adapter."""
import basic_memory.mcp.tools.chatgpt_tools as chatgpt_tools

async def fake_search_notes_fn(*args, **kwargs):
return f"{chatgpt_tools._SERVICE_UNAVAILABLE_HEADING}\n\nReranker temporarily unavailable"

monkeypatch.setattr(chatgpt_tools, "search_notes", fake_search_notes_fn)

context = await _openai_mcp_context(context_state)
result = await chatgpt_tools.search("retryable query", context=context)

content = json.loads(result[0]["text"])
assert content == {
"results": [],
"error": "Search temporarily unavailable",
"error_message": "Search temporarily unavailable, retry shortly",
}


@pytest.mark.asyncio
async def test_search_uses_dynamic_default_search_type(
monkeypatch, client, test_project, context_state
Expand Down Expand Up @@ -315,7 +338,7 @@ async def boom(*args, **kwargs):
assert isinstance(result, list)
content = json.loads(result[0]["text"])
assert content["error"] == "Internal search error"
assert "error_message" in content
assert content["error_message"] == "boom"


@pytest.mark.asyncio
Expand Down
7 changes: 7 additions & 0 deletions tests/repository/test_fastembed_rerank_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import pytest
from requests import Response, exceptions as requests_exceptions

from basic_memory.config_models import DEFAULT_FASTEMBED_RERANK_MODEL
from basic_memory.repository.fastembed_rerank_provider import FastEmbedRerankProvider
from basic_memory.repository.semantic_errors import (
RerankProviderContractError,
Expand Down Expand Up @@ -49,6 +50,12 @@ def _http_error(status_code: int) -> requests_exceptions.HTTPError:
return requests_exceptions.HTTPError(f"HTTP {status_code}", response=response)


def test_constructor_uses_configured_default_model():
provider = FastEmbedRerankProvider()

assert provider.model_name == DEFAULT_FASTEMBED_RERANK_MODEL


@pytest.mark.asyncio
async def test_lazy_loads_once_and_reuses_model(monkeypatch):
_install_stub(monkeypatch)
Expand Down
Loading
Loading