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) |