mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-11-11 08:38:48 +00:00
ff2956a41f
* refactor: share backend search, file responses, and vector digests * refactor(studio): remove unused scaffolding and share workspace updates * test: simplify async tests and share deterministic store fixtures
68 lines
2.8 KiB
Python
68 lines
2.8 KiB
Python
"""Shared execution for single-backend search steps."""
|
|
|
|
from abc import abstractmethod
|
|
from typing import Final
|
|
|
|
from ._dedup import _ToolContextDedupMixin
|
|
from ._source_format import ALL_RETURNED_MESSAGE, NO_RESULTS_MESSAGE, join_chunk_entries
|
|
from ._source_format import merge_session_chunk_intervals, render_chunk_entries
|
|
from ..base_step import BaseStep
|
|
from ...schema import FileChunk
|
|
|
|
_MAX_CANDIDATES: Final = 200
|
|
_CANDIDATE_MULTIPLIER: Final = 10
|
|
|
|
|
|
class _SingleSearchStep(_ToolContextDedupMixin, BaseStep):
|
|
"""Retrieve candidates, then apply the common single-backend search contract."""
|
|
|
|
def __init__(self, *args, seen_ttl_hours: float = 24, include_source: bool = True, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.seen_ttl_hours = seen_ttl_hours
|
|
self.include_source = include_source
|
|
|
|
async def execute(self):
|
|
assert self.context is not None
|
|
raw_query = self.context.get("query")
|
|
query: str = "" if raw_query is None else str(raw_query).strip()
|
|
limit: int = int(self.context.get("limit") or 5)
|
|
min_score: float = float(self.context.get("min_score") or 0.0)
|
|
tool_context_id: str = (self.context.get("tool_context_id", "") or "").strip()
|
|
|
|
if not query:
|
|
self.context.response.success = False
|
|
self.context.response.answer = "Error: query cannot be empty"
|
|
return self.context.response
|
|
assert limit > 0, f"limit must be positive, got {limit}"
|
|
|
|
candidates = min(_MAX_CANDIDATES, max(1, limit * _CANDIDATE_MULTIPLIER))
|
|
results = await self._retrieve(query, candidates)
|
|
self.logger.info(f"[{self.name}] query={query!r} candidates={candidates} hits={len(results)}")
|
|
|
|
if min_score > 0.0:
|
|
results = [chunk for chunk in results if chunk.score >= min_score]
|
|
|
|
pre_dedup_count = 0
|
|
if tool_context_id:
|
|
pre_dedup_count = len(results)
|
|
results, _ = self._dedupe_tool_context(results, tool_context_id, limit)
|
|
else:
|
|
results = results[:limit]
|
|
|
|
session_dir = self.config_value("session_dir")
|
|
entries = render_chunk_entries(
|
|
merge_session_chunk_intervals(results, session_dir),
|
|
session_dir,
|
|
include_source=self.include_source,
|
|
)
|
|
self.context.response.answer = join_chunk_entries(entries)
|
|
if not results:
|
|
self.context.response.answer = ALL_RETURNED_MESSAGE if pre_dedup_count > 0 else NO_RESULTS_MESSAGE
|
|
self.context.response.metadata["results"] = [
|
|
c.model_dump(exclude_none=True, exclude={"embedding"}) for c in results
|
|
]
|
|
return self.context.response
|
|
|
|
@abstractmethod
|
|
async def _retrieve(self, query: str, candidates: int) -> list[FileChunk]:
|
|
"""Fetch ranked candidates from the concrete backend."""
|