Quickly Build an Agent with Llama-Index

Quickly Build an Agent with Llama-Index

by
Meow!
In the previous article, we used Tongyi Qianwen to create an intelligent customer service agent with four major functions through four system-level prompts. This article will build an upgraded agent based on calling Tongyi Qianwen and combining it with Llama-Index.
First, let’s implement the simplest example using ReActAgent and Functional Tool to create a custom model and stream processing for an agent that solves arithmetic problems.
1. Preparation
1. First, prepare the appetizer, which is basically the same style as the previous article.
import os
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
# Read api_key and other appetizers from environment variables
api_key = os.getenv('DASHSCOPE_API_KEY')
base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"
chat_model = "qwen-plus"
emb_model = "embedding-3"
2. Build a custom LLM.
Here we define a class named OurLLM that inherits from the CustomLLM base class. The main purpose of this class is to interact with the OpenAI API and provide custom language model functionality: corresponding to the addition and multiplication below.
Additionally, we define a method named stream_complete that handles streaming responses. This method allows users to receive the language model’s responses in a streaming manner rather than getting all responses at once, meaning you can see the LLM “thinking and typing” at the same time.
from openai import Client
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: Client = 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 = Client(api_key=self.api_key, base_url=self.base_url)  # Initialize client instance using passed 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. Test if this LLM works. (This part can be commented out later~)
response = llm.stream_complete("Who are you?")
for chunk in response:    print(chunk, end="", flush=True)
If called correctly, the LLM will generate a response like:
I am a large-scale language model from Alibaba Cloud, called Tongyi Qianwen.
2. Implement the Agent
1. Import modules
import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
from llama_index.core.agent import ReActAgent
from llama_index.core.tools import FunctionTool
2. Define tools
Addition tool
def multiply(a: float, b: float) -> float:    """Multiply two numbers and returns the product"""    return a * b
Multiplication tool
def add(a: float, b: float) -> float:    """Add two numbers and returns the sum"""    return a + b
3. Create tool instances and ReActAgent instance in the main function
def main():
    multiply_tool = FunctionTool.from_defaults(fn=multiply)
    add_tool = FunctionTool.from_defaults(fn=add)
4. Input arithmetic questions and run!
    response = agent.chat("What is 200 + (24 * 84) * 21? Use tools to calculate each step")
    print(response)
Its response is as follows:
> Running step ef75a574-3f5f-4b70-88f8-c11febe357f6. Step input: What is 200 + (24 * 84) * 21? Use tools to calculate each step
Thought: The current language of the user is: Chinese. I need to use a tool to help me answer the question.
Action: multiply
Action Input: {'a': 24, 'b': 84}
Observation: 2016
> Running step 8884dd41-d22f-469c-9afd-81613d03855c. Step input: None
Thought: I have the result of the multiplication. Next, I need to multiply this result by 21.
Action: multiply
Action Input: {'a': 2016, 'b': 21}
Observation: 42336
> Running step ec91d672-5fe2-4032-b6c7-5f80eb87b511. Step input: None
Thought: I have the result of the second multiplication. Now, I need to add 200 to this result.
Action: add
Action Input: {'a': 200, 'b': 42336}
Observation: 42536
> Running step 883497ad-b9ce-4732-8c77-2a6f4f470b67. Step input: None
Thought: I can answer without using any more tools. I'll use the user's language to answer.
Answer: 200 + (24 * 84) * 21 equals 42536. 200 + (24 * 84) * 21 equals 42536.

You can see that it prioritized using custom functions add and multiply to complete the calculations in the question, rather than relying on the large model. At the end, it summarized its thinking with “I can answer without using any more tools.” Interesting, right? We can customize part of the processing flow of the Agent, and besides using prompts, we have greater control.

Next, let’s customize a function for weather queries for the agent just created, and provide some fake data for return. (This is to review the custom knowledge points, not to call external tools for queries. Calling external tools will be mentioned in later articles.)

def get_weather(city: str) -> int:    """    Gets the weather temperature of a specified city.
    Args:    city (str): The name or abbreviation of the city.
    Returns:    int: The temperature of the city. Returns 20 for 'NY' (New York),         30 for 'BJ' (Beijing), and -1 for unknown cities.    """
    # Convert the input city to uppercase to handle case-insensitive comparisons    city = city.upper()
    # Check if the city is New York ('NY')    if city == "NY":        return 20  # Return 20°C for New York
    # Check if the city is Beijing ('BJ')    elif city == "BJ":        return 30  # Return 30°C for Beijing
    # If the city is neither 'NY' nor 'BJ', return -1 to indicate unknown city    else:        return -1
weather_tool = FunctionTool.from_defaults(fn=get_weather)

In the main function, unlike the previous text, we added the weather_tool tool object.

agent = ReActAgent.from_tools([multiply_tool, add_tool, weather_tool], llm=llm, verbose=True)

Test the functionality! (Place this before the main function)

response = agent.chat("How is the weather in New York?")
print(response)

Its response is as follows:

> Running step e0002ada-fd52-44b9-9272-0fff7a4cb313. Step input: How is the weather in New York?
Thought: The current language of the user is: Chinese. I need to use a tool to help me answer the question.
Action: get_weather
Action Input: {'city': 'NY'}
Observation: 20
> Running step 81ea310a-3c35-4f92-9ed5-ee662394801b. Step input: None
Thought: I can answer without using any more tools. I'll use the user's language to answer.
Answer: The temperature in New York is 20 degrees.

You can see that the model’s reasoning ability is strong, as it called the function get_weather we defined, and completed the requirement of capitalizing the first letter of the city in the function’s prompt. Finally, it summarized the answer and output the result in natural language.

The ReActAgent realizes the possibility of automatically converting business into code, requiring only an API model to call, suitable for various business scenarios. In addition, LlamaIndex provides open-source tools, for more details, please refer to the official website.

Just a thought

Every agent is like an employee of a company, each with its own duties and unable to complete all tasks. Different agents are like different departments, each performing their roles and collaborating to complete complex tasks. Mastering the usage of different agents allows for combinatorial use to achieve different functionalities and improve control over LLMs.

Imitation, learning, and practical application are essential pathways for acquiring new knowledge. Don’t panic when encountering bugs; others have encountered similar errors. Search for solutions more, seek guidance from peers, and give yourself some patience. You will eventually run it through and learn it~

Quickly Build an Agent with Llama-Index

The rabbit’s mooncake factory (image source: internet)

Every little rabbit is like an agent, working together to make mooncakes.

This experience note is based on the Datawhale project “WowAgent” and contains no advertisements or related interests.

-End-

Quickly Build an Agent with Llama-Index

Leave a Comment