【LLM-agent】(task4)搜索引擎Agent

note

• 新增工具：搜索引擎Agent

一、搜索引擎Agent

import os
from dotenv import load_dotenv# 加载环境变量
load_dotenv()
# 初始化变量
base_url = None
chat_model = None
api_key = None# 使用with语句打开文件，确保文件使用完毕后自动关闭
env_path = "/Users/guomiansheng/Desktop/LLM/llm_app/wow-agent/.env.txt"
with open(env_path, 'r') as file:# 逐行读取文件for line in file:# 移除字符串头尾的空白字符（包括'\n'）line = line.strip()# 检查并解析变量if "base_url" in line:base_url = line.split('=', 1)[1].strip().strip('"')elif "chat_model" in line:chat_model = line.split('=', 1)[1].strip().strip('"')elif "ZHIPU_API_KEY" in line:api_key = line.split('=', 1)[1].strip().strip('"')elif "BOCHA_API_KEY" in line:BOCHA_API_KEY = line.split('=', 1)[1].strip().strip('"')# 打印变量以验证
print(f"base_url: {base_url}")
print(f"chat_model: {chat_model}")
print(f"ZHIPU_API_KEY: {api_key}")from openai import OpenAI
client = OpenAI(api_key = api_key,base_url = base_url
)
print(client)def get_completion(prompt):response = client.chat.completions.create(model="glm-4-flash",  # 填写需要调用的模型名称messages=[{"role": "user", "content": prompt},],)return response.choices[0].message.content# 一、定义上个task的llm
from openai import OpenAI
from pydantic import Field  # 导入Field，用于Pydantic模型中定义字段的元数据
from llama_index.core.llms import (CustomLLM,CompletionResponse,LLMMetadata,
)
from llama_index.core.embeddings import BaseEmbedding
from llama_index.core.llms.callbacks import llm_completion_callback
from typing import List, Any, Generator# 定义OurLLM类，继承自CustomLLM基类
class OurLLM(CustomLLM):api_key: str = Field(default=api_key)base_url: str = Field(default=base_url)model_name: str = Field(default=chat_model)client: OpenAI = Field(default=None, exclude=True)  # 显式声明 client 字段def __init__(self, api_key: str, base_url: str, model_name: str = chat_model, **data: Any):super().__init__(**data)self.api_key = api_keyself.base_url = base_urlself.model_name = model_nameself.client = OpenAI(api_key=self.api_key, base_url=self.base_url)  # 使用传入的api_key和base_url初始化 client 实例@propertydef metadata(self) -> LLMMetadata:"""Get LLM metadata."""return LLMMetadata(model_name=self.model_name,)@llm_completion_callback()def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:response = self.client.chat.completions.create(model=self.model_name, messages=[{"role": "user", "content": prompt}])if hasattr(response, 'choices') and len(response.choices) > 0:response_text = response.choices[0].message.contentreturn CompletionResponse(text=response_text)else:raise Exception(f"Unexpected response format: {response}")@llm_completion_callback()def stream_complete(self, prompt: str, **kwargs: Any) -> Generator[CompletionResponse, None, None]:response = self.client.chat.completions.create(model=self.model_name,messages=[{"role": "user", "content": prompt}],stream=True)try:for chunk in response:chunk_message = chunk.choices[0].deltaif not chunk_message.content:continuecontent = chunk_message.contentyield CompletionResponse(text=content, delta=content)except Exception as e:raise Exception(f"Unexpected response format: {e}")llm = OurLLM(api_key=api_key, base_url=base_url, model_name=chat_model)
# print(llm)
# 测试模型是否能正常回答
response = llm.stream_complete("你是谁？")
for chunk in response:print(chunk, end="", flush=True)# 二、搜索工具
from llama_index.core.tools import FunctionTool
import requests
# 需要先把BOCHA_API_KEY填写到.env文件中去。
# BOCHA_API_KEY = os.getenv('BOCHA_API_KEY')# 定义Bocha Web Search工具
def bocha_web_search_tool(query: str, count: int = 8) -> str:"""使用Bocha Web Search API进行联网搜索，返回搜索结果的字符串。参数:- query: 搜索关键词- count: 返回的搜索结果数量返回:- 搜索结果的字符串形式"""url = 'https://api.bochaai.com/v1/web-search'headers = {'Authorization': f'Bearer {BOCHA_API_KEY}',  # 请替换为你的API密钥'Content-Type': 'application/json'}data = {"query": query,"freshness": "noLimit", # 搜索的时间范围，例如 "oneDay", "oneWeek", "oneMonth", "oneYear", "noLimit""summary": True, # 是否返回长文本摘要总结"count": count}response = requests.post(url, headers=headers, json=data)if response.status_code == 200:# 返回给大模型的格式化的搜索结果文本# 可以自己对博查的搜索结果进行自定义处理return str(response.json())else:raise Exception(f"API请求失败，状态码: {response.status_code}, 错误信息: {response.text}")search_tool = FunctionTool.from_defaults(fn=bocha_web_search_tool)
from llama_index.core.agent import ReActAgent
agent = ReActAgent.from_tools([search_tool], llm=llm, verbose=True, max_iterations=10)# 测试用例
query = "阿里巴巴2024年的ESG报告主要讲了哪些内容？"
response = agent.chat(f"请帮我搜索以下内容：{query}")
print(response)

Reference

[1] https://github.com/datawhalechina/wow-agent
[2] https://www.datawhale.cn/learn/summary/86
[3] https://open.bochaai.com/
[4] https://github.com/run-llama/llama_index/issues/14843
[5] 官方文档：https://docs.cloud.llamaindex.ai/