Background
The engineering team at to improve information retrieval (similar to its popular commercial competitor, , for one) or checking out some of the open source tools out there ( of open source AI Agents that you could start using tomorrow.
What’s your integration strategy?
🤔 Considerations: As mentioned above, the biggest consideration here is: what sort of information do you need your Agent to have access to? You could maybe get away with simply integrating it with third-party providers via an API, but if you need the integration to be more specific to your needs then you’ll need to be more thoughtful with how your integrations work.
By carefully considering what you’ll need to integrate with before you start building, you’ll save yourself some headache later on. Do you need your Agent to be able to execute custom scripts to query your databases? Do you need real-time retrieval of logs and metrics, and how will you design the Agent to retrieve that information? Will it return the link to the source? Will it return a chunk of lines of logs that you still have to manually sift through, or will it be able to deduce where the anomaly may be?
⚒️ What we built: At its core, Aptible AI is built on a series of integrations. An integration is more than just a connection to a third-party provider, it’s also a collection of configurations that are unique to how our team uses that provider. For example, Aptible AI supports multiple integrations for the same provider since we may want to use that provider in different ways. Different teams use Datadog differently and care about different metrics or use different tags, so each team can use the integration to the same tool in the way that they need.
Aptible AI supports a range of common SRE tooling, including:
- Chat and other highly synchronous communications
- Documentation and other knowledge repositories
- Observability
- Alerting
The actual implementation of these integrations fits into one of three categories of customizability:
For starters, you have a basic integration that requires no customization (PagerDuty is one example). Since it’s just pulling data from PagerDuty and adding it to the AI’s context, every single team that leverages the PagerDuty integration uses it in the same way.
Next, we have more customizable integrations (like the Datadog example from before) that are built on top of a generic InfluxDB integration but customized to the specific use cases of looking up container metrics and looking up restart activity.
Finally, there are fully custom tools that would likely make no sense to anyone outside of Aptible (an example here would be our integration that gets containers for an application). These are entirely specific to how we run our infrastructure and can be implemented either by a lightweight PubSub interface or a websocket-based “safe” proxy.
💡 Pro tip: Less is more! If you give the model too many tools to choose from, it can start choosing incorrect tools and confuse itself. More on that in the next section
So many models, how do you pick one?!
🤔 Considerations: Here’s the thing with models… new ones pop up every day, and there are several considerations to keep in mind when choosing one (mainly to do with your specific use cases). Should you self-host? Do you need your Agent to be conversational or task-based or both? Will it be conducting simple or complex tasks? Do you need real-time performance?
There’s no need for us to go through all the models that exist since that content is already all over the place (if you want a deep dive, ). The better your prompt engineering, the better your Agent will be.
For context, here are a few that we considered (over time) when building Aptible AI:
: this is what slightly-more-experienced people do when talking to ChatGPT; they ask it a question and include examples of the output they want. You might use zero- and/or few-shot prompting for very simple tasks that the underlying model already knows how to do.
: this technique allows an agent to generate “thoughts” and take “actions” in an iterative way to solve a problem, most similar to human reasoning. ReAct is great for moderately complex problems, like navigating references through documentation and tools in real time to compose an answer.
An important thing to keep in mind is that you can mix and match with these techniques (we’ll cover the multi-agent approach next). Here’s what we did…
⚒️ What we built: Because Aptible AI has a multi-agent structure (more on that later), we’ve implemented a mix of ReAct and RAG depending on the complexity of the task/question.
Can’t forget about security...
🤔 Considerations: Here’s a topic that comes up a lot when we chat with Aptible AI early users. Most engineering teams eventually have to face their security team when it comes to implementing new tools, and it’s critical to ensure that the data is safe (especially if you’re working in a highly regulated industry). So the first thing you have to do is to know your organization’s AI security policy, then there are a few things you can do to protect against potential data leaks or external threats.
⚒️ What we built: For starters, we use a model that doesn’t train on our data. We're still doing a lot of discovery around what customers need regarding security, whether that's self-hosting or something else! Stay tuned.
💡 Pro tip: to ensure what’s passed to the LLM and end users is sanitized
Oh, and of course, it needs to be usable!
🤔 Considerations: How do you plan to use your Agent? Does it need to have a UI? Will it be used across the organization?
You likely don’t need to spend time reinventing the wheel when it comes to the UX around your bot. Frameworks like
2. Make your application smarter by connecting an LLM
With our Chainlit app scaffolded, we can connect it to an LLM so that we can talk to it and get a human-like response.
We’ll use OpenAI’s hosted gpt-4o model for simplicity, but using another provider is just a matter of syntax.
The Goal
By the end of this article, you’ll be able to prompt the gpt-4o model and get a response, similar to how you’d interact with ChatGPT. We’ll also make sure that the bot maintains conversation context so that you can ask follow-up questions.
Prerequisites
Before you get started, you’ll need:
An OpenAI account and an
Curing Amnesia
If you’ve played around a bit and asked follow-up questions, you may have noticed that the bot doesn’t “remember” anything you’ve talked about. For example:
3. Faster feedback 🏎️
After completing the first few steps from Part 1, you may have noticed that when you ask questions that require a long response, there’s a delay before you see anything.
This can make for a poor user experience (especially later in part 3, when we start adding long-running tool calls) so let’s fix that.
The Goal
At the end of this step, you’ll be able to see your bot “type” in real-time, similar to ChatGPT.
Stream it
To get real-time message updates, we need to update our implementation to use a “stream”. Basically, whenever we receive a message, we’ll respond immediately with an empty message, start a stream with the LLM, and update our empty message every time we receive a new chunk of the response from the stream.
This might sound complicated, but it’s surprisingly easy! Update your handle_message function as follows:
# ...
@cl.on_message
async def handle_message(message: cl.Message) -> None:
# Send an empty initial message that we can update with a streaming
# response.
message = cl.Message(content="")
await message.send()
# Stream the response from the LLM
stream = await client.chat.completions.create(
messages=[
# Prepend all previous messages to maintain the conversation.
*cl.chat_context.to_openai(),
{"content": message.content, "role": "user"}
],
model="gpt-4o",
stream=True,
)
# Update the existing (initially-empty) message with new content
# from each "chunk" in the stream.
async for chunk in stream:
if token := chunk.choices[0].delta.content:
await message.stream_token(token)
# Send a final update to let the message know it's complete.
await message.update()
🧑💻 So, here's the complete code so far:
import os
import chainlit as cl
from openai import AsyncOpenAI
##
# Settings
#
try:
OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
except KeyError as ex:
raise LookupError(f"Missing required environment variable: {ex}")
client = AsyncOpenAI(api_key=OPENAI_API_KEY)
@cl.on_message
async def handle_message(message: cl.Message) -> None:
# Send an empty initial message that we can update with a streaming
# response.
message = cl.Message(content="")
await message.send()
# Stream the response from the LLM
stream = await client.chat.completions.create(
messages=[
# Prepend all previous messages to maintain the conversation.
*cl.chat_context.to_openai(),
{"content": message.content, "role": "user"}
],
model="gpt-4o",
stream=True,
)
# Update the existing (initially-empty) message with the content
# from each "chunk" in the stream.
Try it out
Now, when you ask a question, you should see your bot “typing” in real-time!
, so feel free to update the glob pattern to whatever makes sense for your use case!
This will automatically upload all of the files in the ./docs folder and add them to our vector store.
Add an indicator
File search can sometimes take a while, especially for larger datasets. In those cases, you’ll probably want to let the user know what’s going on so they don’t get frustrated.
Luckily, Chainlit makes this easy by providing a Step class that we can use to tell the user that something’s happening in the background. We can use the Step class in conjunction with the MessageEventHandler we built earlier, and add an indicator any time a tool is called.
Add the following to your MessageEventHandler:
class MessageEventHandler(AsyncAssistantEventHandler):
# ...
@override
async def on_tool_call_created(self, tool_call: ToolCall) -> None:
"""Create a new step in the conversation to indicate that a tool is being used."""
async with cl.Step(tool_call.type) as step:
self.step = step
Try it out
Now that you’ve uploaded some of your own documentation, try asking some questions that are more specific to your use case, and see what you get!
For our test case, it correctly referenced our runbook when asked about high CPU utilization on a customer database:
🧑💻 Here's the complete code:
from datetime import datetime, timedelta
import json
import os
from pathlib import Path
from typing import override
import chainlit as cl
from openai import AsyncOpenAI, AsyncAssistantEventHandler
from openai.types.beta.threads import Message, TextDelta, Text
from openai.types.beta import AssistantStreamEvent
from openai.types.beta.threads.runs.tool_call import ToolCall
from openai.types.beta.threads import Run
##
# Settings
#
try:
OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
except KeyError as ex:
raise LookupError(f"Missing required environment variable: {ex}")
# Give your assistant a name.
OPENAI_ASSISTANT_NAME = "roger"
# Give your assistant some custom instructions.
OPENAI_ASSISTANT_INSTRUCTIONS = """
You are an expert Site Reliability Engineer, tasked with helping
the SRE team respond to and resolve incidents.
If you are presented with a question that does not seem like it
could be related to infrastructure, begin your response with a polite
reminder that your primary responsibilities are to help with incident
response, before fully answering the question to the best of your ability.
Use the provided tools to gather additional context about the incident, if
applicable.
"""
client = AsyncOpenAI(api_key=OPENAI_API_KEY)
class MessageEventHandler(AsyncAssistantEventHandler):
"""An event handler for updating a Chainlit message while streaming an OpenAI response."""
message: cl.Message
@override
async def on_text_created(self, text: Text) -> None:
"""Create a new message so that we can update it."""
self.message = cl.Message(content="")
await self.message.send()
@override
async def on_text_delta(self, delta: TextDelta, snapshot: Text) -> None:
"""Update the message with the latest text delta streamed to us."""
await self.message.stream_token(delta.value)
@override
async def on_message_done(self, message: Message) -> None:
"""Update the message with the final text when the stream completes."""
await self.message.update()
@override
async def on_tool_call_created(self, tool_call: ToolCall) -> None:
"""Create a new step in the conversation to indicate that a tool is being used."""
async with cl.Step(tool_call.type) as step:
self.step = step
@override
async def on_event(self, event: AssistantStreamEvent) -> None:
"""Handle specific events."""
if event.event == "thread.run.requires_action":
run_id = event.data.id # Retrieve the run ID from the event data
self.current_run.id = run_id
await self.handle_requires_action(event.data, run_id)
async def handle_requires_action(self, run: Run, run_id: str) -> None:
"""Handle events that require an action to be taken, like tool calls."""
tool_outputs = []
# Execute each tool call and collect the result.
for tool in run.required_action.submit_tool_outputs.tool_calls:
func_name = tool.function.name
func_args = tool.function.arguments
# TODO: Build a function map with a decorator instead of looking up
# functions in globals.
if func_to_call := globals()[func_name]:
try:
# Parse the func_args JSON string to a dictionary
tool_outputs.append(
{
"tool_call_id": tool.id,
"output": await func_to_call(**json.loads(func_args)),
}
)
except TypeError as ex:
print(f"Error calling function {func_name!r}: {str(ex)}")
else:
print(f"Function {func_name!r} not found")
# Submit tool outputs to the conversation thread.
async with client.beta.threads.runs.submit_tool_outputs_stream(
thread_id=self.current_run.thread_id,
run_id=run_id,
tool_outputs=tool_outputs,
event_handler=MessageEventHandler(),
) as stream:
await stream.until_done()
@cl.on_chat_start
async def handle_chat_start() -> str:
vector_store = None
# Try to find an existing vector store so we don't create duplicates.
async for existing_vector_store in await client.beta.vector_stores.list():
if existing_vector_store.name == OPENAI_ASSISTANT_NAME:
vector_store = existing_vector_store
break
# Create a vector store if we didn't find an existing one.
vector_store = vector_store or await client.beta.vector_stores.create(
name=OPENAI_ASSISTANT_NAME,
)
if documents := list(Path("./docs").glob("**/*.md")):
await client.beta.vector_stores.file_batches.upload_and_poll(
vector_store_id=vector_store.id,
files=(f.open("rb") for f in documents),
)
assistant = None
# Try to find an existing assistant so we don't create duplicates.
async for existing_assistant in await client.beta.assistants.list():
if existing_assistant.name == OPENAI_ASSISTANT_NAME:
assistant = existing_assistant
break
# Create an assistant if we didn't find an existing one.
assistant = assistant or await client.beta.assistants.create(
name=OPENAI_ASSISTANT_NAME,
model="gpt-4o",
)
# Update the assistant so that it always has the latest instructions
assistant = await client.beta.assistants.update(
assistant_id=assistant.id,
instructions=OPENAI_ASSISTANT_INSTRUCTIONS,
tools=[
# Our existing file search tool
{"type": "file_search"},
# Our new pagerduty alert tool
{
"type": "function",
"function": {
"name": get_pagerduty_alert_details.__name__,
"description": get_pagerduty_alert_details.__doc__,
"parameters": {
"type": "object",
"properties": {
"pagerduty_alert_url": {
"type": "string",
"format": "uri",
"description": "The PagerDuty alert URL. For example, 'https://example.pagerduty.com/alerts/Q3YDP8VKEZ9THL'.",
},
},
"required": ["pagerduty_alert_url"],
"additionalProperties": False,
},
},
},
],
tool_resources={
"file_search": {
"vector_store_ids": [vector_store.id],
}
},
)
# Create a thread for the conversation
thread = await client.beta.threads.create()
# Add the assistant and the new thread to the user session so that
# we can reference it in other handlers.
cl.user_session.set("assistant", assistant)
cl.user_session.set("thread", thread)
@cl.on_message
async def handle_message(message: cl.Message) -> None:
# Retrieve our Assistant and Thread from our user session.
assistant = cl.user_session.get("assistant")
thread = cl.user_session.get("thread")
# Add the latest message to the thread.
await client.beta.threads.messages.create(
thread_id=thread.id,
role="user",
content=message.content,
)
# Stream a response to the Thread (called a "Run") using our Assistant.
async with client.beta.threads.runs.stream(
assistant_id=assistant.id,
thread_id=thread.id,
# Use our custom message handler.
event_handler=MessageEventHandler(),
) as stream:
await stream.until_done()
async def get_pagerduty_alert_details(pagerduty_alert_url: str) -> dict:
"""Return the details of the PagerDuty alert at the given URL."""
# TODO: Make this implementation real!
return json.dumps(
{
"alert": {
"id": "PT4KHLK",
"type": "alert",
"summary": "A customer database is experiencing high CPU usage.",
"self": pagerduty_alert_url,
"html_url": pagerduty_alert_url,
"created_at": (datetime.now() - timedelta(minutes=5)).isoformat(),
"status": "resolved",
"alert_key": "baf7cf21b1da41b4b0221008339ff357",
"suppressed": False,
"severity": "critical",
}
}
)
Wrapping up
Now you're all set to build a useful AI Agent for your SRE team! If you have any questions about anything we've covered in this guide, please reach out, and we'll be happy to help. In the meantime, if there is anything missing or any other AI Agent-related thing you'd like to learn, let us know!
If you're curious to try out Aptible AI for yourself rather than building your own Agent, you can visit www.aptible.ai to sign up.
SOCIAL SHARE CARD GENERATOR