【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_key
        self.base_url = base_url
        self.model_name = model_name
        self.client = OpenAI(api_key=self.api_key, base_url=self.base_url)  # 使用传入的api_key和base_url初始化 client 实例

    @property
    def 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.content
            return 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].delta
                if not chunk_message.content:
                    continue
                content = chunk_message.content
                yield 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/