130 lines
4.8 KiB
Python
130 lines
4.8 KiB
Python
|
import requests
|
|||
|
|
import json
|
|||
|
|
import dashscope
|
|||
|
|
from http import HTTPStatus
|
|||
|
|
from typing import List, Dict
|
|||
|
|
|
|||
|
|
# ================= 配置区域 =================
|
|||
|
|
# 1. 设置 DashScope API Key (这里填入你提供的 Key)
|
|||
|
|
dashscope.api_key = "sk-8b091493de594c5e9eb42f12f1cc5805"
|
|||
|
|
|
|||
|
|
# 2. 本地后端地址 (刚才写的 FastAPI)
|
|||
|
|
BACKEND_SEARCH_URL = "http://127.0.0.1:8000/api/v2/search"
|
|||
|
|
|
|||
|
|
# 3. 选择模型 (qwen-turbo, qwen-plus, qwen-max)
|
|||
|
|
MODEL_NAME = dashscope.Generation.Models.qwen_plus
|
|||
|
|
# ===========================================
|
|||
|
|
|
|||
|
|
class WikiBot:
|
|||
|
|
def __init__(self):
|
|||
|
|
self.history = [] #以此保存多轮对话上下文(可选)
|
|||
|
|
|
|||
|
|
def search_knowledge_base(self, query: str, top_k: int = 5) -> List[Dict]:
|
|||
|
|
"""调用本地后端接口检索相关知识"""
|
|||
|
|
try:
|
|||
|
|
payload = {
|
|||
|
|
"query": query,
|
|||
|
|
"limit": top_k
|
|||
|
|
}
|
|||
|
|
# 调用 /api/v2/search，后端会自动做 embedding
|
|||
|
|
resp = requests.post(BACKEND_SEARCH_URL, json=payload)
|
|||
|
|
|
|||
|
|
if resp.status_code == 200:
|
|||
|
|
data = resp.json()
|
|||
|
|
if data.get("code") == 1:
|
|||
|
|
return data.get("data", [])
|
|||
|
|
|
|||
|
|
print(f"[Warning] 检索失败: {resp.text}")
|
|||
|
|
return []
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"[Error] 连接后端失败: {e}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
def build_prompt(self, query: str, context_chunks: List[Dict]) -> str:
|
|||
|
|
"""构建 RAG 提示词"""
|
|||
|
|
|
|||
|
|
if not context_chunks:
|
|||
|
|
return f"用户问题：{query}\n\n当前知识库中没有找到相关信息，请直接告知用户无法回答。"
|
|||
|
|
|
|||
|
|
# 拼接参考资料
|
|||
|
|
context_str = ""
|
|||
|
|
for idx, item in enumerate(context_chunks):
|
|||
|
|
# 这里把 source_url 也带上，方便 AI 引用来源
|
|||
|
|
source = item.get('source_url', '未知来源')
|
|||
|
|
content = item.get('content', '').strip()
|
|||
|
|
context_str += f"【参考资料 {idx+1}】(来源: {source}):\n{content}\n\n"
|
|||
|
|
|
|||
|
|
# 系统提示词 (System Prompt)
|
|||
|
|
prompt = f"""你是一个专业的 Wiki 知识库助手。
|
|||
|
|
请严格根据下方的【参考上下文】来回答用户的【问题】。
|
|||
|
|
|
|||
|
|
要求：
|
|||
|
|
1. 回答要准确、简洁，并整合不同参考资料中的信息。
|
|||
|
|
2. 如果【参考上下文】中包含答案，请用自己的话回答，并在句尾标注来源，例如 [参考资料 1]。
|
|||
|
|
3. 如果【参考上下文】与问题无关或不包含答案，请直接回答：“知识库中暂未收录相关信息”，不要编造答案。
|
|||
|
|
4. 保持回答格式清晰（可以使用 Markdown）。
|
|||
|
|
|
|||
|
|
====== 参考上下文 开始 ======
|
|||
|
|
{context_str}
|
|||
|
|
====== 参考上下文 结束 ======
|
|||
|
|
|
|||
|
|
用户问题：{query}
|
|||
|
|
"""
|
|||
|
|
return prompt
|
|||
|
|
|
|||
|
|
def chat(self, query: str):
|
|||
|
|
"""主对话逻辑"""
|
|||
|
|
print(f"\n🔍 正在检索知识库...")
|
|||
|
|
|
|||
|
|
# 1. 检索
|
|||
|
|
chunks = self.search_knowledge_base(query)
|
|||
|
|
print(f"✅ 找到 {len(chunks)} 条相关资料")
|
|||
|
|
|
|||
|
|
# 2. 构建 Prompt
|
|||
|
|
prompt = self.build_prompt(query, chunks)
|
|||
|
|
|
|||
|
|
# (可选) 调试时打印 prompt 看看给 AI 喂了什么
|
|||
|
|
# print(f"DEBUG PROMPT:\n{prompt}\n")
|
|||
|
|
|
|||
|
|
print("🤖 Wiki助手正在思考...\n" + "-"*30)
|
|||
|
|
|
|||
|
|
# 3. 调用 DashScope 生成 (流式输出)
|
|||
|
|
responses = dashscope.Generation.call(
|
|||
|
|
model=MODEL_NAME,
|
|||
|
|
messages=[
|
|||
|
|
{'role': 'system', 'content': 'You are a helpful assistant.'},
|
|||
|
|
{'role': 'user', 'content': prompt}
|
|||
|
|
],
|
|||
|
|
result_format='message', # 设置输出为 message 格式
|
|||
|
|
stream=True, # 开启流式输出
|
|||
|
|
incremental_output=True # 增量输出
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
full_content = ""
|
|||
|
|
for response in responses:
|
|||
|
|
if response.status_code == HTTPStatus.OK:
|
|||
|
|
text = response.output.choices[0]['message']['content']
|
|||
|
|
full_content += text
|
|||
|
|
print(text, end='', flush=True)
|
|||
|
|
else:
|
|||
|
|
print(f"\nRequest id: {response.request_id}, Status code: {response.status_code}, error code: {response.code}, error message: {response.message}")
|
|||
|
|
|
|||
|
|
print("\n" + "-"*30 + "\n")
|
|||
|
|
|
|||
|
|
# ================= 运行入口 =================
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
bot = WikiBot()
|
|||
|
|
print("✨ Wiki 知识库助手已启动 (输入 'q' 或 'exit' 退出)")
|
|||
|
|
print("⚠️ 请确保后端服务 (main.py) 正在 localhost:8000 运行")
|
|||
|
|
|
|||
|
|
while True:
|
|||
|
|
user_input = input("\n🙋 请输入问题: ").strip()
|
|||
|
|
|
|||
|
|
if user_input.lower() in ['q', 'exit', 'quit']:
|
|||
|
|
print("再见！")
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
if not user_input:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
bot.chat(user_input)
|