RAG Agent

Build a retrieval-augmented generation agent with embeddings and search tools

Build an AI agent that retrieves relevant information before answering, reducing hallucinations and providing accurate, sourced responses.

The Pattern

Give the AI a search tool backed by vector embeddings to retrieve semantically similar information, then let it synthesize a response from the results.

from ai_query import generate_text, embed, embed_many
from ai_query.providers import openai, tool, Field
import numpy as np

# Pre-compute document embeddings
documents = [
    "Our refund policy allows returns within 30 days of purchase.",
    "Shipping takes 3-5 business days for domestic orders.",
    "Premium members get free shipping on all orders.",
]

# Embed documents at startup
doc_embeddings = None

async def init_embeddings():
    global doc_embeddings
    result = await embed_many(
        model=openai.embedding("text-embedding-3-small"),
        values=documents
    )
    doc_embeddings = result.embeddings

def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

async def search_documents(query: str, top_k: int = 2) -> list[str]:
    # Embed the query
    query_result = await embed(
        model=openai.embedding("text-embedding-3-small"),
        value=query
    )

    # Find most similar documents
    similarities = [
        (i, cosine_similarity(query_result.embedding, doc_emb))
        for i, doc_emb in enumerate(doc_embeddings)
    ]
    similarities.sort(key=lambda x: x[1], reverse=True)

    return [documents[i] for i, _ in similarities[:top_k]]

@tool(description="Search the knowledge base for relevant information")
async def search(query: str = Field(description="Search query")) -> str:
    results = await search_documents(query)
    return "\n\n".join([f"[Result {i+1}]: {r}" for i, r in enumerate(results)])

# Initialize embeddings before use
await init_embeddings()

result = await generate_text(
    model=openai("gpt-4o"),
    system="""You are a helpful assistant with access to a knowledge base.
    Always search for information before answering questions.
    Cite your sources in your response.""",
    prompt="What is your refund policy?",
    tools={"search": search},
)
print(result.text)

Complete Example: Documentation Assistant

A full RAG agent with vector search that answers questions about your documentation:

import asyncio
import numpy as np
from ai_query import generate_text, embed, embed_many
from ai_query.providers import openai, tool, Field, step_count_is

# Document store with embeddings
class VectorStore:
    def __init__(self):
        self.documents: list[dict] = []
        self.embeddings: list[list[float]] = []
        self.model = openai.embedding("text-embedding-3-small")

    async def add_documents(self, docs: list[dict]):
        """Add documents with title and content."""
        texts = [f"{d['title']}: {d['content']}" for d in docs]
        result = await embed_many(model=self.model, values=texts)

        self.documents.extend(docs)
        self.embeddings.extend(result.embeddings)

    async def search(self, query: str, top_k: int = 3) -> list[dict]:
        """Search for similar documents."""
        query_result = await embed(model=self.model, value=query)

        # Calculate similarities
        similarities = []
        for i, emb in enumerate(self.embeddings):
            sim = np.dot(query_result.embedding, emb) / (
                np.linalg.norm(query_result.embedding) * np.linalg.norm(emb)
            )
            similarities.append((sim, self.documents[i]))

        # Return top-k results
        similarities.sort(reverse=True, key=lambda x: x[0])
        return [doc for _, doc in similarities[:top_k]]


class DocsAssistant:
    def __init__(self, store: VectorStore):
        self.store = store

    async def answer(self, question: str) -> str:
        @tool(description="Search the documentation for relevant information")
        async def search_docs(
            query: str = Field(description="Search query for documentation"),
        ) -> str:
            results = await self.store.search(query)
            if not results:
                return "No relevant documentation found."
            return "\n\n".join([
                f"**{doc['title']}**: {doc['content']}"
                for doc in results
            ])

        result = await generate_text(
            model=openai("gpt-4o"),
            system="""You are a documentation assistant for ai-query.

            Instructions:
            1. Search the docs to find relevant information
            2. Answer the user's question based on the search results
            3. If you can't find the answer, say so
            4. Cite which docs you used""",
            prompt=question,
            tools={"search_docs": search_docs},
            stop_when=step_count_is(3),
        )

        return result.text


