增加测试脚本

2025-12-30 16:19:58 +08:00
parent 4d35626b90
commit 8972246445
3 changed files with 143 additions and 49 deletions
--- a/scripts/init_db.py
+++ b/scripts/init_db.py
@@ -1,8 +0,0 @@
-import random
-
-# 生成1536 8位随机向量
-def generate_random_vector(dim=1536):
-    return [round(random.uniform(-1, 1), 8) for _ in range(dim)]
-
-data = [generate_random_vector() for _ in range(1000)]
-print(data[0])
--- a/scripts/rob.py
+++ b/scripts/rob.py
@@ -0,0 +1,130 @@
+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)