Pydantic AI와 Mistral AI로 은행 지원 에이전트 만들기

Pydantic AI와 Mistral AI로 은행 지원 에이전트 만들기

Mistral AI와 PydanticAI로 은행 고객 지원 에이전트를 만드는 쿡북이에요. 구조화된 응답, 외부 의존성 주입(데이터베이스), 동적 시스템 프롬프트, 도구 통합이라는 네 가지 핵심 기능을 순서대로 구현해 봅니다.

출처: 문서

본문

이 쿡북의 예제는 ai.pydantic.dev에서 차용했어요. PydanticAI가 제공하는 기능은 다음과 같습니다.

  • 구조화된 응답(Structured Responses): Pydantic이 출력이 사전 정의된 스키마를 따르도록 보장해 일관되고 검증된 응답을 제공합니다.
  • 외부 의존성(External Dependencies): 타입 안전한 의존성 주입 시스템을 통해 데이터베이스 같은 외부 의존성을 AI 상호작용에 통합합니다.
  • 동적 문맥(Dynamic Context): 시스템 프롬프트 함수로 고객 이름 같은 런타임 정보를 에이전트의 문맥에 주입해 개인화된 상호작용을 가능하게 합니다.
  • 도구 통합(Tool Integration): 에이전트가 실시간 정보 검색을 위해 도구를 호출할 수 있어, 정적 응답을 넘어서는 능력을 갖춥니다.

필요한 패키지를 설치하고 설정합니다.

!pip install pydantic-ai==0.0.14 nest_asyncio

Jupyter 노트북이나 Colab에서 pydantic-ai를 실행한다면, Jupyter의 이벤트 루프와 pydantic-ai의 이벤트 루프 사이의 충돌을 관리하려면 nest-asyncio가 필요해요.

import nest_asyncio
nest_asyncio.apply()

import os
from getpass import getpass

os.environ["MISTRAL_API_KEY"] = getpass("Type your API Key")

예제 1: Mistral로 기본 Q&A

Mistral로 기본 Q&A를 하는 간단한 예제부터 시작합니다. 간결한 응답을 보장하는 시스템 프롬프트를 가진 에이전트를 정의하고, "hello world"의 기원을 물어보면 한 문장으로 답하게 합니다.

from pydantic_ai import Agent
from pydantic_ai.models.mistral import MistralModel

model = MistralModel('mistral-small-latest')

agent = Agent(
    model,
    system_prompt='Be concise, reply with one sentence.',
)

result = agent.run_sync('Where does "hello world" come from?')
print(result.data)

은행 지원 에이전트 정의하기

더 고급 AI 워크플로에서는 모델이 데이터베이스 정보 같은 외부 데이터를 필요로 할 수 있어요. 여기서는 고객의 이름과 잔액을 가져오는 가상의 데이터베이스 클래스를 정의합니다. 실제 상황에서는 이 클래스가 라이브 데이터베이스에 연결될 수 있어요.

from dataclasses import dataclass
from pydantic import BaseModel, Field
from pydantic_ai import Agent, RunContext

class DatabaseConn:
    """This is a fake database for example purposes.
    In reality, you'd be connecting to an external database
    (e.g. PostgreSQL) to get information about customers.
    """

    @classmethod
    async def customer_name(cls, *, id: int) -> str | None:
        if id == 123:
            return 'John'

    @classmethod
    async def customer_balance(cls, *, id: int, include_pending: bool) -> float:
        if id == 123:
            return 123.45
        else:
            raise ValueError('Customer not found')

이 단계에서는 지원 에이전트의 입력 의존성과 예상 출력 형식을 설정해 작동 방식을 정의합니다. 코드를 나눠 보면 이렇습니다.

1. 입력 의존성 (Input Dependencies)

  • SupportDependencies는 에이전트가 동작하는 데 필요한 것을 지정합니다.
    • customer_id: 도움을 받는 고객의 ID.
    • db: 데이터베이스 연결.

