""" RAG搜索工具模块 包含知识库搜索、文件处理等功能 """ import os import json import base64 import asyncio import httpx from typing import Optional, Dict, Any from langchain_core.tools import tool from config import RAG_CONFIG from utils.function_tracker import track_function_calls @track_function_calls async def rag_search( query: str, top_k: Optional[int] = 5, file_id: Optional[int] = None, kb_id: Optional[str] = None ) -> Dict[str, Any]: """ 知识库搜索工具 调用知识库搜索API，检索与query相关的文档内容 Args: query: 搜索关键词（必填） top_k: 返回结果数量（默认5） file_id: 文件ID（可选） kb_id: 知识库ID（可选） Returns: 搜索结果，包含sourceCitation格式的数据 """ if not query or not query.strip(): return {"success": False, "error": "查询文本不能为空"} SEARCH_URL = RAG_CONFIG["search_endpoint"] params = {"query": query.strip()} if file_id is None: config_file_id = RAG_CONFIG.get("file-id") if config_file_id is not None: try: file_id = int(config_file_id) if isinstance(config_file_id, str) else config_file_id except (ValueError, TypeError): file_id = None if kb_id is None: config_kb_id = RAG_CONFIG.get("kb-id") if config_kb_id is not None: kb_id = str(config_kb_id) if config_kb_id is not None else None if file_id is not None and file_id != "": params["file_id"] = file_id if kb_id is not None and kb_id != "": params["kb_id"] = kb_id print(111111111111111111) print(RAG_CONFIG["x-user-id"]) print(RAG_CONFIG["x-user-name"]) print(RAG_CONFIG["x-role"]) headers = { "accept": "application/json", "x-user-id": RAG_CONFIG["x-user-id"], "x-role": RAG_CONFIG["x-role"] } print(f"RAG params: {params}") max_retries = 1 retry_delay = 2 response = None last_exception = None for attempt in range(max_retries + 1): try: async with httpx.AsyncClient(timeout=180.0, verify=False) as client: response = await client.get(SEARCH_URL, params=params, headers=headers) break except (httpx.ReadTimeout, httpx.ConnectTimeout, httpx.RequestError) as e: last_exception = e if attempt < max_retries: print(f"RAG请求超时，第 {attempt+1} 次重试，2秒后重新请求...") await asyncio.sleep(retry_delay) else: return {"success": False, "error": f"请求超时: {str(e)}"} try: if response.status_code != 200: return {"success": False, "error": f"HTTP {response.status_code}"} result = response.json() data = result.get("results", []) filtered_data = [item for item in data if item.get("score", 0.0) >= 0.4] sorted_data = sorted( filtered_data, key=lambda x: x.get("score", 0.0), reverse=True )[:top_k] source_citation = {} for idx, item in enumerate(sorted_data, start=1): if not isinstance(item, dict): continue file_info = item.get("file") kb_info = item.get("kb") filename = file_info.get("filename", "") if file_info else "" resource = filename or (kb_info.get("name", "未分组文档") if kb_info else "未分组文档") positions = item.get("positions", []) page_idx = positions[0].get("page_idx", 0) if positions else 0 citation_item = { "index": idx, "id": item.get("id", ""), "text": item.get("content", ""), "resource": resource, "filename": filename, "score": item.get("score", 0.0), "page_idx": page_idx, "file_id": item.get("file_id", ""), "kb_id": kb_info.get("id", "") if kb_info else "", "kb_name": kb_info.get("name", "") if kb_info else "", "positions": positions } if resource not in source_citation: source_citation[resource] = [] source_citation[resource].append(citation_item) return { "success": True, "sourceCitation": source_citation, "total": len(sorted_data), "total_available": len(data) } except Exception as e: return {"success": False, "error": str(e)} @tool async def rag_search_tool(query: str, top_k: Optional[int] = 5) -> str: """ RAG检索工具（供工作流调用） 在知识库中检索与query相关的文档内容 Args: query: 检索查询 top_k: 返回结果数量 Returns: 检索结果文本 """ result = await rag_search(query, top_k=top_k) if not result.get("success", False): return "" source_citation = result.get("sourceCitation", {}) if not source_citation: return "" results_parts = [] for resource, items in source_citation.items(): for item in items: text = item.get("text", "") if text: results_parts.append(f"【{resource}】{text}") results_text = "\n\n".join(results_parts) if results_parts else "" if len(results_text) > 2000: results_text = results_text[:2000] + "..." return results_text @track_function_calls async def file_to_text(file_path: str) -> Dict[str, Any]: """ 文件转文本工具 支持多种文件格式：PDF、Word、TXT、图片等 Args: file_path: 文件路径 Returns: 转换结果，包含文本内容 """ if not file_path or not file_path.strip(): return {"success": False, "error": "文件路径不能为空"} file_path = file_path.strip() if not os.path.exists(file_path): return {"success": False, "error": f"文件不存在: {file_path}"} try: file_ext = os.path.splitext(file_path)[1].lower() if file_ext == ".txt": with open(file_path, "r", encoding="utf-8") as f: text = f.read() return {"success": True, "text": text} elif file_ext == ".pdf": import fitz doc = fitz.open(file_path) text_parts = [] for page in doc: text_parts.append(page.get_text()) text = "\n".join(text_parts) doc.close() return {"success": True, "text": text} elif file_ext in [".docx", ".doc"]: from docx import Document doc = Document(file_path) text_parts = [] for para in doc.paragraphs: text_parts.append(para.text) text = "\n".join(text_parts) return {"success": True, "text": text} elif file_ext in [".jpg", ".jpeg", ".png", ".bmp", ".gif"]: with open(file_path, "rb") as f: image_data = f.read() base64_image = base64.b64encode(image_data).decode("utf-8") base64_url = f"data:image/{file_ext[1:]};base64,{base64_image}" from modelsAPI.model_api import OpenaiAPI text = await OpenaiAPI.vlm_chat( image_path=base64_url, prompt="请描述这张图片的内容" ) return {"success": True, "text": text} else: return {"success": False, "error": f"不支持的文件格式: {file_ext}"} except Exception as e: return {"success": False, "error": f"文件处理失败: {str(e)}"} @tool async def file_to_text_tool(file_path: str) -> str: """ 文件转文本工具（供工作流调用） 支持多种文件格式：PDF、Word、TXT、图片等 Args: file_path: 文件路径 Returns: 文件文本内容 """ result = await file_to_text(file_path) if not result.get("success", False): return f"文件处理失败: {result.get('error', '未知错误')}" return result.get("text", "")