"""
MMORE API Client for multimodal document retrieval.
Provides interface to mmore service for advanced RAG capabilities.
Targets the MMORE v2 API. v2 ingests asynchronously: ``POST /v1/files``
returns ``202`` with a job id and indexes on a background queue, rather than
blocking until the document is searchable as v1 did. :meth:`MMOREClient.upload_file`
polls that job to completion so callers keep the synchronous semantics they
rely on. See ``services/mmore/frozen-v1/`` for the pre-v2 stack.
"""
import logging
import os
import time
from collections.abc import Callable
from pathlib import Path
from typing import Any
import requests
from langchain_core.documents import Document
logger = logging.getLogger(__name__)
# HTTP status codes
HTTP_OK = 200
HTTP_ACCEPTED = 202
HTTP_CONFLICT = 409
# Terminal states of a MMORE indexing job. Upstream's JobStatus enum is
# queued | processing | done | failed; the extra spellings are defensive, since
# the enum is not part of the documented HTTP contract. Anything NOT listed
# here is treated as "still running", so a state added upstream degrades to a
# timeout rather than a false success.
_JOB_SUCCEEDED = frozenset(
{"done", "completed", "complete", "succeeded", "success", "finished"}
)
_JOB_FAILED = frozenset({"failed", "error", "cancelled", "canceled"})
# Job polling
DEFAULT_INDEX_TIMEOUT = 900.0 # seconds; large scanned PDFs on CPU are slow
_POLL_INITIAL_INTERVAL = 1.0
_POLL_BACKOFF = 1.5
_POLL_MAX_INTERVAL = 5.0
# Floor the backoff so it cannot stay at zero. Multiplying is not enough:
# 0 * 1.5 is still 0, so a caller passing poll_interval=0 would busy-spin for
# the whole timeout -- 900s of back-to-back requests against a live server.
# The caller's interval is still honoured for the first sleep, so 0 keeps
# meaning "re-poll immediately once"; every sleep after that is bounded.
_POLL_MIN_INTERVAL = 0.05
# Retrieval
DEFAULT_MIN_SIMILARITY = 0.3
# How many consecutive transport errors to absorb while polling a job before
# giving up. Job records live in memory and are lost on restart, so a service
# bounce mid-index otherwise turns a healthy upload into a hard failure.
_POLL_MAX_TRANSIENT_ERRORS = 3
# Progress callback - set by UI to receive progress updates
_progress_callback: Callable[[str, str], None] | None = None
[docs]
def empty_rag_enabled() -> bool:
"""Whether the empty-RAG control condition is active.
In this mode the RAG tools are still created and offered to the agent, but
retrieval and listing return nothing, so the benchmark can measure whether
the mere availability of a retrieval tool changes agent behavior.
Read at call time, not cached: evaluation harnesses flip the variable after
this module is imported.
Note this deliberately makes the client *lie* about the index. Tooling that
needs the truth -- anything that decides what to ingest, or asserts that
documents are present -- must refuse to run while it is set rather than
trust the answer. See scripts/build_mmore_index.py.
"""
return os.getenv("MMORE_EMPTY_RAG", "false").lower() == "true"
[docs]
def default_min_similarity() -> float:
"""Score floor for agent-facing document search.
Overridable via MMORE_MIN_SIMILARITY. The floor is worth tuning if an
upstream bump shifts the similarity scale: too high and retrieval silently
returns nothing while the service looks perfectly healthy.
"""
raw = os.getenv("MMORE_MIN_SIMILARITY")
if raw is None or not raw.strip():
return DEFAULT_MIN_SIMILARITY
try:
return float(raw)
except ValueError:
logger.warning(
f"MMORE_MIN_SIMILARITY={raw!r} is not a number; "
f"using {DEFAULT_MIN_SIMILARITY}"
)
return DEFAULT_MIN_SIMILARITY
[docs]
def set_mmore_progress_callback(callback: Callable[[str, str], None] | None) -> None:
"""
Set the progress callback for MMORE operations.
Args:
callback: Callback function(step: str, message: str) or None to clear
"""
global _progress_callback # noqa: PLW0603
_progress_callback = callback
[docs]
def get_mmore_progress_callback() -> Callable[[str, str], None] | None:
"""Get the current MMORE progress callback."""
return _progress_callback
def _report_progress(step: str, message: str) -> None:
"""Report progress if callback is set."""
callback = get_mmore_progress_callback()
if callback:
callback(step, message)
[docs]
class MMOREFileNotFoundError(FileNotFoundError):
"""Exception raised when a file is not found for MMORE operations."""
[docs]
class MMOREDuplicateFileError(requests.HTTPError):
"""Raised when MMORE already holds a file with the requested ``fileId``.
MMORE v2 answers a duplicate upload with ``409``. Subclasses
:class:`requests.HTTPError` so existing broad ``except requests.HTTPError``
handlers keep working.
"""
[docs]
def __init__(self, file_id: str, *args: Any, **kwargs: Any) -> None:
self.file_id = file_id
super().__init__(
f"MMORE already holds a file with id '{file_id}'", *args, **kwargs
)
[docs]
class MMOREIndexingError(RuntimeError):
"""Raised when a background indexing job reaches a terminal failure state."""
[docs]
class MMOREIndexingTimeoutError(TimeoutError):
"""Raised when a background indexing job does not finish in time."""
[docs]
class MMOREClient:
"""Client for interacting with MMORE RAG service API."""
[docs]
def __init__(
self,
base_url: str | None = None,
*,
index_timeout: float | None = None,
poll_interval: float | None = None,
):
"""Initialize MMORE client.
Args:
base_url: Base URL of mmore service (defaults to MMORE_RAG_URL env var)
index_timeout: Seconds to wait for a background indexing job before
giving up (defaults to the MMORE_INDEX_TIMEOUT env var, then 900)
poll_interval: Initial delay between job polls, in seconds. Grows
geometrically up to 5s.
"""
base_url_str: str = base_url or (
os.getenv("MMORE_RAG_URL") or "http://localhost:8000"
)
self.base_url = base_url_str.rstrip("/")
self.index_timeout = (
index_timeout
if index_timeout is not None
else float(os.getenv("MMORE_INDEX_TIMEOUT", str(DEFAULT_INDEX_TIMEOUT)))
)
self.poll_interval = (
poll_interval if poll_interval is not None else _POLL_INITIAL_INTERVAL
)
logger.info(f"MMORE client initialized with URL: {self.base_url}")
[docs]
def health_check(self) -> bool:
"""Check if mmore service is healthy.
Probes ``GET /``, which is the health route on MMORE's live-retrieval
app. That app does not expose ``/health`` -- that belongs to the
separate ``mmore rag`` generation server, which this project does not run.
Returns:
True if service is reachable, False otherwise
"""
try:
response = requests.get(f"{self.base_url}/", timeout=5)
except requests.RequestException as e:
logger.warning(f"MMORE health check failed: {e}")
return False
else:
return response.status_code == HTTP_OK
[docs]
def upload_file(
self,
file_path: str,
file_id: str | None = None,
original_name: str | None = None,
*,
wait: bool = True,
) -> dict[str, Any]:
"""Upload a file to mmore and, by default, block until it is indexed.
MMORE v2 accepts the upload with ``202`` and indexes it on a background
queue, so the response arrives well before the document is searchable.
This method polls ``GET /v1/jobs/{jobId}`` to preserve the synchronous
"returns implies searchable" contract callers depend on.
Args:
file_path: Path to the file to upload
file_id: Optional unique identifier (defaults to filename)
original_name: Optional original file name for display (uses file_path name if not provided)
wait: If False, return as soon as the upload is accepted, without
waiting for indexing to finish.
Returns:
The terminal job snapshot, always carrying at least ``fileId`` and
``jobId``. With ``wait=False``, the immediate acceptance payload.
Raises:
MMOREFileNotFoundError: The local file does not exist
MMOREDuplicateFileError: MMORE already holds this fileId (409)
MMOREIndexingError: The indexing job failed
MMOREIndexingTimeoutError: The job did not finish within index_timeout
requests.HTTPError: Any other non-2xx response
"""
file_path_obj = Path(file_path)
if not file_path_obj.exists():
raise MMOREFileNotFoundError(file_path)
file_id = file_id or file_path_obj.stem
display_name = original_name or file_path_obj.name
# Report progress: preparing
file_size_mb = file_path_obj.stat().st_size / (1024 * 1024)
_report_progress(
"prepare", f"Preparing to upload '{display_name}' ({file_size_mb:.1f} MB)"
)
with file_path_obj.open("rb") as f:
files = {"file": (file_path_obj.name, f)}
data = {"fileId": file_id}
# Report progress: uploading
_report_progress(
"upload", f"Uploading '{display_name}' to MMORE service..."
)
# This timeout now covers only the body transfer: indexing happens
# afterwards, on the server's background queue.
response = requests.post(
f"{self.base_url}/v1/files",
files=files,
data=data,
timeout=300,
)
if response.status_code == HTTP_CONFLICT:
raise MMOREDuplicateFileError(file_id, response=response)
response.raise_for_status()
payload = response.json()
job_id = payload.get("jobId")
if job_id is None:
# A server that indexed inline (e.g. the frozen v1 stack) has
# nothing to poll; the upload is already complete.
logger.warning(
f"MMORE accepted {file_id} without a jobId; assuming synchronous indexing"
)
return dict(payload)
logger.info(f"Uploaded file {file_id} to mmore (job {job_id})")
if not wait:
return dict(payload)
snapshot = self._await_job(job_id, display_name)
snapshot.setdefault("fileId", file_id)
snapshot.setdefault("jobId", job_id)
return snapshot
[docs]
def get_job(self, job_id: str) -> dict[str, Any]:
"""Fetch a snapshot of a background indexing job.
Args:
job_id: The job identifier returned by an upload
Returns:
Job snapshot: ``{jobId, fileId, filename, status, device, result, error}``.
On success ``result`` is ``{"chunks": <int>}``.
Raises:
requests.HTTPError: If the job is unknown or expired (404). MMORE
keeps job records in memory for two hours only.
"""
response = requests.get(f"{self.base_url}/v1/jobs/{job_id}", timeout=30)
response.raise_for_status()
data: dict[str, Any] = response.json()
return data
def _await_job(self, job_id: str, display_name: str) -> dict[str, Any]:
"""Poll an indexing job until it reaches a terminal state.
Progress is reported on status *transitions* rather than on every poll,
so a long job does not flood the UI with one message per second.
A few consecutive transport errors are absorbed rather than aborting
the wait. Indexing a large PDF takes minutes, and treating one dropped
connection as fatal would fail an upload that is in fact progressing
fine. Errors that persist still surface, and the overall deadline still
applies.
"""
deadline = time.monotonic() + self.index_timeout
interval = self.poll_interval
last_status: str | None = None
transient_errors = 0
while True:
try:
snapshot = self.get_job(job_id)
except requests.RequestException as e:
transient_errors += 1
if (
transient_errors > _POLL_MAX_TRANSIENT_ERRORS
or time.monotonic() >= deadline
):
raise
logger.warning(
f"MMORE job {job_id} poll failed "
f"({transient_errors}/{_POLL_MAX_TRANSIENT_ERRORS}): {e}"
)
time.sleep(min(max(interval, _POLL_MIN_INTERVAL), _POLL_MAX_INTERVAL))
continue
transient_errors = 0
status = str(snapshot.get("status", "")).strip().lower()
if status != last_status:
last_status = status
_report_progress("process", f"Indexing '{display_name}': {status}")
if status in _JOB_SUCCEEDED:
chunks = (snapshot.get("result") or {}).get("chunks")
logger.info(
f"MMORE job {job_id} finished ({status})"
+ (f", {chunks} chunks" if chunks is not None else "")
)
return dict(snapshot)
if status in _JOB_FAILED:
detail = snapshot.get("error") or "no error detail reported"
raise MMOREIndexingError(
f"MMORE job {job_id} for '{display_name}' {status}: {detail}"
)
remaining = deadline - time.monotonic()
if remaining <= 0:
raise MMOREIndexingTimeoutError(
f"MMORE job {job_id} for '{display_name}' still '{status}' "
f"after {self.index_timeout:.0f}s"
)
time.sleep(min(interval, remaining))
interval = min(
max(interval * _POLL_BACKOFF, _POLL_MIN_INTERVAL),
_POLL_MAX_INTERVAL,
)
[docs]
def replace_file(
self,
file_path: str,
file_id: str,
original_name: str | None = None,
*,
wait: bool = True,
) -> dict[str, Any]:
"""Replace an already-indexed file, reusing its ``fileId``.
Uses ``PUT /v1/files/{fileId}``, which avoids the delete-then-upload
dance and the 409 that a plain re-upload would raise.
Args:
file_path: Path to the replacement file
file_id: ID of the existing document to replace
original_name: Optional display name for progress messages
wait: If False, return without waiting for re-indexing to finish
Returns:
The terminal job snapshot (or the acceptance payload if wait=False)
Raises:
MMOREFileNotFoundError: The local file does not exist
requests.HTTPError: If MMORE holds no file with this ID (404)
MMOREIndexingError: The indexing job failed
MMOREIndexingTimeoutError: The job did not finish within index_timeout
"""
file_path_obj = Path(file_path)
if not file_path_obj.exists():
raise MMOREFileNotFoundError(file_path)
display_name = original_name or file_path_obj.name
_report_progress("upload", f"Replacing '{display_name}' in MMORE service...")
with file_path_obj.open("rb") as f:
response = requests.put(
f"{self.base_url}/v1/files/{file_id}",
files={"file": (file_path_obj.name, f)},
timeout=300,
)
response.raise_for_status()
payload = response.json()
job_id = payload.get("jobId")
if job_id is None or not wait:
return dict(payload)
snapshot = self._await_job(job_id, display_name)
snapshot.setdefault("fileId", file_id)
snapshot.setdefault("jobId", job_id)
return snapshot
[docs]
def retrieve(
self,
query: str,
file_ids: list[str] | None = None,
max_matches: int = 5,
min_similarity: float = 0.0,
collection_name: str | None = None,
) -> list[Document]:
"""Retrieve relevant documents for a query.
Args:
query: Search query
file_ids: Optional list of file IDs to restrict search to
max_matches: Maximum number of results to return
min_similarity: Minimum similarity score threshold. Always sent
explicitly, so MMORE's own default never applies.
collection_name: Optional Milvus collection override
Returns:
List of LangChain Document objects with retrieved content
Raises:
requests.HTTPError: If retrieval fails
"""
# Empty-RAG evaluation mode: tools are available but index is empty.
if empty_rag_enabled():
logger.debug("MMORE_EMPTY_RAG=true: returning empty results")
return []
payload: dict[str, Any] = {
"query": query,
"fileIds": file_ids or [],
"maxMatches": max_matches,
"minSimilarity": min_similarity,
}
if collection_name is not None:
payload["collectionName"] = collection_name
response = requests.post(
f"{self.base_url}/v1/retrieve",
json=payload,
timeout=30,
)
response.raise_for_status()
data = response.json()
# MMORE returns a list directly, not a dict with 'docs' key
docs_list = data if isinstance(data, list) else data.get("docs", [])
logger.debug(f"Retrieved {len(docs_list)} documents from mmore")
# Convert to LangChain documents
documents = []
for doc_info in docs_list:
metadata: dict[str, Any] = {
"source": doc_info.get("fileId", "unknown"),
"chunk_id": doc_info.get("chunkId"),
"score": doc_info.get("similarity", doc_info.get("score", 0.0)),
}
if doc_info.get("filePath") is not None:
metadata["file_path"] = doc_info["filePath"]
metadata.update(_flatten_chunk_metadata(doc_info.get("metadata") or {}))
documents.append(
Document(page_content=doc_info.get("content", ""), metadata=metadata)
)
return documents
[docs]
def delete_file(self, file_id: str) -> dict[str, Any]:
"""Delete a file from mmore index.
Args:
file_id: ID of the file to delete
Returns:
Response from mmore API
Raises:
requests.HTTPError: If deletion fails
"""
response = requests.delete(
f"{self.base_url}/v1/files/{file_id}",
timeout=30,
)
response.raise_for_status()
logger.info(f"Deleted file {file_id} from mmore")
data: dict[str, Any] = response.json()
return data
[docs]
def list_files(
self, collection_name: str = "my_docs", limit: int | None = None
) -> list[str]:
"""List all files stored in the mmore index.
Args:
collection_name: Name of the Milvus collection (defaults to 'my_docs')
limit: Optional cap on how many entries to return (MMORE defaults
to 16000, and rejects anything above 100000)
Returns:
List of file IDs
Raises:
requests.HTTPError: If listing fails
"""
# Empty-RAG evaluation mode: tools are available but index is empty.
if empty_rag_enabled():
logger.debug("MMORE_EMPTY_RAG=true: returning empty file list")
return []
params: dict[str, Any] = {"collection_name": collection_name}
if limit is not None:
params["limit"] = limit
response = requests.get(
f"{self.base_url}/list_files",
params=params,
timeout=30,
)
response.raise_for_status()
data = response.json()
# Entries are normally {"id": ..., "filename": ...}; tolerate a bare
# string list and the alternative "fileId" spelling.
file_ids: list[str] = []
for item in data or []:
candidate = (
item
if isinstance(item, str)
else (item.get("id") or item.get("fileId"))
)
if candidate:
file_ids.append(candidate)
logger.info(f"Listed {len(file_ids)} files from mmore")
return file_ids
def _flatten_chunk_metadata(raw: dict[str, Any]) -> dict[str, Any]:
"""Keep MMORE v2's nested chunk metadata while restoring the flat keys.
v2 reports positions as ``{"first": {"page", "paragraph"}, "last": {...}}``.
Splatting that directly would bury dicts under ``first``/``last`` and drop
the flat ``page`` key that v1 consumers read, so re-derive them from
``first`` while preserving the raw structure alongside.
"""
flat = dict(raw)
first = raw.get("first")
if isinstance(first, dict):
if "page" in first:
flat.setdefault("page", first["page"])
if "paragraph" in first:
flat.setdefault("paragraph", first["paragraph"])
return flat