async def main():
    # Initialize vector store
    store = VectorStore()

    # Add documentation
    await store.add_documents([
        {"title": "Installation", "content": "To install ai-query, run: pip install ai-query"},
        {"title": "Quickstart", "content": "Import with: from ai_query import generate_text"},
        {"title": "Streaming", "content": "Use stream_text() for streaming responses"},
        {"title": "Tools", "content": "Define tools with @tool decorator and Field() for parameters"},
        {"title": "Agents", "content": "Use Agent class for stateful, persistent AI agents"},
        {"title": "Embeddings", "content": "Use embed() and embed_many() for vector embeddings"},
    ])

    assistant = DocsAssistant(store)

    # Ask questions
    questions = [
        "How do I install this library?",
        "How do I create streaming responses?",
        "How do I generate embeddings?",
    ]

    for q in questions:
        print(f"\n**Q: {q}**")
        answer = await assistant.answer(q)
        print(f"A: {answer}\n")
        print("-" * 50)

asyncio.run(main())

Using Different Embedding Providers

You can use any supported embedding provider:

from ai_query import embed, embed_many
from ai_query.providers import openai, google

# OpenAI embeddings (1536 dimensions)
openai_result = await embed(
    model=openai.embedding("text-embedding-3-small"),
    value="Hello world"
)

# OpenAI with reduced dimensions (faster, less storage)
openai_small = await embed(
    model=openai.embedding("text-embedding-3-small"),
    value="Hello world",
    dimensions=256
)

# Google embeddings (flexible dimensions, default 3072)
google_result = await embed(
    model=google.embedding("gemini-embedding-001"),
    value="Hello world"
)

Multi-Source RAG

Search multiple sources and combine results:

from ai_query import generate_text, embed
from ai_query.providers import openai, tool, Field, step_count_is
import numpy as np

# Multiple vector stores for different sources
docs_store = VectorStore()  # Documentation
tickets_store = VectorStore()  # Support tickets
faq_store = VectorStore()  # FAQ

@tool(description="Search internal documentation")
async def search_docs(query: str = Field(description="Search query")) -> str:
    results = await docs_store.search(query)
    return "\n".join([f"[DOC] {r['title']}: {r['content']}" for r in results])

@tool(description="Search support tickets and past issues")
async def search_tickets(query: str = Field(description="Search query")) -> str:
    results = await tickets_store.search(query)
    return "\n".join([f"[TICKET] {r['title']}: {r['content']}" for r in results])

@tool(description="Search the FAQ database")
async def search_faq(query: str = Field(description="Search query")) -> str:
    results = await faq_store.search(query)
    return "\n".join([f"[FAQ] {r['title']}: {r['content']}" for r in results])

result = await generate_text(
    model=openai("gpt-4o"),
    system="""You are a support agent. Use the available search tools to find
    relevant information before answering. You can search multiple sources.""",
    prompt="How do I reset my password?",
    tools={
        "search_docs": search_docs,
        "search_tickets": search_tickets,
        "search_faq": search_faq,
    },
    stop_when=step_count_is(5),
)

With Citations

Track which documents were used and include citations:

from ai_query import generate_text, embed
from ai_query.providers import openai, tool, Field, has_tool_call
import numpy as np

class CitableVectorStore:
    def __init__(self):
        self.documents = []
        self.embeddings = []
        self.model = openai.embedding("text-embedding-3-small")

    async def add(self, doc_id: str, title: str, content: str, url: str):
        result = await embed(model=self.model, value=content)
        self.documents.append({
            "id": doc_id,
            "title": title,
            "content": content,
            "url": url
        })
        self.embeddings.append(result.embedding)

    async def search(self, query: str, top_k: int = 3) -> list[dict]:
        query_result = await embed(model=self.model, value=query)

        similarities = [
            (np.dot(query_result.embedding, emb) /
             (np.linalg.norm(query_result.embedding) * np.linalg.norm(emb)), doc)
            for emb, doc in zip(self.embeddings, self.documents)
        ]
        similarities.sort(reverse=True, key=lambda x: x[0])
        return [doc for _, doc in similarities[:top_k]]


