RAG

RAG

RAG(검색 증강 생성) 검색 예시를 만들어 볼게요. 이 데모에서는 Logfire 문서의 2024년 10월 스냅샷에 질문을 던질 수 있어요. 문서의 각 섹션을 데이터베이스에 저장해 두고, 검색 도구를 Pydantic AI 에이전트에 등록하는 방식이에요. 마크다운 파일에서 섹션을 추출하는 로직과 해당 데이터가 담긴 JSON 파일은 이 gist에서 확인할 수 있어요.

출처: 공식문서

예시가 보여주는 것

• tools
• agent dependencies
• RAG 검색

준비하기

검색 데이터베이스로는 pgvector가 붙은 PostgreSQL을 써요. pgvector를 내려받아 실행하는 가장 쉬운 방법은 Docker예요.

mkdir postgres-data
docker run --rm \
  -e POSTGRES_PASSWORD=postgres \
  -p 54320:5432 \
  -v `pwd`/postgres-data:/var/lib/postgresql/data \
  pgvector/pgvector:pg17

SQL gen 예시처럼 포트 54320을 써서, 이미 실행 중인 다른 postgres 인스턴스와 충돌하지 않게 해요. 또 PostgreSQL의 data 디렉토리를 로컬에 마운트해서, 컨테이너를 멈췄다 다시 켜도 데이터가 유지되게 했어요.

PostgreSQL이 실행 중이고 의존성 설치와 환경 변수 설정이 끝났다면, 검색 데이터베이스를 만들 수 있어요. (주의: 이 과정은 OPENAI_API_KEY 환경 변수가 필요하고, 문서의 각 섹션에 대한 임베딩을 만들기 위해 OpenAI embedding API를 약 300회 호출해요.)

python -m pydantic_ai_examples.rag build

uv run -m pydantic_ai_examples.rag build

데이터베이스를 만드는 과정은 지금은 Pydantic AI를 쓰지 않고 OpenAI SDK를 직접 사용해요.

임베딩 모델과 인덱스 스키마

이 예시는 문서와 쿼리 모두에 text-embedding-3-small을 사용하고, 그 1,536차원 출력을 vector(1536) 컬럼에 저장해요. 모델이나 차원을 바꾸고 싶다면 PostgreSQL을 멈추고, 예시의 postgres-data 디렉토리를 지우고(로컬 예시 데이터베이스 전체가 사라져요), 필요하면 DB_SCHEMA를 수정한 뒤 PostgreSQL을 다시 시작하고 build를 다시 실행하면 돼요. pgvector의 HNSW vector 인덱스는 최대 2,000차원을 지원해요.

질문하기

검색 데이터베이스가 준비되면 에이전트에 질문할 수 있어요.

python -m pydantic_ai_examples.rag search "How do I configure logfire to work with FastAPI?"

uv run -m pydantic_ai_examples.rag search "How do I configure logfire to work with FastAPI?"

예시 코드

rag.py

from __future__ import annotations as _annotations

import asyncio
import re
import sys
import unicodedata
from contextlib import asynccontextmanager
from dataclasses import dataclass

import asyncpg
import httpx
import logfire
import pydantic_core
from anyio import create_task_group
from openai import AsyncOpenAI
from pydantic import TypeAdapter
from typing_extensions import AsyncGenerator

from pydantic_ai import Agent, RunContext

# 'if-token-present' means nothing will be sent (and the example will work) if you don't have logfire configured
logfire.configure(send_to_logfire='if-token-present')
logfire.instrument_asyncpg()
logfire.instrument_pydantic_ai()


@dataclass
class Deps:
    openai: AsyncOpenAI
    pool: asyncpg.Pool


agent = Agent('openai:gpt-5.2', deps_type=Deps)


@agent.tool
async def retrieve(context: RunContext[Deps], search_query: str) -> str:
    """Retrieve documentation sections based on a search query.

    Args:
        context: The call context.
        search_query: The search query.
    """
    with logfire.span(
        'create embedding for {search_query=}', search_query=search_query
    ):
        embedding = await context.deps.openai.embeddings.create(
            input=search_query,
            model='text-embedding-3-small',
        )

    assert len(embedding.data) == 1, (
        f'Expected 1 embedding, got {len(embedding.data)}, doc query: {search_query!r}'
    )
    embedding = embedding.data[0].embedding
    embedding_json = pydantic_core.to_json(embedding).decode()
    rows = await context.deps.pool.fetch(
        'SELECT url, title, content FROM doc_sections ORDER BY embedding <-> $1 LIMIT 8',
        embedding_json,
    )
    return '\n\n'.join(
        f'# {row["title"]}\nDocumentation URL:{row["url"]}\n\n{row["content"]}\n'
        for row in rows
    )


async def run_agent(question: str):
    """Entry point to run the agent and perform RAG based question answering."""
    openai = AsyncOpenAI()
    logfire.instrument_openai(openai)

    logfire.info('Asking "{question}"', question=question)

    async with database_connect(False) as pool:
        deps = Deps(openai=openai, pool=pool)
        answer = await agent.run(question, deps=deps)
    print(answer.output)


#######################################################
# The rest of this file is dedicated to preparing the #
# search database, and some utilities.                #
#######################################################

# JSON document from
# https://gist.github.com/samuelcolvin/4b5bb9bb163b1122ff17e29e48c10992
DOCS_JSON = (
    'https://gist.githubusercontent.com/'
    'samuelcolvin/4b5bb9bb163b1122ff17e29e48c10992/raw/'
    '80c5925c42f1442c24963aaf5eb1a324d47afe95/logfire_docs.json'
)


