mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-30 20:35:22 +08:00
249 lines
9.3 KiB
Python
249 lines
9.3 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from typing import Any
|
|
|
|
from langbot_plugin.api.definition.components.tool.tool import Tool
|
|
from langbot_plugin.api.entities.builtin.provider import session as provider_session
|
|
from langbot_plugin.api.proxies.query_based_api import QueryBasedAPIProxy
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class QueryKnowledge(Tool):
|
|
@staticmethod
|
|
def _preview_text(value: str, max_len: int = 120) -> str:
|
|
text = value.strip().replace("\n", " ")
|
|
if len(text) <= max_len:
|
|
return text
|
|
return f"{text[:max_len]}..."
|
|
|
|
@staticmethod
|
|
def _normalize_kb_ids(params: dict[str, Any]) -> tuple[list[str] | None, str | None]:
|
|
kb_id = params.get("kb_id")
|
|
kb_ids = params.get("kb_ids")
|
|
|
|
normalized_ids: list[str] = []
|
|
|
|
if kb_id is not None:
|
|
if not isinstance(kb_id, str):
|
|
return None, "kb_id must be a string."
|
|
if not kb_id.strip():
|
|
return None, "kb_id must be a non-empty string."
|
|
normalized_ids.append(kb_id.strip())
|
|
|
|
if kb_ids is not None:
|
|
if not isinstance(kb_ids, list):
|
|
return None, "kb_ids must be an array of strings."
|
|
if not kb_ids:
|
|
if normalized_ids:
|
|
kb_ids = None
|
|
else:
|
|
return None, "kb_ids must be a non-empty array of strings. For a single knowledge base, use kb_id instead."
|
|
|
|
if kb_ids is not None:
|
|
for item in kb_ids:
|
|
if not isinstance(item, str) or not item.strip():
|
|
return None, "kb_ids must contain only non-empty strings."
|
|
normalized_ids.append(item.strip())
|
|
|
|
if not normalized_ids:
|
|
return None, "Either kb_id or kb_ids is required for action 'query'."
|
|
|
|
deduped_ids: list[str] = []
|
|
seen_ids: set[str] = set()
|
|
for item in normalized_ids:
|
|
if item not in seen_ids:
|
|
seen_ids.add(item)
|
|
deduped_ids.append(item)
|
|
|
|
return deduped_ids, None
|
|
|
|
@staticmethod
|
|
def _sort_key(result: dict[str, Any]) -> tuple[int, float]:
|
|
distance = result.get("distance")
|
|
if isinstance(distance, (int, float)):
|
|
return (0, float(distance))
|
|
|
|
score = result.get("score")
|
|
if isinstance(score, (int, float)):
|
|
return (1, -float(score))
|
|
|
|
return (2, 0.0)
|
|
|
|
@staticmethod
|
|
def _attach_kb_metadata(result: dict[str, Any], kb_id: str) -> dict[str, Any]:
|
|
enriched = dict(result)
|
|
metadata = enriched.get("metadata")
|
|
if not isinstance(metadata, dict):
|
|
metadata = {}
|
|
else:
|
|
metadata = dict(metadata)
|
|
|
|
metadata.setdefault("knowledge_base_id", kb_id)
|
|
enriched["metadata"] = metadata
|
|
return enriched
|
|
|
|
async def call(self, params: dict[str, Any], session: provider_session.Session, query_id: int) -> str:
|
|
api = QueryBasedAPIProxy(
|
|
query_id=query_id,
|
|
plugin_runtime_handler=self.plugin.plugin_runtime_handler,
|
|
)
|
|
|
|
action = params.get("action", "list")
|
|
logger.info(
|
|
"[AgenticRAG] tool call started: query_id=%s action=%s params_keys=%s",
|
|
query_id,
|
|
action,
|
|
sorted(params.keys()),
|
|
)
|
|
|
|
if action == "list":
|
|
logger.info("[AgenticRAG] listing knowledge bases: query_id=%s", query_id)
|
|
knowledge_bases = await api.list_pipeline_knowledge_bases()
|
|
logger.info(
|
|
"[AgenticRAG] knowledge base list loaded: query_id=%s kb_count=%s",
|
|
query_id,
|
|
len(knowledge_bases),
|
|
)
|
|
if not knowledge_bases:
|
|
return "No knowledge bases are configured for the current pipeline."
|
|
return json.dumps(knowledge_bases, ensure_ascii=False)
|
|
|
|
elif action == "query":
|
|
query_text = params.get("query_text")
|
|
if not isinstance(query_text, str) or not query_text.strip():
|
|
logger.warning(
|
|
"[AgenticRAG] invalid query_text: query_id=%s query_text=%r",
|
|
query_id,
|
|
query_text,
|
|
)
|
|
return "query_text must be a non-empty string."
|
|
|
|
top_k = params.get("top_k", 5)
|
|
if not isinstance(top_k, int) or top_k <= 0:
|
|
logger.warning(
|
|
"[AgenticRAG] invalid top_k: query_id=%s top_k=%r",
|
|
query_id,
|
|
top_k,
|
|
)
|
|
return "top_k must be a positive integer."
|
|
|
|
kb_ids, kb_error = self._normalize_kb_ids(params)
|
|
if kb_error:
|
|
logger.warning(
|
|
"[AgenticRAG] invalid kb params: query_id=%s error=%s",
|
|
query_id,
|
|
kb_error,
|
|
)
|
|
return kb_error
|
|
|
|
# Keep the agent-facing tool surface minimal here. Although the
|
|
# underlying runtime supports metadata filters, different knowledge
|
|
# engines do not expose a consistent, query-time filter schema to
|
|
# the agent, so passing raw filters would be unreliable.
|
|
query_text = query_text.strip()
|
|
logger.info(
|
|
"[AgenticRAG] retrieval requested: query_id=%s kb_ids=%s kb_count=%s top_k=%s query=%r",
|
|
query_id,
|
|
kb_ids,
|
|
len(kb_ids),
|
|
top_k,
|
|
self._preview_text(query_text),
|
|
)
|
|
|
|
async def _query_single_kb(target_kb_id: str) -> dict[str, Any]:
|
|
try:
|
|
logger.info(
|
|
"[AgenticRAG] querying knowledge base: query_id=%s kb_id=%s top_k=%s",
|
|
query_id,
|
|
target_kb_id,
|
|
top_k,
|
|
)
|
|
kb_results = await api.retrieve_knowledge(
|
|
kb_id=target_kb_id,
|
|
query_text=query_text,
|
|
top_k=top_k,
|
|
)
|
|
logger.info(
|
|
"[AgenticRAG] knowledge base query succeeded: query_id=%s kb_id=%s result_count=%s",
|
|
query_id,
|
|
target_kb_id,
|
|
len(kb_results),
|
|
)
|
|
return {
|
|
"kb_id": target_kb_id,
|
|
"results": [
|
|
self._attach_kb_metadata(result, target_kb_id)
|
|
for result in kb_results
|
|
if isinstance(result, dict)
|
|
],
|
|
"error": None,
|
|
}
|
|
except Exception as exc:
|
|
logger.exception(
|
|
"[AgenticRAG] knowledge base query failed: query_id=%s kb_id=%s error=%s",
|
|
query_id,
|
|
target_kb_id,
|
|
exc,
|
|
)
|
|
return {
|
|
"kb_id": target_kb_id,
|
|
"results": [],
|
|
"error": str(exc),
|
|
}
|
|
|
|
retrieval_outcomes = await asyncio.gather(*[_query_single_kb(kb_id) for kb_id in kb_ids])
|
|
merged_results = [
|
|
result
|
|
for outcome in retrieval_outcomes
|
|
for result in outcome["results"]
|
|
]
|
|
failed_kbs = [
|
|
{"kb_id": outcome["kb_id"], "error": outcome["error"]}
|
|
for outcome in retrieval_outcomes
|
|
if outcome["error"] is not None
|
|
]
|
|
logger.info(
|
|
"[AgenticRAG] retrieval completed: query_id=%s merged_results=%s failed_kbs=%s",
|
|
query_id,
|
|
len(merged_results),
|
|
len(failed_kbs),
|
|
)
|
|
|
|
if not merged_results:
|
|
if failed_kbs:
|
|
logger.info(
|
|
"[AgenticRAG] returning failure summary: query_id=%s failed_kbs=%s",
|
|
query_id,
|
|
failed_kbs,
|
|
)
|
|
return json.dumps({"results": [], "failed_kbs": failed_kbs}, ensure_ascii=False)
|
|
logger.info("[AgenticRAG] no relevant documents found: query_id=%s", query_id)
|
|
return "No relevant documents found."
|
|
|
|
merged_results.sort(key=self._sort_key)
|
|
truncated_results = merged_results[:top_k]
|
|
logger.info(
|
|
"[AgenticRAG] returning retrieval response: query_id=%s returned_results=%s truncated=%s",
|
|
query_id,
|
|
len(truncated_results),
|
|
len(merged_results) > len(truncated_results),
|
|
)
|
|
if failed_kbs:
|
|
return json.dumps(
|
|
{"results": truncated_results, "failed_kbs": failed_kbs},
|
|
ensure_ascii=False,
|
|
)
|
|
return json.dumps(truncated_results, ensure_ascii=False)
|
|
|
|
else:
|
|
logger.warning(
|
|
"[AgenticRAG] unknown action: query_id=%s action=%r",
|
|
query_id,
|
|
action,
|
|
)
|
|
return f"Unknown action: {action}. Use 'list' or 'query'."
|