wehub-resource-sync
143 lines
5.2 KiB
Python
143 lines
5.2 KiB
Python
"""知识库示例问题生成工具。"""
|
||
|
||
import json
|
||
import textwrap
|
||
from typing import Any
|
||
|
||
from fastapi import HTTPException
|
||
|
||
from yuxi import config, knowledge_base
|
||
from yuxi.knowledge.factory import KnowledgeBaseFactory
|
||
from yuxi.models import select_model
|
||
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
||
from yuxi.utils import logger
|
||
|
||
SAMPLE_QUESTIONS_SYSTEM_PROMPT = """你是一个专业的知识库问答测试专家。
|
||
|
||
你的任务是根据知识库中的文件列表，生成有价值的测试问题。
|
||
|
||
要求：
|
||
1. 问题要具体、有针对性，基于文件名称和类型推测可能的内容
|
||
2. 问题要涵盖不同方面和难度
|
||
3. 问题要简洁明了，适合用于检索测试
|
||
4. 问题要多样化，包括事实查询、概念解释、操作指导等
|
||
5. 问题长度控制在10-30字之间
|
||
6. 直接返回JSON数组格式，不要其他说明
|
||
|
||
返回格式：
|
||
```json
|
||
{
|
||
"questions": [
|
||
"问题1？",
|
||
"问题2？",
|
||
"问题3？"
|
||
]
|
||
}
|
||
```
|
||
"""
|
||
|
||
|
||
def build_sample_question_file_list(files: dict[str, dict[str, Any]]) -> list[dict[str, str]]:
|
||
return [
|
||
{
|
||
"filename": file_info.get("filename", ""),
|
||
"type": file_info.get("type") or file_info.get("file_type", ""),
|
||
}
|
||
for file_info in files.values()
|
||
]
|
||
|
||
|
||
def build_sample_questions_user_message(db_name: str, files_info: list[dict[str, str]], count: int) -> str:
|
||
files_text = "\n".join([f"- {file_info['filename']} ({file_info['type']})" for file_info in files_info[:20]])
|
||
file_count_text = f"（共{len(files_info)}个文件）" if len(files_info) > 20 else ""
|
||
|
||
return textwrap.dedent(f"""请为知识库\"{db_name}\"生成{count}个测试问题。
|
||
|
||
知识库文件列表{file_count_text}：
|
||
{files_text}
|
||
|
||
请根据这些文件的名称和类型，生成{count}个有价值的测试问题。""")
|
||
|
||
|
||
def parse_sample_questions_content(content: str) -> list[str]:
|
||
if "```json" in content:
|
||
json_start = content.find("```json") + 7
|
||
json_end = content.find("```", json_start)
|
||
if json_end == -1:
|
||
raise ValueError("AI返回的JSON代码块不完整")
|
||
content = content[json_start:json_end].strip()
|
||
elif "```" in content:
|
||
json_start = content.find("```") + 3
|
||
json_end = content.find("```", json_start)
|
||
if json_end == -1:
|
||
raise ValueError("AI返回的代码块不完整")
|
||
content = content[json_start:json_end].strip()
|
||
|
||
questions_data = json.loads(content)
|
||
questions = questions_data.get("questions", []) if isinstance(questions_data, dict) else []
|
||
if not questions or not isinstance(questions, list):
|
||
raise ValueError("AI返回的问题格式不正确")
|
||
return questions
|
||
|
||
|
||
async def generate_database_sample_questions(kb_id: str, count: int = 10) -> dict[str, Any]:
|
||
db_info = await knowledge_base.get_database_info(kb_id, include_files=True)
|
||
if not db_info:
|
||
raise HTTPException(status_code=404, detail=f"知识库 {kb_id} 不存在")
|
||
|
||
kb_type = (db_info.get("kb_type") or "").lower()
|
||
if not KnowledgeBaseFactory.get_kb_class(kb_type).supports_documents:
|
||
raise HTTPException(status_code=400, detail=f"{db_info.get('name') or kb_type} 不支持基于文件生成测试问题")
|
||
|
||
db_name = db_info.get("name", "")
|
||
all_files = db_info.get("files", {})
|
||
if not all_files:
|
||
raise HTTPException(status_code=400, detail="知识库中没有文件")
|
||
|
||
files_info = build_sample_question_file_list(all_files)
|
||
logger.info(f"开始生成知识库问题，知识库: {db_name}, 文件数量: {len(files_info)}, 问题数量: {count}")
|
||
|
||
model = select_model(model_spec=config.default_model)
|
||
messages = [
|
||
{"role": "system", "content": SAMPLE_QUESTIONS_SYSTEM_PROMPT},
|
||
{"role": "user", "content": build_sample_questions_user_message(db_name, files_info, count)},
|
||
]
|
||
response = await model.call(messages, stream=False)
|
||
content = response.content if hasattr(response, "content") else str(response)
|
||
|
||
try:
|
||
questions = parse_sample_questions_content(content)
|
||
except (json.JSONDecodeError, ValueError) as e:
|
||
logger.error(f"AI返回的JSON解析失败: {e}, 原始内容: {content}")
|
||
raise HTTPException(status_code=500, detail=f"AI返回格式错误: {str(e)}") from e
|
||
|
||
logger.info(f"成功生成{len(questions)}个问题")
|
||
|
||
try:
|
||
await KnowledgeBaseRepository().update(kb_id, {"sample_questions": questions})
|
||
logger.info(f"成功保存 {len(questions)} 个问题到知识库 {kb_id}")
|
||
except Exception as save_error:
|
||
logger.error(f"保存问题失败: {save_error}")
|
||
|
||
return {
|
||
"message": "success",
|
||
"questions": questions,
|
||
"count": len(questions),
|
||
"kb_id": kb_id,
|
||
"db_name": db_name,
|
||
}
|
||
|
||
|
||
async def get_database_sample_questions(kb_id: str) -> dict[str, Any]:
|
||
kb = await KnowledgeBaseRepository().get_by_kb_id(kb_id)
|
||
if kb is None:
|
||
raise HTTPException(status_code=404, detail=f"知识库 {kb_id} 不存在")
|
||
|
||
questions = kb.sample_questions or []
|
||
return {
|
||
"message": "success",
|
||
"questions": questions,
|
||
"count": len(questions),
|
||
"kb_id": kb_id,
|
||
}
|