Using LlamaIndex to Create Custom Agent Functions

Using LlamaIndex to Create Custom Agent Functions

by

Overview

This article introduces how to use LlamaIndex to write your own Agent handling functions. Note that this article uses a locally deployed LLM supported by Ollama for practical implementation, rather than remotely calling the OpenAI API.

The goal of this article is to save the output content to a PDF file and then stop the entire process.

Implementation Steps:

(1) Define the Agent function: Define a regular function that can implement any functionality you need.

(2) Wrap the function with FunctionTool: Note that the name of the wrapped function needs to be the same as the defined function name.

(3) Create an LLM object: Here I am using Ollama to define an LLM object. Different LLM models can be used here, and the results may vary. You can set the value of the model parameter based on your needs to specify the model name.

(4) Create a ReActAgent object to create an agent object.

(5) Use agent.chat to converse with the agent, allowing the agent and LLM to complete the task.

Complete Code

from llama_index.llms.ollama import Ollama
from llama_index.core.agent import ReActAgent
from llama_index.core.tools import FunctionTool
from fpdf import FPDF  # Using fpdf library to generate PDF files
from typing import Any

# Define a tool function to write the response content to PDF
def save_response_to_pdf(response: str, **kwargs) -> str:
    """Save the response to a PDF file."""
    # Initialize PDF
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)
    
    # Write the response content to PDF
    pdf.cell(200, 10, txt="Agent Response:", ln=True, align='L')
    pdf.multi_cell(0, 10, txt=response)
    
    # Save PDF file
    pdf_filename = "agent_response.pdf"
    pdf.output(pdf_filename)
    
    print(f"Response saved to {pdf_filename}.")
    return f"Response saved to PDF as {pdf_filename}."

# Wrap the tool function as FunctionTool
save_response_to_pdf_tool = FunctionTool.from_defaults(fn=save_response_to_pdf)

# Set LLM and Agent
llm = Ollama(model="llama3.2", request_timeout=360)
agent = ReActAgent.from_tools([save_response_to_pdf_tool], llm=llm, verbose=True)

# Use the Agent to converse and save the response to PDF
response = agent.chat("Please answer the question: who are you? and then save the answer to pdf file!")
print(str(response))

Code Output:

> Running step 00a441f1-302b-4e7b-afec-18a4e471a138. Step input: Please answer the question: who are you? and then save the answer to pdf file!
Thought: The current language of the user is: English. I need to use a tool to help me answer the question.
Action: save_response_to_pdf
Action Input: {'response': 'I am an artificial intelligence designed to simulate human-like conversations, answer questions, and provide information on a wide range of topics.', 'title': 'Answer'}
Response saved to agent_response.pdf.
Observation: Response saved to PDF as agent_response.pdf.
> Running step dd52ab6a-40dc-481c-a7a4-245fce133119. Step input: None
Thought: I can answer without using any more tools. I'll use the user's language to answer
Answer: I am an artificial intelligence designed to simulate human-like conversations, answer questions, and provide information on a wide range of topics.
I am an artificial intelligence designed to simulate human-like conversations, answer questions, and provide information on a wide range of topics.

As can be seen, the Agent understood my goals and requirements, completing the task through the steps of: Thinking (Thought), Action (Action) automatically finding the function to execute, Input (Action Input), and Observation (Observation).

Summary

This article demonstrated a simple agent functionality (saving content to a PDF document) using the LlamaIndex framework. The agent here only implements one functionality with a single step. Later, I will show how to enable the Agent to call multiple tool functions.

Leave a Comment