RAG agent example walkthrough

START -> RETRIEVE -> [RERANK] -> SYNTHESIZE -> END

pip install "10xgraph[openai]"
export OPENAI_API_KEY="sk-..."

import asyncio

from tenxgraph.core.graph.agent import Agent
from tenxgraph.core.state.message import Message
from tenxgraph.prebuilt.agent.rag import RAGAgent
from tenxgraph.storage.store.base_store import BaseStore
from tenxgraph.storage.store.store_schema import MemorySearchResult

# The knowledge base: a few support-policy chunks
DOCS = [
    "Our refund policy allows returns within 30 days of purchase with a receipt.",
    "Products must be in original packaging to qualify for a full refund.",
    "Digital downloads are non-refundable once accessed.",
    "Shipping costs are non-refundable unless the return is due to our error.",
    "Contact [email protected] with your order number to start a return.",
]


class StubStore(BaseStore):
    """Demo store: ranks chunks by word overlap with the query."""

    async def asetup(self):
        pass

    async def asearch(self, config, query: str, limit: int = 5, **kwargs):
        query_words = set(query.lower().split())
        scored = sorted(
            ((len(query_words & set(chunk.lower().split())), chunk) for chunk in DOCS),
            reverse=True,
        )
        return [
            MemorySearchResult(content=chunk, score=float(score))
            for score, chunk in scored[:limit]
            if score > 0
        ]

    # RAGAgent only calls asearch. The rest satisfy the abstract BaseStore API.
    async def astore(self, config, content, **kwargs):
        return "id"

    async def aget(self, config, memory_id, **kwargs):
        return None

    async def aget_all(self, config, **kwargs):
        return []

    async def aupdate(self, config, memory_id, content, **kwargs):
        pass

    async def adelete(self, config, memory_id, **kwargs):
        pass

    async def aforget_memory(self, config, **kwargs):
        pass


# The agent that writes the answer from the retrieved context
answerer = Agent(
    model="gpt-4o-mini",
    provider="openai",
    system_prompt=[
        {
            "role": "system",
            "content": (
                "You are a customer-support agent. "
                "Answer using ONLY the provided context. "
                "If the context does not contain the answer, say so."
            ),
        }
    ],
)

# Retrieve the 4 best chunks for each question
rag = RAGAgent(store=StubStore(), agent=answerer, top_k=4)
app = rag.compile()


async def main():
    question = "Can I return a digital product?"
    result = await app.ainvoke(
        {"messages": [Message.text_message(question, role="user")]},
        config={"thread_id": "demo"},
    )
    answer = next(m for m in reversed(result["messages"]) if m.role == "assistant")
    print(f"Q: {question}")
    print(f"A: {answer.text()}")


if __name__ == "__main__":
    asyncio.run(main())

pip install "10xgraph[openai,qdrant]"

import asyncio

from tenxgraph.core.graph.agent import Agent
from tenxgraph.core.state.message import Message
from tenxgraph.prebuilt.agent.rag import RAGAgent
from tenxgraph.storage import create_local_qdrant_store
from tenxgraph.storage.store.embedding import OpenAIEmbedding

# Replace with your own chunks (split by headings or paragraphs, not fixed size)
CHUNKS = [
    "Digital downloads are non-refundable once accessed.",
    "Returns are accepted within 30 days with a receipt.",
]

# Search and indexing must use the same user_id, because QdrantStore filters by it
CONFIG = {"user_id": "kb"}

store = create_local_qdrant_store(
    path="./knowledge_base",
    embedding=OpenAIEmbedding(model="text-embedding-3-small"),
)


async def main():
    # Index once; skip this step on later runs
    for chunk in CHUNKS:
        await store.astore(config=CONFIG, content=chunk)

    rag = RAGAgent(
        store=store,
        agent=Agent(model="gpt-4o-mini", provider="openai"),
        top_k=5,
        store_config=CONFIG,  # passed to every store.asearch call
    )
    app = rag.compile()
    result = await app.ainvoke(
        {"messages": [Message.text_message("Can I return a download?", role="user")]},
        config={"thread_id": "kb-demo"},
    )
    print(result["messages"][-1].text())


asyncio.run(main())

import os

from tenxgraph.core.graph.agent import Agent
from tenxgraph.prebuilt.agent.rag import CohereReranker, CrossEncoderReranker, RAGAgent

# Hosted reranking: 20 candidates in, best 5 go to the model
rag = RAGAgent(
    store=store,  # the QdrantStore from the previous section
    agent=Agent(model="gpt-4o", provider="openai"),
    reranker=CohereReranker(api_key=os.environ["COHERE_API_KEY"]),
    top_k=20,
    top_n=5,
)

# Local reranking: 15 candidates in, best 4 go to the model
rag = RAGAgent(
    store=store,
    agent=Agent(model="gpt-4o-mini", provider="openai"),
    reranker=CrossEncoderReranker(),
    top_k=15,
    top_n=4,
)

import logging

from tenxgraph.utils import ResponseGranularity

# Logs the query, candidate count and rerank step
logging.getLogger("tenxgraph.prebuilt.rag").setLevel(logging.DEBUG)

# Inside an async function, with `app` compiled as above
result = await app.ainvoke(
    {"messages": [Message.text_message("Can I return a download?", role="user")]},
    config={"thread_id": "debug"},
    response_granularity=ResponseGranularity.FULL,
)
docs = result["state"].execution_meta.internal_data.get("rag_docs", [])
print(f"{len(docs)} documents reached the model")

from tenxgraph.core.state.message import Message
from tenxgraph.storage.checkpointer import InMemoryCheckpointer

app = rag.compile(checkpointer=InMemoryCheckpointer())
config = {"thread_id": "user_123"}

# Inside an async function
first = await app.ainvoke(
    {"messages": [Message.text_message("What is the refund window?", role="user")]},
    config=config,
)
second = await app.ainvoke(
    {"messages": [Message.text_message("Does that apply to downloads?", role="user")]},
    config=config,
)