Stream Markdown

Stream Markdown

에이전트가 만들어내는 마크다운을 실시간으로 스트리밍하고, 터미널에서 rich 라이브러리로 하이라이트해서 보여주는 예시예요. 이 예시는 OpenAI와 Google Gemini 모델 두 가지 모두에 대해, 필요한 환경 변수가 설정되어 있으면 각각 실행해요.

출처: 공식문서

예시가 보여주는 것

예시 실행하기

의존성을 설치하고 환경 변수를 설정했다면, 실행해요.

python -m pydantic_ai_examples.stream_markdown

uv run -m pydantic_ai_examples.stream_markdown

예시 코드

stream_markdown.py

import asyncio
import os

import logfire
from rich.console import Console, ConsoleOptions, RenderResult
from rich.live import Live
from rich.markdown import CodeBlock, Markdown
from rich.syntax import Syntax
from rich.text import Text

from pydantic_ai import Agent
from pydantic_ai.models import KnownModelName

# '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_pydantic_ai()

agent = Agent()

# models to try, and the appropriate env var
models: list[tuple[KnownModelName, str]] = [
    ('google:gemini-3-flash-preview', 'GEMINI_API_KEY'),
    ('openai:gpt-5-mini', 'OPENAI_API_KEY'),
    ('groq:llama-3.3-70b-versatile', 'GROQ_API_KEY'),
]


async def main():
    prettier_code_blocks()
    console = Console()
    prompt = 'Show me a short example of using Pydantic.'
    console.log(f'Asking: {prompt}...', style='cyan')
    for model, env_var in models:
        if env_var in os.environ:
            console.log(f'Using model: {model}')
            with Live('', console=console, vertical_overflow='visible') as live:
                async with agent.run_stream(prompt, model=model) as result:
                    async for message in result.stream_output():
                        live.update(Markdown(message))
            console.log(result.usage)
        else:
            console.log(f'{model} requires {env_var} to be set.')


def prettier_code_blocks():
    """Make rich code blocks prettier and easier to copy.

    From https://github.com/samuelcolvin/aicli/blob/v0.8.0/samuelcolvin_aicli.py#L22
    """

    class SimpleCodeBlock(CodeBlock):
        def __rich_console__(
            self, console: Console, options: ConsoleOptions
        ) -> RenderResult:
            code = str(self.text).rstrip()
            yield Text(self.lexer_name, style='dim')
            yield Syntax(
                code,
                self.lexer_name,
                theme=self.theme,
                background_color='default',
                word_wrap=True,
            )
            yield Text(f'/{self.lexer_name}', style='dim')

    Markdown.elements['fence'] = SimpleCodeBlock


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

핵심은 이 부분이에요. agent = Agent()처럼 모델 없이 에이전트를 만들고, run_stream(prompt, model=model)을 호출할 때마다 모델을 넘기는 방식을 써요. result.stream_output()로 마크다운 조각을 하나씩 받아, Live 컨텍스트 안에서 Markdown(message)로 실시간 렌더링해요.

prettier_code_blocks()는 rich의 마크다운 코드블록 렌더링을 개선해서, 코드 블록에 언어 이름과 구문 하이라이트를 더 보기 좋게 보여줘요. 이 함수는 rich의 CodeBlock 클래스를 상속받아 커스텀 렌더링 메서드를 정의하고, Markdown.elements['fence']에 등록해요.

환경 변수가 있는 모델만 순서대로 시도하므로, 키가 없는 모델은 건너뛰고 로그만 남겨요.

더 알아보기 (Learn more)