0Pricing
AI Engineering Academy · Lesson

Implementing Core RAG and Agent Features

Build the document ingestion pipeline, vector store indexing, retrieval with re-ranking, and agent tool integrations following the patterns learned throughout the track.

Implementing Core RAG and Agent Features is a free AI Engineering Academy lesson on CoddyKit — lesson 2 of 4. You can read the complete lesson below for free — then practise it hands-on in the browser with a built-in code editor and a 24/7 AI tutor. It is part of the AI Engineering Academy learning path, one of 4 lessons in the course, and your progress syncs across the web and the CoddyKit app.

Implementation Order and Dependencies

Build your system bottom-up: start with the components that have no external dependencies, then layer on components that depend on them. For a RAG + agent system, the order is: (1) vector store schema, (2) ingestion pipeline, (3) retriever, (4) basic Q&A chain, (5) streaming endpoint, (6) agent with tools, (7) caching layer, (8) tracing instrumentation. Test each component in isolation before integrating it into the pipeline.

# Build order:
IMPL_ORDER = [
    'pgvector_schema',      # prerequisite for everything
    'document_ingestion',   # populate the vector store
    'hybrid_retriever',     # test retrieval in isolation
    'qa_chain_basic',       # integrate LLM with retrieval
    'streaming_endpoint',   # expose via API
    'function_calling',     # add agent tool calls
    'semantic_cache',       # reduce repeat API calls
    'langsmith_tracing',    # add after core works
    'injection_filter',     # harden before load testing
]

Setting Up the Vector Store

Create the pgvector table with the right schema before ingesting documents. Include columns for the vector, the chunk text, all metadata fields, and an updated_at timestamp for selective re-indexing. Create a vector index (HNSW or IVFFlat) on the embedding column immediately — adding the index after millions of rows is much slower than adding it upfront on an empty table.

-- PostgreSQL schema with pgvector
CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE document_chunks (
    id          BIGSERIAL PRIMARY KEY,
    doc_id      TEXT NOT NULL,
    chunk_text  TEXT NOT NULL,
    embedding   VECTOR(1536) NOT NULL,
    source_file TEXT,
    page_number INT,
    section     TEXT,
    doc_type    TEXT,
    tenant_id   TEXT NOT NULL,  -- for data isolation
    created_at  TIMESTAMPTZ DEFAULT NOW()
);

CREATE INDEX idx_chunks_hnsw ON document_chunks
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

CREATE INDEX idx_chunks_tenant ON document_chunks(tenant_id);

Building the Ingestion Pipeline

Implement ingestion as a single async function that accepts a file path or URL and returns the number of chunks indexed. Use LangChain's document loaders for different file types and the recursive character text splitter for chunking. Batch embedding calls to stay within the 2048-token input limit and to reduce API calls from thousands to dozens.

from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from openai import AsyncOpenAI

async def ingest_document(file_path: str, doc_id: str, tenant_id: str) -> int:
    loader = PyPDFLoader(file_path)
    pages = loader.load()
    splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
    chunks = splitter.split_documents(pages)
    # Batch embed
    texts = [c.page_content for c in chunks]
    client = AsyncOpenAI()
    embeddings_response = await client.embeddings.create(
        model='text-embedding-3-small',
        input=texts
    )
    embeddings = [e.embedding for e in embeddings_response.data]
    await insert_chunks_to_pgvector(chunks, embeddings, doc_id, tenant_id)
    return len(chunks)

Building the Hybrid Retriever

Combine dense vector search with BM25 keyword search and merge results using reciprocal rank fusion. Implement the retriever as a class with a single retrieve(query, tenant_id, top_k) method. Inside, run both searches concurrently with asyncio.gather, merge the ranked lists with RRF, deduplicate by chunk ID, and return the top_k results with their source metadata.

import asyncio
from rank_bm25 import BM25Okapi

class HybridRetriever:
    def __init__(self, pool, k_rrf: int = 60):
        self.pool = pool
        self.k_rrf = k_rrf

    async def retrieve(self, query: str, tenant_id: str, top_k: int = 10) -> list:
        dense_results, sparse_results = await asyncio.gather(
            self._dense_search(query, tenant_id, top_k * 3),
            self._bm25_search(query, tenant_id, top_k * 3)
        )
        merged = self._rrf_merge(dense_results, sparse_results)
        return merged[:top_k]

    def _rrf_score(self, rank: int) -> float:
        return 1.0 / (self.k_rrf + rank + 1)

Implementing the Core QA Chain

Build the core Q&A chain using LangChain LCEL. The chain takes a question and retrieved chunks, formats an augmented prompt with instructions to use only the provided context, and streams the response. Add explicit instructions for the model to cite sources by document name and page number, and to say 'I don't know' when the answer is not in the retrieved context.

