mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-20 22:51:06 +08:00
454dea686e
### Summary Agentic search <img width="1149" height="1575" alt="image" src="https://github.com/user-attachments/assets/bce9a3e7-0517-4fb2-80a2-5d2a81a4da78" /> --------- Co-authored-by: Yingfeng Zhang <yingfeng.zhang@gmail.com>
88 lines
3.0 KiB
Python
88 lines
3.0 KiB
Python
"""Inspector tools: operate on already-returned results."""
|
|
|
|
import logging
|
|
from copy import deepcopy
|
|
|
|
_LOG = logging.getLogger(__name__)
|
|
|
|
|
|
async def open_context(tools, chunk_id: str, width: int = 500) -> dict:
|
|
"""Expand context around a chunk by looking up adjacent chunks in kbinfos."""
|
|
chunks = tools.kbinfos.get("chunks", [])
|
|
idx = _find_chunk_index(chunks, chunk_id)
|
|
if idx is None:
|
|
return {"chunks": [], "doc_aggs": []}
|
|
start = max(0, idx - 2)
|
|
end = min(len(chunks), idx + 2)
|
|
context = chunks[start:end]
|
|
_LOG.info("open_context: chunk=%s index=%d context=%d chunks", chunk_id, idx, len(context))
|
|
return {"chunks": context, "doc_aggs": _collect_doc_aggs(tools, context)}
|
|
|
|
|
|
async def compare_sources(tools, chunk_ids: list[str]) -> dict:
|
|
"""Find chunks in kbinfos and list the document sources they come from."""
|
|
if not chunk_ids:
|
|
return {"chunks": [], "doc_aggs": []}
|
|
chunks = tools.kbinfos.get("chunks", [])
|
|
matched = [c for c in chunks if _chunk_id(c) in chunk_ids]
|
|
return {"chunks": matched, "doc_aggs": _collect_doc_aggs(tools, matched)}
|
|
|
|
|
|
async def grep_within(tools, doc_id: str, pattern: str) -> dict:
|
|
"""Find a keyword within a document and return its chunks narrowed to the
|
|
matching sentences (+/- 1 neighbour).
|
|
|
|
``pattern`` is treated as the keyword string (comma-separate for several).
|
|
Delegates to :func:`_narrow_by_keywords`, which keeps only keyword-bearing
|
|
sentences and drops chunks with no match. Operates on copies so the shared
|
|
``kbinfos`` citation pool is never mutated.
|
|
"""
|
|
from rag.advanced_rag.harness.tools.search import _narrow_by_keywords
|
|
|
|
chunks = [deepcopy(c) for c in tools.kbinfos.get("chunks", []) if c.get("doc_id") == doc_id]
|
|
matched = _narrow_by_keywords(chunks, pattern)
|
|
return {"chunks": matched, "doc_aggs": _collect_doc_aggs(tools, matched)}
|
|
|
|
|
|
async def request_adjacent(tools, chunk_id: str, direction: str = "next", count: int = 3) -> dict:
|
|
"""Get adjacent entries before or after a chunk."""
|
|
chunks = tools.kbinfos.get("chunks", [])
|
|
idx = _find_chunk_index(chunks, chunk_id)
|
|
if idx is None:
|
|
return {"chunks": [], "doc_aggs": []}
|
|
|
|
if direction == "prev":
|
|
start = max(0, idx - count)
|
|
end = idx
|
|
else:
|
|
start = idx + 1
|
|
end = min(len(chunks), start + count)
|
|
|
|
adjacent = chunks[start:end]
|
|
return {"chunks": adjacent, "doc_aggs": _collect_doc_aggs(tools, adjacent)}
|
|
|
|
|
|
# Helpers
|
|
|
|
|
|
def _find_chunk_index(chunks: list[dict], chunk_id: str) -> int | None:
|
|
for i, c in enumerate(chunks):
|
|
if _chunk_id(c) == chunk_id:
|
|
return i
|
|
return None
|
|
|
|
|
|
def _chunk_id(ck: dict) -> str:
|
|
return str(ck.get("chunk_id") or ck.get("id") or "")
|
|
|
|
|
|
def _collect_doc_aggs(tools, chunks: list[dict]) -> list[dict]:
|
|
seen = set()
|
|
aggs = []
|
|
for c in chunks:
|
|
doc_id = c.get("doc_id")
|
|
if doc_id and doc_id not in seen:
|
|
seen.add(doc_id)
|
|
aggs.append({"doc_id": doc_id, "doc_name": c.get("docnm_kwd", "")})
|
|
return aggs
|