> ## Documentation Index
> Fetch the complete documentation index at: https://thethirdpenco-feat-tool-lifecycle-events.mintlify.site/llms.txt
> Use this file to discover all available pages before exploring further.

# 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.

```python theme={null}
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:

```python theme={null}
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:

```python theme={null}
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:

```python theme={null}
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:

```python theme={null}
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:

```python theme={null}
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

<AccordionGroup>
  <Accordion title="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.
  </Accordion>

  <Accordion title="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:

    ```python theme={null}
    result = await embed(
        model=openai.embedding("text-embedding-3-small"),
        value="text",
        dimensions=256
    )
    ```
  </Accordion>

  <Accordion title="Chunk long documents">
    Split long documents into smaller chunks (e.g., 500-1000 tokens) before
    embedding. This improves retrieval precision.
  </Accordion>

  <Accordion title="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.
  </Accordion>

  <Accordion title="Add source metadata">
    Include document titles, URLs, and IDs so the AI can cite sources.
  </Accordion>

  <Accordion title="Consider hybrid search">
    Combine vector similarity with keyword matching for better results,
    especially for exact term matching.
  </Accordion>
</AccordionGroup>

## Next Steps

<CardGroup cols={2}>
  <Card title="Embeddings Guide" icon="cube" href="/core/embeddings">
    Learn more about generating embeddings
  </Card>

  <Card title="embed() Reference" icon="code" href="/reference/embed">
    Full API reference for embed functions
  </Card>
</CardGroup>