From 454dea686ec2a55b1b762ed12d21f97f2ea288dc Mon Sep 17 00:00:00 2001 From: Kevin Hu Date: Wed, 15 Jul 2026 23:46:23 +0800 Subject: [PATCH] Feat: agentic search framework (#16859) ### Summary Agentic search image --------- Co-authored-by: Yingfeng Zhang --- api/apps/restful_apis/chat_api.py | 8 +- api/db/services/dialog_service.py | 178 ++++- api/db/services/llm_service.py | 2 +- pyproject.toml | 1 + rag/advanced_rag/__init__.py | 13 +- rag/advanced_rag/agentic_rag.py | 614 ++++++++++++++++++ rag/advanced_rag/agentic_rag_graph.py | 346 ++++++++++ rag/advanced_rag/harness/__init__.py | 1 + rag/advanced_rag/harness/agent.py | 345 ++++++++++ rag/advanced_rag/harness/config.py | 103 +++ .../harness/orchestrator/__init__.py | 35 + .../harness/orchestrator/agentic.py | 241 +++++++ .../harness/orchestrator/decompose.py | 112 ++++ .../harness/orchestrator/direct.py | 49 ++ rag/advanced_rag/harness/pipeline.py | 125 ++++ rag/advanced_rag/harness/planner.py | 147 +++++ rag/advanced_rag/harness/prompts/__init__.py | 1 + .../harness/prompts/decompose_prompts.py | 121 ++++ .../harness/prompts/report_prompt.py | 16 + .../harness/prompts/research_agent_prompt.py | 58 ++ .../harness/prompts/route_prompt.py | 19 + .../harness/prompts/sufficiency_prompt.py | 31 + rag/advanced_rag/harness/route.py | 77 +++ rag/advanced_rag/harness/sufficiency.py | 188 ++++++ rag/advanced_rag/harness/tools/__init__.py | 46 ++ rag/advanced_rag/harness/tools/exploration.py | 32 + rag/advanced_rag/harness/tools/gating.py | 125 ++++ rag/advanced_rag/harness/tools/inspector.py | 87 +++ rag/advanced_rag/harness/tools/navigation.py | 36 + rag/advanced_rag/harness/tools/registry.py | 141 ++++ rag/advanced_rag/harness/tools/search.py | 289 +++++++++ rag/advanced_rag/harness/types.py | 162 +++++ rag/advanced_rag/think_log.py | 94 +++ rag/llm/chat_model.py | 35 +- rag/llm/tool_decorator.py | 152 ++++- rag/prompts/generator.py | 17 + rag/prompts/sufficiency_select.md | 28 + rag/utils/es_conn.py | 2 +- test/testcases/restful_api/test_chats.py | 1 + .../test_user_tenant_routes_unit.py | 1 + uv.lock | 171 ++++- 41 files changed, 4217 insertions(+), 33 deletions(-) create mode 100644 rag/advanced_rag/agentic_rag.py create mode 100644 rag/advanced_rag/agentic_rag_graph.py create mode 100644 rag/advanced_rag/harness/__init__.py create mode 100644 rag/advanced_rag/harness/agent.py create mode 100644 rag/advanced_rag/harness/config.py create mode 100644 rag/advanced_rag/harness/orchestrator/__init__.py create mode 100644 rag/advanced_rag/harness/orchestrator/agentic.py create mode 100644 rag/advanced_rag/harness/orchestrator/decompose.py create mode 100644 rag/advanced_rag/harness/orchestrator/direct.py create mode 100644 rag/advanced_rag/harness/pipeline.py create mode 100644 rag/advanced_rag/harness/planner.py create mode 100644 rag/advanced_rag/harness/prompts/__init__.py create mode 100644 rag/advanced_rag/harness/prompts/decompose_prompts.py create mode 100644 rag/advanced_rag/harness/prompts/report_prompt.py create mode 100644 rag/advanced_rag/harness/prompts/research_agent_prompt.py create mode 100644 rag/advanced_rag/harness/prompts/route_prompt.py create mode 100644 rag/advanced_rag/harness/prompts/sufficiency_prompt.py create mode 100644 rag/advanced_rag/harness/route.py create mode 100644 rag/advanced_rag/harness/sufficiency.py create mode 100644 rag/advanced_rag/harness/tools/__init__.py create mode 100644 rag/advanced_rag/harness/tools/exploration.py create mode 100644 rag/advanced_rag/harness/tools/gating.py create mode 100644 rag/advanced_rag/harness/tools/inspector.py create mode 100644 rag/advanced_rag/harness/tools/navigation.py create mode 100644 rag/advanced_rag/harness/tools/registry.py create mode 100644 rag/advanced_rag/harness/tools/search.py create mode 100644 rag/advanced_rag/harness/types.py create mode 100644 rag/advanced_rag/think_log.py create mode 100644 rag/prompts/sufficiency_select.md diff --git a/api/apps/restful_apis/chat_api.py b/api/apps/restful_apis/chat_api.py index 9398676a30..12a2662795 100644 --- a/api/apps/restful_apis/chat_api.py +++ b/api/apps/restful_apis/chat_api.py @@ -30,7 +30,7 @@ from api.apps.restful_apis._generation_params import merge_generation_config, po from api.db.joint_services.tenant_model_service import get_api_key, get_tenant_default_model_by_type, resolve_model_config from api.db.services.chunk_feedback_service import ChunkFeedbackService from api.db.services.conversation_service import ConversationService, structure_answer -from api.db.services.dialog_service import DialogService, async_chat, gen_mindmap +from api.db.services.dialog_service import DialogService, async_chat, gen_mindmap, rag_agent from api.db.services.knowledgebase_service import KnowledgebaseService, validate_dataset_embedding_models from api.db.services.llm_service import LLMBundle from api.db.services.search_service import SearchService @@ -1256,7 +1256,7 @@ async def session_completion(chat_id_in_arg=""): # start_to_think/end_to_think events. legacy_answer = "" final_answer = None - async for ans in async_chat(dia, msg, True, session_id=session_id, **req): + async for ans in rag_agent(dia, msg, True, session_id=session_id, **req): ans = _format_answer(ans) if ans.get("final"): final_answer = ans @@ -1293,7 +1293,7 @@ async def session_completion(chat_id_in_arg=""): payload = _sanitize_json_floats({"code": 0, "message": "", "data": final_chunk}) yield "data:" + json.dumps(payload, ensure_ascii=False) + "\n\n" else: - async for ans in async_chat(dia, msg, True, session_id=session_id, **req): + async for ans in rag_agent(dia, msg, True, session_id=session_id, **req): ans = _format_answer(ans) payload = _sanitize_json_floats({"code": 0, "message": "", "data": ans}) yield "data:" + json.dumps(payload, ensure_ascii=False) + "\n\n" @@ -1313,7 +1313,7 @@ async def session_completion(chat_id_in_arg=""): return resp answer = None - async for ans in async_chat(dia, msg, False, session_id=session_id, **req): + async for ans in rag_agent(dia, msg, False, session_id=session_id, **req): answer = _format_answer(ans) if conv is not None: await thread_pool_exec(ConversationService.update_by_id, conv.id, conv.to_dict()) diff --git a/api/db/services/dialog_service.py b/api/db/services/dialog_service.py index 8572211a1a..b99207e073 100644 --- a/api/db/services/dialog_service.py +++ b/api/db/services/dialog_service.py @@ -19,6 +19,7 @@ import re import time import uuid from copy import deepcopy +from rag.advanced_rag.agentic_rag import RAGTools logger = logging.getLogger(__name__) from datetime import datetime @@ -715,7 +716,8 @@ async def async_chat(dialog, messages, stream=True, **kwargs): logging.debug("Proceeding with retrieval") tenant_ids = list(set([kb.tenant_id for kb in kbs])) knowledges = [] - if prompt_config.get("reasoning", False) or kwargs.get("reasoning"): + # replaced by extension of reasoning: 0, 1, 2 + if False: # prompt_config.get("reasoning", False) or kwargs.get("reasoning"): reasoner = DeepResearcher( chat_mdl, prompt_config, @@ -1835,3 +1837,177 @@ async def gen_mindmap(question, kb_ids, tenant_id, search_config={}): mindmap = MindMapExtractor(chat_mdl) mind_map = await mindmap([c["content_with_weight"] for c in ranks["chunks"]]) return mind_map.output + + +async def rag_agent(dialog, messages, stream=True, **kwargs): + logging.debug("Begin rag_agent") + assert messages[-1]["role"] == "user", "The last content of this conversation is not from user." + prompt_config = dialog.prompt_config + if not prompt_config.get("reasoning", 0) and not kwargs.get("reasoning"): + async for ans in async_chat(dialog, messages, stream, **kwargs): + yield ans + return + kbs, embd_mdl, rerank_mdl, chat_mdl, tts_mdl = get_models(dialog) + use_web_search = _should_use_web_search(prompt_config, kwargs.get("internet")) + logging.debug("web_search kb=%s tavily=%s internet=%r enabled=%s", bool(dialog.kb_ids), bool(dialog.prompt_config.get("tavily_api_key")), kwargs.get("internet"), use_web_search) + tenant_ids = list(set([kb.tenant_id for kb in kbs])) + # "reasoning" arrives as "1".."4" mapping to the ordered THINKING_MODES + # (low, medium, high, ultra); fall back to "medium" on anything else. + from rag.advanced_rag.harness.config import THINKING_MODES + + _mode_labels = list(THINKING_MODES.keys()) + try: + _n = int(str(kwargs.get("reasoning")).strip()) + thinking_mode = _mode_labels[_n - 1] if 1 <= _n <= len(_mode_labels) else "medium" + except (TypeError, ValueError): + thinking_mode = "medium" + + rag_tools = RAGTools( + tenant_ids, + chat_mdl, + embed_mdl=embd_mdl, + kb_ids=dialog.kb_ids, + tav=Tavily(prompt_config["tavily_api_key"]) if use_web_search else None, + do_refer=False, + thinking_mode=thinking_mode, + ) + + async def decorate_answer(answer): + nonlocal rag_tools, messages + + refs = [] + ans = answer.split("") + think = "" + if len(ans) == 2: + think = ans[0] + "" + answer = ans[1] + + idx = set([]) + normalized_answer = normalize_arabic_digits(answer) or "" + for match in CITATION_MARKER_PATTERN.finditer(normalized_answer): + i = int(match.group(1)) + if i < len(rag_tools.kbinfos["chunks"]): + idx.add(i) + + answer, idx = repair_bad_citation_formats(answer, rag_tools.kbinfos, idx) + + doc_ids = set() + for citation in idx: + try: + chunk_index = int(citation) + except (TypeError, ValueError): + if citation: + doc_ids.add(str(citation)) + continue + if 0 <= chunk_index < len(rag_tools.kbinfos["chunks"]): + doc_id = rag_tools.kbinfos["chunks"][chunk_index].get("doc_id") + if doc_id: + doc_ids.add(doc_id) + + recall_docs = [d for d in rag_tools.kbinfos["doc_aggs"] if d["doc_id"] in doc_ids] + if not recall_docs: + recall_docs = rag_tools.kbinfos["doc_aggs"] + rag_tools.kbinfos["doc_aggs"] = recall_docs + + refs = deepcopy(rag_tools.kbinfos) if doc_ids else [] + for c in refs.get("chunks", []) if isinstance(refs, dict) else []: + if c.get("vector"): + del c["vector"] + + if answer.lower().find("invalid key") >= 0 or answer.lower().find("invalid api") >= 0: + answer += " Please set LLM API-Key in 'User Setting -> Model providers -> API-Key'" + + return {"answer": think + answer, "reference": refs, "prompt": "", "created_at": time.time()} + + # The agentic-search graph composes the final cited answer itself, so we + # stream its tokens straight to the client instead of relaying a tool + # result through a second outer-LLM pass. + + chat_mdl.bind_tools(None, rag_tools.tools) + # `rag` composes the full cited answer itself, so treat it as terminal: once + # the model calls it, stream its result and stop — otherwise the model would + # have to relay the (citation-bearing) answer through another round, which + # small models mangle or drop, so the client receives nothing. + if getattr(chat_mdl, "mdl", None) is not None: + chat_mdl.mdl.terminal_tools = {"rag"} + gen_conf = dialog.llm_setting + if stream: + # Surface the agentic pipeline's bracket-tagged progress logs to the + # client as content, interleaved with the real token stream. + from rag.advanced_rag.think_log import install_think_log_handler, set_think_log_sink, reset_think_log_sink + + install_think_log_handler() + event_queue: asyncio.Queue = asyncio.Queue() + loop = asyncio.get_running_loop() + + def _log_sink(msg): + try: + loop.call_soon_threadsafe(event_queue.put_nowait, ("log", msg)) + except RuntimeError: + pass + + async def _drive_stream(): + try: + stream_iter = chat_mdl.async_chat_streamly_delta(rag_tools.sys_prompt(), messages, gen_conf) + async for kind, value, state in _stream_with_think_delta(stream_iter): + event_queue.put_nowait(("stream", kind, value, state)) + except Exception: + logging.exception("rag_agent: agentic stream failed") + finally: + event_queue.put_nowait(("stream_done",)) + + token = set_think_log_sink(_log_sink) + drive = asyncio.create_task(_drive_stream()) + last_state = None + log_think_open = False + try: + while True: + item = await event_queue.get() + if item[0] == "log": + if not log_think_open: + yield {"answer": "", "reference": {}, "audio_binary": None, "final": False, "start_to_think": True} + log_think_open = True + yield {"answer": item[1] + "\n", "reference": {}, "audio_binary": None, "final": False} + continue + if item[0] == "stream_done": + break + _, kind, value, state = item + if state is not None: + last_state = state + # A real stream event follows the logs -> close the log think block. + if log_think_open: + yield {"answer": "", "reference": {}, "audio_binary": None, "final": False, "end_to_think": True} + log_think_open = False + if kind == "marker": + flags = {"start_to_think": True} if value == "" else {"end_to_think": True} + yield {"answer": "", "reference": {}, "audio_binary": None, "final": False, **flags} + continue + yield {"answer": value, "reference": {}, "audio_binary": tts(tts_mdl, value), "final": False} + if log_think_open: + yield {"answer": "", "reference": {}, "audio_binary": None, "final": False, "end_to_think": True} + log_think_open = False + finally: + reset_think_log_sink(token) + if not drive.done(): + drive.cancel() + try: + await drive + except asyncio.CancelledError: + pass + except Exception: + logging.exception("rag_agent: drive task error") + + full_answer = last_state.full_text if last_state else "" + if full_answer: + final = await decorate_answer(_extract_visible_answer(full_answer)) + final["final"] = True + final["audio_binary"] = None + yield final + else: + answer = await chat_mdl.async_chat(rag_tools.sys_prompt(), messages, gen_conf) + user_content = messages[-1].get("content", "[content not available]") + logging.debug("User: {}|Assistant: {}".format(user_content, answer)) + res = await decorate_answer(answer) + res["audio_binary"] = tts(tts_mdl, answer) + yield res + return diff --git a/api/db/services/llm_service.py b/api/db/services/llm_service.py index c1fe749bee..2a41918be9 100644 --- a/api/db/services/llm_service.py +++ b/api/db/services/llm_service.py @@ -412,7 +412,7 @@ class LLMBundle(LLM4Tenant): return queue async def async_chat(self, system: str, history: list, gen_conf: dict = {}, **kwargs): - if self.is_tools and getattr(self.mdl, "is_tools", False) and hasattr(self.mdl, "async_chat_with_tools"): + if self.is_tools and hasattr(self.mdl, "async_chat_with_tools"): base_fn = self.mdl.async_chat_with_tools elif hasattr(self.mdl, "async_chat"): base_fn = self.mdl.async_chat diff --git a/pyproject.toml b/pyproject.toml index bbced889bf..32bb849fbb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -132,6 +132,7 @@ dependencies = [ "yfinance==0.2.65", "zhipuai==2.0.1", "peewee>=3.17.1,<4.0.0", + "langgraph==1.2.0", # following modules aren't necessary # "nltk==3.9.1", # "numpy>=1.26.0,<2.0.0", diff --git a/rag/advanced_rag/__init__.py b/rag/advanced_rag/__init__.py index c294e8898d..d2b3ddc294 100644 --- a/rag/advanced_rag/__init__.py +++ b/rag/advanced_rag/__init__.py @@ -15,6 +15,17 @@ # from .tree_structured_query_decomposition_retrieval import TreeStructuredQueryDecompositionRetrieval as DeepResearcher +from .harness.config import THINKING_MODES, get_mode +from .harness.types import RouteDecision, ExecutionStrategy, ClaimTarget, WorkflowPlan, SufficiencyVerdict -__all__ = ["DeepResearcher"] +__all__ = [ + "DeepResearcher", + "THINKING_MODES", + "get_mode", + "RouteDecision", + "ExecutionStrategy", + "ClaimTarget", + "WorkflowPlan", + "SufficiencyVerdict", +] diff --git a/rag/advanced_rag/agentic_rag.py b/rag/advanced_rag/agentic_rag.py new file mode 100644 index 0000000000..01eada141e --- /dev/null +++ b/rag/advanced_rag/agentic_rag.py @@ -0,0 +1,614 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Agentic-RAG capability layer. + +``RAGTools`` bundles every retrieval primitive the agentic-search graph +(:mod:`rag.advanced_rag.agentic_rag_graph`) needs — question formalisation, +document scoping, keyword analysis, KB / web / structured retrieval, a +sufficiency judge and follow-up-question generation — plus the two things +the *outer* LLM is ever allowed to call as tools: ``rag`` (run the whole +agentic-search graph) and ``summarize_document`` (dump one document for an +explicit summary request). + +The individual search steps are deliberately NOT ``@tool``-decorated: the +graph orchestrates them itself, so ``chat_mdl`` stays a plain reasoning +model (no tool schema is bound onto it) and its ``async_chat*`` calls take +the fast non-tool-calling path. +""" + +from copy import deepcopy +import logging +import re +from typing import Any, List + +import json_repair +from api.db.services.doc_metadata_service import DocMetadataService +from api.db.services.document_service import DocumentService +from api.db.services.knowledgebase_service import KnowledgebaseService +from api.db.services.llm_service import LLMBundle +from common import settings +from common.misc_utils import thread_pool_exec +from common.token_utils import num_tokens_from_string +from rag.advanced_rag.agentic_rag_graph import _strip_think_stream +from rag.app.tag import label_question +from rag.llm.tool_decorator import tool +from rag.prompts.generator import ( + citation_prompt, + form_message, + gen_meta_filter, + kb_prompt, + message_fit_in, + multi_queries_gen, + sufficiency_select, +) +from api.db.db_models import Document, Knowledgebase +from rag.utils.tavily_conn import Tavily + + +# Tokens held back from the model's context when fitting retrieved evidence +# into the sufficiency / follow-up prompts. The evidence sits in the MIDDLE of +# those templates (question first, JSON output rules last), so if the combined +# prompt overflows the downstream trimmer eats the output rules, not the +# evidence. Reserving headroom for the template skeleton + question + output +# lets us trim the evidence up front instead. +_EVIDENCE_PROMPT_RESERVE_TOKENS = 1024 + + +class RAGTools: + def __init__( + self, + tenant_ids: list[str], + chat_mdl: LLMBundle, + embed_mdl: LLMBundle | None = None, + kb_ids: List[str] | None = None, + kbs: list[Knowledgebase] | None = None, + tav: Tavily | None = None, + meta_data_filter: dict | None = None, + user_defined_prompts: dict | None = None, + do_refer: bool | None = True, + thinking_mode: str = "medium", + ): + self.tenant_ids = tenant_ids + self.chat_mdl = deepcopy(chat_mdl) + self.embed_mdl = embed_mdl + self.thinking_mode = thinking_mode + self.field_map = {} + self.sql_kbs = [] + self.kbs = [] + self.kb_ids = [] + + def _exclude_sql_kb(kb): + if kb.parser_config and "field_map" in kb.parser_config: + self.field_map.update(kb.parser_config["field_map"]) + self.sql_kbs.append(kb) + else: + self.kbs.append(kb) + self.kb_ids.append(kb.id) + + if kb_ids: + for kb in KnowledgebaseService.get_by_ids(kb_ids): + _exclude_sql_kb(kb) + elif kbs: + for kb in kbs: + _exclude_sql_kb(kb) + + self.tav = tav + self.meta_data_filter = meta_data_filter + self.user_defined_prompts = user_defined_prompts or {} + self.kbinfos = {"chunks": [], "doc_aggs": []} + self.do_refer = do_refer + # Citation pool shared with the final-answer node: the graph publishes + # the chunks it actually used here (in the SAME order the answer's + # ``[ID:n]`` markers index), so the caller can resolve references. + self.kbinfos: dict[str, list] = {"chunks": [], "doc_aggs": []} + + # The two tools the outer LLM may bind. They are NOT auto-bound here — + # the agentic-search flow drives the graph directly — but callers that + # want a tool surface can do ``chat_mdl.bind_tools(tools=rag_tools.tools)``. + self.tools = [self.rag, self.summarize_document] + + # ------------------------------------------------------------------ # + # Capability flags / cheap introspection + # ------------------------------------------------------------------ # + def has_unstructured(self) -> bool: + return bool(self.kb_ids) + + def has_structured(self) -> bool: + return bool(self.sql_kbs and self.field_map) + + def has_web(self) -> bool: + return self.tav is not None + + def has_llm(self) -> bool: + return self.chat_mdl is not None + + async def _fit_messages(self, system: str, user: str) -> list: + """Fit system+user messages into the model's context window.""" + from rag.prompts.generator import form_message, message_fit_in + + _, msg = message_fit_in(form_message(system, user), self.chat_mdl.max_length) + return msg + + def get_citation_guidelines(self) -> str: + """Return the citation guidelines the final answer must follow.""" + return citation_prompt(self.user_defined_prompts) + + def sys_prompt(self) -> str: + """Thin router prompt for callers that bind ``self.tools``. + + The heavy workflow now lives inside the ``rag`` graph, so the outer + model only has to decide between answering-with-retrieval (``rag``) + and an explicit single-document summary (``summarize_document``). + """ + summarize_line = ( + "- Call `summarize_document` ONLY when the user explicitly asks to summarise a specific document ('summarise the security audit', 'tldr the onboarding guide'). It needs a document ID.\n" + if self.has_unstructured() + else "" + ) + return ( + "You are a smart agent. For any question that needs " + "evidence from the knowledge bases or the web, call the `rag` tool " + "with a self-contained question — it runs the full search-and-answer " + "pipeline and returns a cited answer.\n" + "After the `rag` tool returns, do not call `rag` again for the same " + "user question. Use the returned cited answer as the final answer " + "unless the user explicitly asks a new question.\n" + f"{summarize_line}" + "Do not invent facts and do not fabricate document IDs." + ) + + # ------------------------------------------------------------------ # + # Graph node helpers (plain async methods — never exposed as tools) + # ------------------------------------------------------------------ # + async def formalize(self, messages: List[Any]) -> tuple[str, str]: + """Rewrite the latest user message into a standalone question AND derive + its search keywords (each with close synonyms), in one LLM call. + + ``messages`` may be a list of role dicts (``{"role", "content"}``) or + pre-formatted ``"Speaker: text"`` strings. + + Returns ``(question, keywords)`` where ``keywords`` is a comma-separated + string of the question's key terms plus 1-2 close synonyms / alternative + phrasings for each, in the same language as the question. + """ + if not messages: + return "", "" + + lines: list[str] = [] + last_user = "" + for m in messages: + if isinstance(m, str): + lines.append(m) + last_user = m + continue + role = m.get("role", "user") + content = m.get("content", "") or "" + if role == "user": + last_user = content + prefix = "User" if role == "user" else ("Assistant" if role == "assistant" else str(role).capitalize()) + lines.append(f"{prefix}: {content}") + transcript = "\n".join(lines) + + system = ( + "You are given a conversation. Do BOTH of the following and return JSON only:\n" + "1. Rewrite the LAST user message into a single, self-contained question that can be " + "understood without seeing the prior conversation — resolve pronouns, ellipses and " + "follow-up shortcuts using earlier turns. Preserve the original language of the last " + "user message. If it is already a complete standalone question, keep it unchanged.\n" + "2. Extract keywords ONLY from the wording of the STANDALONE QUESTION itself — the " + "salient content words and phrases that literally appear in it (key nouns, named " + "entities, domain terms). Do NOT answer the question, and do NOT include any term that " + "would be part of the answer or is not present in the question. Then, for each extracted " + "term, you MAY add 1-2 close synonyms or alternative phrasings OF THAT SAME TERM. Output " + "them all together as one comma-separated list, in the SAME language as the question.\n" + ' Example — question "In which year did Apple acquire Beats?": keywords = ' + '"Apple, Apple Inc., acquire, acquisition, Beats" (terms from the question + synonyms; ' + 'the year is the ANSWER, so it must NOT appear).\n\n' + 'Output ONLY JSON, no prose, no code fences: ' + '{"question": "", "keywords": ""}' + ) + user = f"Conversation:\n{transcript}\n\nOutput JSON:" + _, msg = message_fit_in(form_message(system, user), self.chat_mdl.max_length) + ans = await self.chat_mdl.async_chat(msg[0]["content"], msg[1:], {"temperature": 0.1}) + if isinstance(ans, tuple): + ans = ans[0] + cleaned = re.sub(r"^.*", "", ans, flags=re.DOTALL) + cleaned = re.sub(r"```(?:json)?\s*|\s*```", "", cleaned).strip() + try: + data = json_repair.loads(cleaned) + except Exception as e: + logging.warning(f"formalize could not parse LLM output: {e!r} raw={ans[:200]!r}") + data = {} + if not isinstance(data, dict): + data = {} + + question = str(data.get("question") or "").strip().strip('"').strip("'") + if not question: + # Fall back to the raw last user message rather than an empty question. + question = (last_user or "").strip() + + keywords = data.get("keywords") or "" + if isinstance(keywords, list): + keywords = ", ".join(str(k).strip() for k in keywords if str(k).strip()) + keywords = str(keywords).strip() + return question, keywords + + async def pick_documents(self, question: str) -> List[str] | None: + """Narrow the search to a document subset for ``question``. + + Uses document metadata when the bound KBs carry any (mirrors the old + ``filter_docs_by_metadata``); otherwise asks an LLM to pick relevant + titles (mirrors the old ``select_documents``). Returns ``None`` when + no useful scope can be derived, meaning "search everything". + """ + return None + if not self.kb_ids: + return None + + metas = await self._get_cached_metas() + if metas: + ids = await self._filter_by_metadata(question, metas) + return ids or None + + ids = await self._select_by_titles(question) + return ids or None + + async def _filter_by_metadata(self, question: str, metas: dict) -> List[str]: + filters = await gen_meta_filter(self.chat_mdl, metas, question) + logging.debug(f"Metadata filter(auto) generated: {filters}") + conditions = filters.get("conditions") or [] + if not conditions: + return [] + logic = filters.get("logic", "and") + try: + doc_ids = await thread_pool_exec( + DocMetadataService.filter_doc_ids_by_meta_pushdown, + self.kb_ids, + conditions, + logic, + ) + except Exception as e: + logging.error(f"Metadata filter push down errored: {e}") + return [] + return doc_ids or [] + + async def _select_by_titles(self, question: str, max_docs: int = 512) -> List[str]: + docs = await thread_pool_exec(self._collect_doc_titles, max_docs) + if not docs: + return [] + + catalogue = "\n".join(f"docID: {doc_id}, title: {title}" for doc_id, title in docs) + system = ( + "You filter a document catalogue to find which documents are relevant " + "to a user's question. Use ONLY the titles in the catalogue — do not " + "invent docIDs. " + "Output ONLY a JSON array of the docIDs you consider relevant, e.g. " + '["abc123", "def456"]. If no document is clearly relevant, output []. ' + "No explanations, no Markdown, no code fences, no prose around the array." + ) + user = f"Question:\n{question}\n\nDocuments:\n{catalogue}\n\nRelevant docIDs (JSON array):" + _, msg = message_fit_in(form_message(system, user), self.chat_mdl.max_length) + ans = await self.chat_mdl.async_chat(msg[0]["content"], msg[1:], {"temperature": 0.1}) + if isinstance(ans, tuple): + ans = ans[0] + cleaned = re.sub(r"^.*", "", ans, flags=re.DOTALL) + cleaned = re.sub(r"```(?:json)?\s*|\s*```", "", cleaned).strip() + try: + ids = json_repair.loads(cleaned) + except Exception as e: + logging.warning(f"select_by_titles could not parse LLM output: {e!r} raw={ans[:200]!r}") + return [] + if not isinstance(ids, list): + return [] + known = {doc_id for doc_id, _ in docs} + return [doc_id for doc_id in ids if isinstance(doc_id, str) and doc_id in known] + + async def extract_keywords(self, question: str) -> str: + """Produce a compact keyword string (terms + a few close synonyms). + + Replaces the keywords the outer LLM used to hand to the retrieval + tool. Falls back to the question itself when extraction fails. + """ + if not question: + return "" + system = ( + "Extract the search terms for a knowledge-base query from the " + "question below. Output 3-8 of the most important content terms, " + "plus 1-2 close synonyms or alternative phrasings for any ambiguous " + "term. Single words or short noun phrases, space-separated, in the " + "SAME language as the question. Output ONLY the terms — no labels, " + "no punctuation lists, no explanation." + ) + try: + _, msg = message_fit_in(form_message(system, question), self.chat_mdl.max_length) + ans = await self.chat_mdl.async_chat(msg[0]["content"], msg[1:], {"temperature": 0.2}) + if isinstance(ans, tuple): + ans = ans[0] + ans = re.sub(r"^.*", "", ans, flags=re.DOTALL).strip() + except Exception: + logging.exception("extract_keywords failed") + ans = "" + return ans or question + + async def retrieve( + self, + question: str, + keywords: str | list = "", + doc_scope: List[str] | None = None, + top_n: int = 6, + similarity_threshold: float = 0.2, + using_embedding: bool = False, + ) -> dict[str, list]: + """Retrieve chunks from the unstructured KBs for one question. + + Returns a raw ``{"chunks": [...], "doc_aggs": [...]}`` dict — no + citation stamping, no accumulation onto ``self.kbinfos`` (the graph + owns merging so parallel per-question retrieval stays race-free). + """ + if not self.kb_ids: + return {"chunks": [], "doc_aggs": []} + if isinstance(keywords, list): + keywords = ",".join(keywords) + logging.info(f"@retrieve: {question}@{keywords}") + + if doc_scope: + candidates = [d for d in doc_scope if isinstance(d, str)] + known = await thread_pool_exec(self._filter_known_doc_ids, candidates) + valid = [d for d in candidates if d in known] + if valid: + doc_scope = valid + else: + if candidates: + logging.warning("retrieve: every supplied doc ID was unknown; falling back to unfiltered retrieval") + doc_scope = None + + search_terms = keywords.strip() if keywords else "" + if not search_terms: + search_terms = question + else: + question = question + " " + search_terms + + embd_mdl = self.embed_mdl if using_embedding else None + vector_weight = 0.7 if embd_mdl else 0 + kbinfos = await settings.retriever.retrieval( + question, + embd_mdl, + self.tenant_ids, + self.kb_ids, + 1, + top_n, + similarity_threshold, + vector_similarity_weight=vector_weight, + aggs=True, + highlight=True, + doc_ids=doc_scope, + rank_feature=label_question(question, self.kbs), + ) + if not kbinfos: + return {"chunks": [], "doc_aggs": []} + kbinfos["chunks"] = settings.retriever.retrieval_by_children(kbinfos.get("chunks", []), self.tenant_ids) + return {"chunks": kbinfos.get("chunks", []), "doc_aggs": kbinfos.get("doc_aggs", [])} + + async def web_retrieve(self, query: str) -> dict[str, list]: + """Retrieve chunks from the public web (Tavily). Raw kbinfos shape.""" + if self.tav is None: + return {"chunks": [], "doc_aggs": []} + try: + tav_res = await thread_pool_exec(self.tav.retrieve_chunks, query) + except Exception: + logging.exception("web_retrieve failed") + return {"chunks": [], "doc_aggs": []} + return {"chunks": tav_res.get("chunks", []), "doc_aggs": tav_res.get("doc_aggs", [])} + + async def structured_retrieve(self, question: str) -> dict[str, Any]: + """Query the structured (tabular) KBs by translating to SQL. + + Returns ``{"answer": str, "chunks": [...], "doc_aggs": [...]}``. The + answer is the natural-language SQL result the final node can weave in; + the chunks/doc_aggs feed the shared citation pool. + """ + if not self.has_structured(): + return {"answer": "", "chunks": [], "doc_aggs": []} + + # Lazy import — dialog_service constructs RAGTools. + from api.db.services.dialog_service import use_sql + + sql_kb_ids = [kb.id for kb in self.sql_kbs] + tenant_id = self.sql_kbs[0].tenant_id + try: + ans = await use_sql(question, self.field_map, tenant_id, self.chat_mdl, quota=True, kb_ids=sql_kb_ids) + except Exception as e: + logging.exception(f"structured_retrieve: use_sql failed: {e}") + return {"answer": "", "chunks": [], "doc_aggs": []} + if not ans: + return {"answer": "", "chunks": [], "doc_aggs": []} + reference = ans.get("reference") or {} + return { + "answer": ans.get("answer", "") or "", + "chunks": reference.get("chunks") or [], + "doc_aggs": reference.get("doc_aggs") or [], + } + + def _fit_evidence(self, question: str, evidence_md: str) -> str: + """Trim ``evidence_md`` so ``question`` + evidence + the prompt template + stay inside the model's context window. + + ``message_fit_in`` keeps the small side (the question) whole and trims + the large side (the evidence); we shrink the budget by a reserve so the + template skeleton and JSON output rules still fit afterwards. + """ + if not evidence_md: + return evidence_md + budget = max(256, self.chat_mdl.max_length - _EVIDENCE_PROMPT_RESERVE_TOKENS) + _, msg = message_fit_in(form_message(question, evidence_md), budget) + return msg[-1]["content"] + + async def judge_sufficiency(self, question: str, evidence_md: str) -> dict: + """Judge whether ``evidence_md`` answers ``question`` and pick useful chunks. + + ``evidence_md`` must carry ``ID: n`` markers per chunk (as produced by + ``kb_prompt``). Returns the verdict dict: + ``{"is_sufficient": bool, "reasoning": str, "missing_information": [...], + "useful_chunk_ids": [int, ...]}``. + """ + evidence_md = self._fit_evidence(question, evidence_md) + try: + return await sufficiency_select(self.chat_mdl, question, evidence_md) or {} + except Exception: + logging.exception("judge_sufficiency failed") + return {} + + async def gen_followups(self, question: str, query: str, missing: List[str], evidence_md: str) -> List[dict]: + """Generate complementary follow-up (question, query) pairs for gaps.""" + evidence_md = self._fit_evidence(question, evidence_md) + try: + res = await multi_queries_gen(self.chat_mdl, question, query or question, missing or [], evidence_md) or {} + except Exception: + logging.exception("gen_followups failed") + return [] + qs = res.get("questions") or [] + return [q for q in qs if isinstance(q, dict) and (q.get("question") or "").strip()] + + async def fetch_full_document(self, doc_id: str) -> dict[str, list]: + """Fetch a whole document's chunks in reading order (raw kbinfos).""" + if not self.kb_ids: + return {"chunks": [], "doc_aggs": []} + resolved = await thread_pool_exec(self._resolve_doc_tenant, doc_id) + if resolved is None: + logging.warning(f"fetch_full_document: doc_id {doc_id!r} not in any bound KB — refusing to fetch") + return {"chunks": [], "doc_aggs": []} + kb_id, tenant_id = resolved + + cks = [] + tokens = 0 + for offset in range(0, 10000, 128): + chunks = await thread_pool_exec( + settings.retriever.chunk_list, + doc_id, + tenant_id, + [kb_id], + max_count=offset + 128, + offset=offset, + fields=["content_with_weight", "docnm_kwd", "doc_id"], + sort_by_position=True, + retrieve_all=False, + ) + if not chunks: + break + for ck in chunks: + num = num_tokens_from_string(str(ck["content_with_weight"])) + if tokens + num > self.chat_mdl.max_length: + break + tokens += num + cks.append(ck) + if not cks: + return {"chunks": [], "doc_aggs": []} + doc_name = next((c.get("docnm_kwd") or "" for c in cks if c.get("docnm_kwd")), "") + return { + "chunks": cks, + "doc_aggs": [{"doc_name": doc_name, "doc_id": doc_id, "count": len(cks)}], + } + + # ------------------------------------------------------------------ # + # Bound tools + # ------------------------------------------------------------------ # + @tool(timeout=600) + async def rag(self, question: str) -> str: + """Answer a question with evidence from the knowledge bases and the web. + + Runs the full agentic-search pipeline: it formalises the question, + narrows the document scope, analyses keywords, retrieves evidence, + checks whether the evidence is sufficient (looping with follow-up + searches when it is not), and finally composes a cited answer. + + :param question: a self-contained natural-language question. + + :returns: the composed answer with inline citation markers. + """ + from rag.advanced_rag.agentic_rag_graph import run_agentic_rag + + messages = [{"role": "user", "content": question}] if question else [] + final = "" + async for delta in _strip_think_stream(run_agentic_rag(self, messages)): + if isinstance(delta, str): + final += delta + for p, r in [(r"\(\**(ID:\d)\**\)", "[\1]")]: + final = re.sub(p, r, final) + return final + + @tool + async def summarize_document(self, doc_id: str) -> list[str]: + """Return a single document's content, position-ordered, ready to summarise. + + Call ONLY for an explicit summary request about a specific document. + For general Q&A use the `rag` tool instead. + + :param doc_id: a 32-character lowercase hex document ID that some + other tool returned in this turn. Never invent one. + + :returns: formatted chunk blocks (document order) fitting the model's + context budget, prefixed with the citation rules to apply. + """ + kbinfos = await self.fetch_full_document(doc_id) + if not kbinfos["chunks"]: + return [] + start_idx = len(self.kbinfos.get("chunks", [])) + self.kbinfos["chunks"].extend(kbinfos["chunks"]) + self.kbinfos["doc_aggs"].extend(kbinfos["doc_aggs"]) + blocks = kb_prompt(self.kbinfos, self.chat_mdl.max_length) + if not self.do_refer: + return blocks[start_idx:] if start_idx else blocks + header = "# Citation rules\nApply the following rules VERBATIM to your final answer.\n\n" + citation_prompt(self.user_defined_prompts).strip() + "\n\n----\n\n" + return [header] + (blocks[start_idx:] if start_idx else blocks) + + # ------------------------------------------------------------------ # + # Low-level DB helpers (sync — wrap in thread_pool_exec at call sites) + # ------------------------------------------------------------------ # + async def _get_cached_metas(self) -> dict: + cached = getattr(self, "_metas_cache", None) + if cached is not None: + return cached + if not self.kb_ids: + self._metas_cache = {} + return self._metas_cache + self._metas_cache = await thread_pool_exec(DocMetadataService.get_flatted_meta_by_kbs, self.kb_ids) + return self._metas_cache or {} + + def _collect_doc_titles(self, max_docs: int = 512) -> list[tuple[str, str]] | None: + result: list[tuple[str, str]] = [] + for kb_id in self.kb_ids: + for doc in DocumentService.query(kb_id=kb_id): + result.append((doc.id, doc.name)) + if len(result) >= max_docs: + return None + return result + + def _filter_known_doc_ids(self, candidate_ids: list[str]) -> set[str]: + if not candidate_ids or not self.kb_ids: + return set() + rows = Document.select(Document.id).where((Document.id.in_(list(candidate_ids))) & (Document.kb_id.in_(self.kb_ids))) + return {row.id for row in rows} + + def _resolve_doc_tenant(self, doc_id: str) -> tuple[str, str] | None: + rows = list(Document.select(Document.kb_id).where((Document.id == doc_id) & (Document.kb_id.in_(self.kb_ids)))) + if not rows: + return None + kb_id = rows[0].kb_id + for kb in self.kbs: + if kb.id == kb_id: + return kb_id, kb.tenant_id + return None diff --git a/rag/advanced_rag/agentic_rag_graph.py b/rag/advanced_rag/agentic_rag_graph.py new file mode 100644 index 0000000000..fb87300cc1 --- /dev/null +++ b/rag/advanced_rag/agentic_rag_graph.py @@ -0,0 +1,346 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""LangGraph agentic-search graph — 4 nodes. + +Architecture: + + formalize_question → route → planner → orchestrator_loop → formalize_answer + +The ``orchestrator_loop`` node internally dispatches to one of three execution +strategies based on the thinking mode: + + low: direct_search — single hybrid search, no decomposition + medium: decompose_and_search — decompose → parallel search → sufficiency + high: agentic_research — two-level loop (orchestrator + research agent) + ultra: deep_research — same as high + dynamic claim expansion + replan +""" + +from __future__ import annotations + +import asyncio +import json +import logging +from typing import Any, TypedDict + +from langgraph.graph import END, START, StateGraph + +from rag.prompts.generator import form_message, kb_prompt, message_fit_in + +_LOG = logging.getLogger(__name__) + + +def _snip(value: Any, limit: int = 240) -> str: + try: + s = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False, default=str) + except Exception: + s = str(value) + s = " ".join(s.split()) + if len(s) > limit: + s = s[:limit] + f"...(+{len(s) - limit} chars)" + return s + + +class AgenticState(TypedDict, total=False): + messages: list + question: str + keywords: str # search keywords + close synonyms for the formalized question + seed_chunks: list # preliminary hybrid_search chunks used to ground the plan + route: dict # RouteDecision serialized + plan: dict # WorkflowPlan serialized + claims: list # ClaimTarget[] serialized + kbinfos: dict # accumulated chunks & doc_aggs + verdict: dict # SufficiencyVerdict serialized + partial_answer: bool + abstain: bool + empty_result: bool + final_answer: str + loop: int + feedback: str # replanning feedback + + +# ── Think tag helpers ── + +_THINK_OPEN = "" +_THINK_CLOSE = "" + + +def _partial_tag_tail(s: str, tag: str) -> int: + for k in range(min(len(s), len(tag) - 1), 0, -1): + if s.endswith(tag[:k]): + return k + return 0 + + +async def _strip_think_stream(stream): + """Strip ... spans from a token stream.""" + buf = "" + in_think = False + async for token in stream: + if not isinstance(token, str): + yield token + continue + buf += token + out = [] + while buf: + if not in_think: + idx = buf.find(_THINK_OPEN) + if idx == -1: + hold = _partial_tag_tail(buf, _THINK_OPEN) + if hold: + out.append(buf[: len(buf) - hold]) + buf = buf[len(buf) - hold :] + else: + out.append(buf) + buf = "" + break + out.append(buf[:idx]) + buf = buf[idx + len(_THINK_OPEN) :] + in_think = True + else: + idx = buf.find(_THINK_CLOSE) + if idx != -1: + buf = buf[idx + len(_THINK_CLOSE) :] + in_think = False + continue + hold = _partial_tag_tail(buf, _THINK_CLOSE) + buf = buf[len(buf) - hold :] if hold else "" + break + piece = "".join(out) + if piece: + yield piece + if buf and not in_think: + yield buf + + +# ── Graph construction ── + + +def _merge_result_into_kbinfos(tools, result: dict) -> None: + """Merge a search result's chunks/doc_aggs into ``tools.kbinfos``, deduped. + + Mirrors the orchestrators' merge so seed evidence and orchestrator evidence + share one deduplicated pool. + """ + if not result or not result.get("chunks"): + return + kb = tools.kbinfos + seen = {c.get("chunk_id") or c.get("id") or id(c) for c in kb.get("chunks", [])} + for c in result.get("chunks", []): + k = c.get("chunk_id") or c.get("id") or id(c) + if k in seen: + continue + seen.add(k) + kb.setdefault("chunks", []).append(c) + dseen = {d.get("doc_id") for d in kb.get("doc_aggs", [])} + for d in result.get("doc_aggs", []): + if d.get("doc_id") in dseen: + continue + dseen.add(d.get("doc_id")) + kb.setdefault("doc_aggs", []).append(d) + + +def build_agentic_graph(tools, token_queue: asyncio.Queue, gen_conf: dict | None = None): + """Compile the 4-node agentic-search graph.""" + answer_conf = dict(gen_conf) if gen_conf else {"temperature": 0.3} + + # ── Node: formalize_question ── + async def formalize_question(state: AgenticState) -> dict: + msgs = state.get("messages") or [] + _LOG.info("[formalize_question] IN | %d msg(s)", len(msgs)) + q, kw = await tools.formalize(msgs) + q = (q or "").strip() + kw = (kw or "").strip() + _LOG.info("[formalize_question] OUT | question=%s | keywords=%s", _snip(q), _snip(kw)) + return { + "question": q, + "keywords": kw, + "kbinfos": {"chunks": [], "doc_aggs": []}, + "loop": 0, + "partial_answer": False, + "abstain": False, + } + + # ── Node: route ── + async def route(state: AgenticState) -> dict: + from rag.advanced_rag.harness.route import route_node + + return await route_node(state, tools) + + # ── Node: pre_search ── + async def pre_search(state: AgenticState) -> dict: + """Preliminary hybrid_search to ground the planner's decomposition. + + Only runs for decomposition modes (direct/low mode retrieves in + orchestrator_loop anyway, so we skip the duplicate search). The result + is narrowed by keywords inside ``hybrid_search`` and merged into the + shared citation pool so it also enriches the final answer. + """ + route = state.get("route") + if not route or not getattr(route, "requires_decomposition", False): + _LOG.info("[pre_search] SKIP | direct/low mode (no decomposition)") + return {"seed_chunks": []} + + from rag.advanced_rag.harness.tools.search import hybrid_search + + q = state.get("question", "") + kw = state.get("keywords", "") + _LOG.info("[pre_search] IN | question=%s | keywords=%s", _snip(q), _snip(kw)) + try: + result = await hybrid_search(tools, query=q, keywords=kw) + except Exception: + _LOG.exception("[pre_search] hybrid_search failed") + return {"seed_chunks": []} + + chunks = result.get("chunks", []) or [] + _merge_result_into_kbinfos(tools, result) + _LOG.info("[pre_search] OUT | %d seed chunk(s), kbinfos now %d", len(chunks), len(tools.kbinfos.get("chunks", []))) + return {"seed_chunks": chunks} + + # ── Node: planner ── + async def planner(state: AgenticState) -> dict: + from rag.advanced_rag.harness.planner import planner_node + + return await planner_node(state, tools) + + # ── Node: orchestrator_loop ── + async def orchestrator_loop(state: AgenticState) -> dict: + from rag.advanced_rag.harness.orchestrator import orchestrator_loop as _run + + return await _run(state, tools) + + # ── Node: formalize_answer ── + async def formalize_answer(state: AgenticState) -> dict: + kbinfos = state.get("kbinfos") or {"chunks": [], "doc_aggs": []} + question = state.get("question") or "" + partial = state.get("partial_answer", False) + abstain = state.get("abstain", False) + empty_result = state.get("empty_result", False) + + _LOG.info("[formalize_answer] IN | question=%s | chunks=%d | partial=%s | abstain=%s", _snip(question), len(kbinfos["chunks"]), partial, abstain) + + tools.kbinfos = kbinfos + + # Abstain + if abstain: + msg = "I cannot answer this question based on the available information." + token_queue.put_nowait(msg) + return {"final_answer": msg} + + # Empty result + if empty_result or not kbinfos["chunks"]: + msg = "I don't have enough information based on the available sources." + token_queue.put_nowait(msg) + return {"final_answer": msg} + + # Build evidence + evidence = kb_prompt(kbinfos, tools.chat_mdl.max_length) + parts = [f"Question:\n{question}\n"] + + # Include pre_summary from agent results if available + pre_summary = kbinfos.get("pre_summary") + if pre_summary: + parts.append(f"Research Summary:\n{pre_summary}\n") + + if partial: + from rag.advanced_rag.harness.prompts.report_prompt import PARTIAL_ANSWER_PREAMBLE + + parts.append(f"{PARTIAL_ANSWER_PREAMBLE}\n") + + from rag.advanced_rag.harness.prompts.report_prompt import FINAL_ANSWER_SYSTEM + from rag.prompts.generator import citation_prompt as cp + + rules = cp(tools.user_defined_prompts).strip() + system = FINAL_ANSWER_SYSTEM.format(cite_rules=rules) + + parts.append(f"Evidence:\n{evidence}") + user_content = "\n".join(parts) + + _, msg = message_fit_in(form_message(system, user_content), tools.chat_mdl.max_length) + try: + async for tok in tools.chat_mdl.async_chat_streamly_delta(msg[0]["content"], msg[1:], answer_conf): + token_queue.put_nowait(tok) + except Exception: + _LOG.exception("formalize_answer: stream failed") + token_queue.put_nowait("I'm sorry, I encountered an error while composing the answer.") + + return {"final_answer": ""} + + # ── Build graph ── + g = StateGraph(AgenticState) + g.add_node("formalize_question", formalize_question) + g.add_node("route", route) + g.add_node("pre_search", pre_search) + g.add_node("planner", planner) + g.add_node("orchestrator_loop", orchestrator_loop) + g.add_node("formalize_answer", formalize_answer) + + g.add_edge(START, "formalize_question") + g.add_edge("formalize_question", "route") + g.add_edge("route", "pre_search") + g.add_edge("pre_search", "planner") + g.add_edge("planner", "orchestrator_loop") + g.add_edge("orchestrator_loop", "formalize_answer") + g.add_edge("formalize_answer", END) + + return g.compile() + + +# ── Runner ── + + +async def run_agentic_rag(tools, messages: list, max_loops: int = 3, gen_conf: dict | None = None): + """Drive the agentic-search graph, yielding answer-token strings.""" + _LOG.info("[agentic-rag] RUN START | %d message(s), max_loops=%d", len(messages or []), max_loops) + + token_queue: asyncio.Queue = asyncio.Queue() + graph = build_agentic_graph(tools, token_queue, gen_conf=gen_conf) + _SENTINEL = object() + holder: dict[str, Any] = {} + + async def _drive(): + try: + holder["state"] = await graph.ainvoke( + {"messages": messages}, + {"recursion_limit": max(25, max_loops * 8)}, + ) + except Exception: + logging.exception("run_agentic_rag: graph execution failed") + holder["error"] = True + finally: + token_queue.put_nowait(_SENTINEL) + + task = asyncio.create_task(_drive()) + produced = False + try: + while True: + item = await token_queue.get() + if item is _SENTINEL: + break + produced = True + yield item + finally: + await task + + state = holder.get("state") or {} + final_kb = state.get("kbinfos") + if isinstance(final_kb, dict) and final_kb.get("chunks"): + tools.kbinfos = final_kb + + _LOG.info("[agentic-rag] RUN END | streamed=%s, loops=%d, chunks=%d", produced, state.get("loop", 0), len((state.get("kbinfos") or {}).get("chunks", []))) + + if not produced and holder.get("error"): + yield "I couldn't complete the search due to an internal error." diff --git a/rag/advanced_rag/harness/__init__.py b/rag/advanced_rag/harness/__init__.py new file mode 100644 index 0000000000..af1a61737d --- /dev/null +++ b/rag/advanced_rag/harness/__init__.py @@ -0,0 +1 @@ +"""Harness: Agentic RAG orchestration layer.""" diff --git a/rag/advanced_rag/harness/agent.py b/rag/advanced_rag/harness/agent.py new file mode 100644 index 0000000000..31061e443e --- /dev/null +++ b/rag/advanced_rag/harness/agent.py @@ -0,0 +1,345 @@ +"""Research Agent — inner tool-calling loop for high/ultra modes. + +Native tool-calling: a chat model deep-copied from ``tools.chat_mdl`` is bound +(via ``bind_tools``) to the phase-gated tool schemas plus ``think_tool`` / +``generate_report``, and a lightweight session routes each tool call to the +harness pipeline. Binding onto a *copy* keeps the shared ``tools.chat_mdl`` +(used by the other graph nodes) free of any tool schema. + +Models without native tool-calling fall back to prompt-based tool selection: +the tools are described in the prompt and the model emits ```` JSON +that the loop parses. +""" + +import json +import logging +import re +from copy import deepcopy + +from rag.advanced_rag.harness.types import ClaimTarget, ExecutionStrategy, ToolResult +from rag.advanced_rag.harness.pipeline import Pipeline +from rag.advanced_rag.harness.tools.gating import ( + get_gated_tools, + determine_current_phase, + SEARCH_PHASES, +) +from rag.advanced_rag.harness.tools.registry import _generate_report_schema, _think_schema +from rag.advanced_rag.harness.prompts.research_agent_prompt import ( + RESEARCH_AGENT_PROMPT, + RESEARCH_AGENT_TEXT_PROMPT, +) + +_LOG = logging.getLogger(__name__) + + +class ResearchToolSession: + """ToolCallSession adapter routing native tool calls to the harness pipeline. + + - regular tools run through :func:`execute_with_fallback`; + - ``think_tool`` is a no-op reasoning step that just lets the loop continue; + - ``generate_report`` is *captured* (not executed) so the agent loop can + return its structured arguments as the claim result. + """ + + def __init__(self, pipeline: Pipeline, phase: str): + self.pipeline = pipeline + self.phase = phase + self.report: dict | None = None + self.got_evidence = False + self.evidence_ids: list[int] = [] + self._seen_evidence_ids: set[int] = set() + + async def tool_call_async(self, name: str, arguments: dict, request_timeout: float | int = 300): + arguments = arguments or {} + if name == "generate_report": + self.report = self._normalize_report(arguments) + return "Report received. Stop calling tools now." + if name == "think_tool": + return "Noted. Proceed with the next tool call." + result = await execute_with_fallback(self.pipeline, name, self.phase, **arguments) + if result.chunks: + self.got_evidence = True + self._record_evidence_ids(result.chunks) + return _fmt_tool_result(result) + + def _normalize_report(self, report: dict) -> dict: + normalized = dict(report) + evidence_ids = [] + for eid in normalized.get("evidence_ids") or []: + try: + idx = int(eid) + except (TypeError, ValueError): + continue + if idx not in evidence_ids: + evidence_ids.append(idx) + if not evidence_ids and self.evidence_ids: + evidence_ids = list(self.evidence_ids) + normalized["evidence_ids"] = evidence_ids + return normalized + + def _record_evidence_ids(self, chunks: list[dict]) -> None: + all_chunks = self.pipeline.tools.kbinfos.get("chunks", []) + index_by_key = {} + for idx, chunk in enumerate(all_chunks): + index_by_key[_chunk_key(chunk)] = idx + + for chunk in chunks: + idx = index_by_key.get(_chunk_key(chunk)) + if idx is None: + idx = next((i for i, existing in enumerate(all_chunks) if existing is chunk), None) + if idx is None or idx in self._seen_evidence_ids: + continue + self._seen_evidence_ids.add(idx) + self.evidence_ids.append(idx) + + +def _chunk_key(chunk: dict) -> object: + return chunk.get("chunk_id") or chunk.get("id") or id(chunk) + + +def _build_tool_schemas(gated_defs: list[dict]) -> list[dict]: + """Phase-gated schemas (minus harness-only ``x_*`` keys) + the control tools.""" + schemas: list[dict] = [] + for d in gated_defs: + schemas.append({k: v for k, v in d.items() if not k.startswith("x_")}) + schemas.append(_think_schema()) + schemas.append(_generate_report_schema()) + return schemas + + +async def research_agent_loop( + claim: ClaimTarget, + tools, + pipeline: Pipeline, + context, + mode: ExecutionStrategy, + compilation_map: dict, +) -> dict: + """Inner loop for a single claim — native tool-calling with a text fallback.""" + phase = determine_current_phase(context) + phase_config = SEARCH_PHASES.get(phase, {}) + gated_defs = get_gated_tools( + phase=phase, + available_tools=mode.available_tools, + compilation_map=compilation_map, + context=context, + ) + + # Deep-copy so binding tools never leaks onto the shared chat model. + agent_mdl = deepcopy(tools.chat_mdl) + if getattr(agent_mdl, "is_tools", False): + return await _research_native(claim, agent_mdl, pipeline, phase, phase_config, gated_defs, mode) + + _LOG.info("research_agent: model lacks native tool support; falling back to text-based tool selection") + return await _research_text(claim, tools, pipeline, phase, phase_config, gated_defs, mode) + + +async def _research_native( + claim: ClaimTarget, + agent_mdl, + pipeline: Pipeline, + phase: str, + phase_config: dict, + gated_defs: list[dict], + mode: ExecutionStrategy, +) -> dict: + """Bind tools onto ``agent_mdl`` and let its native tool loop drive research.""" + schemas = _build_tool_schemas(gated_defs) + session = ResearchToolSession(pipeline, phase) + agent_mdl.bind_tools(session, schemas) + # Bound the model's internal tool loop to the mode's agent-cycle budget. + if hasattr(agent_mdl, "mdl") and hasattr(agent_mdl.mdl, "max_rounds"): + agent_mdl.mdl.max_rounds = max(1, mode.max_agent_cycles) + + system = RESEARCH_AGENT_PROMPT.format( + claim_description=claim.description, + phase=phase, + phase_hint=phase_config.get("tool_hint", ""), + max_cycles=mode.max_agent_cycles, + ) + history = [{"role": "user", "content": f"Research task: {claim.description}\nBegin."}] + + final_text = "" + try: + final_text = await agent_mdl.async_chat(system, history, {"temperature": 0.3}) + if isinstance(final_text, tuple): + final_text = final_text[0] + except Exception: + _LOG.exception("research_agent(native): tool loop failed") + + if session.report is not None: + return session.report + + # The model finished without calling generate_report — synthesize a report + # from its final free-text turn so the claim still yields something usable. + _LOG.info("research_agent(native): no generate_report call; using final text as report") + return { + "report": (final_text or "").strip(), + "is_verified": session.got_evidence, + "confidence": 0.5 if session.got_evidence else 0.0, + "evidence_ids": list(session.evidence_ids), + "gaps": [] if session.got_evidence else ["no generate_report emitted"], + "discovered_claims": [], + } + + +async def _research_text( + claim: ClaimTarget, + tools, + pipeline: Pipeline, + phase: str, + phase_config: dict, + gated_defs: list[dict], + mode: ExecutionStrategy, +) -> dict: + """Fallback: prompt-based tool selection for models without native tools.""" + system = RESEARCH_AGENT_TEXT_PROMPT.format( + claim_description=claim.description, + phase=phase, + phase_hint=phase_config.get("tool_hint", ""), + tool_list=_fmt_tool_list(gated_defs), + max_cycles=mode.max_agent_cycles, + ) + + history: list[dict] = [] + + for cycle in range(mode.max_agent_cycles): + try: + ans = await tools.chat_mdl.async_chat(system, history, {"temperature": 0.3}) + if isinstance(ans, tuple): + ans = ans[0] + except Exception: + _LOG.exception("research_agent(text): LLM call failed cycle %d", cycle) + continue + + history.append({"role": "assistant", "content": ans}) + + tool_call = _parse_tool_call(ans) + if not tool_call: + history.append({"role": "user", "content": "Please call a tool. Do not output plain text."}) + continue + + if tool_call.get("name") == "generate_report": + return tool_call.get("arguments", {}) + + if tool_call.get("name") == "think_tool": + history.append({"role": "user", "content": "[continue]"}) + continue + + args = tool_call.get("arguments", {}) + result = await execute_with_fallback(pipeline, tool_call["name"], phase, **args) + history.append({"role": "user", "content": _fmt_tool_result(result)}) + + return await _force_generate_report(history, tools, claim.claim_id) + + +def _parse_tool_call(text: str) -> dict | None: + """Parse tool call from LLM response text (text-fallback path).""" + m = re.search(r"(.*?)", text, re.DOTALL) + if m: + try: + return json.loads(m.group(1).strip()) + except Exception: + pass + + m = re.search(r"```(?:json)?\s*({.*?})\s*```", text, re.DOTALL) + if m: + try: + return json.loads(m.group(1).strip()) + except Exception: + pass + + m = re.search(r'\{\s*"name"\s*:', text) + if m: + try: + import json_repair + + return json_repair.loads(text) + except Exception: + pass + + return None + + +async def execute_with_fallback( + pipeline: Pipeline, + tool_name: str, + phase: str, + **kwargs, +) -> ToolResult: + """Execute tool; if empty, fall back along phase priority.""" + result = await pipeline.execute(tool_name, **kwargs) + + if result.chunks or result.error: + return result + + phase_config = SEARCH_PHASES.get(phase, {}) + priority = phase_config.get("tools_priority", []) + current_idx = next( + (i for i, t in enumerate(priority) if t == tool_name), + -1, + ) + for fallback_name in priority[current_idx + 1 :]: + fallback_result = await pipeline.execute(fallback_name, **kwargs) + if fallback_result.chunks: + _LOG.info("fallback: %s empty → %s found %d chunks", tool_name, fallback_name, len(fallback_result.chunks)) + fallback_result.metadata["was_fallback"] = True + fallback_result.metadata["fallback_from"] = tool_name + return fallback_result + if fallback_result.error: + break + return result + + +async def _force_generate_report( + history: list, + tools, + claim_id: str, +) -> dict: + """Force generate report when max cycles reached (text-fallback path).""" + try: + ans = await tools.chat_mdl.async_chat( + "", + history + [{"role": "user", "content": "We've reached the research limit. Please output a final report as JSON."}], + {"temperature": 0.3}, + ) + if isinstance(ans, tuple): + ans = ans[0] + text = re.sub(r"```(?:json)?\s*|\s*```", "", ans).strip() + import json_repair + + return json_repair.loads(text) + except Exception: + _LOG.exception("force_generate_report failed") + return { + "report": "", + "is_verified": False, + "confidence": 0.0, + "evidence_ids": [], + "gaps": ["forced report — data may be incomplete"], + "discovered_claims": [], + } + + +def _fmt_tool_list(defs: list[dict]) -> str: + lines = [] + for d in defs: + func = d.get("function", d) + name = func.get("name", "?") + desc = func.get("description", "") + params = func.get("parameters", {}).get("properties", {}) + params_text = ", ".join(f"{k}: {v.get('description', '')}" for k, v in params.items()) + lines.append(f"- {name}: {desc}") + if params_text: + lines.append(f" Parameters: {params_text}") + return "\n".join(lines) + + +def _fmt_tool_result(result: ToolResult) -> str: + if result.error: + return f"[tool error] {result.error}" + chunks = result.chunks[:3] + texts = [c.get("content_with_weight", c.get("text", ""))[:300] for c in chunks] + if not texts: + return "[no results found]" + return "\n\n".join(texts) diff --git a/rag/advanced_rag/harness/config.py b/rag/advanced_rag/harness/config.py new file mode 100644 index 0000000000..af9ffbd42e --- /dev/null +++ b/rag/advanced_rag/harness/config.py @@ -0,0 +1,103 @@ +"""Thinking mode configurations.""" + +from rag.advanced_rag.harness.types import ExecutionStrategy + +THINKING_MODES: dict[str, ExecutionStrategy] = { + "low": ExecutionStrategy( + label="low", + execution_strategy="direct_search", + requires_decomposition=False, + requires_agent_loop=False, + requires_sufficiency_judge=False, + requires_selective_gen=False, + allows_dynamic_claims=False, + allows_replan=False, + max_orchestrator_cycles=1, + max_agent_cycles=0, + max_parallel_agents=1, + available_tools=["hybrid_search"], + sufficiency_threshold=0.85, + partial_threshold=0.50, + fallback_to_direct_llm=False, + ), + "medium": ExecutionStrategy( + label="medium", + execution_strategy="decompose_and_search", + requires_decomposition=True, + requires_agent_loop=False, + requires_sufficiency_judge=True, + requires_selective_gen=True, + allows_dynamic_claims=False, + allows_replan=False, + max_orchestrator_cycles=3, + max_agent_cycles=0, + max_parallel_agents=1, + available_tools=["hybrid_search"], + sufficiency_threshold=0.75, + partial_threshold=0.40, + fallback_to_direct_llm=False, + ), + "high": ExecutionStrategy( + label="high", + execution_strategy="agentic_research", + requires_decomposition=True, + requires_agent_loop=True, + requires_sufficiency_judge=True, + requires_selective_gen=True, + allows_dynamic_claims=False, + allows_replan=False, + max_orchestrator_cycles=3, + max_agent_cycles=2, + max_parallel_agents=2, + available_tools=[ + "hybrid_search", + "web_search", + "toc_navigate", + "page_index_navigate", + "graph_explore", + "inspector_open_context", + "inspector_compare", + ], + sufficiency_threshold=0.65, + partial_threshold=0.30, + fallback_to_direct_llm=False, + ), + "ultra": ExecutionStrategy( + label="ultra", + execution_strategy="deep_research", + requires_decomposition=True, + requires_agent_loop=True, + requires_sufficiency_judge=True, + requires_selective_gen=True, + allows_dynamic_claims=True, + allows_replan=True, + max_orchestrator_cycles=4, + max_agent_cycles=2, + max_parallel_agents=3, + available_tools=[ + "hybrid_search", + "bm25_search", + "web_search", + "structured_query", + "toc_navigate", + "page_index_navigate", + "mindmap_navigate", + "graph_explore", + "wiki_query", + "inspector_open_context", + "inspector_compare", + "inspector_grep_within", + "inspector_request_adjacent", + ], + sufficiency_threshold=0.55, + partial_threshold=0.20, + fallback_to_direct_llm=True, + ), +} + + +def get_mode(label: str) -> ExecutionStrategy: + mode = THINKING_MODES.get(label) + if not mode: + raise ValueError(f"Unknown thinking mode: {label}. Available: {list(THINKING_MODES.keys())}") + return mode diff --git a/rag/advanced_rag/harness/orchestrator/__init__.py b/rag/advanced_rag/harness/orchestrator/__init__.py new file mode 100644 index 0000000000..66b6e35497 --- /dev/null +++ b/rag/advanced_rag/harness/orchestrator/__init__.py @@ -0,0 +1,35 @@ +"""Orchestrator loop — dispatches to execution strategy based on thinking mode.""" + +import logging + +from rag.advanced_rag.harness.config import get_mode +from rag.advanced_rag.harness.orchestrator.direct import direct_search +from rag.advanced_rag.harness.orchestrator.decompose import decompose_and_search +from rag.advanced_rag.harness.orchestrator.agentic import agentic_research + +_LOG = logging.getLogger(__name__) + + +async def orchestrator_loop(state: dict, tools) -> dict: + """Main orchestrator — dispatch to strategy based on thinking mode.""" + route = state.get("route") + if not route: + _LOG.warning("orchestrator: no route, using direct_search") + return await direct_search(state, tools) + + mode_label = route.thinking_mode if isinstance(route, dict) else route.thinking_mode + mode = get_mode(mode_label) + + _LOG.info("[orchestrator] strategy=%s mode=%s", mode.execution_strategy, mode_label) + + if mode.execution_strategy == "direct_search": + return await direct_search(state, tools) + + if mode.execution_strategy == "decompose_and_search": + return await decompose_and_search(state, tools) + + if mode.execution_strategy in ("agentic_research", "deep_research"): + return await agentic_research(state, tools) + + _LOG.warning("orchestrator: unknown strategy %s, fallback to direct", mode.execution_strategy) + return await direct_search(state, tools) diff --git a/rag/advanced_rag/harness/orchestrator/agentic.py b/rag/advanced_rag/harness/orchestrator/agentic.py new file mode 100644 index 0000000000..844e0306f8 --- /dev/null +++ b/rag/advanced_rag/harness/orchestrator/agentic.py @@ -0,0 +1,241 @@ +"""High/Ultra: two-level loop — orchestrator assigns claims, agent researches, sufficiency checks.""" + +import asyncio +import logging + +from rag.advanced_rag.harness.types import ( + ClaimTarget, + AgentResult, + OrchestratorContext, +) +from rag.advanced_rag.harness.config import get_mode +from rag.advanced_rag.harness.pipeline import Pipeline +from rag.advanced_rag.harness.agent import research_agent_loop +from rag.advanced_rag.harness.sufficiency import ( + cross_check_claim, + compute_fusion_score, + route_sufficiency_verdict, +) + +_LOG = logging.getLogger(__name__) +CLAIM_RESEARCH_TIMEOUT_SECONDS = 180 + + +def _snip(text: str, limit: int = 160) -> str: + text = (text or "").replace("\n", " ").strip() + return text if len(text) <= limit else text[: limit - 3] + "..." + + +async def agentic_research(state: dict, tools) -> dict: + """Two-level loop for high/ultra modes.""" + question = state.get("question", "") + claims_raw = state.get("claims", []) + route = state.get("route", {}) + mode_label = route.thinking_mode if route else "high" + mode = get_mode(mode_label) + + # Resolve compilation map + compilation_map = await _get_compilation_map(tools) + + claims = [ClaimTarget(**c) if isinstance(c, dict) else c for c in claims_raw] + ctx = OrchestratorContext(question=question, claims=claims, mode=mode_label) + pipeline = Pipeline(tools, compilation_map) + + for cycle in range(mode.max_orchestrator_cycles): + ctx.iteration = cycle + _LOG.info("[agentic] cycle %d/%d, claims: %d unverified", cycle + 1, mode.max_orchestrator_cycles, sum(1 for c in ctx.claims if not c.is_verified)) + + # ── Step A: Research unverified claims (parallel if mode allows) ── + unverified = [c for c in ctx.claims if not c.is_verified] + + if unverified: + # Process in batches of max_parallel_agents + batch_size = mode.max_parallel_agents + for i in range(0, len(unverified), batch_size): + batch = unverified[i : i + batch_size] + _LOG.info( + "[agentic] batch start cycle=%d offset=%d size=%d claim_ids=%s", + cycle + 1, + i, + len(batch), + [c.claim_id for c in batch], + ) + tasks = [_run_claim_research(c, tools, pipeline, ctx, mode, compilation_map) for c in batch] + agent_results = await asyncio.gather(*tasks) + _LOG.info( + "[agentic] batch done cycle=%d offset=%d size=%d", + cycle + 1, + i, + len(agent_results), + ) + + for c, result in zip(batch, agent_results): + is_verified = result.get("is_verified", False) + c.is_verified = is_verified + c.confidence = result.get("confidence", 0.0) + c.agent_result = AgentResult( + claim_id=c.claim_id, + report=result.get("report", ""), + is_verified=is_verified, + confidence=c.confidence, + evidence_ids=result.get("evidence_ids", []), + gaps=result.get("gaps", []), + discovered_claims=result.get("discovered_claims", []), + ) + + # Ultra: dynamic claim expansion + if mode.allows_dynamic_claims and result.get("discovered_claims"): + for dc in result["discovered_claims"]: + if dc and dc not in [cc.description for cc in ctx.claims]: + ctx.claims.append( + ClaimTarget( + claim_id=f"c_dyn_{len(ctx.claims)}", + description=dc, + ) + ) + _LOG.info("[agentic] discovered new claim: %s", dc) + + # ── Step B: Sufficiency Check ── + all_chunks = {i: c for i, c in enumerate(tools.kbinfos.get("chunks", []))} + agent_results_list = [c.agent_result for c in ctx.claims if c.agent_result] + cross_results = [cross_check_claim(r, all_chunks) for r in agent_results_list] + + verdict = compute_fusion_score(agent_results_list, cross_results, mode) + ctx.verdict = verdict + + action, should_continue = route_sufficiency_verdict( + verdict, + mode_label, + cycle, + mode.max_orchestrator_cycles, + ) + + _LOG.info("[agentic] cycle=%d verdict=%s score=%.2f action=%s", cycle, verdict.status, verdict.score, action) + + if action == "ANSWER": + return _finalize(ctx, tools, partial=False) + if action == "ANSWER_PARTIAL": + return _finalize(ctx, tools, partial=True) + if action == "ABSTAIN": + tools.kbinfos["chunks"] = [] + return {"verdict": verdict.__dict__, "abstain": True} + if action == "REPLAN": + # Ultra: re-plan on low score + from rag.advanced_rag.harness.planner import planner_node + + state["feedback"] = verdict.feedback + state["route"] = route + new_plan = await planner_node(state, tools) + ctx.claims = new_plan.get("claims", ctx.claims) + if action == "FALLBACK_LLM": + return _finalize(ctx, tools, partial=True, fallback=True) + + # Max cycles reached + return _finalize(ctx, tools, partial=True) + + +async def _run_claim_research( + claim: ClaimTarget, + tools, + pipeline: Pipeline, + ctx: OrchestratorContext, + mode, + compilation_map: dict, +) -> dict: + _LOG.info("[agentic] claim start id=%s desc=%s", claim.claim_id, _snip(claim.description)) + try: + result = await asyncio.wait_for( + research_agent_loop(claim, tools, pipeline, ctx, mode, compilation_map), + timeout=CLAIM_RESEARCH_TIMEOUT_SECONDS, + ) + except asyncio.CancelledError: + raise + except asyncio.TimeoutError: + _LOG.warning( + "[agentic] claim timeout id=%s timeout=%ss desc=%s", + claim.claim_id, + CLAIM_RESEARCH_TIMEOUT_SECONDS, + _snip(claim.description), + ) + return { + "report": "", + "is_verified": False, + "confidence": 0.0, + "evidence_ids": [], + "gaps": [f"claim research timeout after {CLAIM_RESEARCH_TIMEOUT_SECONDS}s"], + "discovered_claims": [], + } + except Exception: + _LOG.exception("[agentic] claim failed id=%s desc=%s", claim.claim_id, _snip(claim.description)) + return { + "report": "", + "is_verified": False, + "confidence": 0.0, + "evidence_ids": [], + "gaps": ["claim research failed"], + "discovered_claims": [], + } + + _LOG.info( + "[agentic] claim done id=%s verified=%s confidence=%.2f evidence=%d gaps=%d", + claim.claim_id, + result.get("is_verified", False), + float(result.get("confidence") or 0.0), + len(result.get("evidence_ids") or []), + len(result.get("gaps") or []), + ) + return result + + +def _finalize(ctx: OrchestratorContext, tools, partial: bool = False, fallback: bool = False) -> dict: + """Merge agent results into kbinfos and return.""" + _merge_agent_results(ctx, tools) + return { + "verdict": ctx.verdict.__dict__ if ctx.verdict else None, + "partial_answer": partial or fallback, + "kbinfos": tools.kbinfos, + } + + +def _merge_agent_results(ctx: OrchestratorContext, tools): + """Merge agent result reports into kbinfos as a pre_summary.""" + combined = [] + seen_evidence = set() + + for c in ctx.claims: + if c.agent_result and c.agent_result.report: + status = "✅" if c.is_verified else "❌" + combined.append(f"【{c.claim_id}】{status} {c.agent_result.report[:500]}") + + if combined: + tools.kbinfos["pre_summary"] = "\n\n".join(combined) + + # Collect evidence chunks from agent results + for c in ctx.claims: + if c.agent_result and c.agent_result.evidence_ids: + for eid in c.agent_result.evidence_ids: + if eid not in seen_evidence: + seen_evidence.add(eid) + + +async def _get_compilation_map(tools) -> dict[str, set[str]]: + """Build compilation map from RAGTools - check which KBs have compilation artifacts.""" + result = {} + if not tools.kbs: + return result + for kb in tools.kbs: + comps = set() + parser_config = getattr(kb, "parser_config", None) or {} + if parser_config.get("toc"): + comps.add("toc") + if parser_config.get("knowledge_graph"): + comps.add("knowledge_graph") + if parser_config.get("wiki"): + comps.add("wiki") + if parser_config.get("mindmap"): + comps.add("mindmap") + if parser_config.get("page_index"): + comps.add("page_index") + if comps: + result[kb.id] = comps + return result diff --git a/rag/advanced_rag/harness/orchestrator/decompose.py b/rag/advanced_rag/harness/orchestrator/decompose.py new file mode 100644 index 0000000000..ea23a17cf7 --- /dev/null +++ b/rag/advanced_rag/harness/orchestrator/decompose.py @@ -0,0 +1,112 @@ +"""Medium mode: decompose → parallel search → sufficiency check.""" + +import asyncio +import logging + +from rag.advanced_rag.harness.types import ClaimTarget, AgentResult, OrchestratorContext +from rag.advanced_rag.harness.config import get_mode +from rag.advanced_rag.harness.sufficiency import ( + cross_check_claim, + compute_fusion_score, + route_sufficiency_verdict, +) +from rag.advanced_rag.harness.tools.search import hybrid_search + +_LOG = logging.getLogger(__name__) + + +async def decompose_and_search(state: dict, tools) -> dict: + """Decompose → parallel search → merge → sufficiency check → iterate.""" + question = state.get("question", "") + keywords = state.get("keywords", "") + claims_raw = state.get("claims", []) + mode_label = state.get("route", {}).thinking_mode if state.get("route") else "medium" + mode = get_mode(mode_label) + + claims = [ClaimTarget(**c) if isinstance(c, dict) else c for c in claims_raw] + ctx = OrchestratorContext(question=question, claims=claims, mode=mode_label) + + for cycle in range(mode.max_orchestrator_cycles): + ctx.iteration = cycle + unverified = [c for c in ctx.claims if not c.is_verified] + if not unverified: + break + + # Parallel search on unverified claims + tasks = [] + for c in unverified: + tasks.append(hybrid_search(tools, query=c.description, keywords=keywords)) + results = await asyncio.gather(*tasks) + + for c, result in zip(unverified, results): + if result.get("chunks"): + c.is_verified = True + c.confidence = 0.8 + c.agent_result = AgentResult( + claim_id=c.claim_id, + report=_summarize(result), + is_verified=True, + confidence=0.8, + evidence_ids=list(range(len(result.get("chunks", [])))), + ) + _merge_kbinfos(tools, result) + else: + c.agent_result = AgentResult( + claim_id=c.claim_id, + report="", + is_verified=False, + confidence=0.0, + ) + + all_chunks = {i: c for i, c in enumerate(tools.kbinfos.get("chunks", []))} + agent_results = [c.agent_result for c in ctx.claims if c.agent_result] + cross_results = [cross_check_claim(r, all_chunks) for r in agent_results] + + verdict = compute_fusion_score(agent_results, cross_results, mode) + + action, should_continue = route_sufficiency_verdict( + verdict, + mode_label, + cycle, + mode.max_orchestrator_cycles, + ) + + if action in ("ANSWER", "ANSWER_PARTIAL"): + return { + "verdict": verdict.__dict__, + "partial_answer": action == "ANSWER_PARTIAL", + "kbinfos": tools.kbinfos, + } + if action == "ABSTAIN": + tools.kbinfos["chunks"] = [] + return {"verdict": verdict.__dict__, "abstain": True} + + return {"kbinfos": tools.kbinfos} + + +def _merge_kbinfos(tools, result: dict): + if not result or not result.get("chunks"): + return + seen = {_chunk_key(c) for c in tools.kbinfos.get("chunks", [])} + for c in result.get("chunks", []): + k = _chunk_key(c) + if k in seen: + continue + seen.add(k) + tools.kbinfos.setdefault("chunks", []).append(c) + dseen = {d.get("doc_id") for d in tools.kbinfos.get("doc_aggs", [])} + for d in result.get("doc_aggs", []): + if d.get("doc_id") in dseen: + continue + dseen.add(d.get("doc_id")) + tools.kbinfos.setdefault("doc_aggs", []).append(d) + + +def _chunk_key(ck: dict) -> str: + return ck.get("chunk_id") or ck.get("id") or str(id(ck)) + + +def _summarize(result: dict) -> str: + chunks = result.get("chunks", []) + texts = [c.get("content_with_weight", "")[:200] for c in chunks[:3]] + return " | ".join(texts) diff --git a/rag/advanced_rag/harness/orchestrator/direct.py b/rag/advanced_rag/harness/orchestrator/direct.py new file mode 100644 index 0000000000..0059771a79 --- /dev/null +++ b/rag/advanced_rag/harness/orchestrator/direct.py @@ -0,0 +1,49 @@ +"""Low mode: direct single-pass search.""" + +import logging + +from rag.advanced_rag.harness.tools.search import hybrid_search + +_LOG = logging.getLogger(__name__) + + +async def direct_search(state: dict, tools) -> dict: + """Single hybrid search → merge into kbinfos.""" + question = state.get("question", "") + keywords = state.get("keywords", "") + _LOG.info("[direct_search] question=%s | keywords=%s", question, keywords) + + result = await hybrid_search(tools, query=question, keywords=keywords) + _merge_kbinfos(tools, result) + + if not _has_chunks(tools): + _LOG.info("[direct_search] no results found") + return {"empty_result": True, "kbinfos": tools.kbinfos} + + return {"kbinfos": tools.kbinfos} + + +def _merge_kbinfos(tools, result: dict): + if not result or not result.get("chunks"): + return + seen = {_chunk_key(c) for c in tools.kbinfos.get("chunks", [])} + for c in result.get("chunks", []): + k = _chunk_key(c) + if k in seen: + continue + seen.add(k) + tools.kbinfos.setdefault("chunks", []).append(c) + dseen = {d.get("doc_id") for d in tools.kbinfos.get("doc_aggs", [])} + for d in result.get("doc_aggs", []): + if d.get("doc_id") in dseen: + continue + dseen.add(d.get("doc_id")) + tools.kbinfos.setdefault("doc_aggs", []).append(d) + + +def _chunk_key(ck: dict) -> str: + return ck.get("chunk_id") or ck.get("id") or str(id(ck)) + + +def _has_chunks(tools) -> bool: + return bool(tools.kbinfos.get("chunks")) diff --git a/rag/advanced_rag/harness/pipeline.py b/rag/advanced_rag/harness/pipeline.py new file mode 100644 index 0000000000..a099436aea --- /dev/null +++ b/rag/advanced_rag/harness/pipeline.py @@ -0,0 +1,125 @@ +"""Pipeline — unified tool execution dispatcher.""" + +import time +import logging +from typing import Any + +from rag.advanced_rag.harness.types import ToolResult +from rag.advanced_rag.harness.tools.registry import TOOL_REGISTRY + +_LOG = logging.getLogger(__name__) + + +class Pipeline: + """Unified tool execution layer. + + - execute(tool_name, **kwargs): dispatch to registered tool, normalize result + - available_tools(mode_tools): return LLM-visible tool definitions (compilation-filtered) + - get_chunks(evidence_ids): retrieve raw chunks for sufficiency cross-check + - trace: execution history for auditing + """ + + def __init__(self, rag_tools, compilation_map: dict[str, set[str]] | None = None): + self.tools = rag_tools + self.compilation_map = compilation_map or {} + self.trace: list[dict] = [] + + async def execute(self, tool_name: str, **kwargs) -> ToolResult: + """Execute a registered tool by name.""" + tool = TOOL_REGISTRY.get(tool_name) + if not tool: + return ToolResult(chunks=[], metadata={}, error=f"Unknown tool: {tool_name}") + + fn = tool.get("fn") + if not fn: + return ToolResult(chunks=[], metadata={}, error=f"Tool {tool_name} has no executor") + + start = time.time() + try: + raw = await fn(self.tools, **kwargs) + elapsed = time.time() - start + self.trace.append({"tool": tool_name, "args": kwargs, "elapsed": elapsed, "success": True}) + result = self._normalize(raw) + # Feed the shared citation pool: agent searches go through the + # pipeline, so without this their evidence never reaches kbinfos and + # the final answer has nothing to cite. + self._merge_into_kbinfos(result) + return result + except Exception as e: + elapsed = time.time() - start + _LOG.exception("Pipeline.execute(%s) failed", tool_name) + self.trace.append({"tool": tool_name, "args": kwargs, "elapsed": elapsed, "success": False, "error": str(e)}) + return ToolResult(chunks=[], metadata={}, error=str(e)) + + def available_tools(self, mode_tools: list[str]) -> list[dict]: + """Return LLM-visible tool definitions, filtered by compilation availability.""" + names = filter_available_tools(mode_tools, self.compilation_map) + defs = [] + for name in names: + tool = TOOL_REGISTRY.get(name) + if tool and tool.get("function_schema"): + defs.append(tool["function_schema"]) + return defs + + def get_chunks(self, evidence_ids: list[int]) -> dict[int, dict]: + """Retrieve raw chunks by ID from current kbinfos.""" + result = {} + chunks = self.tools.kbinfos.get("chunks", []) + for eid in evidence_ids: + if 0 <= eid < len(chunks): + result[eid] = chunks[eid] + return result + + def get_trace(self) -> list[dict]: + return list(self.trace) + + # ── Private ── + + def _merge_into_kbinfos(self, result: ToolResult) -> None: + """Merge a tool result's chunks/doc_aggs into ``tools.kbinfos``, deduped.""" + if not result or not result.chunks: + return + kb = self.tools.kbinfos + seen = {c.get("chunk_id") or c.get("id") or id(c) for c in kb.get("chunks", [])} + for c in result.chunks: + k = c.get("chunk_id") or c.get("id") or id(c) + if k in seen: + continue + seen.add(k) + kb.setdefault("chunks", []).append(c) + aggs = result.metadata.get("aggs") if isinstance(result.metadata, dict) else None + if aggs: + dseen = {d.get("doc_id") for d in kb.get("doc_aggs", [])} + for d in aggs: + if d.get("doc_id") in dseen: + continue + dseen.add(d.get("doc_id")) + kb.setdefault("doc_aggs", []).append(d) + + @staticmethod + def _normalize(raw: Any) -> ToolResult: + if isinstance(raw, ToolResult): + return raw + if isinstance(raw, dict): + return ToolResult( + chunks=raw.get("chunks", []), + metadata={"aggs": raw.get("doc_aggs", []), "answer": raw.get("answer", "")}, + ) + if isinstance(raw, list): + return ToolResult(chunks=raw, metadata={}) + return ToolResult(chunks=[], metadata={"raw": str(raw)}) + + +def filter_available_tools(tool_names: list[str], compilation_map: dict[str, set[str]]) -> list[str]: + """Filter tool list by compilation artifact availability.""" + available = [] + for name in tool_names: + tool = TOOL_REGISTRY.get(name) + if not tool: + continue + if tool.get("requires_compilation"): + comp_type = tool.get("compilation_type") + if comp_type and not any(comp_type in comps for comps in compilation_map.values()): + continue + available.append(name) + return available diff --git a/rag/advanced_rag/harness/planner.py b/rag/advanced_rag/harness/planner.py new file mode 100644 index 0000000000..072db5b739 --- /dev/null +++ b/rag/advanced_rag/harness/planner.py @@ -0,0 +1,147 @@ +"""Planner node — question-type-aware claim decomposition.""" + +import json +import logging +import re + +from rag.advanced_rag.agentic_rag_graph import _snip +from rag.advanced_rag.harness.types import ClaimTarget, WorkflowPlan, RouteDecision +from rag.advanced_rag.harness.config import get_mode +from rag.advanced_rag.harness.prompts.decompose_prompts import ( + DECOMPOSE_FACTUAL, + DECOMPOSE_COMPARATIVE, + DECOMPOSE_PROCEDURAL, + DECOMPOSE_EXPLORATORY, +) + +_LOG = logging.getLogger(__name__) + + +def _extract_json(text: str) -> dict: + text = re.sub(r"^.*", "", text, flags=re.DOTALL).strip() + text = re.sub(r"```(?:json)?\s*|\s*```", "", text).strip() + try: + import json_repair + + return json_repair.loads(text) + except Exception: + try: + return json.loads(text) + except Exception: + _LOG.warning("planner: failed to parse LLM output: %s", text[:200]) + return {} + + +async def planner_node(state: dict, tools) -> dict: + """Planner node — decompose question into claims based on question type.""" + route: RouteDecision = state.get("route") + if not route: + _LOG.warning("planner: no route found, using defaults") + return _default_plan(state.get("question", "")) + + _LOG.info("[Planner] IN | question=%s type=%s", _snip(route.question), route.question_type) + if not route.requires_decomposition: + # Direct mode: single coarse claim + return _direct_plan(route.question) + + # Select decompose prompt by question type + prompt_map = { + "factual": DECOMPOSE_FACTUAL, + "comparative": DECOMPOSE_COMPARATIVE, + "procedural": DECOMPOSE_PROCEDURAL, + "analytical": DECOMPOSE_EXPLORATORY, + "exploratory": DECOMPOSE_EXPLORATORY, + } + decompose_prompt = prompt_map.get(route.question_type, DECOMPOSE_FACTUAL) + + mode = get_mode(route.thinking_mode) + max_claims = _get_max_claims(mode.label) + detail_level = _get_detail_level(mode.label) + retrieved = _format_seed_chunks(state.get("seed_chunks"), tools) + + try: + prompt = decompose_prompt.format( + question=route.question, + max_claims=max_claims, + detail_level=detail_level, + retrieved=retrieved, + ) + system, user = prompt.split("Output format", 1) + system = system.strip() + user = "Output format" + user + msg = await tools._fit_messages(system, user) + ans = await tools.chat_mdl.async_chat(msg[0]["content"], msg[1:], {"temperature": 0.2}) + if isinstance(ans, tuple): + ans = ans[0] + result = _extract_json(ans) + except Exception: + _LOG.exception("planner_node failed") + return _direct_plan(route.question) + + claims_raw = result.get("claims", []) + plan_type = { + "factual": "fact_decomposition", + "comparative": "comparative_decomposition", + "procedural": "procedural_decomposition", + }.get(route.question_type, "exploratory_decomposition") + + claims = [] + for i, c in enumerate(claims_raw): + if isinstance(c, dict) and c.get("description"): + claims.append( + ClaimTarget( + claim_id=c.get("claim_id", f"c{i}"), + description=c["description"], + priority=c.get("priority", 0), + suggested_tools=c.get("suggested_tools", []), + ) + ) + + if not claims: + return _direct_plan(route.question) + + plan = WorkflowPlan( + plan_type=plan_type, + claims=claims, + max_iterations=mode.max_orchestrator_cycles, + ) + _LOG.info("[Planner] OUT | plan type=%s | claims=%d", plan_type, len(plan.claims)) + + return {"plan": plan, "claims": plan.claims} + + +def _format_seed_chunks(seed_chunks, tools) -> str: + """Render preliminary-search chunks as grounding context for the planner.""" + if not seed_chunks: + return "(no preliminary results)" + try: + from rag.prompts.generator import kb_prompt + + blocks = kb_prompt({"chunks": seed_chunks, "doc_aggs": []}, tools.chat_mdl.max_length) + text = "\n".join(blocks).strip() + return text or "(no preliminary results)" + except Exception: + _LOG.exception("planner: failed to format seed chunks") + return "(no preliminary results)" + + +def _direct_plan(question: str) -> dict: + """Single-claim plan for non-decomposed mode.""" + plan = WorkflowPlan( + plan_type="direct", + claims=[ClaimTarget(claim_id="c0", description=question, priority=0)], + max_iterations=1, + ) + return {"plan": plan, "claims": plan.claims} + + +def _default_plan(question: str) -> dict: + return _direct_plan(question) + + +def _get_max_claims(mode_label: str) -> int: + return {"low": 1, "medium": 3, "high": 5, "ultra": 8}.get(mode_label, 3) + + +def _get_detail_level(mode_label: str) -> str: + return {"low": "coarse", "medium": "normal", "high": "fine", "ultra": "extra_fine"}.get(mode_label, "normal") diff --git a/rag/advanced_rag/harness/prompts/__init__.py b/rag/advanced_rag/harness/prompts/__init__.py new file mode 100644 index 0000000000..2c759e9a86 --- /dev/null +++ b/rag/advanced_rag/harness/prompts/__init__.py @@ -0,0 +1 @@ +"""Prompt templates for Agentic RAG harness.""" diff --git a/rag/advanced_rag/harness/prompts/decompose_prompts.py b/rag/advanced_rag/harness/prompts/decompose_prompts.py new file mode 100644 index 0000000000..db5e15548d --- /dev/null +++ b/rag/advanced_rag/harness/prompts/decompose_prompts.py @@ -0,0 +1,121 @@ +"""Planner decompose prompts: one per question type. + +Each prompt is grounded in a preliminary hybrid search: ``{retrieved}`` carries +the (keyword-narrowed) chunks retrieved for the user's question, so the +decomposition reflects what the corpus actually contains rather than guessing. +""" + +DECOMPOSE_FACTUAL = """This is a factual question. List all atomic facts that need to be retrieved. +If there are multiple facts, list them one by one. If there is only one fact, output exactly one item. +Base the decomposition on BOTH the question and the preliminary retrieved context below. + +Question: {question} +Maximum number of claims: {max_claims} +Detail level: {detail_level} + +Preliminary retrieved context for this question (use it to ground each claim in what the corpus actually contains; do not invent facts it cannot support): +{retrieved} + +Output format (JSON): +{{ + "claims": [ + {{ + "claim_id": "c1", + "description": "The year Apple acquired Beats", + "priority": 1 + }} + ] +}} +""" + + +DECOMPOSE_COMPARATIVE = """This is a comparative question. It needs to be decomposed into: +1. Information about entity A for the comparison dimension. +2. Information about entity B for the comparison dimension. +3. Optional information that directly compares the two entities. +Base the decomposition on BOTH the question and the preliminary retrieved context below. + +Question: {question} +Maximum number of claims: {max_claims} +Detail level: {detail_level} + +Preliminary retrieved context for this question (use it to ground each claim in what the corpus actually contains; do not invent facts it cannot support): +{retrieved} + +Output format (JSON): +{{ + "claims": [ + {{ + "claim_id": "c1", + "description": "The distance from Hangzhou to Beijing", + "priority": 1 + }}, + {{ + "claim_id": "c2", + "description": "The distance from Shanghai to Beijing", + "priority": 1 + }}, + {{ + "claim_id": "c3", + "description": "Which city is closer to Beijing", + "priority": 2 + }} + ] +}} +""" + + +DECOMPOSE_PROCEDURAL = """This is a procedural question. Decompose it into the information needed for each step required to complete the operation. +Base the decomposition on BOTH the question and the preliminary retrieved context below. + +Question: {question} +Maximum number of claims: {max_claims} +Detail level: {detail_level} + +Preliminary retrieved context for this question (use it to ground each claim in what the corpus actually contains; do not invent facts it cannot support): +{retrieved} + +Output format (JSON): +{{ + "claims": [ + {{ + "claim_id": "c1", + "description": "Information needed for the first step", + "priority": 1 + }}, + {{ + "claim_id": "c2", + "description": "Information needed for the second step", + "priority": 2 + }} + ] +}} +""" + + +DECOMPOSE_EXPLORATORY = """This is an analytical or exploratory question. Decompose it into the main aspects or dimensions that need to be researched. +Base the decomposition on BOTH the question and the preliminary retrieved context below. + +Question: {question} +Maximum number of claims: {max_claims} +Detail level: {detail_level} + +Preliminary retrieved context for this question (use it to ground each aspect in what the corpus actually contains; do not invent aspects it cannot support): +{retrieved} + +Output format (JSON): +{{ + "claims": [ + {{ + "claim_id": "c1", + "description": "The first aspect that needs to be researched", + "priority": 1 + }}, + {{ + "claim_id": "c2", + "description": "The second aspect that needs to be researched", + "priority": 2 + }} + ] +}} +""" diff --git a/rag/advanced_rag/harness/prompts/report_prompt.py b/rag/advanced_rag/harness/prompts/report_prompt.py new file mode 100644 index 0000000000..f6de3d191c --- /dev/null +++ b/rag/advanced_rag/harness/prompts/report_prompt.py @@ -0,0 +1,16 @@ +"""Report synthesis prompts.""" + +FINAL_ANSWER_SYSTEM = """You are a smart agent. Answer the user's question using ONLY the evidence provided below. Do not invent facts: if the evidence cannot support a claim, say so plainly instead of guessing. + +# Citation rules +{cite_rules} + +# Language +Answer in the SAME language as the question. Translate retrieved evidence into that language as part of composing the answer; only verbatim quoted snippets may stay in their source language. + +# Fallback +If the evidence does not answer the question, reply with a clear statement that you don't have enough information based on the available sources (in the user's language). +""" + + +PARTIAL_ANSWER_PREAMBLE = "Note: the following answer is based on partial information and may be incomplete." diff --git a/rag/advanced_rag/harness/prompts/research_agent_prompt.py b/rag/advanced_rag/harness/prompts/research_agent_prompt.py new file mode 100644 index 0000000000..92d3cb84d4 --- /dev/null +++ b/rag/advanced_rag/harness/prompts/research_agent_prompt.py @@ -0,0 +1,58 @@ +"""Research Agent prompts. + +``RESEARCH_AGENT_PROMPT`` — native tool-calling: the tool schemas are bound + onto the chat model via ``bind_tools``, so the + prompt only describes the task, not the tools. +``RESEARCH_AGENT_TEXT_PROMPT`` — fallback for models without native tool-calling: + the tools are described in-prompt and the model + emits ```` JSON that the loop parses. +""" + +RESEARCH_AGENT_PROMPT = """You are a research assistant. Investigate the given research task by calling the provided tools. + +Research task: {claim_description} + +Current phase: {phase} +Phase hint: {phase_hint} + +Rules: +1. Prefer search / navigation tools to gather evidence for the task. +2. Use think_tool to analyze results and plan the next step after each search. +3. When you have gathered enough evidence, call generate_report with your findings + (report, is_verified, confidence, evidence_ids, gaps, discovered_claims). + +You have at most {max_cycles} tool-calling rounds. Call exactly one tool per round, +and do not write a plain-text answer until you call generate_report. +""" + + +RESEARCH_AGENT_TEXT_PROMPT = """You are a research assistant. For the given research task, use the available tools to search for information. + +Research task: {claim_description} + +Current phase: {phase} +Phase hint: {phase_hint} + +Available tools: +{tool_list} + +Rules: +1. Prefer search tools to gather information. +2. Use think_tool to analyze results after each search. +3. When you are confident enough to answer the research task, call generate_report. + +Tool call format: output exactly one JSON tool call per round: +{{"name": "tool_name", "arguments": {{"parameter_name": "value"}} }} + +generate_report argument format: +{{ + "report": "Research result report, factual and unformatted", + "is_verified": true/false, + "confidence": 0.0-1.0, + "evidence_ids": [0, 3], + "gaps": ["Information that was not found"], + "discovered_claims": ["New research directions discovered during research"] +}} + +Maximum {max_cycles} rounds. Output one tag in each round and no other text. +""" diff --git a/rag/advanced_rag/harness/prompts/route_prompt.py b/rag/advanced_rag/harness/prompts/route_prompt.py new file mode 100644 index 0000000000..ea3af15b4e --- /dev/null +++ b/rag/advanced_rag/harness/prompts/route_prompt.py @@ -0,0 +1,19 @@ +"""Route node prompt: classify query type.""" + +ROUTE_PROMPT = """Analyze the following question and output a structured query analysis. + +Question: {question} + +Analyze it across these dimensions: +1. Question type: factual / comparative / analytical / procedural / exploratory / verification / summarization. +2. Whether it needs decomposition into atomic facts, meaning whether multiple independent pieces of information must be retrieved separately before answering: true/false. +3. Suggested knowledge compilation tool: null (none) / toc (document table of contents) / graph (knowledge graph) / wiki (compiled domain knowledge). + +Output format (JSON): +{{ + "question_type": "comparative", + "requires_decomposition": true, + "suggests_compilation": null, + "reasoning": "This is a comparative question, so it needs to be decomposed into two independent facts and one comparison relation." +}} +""" diff --git a/rag/advanced_rag/harness/prompts/sufficiency_prompt.py b/rag/advanced_rag/harness/prompts/sufficiency_prompt.py new file mode 100644 index 0000000000..a941e62c47 --- /dev/null +++ b/rag/advanced_rag/harness/prompts/sufficiency_prompt.py @@ -0,0 +1,31 @@ +"""Sufficiency judge prompt: verdict with claim-level assessment.""" + +SUFFICIENCY_JUDGE_PROMPT = """You are an expert judge of information retrieval sufficiency. Decide whether the currently collected evidence is sufficient to answer the question. + +Question: {question} + +Claim-level evidence: +{evidence_summary} + +Judgment tasks: +1. Evaluate each claim one by one and decide whether it has been sufficiently verified. +2. Make an overall judgment about whether the evidence is sufficient to answer the user's question. +3. If it is not sufficient, provide targeted feedback. + +Output format (JSON): +{{ + "status": "SUFFICIENT" | "USEFUL_BUT_INCOMPLETE" | "INSUFFICIENT" | "UNANSWERABLE", + "score": 0.85, + "claim_assessments": [ + {{ + "claim_id": "c1", + "is_verified": true, + "confidence": 0.95, + "reason": "Consistent data was found in three chunks." + }} + ], + "missing": ["Some data for c2 was not found."], + "feedback": "Use web_search for c2 to supplement the latest data.", + "overall_reason": "The main facts are covered, but some details still need supplementation." +}} +""" diff --git a/rag/advanced_rag/harness/route.py b/rag/advanced_rag/harness/route.py new file mode 100644 index 0000000000..d3d4cbb7bf --- /dev/null +++ b/rag/advanced_rag/harness/route.py @@ -0,0 +1,77 @@ +"""Route node — query classification (one-time, no KB dependency).""" + +import json +import logging +import re + +from rag.advanced_rag.harness.types import RouteDecision +from rag.advanced_rag.harness.config import get_mode +from rag.advanced_rag.harness.prompts.route_prompt import ROUTE_PROMPT + +_LOG = logging.getLogger(__name__) + + +def _extract_json(text: str) -> dict: + """Extract JSON from LLM response, handling markdown fences and think tags.""" + text = re.sub(r"^.*", "", text, flags=re.DOTALL).strip() + text = re.sub(r"```(?:json)?\s*|\s*```", "", text).strip() + try: + import json_repair + + return json_repair.loads(text) + except Exception: + try: + return json.loads(text) + except Exception: + _LOG.warning("route: failed to parse LLM output: %s", text[:200]) + return {} + + +async def route_node(state: dict, tools) -> dict: + """Route node — analyze the question, produce RouteDecision.""" + question = state.get("question", "") + if not question: + return _fallback_route(question) + + mode_label = getattr(tools, "thinking_mode", "medium") + mode = get_mode(mode_label) + + try: + system = ROUTE_PROMPT.format(question=question) + msg = await tools._fit_messages(system, question) + ans = await tools.chat_mdl.async_chat(msg[0]["content"], msg[1:], {"temperature": 0.1}) + if isinstance(ans, tuple): + ans = ans[0] + result = _extract_json(ans) + except Exception: + _LOG.exception("route_node failed") + result = {} + + question_type = result.get("question_type", "factual") + requires_decomp = result.get("requires_decomposition", True) + suggests_comp = result.get("suggests_compilation") + + route = RouteDecision( + question=question, + thinking_mode=mode_label, + question_type=question_type, + requires_decomposition=mode.requires_decomposition and requires_decomp, + suggests_compilation=suggests_comp, + execution_strategy=mode.execution_strategy, + reasoning=result.get("reasoning", ""), + ) + + return {"route": route} + + +def _fallback_route(question: str) -> dict: + route = RouteDecision( + question=question, + thinking_mode="medium", + question_type="factual", + requires_decomposition=False, + suggests_compilation=None, + execution_strategy="direct_search", + reasoning="fallback: empty question", + ) + return {"route": route} diff --git a/rag/advanced_rag/harness/sufficiency.py b/rag/advanced_rag/harness/sufficiency.py new file mode 100644 index 0000000000..9dab9d2947 --- /dev/null +++ b/rag/advanced_rag/harness/sufficiency.py @@ -0,0 +1,188 @@ +"""Sufficiency check — cross-check + fusion score + 5-way verdict.""" + +import logging + +from rag.advanced_rag.harness.types import ( + AgentResult, + ClaimCrossCheckResult, + SufficiencyVerdict, + ExecutionStrategy, +) +from rag.advanced_rag.harness.config import get_mode + +_LOG = logging.getLogger(__name__) + + +# ═══════════════════════════════════════════════════════════════ +# Cross-check: code-only +# ═══════════════════════════════════════════════════════════════ + +import re + + +def extract_numbers(text: str) -> list[float]: + """Extract numeric values from text.""" + return [float(m) for m in re.findall(r"\d+\.?\d*", text)] + + +def extract_named_entities(text: str) -> list[str]: + """Simple entity extraction — looks for capitalized multi-word sequences.""" + entities = re.findall(r"\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*\b", text) + return list(set(entities)) + + +def cross_check_claim(agent_result: AgentResult, all_chunks: dict) -> ClaimCrossCheckResult: + """Code-level cross-check: number matching + entity presence.""" + report = agent_result.report + claimed = agent_result.is_verified + + if not claimed: + return ClaimCrossCheckResult( + claim_id=agent_result.claim_id, + cross_check_passed=False, + cross_check_score=0.0, + mismatches=["agent self-reported as unverified"], + ) + + numbers = extract_numbers(report) + entities = extract_named_entities(report) + + mismatches = [] + matches = [] + + for eid in agent_result.evidence_ids or []: + chunk = all_chunks.get(eid) + if not chunk: + mismatches.append(f"evidence_id={eid}: chunk not found") + continue + text = chunk.get("content_with_weight", chunk.get("text", "")) + text_lower = text.lower() + + for num in numbers: + if str(num) not in text_lower: + mismatches.append(f"number {num} not found in chunk {eid}") + else: + matches.append(f"number {num} found in chunk {eid}") + + for ent in entities: + if ent.lower() not in text_lower: + mismatches.append(f"entity '{ent}' not found in chunk {eid}") + + total = len(matches) + len(mismatches) + cross_score = len(matches) / max(total, 1) if total > 0 else 0.0 + cross_passed = len(mismatches) < len(matches) * 0.5 + + return ClaimCrossCheckResult( + claim_id=agent_result.claim_id, + cross_check_passed=cross_passed, + cross_check_score=cross_score, + evidence_matches=matches, + mismatches=mismatches, + ) + + +# ═══════════════════════════════════════════════════════════════ +# Fusion score +# ═══════════════════════════════════════════════════════════════ + + +def compute_fusion_score( + agent_results: list[AgentResult], + cross_check_results: list[ClaimCrossCheckResult], + mode: ExecutionStrategy, +) -> SufficiencyVerdict: + """Dual-signal fusion: agent confidence + cross-check pass rate.""" + # Signal A: agent self-assessment + verified_count = sum(1 for r in agent_results if r.is_verified) + agent_score = verified_count / max(len(agent_results), 1) + + # Signal B: cross-check + passed_count = sum(1 for r in cross_check_results if r.cross_check_passed) + cross_score = passed_count / max(len(cross_check_results), 1) + + # Fusion strategy by mode + fusion = { + "ultra": lambda a, c: min(a, c), + "high": lambda a, c: (a + c) / 2, + "medium": lambda a, c: max(a, c), + "low": lambda a, c: max(a, c), + }.get(mode.label, lambda a, c: max(a, c)) + fusion_score = fusion(agent_score, cross_score) + + # Conflict detection + has_conflicts = any(len(r.mismatches) > 0 for r in cross_check_results) + + # 5-way verdict + if has_conflicts and fusion_score < mode.partial_threshold: + status = "CONFLICTING" + elif fusion_score >= mode.sufficiency_threshold: + status = "SUFFICIENT" + elif fusion_score >= mode.partial_threshold: + status = "USEFUL_BUT_INCOMPLETE" + elif not any(r.cross_check_passed for r in cross_check_results): + status = "UNANSWERABLE" + else: + status = "INSUFFICIENT" + + missing = [r.claim_id for r in cross_check_results if not r.cross_check_passed] + + return SufficiencyVerdict( + status=status, + score=fusion_score, + agent_score=agent_score, + cross_score=cross_score, + claim_assessments=[{"claim_id": r.claim_id, "is_verified": r.cross_check_passed, "score": r.cross_check_score, "mismatches": r.mismatches} for r in cross_check_results], + has_conflicts=has_conflicts, + missing_claims=missing, + feedback=_build_feedback(missing, cross_check_results), + overall_reason=_format_reason(status, fusion_score, missing), + ) + + +# ═══════════════════════════════════════════════════════════════ +# Helpers +# ═══════════════════════════════════════════════════════════════ + + +def _build_feedback(missing: list[str], results: list[ClaimCrossCheckResult]) -> str: + if not missing: + return "all claims verified" + hints = [] + for r in results: + if not r.cross_check_passed: + hints.append(f"claim {r.claim_id}: {len(r.mismatches)} mismatch(es)") + return "missing: " + "; ".join(hints) + + +def _format_reason(status: str, score: float, missing: list[str]) -> str: + return f"{status} score={score:.2f} missing={missing}" + + +def route_sufficiency_verdict(verdict: SufficiencyVerdict, mode_label: str, cycle: int, max_cycles: int) -> tuple: + """Return (action, should_continue).""" + mode = get_mode(mode_label) + + if verdict.status == "SUFFICIENT": + return ("ANSWER", False) + + if verdict.status == "USEFUL_BUT_INCOMPLETE": + if mode.requires_selective_gen: + return ("ANSWER_PARTIAL", False) + return ("CONTINUE", False) + + if verdict.status == "INSUFFICIENT": + if cycle >= max_cycles * 0.8: + return ("ANSWER_PARTIAL", False) + return ("CONTINUE", True) + + if verdict.status == "CONFLICTING": + if mode.allows_replan and cycle < max_cycles * 0.5: + return ("REPLAN", True) + return ("ANSWER_PARTIAL", False) + + if verdict.status == "UNANSWERABLE": + if mode.fallback_to_direct_llm: + return ("FALLBACK_LLM", False) + return ("ABSTAIN", False) + + return ("CONTINUE", True) diff --git a/rag/advanced_rag/harness/tools/__init__.py b/rag/advanced_rag/harness/tools/__init__.py new file mode 100644 index 0000000000..5be55f746c --- /dev/null +++ b/rag/advanced_rag/harness/tools/__init__.py @@ -0,0 +1,46 @@ +"""Tool system: register all tools with the registry on import.""" + +from rag.advanced_rag.harness.tools.registry import register_tool, _search_schema, _navigate_schema, _inspector_schema + +# Register tools + +# Search tools +from rag.advanced_rag.harness.tools.search import hybrid_search, vector_search, bm25_search, web_search, structured_query + +register_tool("hybrid_search", _search_schema("hybrid_search", "Embedding + Keywords search"), hybrid_search) +register_tool("vector_search", _search_schema("vector_search", "Embedding search"), vector_search) +register_tool("bm25_search", _search_schema("bm25_search", "Keywords search"), bm25_search) +register_tool("web_search", _search_schema("web_search", "Internet search"), web_search) +register_tool("structured_query", _search_schema("structured_query", "SQL search"), structured_query) + +# Navigation tools (require compilation) +from rag.advanced_rag.harness.tools.navigation import toc_navigate, page_index_navigate, mindmap_navigate + +register_tool("toc_navigate", _navigate_schema("toc_navigate", "Navigation with table of content"), toc_navigate, requires_compilation=True, compilation_type="toc") +register_tool("page_index_navigate", _navigate_schema("page_index_navigate", "Navigation with extracted titles"), page_index_navigate, requires_compilation=True, compilation_type="page_index") +register_tool("mindmap_navigate", _navigate_schema("mindmap_navigate", "Navigate by mindmap"), mindmap_navigate, requires_compilation=True, compilation_type="mindmap") + +# Exploration tools (require compilation) +from rag.advanced_rag.harness.tools.exploration import graph_explore, wiki_query + +register_tool("graph_explore", _search_schema("graph_explore", "Knowledge graph exploration"), graph_explore, requires_compilation=True, compilation_type="knowledge_graph") +register_tool("wiki_query", _search_schema("wiki_query", "Wiki search"), wiki_query, requires_compilation=True, compilation_type="wiki") + +# Inspector tools +from rag.advanced_rag.harness.tools.inspector import open_context, compare_sources, grep_within, request_adjacent + +register_tool( + "inspector_open_context", + _inspector_schema("open_context", "Expand the original context around a chunk", {"chunk_id": {"type": "string"}, "width": {"type": "integer", "description": "Number of characters to expand"}}), + open_context, +) +register_tool("inspector_compare", _inspector_schema("compare_sources", "Compare multiple chunks to find common points and contradictions", {"chunk_ids": {"type": "array", "items": {"type": "string"}}}), compare_sources) +register_tool("inspector_grep_within", _inspector_schema("grep_within", "Search for an exact keyword or pattern within a document", {"doc_id": {"type": "string"}, "pattern": {"type": "string"}}), grep_within) +register_tool( + "inspector_request_adjacent", + _inspector_schema("request_adjacent", "Get adjacent entries before or after a chunk", {"chunk_id": {"type": "string"}, "direction": {"type": "string", "enum": ["prev", "next"]}, "count": {"type": "integer"}}), + request_adjacent, +) + +# Built-in agent tools +# (generate_report and think_tool are handled by the agent loop itself, not by Pipeline) diff --git a/rag/advanced_rag/harness/tools/exploration.py b/rag/advanced_rag/harness/tools/exploration.py new file mode 100644 index 0000000000..9969b30a6a --- /dev/null +++ b/rag/advanced_rag/harness/tools/exploration.py @@ -0,0 +1,32 @@ +"""Exploration tools: knowledge graph and wiki lookup.""" + +import logging + +_LOG = logging.getLogger(__name__) + + +async def graph_explore(tools, entity: str, relation: str | None = None, depth: int = 1) -> dict: + """Explore relationships in the knowledge graph. + + This is currently a placeholder. The final implementation should call the + knowledge graph store. + """ + _LOG.info("graph_explore: entity=%s relation=%s depth=%d", entity, relation, depth) + # TODO: implement actual KG walk + from rag.advanced_rag.harness.tools.search import hybrid_search + + query = f"{entity} {relation or ''}" + return await hybrid_search(tools, query=query.strip()) + + +async def wiki_query(tools, topic: str) -> dict: + """Query compiled wiki knowledge. + + This is currently a placeholder. The final implementation should call the + compiled wiki store. + """ + _LOG.info("wiki_query: topic=%s", topic) + # TODO: implement actual wiki lookup + from rag.advanced_rag.harness.tools.search import hybrid_search + + return await hybrid_search(tools, query=topic) diff --git a/rag/advanced_rag/harness/tools/gating.py b/rag/advanced_rag/harness/tools/gating.py new file mode 100644 index 0000000000..22ee3292a7 --- /dev/null +++ b/rag/advanced_rag/harness/tools/gating.py @@ -0,0 +1,125 @@ +"""Tool selection gating: phase-based filtering and fallback chain.""" + +from rag.advanced_rag.harness.types import OrchestratorContext +from rag.advanced_rag.harness.tools.registry import TOOL_REGISTRY + + +# Search phase definitions + +SEARCH_PHASES = { + "locate": { + "goal": "Locate documents or regions that may contain the answer.", + "tools_priority": [ + "toc_navigate", + "mindmap_navigate", + "page_index_navigate", + "wiki_query", + "hybrid_search", + "bm25_search", + ], + "max_returned": 5, + "tool_hint": "Prefer navigation tools to locate document regions before directly searching keywords.", + }, + "explore": { + "goal": "Explore deeply within the already located region.", + "tools_priority": [ + "hybrid_search", + "vector_search", + "bm25_search", + "graph_explore", + "inspector_open_context", + "inspector_request_adjacent", + ], + "max_returned": 4, + "tool_hint": "Prefer retrieval tools to gather detailed information within the located region.", + }, + "verify": { + "goal": "Verify consistency across multiple sources.", + "tools_priority": [ + "inspector_open_context", + "inspector_compare", + "inspector_grep_within", + "hybrid_search", + "web_search", + ], + "max_returned": 4, + "tool_hint": "Prefer inspector tools to compare existing evidence before searching for new content.", + }, + "cross_domain": { + "goal": "Explore cross-domain relationships for discovered entities.", + "tools_priority": [ + "graph_explore", + "wiki_query", + "hybrid_search", + "web_search", + ], + "max_returned": 3, + "tool_hint": "Prefer walking the graph to discover cross-domain relationships.", + }, +} + + +def compilation_available(tool_name: str, compilation_map: dict) -> bool: + """Check if any KB provides the required compilation artifact.""" + tool = TOOL_REGISTRY.get(tool_name) + if not tool or not tool.get("requires_compilation"): + return True + comp_type = tool["compilation_type"] + if not compilation_map: + return False + return any(comp_type in comps for comps in compilation_map.values()) + + +def tool_fits_context(tool_name: str, context: OrchestratorContext) -> bool: + """Check if a tool is sensible given current search context.""" + if tool_name.startswith("inspector_") and not context.has_any_chunks(): + return False + if tool_name == "toc_navigate" and not context.current_claim: + return False + if tool_name == "graph_explore" and not context.last_entity: + return False + if tool_name == "mindmap_navigate" and not context.current_claim: + return False + return True + + +def get_gated_tools( + phase: str, + available_tools: list[str], + compilation_map: dict[str, set[str]], + context: OrchestratorContext, +) -> list[dict]: + """Filter, sort, and gate tools by phase priority and context.""" + phase_config = SEARCH_PHASES.get(phase) + if not phase_config: + return _default_defs(available_tools) + + sorted_tools = [] + for tool_name in phase_config["tools_priority"]: + if tool_name not in available_tools: + continue + if not compilation_available(tool_name, compilation_map): + continue + if not tool_fits_context(tool_name, context): + continue + sorted_tools.append(tool_name) + + selected = sorted_tools[: phase_config["max_returned"]] + defs = [TOOL_REGISTRY[n]["function_schema"] for n in selected if n in TOOL_REGISTRY] + for d in defs: + d["x_phase"] = phase + d["x_phase_hint"] = phase_config["tool_hint"] + return defs + + +def _default_defs(tool_names: list[str]) -> list[dict]: + return [TOOL_REGISTRY[n]["function_schema"] for n in tool_names if n in TOOL_REGISTRY] + + +def determine_current_phase(context: OrchestratorContext) -> str: + """Determine the current search phase based on context.""" + if not context.has_any_chunks(): + return "locate" + if context.verdict and context.verdict.has_conflicts: + return "verify" + return "explore" diff --git a/rag/advanced_rag/harness/tools/inspector.py b/rag/advanced_rag/harness/tools/inspector.py new file mode 100644 index 0000000000..cc7fcaf067 --- /dev/null +++ b/rag/advanced_rag/harness/tools/inspector.py @@ -0,0 +1,87 @@ +"""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 diff --git a/rag/advanced_rag/harness/tools/navigation.py b/rag/advanced_rag/harness/tools/navigation.py new file mode 100644 index 0000000000..bb2ce25a9b --- /dev/null +++ b/rag/advanced_rag/harness/tools/navigation.py @@ -0,0 +1,36 @@ +"""Navigation tools: TOC, page index, and mindmap.""" + +import logging + +_LOG = logging.getLogger(__name__) + + +async def toc_navigate(tools, topic: str, doc_scope: list[str] | None = None) -> dict: + """Locate by document table of contents. + + Requires a KB with a compiled TOC artifact. This is currently a placeholder. + Replace it with the real implementation when the knowledge compilation TOC + tool is available. Falls back to hybrid search for now. + """ + _LOG.info("toc_navigate called for topic=%s", topic) + # TODO: replace with actual TOC navigation when compilation layer is ready + # For now, fallback to hybrid search + from rag.advanced_rag.harness.tools.search import hybrid_search + + return await hybrid_search(tools, query=topic, doc_scope=doc_scope) + + +async def page_index_navigate(tools, topic: str, kb_ids: list[str] | None = None) -> dict: + """Navigate by page index.""" + _LOG.info("page_index_navigate called for topic=%s", topic) + from rag.advanced_rag.harness.tools.search import hybrid_search + + return await hybrid_search(tools, query=topic, kb_ids=kb_ids) + + +async def mindmap_navigate(tools, concept: str, kb_ids: list[str] | None = None) -> dict: + """Locate by concept mindmap.""" + _LOG.info("mindmap_navigate called for concept=%s", concept) + from rag.advanced_rag.harness.tools.search import hybrid_search + + return await hybrid_search(tools, query=concept, kb_ids=kb_ids) diff --git a/rag/advanced_rag/harness/tools/registry.py b/rag/advanced_rag/harness/tools/registry.py new file mode 100644 index 0000000000..072e0b2ba0 --- /dev/null +++ b/rag/advanced_rag/harness/tools/registry.py @@ -0,0 +1,141 @@ +"""Tool registry: all available tools with metadata and function schemas.""" + +from typing import Any + +# Tool registry: tool_name -> {metadata, function_schema, fn} +# 'fn' filled at registration time; schema used for LLM tool definitions. +TOOL_REGISTRY: dict[str, dict[str, Any]] = {} + +# Executor interface +# Each tool registers a callable with signature: +# async def fn(tools, **kwargs) -> dict # {"chunks": [...], ...} + + +def register_tool(name: str, schema: dict, fn: callable, requires_compilation: bool = False, compilation_type: str | None = None, processing_time: str = "fast") -> None: + TOOL_REGISTRY[name] = { + "name": name, + "function_schema": schema, + "fn": fn, + "requires_compilation": requires_compilation, + "compilation_type": compilation_type, + "processing_time": processing_time, + } + + +def get_tool(tool_name: str) -> dict | None: + return TOOL_REGISTRY.get(tool_name) + + +def get_function_schemas(tool_names: list[str]) -> list[dict]: + """Return function schemas for the given tool names, if registered.""" + return [TOOL_REGISTRY[n]["function_schema"] for n in tool_names if n in TOOL_REGISTRY] + + +# Common schema builders + +def _search_schema(name: str, desc: str) -> dict: + return { + "type": "function", + "function": { + "name": name, + "description": desc, + "parameters": { + "type": "object", + "properties": { + "query": {"type": "string", "description": "the original user's question."}, + "keywords": {"type": "string", "description": "the keywords used for searching split by space or ','."}, + }, + "required": ["query"], + }, + }, + } + + +def _navigate_schema(name: str, desc: str) -> dict: + return { + "type": "function", + "function": { + "name": name, + "description": desc, + "parameters": { + "type": "object", + "properties": { + "topic": {"type": "string", "description": "navigate by topic"}, + }, + "required": ["topic"], + }, + }, + } + + +def _inspector_schema(name: str, desc: str, props: dict = None) -> dict: + schema = { + "type": "function", + "function": { + "name": name, + "description": desc, + "parameters": { + "type": "object", + "properties": props + or { + "chunk_id": {"type": "string", "description": "chunk ID"}, + }, + "required": list((props or {"chunk_id": {}}).keys()), + }, + }, + } + return schema + + +def _think_schema() -> dict: + return { + "type": "function", + "function": { + "name": "think_tool", + "description": "Internal reasoning. Analyze the collected results and plan the next step. Do not output final user-facing content while reasoning.", + "parameters": { + "type": "object", + "properties": { + "reasoning": { + "type": "string", + "description": "Reasoning content: what has been found, what is still missing, and what to do next.", + }, + }, + "required": ["reasoning"], + }, + }, + } + + +def _generate_report_schema() -> dict: + return { + "type": "function", + "function": { + "name": "generate_report", + "description": "Call when the research is complete. Output the research report and claim-level verification results.", + "parameters": { + "type": "object", + "properties": { + "report": {"type": "string", "description": "Research result report, factual and unformatted."}, + "is_verified": {"type": "boolean", "description": "Whether sufficient evidence was found."}, + "confidence": {"type": "number", "description": "Confidence from 0 to 1."}, + "evidence_ids": { + "type": "array", + "items": {"type": "integer"}, + "description": "Referenced chunk IDs.", + }, + "gaps": { + "type": "array", + "items": {"type": "string"}, + "description": "Information that was not found.", + }, + "discovered_claims": { + "type": "array", + "items": {"type": "string"}, + "description": "New research directions discovered during research.", + }, + }, + "required": ["report", "is_verified", "confidence"], + }, + }, + } diff --git a/rag/advanced_rag/harness/tools/search.py b/rag/advanced_rag/harness/tools/search.py new file mode 100644 index 0000000000..4014cf190c --- /dev/null +++ b/rag/advanced_rag/harness/tools/search.py @@ -0,0 +1,289 @@ +"""Search tools: hybrid, vector, BM25, web, structured.""" + +import logging +import re +import hashlib +from common import settings + +_LOG = logging.getLogger(__name__) + + +# Sentence terminators: Chinese 。！？；, English ! ? ;, newline, and a +# digit-guarded English period (so "3.14" / "v1.2" don't split). +_SENT_END = re.compile(r"[。！？；!?;]+|(?]*>.*?", re.IGNORECASE | re.DOTALL) +# Markdown table: a header row with a pipe, a separator row of dashes/colons/ +# pipes, then zero+ body rows with a pipe. +_MD_TABLE = re.compile( + r"^[ \t]*\|?[^\n]*\|[^\n]*\r?\n" + r"[ \t]*\|?[ \t]*:?-{1,}:?[ \t]*(?:\|[ \t]*:?-{1,}:?[ \t]*)+\|?[ \t]*\r?\n" + r"(?:[ \t]*\|?[^\n]*\|[^\n]*\r?\n?)*", + re.MULTILINE, +) + + +def _protected_spans(text: str) -> list[tuple[int, int]]: + """Non-overlapping ``(start, end)`` spans of table blocks, in order.""" + spans = [(m.start(), m.end()) for m in _HTML_TABLE.finditer(text)] + spans += [(m.start(), m.end()) for m in _MD_TABLE.finditer(text)] + spans.sort() + merged: list[tuple[int, int]] = [] + last_end = -1 + for s, e in spans: + if s < last_end: # overlaps an already-kept span -> skip + continue + merged.append((s, e)) + last_end = e + return merged + + +def _split_plain(text: str) -> list[str]: + """Terminator-based sentence split, keeping each terminator attached.""" + sents: list[str] = [] + start = 0 + for m in _SENT_END.finditer(text): + end = m.end() + seg = text[start:end] + if seg.strip(): + sents.append(seg) + start = end + if start < len(text): + tail = text[start:] + if tail.strip(): + sents.append(tail) + return sents + + +def _split_sentences(text: str) -> list[str]: + """Split ``text`` into sentences, keeping each terminator attached. + + Table blocks — HTML ``...
`` and markdown tables — are treated + as a single atomic sentence and are never split internally. + """ + if not text: + return [] + spans = _protected_spans(text) + if not spans: + return _split_plain(text) + + sents: list[str] = [] + pos = 0 + for s, e in spans: + if s > pos: + sents.extend(_split_plain(text[pos:s])) + block = text[s:e] + if block.strip(): + sents.append(block) + pos = e + if pos < len(text): + sents.extend(_split_plain(text[pos:])) + return sents + + +def _narrow_content(content: str, kwds: list[str]) -> str | None: + """Return ``content`` narrowed to keyword sentences +/- 1 neighbour. + + Returns ``None`` when no keyword occurs anywhere in ``content``. + """ + sents = _split_sentences(content) + if not sents: + return None + keep: set[int] = set() + matched = False + for i, s in enumerate(sents): + low = s.lower() + if any(kw in low for kw in kwds): + matched = True + if i > 0: + keep.add(i - 1) + keep.add(i) + if i + 1 < len(sents): + keep.add(i + 1) + if not matched: + return None + narrowed = "".join(sents[i] for i in sorted(keep)).strip() + return "..." + _highlight_keywords(narrowed, kwds) + "..." + + +def _highlight_keywords(text: str, kwds: list[str]) -> str: + terms = sorted({kw for kw in kwds if kw}, key=len, reverse=True) + if not terms: + return text + pattern = re.compile("|".join(re.escape(term) for term in terms), re.IGNORECASE) + return pattern.sub(lambda m: f"{m.group(0)}", text) + + +def _narrow_by_keywords(chunks: list[dict], keywords: str) -> list[dict]: + """Narrow each chunk to its keyword-bearing sentences (+/- 1 neighbour) and + drop keyword-less chunks. + + Keywords are the comma-separated terms (with close synonyms) produced by + ``formalize``; matching is case-insensitive substring. + """ + kwds = [k.strip().lower() for k in (keywords or "").split(",") if k.strip()] + if not kwds or not chunks: + return chunks + if len(kwds) < 3: + kwds = [k.strip().lower() for k in (keywords or "").split(" ") if k.strip()] + _kwds = [] + for i in range(len(kwds)-1): + _kwds.append(kwds[i] + " "+ kwds[i+1]) + kwds = _kwds + + scored = [(ck, _narrow_content(ck.get("content_with_weight") or ck.get("content") or "", kwds)) for ck in chunks] + out: list[dict] = [] + dedup: set[str] = set() + for ck, nc in scored: + if nc is not None: + nc_hash = hashlib.md5(nc.encode("utf-8")).hexdigest() + if nc_hash in dedup: + continue + dedup.add(nc_hash) + ck["content_with_weight"] = nc + if "content" in ck: + ck["content"] = nc + ck.pop("highlight", None) + out.append(ck) + return out + + +def _normalize(kbinfos: dict, tenant_ids: list[str] | str | None) -> dict: + if not kbinfos: + return {"chunks": [], "doc_aggs": []} + if not tenant_ids: + _LOG.warning("search: skip child retrieval because tenant_ids is empty") + return kbinfos + if isinstance(tenant_ids, str): + tenant_ids = [tenant_ids] + kbinfos["chunks"] = settings.retriever.retrieval_by_children( + kbinfos.get("chunks", []), + tenant_ids, + ) + return kbinfos + + +async def hybrid_search(tools, query: str, kb_ids: list[str] | None = None, top_n: int = 12, doc_scope: list[str] | None = None, keywords: str = "") -> dict: + if not tools.kb_ids and not kb_ids: + return {"chunks": [], "doc_aggs": []} + target_ids = kb_ids or tools.kb_ids + _LOG.info(f"[Hybrid search]: {query} -> {keywords}") + + # Query expansion: append the formalized-question keywords + close synonyms + # so hybrid/BM25 retrieval gets extra recall signal. + effective_query = f"{query} {keywords}".strip() if keywords else query + + embd_mdl = tools.embed_mdl + vector_weight = 0.3 if embd_mdl else 0 + + kbinfos = await settings.retriever.retrieval( + effective_query, + embd_mdl, + tools.tenant_ids, + target_ids, + 1, + top_n, + 0.2, + vector_similarity_weight=vector_weight, + aggs=True, + highlight=False, + doc_ids=doc_scope, + ) + kbinfos = _normalize(kbinfos, tools.tenant_ids) + if keywords: + length = len(kbinfos["chunks"]) + kbinfos["chunks"] = _narrow_by_keywords(kbinfos.get("chunks", []), keywords) + _LOG.info(f"[Hybrid search]({keywords}): snippet {length} -> {len(kbinfos['chunks'])}") + return kbinfos + + +async def vector_search(tools, query: str, kb_ids: list[str] | None = None, top_n: int = 12, keywords: str = "") -> dict: + if not tools.embed_mdl: + _LOG.warning("vector_search: no embed_mdl available") + return {"chunks": [], "doc_aggs": []} + + _LOG.info(f"[Vector search]: {query} -> {keywords}") + effective_query = f"{query} {keywords}".strip() if keywords else query + target_ids = kb_ids or tools.kb_ids + kbinfos = await settings.retriever.retrieval( + effective_query, + tools.embed_mdl, + tools.tenant_ids, + target_ids, + 1, + top_n, + 0.2, + vector_similarity_weight=1.0, + aggs=False, + highlight=False, + ) + if keywords: + length = len(kbinfos["chunks"]) + kbinfos["chunks"] = _narrow_by_keywords(kbinfos.get("chunks", []), keywords) + _LOG.info(f"[Vector search]({keywords}): snippet {length} -> {len(kbinfos['chunks'])}") + return _normalize(kbinfos, tools.tenant_ids) + + +async def bm25_search(tools, query: str, kb_ids: list[str] | None = None, top_n: int = 12, keywords: str = "") -> dict: + _LOG.info(f"[BM25 search]: {query} -> {keywords}") + target_ids = kb_ids or tools.kb_ids + effective_query = f"{query} {keywords}".strip() if keywords else query + kbinfos = await settings.retriever.retrieval( + effective_query, + None, + tools.tenant_ids, + target_ids, + 1, + top_n, + 0.0, + vector_similarity_weight=0, + aggs=False, + highlight=False, + ) + if keywords: + length = len(kbinfos["chunks"]) + kbinfos["chunks"] = _narrow_by_keywords(kbinfos.get("chunks", []), keywords) + _LOG.info(f"[BM25 search]({keywords}): snippet {length} -> {len(kbinfos['chunks'])}") + return _normalize(kbinfos, tools.tenant_ids) + + +async def web_search(tools, query: str, keywords: str = "") -> dict: + if not tools.has_web(): + return {"chunks": [], "doc_aggs": []} + + _LOG.info(f"[Web search]: {query} -> {keywords}") + try: + from common.misc_utils import thread_pool_exec + + effective_query = f"{query} {keywords}".strip() if keywords else query + tav_res = await thread_pool_exec(tools.tav.retrieve_chunks, effective_query) + return {"chunks": tav_res.get("chunks", []), "doc_aggs": tav_res.get("doc_aggs", [])} + except Exception: + _LOG.exception("web_search failed") + return {"chunks": [], "doc_aggs": []} + + +async def structured_query(tools, question: str, kb_ids: list[str] | None = None) -> dict: + _LOG.info(f"[Structured search]: {question}") + sql_kbs = [kb for kb in tools.sql_kbs if kb_ids is None or kb.id in kb_ids] + if not sql_kbs: + return {"answer": "", "chunks": [], "doc_aggs": []} + from api.db.services.dialog_service import use_sql + + tenant_id = sql_kbs[0].tenant_id + sql_kb_ids = [kb.id for kb in sql_kbs] + try: + ans = await use_sql(question, tools.field_map, tenant_id, tools.chat_mdl, quota=True, kb_ids=sql_kb_ids) + except Exception: + _LOG.exception("structured_query failed") + return {"answer": "", "chunks": [], "doc_aggs": []} + if not ans: + return {"answer": "", "chunks": [], "doc_aggs": []} + ref = ans.get("reference") or {} + return { + "answer": ans.get("answer", "") or "", + "chunks": ref.get("chunks") or [], + "doc_aggs": ref.get("doc_aggs") or [], + } diff --git a/rag/advanced_rag/harness/types.py b/rag/advanced_rag/harness/types.py new file mode 100644 index 0000000000..3bd2328b39 --- /dev/null +++ b/rag/advanced_rag/harness/types.py @@ -0,0 +1,162 @@ +"""Data types for Agentic RAG harness.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Literal + + +# ═══════════════════════════════════════════════════════════════ +# Route +# ═══════════════════════════════════════════════════════════════ + + +@dataclass +class RouteDecision: + question: str + thinking_mode: str + question_type: str # factual | comparative | analytical | procedural | exploratory | verification | summarization + requires_decomposition: bool + suggests_compilation: str | None + execution_strategy: str + reasoning: str = "" + + +# ═══════════════════════════════════════════════════════════════ +# Thinking Mode +# ═══════════════════════════════════════════════════════════════ + + +@dataclass +class ExecutionStrategy: + label: Literal["low", "medium", "high", "ultra"] + execution_strategy: Literal["direct_search", "decompose_and_search", "agentic_research", "deep_research"] + requires_decomposition: bool + requires_agent_loop: bool + requires_sufficiency_judge: bool + requires_selective_gen: bool + allows_dynamic_claims: bool + allows_replan: bool + max_orchestrator_cycles: int + max_agent_cycles: int + max_parallel_agents: int + available_tools: list[str] + sufficiency_threshold: float + partial_threshold: float + fallback_to_direct_llm: bool + + +# ═══════════════════════════════════════════════════════════════ +# Plan & Claims +# ═══════════════════════════════════════════════════════════════ + + +@dataclass +class ClaimTarget: + claim_id: str + description: str + priority: int = 0 + is_verified: bool = False + confidence: float = 0.0 + suggested_tools: list[str] = field(default_factory=list) + agent_result: dict | None = None + + +@dataclass +class WorkflowPlan: + plan_type: str # direct | fact_decomposition | comparative_decomposition | procedural_decomposition | exploratory_decomposition + claims: list[ClaimTarget] + max_iterations: int + + +# ═══════════════════════════════════════════════════════════════ +# Agent Result +# ═══════════════════════════════════════════════════════════════ + + +@dataclass +class AgentResult: + claim_id: str + report: str + is_verified: bool + confidence: float + evidence_ids: list[int] = field(default_factory=list) + gaps: list[str] = field(default_factory=list) + discovered_claims: list[str] = field(default_factory=list) + + +# ═══════════════════════════════════════════════════════════════ +# Cross Check +# ═══════════════════════════════════════════════════════════════ + + +@dataclass +class ClaimCrossCheckResult: + claim_id: str + cross_check_passed: bool + cross_check_score: float + evidence_matches: list[str] = field(default_factory=list) + mismatches: list[str] = field(default_factory=list) + + +# ═══════════════════════════════════════════════════════════════ +# Sufficiency Verdict +# ═══════════════════════════════════════════════════════════════ + + +@dataclass +class SufficiencyVerdict: + status: str # SUFFICIENT | USEFUL_BUT_INCOMPLETE | INSUFFICIENT | CONFLICTING | UNANSWERABLE + score: float + agent_score: float + cross_score: float + claim_assessments: list[dict] + has_conflicts: bool + missing_claims: list[str] + feedback: str + overall_reason: str + + +# ═══════════════════════════════════════════════════════════════ +# Pipeline +# ═══════════════════════════════════════════════════════════════ + + +@dataclass +class ToolResult: + chunks: list[dict] + metadata: dict = field(default_factory=dict) + error: str | None = None + + +# ═══════════════════════════════════════════════════════════════ +# Orchestrator Context +# ═══════════════════════════════════════════════════════════════ + + +@dataclass +class OrchestratorContext: + question: str + claims: list[ClaimTarget] + mode: str + iteration: int = 0 + current_phase: str = "locate" + agent_results: dict[str, Any] = field(default_factory=dict) + verdict: SufficiencyVerdict | None = None + history: list[dict] = field(default_factory=list) + _last_entity: str | None = None + + @property + def last_entity(self) -> str | None: + return self._last_entity + + @property + def current_claim(self) -> str | None: + unverified = [c for c in self.claims if not c.is_verified] + return unverified[0].description if unverified else None + + def has_any_chunks(self) -> bool: + return any(r.get("evidence_ids") for r in self.agent_results.values()) + + def record_fallback(self, tool_name: str, fallback_from: str | None = None): + pass diff --git a/rag/advanced_rag/think_log.py b/rag/advanced_rag/think_log.py new file mode 100644 index 0000000000..87790501b3 --- /dev/null +++ b/rag/advanced_rag/think_log.py @@ -0,0 +1,94 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Surface selected internal INFO logs to the client as ```` content. + +During an agentic ``rag_agent`` turn we attach a context-scoped logging sink so +the pipeline's bracket-tagged progress logs — ``[agentic-rag]``, +``[formalize_question]``, ``[pre_search]``, ``[Planner]``, ``[orchestrator]``, +``[agentic]``, ``[Hybrid search]``, ``[BM25 search]``, ``[Web search]``, +``[ToolLoop]``, ``[FunctionTool]`` … — can be streamed to the front end as +reasoning without instrumenting every call site. + +The sink is stored in a :class:`contextvars.ContextVar`, so only the async task +tree of the current request (which inherits the context) forwards its logs — +concurrent requests stay isolated. Logs emitted from ``thread_pool_exec`` +workers that did not inherit the context are simply skipped. +""" + +from __future__ import annotations + +import contextvars +import logging +from typing import Callable + +# Per-request sink: a callable(str) that forwards one log line, or None when no +# agentic turn is streaming in the current context. +_think_log_sink: contextvars.ContextVar[Callable[[str], None] | None] = contextvars.ContextVar( + "think_log_sink", default=None +) + +# Only bracket-tagged INFO lines from these logger namespaces are surfaced. +_SCOPED_PREFIXES = ("rag.advanced_rag", "rag.llm.chat_model", "rag.llm.tool_decorator") + +_installed = False + + +class ThinkLogHandler(logging.Handler): + """Forward in-scope, bracket-tagged records to the active per-request sink.""" + + def emit(self, record: logging.LogRecord) -> None: + sink = _think_log_sink.get() + if sink is None: + return + name = record.name or "" + if not name.startswith(_SCOPED_PREFIXES): + return + try: + msg = record.getMessage() + except Exception: + return + # Only the bracket-tagged progress lines ("[Hybrid search] ..."). + if not msg or not msg.lstrip().startswith("["): + return + try: + sink("
"+msg.strip()) + except Exception: + # Never let think-log forwarding break the request or the logging + # subsystem itself. + pass + + +def install_think_log_handler() -> None: + """Install the forwarding handler on the root logger exactly once.""" + global _installed + if _installed: + return + handler = ThinkLogHandler(level=logging.INFO) + logging.getLogger().addHandler(handler) + _installed = True + + +def set_think_log_sink(sink: Callable[[str], None] | None): + """Activate ``sink`` for the current context; returns the reset token.""" + return _think_log_sink.set(sink) + + +def reset_think_log_sink(token) -> None: + try: + _think_log_sink.reset(token) + except Exception: + pass diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py index 55a81d414a..9a9161a416 100644 --- a/rag/llm/chat_model.py +++ b/rag/llm/chat_model.py @@ -679,8 +679,22 @@ class Base(ABC): args = json_repair.loads(tc.function.arguments) except Exception: args = {} - yield self._verbose_tool_use(tc.function.name, args, "Begin to call...") + yield f"Executing {tc.function.name} with args: {tc.function.arguments}" results = await asyncio.gather(*[_exec_tool(tc) for tc in tcs]) + + # Terminal-tool short-circuit: stream a terminal tool's + # result (already the final answer) and stop the loop. + _terminal = getattr(self, "terminal_tools", None) + if _terminal: + for tc, name, args, result, err in results: + if name in _terminal and not err: + logging.info(f"[ToolLoop] terminal tool {name!r} called; streaming result and stopping") + out = result if isinstance(result, str) else json.dumps(result, ensure_ascii=False) + if out: + yield out + yield total_tokens + return + history = self._append_history_batch(history, results) for tc, name, args, result, err in results: yield self._verbose_tool_use(name, args, err if err else result) @@ -1925,7 +1939,7 @@ class LiteLLMBase(ABC): history = deepcopy(hist) try: for _ in range(self.max_rounds + 1): - logging.info(f"{self.tools=}") + logging.info(f"HAS TOOL:{len(self.tools)}\n{history=}") completion_args = self._construct_completion_args(history=history, stream=False, tools=True, **gen_conf) response = await litellm.acompletion( @@ -2122,8 +2136,23 @@ class LiteLLMBase(ABC): args = json_repair.loads(tc.function.arguments) except Exception: args = {} - yield self._verbose_tool_use(tc.function.name, args, "Begin to call...") + yield f"Executing {tc.function.name} with args: {tc.function.arguments}" results = await asyncio.gather(*[_exec_tool(tc) for tc in tcs]) + + # Terminal-tool short-circuit: a terminal tool already + # produces the final answer, so stream its result and stop + # instead of feeding it back for another LLM round. + _terminal = getattr(self, "terminal_tools", None) + if _terminal: + for tc, name, args, result, err in results: + if name in _terminal and not err: + logging.info(f"[ToolLoop] terminal tool {name!r} called; streaming result and stopping") + out = result if isinstance(result, str) else json.dumps(result, ensure_ascii=False) + if out: + yield out + yield total_tokens + return + history = self._append_history_batch( history, results, diff --git a/rag/llm/tool_decorator.py b/rag/llm/tool_decorator.py index 2870f94ed8..3f1a421eb0 100644 --- a/rag/llm/tool_decorator.py +++ b/rag/llm/tool_decorator.py @@ -91,14 +91,33 @@ def _py_type_to_json(py_type: Any) -> dict[str, Any]: return {"type": "string"} -_PARAM_RE = re.compile(r"^\s*:param\s+(?P\w+)\s*:\s*(?P.+?)\s*$") +_PARAM_RE = re.compile(r"^\s*:param\s+(?P\w+)\s*:\s*(?P.*?)\s*$") +_GOOGLE_ARGS_HDR_RE = re.compile(r"^(Args|Arguments|Parameters)\s*:\s*$") +_GOOGLE_SECTION_HDR_RE = re.compile( + r"^(Returns?|Yields?|Raises|Notes?|Examples?|Attributes?|Todo|See Also|Warning|Warnings|Tip)\s*:\s*$" +) +# Google-style parameter line: leading indent, identifier, optional ``(type)``, +# then ``: description``. The description can be empty (continuation lines fill it). +_GOOGLE_PARAM_RE = re.compile( + r"^(?P\s+)(?P\w+)\s*(?:\([^)]*\))?\s*:\s*(?P.*)$" +) def _parse_param_docs(docstring: str | None) -> tuple[str, dict[str, str]]: - """Pull a short function description and ``:param name:`` lines out of a docstring. + """Pull a function description and per-parameter descriptions out of a docstring. - Intentionally minimal — Google/NumPy styles are not parsed. Anything - before the first ``:param`` line becomes the function description. + Recognises two conventions and handles multi-line descriptions in both: + + * **reST / Sphinx**: ``:param name: description`` followed by deeper-indented + continuation lines. + * **Google**: an ``Args:`` (or ``Arguments:`` / ``Parameters:``) section + whose body is `` name: description`` lines, with deeper-indented + continuation lines folded onto the same entry. Other Google sections + (``Returns:``, ``Raises:``, ...) terminate the description but are + otherwise dropped — they aren't sent to the LLM. + + Both styles can co-exist in one docstring. Anything before the first + parameter entry / section header becomes the function description. """ if not docstring: return "", {} @@ -106,12 +125,65 @@ def _parse_param_docs(docstring: str | None) -> tuple[str, dict[str, str]]: lines = inspect.cleandoc(docstring).splitlines() desc_lines: list[str] = [] param_docs: dict[str, str] = {} + state = "desc" # "desc" | "rst_param" | "google_args" | "other_section" + current_param: str | None = None + current_indent = 0 + after_first_param = False + + def _append_continuation(name: str, text: str) -> None: + param_docs[name] = (param_docs[name] + " " + text).strip() if param_docs.get(name) else text + for line in lines: + stripped = line.strip() + line_indent = len(line) - len(line.lstrip()) + + # reST :param: line — works in any state, resets it. m = _PARAM_RE.match(line) if m: - param_docs[m.group("name")] = m.group("desc") - elif not param_docs: + current_param = m.group("name") + current_indent = line_indent + param_docs[current_param] = m.group("desc").strip() + state = "rst_param" + after_first_param = True + continue + + # Google section headers. + if _GOOGLE_ARGS_HDR_RE.match(stripped): + state = "google_args" + current_param = None + after_first_param = True + continue + if _GOOGLE_SECTION_HDR_RE.match(stripped): + state = "other_section" + current_param = None + after_first_param = True + continue + + # Google `` name: desc`` entry inside an Args block. + if state == "google_args": + gm = _GOOGLE_PARAM_RE.match(line) + if gm: + current_param = gm.group("name") + current_indent = line_indent + param_docs[current_param] = gm.group("desc").strip() + continue + + # Continuation line for the most recent reST or Google param. + if state in ("rst_param", "google_args") and current_param and stripped: + if line_indent > current_indent: + _append_continuation(current_param, stripped) + continue + + # Blank line ends the current param's continuation but stays in-state. + if not stripped: + current_param = None + continue + + # Lines outside any param block, before the first param/section, + # accumulate as the function description. + if not after_first_param: desc_lines.append(line) + return "\n".join(desc_lines).strip(), param_docs @@ -153,19 +225,58 @@ def _build_openai_schema(fn: Callable[..., Any]) -> dict[str, Any]: } -def tool(fn: Callable[..., Any]) -> Callable[..., Any]: +# Sentinel separating "caller did not pass a timeout" from "caller passed None +# (= run forever)". Plain ``None`` is a legal value for the kwarg. +_TIMEOUT_UNSET: Any = object() + + +def tool( + fn: Callable[..., Any] | None = None, + *, + timeout: float | int | None = _TIMEOUT_UNSET, +) -> Callable[..., Any]: """Mark ``fn`` as an LLM tool and attach an OpenAI-format schema to it. - The wrapped callable is the same callable — we only set two attributes: + Usable in two styles: + + * Bare: ``@tool`` — no per-tool timeout; the session + falls back to its caller-supplied + ``request_timeout`` (default 10s). + * Parameterised: ``@tool(timeout=60)`` — 60s timeout, overrides the + session's default for this tool. + Pass ``timeout=None`` to disable + the timeout entirely (the tool + runs until it completes). + + The wrapped callable is the same callable — we only set attributes on it: * ``fn._is_tool = True`` — sentinel so :meth:`Base.bind_tools` can tell a ``@tool`` callable apart from a raw schema dict. * ``fn.openai_schema`` — the schema dict passed verbatim to the LLM provider in the ``tools=[...]`` request field. + * ``fn._tool_timeout`` (only when ``timeout=`` was passed) — read by + :class:`FunctionToolSession` to override its default timeout for this + tool. May be ``None`` to mean "no timeout". """ - fn.openai_schema = _build_openai_schema(fn) # type: ignore[attr-defined] - fn._is_tool = True # type: ignore[attr-defined] - return fn + + def decorate(f: Callable[..., Any]) -> Callable[..., Any]: + f.openai_schema = _build_openai_schema(f) # type: ignore[attr-defined] + f._is_tool = True # type: ignore[attr-defined] + if timeout is not _TIMEOUT_UNSET: + f._tool_timeout = timeout # type: ignore[attr-defined] + return f + + # ``@tool`` (no parens) — ``fn`` is the function being decorated. + if fn is not None: + if not callable(fn): + raise TypeError( + "@tool used incorrectly. Use `@tool` or `@tool(timeout=N)`; " + f"got first positional argument of type {type(fn).__name__}." + ) + return decorate(fn) + + # ``@tool(timeout=N)`` — return the decorator that will receive the function. + return decorate def is_tool(obj: Any) -> bool: @@ -188,21 +299,25 @@ class FunctionToolSession: self.tools_map: dict[str, Callable[..., Any]] = {} for fn in tools: if not is_tool(fn): - raise TypeError(f"{getattr(fn, '__name__', fn)!r} is not a @tool-decorated callable") + raise TypeError( + f"{getattr(fn, '__name__', fn)!r} is not a @tool-decorated callable" + ) self.tools_map[fn.openai_schema["function"]["name"]] = fn @property def schemas(self) -> list[dict[str, Any]]: return [fn.openai_schema for fn in self.tools_map.values()] - def tool_call(self, name: str, arguments: dict[str, Any], timeout: float | int = 10) -> Any: + def tool_call(self, name: str, arguments: dict[str, Any], timeout: float | int = 300) -> Any: return asyncio.run(self.tool_call_async(name, arguments, request_timeout=timeout)) - async def tool_call_async(self, name: str, arguments: dict[str, Any], request_timeout: float | int = 10) -> Any: + async def tool_call_async(self, name: str, arguments: dict[str, Any], request_timeout: float | int = 300) -> Any: if name not in self.tools_map: raise KeyError(f"Tool {name!r} is not registered") if not isinstance(arguments, Mapping): - raise TypeError(f"Tool arguments for {name} must be an object, got {type(arguments).__name__}") + raise TypeError( + f"Tool arguments for {name} must be an object, got {type(arguments).__name__}" + ) fn = self.tools_map[name] logging.info(f"[FunctionTool] invoke name={name} args={str(arguments)[:200]}") if asyncio.iscoroutinefunction(fn): @@ -214,4 +329,9 @@ class FunctionToolSession: # background until it returns. Callers should treat sync tools # that block on I/O accordingly. coro = thread_pool_exec(fn, **arguments) - return await asyncio.wait_for(coro, timeout=request_timeout) + # Per-tool timeout set via ``@tool(timeout=N)`` overrides the + # session-default. ``None`` is a legal explicit choice meaning + # "wait forever" — ``asyncio.wait_for(..., timeout=None)`` handles it. + configured = getattr(fn, "_tool_timeout", _TIMEOUT_UNSET) + effective_timeout = request_timeout if configured is _TIMEOUT_UNSET else configured + return await asyncio.wait_for(coro, timeout=effective_timeout) \ No newline at end of file diff --git a/rag/prompts/generator.py b/rag/prompts/generator.py index 30a5848b19..2ef35fdbdb 100644 --- a/rag/prompts/generator.py +++ b/rag/prompts/generator.py @@ -967,6 +967,23 @@ async def sufficiency_check(chat_mdl, question: str, ret_content: str): return {} +SUFFICIENCY_SELECT = load_prompt("sufficiency_select") + + +async def sufficiency_select(chat_mdl, question: str, ret_content: str): + """Sufficiency judgement that also returns the IDs of the useful chunks. + + ``ret_content`` must label each chunk with an ``ID: n`` marker (as + :func:`kb_prompt` does). Returns a dict with ``is_sufficient``, + ``reasoning``, ``missing_information`` and ``useful_chunk_ids``. + """ + try: + return await gen_json(PROMPT_JINJA_ENV.from_string(SUFFICIENCY_SELECT).render(question=question, retrieved_docs=ret_content), "Output:\n", chat_mdl) + except Exception as e: + logging.exception(e) + return {} + + MULTI_QUERIES_GEN = load_prompt("multi_queries_gen") diff --git a/rag/prompts/sufficiency_select.md b/rag/prompts/sufficiency_select.md new file mode 100644 index 0000000000..c0435b9190 --- /dev/null +++ b/rag/prompts/sufficiency_select.md @@ -0,0 +1,28 @@ +You are an information retrieval evaluation expert. Assess whether the currently retrieved content is sufficient to answer the user's question(s), and identify exactly which retrieved chunks are useful. + +Each retrieved chunk is labeled with an integer ID on a line like `ID: 3`. + +User question(s): +{{ question }} + +Retrieved content: +{{ retrieved_docs }} + +Determine whether these contents are sufficient to answer the user's question(s), and list the IDs of the chunks that actually contribute useful information toward answering them. + +Output format (JSON): +```json +{ + "is_sufficient": true/false, + "reasoning": "Your reasoning for the judgment", + "missing_information": ["Missing information 1", "Missing information 2"], + "useful_chunk_ids": [0, 3, 7] +} +``` + +Requirements: +1. If the retrieved content contains the key information needed to answer the question(s), judge as sufficient (true). +2. If key information is missing, judge as insufficient (false), and list the missing information. +3. `useful_chunk_ids` must contain ONLY the integer IDs (taken from the `ID:` labels above) of chunks that provide information useful for answering the question(s). Exclude irrelevant or redundant chunks. Use an empty array when none are useful. +4. The `missing_information` should only be filled when insufficient, otherwise an empty array. +5. The `reasoning` should be concise and clear. diff --git a/rag/utils/es_conn.py b/rag/utils/es_conn.py index b69cca2570..99239c6de5 100644 --- a/rag/utils/es_conn.py +++ b/rag/utils/es_conn.py @@ -230,7 +230,7 @@ class ESConnection(ESConnectionBase): if bool_query: s = s.query(bool_query) for field in highlight_fields: - s = s.highlight(field) + s = s.highlight(field, fragment_size=50, number_of_fragments=5) if order_by: orders = list() diff --git a/test/testcases/restful_api/test_chats.py b/test/testcases/restful_api/test_chats.py index b86fda7830..5d5a9b0882 100644 --- a/test/testcases/restful_api/test_chats.py +++ b/test/testcases/restful_api/test_chats.py @@ -702,6 +702,7 @@ def _load_chat_routes_unit_module(monkeypatch): dialog_service_mod.async_ask = lambda *_args, **_kwargs: None dialog_service_mod.async_chat = lambda *_args, **_kwargs: None dialog_service_mod.gen_mindmap = lambda *_args, **_kwargs: None + dialog_service_mod.rag_agent = lambda *_args, **_kwargs: None monkeypatch.setitem(sys.modules, "api.db.services.dialog_service", dialog_service_mod) conversation_service_mod = ModuleType("api.db.services.conversation_service") diff --git a/test/testcases/restful_api/test_user_tenant_routes_unit.py b/test/testcases/restful_api/test_user_tenant_routes_unit.py index f983f40fc3..1076acd4ac 100644 --- a/test/testcases/restful_api/test_user_tenant_routes_unit.py +++ b/test/testcases/restful_api/test_user_tenant_routes_unit.py @@ -1471,6 +1471,7 @@ def _load_chat_routes_unit_module(monkeypatch): dialog_service_mod.async_ask = lambda *_args, **_kwargs: None dialog_service_mod.async_chat = lambda *_args, **_kwargs: None dialog_service_mod.gen_mindmap = lambda *_args, **_kwargs: None + dialog_service_mod.rag_agent = lambda *_args, **_kwargs: None monkeypatch.setitem(sys.modules, "api.db.services.dialog_service", dialog_service_mod) conversation_service_mod = ModuleType("api.db.services.conversation_service") diff --git a/uv.lock b/uv.lock index 478f174c34..1e0db4f334 100644 --- a/uv.lock +++ b/uv.lock @@ -3114,12 +3114,33 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/32/1f/2a2b5eea8ef5762a86ad3f8fddddaaba2c0d76dd44e644b9158900868bec/json_repair-0.60.1-py3-none-any.whl", hash = "sha256:ba6ff974f2a8bef2f7768144a7f03f870a816443f03da27a49cdd0ec31a78049" }, ] +[[package]] +name = "jsonpatch" +version = "1.33" +source = { registry = "https://mirrors.aliyun.com/pypi/simple" } +dependencies = [ + { name = "jsonpointer" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/42/78/18813351fe5d63acad16aec57f94ec2b70a09e53ca98145589e185423873/jsonpatch-1.33.tar.gz", hash = "sha256:9fcd4009c41e6d12348b4a0ff2563ba56a2923a7dfee731d004e212e1ee5030c" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/73/07/02e16ed01e04a374e644b575638ec7987ae846d25ad97bcc9945a3ee4b0e/jsonpatch-1.33-py2.py3-none-any.whl", hash = "sha256:0ae28c0cd062bbd8b8ecc26d7d164fbbea9652a1a3693f3b956c1eae5145dade" }, +] + [[package]] name = "jsonpath" version = "0.82.2" source = { registry = "https://mirrors.aliyun.com/pypi/simple" } sdist = { url = "https://mirrors.aliyun.com/pypi/packages/cf/a1/693351acd0a9edca4de9153372a65e75398898ea7f8a5c722ab00f464929/jsonpath-0.82.2.tar.gz", hash = "sha256:d87ef2bcbcded68ee96bc34c1809b69457ecec9b0c4dd471658a12bd391002d1" } +[[package]] +name = "jsonpointer" +version = "3.1.1" +source = { registry = "https://mirrors.aliyun.com/pypi/simple" } +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/18/c7/af399a2e7a67fd18d63c40c5e62d3af4e67b836a2107468b6a5ea24c4304/jsonpointer-3.1.1.tar.gz", hash = "sha256:0b801c7db33a904024f6004d526dcc53bbb8a4a0f4e32bfd10beadf60adf1900" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/9e/6a/a83720e953b1682d2d109d3c2dbb0bc9bf28cc1cbc205be4ef4be5da709d/jsonpointer-3.1.1-py3-none-any.whl", hash = "sha256:8ff8b95779d071ba472cf5bc913028df06031797532f08a7d5b602d8b2a488ca" }, +] + [[package]] name = "jsonschema" version = "4.26.0" @@ -3189,6 +3210,38 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/2a/8f/8f6f491d595a9e5912971f3f863d81baddccc8a4d0c3749d6a0dd9ffc9df/kiwisolver-1.4.9-cp313-cp313t-win_arm64.whl", hash = "sha256:0749fd8f4218ad2e851e11cc4dc05c7cbc0cbc4267bdfdb31782e65aace4ee9c" }, ] +[[package]] +name = "langchain-core" +version = "1.4.9" +source = { registry = "https://mirrors.aliyun.com/pypi/simple" } +dependencies = [ + { name = "jsonpatch" }, + { name = "langchain-protocol" }, + { name = "langsmith" }, + { name = "packaging" }, + { name = "pydantic" }, + { name = "pyyaml" }, + { name = "tenacity" }, + { name = "typing-extensions" }, + { name = "uuid-utils" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/2a/b9/e937d0a90b26540bff07e7a7c64349f3b29c2dcc36257cd1cd3fdce17f2a/langchain_core-1.4.9.tar.gz", hash = "sha256:f8078901145bed0466755277500a5a22822a7b628808c4c0a28d4fc88895fcf2" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/84/70/ade2fada52772798ef815b6352b59e71b116aa0c32c3aef5be3dc2cbed12/langchain_core-1.4.9-py3-none-any.whl", hash = "sha256:28e3909e2a10cc81504952d795ac0a9e014c0018121ef89d48dd396fa09ec624" }, +] + +[[package]] +name = "langchain-protocol" +version = "0.0.18" +source = { registry = "https://mirrors.aliyun.com/pypi/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/d2/59/b5959aea96faa9146e2e49a7a22882b3528c62efafe9a6a95beab30c2305/langchain_protocol-0.0.18.tar.gz", hash = "sha256:ec3e11782f1ed0c9db38e5a9ed01b0e7a0d3fba406faa8aef6594b73c56a63e6" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/99/2e/d82db9eec13ad0f72e7aaad5c4bc730ab111934fdc83c85523206eb9b0a0/langchain_protocol-0.0.18-py3-none-any.whl", hash = "sha256:70b53a86fbf9cedc863555effe44da192ab02d556ddbf2cf95b8873adcf41b5a" }, +] + [[package]] name = "langfuse" version = "4.7.1" @@ -3208,6 +3261,87 @@ wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/9f/9a/bd3368f46b6c72ee2068b80536826b02ae86df53eff1c79941344503098f/langfuse-4.7.1-py3-none-any.whl", hash = "sha256:a4e59c81ad5e5b16a65d3849f4923ebc3ad6e67ec803ada83d50c0cb66149490" }, ] +[[package]] +name = "langgraph" +version = "1.2.0" +source = { registry = "https://mirrors.aliyun.com/pypi/simple" } +dependencies = [ + { name = "langchain-core" }, + { name = "langgraph-checkpoint" }, + { name = "langgraph-prebuilt" }, + { name = "langgraph-sdk" }, + { name = "pydantic" }, + { name = "xxhash" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/58/61/d5d25e783035aa307d289b37e082258a6061c0fb4caa4a284f3bf1e87169/langgraph-1.2.0.tar.gz", hash = "sha256:4a9baaf62afc5d5f63144a50095140a34b9aa9b7cea695d25326d564775348e7" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/f6/e8/e3304ac0015c2bdb04ad9785e4ed65c788855ce7857ce6104dd2f5d322db/langgraph-1.2.0-py3-none-any.whl", hash = "sha256:03fd5895a8d4b70db1ff63ebc3bacead29dd20cd794a8b1a483e7ec9018f7a65" }, +] + +[[package]] +name = "langgraph-checkpoint" +version = "4.1.1" +source = { registry = "https://mirrors.aliyun.com/pypi/simple" } +dependencies = [ + { name = "langchain-core" }, + { name = "ormsgpack" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/83/47/886af6f886f0bff2273164a45f008694e48a96ff3cd25ff0228f2aa9480e/langgraph_checkpoint-4.1.1.tar.gz", hash = "sha256:6c2bdb530c91f91d7d9c1bd100925d0fc4f498d418c17f3587d1526279482a25" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/bd/b4/71425e3e38be92611300b9cc5e46a5bf98ab23f5ea8a75b73d02a2f1413c/langgraph_checkpoint-4.1.1-py3-none-any.whl", hash = "sha256:25d29144b082827218e7bc3f1e9b0566a4bb007895cd6cc26f66a8428739f56e" }, +] + +[[package]] +name = "langgraph-prebuilt" +version = "1.1.0" +source = { registry = "https://mirrors.aliyun.com/pypi/simple" } +dependencies = [ + { name = "langchain-core" }, + { name = "langgraph-checkpoint" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/29/66/ed9b93f56bc17ef22d551892f0ac2b225a97fe0fcf23a511b857f70d590b/langgraph_prebuilt-1.1.0.tar.gz", hash = "sha256:3c579cf6eed2d17f9c157c2d0fcaddcd8688524e7022d3b22b37a3bf4589d528" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/e9/43/3fe1a700b8490ed02679cdbbc8c915eb23a092faf496c9c1118abcd10be3/langgraph_prebuilt-1.1.0-py3-none-any.whl", hash = "sha256:51e311747d755b751d5c6b39b0c1446124d3a7643d2515017e6714b323508fc9" }, +] + +[[package]] +name = "langgraph-sdk" +version = "0.3.6" +source = { registry = "https://mirrors.aliyun.com/pypi/simple" } +dependencies = [ + { name = "httpx" }, + { name = "orjson" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/3e/ec/477fa8b408f948b145d90fd935c0a9f37945fa5ec1dfabfc71e7cafba6d8/langgraph_sdk-0.3.6.tar.gz", hash = "sha256:7650f607f89c1586db5bee391b1a8754cbe1fc83b721ff2f1450f8906e790bd7" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/d8/61/12508e12652edd1874327271a5a8834c728a605f53a1a1c945f13ab69664/langgraph_sdk-0.3.6-py3-none-any.whl", hash = "sha256:7df2fd552ad7262d0baf8e1f849dce1d62186e76dcdd36db9dc5bdfa5c3fc20f" }, +] + +[[package]] +name = "langsmith" +version = "0.10.5" +source = { registry = "https://mirrors.aliyun.com/pypi/simple" } +dependencies = [ + { name = "anyio" }, + { name = "distro" }, + { name = "httpx" }, + { name = "orjson", marker = "platform_python_implementation != 'PyPy'" }, + { name = "packaging" }, + { name = "pydantic" }, + { name = "requests" }, + { name = "requests-toolbelt" }, + { name = "sniffio" }, + { name = "typing-extensions" }, + { name = "uuid-utils" }, + { name = "websockets" }, + { name = "xxhash" }, + { name = "zstandard" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/1f/65/3867765976e4d43b98a4ea6a41c0712dda17a600ad998e02976b445874d7/langsmith-0.10.5.tar.gz", hash = "sha256:60053c1d88dc332a002cbac38601cc8b912466e7fc2a86bc9e690fa4d5bc1c78" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/a9/3e/213d9bb122f97d89987bd4c175cc4be9f2fa090e868ac8b5156c3265d8dd/langsmith-0.10.5-py3-none-any.whl", hash = "sha256:116adf2c30dfc1d0daf16919879b90c4093aad6122f44b80cc1f035b874dc9d6" }, +] + [[package]] name = "lark" version = "1.3.1" @@ -3905,12 +4039,12 @@ name = "onnxruntime-gpu" version = "1.23.2" source = { registry = "https://mirrors.aliyun.com/pypi/simple" } dependencies = [ - { name = "coloredlogs" }, - { name = "flatbuffers" }, - { name = "numpy" }, - { name = "packaging" }, - { name = "protobuf" }, - { name = "sympy" }, + { name = "coloredlogs", marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux')" }, + { name = "flatbuffers", marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux')" }, + { name = "numpy", marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux')" }, + { name = "packaging", marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux')" }, + { name = "protobuf", marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux')" }, + { name = "sympy", marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux')" }, ] wheels = [ { url = "https://mirrors.aliyun.com/pypi/packages/03/05/40d561636e4114b54aa06d2371bfbca2d03e12cfdf5d4b85814802f18a75/onnxruntime_gpu-1.23.2-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1e8f75af5da07329d0c3a5006087f4051d8abd133b4be7c9bae8cdab7bea4c26" }, @@ -7940,6 +8074,7 @@ dependencies = [ { name = "jira" }, { name = "json-repair" }, { name = "langfuse" }, + { name = "langgraph" }, { name = "lark-oapi" }, { name = "line-bot-sdk" }, { name = "litellm" }, @@ -8095,6 +8230,7 @@ requires-dist = [ { name = "jira", specifier = "==3.10.5" }, { name = "json-repair", specifier = "==0.60.1" }, { name = "langfuse", specifier = ">=4.0.1" }, + { name = "langgraph", specifier = "==1.2.0" }, { name = "lark-oapi", specifier = ">=1.2.0" }, { name = "line-bot-sdk", specifier = ">=3.0.0" }, { name = "litellm", specifier = "==1.82.5" }, @@ -9537,6 +9673,29 @@ socks = [ { name = "pysocks" }, ] +[[package]] +name = "uuid-utils" +version = "0.17.0" +source = { registry = "https://mirrors.aliyun.com/pypi/simple" } +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/e7/91/63938e0e7e7876658e5e40178e7c0735b53527886fe11797a11699c55edd/uuid_utils-0.17.0.tar.gz", hash = "sha256:abb5667a36119019b3fa320c4d10c21ebccfcc87c8a739e6a0056cee7f48dde2" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/d2/dd/614fb9912157ac0128e6050859ccf06d9f13df9a944a803e8f80f6157e38/uuid_utils-0.17.0-cp313-cp313-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:d11a7bc1e02da8984d32e6de9e0826c6edac00eac17de270f372bf32f9a0af63" }, + { url = "https://mirrors.aliyun.com/pypi/packages/3e/11/d072711704de3d21bec08b6c2f36a215200ca1d5e01a390ea1ac434080a0/uuid_utils-0.17.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:7a49f47ac26df3e431c56b825c1bae8e6d3d591fdbb7438c227cc9845a7e3d73" }, + { url = "https://mirrors.aliyun.com/pypi/packages/18/6d/8a63e5eb2d5a6ba69a6c2036e305075bd6f5a022e7ea25fc6ce0eb7c51d2/uuid_utils-0.17.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:32df1944808877702ceea398c103881c09a679bb672a215e01c2a84231266bf9" }, + { url = "https://mirrors.aliyun.com/pypi/packages/f7/2d/bdc2caf9719d9090d7c46043242ae6136cba4f7a7ee384992ab905ad9aa1/uuid_utils-0.17.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:98c88d3edd08e7245562e9815996dbc6f0bd4745e1c76462f24af5ae4e187dd1" }, + { url = "https://mirrors.aliyun.com/pypi/packages/b6/33/9219d09d51ead282b578b2a4e0a515c2cce3ec52076cada8bfb7e35727d5/uuid_utils-0.17.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5a4370089c8b2e42f1db51d76408c7fa8eaa2934bf854d17983d16179c07c098" }, + { url = "https://mirrors.aliyun.com/pypi/packages/d8/79/e8e0f8b3955f2081c116157119d87659937893242eb834aa170da04d660b/uuid_utils-0.17.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:09a55b7a5ae764985cb46467496a1787678d0a1400356157a080ad95b1a36869" }, + { url = "https://mirrors.aliyun.com/pypi/packages/d5/5e/d1ceddc430ff04b6e21704b2030d4438074a2f478b265dab43da957791c1/uuid_utils-0.17.0-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:56aa6488b931246fae11924e4bd0e2b32677e63945eecb71c29e3c2ca0dc3131" }, + { url = "https://mirrors.aliyun.com/pypi/packages/d5/62/89438e12f389a843e626b7e37691319a057b3d6b80914609106891faadda/uuid_utils-0.17.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:309a35f12d99dde19032bc2259cda6431c85eeac0879134dc777cc3087d7e1cb" }, + { url = "https://mirrors.aliyun.com/pypi/packages/87/d2/eedcd99f522d60e238ead03844f0d51743ba84d33044959e230b756bf212/uuid_utils-0.17.0-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:21c79b61ff750abcf057163dd764ccb6196cde7a26cda1b31b45cd97769e03b3" }, + { url = "https://mirrors.aliyun.com/pypi/packages/0e/a8/bb1b38aaddd7243b6e562c6694f499bf094800918316192fd8cb2cdc2620/uuid_utils-0.17.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:4134353bfe3026ddab8e886002dc52bc5a0ab04611aabb0eaae23c32e6e57f64" }, + { url = "https://mirrors.aliyun.com/pypi/packages/b4/77/5f7ed930dc105e293845c09e4d5bd84076318a12f45a46783e1af64906d7/uuid_utils-0.17.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:7c89359affecebe2e39e6a116d069b363c936511a9572b308402489a26957d89" }, + { url = "https://mirrors.aliyun.com/pypi/packages/fd/25/1b55697adf6811a6f92cff6340e6b03e31fd6bc51066a5c10698c29b3679/uuid_utils-0.17.0-cp313-cp313-pyemscripten_2025_0_wasm32.whl", hash = "sha256:6a019a31bc4db89a0903a3e4f6b218571f3a6ff0ad4b3d3fe1c8f91a05ff6e3e" }, + { url = "https://mirrors.aliyun.com/pypi/packages/26/bf/cd729343de4684230be8a966bad7bfc2cf10ce3e643b1189a8b5370dbe35/uuid_utils-0.17.0-cp313-cp313-win32.whl", hash = "sha256:b3131a82d0c7611f0aa480a6d36929e001a3f54ba0fc029a8118a5863cce513c" }, + { url = "https://mirrors.aliyun.com/pypi/packages/76/f0/e602ae0a1b139a7826e5189b93d91902564def06d5006324fd2faf82c8fc/uuid_utils-0.17.0-cp313-cp313-win_amd64.whl", hash = "sha256:9e311f908d2f842fca4c7dcebc4f10306b8089b204ef04cf6704b4332c9ff6ff" }, + { url = "https://mirrors.aliyun.com/pypi/packages/1a/52/024ebece265b387154115dc4f1d9727174ef82623069f4bec8b7ed7e73f7/uuid_utils-0.17.0-cp313-cp313-win_arm64.whl", hash = "sha256:c351737e2e65497c7200ab4ffb8af97e9f48be6488309abdd265fe08d66ee92f" }, +] + [[package]] name = "uuid7" version = "0.1.0"