Added agent head

7b_rag.py

@@ -5,6 +5,7 @@ from huggingface_hub import hf_hub_download
 from pynvml import *
 from duckduckgo_search import DDGS
 import re
+import asyncio
 
 # Flag to check if GPU is present
 HAS_GPU = False
@@ -43,10 +44,37 @@ model = RWKV(model=model_path, strategy=MODEL_STRAT)
 from rwkv.utils import PIPELINE, PIPELINE_ARGS
 pipeline = PIPELINE(model, "rwkv_vocab_v20230424")
 
-# Web search function for RAG
+# 理解问题并提取关键词的函数
+async def understanding_question(question: str):
+    # 简单处理：移除常见的问题词，保留关键内容
+    question = question.lower()
+    question = re.sub(r'^(can you|could you|please|tell me about|what is|who is|how to|why is|when did)\s+', '', question)
+    # 返回处理后的问题作为关键词
+    return question
+
+# Web search function for RAG with browser agent HTTP headers
+async def run_duckduckgo_search_tool(question: str):
+    text = await understanding_question(question)
+    
+    keywords = text.split(",")
+    headers = {
+        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:124.0) Gecko/20100101 Firefox/124.0"
+    }
+    
+    results = DDGS(headers=headers).text(keywords[0], max_results=5)
+    print(results)
+    
+    return text
+
+# 修改后的web_search函数，使用run_duckduckgo_search_tool
 def web_search(query, max_results=3):
     try:
-        with DDGS() as ddgs:
+        # 设置浏览器代理HTTP头部
+        headers = {
+            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:124.0) Gecko/20100101 Firefox/124.0"
+        }
+        
+        with DDGS(headers=headers) as ddgs:
             results = list(ddgs.text(query, max_results=max_results))
             if not results:
                 return "No search results found."