RAG Agent

┌──────────────┐     ┌──────────────┐     ┌──────────────┐
│  Documents   │────▶│   Chunking   │────▶│  Embeddings  │
└──────────────┘     └──────────────┘     └──────────────┘
                                                  │
                                                  ▼
┌──────────────┐     ┌──────────────┐     ┌──────────────┐
│    Answer    │◀────│     LLM      │◀────│   Qdrant     │
└──────────────┘     └──────────────┘     └──────────────┘
                            ▲
                            │
                     ┌──────────────┐
                     │   Question   │
                     └──────────────┘

from isa_agent_sdk import Agent
from isa_agent_sdk.rag import DocumentStore
 
# Create a document store
store = DocumentStore(
    collection="my-docs",
    embedding_model="text-embedding-3-small"
)
 
# Ingest documents
await store.ingest("./documents/")
 
# Create RAG agent
agent = Agent(
    name="rag-agent",
    model="claude-sonnet-4-20250514",
    tools=[store.as_tool()]
)
 
# Ask questions
response = await agent.run("What does the documentation say about authentication?")

# Ingest a directory of documents
await store.ingest(
    "./documents/",
    file_types=["pdf", "md", "txt", "docx"],
    metadata_extractor=lambda doc: {
        "source": doc.filename,
        "category": doc.parent_folder
    }
)
 
# Ingest specific files
await store.ingest_file("./important-doc.pdf", metadata={"priority": "high"})
 
# Ingest from URL
await store.ingest_url(
    "https://docs.example.com/api-reference",
    metadata={"type": "api-docs"}
)
 
# Check ingestion status
stats = await store.stats()
print(f"Documents: {stats.document_count}")
print(f"Chunks: {stats.chunk_count}")
print(f"Vector dimensions: {stats.dimensions}")

from isa_agent_sdk import Agent
 
# Create the retrieval tool
retrieval_tool = store.as_tool(
    name="search_docs",
    description="Search the company documentation for relevant information",
    top_k=5,  # Number of chunks to retrieve
    score_threshold=0.7  # Minimum similarity score
)
 
# Create the agent
agent = Agent(
    name="docs-assistant",
    model="claude-sonnet-4-20250514",
    tools=[retrieval_tool],
    system_prompt="""You are a documentation assistant. Answer questions based on the retrieved documentation.
 
    Guidelines:
    - Always search the docs before answering
    - Cite your sources with document names
    - If information isn't in the docs, say so
    - Be precise and accurate
    """
)

from isa_agent_sdk import Agent
from isa_agent_sdk.rag import CitationFormatter
 
agent = Agent(
    name="rag-with-citations",
    model="claude-sonnet-4-20250514",
    tools=[store.as_tool()],
    system_prompt="""Answer questions using the documentation.
    Always cite sources using [Source: filename] format."""
)
 
response = await agent.run("How do I configure authentication?")
print(response.content)
# Output: "To configure authentication, set the AUTH_SECRET environment variable... [Source: auth-guide.md]"
 
# Get structured citations
citations = CitationFormatter.extract(response)
for cite in citations:
    print(f"- {cite.source}: {cite.text}")

from isa_agent_sdk.rag import HybridSearch, Reranker
 
# Hybrid search (vector + keyword)
hybrid_tool = store.as_tool(
    search_strategy=HybridSearch(
        vector_weight=0.7,
        keyword_weight=0.3
    )
)
 
# With reranking for better accuracy
reranked_tool = store.as_tool(
    top_k=20,  # Retrieve more initially
    reranker=Reranker(
        model="cross-encoder",
        top_n=5  # Return top 5 after reranking
    )
)
 
# Filtered search
filtered_tool = store.as_tool(
    filters={
        "category": "api-docs",
        "updated_after": "2025-01-01"
    }
)

from isa_agent_sdk import Agent
from isa_agent_sdk.rag import (
    DocumentStore,
    ChunkingStrategy,
    HybridSearch,
    Reranker,
    CitationFormatter
)
import logging
 
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
 
class RAGAssistant:
    def __init__(self, collection: str):
        # Initialize document store
        self.store = DocumentStore(
            vector_db="qdrant",
            collection=collection,
            embedding_model="text-embedding-3-small",
            chunking=ChunkingStrategy(
                method="semantic",
                chunk_size=512,
                chunk_overlap=50
            )
        )
 
        # Create retrieval tool with hybrid search
        self.retrieval_tool = self.store.as_tool(
            name="search_knowledge_base",
            description="Search the knowledge base for relevant information",
            search_strategy=HybridSearch(
                vector_weight=0.7,
                keyword_weight=0.3
            ),
            top_k=10,
            reranker=Reranker(model="cross-encoder", top_n=5)
        )
 
        # Create the agent
        self.agent = Agent(
            name="rag-assistant",
            model="claude-sonnet-4-20250514",
            tools=[self.retrieval_tool],
            system_prompt="""You are a knowledgeable assistant with access to a documentation database.
 
            Instructions:
            1. Always search the knowledge base before answering
            2. Base your answers strictly on retrieved information
            3. Cite sources using [Source: document_name] format
            4. If the information isn't available, clearly state that
            5. Provide accurate, helpful responses
 
            Remember: accuracy is more important than completeness.
            """
        )
 
    async def ingest_documents(self, path: str):
        """Ingest documents from a directory."""
        logger.info(f"Ingesting documents from {path}")
        await self.store.ingest(path)
        stats = await self.store.stats()
        logger.info(f"Ingested {stats.chunk_count} chunks from {stats.document_count} documents")
 
    async def ask(self, question: str) -> dict:
        """Ask a question and get an answer with citations."""
        try:
            response = await self.agent.run(question)
 
            # Extract citations
            citations = CitationFormatter.extract(response)
 
            return {
                "answer": response.content,
                "citations": [
                    {"source": c.source, "text": c.text}
                    for c in citations
                ],
                "tokens_used": response.usage.total_tokens
            }
        except Exception as e:
            logger.error(f"Error answering question: {e}")
            return {
                "answer": "I encountered an error processing your question.",
                "citations": [],
                "error": str(e)
            }
 
    async def stream_answer(self, question: str):
        """Stream the answer for real-time display."""
        async for chunk in self.agent.stream(question):
            yield chunk.content
 
# Usage
async def main():
    # Initialize
    assistant = RAGAssistant(collection="company-kb")
 
    # Ingest documents
    await assistant.ingest_documents("./docs/")
 
    # Ask questions
    result = await assistant.ask("What are the API rate limits?")
    print(f"Answer: {result['answer']}")
    print(f"Sources: {[c['source'] for c in result['citations']]}")
 
    # Stream response
    print("\nStreaming: ", end="")
    async for chunk in assistant.stream_answer("How do I authenticate?"):
        print(chunk, end="", flush=True)
 
if __name__ == "__main__":
    import asyncio
    asyncio.run(main())