Learn AI Agent Development in 11 Days: Day 6 - Integrating Web Search

Learn AI Agent Development in 11 Days: Day 6 – Integrating Web Search

by
In previous days, we introduced various capabilities for the Agent, including integration with local knowledge bases, database queries, and simple functions. The details are as follows:
  • Learn AI Agent Development in 11 Days, Day 5: Integrating RAG with Agent.
  • Learn AI Agent Development in 11 Days, Day 4: Creating Agent with Llama-index.
  • Learn AI Agent Development in 11 Days, Day 3: Implementing Grading Agent.
  • Learn AI Agent Development in 11 Days, Day 2: Building a Basic Agent from Scratch.
  • Learn AI Agent Development in 11 Days, Day 1: Choosing the Agent Framework – wow-agent.
  • Learn AI Agent Development in 11 Days, Day 0: What is an Agent.
Today, we will attempt to integrate a search engine into the Agent.
1.First, prepare various keys and model names.
import os
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
# Read api_key from environment variables
api_key = os.getenv('WOWRAG_API_KEY')
base_url = "http://114.114.114.114:8008/v1"
chat_model = "glm-4-flash"
emb_model = "embedding-3"
2.Next, let’s build the LLM: any usable LLM will do, we will customize one here.
from openai import OpenAI
from pydantic import Field  # Import Field for defining field metadata in Pydantic models
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

# Define OurLLM class, inheriting from CustomLLM base class
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)  # Explicitly declare client field

    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)  # Initialize client instance with provided api_key and base_url

    @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)
3.Let’s test if this LLM works?
response = llm.stream_complete("Who are you?")
for chunk in response:
    print(chunk, end="", flush=True)

I am an AI assistant named ChatGLM, developed based on a language model jointly trained by Tsinghua University Laboratory and Zhizhu Company in 2024. My task is to provide appropriate responses and support for user questions and requests.

4. Integrating the Search Engine: Using the API service provided by Bochaai.
from llama_index.core.tools import FunctionTool
import requests
# Make sure to fill in BOCHA_API_KEY in the .env file.
BOCHA_API_KEY = os.getenv('BOCHA_API_KEY')

# Define Bocha Web Search tool
def bocha_web_search_tool(query: str, count: int = 8) -> str:
    """
    Use Bocha Web Search API for online searching, returning search results as a string.
    Parameters:
    - query: Search keyword
    - count: Number of search results to return
    Returns:
    - Search results in string form
    """
    url = 'https://api.bochaai.com/v1/web-search'
    headers = {
        'Authorization': f'Bearer {BOCHA_API_KEY}',  # Replace with your API key
        'Content-Type': 'application/json'
    }
    data = {
        "query": query,
        "freshness": "noLimit", # Search time range, e.g., "oneDay", "oneWeek", "oneMonth", "oneYear", "noLimit"
        "summary": True, # Whether to return long text summary
        "count": count
    }
    response = requests.post(url, headers=headers, json=data)
    if response.status_code == 200:
        # Formatted search result text returned to the large model
        # You can customize the processing of Bocha's search results
        return str(response.json())
    else:
        raise Exception(f"API request failed, status code: {response.status_code}, error message: {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)
5. Let’s test if it’s usable
# Test case
query = "What are the main points of Alibaba's ESG report for 2024?"
response = agent.chat(f"Please help me search for the following content: {query}")
print(response)
Learn AI Agent Development in 11 Days: Day 6 - Integrating Web Search
Conclusion: Based on the above, we have implemented the functionality of integrating a search engine API service into the Agent, giving it web search capabilities.

Leave a Comment