SingleStore Notebooks
New
How to Build a Multi-Agent AI App with AutoGen
Notebook
Note
This notebook can be run on a Free Starter Workspace. To create a Free Starter Workspace navigate to Start using the left nav. You can also use your existing Standard or Premium workspace with this Notebook.
What you will build
In this notebook you will build a small team of AI agents that work together to answer a technical question, and you will give one of those agents the ability to look things up in a document you provide.
The scenario. A user asks a question about a software library. A general-purpose chat model has to answer from memory, and if the library is niche or recent it will often guess. Instead, we will store the library's documentation in SingleStore, let an agent search it, and have the team answer from what it actually found.
What you will build
A database in SingleStore to hold the document.
A pipeline that splits a document into chunks and turns each chunk into an embedding — a numeric vector that captures its meaning — then stores those vectors in SingleStore.
A team of four agents: a Boss who asks the question, a Product Manager, a Senior Python Engineer who writes the code, and a Code Reviewer who checks it.
Two runs of the same question, side by side: once where the agents answer from the model's own knowledge, and once where an agent retrieves the relevant documentation first. The difference between those two answers is the point of the whole notebook.
What you need
A SingleStore workspace (the notebook connects to whichever one is selected above).
An OpenAI API key, which you will be prompted for. Nothing is hardcoded, and the key is never written into this file.
Expect the whole notebook to cost a few cents in OpenAI usage and take a few minutes to run.
How the pieces fit together
Two ideas are combined here, and it helps to keep them separate in your head.
Retrieval-augmented generation (RAG). A language model can only reason about text you put in front of it. RAG is the pattern of finding the right text first and pasting it into the prompt. To find it, we convert the question and every chunk of the document into embeddings, then ask the database which chunks sit closest to the question in vector space. SingleStore does that search for us, so the database is doing real work here rather than just holding files.
Multi-agent collaboration. Instead of one model answering in one shot, we create several agents, each with its own instructions and role, and let them talk in a group chat. One writes code, another reviews it, a third keeps the goal in view. This tends to catch mistakes that a single response misses.
The interesting connection between the two is that only one agent needs the retrieval power. The Boss's assistant is the agent wired up to SingleStore; the others simply see the documentation it retrieved appear in the conversation. That is a common and useful shape for agent systems — give one member of the team access to your data, and let the rest benefit from it.
We use AG2 for the agents. AG2 is the maintained continuation of the original
AutoGen autogen library, and it provides the group chat, the agent roles, and the retrieval agent we
will customize.
Step 1: Create a database
Everything we embed needs somewhere to live. The cell below creates a database called autogen.
Shared-tier (free) workspaces do not allow creating databases, so the notebook detects that case and uses the database already selected in the dropdown above instead. Either way, the rest of the notebook behaves the same.
In [1]:
1DATABASE_NAME = "autogen"2 3# Check if the database is running on a shared tier4shared_tier_check = %sql show variables like 'is_shared_tier'5on_shared_tier = bool(shared_tier_check) and shared_tier_check[0][1] == 'ON'6 7# On a shared tier we use the database already selected above; otherwise create a dedicated one.8if not on_shared_tier:9 %sql DROP DATABASE IF EXISTS autogen10 %sql CREATE DATABASE autogen11 12print("shared tier:", on_shared_tier)
Step 2: Install the libraries
Four things get installed, and it is worth knowing what each one is for:
ag2— the agent framework: agent roles, group chat, and the retrieval agent.langchain-singlestore— lets us use SingleStore as a vector store, so storing and searching embeddings is a couple of method calls instead of hand-written SQL.langchain-openai— the client for OpenAI's embedding and chat models.langchain-text-splitters— splits a long document into chunks small enough to embed.
The first cell records the versions of two numerical libraries this environment already relies on, so that installing the packages above leaves them untouched. This is good hygiene in any hosted notebook: installs can otherwise upgrade a dependency the running kernel is already using.
In [2]:
1# Note the versions of numpy and protobuf already present, and hold them there during the install.2import importlib.metadata as md3 4 5def hold(package):6 try:7 return f"{package}=={md.version(package)}"8 except md.PackageNotFoundError:9 return ""10 11 12PINS = " ".join(f'"{p}"' for p in (hold("numpy"), hold("protobuf")) if p)13print("holding:", PINS or "(nothing)")
In [3]:
1!pip install --quiet "ag2[openai]==0.9.10" chromadb langchain-singlestore "langchain-openai>=1.5" langchain-text-splitters {PINS}
Two things you may notice while that runs:
chromadbis in the list even though we never use it. AG2's retrieval agent lists it as a requirement, so it has to be present for the import to succeed. All of our actual retrieval happens in SingleStore.ERROR: pip's dependency resolver ...lines are normal here. Hosted images come with hundreds of preinstalled packages, and some of them disagree with each other regardless of what you install. If the imports in the following cells succeed, you are fine.
The next cell smooths over one environment quirk before we start making API calls.
In [4]:
1# Compatibility shim for hosted environments. Some images ship a response-compression library that is2# older than the OpenAI client expects, and the mismatch makes API calls fail with a misleading3# "connection error". Asking servers to skip that compression format avoids the problem entirely.4# This cell does nothing on an up-to-date environment, and it must run before any API call.5import brotli6import httpx2._client7import httpx2._decoders8 9 10def _compression_is_broken():11 try:12 brotli.Decompressor().process(brotli.compress(b"probe"), output_buffer_limit=1 << 20)13 return False14 except TypeError:15 return True16 17 18if _compression_is_broken():19 httpx2._decoders.SUPPORTED_DECODERS.pop("br", None)20 httpx2._client.ACCEPT_ENCODING = ", ".join(21 k for k in httpx2._decoders.SUPPORTED_DECODERS if k != "identity"22 )23 print("applied compatibility fix; response encodings:", httpx2._client.ACCEPT_ENCODING)24else:25 print("environment is up to date; no fix needed")
Step 3: Choose the document to make searchable
This is the knowledge you are giving your agents. We use a page from the FLAML documentation about running FLAML on Spark — a good test case, because it is specific enough that a chat model answering from memory tends to invent details.
Later, we will ask the team "How to use spark for parallel training in FLAML? Give me sample code." Keep that question in mind as you read the rest of the notebook.
To point this notebook at your own data, change the URL below (or load a local file instead). Any markdown or plain text document works.
In [5]:
1import requests2 3r = requests.get("https://raw.githubusercontent.com/microsoft/FLAML/main/website/docs/Examples/Integrate%20-%20Spark.md")4r.raise_for_status()5open('example.md', 'wb').write(r.content)6print(len(r.content), "bytes downloaded")
Step 4: Split the document and prepare embeddings
Three things happen in the next cell.
Your OpenAI key. You will be prompted for it. Typing it at a prompt rather than pasting it into a code cell means it never gets saved into the notebook file, which matters if you ever share or commit this notebook.
Chunking. Embeddings work best on passages, not whole documents: a single vector for a long page blurs every topic in it together, and a whole page is too much to paste into a prompt anyway. So we cut the document into ~1500-character chunks. The 150-character overlap between neighbouring chunks means a sentence that happens to fall on a boundary still appears intact in one of them.
The embedding model. text-embedding-3-small converts text into a 1536-number vector. The only rule
you must respect is consistency: the same model has to embed both the stored chunks and the incoming
question, or the distances between them are meaningless.
The last line makes one tiny embedding call, so that a problem with your key surfaces here rather than several cells later.
In [6]:
1import getpass2import os3 4from langchain_openai import OpenAIEmbeddings5from langchain_text_splitters import MarkdownTextSplitter6 7# Prompt for the key instead of hardcoding it, so it is never stored in the notebook file.8PLACEHOLDERS = {"", "api-key", "your-api-key", "sk-...", "<your-openai-api-key>"}9 10if os.environ.get("OPENAI_API_KEY", "").strip() in PLACEHOLDERS:11 os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API key: ").strip()12 13if os.environ["OPENAI_API_KEY"] in PLACEHOLDERS:14 raise ValueError("No OpenAI API key was entered -- re-run this cell and paste your key.")15if not os.environ["OPENAI_API_KEY"].startswith("sk-"):16 print("Warning: OpenAI keys normally start with 'sk-'; this one does not.")17 18# The hosted notebook environment exposes `connection_url` for the workspace selected above.19try:20 url = connection_url21except NameError:22 url = "admin:<password>@<host>:3306/<database>"23 24if not on_shared_tier:25 url = url.rsplit("/", 1)[0] + "/" + DATABASE_NAME26 27os.environ["SINGLESTOREDB_URL"] = url28 29# Split the document into chunks small enough to embed, overlapping slightly so that a sentence on a30# chunk boundary still survives intact somewhere.31text = open("example.md", encoding="utf-8").read()32text_splitter = MarkdownTextSplitter(chunk_size=1500, chunk_overlap=150)33docs = text_splitter.create_documents([text], metadatas=[{"source": "example.md"}])34 35# The same model must embed both the stored chunks and the incoming question.36embeddings = OpenAIEmbeddings(model="text-embedding-3-small")37 38# One tiny call, so a key problem shows up here rather than deeper in the notebook.39dims = len(embeddings.embed_query("ping"))40 41print(len(docs), "chunks;", dims, "embedding dimensions")
Note
The connection details for your workspace include a token that expires after roughly an hour. If a later cell fails with 2628: JWT token expired, re-run this cell to pick up fresh credentials, then re-create the vector store below.
Step 5: Store the vectors in SingleStore and search them
from_documents does the whole ingestion in one call: it embeds every chunk, creates a table with a
vector column, and inserts the chunks alongside their vectors.
Then we run a search to prove it works before handing anything to an agent. Notice what we search for:
"spark parallel training" — not a phrase that appears verbatim in the document. That is the advantage
of vector search over keyword matching. It finds passages that mean something similar, so a user's
phrasing does not have to match the document's.
The score next to each result is a distance: lower means more similar. If you are curious, change the query and re-run to see how the ranking shifts.
In [7]:
1from langchain_singlestore import SingleStoreVectorStore2 3singlestore_db = SingleStoreVectorStore.from_documents(4 docs,5 embeddings,6 table_name="notebook2", # use table with a custom name7)8 9# Confirm retrieval works before we hand the store to an agent. Lower score = closer match.10for doc, score in singlestore_db.similarity_search_with_score("spark parallel training", k=2):11 print(round(score, 4), doc.page_content[:80].replace("\n", " "))
Those two results are the passages an agent will be given when it asks about Spark. Re-running this cell
appends the chunks a second time, so if you want a clean slate, drop the table first with
%sql DROP TABLE IF EXISTS notebook2.
Your knowledge base is ready. The rest of the notebook is about the agents that use it.
Step 6: Give an agent access to SingleStore
AG2 ships a RetrieveUserProxyAgent — an agent that looks documents up before answering. Out of the
box it keeps its documents in a local store, so we subclass it and replace one method, retrieve_docs,
with a SingleStore search.
This is the seam worth understanding, because it is how you would connect an agent to any data source
you have: the agent framework does not care where the text comes from, only that retrieve_docs puts
results where it expects to find them. Ours returns each match as a small record with an id, the
content, and its metadata, paired with its similarity score.
In [8]:
1import autogen2from autogen.agentchat.contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent3 4print("autogen", autogen.__version__)
In [9]:
1from typing import Any2 3 4class SingleStoreRetrieveUserProxyAgent(RetrieveUserProxyAgent):5 def __init__(self, singlestore_db: SingleStoreVectorStore, **kwargs: Any):6 super().__init__(**kwargs)7 self.singlestore_db = singlestore_db8 9 def retrieve_docs(self, problem: str, n_results: int = 20, search_string: str = "", **kwargs: Any):10 hits = self.singlestore_db.similarity_search_with_score(query=problem, k=n_results)11 12 # One list of (record, score) pairs per query, which is the shape AG2 reads from.13 self._results = [14 [15 (16 {17 "id": f"{i}",18 "content": doc.page_content,19 "metadata": doc.metadata,20 },21 float(score),22 )23 for i, (doc, score) in enumerate(hits)24 ]25 ]26 print("doc_ids: ", [[doc["id"] for doc, _ in self._results[0]]])
Agents can be configured to execute code they write inside a Docker container. That is not available in a hosted notebook, and we do not need it here — our agents only write code for you to read — so we turn it off.
In [10]:
1import os2 3os.environ["AUTOGEN_USE_DOCKER"] = "False"
Step 7: Assemble the team
Now we create the agents. Each one is just a name, a system message describing its job, and a model — their behaviour comes almost entirely from those instructions.
Boss — a stand-in for you. It poses the question.
human_input_mode="NEVER"means it will not stop to ask you to type anything, so the notebook runs start to finish on its own.Boss_Assistant — the retrieval agent from the previous step, the only one that can search SingleStore.
Product Manager, Senior Python Engineer, Code Reviewer — the specialists. Their system messages are short on purpose; read them and imagine how the conversation changes if you reword them.
Two mechanics worth noticing:
Termination. Agents are told to reply TERMINATE when the work is done, and termination_msg looks
for it. Without a stop condition, a group chat would keep talking until it hit max_round.
Speaker selection. round_robin means the agents take strict turns. The alternative, "auto", lets
the model decide who should speak next based on each agent's description — which is why those
descriptions are written for an audience of other agents rather than for you.
The three functions at the bottom (norag_chat, rag_chat, call_rag_chat) set up three different
team structures. We will run the first two next and compare them.
In [11]:
1from typing import Annotated2 3# Any chat model your key can access will work here.4MODEL = "gpt-4.1-mini"5 6llm_config = {7 "config_list": [{"model": MODEL, "api_key": os.environ["OPENAI_API_KEY"]}],8}9 10 11def termination_msg(x):12 return isinstance(x, dict) and "TERMINATE" == str(x.get("content", ""))[-9:].upper()13 14 15boss = autogen.UserProxyAgent(16 name="Boss",17 is_termination_msg=termination_msg,18 human_input_mode="NEVER",19 code_execution_config=False, # we don't want to execute code in this case.20 default_auto_reply="Reply `TERMINATE` if the task is done.",21 description="The boss who ask questions and give tasks.",22)23 24boss_aid = SingleStoreRetrieveUserProxyAgent(25 name="Boss_Assistant",26 is_termination_msg=termination_msg,27 human_input_mode="NEVER",28 max_consecutive_auto_reply=3,29 retrieve_config={30 "task": "code",31 # Retrieval comes from SingleStore via retrieve_docs, so AG2's built-in local store is unused.32 "vector_db": None,33 "docs_path": None,34 },35 code_execution_config=False, # we don't want to execute code in this case.36 description="Assistant who has extra content retrieval power for solving difficult problems.",37 singlestore_db=singlestore_db,38)39 40coder = autogen.AssistantAgent(41 name="Senior_Python_Engineer",42 is_termination_msg=termination_msg,43 system_message="You are a senior python engineer, you provide python code to answer questions. Reply `TERMINATE` in the end when everything is done.",44 llm_config=llm_config,45 description="Senior Python Engineer who can write code to solve problems and answer questions.",46)47 48pm = autogen.AssistantAgent(49 name="Product_Manager",50 is_termination_msg=termination_msg,51 system_message="You are a product manager. Reply `TERMINATE` in the end when everything is done.",52 llm_config=llm_config,53 description="Product Manager who can design and plan the project.",54)55 56reviewer = autogen.AssistantAgent(57 name="Code_Reviewer",58 is_termination_msg=termination_msg,59 system_message="You are a code reviewer. Reply `TERMINATE` in the end when everything is done.",60 llm_config=llm_config,61 description="Code Reviewer who can review the code.",62)63 64PROBLEM = "How to use spark for parallel training in FLAML? Give me sample code."65 66 67def _reset_agents():68 boss.reset()69 boss_aid.reset()70 coder.reset()71 pm.reset()72 reviewer.reset()73 74 75def rag_chat():76 """The retrieval agent searches SingleStore, then opens the group chat with what it found."""77 _reset_agents()78 groupchat = autogen.GroupChat(79 agents=[boss_aid, pm, coder, reviewer], messages=[], max_round=12, speaker_selection_method="round_robin"80 )81 manager = autogen.GroupChatManager(groupchat=groupchat, llm_config=llm_config)82 83 # `message` builds the opening message, which is where the retrieved documentation gets inserted.84 boss_aid.initiate_chat(85 manager,86 message=boss_aid.message_generator,87 problem=PROBLEM,88 n_results=3,89 )90 91 92def norag_chat():93 """The same team and question, with no access to the document."""94 _reset_agents()95 groupchat = autogen.GroupChat(96 agents=[boss, pm, coder, reviewer],97 messages=[],98 max_round=12,99 speaker_selection_method="auto",100 allow_repeat_speaker=False,101 )102 manager = autogen.GroupChatManager(groupchat=groupchat, llm_config=llm_config)103 104 # Start chatting with the boss as this is the user proxy agent.105 boss.initiate_chat(106 manager,107 message=PROBLEM,108 )109 110 111def call_rag_chat():112 """Retrieval as a tool: agents call it when they decide they need it, rather than up front."""113 _reset_agents()114 115 # Here the retrieval agent is not part of the group chat. Instead it is wrapped in a function that116 # the other agents can call, which lets them look things up mid-conversation and as often as they117 # like -- useful when a single search at the start would not be enough.118 def retrieve_content(119 message: Annotated[120 str,121 "Refined message which keeps the original meaning and can be used to retrieve content for code generation and question answering.",122 ],123 n_results: Annotated[int, "number of results"] = 3,124 ) -> str:125 boss_aid.n_results = n_results # Set the number of results to be retrieved.126 # Check if we need to update the context.127 update_context_case1, update_context_case2 = boss_aid._check_update_context(message)128 if (update_context_case1 or update_context_case2) and boss_aid.update_context:129 boss_aid.problem = message if not hasattr(boss_aid, "problem") else boss_aid.problem130 _, ret_msg = boss_aid._generate_retrieve_user_reply(message)131 else:132 ret_msg = boss_aid.message_generator(boss_aid, None, {"problem": message, "n_results": n_results})133 return ret_msg if ret_msg else message134 135 boss_aid.human_input_mode = "NEVER" # Disable human input for boss_aid since it only retrieves content.136 137 for caller in [pm, coder, reviewer]:138 d_retrieve_content = caller.register_for_llm(139 description="retrieve content for code generation and question answering.", api_style="function"140 )(retrieve_content)141 142 for executor in [boss, pm]:143 executor.register_for_execution()(d_retrieve_content)144 145 groupchat = autogen.GroupChat(146 agents=[boss, pm, coder, reviewer],147 messages=[],148 max_round=12,149 speaker_selection_method="round_robin",150 allow_repeat_speaker=False,151 )152 153 manager = autogen.GroupChatManager(groupchat=groupchat, llm_config=llm_config)154 155 # Start chatting with the boss as this is the user proxy agent.156 boss.initiate_chat(157 manager,158 message=PROBLEM,159 )
Step 8: Run the team without retrieval
First, the baseline. The agents get the question and nothing else, so they must answer from what the model already knows.
Read the code they produce with some suspicion. Do the function and parameter names look right? This is where a model tends to produce something plausible and subtly wrong — an argument that does not exist, or a setting borrowed from a different library. Compare it against the real FLAML documentation we downloaded earlier.
The output is long; that is normal, since you are watching four agents take turns.
In [12]:
1norag_chat()
Step 9: Run the team with retrieval
Now the same question, with Boss_Assistant in the room. Before anyone speaks, it searches SingleStore
and puts the three closest chunks into the opening message.
Watch for two things in the output:
The
doc_ids:line printed by theretrieve_docsmethod you wrote — that is your SingleStore query running.The retrieved documentation appearing at the top of the conversation, and then details from it showing up in the engineer's code.
The answer should now be grounded in the actual document rather than in the model's recollection of it. That difference — same model, same question, same team — is what a vector database adds to an agent system.
In [13]:
1rag_chat()
Where to go from here
A few small changes that teach you the most:
Use your own document. Change the URL in step 3 and the question in
PROBLEM. This is the whole point of the pattern — your data, your agents.Try
call_rag_chat(). The third function treats retrieval as a tool the agents call when they decide they need it, instead of a single search up front. Run it and compare the conversation shape.Change the number of chunks retrieved.
n_results=3inrag_chatcontrols how much context the team gets. More is not automatically better.Rewrite a system message. Give the Code Reviewer a stricter brief and see whether the final code changes.
The cell below drops the database created for this notebook, so nothing is left behind. Skip it if you want to keep the table around to explore.
In [14]:
1if not on_shared_tier:2 %sql DROP DATABASE IF EXISTS autogen
Details
About this Template
Learn how to build a multi-agent group chat with RAG using Autogen and SingleStore
This Notebook can be run in Shared Tier, Standard and Enterprise deployments.
Tags
starterautogenragmultiagentgroupchat
See Notebook in action
Launch this notebook in SingleStore and start executing queries instantly.
License
This Notebook has been released under the Apache 2.0 open source license.