store = CitableVectorStore()
citations = []

@tool(description="Search and retrieve relevant documents")
async def retrieve(query: str = Field(description="Search query")) -> str:
    results = await store.search(query)
    for doc in results:
        citations.append({"id": doc["id"], "title": doc["title"], "url": doc["url"]})
    return "\n".join([f"[{doc['id']}] {doc['content']}" for doc in results])

@tool(description="Submit the final answer with citations")
def answer(
    response: str = Field(description="The answer text"),
    sources: list[str] = Field(description="List of document IDs used"),
) -> str:
    return response

result = await generate_text(
    model=openai("gpt-4o"),
    system="Search for information, then submit your answer with source citations.",
    prompt="How do I configure authentication?",
    tools={"retrieve": retrieve, "answer": answer},
    stop_when=has_tool_call("answer"),
)

# Access final answer and citations
final_answer = result.steps[-1].tool_calls[0].arguments
print(f"Answer: {final_answer['response']}")
print(f"Sources: {final_answer['sources']}")
print(f"Citations: {citations}")

Hybrid Search

Combine keyword and vector search for better results:

from ai_query import embed
from ai_query.providers import openai
import numpy as np

class HybridStore:
    def __init__(self):
        self.documents = []
        self.embeddings = []
        self.model = openai.embedding("text-embedding-3-small")

    async def add_documents(self, docs: list[dict]):
        texts = [d["content"] for d in docs]
        result = await embed_many(model=self.model, values=texts)
        self.documents.extend(docs)
        self.embeddings.extend(result.embeddings)

    def keyword_search(self, query: str, top_k: int = 10) -> list[tuple[int, float]]:
        """Simple keyword matching."""
        query_words = set(query.lower().split())
        scores = []
        for i, doc in enumerate(self.documents):
            doc_words = set(doc["content"].lower().split())
            score = len(query_words & doc_words) / len(query_words)
            scores.append((i, score))
        scores.sort(key=lambda x: x[1], reverse=True)
        return scores[:top_k]

    async def hybrid_search(self, query: str, top_k: int = 5, alpha: float = 0.5):
        """Combine keyword and vector search with alpha weighting."""
        # Keyword scores
        keyword_results = self.keyword_search(query, top_k=len(self.documents))
        keyword_scores = {i: score for i, score in keyword_results}

        # Vector scores
        query_result = await embed(model=self.model, value=query)
        vector_scores = {}
        for i, emb in enumerate(self.embeddings):
            sim = np.dot(query_result.embedding, emb) / (
                np.linalg.norm(query_result.embedding) * np.linalg.norm(emb)
            )
            vector_scores[i] = sim

        # Combine scores
        combined = []
        for i in range(len(self.documents)):
            hybrid_score = (
                alpha * keyword_scores.get(i, 0) +
                (1 - alpha) * vector_scores.get(i, 0)
            )
            combined.append((hybrid_score, self.documents[i]))

        combined.sort(reverse=True, key=lambda x: x[0])
        return [doc for _, doc in combined[:top_k]]

Tips

Pre-compute embeddings for static content

Embed your documents once at startup using embed_many() and cache the results. This avoids re-computing embeddings on every query.

Use reduced dimensions for large collections

OpenAI’s v3 models support custom dimensions. Use 256 or 512 dimensions instead of 1536 to reduce storage and improve search speed:

result = await embed(
    model=openai.embedding("text-embedding-3-small"),
    value="text",
    dimensions=256
)

Chunk long documents

Split long documents into smaller chunks (e.g., 500-1000 tokens) before embedding. This improves retrieval precision.

Let the AI decide what to search

Don’t pre-search—let the AI formulate its own queries based on the user’s question for better results.

Add source metadata

Include document titles, URLs, and IDs so the AI can cite sources.

Consider hybrid search

Combine vector similarity with keyword matching for better results, especially for exact term matching.

Next Steps

Embeddings Guide

Learn more about generating embeddings

embed() Reference

Full API reference for embed functions