from langchain.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain.schema.output_parser import StrOutputParser

RAG_PROMPT = ChatPromptTemplate.from_messages([
    ('system', 'You are a precise assistant. Answer ONLY using the provided context. '
               'Cite sources as [DocName, p.N]. If the answer is not in the context, say "I don\'t have information about that."'),
    ('user', 'Context:\n{context}\n\nQuestion: {question}')
])

llm = ChatOpenAI(model='gpt-4o', temperature=0, streaming=True)
qa_chain = RAG_PROMPT | llm | StrOutputParser()

async def answer_question(question: str, chunks: list) -> str:
    context = '\n\n'.join(f'[{c["source"]}]\n{c["text"]}' for c in chunks)
    return await qa_chain.ainvoke({'context': context, 'question': question})

Adding Function Calling to the Agent

Extend the base Q&A system with function calling to handle queries that require real-time data or computation. Define tools for: searching the web for current information, running SQL queries against a safe read-only database, and looking up specific records by ID. The agent decides which tools to call based on the question, executes them, and incorporates results into the final answer.

from openai import AsyncOpenAI
import json

TOOLS = [
    {
        'type': 'function',
        'function': {
            'name': 'search_knowledge_base',
            'description': 'Search the internal document knowledge base for relevant information',
            'parameters': {
                'type': 'object',
                'properties': {
                    'query': {'type': 'string', 'description': 'Search query'},
                    'top_k': {'type': 'integer', 'default': 5}
                },
                'required': ['query']
            }
        }
    }
]

async def agent_with_tools(question: str, tenant_id: str) -> str:
    client = AsyncOpenAI()
    messages = [{'role': 'user', 'content': question}]
    while True:
        resp = await client.chat.completions.create(
            model='gpt-4o', messages=messages, tools=TOOLS)
        if resp.choices[0].finish_reason != 'tool_calls':
            return resp.choices[0].message.content
        tool_call = resp.choices[0].message.tool_calls[0]
        args = json.loads(tool_call.function.arguments)
        result = await dispatch_tool(tool_call.function.name, args, tenant_id)
        messages.append({'role': 'tool', 'tool_call_id': tool_call.id, 'content': result})

Streaming the Agent Response

Wrap the agent in a FastAPI StreamingResponse using server-sent events so the frontend displays tokens as they arrive. For agent responses with tool calls, stream a progress indicator while the tool executes ('Searching knowledge base...'), then stream the final answer token by token. This prevents the page from appearing frozen during tool execution latency.

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import asyncio

app = FastAPI()

async def event_stream(question: str, tenant_id: str):
    yield 'data: {"type": "start"}\n\n'
    chunks = await retriever.retrieve(question, tenant_id)
    yield 'data: {"type": "retrieving", "count": ' + str(len(chunks)) + '}\n\n'
    async for token in qa_chain.astream({'context': format_context(chunks), 'question': question}):
        yield f'data: {{"type": "token", "content": {repr(token)}}}\n\n'
    yield 'data: {"type": "done"}\n\n'

@app.post('/query/stream')
async def stream_query(question: str, tenant_id: str):
    return StreamingResponse(event_stream(question, tenant_id), media_type='text/event-stream')

Wiring in the Semantic Cache

Add the semantic cache as a pre-retrieval step in the query pipeline. Before hitting the retriever and LLM, embed the user question and check the cache. On a cache hit with similarity above the threshold, return the cached answer immediately with a cached: true flag. Cache misses proceed through the full pipeline, and the resulting answer is stored in the cache for future similar queries.

async def query_pipeline(question: str, tenant_id: str) -> dict:
    # 1. Check semantic cache
    cache_hit = await semantic_cache.lookup(question, tenant_id, threshold=0.92)
    if cache_hit:
        return {'answer': cache_hit.answer, 'cached': True, 'sources': cache_hit.sources}

    # 2. Retrieve
    chunks = await retriever.retrieve(question, tenant_id, top_k=5)
    chunks = await reranker.rerank(question, chunks, top_n=3)

    # 3. Generate
    answer = await answer_question(question, chunks)
    sources = [c['source'] for c in chunks]

    # 4. Cache result
    await semantic_cache.store(question, tenant_id, answer, sources)

    return {'answer': answer, 'cached': False, 'sources': sources}

Adding LangSmith Tracing

Instrument the pipeline with LangSmith by setting two environment variables. Every LangChain chain call is automatically traced with token counts, latency, inputs, outputs, and any errors. For custom non-LangChain code (retrieval, reranking), wrap calls in @traceable decorators to include them in the trace. This gives complete end-to-end visibility into every pipeline step.

