Introduction to Pydantic AI for Knowledge Base Chatbots
Pydantic AI is a relatively new Python framework that brings type safety, structured outputs, and dependency injection to LLM-powered applications. Built by the same team behind Pydantic, it offers a developer-friendly way to build AI agents that integrate cleanly with existing Python codebases. When combined with a retrieval-augmented generation (RAG) pipeline, Pydantic AI becomes a powerful foundation for building knowledge base chatbots that can answer questions accurately based on your own documents.
In this tutorial, you'll learn how to build a complete knowledge base chatbot from scratch. We'll cover document ingestion, vector storage, semantic retrieval, and how to wire everything together using Pydantic AI agents. By the end, you'll have a working chatbot that grounds its answers in your custom knowledge base.
What Is a Knowledge Base Chatbot?
A knowledge base chatbot is an AI assistant that answers user questions using a curated set of documents rather than relying solely on the model's training data. This approach, known as Retrieval-Augmented Generation (RAG), dramatically reduces hallucinations and ensures responses are grounded in authoritative sources.
The typical architecture involves three main components:
- Document Store: Where your knowledge base documents live, usually chunked and embedded into vectors.
- Retriever: A mechanism that finds the most relevant document chunks for a given user query using semantic similarity.
- LLM Agent: The language model that synthesizes retrieved context into a coherent, helpful answer.
Pydantic AI shines in the agent layer. It lets you define typed dependencies, structured outputs, and tool functions with full IDE support and runtime validation. This means your chatbot can reliably call retrieval functions, parse results, and return answers in predictable formats.
Why Pydantic AI Matters
Several frameworks exist for building LLM applications, but Pydantic AI differentiates itself in a few key ways. First, it leverages Python's type system natively, so you get autocomplete, type checking, and validation out of the box. Second, it's model-agnostic — you can switch between OpenAI, Anthropic, Gemini, and local models with minimal code changes. Third, its dependency injection system makes testing straightforward, since you can mock dependencies like database connections or retrieval functions easily.
For knowledge base chatbots specifically, these features translate into real benefits. Structured outputs ensure your chatbot returns consistent response formats. Typed tools make retrieval logic explicit and testable. And the framework's lightweight nature means you're not locked into a heavy abstraction layer — you can always drop down to raw API calls when needed.
Prerequisites and Setup
Before we start coding, make sure you have Python 3.10 or higher installed. You'll also need an API key for an LLM provider — this tutorial uses OpenAI, but the code can be adapted for other providers supported by Pydantic AI.
Create a new project directory and install the required packages:
mkdir kb-chatbot
cd kb-chatbot
python -m venv venv
source venv/bin/activate # On Windows: venv\Scripts\activate
pip install pydantic-ai openai chromadb sentence-transformers python-dotenv
Create a .env file in your project root to store your API key:
OPENAI_API_KEY=your-api-key-here
Now let's set up the project structure:
kb-chatbot/
├── .env
├── main.py
├── knowledge_base.py
├── agent.py
└── data/
└── documents/
├── guide.txt
└── faq.txt
Place a few text files in the data/documents/ directory. These will serve as your knowledge base content. For testing, you can use any text — product manuals, internal wikis, FAQs, or policy documents all work well.
Building the Knowledge Base Layer
The first component we'll build handles document ingestion and retrieval. We'll use ChromaDB as our vector database and sentence-transformers for generating embeddings locally. This keeps costs down and avoids additional API calls for embedding.
Document Chunking and Storage
Create knowledge_base.py with the following content:
import os
from pathlib import Path
from typing import List, Dict
import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer
class KnowledgeBase:
def __init__(self, persist_dir: str = "./chroma_db"):
self.persist_dir = persist_dir
self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
self.client = chromadb.PersistentClient(path=persist_dir)
self.collection = self.client.get_or_create_collection(
name="knowledge_base",
metadata={"hnsw:space": "cosine"}
)
def chunk_text(self, text: str, chunk_size: int = 500, overlap: int = 50) -> List[str]:
"""Split text into overlapping chunks for better retrieval."""
chunks = []
start = 0
while start < len(text):
end = start + chunk_size
chunk = text[start:end]
chunks.append(chunk.strip())
start = end - overlap
return [c for c in chunks if c]
def ingest_directory(self, dir_path: str) -> int:
"""Load and ingest all text files from a directory."""
path = Path(dir_path)
if not path.exists():
raise FileNotFoundError(f"Directory not found: {dir_path}")
total_chunks = 0
for file_path in path.glob("*.txt"):
with open(file_path, "r", encoding="utf-8") as f:
content = f.read()
chunks = self.chunk_text(content)
embeddings = self.embedding_model.encode(chunks).tolist()
ids = [f"{file_path.stem}_{i}" for i in range(len(chunks))]
metadatas = [{"source": file_path.name, "chunk_index": i} for i in range(len(chunks))]
self.collection.add(
ids=ids,
documents=chunks,
embeddings=embeddings,
metadatas=metadatas
)
total_chunks += len(chunks)
print(f"Ingested {len(chunks)} chunks from {file_path.name}")
return total_chunks
def search(self, query: str, n_results: int = 5) -> List[Dict]:
"""Retrieve the most relevant chunks for a given query."""
query_embedding = self.embedding_model.encode([query]).tolist()
results = self.collection.query(
query_embeddings=query_embedding,
n_results=n_results
)
retrieved = []
for i in range(len(results["documents"][0])):
retrieved.append({
"content": results["documents"][0][i],
"source": results["metadatas"][0][i]["source"],
"distance": results["distances"][0][i]
})
return retrieved
This class handles three core operations: chunking text into manageable pieces, ingesting documents into the vector store, and searching for relevant content. The chunking strategy uses a simple sliding window with overlap, which helps preserve context across chunk boundaries.
Understanding the Chunking Strategy
Chunk size matters a lot for retrieval quality. Too small, and you lose context. Too large, and semantic similarity scores become noisy. A chunk size of 500 characters with 50 characters of overlap is a reasonable starting point for general text. For code-heavy or highly structured documents, you might want to chunk by sections or headings instead.
The all-MiniLM-L6-v2 embedding model is a good balance of speed and quality. It produces 384-dimensional vectors and runs efficiently on CPU. If you need higher quality embeddings, consider all-mpnet-base-v2 at the cost of slower inference.
Creating the Pydantic AI Agent
Now we'll create the agent that ties the knowledge base to the LLM. Pydantic AI agents are defined with a system prompt, a model, optional dependencies, and tools. The agent uses these tools to retrieve information and then synthesizes an answer.
Create agent.py:
from dataclasses import dataclass
from pydantic_ai import Agent, RunContext
from knowledge_base import KnowledgeBase
@dataclass
class KBDependencies:
"""Dependencies injected into the agent at runtime."""
knowledge_base: KnowledgeBase
# Define structured output model
from pydantic import BaseModel, Field
class ChatResponse(BaseModel):
answer: str = Field(description="The answer to the user's question")
sources: list[str] = Field(description="List of source files referenced")
confidence: float = Field(description="Confidence score between 0 and 1", ge=0, le=1)
kb_agent = Agent(
model="openai:gpt-4o-mini",
deps_type=KBDependencies,
output_type=ChatResponse,
system_prompt=(
"You are a knowledgeable assistant that answers questions based on a "
"knowledge base. Always use the search_knowledge_base tool to find "
"relevant information before answering. If the retrieved information "
"does not contain the answer, say so honestly. Always cite your sources. "
"Be concise but thorough."
),
)
@kb_agent.tool
async def search_knowledge_base(ctx: RunContext[KBDependencies], query: str) -> str:
"""Search the knowledge base for information relevant to the query.
Args:
query: The search query to look up in the knowledge base.
"""
results = ctx.deps.knowledge_base.search(query, n_results=5)
if not results:
return "No relevant documents found."
formatted = []
for i, result in enumerate(results, 1):
formatted.append(
f"[Result {i}] (Source: {result['source']}, "
f"Relevance: {1 - result['distance']:.2f})\n{result['content']}"
)
return "\n\n---\n\n".join(formatted)
Let's break down what's happening here. The KBDependencies dataclass defines what the agent needs at runtime — in this case, our knowledge base instance. The ChatResponse model defines the structured output format, ensuring every response includes an answer, source citations, and a confidence score.
The @kb_agent.tool decorator registers search_knowledge_base as a tool the LLM can call. When the model decides it needs information, it calls this function with a query string. The function retrieves relevant chunks from the vector store and formats them for the model to read.
How Dependency Injection Works
Pydantic AI's dependency injection system is one of its standout features. The RunContext parameter gives your tool functions access to whatever dependencies you pass when running the agent. This means you can swap out a real knowledge base for a mock during testing, or use different knowledge bases for different users, all without changing the agent definition.
Putting It All Together
Now let's create the main entry point that initializes the knowledge base, ingests documents, and runs the chatbot. Create main.py:
import asyncio
import os
from dotenv import load_dotenv
from knowledge_base import KnowledgeBase
from agent import kb_agent, KBDependencies
load_dotenv()
async def main():
# Initialize the knowledge base
kb = KnowledgeBase(persist_dir="./chroma_db")
# Ingest documents (only needed on first run or when documents change)
doc_dir = "./data/documents"
if os.path.exists(doc_dir):
chunk_count = kb.ingest_directory(doc_dir)
print(f"Total chunks ingested: {chunk_count}")
else:
print(f"Document directory not found: {doc_dir}")
print("Creating it now. Add .txt files and re-run.")
os.makedirs(doc_dir, exist_ok=True)
return
# Create dependencies
deps = KBDependencies(knowledge_base=kb)
print("\n=== Knowledge Base Chatbot ===")
print("Type 'quit' to exit, 'clear' to reset conversation.\n")
# Interactive chat loop
while True:
user_input = input("You: ").strip()
if user_input.lower() == "quit":
print("Goodbye!")
break
if user_input.lower() == "clear":
print("Conversation cleared.\n")
continue
if not user_input:
continue
try:
result = await kb_agent.run(user_input, deps=deps)
response = result.output
print(f"\nBot: {response.answer}")
print(f"Sources: {', '.join(response.sources)}")
print(f"Confidence: {response.confidence:.0%}\n")
except Exception as e:
print(f"\nError: {e}\n")
if __name__ == "__main__":
asyncio.run(main())
Run the chatbot with:
python main.py
If everything is set up correctly, you'll see the ingestion messages followed by an interactive prompt. Type a question related to your documents, and the chatbot will search the knowledge base, retrieve relevant chunks, and generate a grounded answer with source citations.
Adding Conversation Memory
The basic version above treats each question independently. For a more natural chatbot experience, we should maintain conversation history so the model can reference previous exchanges. Pydantic AI supports this through message history.
Update main.py to track conversation state:
import asyncio
import os
from dotenv import load_dotenv
from knowledge_base import KnowledgeBase
from agent import kb_agent, KBDependencies
load_dotenv()
async def main():
kb = KnowledgeBase(persist_dir="./chroma_db")
doc_dir = "./data/documents"
if os.path.exists(doc_dir):
chunk_count = kb.ingest_directory(doc_dir)
print(f"Total chunks ingested: {chunk_count}")
deps = KBDependencies(knowledge_base=kb)
print("\n=== Knowledge Base Chatbot ===")
print("Type 'quit' to exit, 'clear' to reset conversation.\n")
message_history = []
while True:
user_input = input("You: ").strip()
if user_input.lower() == "quit":
print("Goodbye!")
break
if user_input.lower() == "clear":
message_history = []
print("Conversation cleared.\n")
continue
if not user_input:
continue
try:
result = await kb_agent.run(
user_input,
deps=deps,
message_history=message_history
)
# Update history with all messages from this run
message_history = result.all_messages()
response = result.output
print(f"\nBot: {response.answer}")
print(f"Sources: {', '.join(response.sources)}")
print(f"Confidence: {response.confidence:.0%}\n")
except Exception as e:
print(f"\nError: {e}\n")
if __name__ == "__main__":
asyncio.run(main())
By passing message_history to each run and updating it with result.all_messages(), the agent maintains context across turns. Users can now ask follow-up questions like "Can you tell me more about that?" and the model will understand what "that" refers to.
Best Practices
Optimize Your Chunking Strategy
No single chunk size works for all documents. Experiment with different sizes based on your content. For technical documentation, chunk by sections or headers. For narrative text, larger chunks with more overlap tend to work better. Always test retrieval quality with representative queries before deploying.
Handle Edge Cases Gracefully
Not every query will have a good match in your knowledge base. Add logic to detect low-relevance results and respond appropriately. You can modify the search function to return a confidence threshold:
def search_with_threshold(self, query: str, n_results: int = 5, min_relevance: float = 0.3) -> List[Dict]:
results = self.search(query, n_results)
filtered = [r for r in results if (1 - r["distance"]) >= min_relevance]
return filtered
This prevents the chatbot from confidently answering questions when the retrieved context is irrelevant.
Use Streaming for Better UX
For production chatbots, streaming responses improves perceived performance. Pydantic AI supports streaming via run_stream:
async with kb_agent.run_stream(user_input, deps=deps, message_history=message_history) as result:
async for chunk in result.stream_text(delta=True):
print(chunk, end="", flush=True)
Monitor and Log Retrieval Quality
Log the queries, retrieved chunks, and relevance scores to identify gaps in your knowledge base. If users frequently ask questions that return low-relevance results, you may need to add more documents or improve your chunking strategy.
Keep Dependencies Injectable
Always define your dependencies as a dataclass or Pydantic model. This makes testing trivial — you can create a mock knowledge base that returns predetermined results without touching the vector database:
class MockKnowledgeBase:
def search(self, query: str, n_results: int = 5):
return [{"content": "Mock answer content", "source": "test.txt", "distance": 0.1}]
async def test_agent():
deps = KBDependencies(knowledge_base=MockKnowledgeBase())
result = await kb_agent.run("What is the return policy?", deps=deps)
print(result.output)
Re-ingest When Documents Change
ChromaDB persists data to disk, so you don't need to re-ingest on every startup. However, when your source documents change, you should clear the collection and re-ingest. Add a flag or check file modification times to automate this:
def should_reingest(self, doc_dir: str) -> bool:
"""Check if any source files are newer than the last ingestion."""
# Implementation depends on your needs
# Simple approach: store a hash of all files and compare
pass
Conclusion
Building a knowledge base chatbot with Pydantic AI gives you a robust, type-safe, and maintainable foundation for RAG applications. The framework's dependency injection, structured outputs, and tool system make it straightforward to connect an LLM to your own data while keeping the code clean and testable. By combining ChromaDB for vector storage, sentence-transformers for local embeddings, and Pydantic AI's agent system, you get a complete pipeline from document ingestion to grounded responses. As you refine your implementation, focus on chunking strategy, retrieval quality monitoring, and graceful handling of edge cases — these are the factors that separate a toy demo from a production-ready assistant. With the patterns covered in this tutorial, you're well-equipped to extend the chatbot with features like multi-document sources, hybrid search, user feedback loops, and deployment to web or messaging platforms.