Building an LLM Application
Quick Summary
In this section, we will be developing an Agentic RAG medical chatbot to record symptoms, diagnose patients, and schedule appointments.
You may use whatever knowledge base you have to power your RAG Engine, but for the purposes of this tutorial, we'll be using The Gale Encyclopedia of Alternative Medicine.
Setting Up
Begin by installing the necessary packages. We'll use llama-index
as our RAG framework and chromadb
for vector indexing.
pip install llama-index llama-index-vector-stores-chroma doc2text chromadb
Since our chatbot will be recording patient information, we’ll need a structured way to store it. We'll define a pydantic
model with the relevant fields (symptoms, diagnoses, and personal information) to ensure that our agent can accurately store data in the correct format.
from pydantic import BaseModel
from typing import Optional
class MedicalAppointment(BaseModel):
name: Optional[str] = None
email: Optional[str] = None
date: Optional[str] = None
symptoms: Optional[str] = None
diagnosis: Optional[str] = None
Defining the Chatbot
Next, we'll create a MedicalAppointmentSystem
class to represent our agent. This class will store all MedicalAppointment
instances in an appointments
dictionary, with each key representing a unique user.
class MedicalAppointmentSystem:
def __init__(self):
self.appointments = {}
As we progress through this tutorial, we'll gradually enhance this class until it evolves into a fully functional medical chatbot agent.
Indexing the Knowledge Base
Let's start by building our RAG engine, which will handle all patient diagnoses. The first step is to load the relevant medical information chunks from our knowledge base into the system. We'll use the SimpleDirectoryReader
from llama-index
to accomplish this.
from llama_index.core import SimpleDirectoryReader
class MedicalAppointmentSystem:
def __init__(self, data_directory):
self.appointments = {}
self.load_data(data_directory)
def load_data(self, data_directory):
# Load documents from a directory
self.documents = SimpleDirectoryReader(data_directory).load_data()
Then, we'll use LlamaIndex's VectorStoreIndex
to embed our chunks and store them in a chromadb
database. This step is crucial, as our encyclopedia contains over 4,000 pages of dense medical information.
Embedding the data in a vector database ensures fast and accurate retrieval, even with such a large knowledge base.
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext
from llama_index.vector_stores.chroma import ChromaVectorStore
class MedicalAppointmentSystem:
def __init__(self, data_directory, db_path):
self.appointments = {}
self.load_data(data_directory)
self.store_data(db_path)
def load_data(self, data_directory):
...
def store_data(self, db_path):
# Set up the database and store vectorized data
db = chromadb.PersistentClient(path=db_path)
chroma_collection = db.get_or_create_collection("medical_knowledge")
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
self.index = VectorStoreIndex.from_documents(self.documents, storage_context=storage_context)
Building the Tools
Finally, we'll create the tools for our chatbot, which includes our RAG engine and function-calling tools responsible for creating, updating, and managing medical appointments, ensuring the system is both dynamic and interactive.
Assembling the RAG Engine
In llama-index
, a RAG engine is abstracted as a QueryEngineTool
, making it ideal for building agentic applications like this one. This approach also allows you to define additional tools alongside the RAG engine for enhanced functionality.
from llama_index.core.tools import QueryEngineTool
class MedicalAppointmentSystem:
def __init__(self, data_directory, db_path):
self.appointments = {}
self.load_data(data_directory)
self.store_data(db_path)
self.setup_tools()
...
def setup_tools(self):
query_engine = self.index.as_query_engine()
self.medical_diagnosis_tool = QueryEngineTool.from_defaults(
query_engine,
name="medical_diagnosis",
description="A RAG engine for retrieving medical information."
)
Function-calling Tools
We'll also need to define the tools for interacting with the medical appointment system, enabling tasks like creating, updating, and confirming appointments. LlamaIndex's FunctionTool
simplifies this by wrapping functions into callable tools.
from llama_index.core.tools import FunctionTool
class MedicalAppointmentSystem:
def __init__(self, data_directory, db_path):
self.appointments = {}
self.load_data(data_directory)
self.store_data(db_path)
self.setup_tools()
...
def setup_tools(self):
# Configure function tools for the system
self.get_appointment_state_tool = FunctionTool.from_defaults(fn=self.get_appointment_state)
self.update_appointment_tool = FunctionTool.from_defaults(fn=self.update_appointment)
self.create_appointment_tool = FunctionTool.from_defaults(fn=self.create_appointment)
self.record_diagnosis_tool = FunctionTool.from_defaults(fn=self.record_diagnosis)
self.medical_diagnosis_tool = QueryEngineTool.from_defaults(
self.query_engine,
name="medical_diagnosis",
description="A RAG engine for retrieving medical information."
)
Key Tools:
get_appointment_state_tool
: Retrieves the state of a specific appointment.update_appointment_tool
: Updates a property of an appointment.create_appointment_tool
: Creates a new appointment.record_diagnosis_tool
: Records a diagnosis for an appointment.
# Retrieves the current state of an appointment based on its ID.
def get_appointment_state(self, appointment_id: str) -> str:
try:
return str(self.appointments[appointment_id].dict())
except KeyError:
return f"Appointment ID {appointment_id} not found"
# Updates a specific property of an appointment.
def update_appointment(self, appointment_id: str, property: str, value: str) -> str:
appointment = self.appointments.get(appointment_id)
if appointment:
setattr(appointment, property, value)
return f"Appointment ID {appointment_id} updated with {property} = {value}"
return "Appointment not found"
# Creates a new appointment with a unique ID.
def create_appointment(self, appointment_id: str) -> str:
self.appointments[appointment_id] = MedicalAppointment()
return "Appointment created."
# Records a diagnosis for an appointment after symptoms have been noted.
def record_diagnosis(self, appointment_id: str, diagnosis: str) -> str:
appointment: MedicalAppointment = self.appointments.get(appointment_id)
if appointment and appointment.symptoms:
appointment.diagnosis = diagnosis
return f"Diagnosis recorded for Appointment ID {appointment_id}. Diagnosis: {diagnosis}"
return "Diagnosis cannot be recorded. Please tell me more about your symptoms."
Assembling the Chatbot
Now that we have set up the tools and data systems, it's time to assemble the chatbot agent. We'll use LlamaIndex's FunctionCallingAgent
to dynamically manage user interactions and choose the appropriate tool based on the input and context. This involves defining the LLM, system prompt, and tool integrations.
from llama_index.core.agent import FunctionCallingAgent
from llama_index.core.llms import ChatMessage
from llama_index.llms.openai import OpenAI
class MedicalAppointmentSystem:
...
def setup_agent(self):
# Initialize the GPT model and define the agent
gpt = OpenAI(model="gpt-4o", temperature=0.1)
self.agent = FunctionCallingAgent.from_tools(
tools=[
self.get_appointment_state_tool,
self.update_appointment_tool,
self.create_appointment_tool,
self.record_diagnosis_tool,
self.medical_diagnosis_tool,
self.confirm_appointment_tool
],
llm=gpt,
prefix_messages=[
ChatMessage(
role="system",
content=(
"You are an expert in medical diagnosis connected to a patient booking system. "
"First, create the appointment and record the symptoms. Ask for specific details! "
"After recording symptoms, make a precise diagnosis, updating the name, date, and email "
"only with explicitly provided information. Confirm all details at the end."
),
)
],
max_function_calls=10,
allow_parallel_tool_calls=False
)
Setting up the Interactive Session
Finally, we'll create an interactive environment where users can engage with the chatbot. This involves configuring input/output, managing conversation flow, and processing user queries.
class MedicalAppointmentSystem:
...
def interactive_session(self):
print("Welcome to the Medical Diagnosis and Booking System!")
print("Please enter your symptoms or ask about appointment details. Type 'exit' to quit.")
while True:
user_input = input("Your query: ")
if user_input.lower() == 'exit':
break
response = self.agent.chat(user_input)
print("Agent Response:", response.response)
To test your chatbot, run the following code:
if __name__ == "__main__":
system = MedicalAppointmentSystem(data_directory="./data", db_path="./chroma_db")
system.interactive_session()
Congratulations on building your Agentic RAG application! In the next sections, we’ll explore how to evaluate our medical chatbot, from selecting relevant metrics to iterating on the chatbot's hyperparameters for improved performance.