import os
from langsmith import traceable

os.environ['LANGCHAIN_TRACING_V2'] = 'true'
os.environ['LANGCHAIN_API_KEY'] = os.environ['LANGSMITH_API_KEY']
os.environ['LANGCHAIN_PROJECT'] = 'document-qa-production'

# Wrap non-LangChain steps with @traceable
@traceable(name='hybrid_retrieval')
async def traced_retrieval(question: str, tenant_id: str, top_k: int) -> list:
    return await retriever.retrieve(question, tenant_id, top_k)

@traceable(name='cohere_reranking')
async def traced_reranking(question: str, chunks: list) -> list:
    return await reranker.rerank(question, chunks)

# LangChain LCEL chains are automatically traced — no extra code needed

Running Integration Tests

Write integration tests that exercise the complete query pipeline from question input to final answer. Use a small test vector store with known documents so you can write deterministic assertions about which chunks should be retrieved and what the answer should contain. Run integration tests against a staging environment that mirrors production infrastructure but uses a separate vector store and LLM key.

import pytest

@pytest.mark.asyncio
async def test_full_pipeline_returns_grounded_answer():
    # Setup: ingest known document
    await ingest_document('tests/fixtures/policy.pdf', 'policy_v1', 'test_tenant')

    # Query with a question that has a known answer in the document
    result = await query_pipeline(
        question='What is the cancellation policy?',
        tenant_id='test_tenant'
    )

    assert result['answer'] is not None
    assert len(result['answer']) > 50
    assert '24 hours' in result['answer']  # known fact in document
    assert 'policy.pdf' in str(result['sources'])
    assert result['cached'] is False  # fresh query

Measuring Baseline Retrieval Quality

Before optimizing anything, establish a retrieval baseline. Use your eval test set to measure hit rate (does the correct document appear in top-5 results?) and MRR (how high does it rank?). Run this baseline after setting up the initial vector store, before adding reranking or hybrid search. The baseline tells you what improvements each optimization actually delivers so you have evidence for the techniques worth keeping.

async def measure_retrieval_baseline(test_cases: list) -> dict:
    hits = 0
    reciprocal_ranks = []
    for case in test_cases:
        results = await retriever.retrieve(case['question'], case['tenant_id'], top_k=5)
        result_docs = [r['doc_id'] for r in results]
        if case['relevant_doc'] in result_docs:
            hits += 1
            rank = result_docs.index(case['relevant_doc']) + 1
            reciprocal_ranks.append(1.0 / rank)
        else:
            reciprocal_ranks.append(0.0)
    return {
        'hit_rate_at_5': hits / len(test_cases),
        'mrr': sum(reciprocal_ranks) / len(reciprocal_ranks)
    }

Quick Check

Test your understanding of building core RAG and agent features.

Lesson Recap

In this lesson you learned: bottom-up implementation order ensures each component is tested in isolation before integration, hybrid retrieval with concurrent dense and BM25 search combined through RRF gives the best retrieval quality, and LangSmith tracing with @traceable decorators provides complete pipeline visibility. Next up we harden the system with security, caching, and reliability patterns.

Frequently asked questions

Is the “Implementing Core RAG and Agent Features” lesson free?

Yes — the full text of “Implementing Core RAG and Agent Features” is free to read here on the web, and the AI Engineering Academy course includes 4 lessons in total. To practise it interactively (a built-in code editor and a 24/7 AI tutor) and unlock the rest of the AI Engineering Academy course, upgrade to CoddyKit PRO.

What will I learn in “Implementing Core RAG and Agent Features”?

Build the document ingestion pipeline, vector store indexing, retrieval with re-ranking, and agent tool integrations following the patterns learned throughout the track. You practise AI Engineering Academy with hands-on code you run directly in the browser, and a 24/7 AI tutor answers your questions as you work through the lesson.

Do I need any experience to start AI Engineering Academy?

No prior experience is required. AI Engineering Academy on CoddyKit is structured for beginners through advanced learners; this is — lesson 2 of 4, so you can start here or from the beginning and move at your own pace.

How long does the “Implementing Core RAG and Agent Features” lesson take?

Most CoddyKit lessons take about 5–10 minutes. Each one is bite-sized and interactive, so you make steady progress and pick up exactly where you left off across the web and the app.

Can I write and run code in this AI Engineering Academy lesson?

Yes. Every AI Engineering Academy lesson includes a built-in code editor, so you write and run real code right in your browser and get instant AI feedback — no local setup required.

All lessons in this course

  1. Designing the Production Architecture
  2. Implementing Core RAG and Agent Features
  3. Hardening: Security, Caching, and Reliability
  4. Evaluation, Deployment, and Retrospective
← Back to AI Engineering Academy