2. 예상 응답 형식 (Expected Response Format)

  • SupportResult는 응답 구조를 정의해 일관성을 보장합니다.
    • support_advice: 고객에게 주는 조언이 담긴 문자열.
    • block_card: 고객의 카드를 차단할지 여부를 나타내는 불리언.
    • risk: 0~10 사이의 정수로, 평가된 위험 수준.

3. 에이전트 초기화 (Agent Initialization)

  • Agent 클래스로 support_agent를 만듭니다.
    • model: 기본 AI 모델.
    • deps_type: 필요한 입력 의존성을 지정 (SupportDependencies).
    • result_type: 예상 출력 구조를 정의 (SupportResult).
    • system_prompt: 에이전트가 은행 지원 담당자처럼 행동하도록 안내하는 프롬프트. 고객 이름을 포함해 고객 특화 응답을 보장합니다.
@dataclass
class SupportDependencies:
    customer_id: int
    db: DatabaseConn

class SupportResult(BaseModel):
    support_advice: str = Field(description='Advice returned to the customer')
    block_card: bool = Field(description='Whether to block their')
    risk: int = Field(description='Risk level of query', ge=0, le=10)

support_agent = Agent(
    model,
    deps_type=SupportDependencies,
    result_type=SupportResult,
    system_prompt=(
        'You are a support agent in our bank, give the '
        'customer support and judge the risk level of their query. '
        "Reply using the customer's name."
    ),
)

왜 이런 설계가 중요한가요? 입력 의존성과 출력 형식을 정의하면 에이전트가 항상 올바른 데이터를 받고 예측 가능한 결과를 만들 수 있어요. 이는 더 큰 시스템에 통합하기 쉽게 해 주고, 명확하고 실용적인 응답을 지원합니다.

동적 시스템 프롬프트 추가하기

아래 코드는 동적 시스템 프롬프트 함수를 붙입니다. 모델이 사용자 질의를 보기 전에, 고객 이름이 포함된 특별한 시스템 프롬프트를 받게 됩니다.

@support_agent.system_prompt
async def add_customer_name(ctx: RunContext[SupportDependencies]) -> str:
    customer_name = await ctx.deps.db.customer_name(id=ctx.deps.customer_id)
    return f"The customer's name is {customer_name!r}"

에이전트가 사용할 도구 정의하기

customer_balance에 @support_agent.tool을 붙이면 모델이 이 함수를 호출해 고객의 잔액을 가져올 수 있음을 알게 됩니다. 이렇게 하면 모델이 수동적인 텍스트 생성기에서 외부 리소스와 상호작용하는 능동적인 문제 해결자로 바뀌어요.

@support_agent.tool
async def customer_balance(
    ctx: RunContext[SupportDependencies], include_pending: bool
) -> str:
    """Returns the customer's current account balance."""
    balance = await ctx.deps.db.customer_balance(
        id=ctx.deps.customer_id,
        include_pending=include_pending,
    )
    return f'${balance:.2f}'

에이전트 실행하기

고객의 잔액을 물어보면 에이전트가 주입된 의존성과 도구를 사용해 구조화된 응답을 반환합니다.

deps = SupportDependencies(customer_id=123, db=DatabaseConn())

result = support_agent.run_sync('What is my balance?', deps=deps)
print(result.data)

result = support_agent.run_sync('I just lost my card!', deps=deps)
print(result.data)

시스템 프롬프트와 도구 사용을 포함한 결과와 메시지 히스토리를 확인할 수 있어요.

result.__dict__

더 알아보기 (Learn more)

  • Pydantic AI 공식 문서 — 에이전트 개발 프레임워크
  • pydantic_ai.models.mistral.MistralModel — Mistral 모델용 PydanticAI 어댑터
  • Agent / RunContext / deps_type — PydanticAI 에이전트 구성 요소
  • mistral-small-latest — 이 예제에서 사용한 Mistral 모델