async def build_search_db():
    """Build the search database."""
    async with httpx.AsyncClient() as client:
        response = await client.get(DOCS_JSON)
        response.raise_for_status()
    sections = sections_ta.validate_json(response.content)

    openai = AsyncOpenAI()
    logfire.instrument_openai(openai)

    async with database_connect(True) as pool:
        with logfire.span('create schema'):
            async with pool.acquire() as conn:
                async with conn.transaction():
                    await conn.execute(DB_SCHEMA)

        sem = asyncio.Semaphore(10)
        async with create_task_group() as tg:
            for section in sections:
                tg.start_soon(insert_doc_section, sem, openai, pool, section)


async def insert_doc_section(
    sem: asyncio.Semaphore,
    openai: AsyncOpenAI,
    pool: asyncpg.Pool,
    section: DocsSection,
) -> None:
    async with sem:
        url = section.url()
        exists = await pool.fetchval('SELECT 1 FROM doc_sections WHERE url = $1', url)
        if exists:
            logfire.info('Skipping {url=}', url=url)
            return

        with logfire.span('create embedding for {url=}', url=url):
            embedding = await openai.embeddings.create(
                input=section.embedding_content(),
                model='text-embedding-3-small',
            )
        assert len(embedding.data) == 1, (
            f'Expected 1 embedding, got {len(embedding.data)}, doc section: {section}'
        )
        embedding = embedding.data[0].embedding
        embedding_json = pydantic_core.to_json(embedding).decode()
        await pool.execute(
            'INSERT INTO doc_sections (url, title, content, embedding) VALUES ($1, $2, $3, $4)',
            url,
            section.title,
            section.content,
            embedding_json,
        )


@dataclass
class DocsSection:
    id: int
    parent: int | None
    path: str
    level: int
    title: str
    content: str

    def url(self) -> str:
        url_path = re.sub(r'\.md$', '', self.path)
        return (
            f'https://logfire.pydantic.dev/docs/{url_path}/#{slugify(self.title, "-")}'
        )

    def embedding_content(self) -> str:
        return '\n\n'.join((f'path: {self.path}', f'title: {self.title}', self.content))


sections_ta = TypeAdapter(list[DocsSection])


# pyright: reportUnknownMemberType=false
# pyright: reportUnknownVariableType=false
@asynccontextmanager
async def database_connect(
    create_db: bool = False,
) -> AsyncGenerator[asyncpg.Pool, None]:
    server_dsn, database = (
        'postgresql://postgres:***@localhost:54320',
        'pydantic_ai_rag',
    )
    if create_db:
        with logfire.span('check and create DB'):
            conn = await asyncpg.connect(server_dsn)
            try:
                db_exists = await conn.fetchval(
                    'SELECT 1 FROM pg_database WHERE datname = $1', database
                )
                if not db_exists:
                    await conn.execute(f'CREATE DATABASE {database}')
            finally:
                await conn.close()

    pool = await asyncpg.create_pool(f'{server_dsn}/{database}')
    try:
        yield pool
    finally:
        await pool.close()


DB_SCHEMA = """
CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE IF NOT EXISTS doc_sections (
    id serial PRIMARY KEY,
    url text NOT NULL UNIQUE,
    title text NOT NULL,
    content text NOT NULL,
    -- text-embedding-3-small returns a vector of 1536 floats
    embedding vector(1536) NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_doc_sections_embedding ON doc_sections USING hnsw (embedding vector_l2_ops);
"""


def slugify(value: str, separator: str, unicode: bool = False) -> str:
    """Slugify a string, to make it URL friendly."""
    # Taken unchanged from https://github.com/Python-Markdown/markdown/blob/3.7/markdown/extensions/toc.py#L38
    if not unicode:
        # Replace Extended Latin characters with ASCII, i.e. `žlutý` => `zluty`
        value = unicodedata.normalize('NFKD', value)
        value = value.encode('ascii', 'ignore').decode('ascii')
    value = re.sub(r'[^\w\s-]', '', value).strip().lower()
    return re.sub(rf'[{separator}\s]+', separator, value)


if __name__ == '__main__':
    action = sys.argv[1] if len(sys.argv) > 1 else None
    if action == 'build':
        asyncio.run(build_search_db())
    elif action == 'search':
        if len(sys.argv) == 3:
            q = sys.argv[2]
        else:
            q = 'How do I configure logfire to work with FastAPI?'
        asyncio.run(run_agent(q))
    else:
        print(
            'uv run --extra examples -m pydantic_ai_examples.rag build|search',
            file=sys.stderr,
        )
        sys.exit(1)

동작 방식에서 눈여겨볼 핵심이 몇 가지 있어요.

retrieve 도구가 RAG의 검색 핵심이에요. 쿼리 문자열을 받아 OpenAI 임베딩으로 바꾸고(model='text-embedding-3-small'), pgvector의 <->(L2 거리) 연산으로 가장 가까운 문서 섹션 8개를 가져와요. 가져온 섹션은 제목·URL·본문으로 묶어 문자열로 반환해서, LLM이 컨텍스트로 사용해요.

**build_search_db**는 모든 문서 섹션을 순회하며 임베딩을 만들어 DB에 넣어요. asyncio.Semaphore(10)으로 동시 10개씩 처리하고, 이미 있는 URL(url = $1)은 건너뛰어요. 이때 DocsSection를 Pydantic TypeAdapter로 파싱해요.

스키마는 doc_sections 테이블에 embedding vector(1536) 컬럼을 두고, HNSW 인덱스(using hnsw (embedding vector_l2_ops))를 만들어 벡터 검색을 빠르게 해요.

더 알아보기 (Learn more)

• Pydantic AI 공식문서 — RAG 예시
• Pydantic AI — Tools
• Pydantic AI — Dependencies
• pgvector