From 7bdb6ca5ebf684e65939d3737f1041b909b2b26f Mon Sep 17 00:00:00 2001 From: caoqianming Date: Thu, 21 May 2026 15:31:21 +0800 Subject: [PATCH] =?UTF-8?q?feat(skill):=20documents=20skill=20=E6=8E=A5?= =?UTF-8?q?=E5=86=85=E9=83=A8=E6=9D=90=E6=96=99=E5=AD=A6=E7=A7=91=E7=9F=A5?= =?UTF-8?q?=E8=AF=86=E5=BA=93(document=5Fsearch=20API)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - skills/documents/{SKILL.md, client.py} 4 函数 list_kb / search / download / health - 走 https://ai.ctc-zc.com:8100/api Bearer 认证;env DOCUMENT_SEARCH_API_KEY + DOCUMENT_SEARCH_URL(可覆盖) - search 默认返 md_content(整篇 Markdown 50K-200K 字符级),反模式段约束"只 print 前 300 字"防爆上下文 - smoke 实测后校准 SKILL.md:库实质是 7 个材料学科(胶凝/陶瓷/玻璃/晶体/复合/耐火/检验检测,21W+ 文件)预收的英文学术论文 + 跨语言语义检索(原猜"主语料中文"错了) - 与 research(OpenAlex 全网)互补:documents 已 Markdown 化对 LLM 友好,但仅覆盖材料领域 Co-Authored-By: Claude Opus 4.7 (1M context) --- DOCUMENT_SEARCH_API.md | 283 +++++++++++++++++++++++++++++++++++++ PROGRESS.md | 3 +- RUN.md | 6 +- skills/documents/SKILL.md | 137 ++++++++++++++++++ skills/documents/client.py | 158 +++++++++++++++++++++ 5 files changed, 585 insertions(+), 2 deletions(-) create mode 100644 DOCUMENT_SEARCH_API.md create mode 100644 skills/documents/SKILL.md create mode 100644 skills/documents/client.py diff --git a/DOCUMENT_SEARCH_API.md b/DOCUMENT_SEARCH_API.md new file mode 100644 index 0000000..46434b3 --- /dev/null +++ b/DOCUMENT_SEARCH_API.md @@ -0,0 +1,283 @@ +# 文档搜索API使用说明 + +## 概述 + +本文档介绍了文档搜索API接口，该接口提供文档检索、知识库列表查询和文档下载功能。 + +**基础URL**: `https://ai.ctc-zc.com:8100/api` + +**认证方式**: Bearer Token（除健康检查接口外） +``` +Authorization: Bearer +``` + +## API端点 + +### 1. 搜索文档 + +**接口地址**: `POST /document_search/search` + +**需要认证**: 是 + +**请求头**: +``` +Content-Type: application/json +Authorization: Bearer +``` + +**请求体参数**: + +| 参数名 | 类型 | 必填 | 默认值 | 描述 | +|--------|------|------|--------|------| +| query | string | 是 | - | 搜索查询词 | +| classification_ids | array[int] | 否 | [] | 分类ID列表 | +| kb_names | array[string] | 否 | ["mu_34_1740625285897"] | 知识库名称列表，支持多个知识库 | +| max_documents | integer | 否 | 6 | 最大返回文档数 (1-20) | + +**响应格式**: + +```json +{ + "code": 200, + "msg": "success", + "data": { + "documents": [ + { + "id": 371633, + "kb_name": "mu_34_1740625285897", + "file_name": "example.pdf", + "file_ext": ".pdf", + "file_version": 1, + "document_loader": "PyPDFLoader", + "text_splitter": "MarkdownHeaderTextSplitter", + "create_time": "2025-01-01T00:00:00", + "file_mtime": "2025-01-01T00:00:00", + "file_size": 1024000, + "custom_docs": false, + "docs_count": 50, + "character_count": 100000, + "md_filename": "example.md", + "md_content": "文档的Markdown内容...", + "classification_ids": [1, 2], + "url": "https://ai.ctc-zc.com:8100/api/document_search/download_doc?knowledge_base_name=mu_34_1740625285897&file_name=example.pdf" + } + ], + "total_count": 15, + "query": "混凝土材料性能", + "classification_ids": [1, 2, 3], + "kb_names": ["mu_34_1740625285897", "mu_34_1740625318986"] + } +} +``` + +**响应字段说明**: + +| 字段名 | 类型 | 描述 | +|--------|------|------| +| code | integer | 响应状态码 | +| msg | string | 响应消息 | +| data | object | 响应数据 | +| data.documents | array | 文档列表（直接返回 kb_file 对象） | +| data.total_count | integer | 文档总数 | +| data.query | string | 原始查询词 | +| data.classification_ids | array[int] | 使用的分类IDs | +| data.kb_names | array[string] | 使用的知识库名称列表 | + +**Document (kb_file) 字段说明**: + +| 字段名 | 类型 | 描述 | +|--------|------|------| +| id | integer | 文档ID | +| kb_name | string | 知识库名称 | +| file_name | string | 原始文件名 | +| file_ext | string | 文件扩展名 | +| file_version | integer | 文件版本 | +| document_loader | string | 文档加载器名称 | +| text_splitter | string | 文本分割器名称 | +| create_time | datetime | 创建时间 | +| file_mtime | datetime | 文件修改时间 | +| file_size | integer | 文件大小（字节） | +| custom_docs | boolean | 是否自定义文档 | +| docs_count | integer | 文档块数量 | +| character_count | integer | 字符总数 | +| md_filename | string | Markdown文件名 | +| md_content | string | Markdown内容 | +| classification_ids | array[int] | 分类ID列表 | +| url | string | 下载URL（自动拼接） | + +--- + +### 2. 下载知识库文档 + +**接口地址**: `GET /document_search/download_doc` + +**需要认证**: 是 + +**请求头**: +``` +Authorization: Bearer +``` + +**查询参数**: + +| 参数名 | 类型 | 必填 | 默认值 | 描述 | +|--------|------|------|--------|------| +| knowledge_base_name | string | 是 | - | 知识库名称 | +| file_name | string | 是 | - | 文件名称（可以是原始文件名或Markdown文件名） | +| preview | boolean | 否 | false | 是否在线预览（true: 浏览器内预览；false: 下载） | + +**功能说明**: +- 支持通过原始文件名（如 `example.pdf`）或 Markdown 文件名（如 `example.md`）下载 +- 如果直接查找失败，系统会自动从数据库查询文件详情并尝试使用原始文件名 +- 确保文件路径安全，防止目录遍历攻击 + +**请求示例**: +``` +GET /document_search/download_doc?knowledge_base_name=mu_34_1740625285897&file_name=example.pdf&preview=false +``` + +**响应**: +- 成功: 返回文件流（FileResponse） +- 失败: 返回JSON错误信息 + +**错误响应示例**: +```json +{ + "detail": "文件不存在: example.pdf" +} +``` + +--- + +### 3. 获取知识库列表 + +**接口地址**: `GET /document_search/list_knowledge_bases` + +**需要认证**: 是 + +**请求头**: +``` +Authorization: Bearer +``` + +**功能说明**: +- 只返回在ID映射关系中存在的知识库（分类ID 1-7对应的知识库） +- 过滤掉其他无效的知识库 +- 返回有效的kb_name列表供参考 + +**响应格式**: + +```json +{ + "code": 200, + "msg": "success", + "data": { + "knowledge_bases": [ + { + "id": 1, + "kb_name": "mu_34_1740625285897", + "ch_name": "材料分类1", + "kb_info": "知识库描述", + "customer_id": 34, + "classification_type": 0, + "create_time": "2025-01-01T00:00:00", + "file_count": 100 + } + ], + "total_count": 7, + "valid_kb_names": [ + "mu_34_1740625285897", + "mu_34_1740625318986", + "mu_34_1740625346474", + "mu_34_1740625303475", + "mu_34_1740625365079", + "mu_34_1740625355308", + "mu_34_1740625376621" + ] + } +} +``` + +**响应字段说明**: + +| 字段名 | 类型 | 描述 | +|--------|------|------| +| code | integer | 响应状态码 | +| msg | string | 响应消息 | +| data | object | 响应数据 | +| data.knowledge_bases | array | 知识库列表 | +| data.total_count | integer | 知识库总数 | +| data.valid_kb_names | array[string] | 有效的知识库名称列表 | + +**KnowledgeBase 字段说明**: + +| 字段名 | 类型 | 描述 | +|--------|------|------| +| id | integer | 知识库ID | +| kb_name | string | 知识库名称 | +| ch_name | string | 中文名称 | +| kb_info | string | 知识库描述 | +| customer_id | integer | 客户ID | +| classification_type | integer | 分类类型 | +| create_time | string | 创建时间（ISO格式） | +| file_count | integer | 文件数量 | + +--- + +### 4. 健康检查 + +**接口地址**: `GET /document_search/health` + +**需要认证**: 否（公开访问） + +**响应格式**: + +```json +{ + "code": 200, + "msg": "success", + "data": { + "status": "healthy", + "service": "document_search_api" + } +} +``` +``` + +## 错误处理 + +API可能返回以下错误： + +### HTTP状态码 + +| 状态码 | 说明 | +|--------|------| +| 200 | 请求成功 | +| 400 | 请求参数格式错误 | +| 401 | 缺少认证信息或认证格式错误 | +| 403 | API密钥无效或非法访问 | +| 404 | 资源不存在（知识库或文件未找到） | +| 500 | 服务器内部错误 | + +### 错误响应格式 + +```json +{ + "detail": "错误描述信息" +} +``` + +### 常见错误信息 + +| 错误信息 | 说明 | +|---------|------| +| Missing authorization header | 缺少Authorization头 | +| Invalid authorization header format. Expected 'Bearer ' | Authorization头格式错误 | +| API key is required | API密钥为空 | +| Invalid API key for document search service | API密钥无效 | +| Don't attack me | 非法的知识库名称 | +| 未找到知识库 {name} | 知识库不存在 | +| 文件不存在: {filename} | 文件不存在 | +| 文档搜索失败: {error} | 搜索过程中发生错误 | +| 文件下载失败: {error} | 下载过程中发生错误 | + diff --git a/PROGRESS.md b/PROGRESS.md index 77c4853..b64f482 100644 --- a/PROGRESS.md +++ b/PROGRESS.md @@ -2,7 +2,7 @@ > 配合 `DESIGN.md`。本文件只记 phase 状态、决策偏差、文件量、下一步。每条 1-2 句:做了啥 + 关键判断;细节查 `git log` / `git diff` / `DESIGN §7.9`。 -最后更新:2026-05-21(dev SPA SSE 客户端重连 + 后端 stream_events 非活跃态立即吐 done) +最后更新:2026-05-21(新增 documents skill 接 ai.ctc-zc.com:8100 内部知识库 API) --- @@ -23,6 +23,7 @@ ### 2026-05-21 +- **新增 documents skill(内部材料学科知识库 document_search API)**:`skills/documents/{SKILL.md, client.py}`,四函数 `list_kb / search / download / health`;走 `https://ai.ctc-zc.com:8100/api` Bearer 认证,env `DOCUMENT_SEARCH_API_KEY` + `DOCUMENT_SEARCH_URL`(可覆盖);search 默认返 `md_content`(整篇 Markdown 50K-200K 字符级),SKILL.md 反模式约束"只 print 前 300 字"防爆上下文;smoke 验证发现库实质是 7 个材料学科预收的英文学术论文(胶凝/陶瓷/玻璃/晶体/复合/耐火/检验检测,21W+ 文件)+ 跨语言语义检索,SKILL.md 据此校准(原写"主语料中文"是错的);与 research(OpenAlex)互补,documents 已 Markdown 化对 LLM 更友好,但仅覆盖材料领域。 - **dev SPA SSE 客户端重连(覆盖 --reload 抖动)**:`fetchSse` 拆出 `consumeSseStream` + 包重连壳(1s/2s/4s 退避,最多 3 次);reader EOF 未见 done/error 算异常关流触发重连;后端 `stream_events` 入口检 `tasks.run_status`,非 running/cancelling 立即吐 done 关流(否则进程重启后新 broker 内存空,客户端会无限挂 ping)。3 次仍失败 → 卡片末尾红色"连接已断开,请重发"。断开期间 LLM delta 丢失,接受。 - **research skill 三次迭代 fetch_pdf 改走静态直链**:`fetch_pdf` 跟 `fetch_xml` 同范式,从 `paper["pdf_url"]` 流式下载,绕开 paper_pdf_view 路径 bug(disk 路径计算错);smoke 5/5 PASS。 - **research skill 二次迭代 list 端点加 pdf_url / xml_url 直链 + 新增 fetch_xml + pg_trgm GIN 索引**:serializer 后端拼直链(避免 LLM 拿 stale URL),`0006_pg_trgm` 给 title/first_author/institution 加 GIN 把 `?search=xxx` 从 30s timeout 降到几十 ms;SKILL.md 加"XML 优先 PDF"原则(XML 已结构化免 OCR)。 diff --git a/RUN.md b/RUN.md index e096fa4..0f9d472 100644 --- a/RUN.md +++ b/RUN.md @@ -2,7 +2,7 @@ > 怎么把 zcbot 跑起来。env / 常用命令 / 故障兜底。设计看 `DESIGN.md`,进度看 `PROGRESS.md`。 -最后更新:2026-05-21(dev SPA SSE 客户端重连 3 次退避;`/v1/tasks/{id}/events` 非活跃 task 立即吐 done) +最后更新:2026-05-21(新增 documents skill 的 env:`DOCUMENT_SEARCH_API_KEY` / `DOCUMENT_SEARCH_URL`) --- @@ -16,6 +16,10 @@ ZHIPUAI_API_KEY=... # 豆包(火山方舟)图像/视频生成:可选。设了就挂上 seedream tool(0.22 元/张);未设 tool 不出现 ARK_API_KEY=... + # documents skill(内部知识库 document_search API):可选。设了 documents skill 才能用,未设调用立即抛 RuntimeError + DOCUMENT_SEARCH_API_KEY=... + # 可选:覆盖默认 base_url(默认 https://ai.ctc-zc.com:8100/api) + # DOCUMENT_SEARCH_URL=https://ai.ctc-zc.com:8100/api ZCBOT_DB_URL=postgresql://user:pass@host:5432/zcbot # main.py web 必填(probe/db/user 不验) PLATFORM_KEY=<≥16 字符随机串,platform 机器对机器入口校验> diff --git a/skills/documents/SKILL.md b/skills/documents/SKILL.md new file mode 100644 index 0000000..ac9dc1d --- /dev/null +++ b/skills/documents/SKILL.md @@ -0,0 +1,137 @@ +--- +name: documents +description: 查内部材料学科知识库(document_search API,7 个学科:胶凝 / 陶瓷 / 玻璃 / 晶体 / 复合 / 耐火 / 检验检测,21W+ 英文学术论文 Markdown 化,跨语言语义检索)。用户找材料领域文献、特定学科论文、材料性能数据时使用;与 research(OpenAlex 外部库)互补,可并用 / 同时试。 +--- + +# Documents + +部署在 `https://ai.ctc-zc.com:8100/api` 的文档检索 API。后端按 `kb_name` 分库存储 7 个材料学科库(中文命名:胶凝 / 陶瓷基 / 玻璃基 / 晶体材料 / 复合材料 / 耐火材料 / 检验检测,共 21W+ 文件),**文档主体是英文学术论文**(Elsevier 期刊为主,DOI 前缀文件名),每个文档带 `md_content`(整篇 Markdown,LLM 友好)+ 可选的原 PDF 下载。**API 后端有跨语言语义检索**,中英文 query 都能命中英文文档。本 skill 给四个 helper(`list_kb` / `search` / `download` / `health`),用 `run_python` 调用,**不要**自己 `httpx` 裸调。 + +## 何时用 + +- 用户要查材料领域文献(7 个学科:胶凝 / 陶瓷 / 玻璃 / 晶体 / 复合 / 耐火 / 检验检测) +- 用户要查特定材料性能 / 工艺数据(实验数据 / 表征结果 / 公式 / 表格) +- 写申报书 / 方案 / 报告需要"国内外现状"段,本库 Markdown 化的论文比 research 拿到的裸 PDF/XML 更直接可用 +- **跟 research 并列**:都是文献检索 —— research 走 OpenAlex 元数据 + Sci-Hub PDF,搜全网;documents 是本地预收的材料学科子集,**已转 Markdown**(LLM 直接读,免 OCR / XML 解析)+ **跨语言语义检索**(中文 query 也能命中英文论文)。找材料类文献优先 documents,找其他学科或要 DOI 走 research,**两者命中不重叠时可并用** + +## 何时不用 + +- 用户只问通识(直接答) +- 用户已经给了具体内部文档路径(直接读,不要二次校验) + +## 准备 + +```python +from skills.documents.client import list_kb, search, download, health +``` + +(import 路径由 `run_python` 注入的 `PYTHONPATH` 提供,直接写就行) + +API key 走 env `DOCUMENT_SEARCH_API_KEY`,未设会抛 `RuntimeError`。base_url 走 env `DOCUMENT_SEARCH_URL`(默认 `https://ai.ctc-zc.com:8100/api`)。 + +## 四个函数 + +### `list_kb() -> list[dict]` + +列所有有效知识库(分类 1-7)。每条含 `id` / `kb_name` / `ch_name` / `kb_info` / `file_count` 等。 + +```python +kbs = list_kb() +for kb in kbs: + print(kb["id"], kb["kb_name"], kb["ch_name"], kb["file_count"]) +``` + +**用途**:用户没指定库 → 先 `list_kb()` 看有哪些库(中文名 `ch_name` 看分类),再选 `kb_names` / `classification_ids` 缩窄 search 范围。 + +### `search(query, kb_names=None, classification_ids=None, max_documents=6) -> list[dict]` + +搜文档,返回精简列表,每条带 **`md_content`**(整篇 Markdown 文本)。 + +- `query`:搜索词。**中英文均可** —— 文档主体是英文学术论文,但 API 后端有跨语言语义检索;复杂技术术语用**英文**更精准(`cement hydration` > `水泥水化`),日常概念中文 OK +- `kb_names`:知识库白名单(从 `list_kb()` 选);`None` 走 server 默认(单库 `mu_34_1740625285897` 胶凝)。**多库联查就显式传**,如 `kb_names=["mu_34_1740625285897", "mu_34_1740625303475"]` +- `classification_ids`:分类 ID 白名单(1-7,对应 7 个学科库);`None` 不过滤 +- `max_documents`:1-20,默认 6 + +**学科库 → kb_name 速查**(`list_kb()` 拿全量,这里只列常用): + +| 学科 | kb_name | +|---|---| +| 胶凝材料(水泥 / 混凝土 / 砂浆) | `mu_34_1740625285897` | +| 陶瓷基材料 | `mu_34_1740625303475` | +| 玻璃基材料 | `mu_34_1740625318986` | +| 晶体材料 | `mu_34_1740625346474` | +| 复合材料 | `mu_34_1740625355308` | +| 耐火材料 | `mu_34_1740625365079` | +| 检验检测 | `mu_34_1740625376621` | + +```python +# 全库搜(走 server 默认单库:胶凝) +docs = search(query="cement hydration", max_documents=10) +for d in docs: + print(d["file_name"], d["character_count"]) + # 只看前 300 字判断切题 ——— md_content 整体动辄几十 K(实测命中文档常 50K-200K 字符) + print((d["md_content"] or "")[:300]) +``` + +```python +# 跨多个学科库联查(如同时找胶凝 + 陶瓷) +docs = search( + query="concrete carbonation", + kb_names=["mu_34_1740625285897", "mu_34_1740625303475"], + max_documents=6, +) +``` + +### `download(file_name, kb_name, working_dir, preview=False) -> str` + +下载原始文档(PDF / Word / ...)到 `/documents/`,返回相对路径。已存在跳过下载直接复用。`file_name` 支持原始文件名(`example.pdf`)或 Markdown 名(`example.md`),server 自动回退。 + +```python +rel = download( + file_name="材料性能手册.pdf", + kb_name="mu_34_1740625285897", + working_dir=r"D:/projects/zcbot/workspace/users//", +) +# rel == "documents/材料性能手册.pdf" +``` + +### `health() -> dict` + +健康检查,公开端点(无需 API key)。主要给 smoke / 排障用:`{"status": "healthy", "service": "document_search_api"}`。 + +## 标准工作流 + +1. **(可选)`list_kb()`** —— 用户没指定库 / 不确定分类时看一下有哪些 +2. **`search(query=..., max_documents=6)`** —— 中英文均可,专业技术术语优先英文 +3. **看返回**: + - 用 `file_name + character_count + md_content[:300]` 判断切题 + - 切题 → 直接用 `md_content` 给 LLM 引用(已结构化 Markdown,不需要再下载原件) + - 需要看图表 / 表格原貌 / 给用户附件 → `download(file_name, kb_name, working_dir)` 拿原文档,然后用主 agent 的 `read` 工具读(zcbot 已内置 PDF/Word 文本抽取) +4. **写产出**:把 md_content 关键段落引到申报书 / 方案里,标注来源文件名 + +## md_content 优先 vs 原件下载 + +- **绝大多数场景用 md_content 就够** —— API 已把原文档转成 Markdown,LLM 直接读,无需 OCR 或 PDF 抽取 +- **仅以下场景下载原件**: + - 用户明说"要原文件"(给客户附件 / 存档 / 引用页码) + - md_content 里图表 / 表格 / 公式信息丢失(Markdown 转换无损不了) + - 文档过大(`character_count` > 10 万),想用 PDF reader 跳页抽取局部 + +## 错误处理 + +- `DOCUMENT_SEARCH_API_KEY` 未设:`RuntimeError`(client 启动时立即报,而不是裸 401) +- 401 / 403 `Invalid API key`:`httpx.HTTPStatusError` —— key 错或失效,告诉用户检查 env +- 404 `未找到知识库`:`kb_names` 拼写错或库已下线,改 `list_kb()` 看当前有效列表 +- 404 `文件不存在: xxx`:`download` 时常见,可能 server 侧文件丢失或 `file_name` 拼写错 +- search 命中 0 条:同义词 / 切换中英文 / 缩短 query / 放宽 `classification_ids` 再试 2-3 次,还是 0 条就明确告诉用户"本库没覆盖,改走 research 或换关键词",**不要凭训练数据脑补文献** +- 网络超时 / server 不可达:`httpx.ConnectError` / `httpx.TimeoutException` —— 告诉用户"document_search 暂时连不上",不要重试堆栈刷屏 + +## 反模式 + +- 用 `httpx` / `requests` 裸调 API(走 helper,免得 base_url / auth / 字段名漂移时四处改) +- `search(max_documents=20)` 一次拉满后 print 全部 `md_content`(单条就可能几十 K,20 条直接爆 LLM 上下文)—— 只 print `file_name + character_count + md_content[:300]`,要看全文用 `docs[i]['md_content']` +- 看到 md_content 切题还 `download` 一遍原件(md_content 已是 LLM 友好的 Markdown,大多数引用场景够用) +- 凭 `ch_name`("胶凝材料学科知识库")就以为 query 要用中文 —— 文档主体是英文,复杂术语用英文更精准 +- 编造 file_name / kb_name —— 不在 `list_kb()` / `search` 返回里就**明确告诉用户"未命中"**,不要瞎传 ID +- 把 `download` 返回的相对路径当绝对路径用(它是相对 `working_dir` 的) +- 不在合适的 task working_dir 里 `download`(原文档应该落到 task 目录,不要污染 repo) diff --git a/skills/documents/client.py b/skills/documents/client.py new file mode 100644 index 0000000..567c3ae --- /dev/null +++ b/skills/documents/client.py @@ -0,0 +1,158 @@ +"""document_search API 客户端 helper。 + +base_url / api_key 走 env: + DOCUMENT_SEARCH_URL 默认 https://ai.ctc-zc.com:8100/api + DOCUMENT_SEARCH_API_KEY 必填,缺失时调用立即抛 RuntimeError(而不是裸 401) +""" +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any, Optional + +import httpx + +_BASE_URL = os.environ.get( + "DOCUMENT_SEARCH_URL", "https://ai.ctc-zc.com:8100/api" +).rstrip("/") +_API = f"{_BASE_URL}/document_search" +_TIMEOUT = 30.0 +_DOWNLOAD_TIMEOUT = 120.0 + +# search 返回字段(剥掉项目里不常用的 file_version / document_loader / text_splitter / custom_docs +# 等,但保留 md_content —— 这是接口最大价值) +_LIST_FIELDS = ( + "id", + "kb_name", + "file_name", + "file_ext", + "create_time", + "file_mtime", + "file_size", + "docs_count", + "character_count", + "md_filename", + "md_content", + "classification_ids", + "url", +) + + +def _api_key() -> str: + key = os.environ.get("DOCUMENT_SEARCH_API_KEY", "").strip() + if not key: + raise RuntimeError( + "DOCUMENT_SEARCH_API_KEY env 未设置 —— 配置后再调 documents skill" + ) + return key + + +def _auth_headers(extra: Optional[dict] = None) -> dict: + h = {"Authorization": f"Bearer {_api_key()}"} + if extra: + h.update(extra) + return h + + +def _safe_name(name: str) -> str: + # 防目录穿越;保留扩展名 + return name.replace("/", "_").replace("\\", "_").replace("..", "_") + + +def list_kb() -> list[dict]: + """列出所有有效知识库(对应 GET /list_knowledge_bases)。 + + 返回每条含 id / kb_name / ch_name / kb_info / customer_id / classification_type + / create_time / file_count。只返回 ID 映射里有效的(分类 1-7)。 + """ + r = httpx.get(f"{_API}/list_knowledge_bases", headers=_auth_headers(), timeout=_TIMEOUT) + r.raise_for_status() + payload = r.json() + data = payload.get("data") or {} + return list(data.get("knowledge_bases") or []) + + +def search( + query: str, + kb_names: Optional[list[str]] = None, + classification_ids: Optional[list[int]] = None, + max_documents: int = 6, +) -> list[dict]: + """搜文档(对应 POST /search)。返回精简列表,每条带 md_content(整篇 Markdown)。 + + query: 搜索词 + kb_names: 知识库白名单(从 list_kb() 选);None 走 server 默认值 + classification_ids: 分类 ID 白名单(1-7);None 不过滤 + max_documents: 1-20,默认 6 + + SKILL.md 反模式:不要 print 整个 md_content(动辄几十 K),只打前 200-400 字判断切题。 + """ + if not query: + raise ValueError("query 不可为空") + if max_documents < 1: + max_documents = 1 + if max_documents > 20: + max_documents = 20 + body: dict[str, Any] = {"query": query, "max_documents": max_documents} + if kb_names: + body["kb_names"] = kb_names + if classification_ids: + body["classification_ids"] = classification_ids + r = httpx.post( + f"{_API}/search", + headers=_auth_headers({"Content-Type": "application/json"}), + json=body, + timeout=_TIMEOUT, + ) + r.raise_for_status() + payload = r.json() + data = payload.get("data") or {} + docs = data.get("documents") or [] + return [{k: d.get(k) for k in _LIST_FIELDS} for d in docs] + + +def download( + file_name: str, + kb_name: str, + working_dir: str, + preview: bool = False, +) -> str: + """下载原始文档到 /documents/,返回相对路径。 + + file_name 支持原始文件名(example.pdf)或 Markdown 名(example.md)—— server 会 + 回退查 DB 拿原始名。已存在跳过下载直接复用。 + + preview=True 给浏览器内预览 header(Content-Disposition: inline),通常 agent + 不需要(我们要落盘读取),保留参数透传。 + """ + if not file_name or not kb_name: + raise ValueError("file_name / kb_name 不可为空") + rel = f"documents/{_safe_name(file_name)}" + dest = Path(working_dir) / rel + if dest.exists() and dest.stat().st_size > 0: + return rel + dest.parent.mkdir(parents=True, exist_ok=True) + params = { + "knowledge_base_name": kb_name, + "file_name": file_name, + "preview": "true" if preview else "false", + } + with httpx.stream( + "GET", + f"{_API}/download_doc", + headers=_auth_headers(), + params=params, + timeout=_DOWNLOAD_TIMEOUT, + ) as resp: + resp.raise_for_status() + with open(dest, "wb") as f: + for chunk in resp.iter_bytes(chunk_size=64 * 1024): + f.write(chunk) + return rel + + +def health() -> dict: + """健康检查(公开,不需要认证)。""" + r = httpx.get(f"{_API}/health", timeout=_TIMEOUT) + r.raise_for_status() + return r.json().get("data") or {}