AI

[OpenAI Agent] 원하는 데이터 타입으로 깔끔하게 응답 받기

baecode 2025. 4. 8. 11:35
반응형
OpenAI Agent with Pydantic Example

Open AI 에서 제공하는 Agent를 사용함에 있어서 원하는 데이터 타입으로 응답 받게 해볼거에용.

코드단에서 우리가 원하는 형태의 데이터로 응답을 받으면, 코드에서 데이터를 정제하는 작업을 하지 않아도 되어서 코드의 양이 약간 줄어듭니다. (이득)

그러면 이전 게시물에 있던 Pydantic의 BaseModel을 사용해서 검증된 데이터를 내려 받아 보도록 합시당.

먼저 OpenAI Agent와 프롬프트를 한번 작성해보겠습니다.

# 저는 다른 기능을 추가로 붙여 쓰기위해 Agent를 상속 받아서 사용했습니다!

class CleaningProAgent(Agent):
    def __init__(self):
        super().__init__(
            name="청소 전문가 에이전트",
            instructions="너는 청소 전문가 에이전트야, 전문 청소지식으로 청소의 올바른 방법을 알려주고 잘못된 방법을 지적해야해",
        )
    async def excute(prompt: str):
        return await Runner.run(self, input=prompt)

def test_cleaning_prompt(message: str):
    return f"""
    너는 청소 전문가야. 사용자에게 청소 관련 질문을 받았을 때, 장소/오염 정도/도구 등을 고려해서 간단하고 실용적인 청소 팁을 제공해야 해.

    ##################
    사용자 메시지:
    {message}
    ##################
    """

물론 위 내용들을 영어로 적으면 2번 좋고 3번 좋습니다.

자 이제 실행을 해볼께용

user_message = "어제 고기 구워 먹고 나서 거실 바닥이 기름기로 미끌미끌해졌어요. 물티슈로는 안 닦이네요. 어떻게 해야 깨끗해질까요?"

clean_agent = CleaningProAgent()

prompt = test_cleaning_prompt(message=user_message)
clean_agent.excute(prompt=prompt)

# Response

"""
거실 바닥에 기름기가 남아 있을 때는 물티슈로는 충분하지 않아요. 아래 방법을 따라 해보세요!

1. **주방세제 희석액 준비**: 따뜻한 물 1L에 주방세제 한두 방울을 섞어주세요.
2. **마른 걸레보단 극세사 걸레 사용**: 부드러운 극세사 걸레나 스폰지를 사용해서 문질러 주세요.
3. **2차 헹굼**: 기름 제거 후, 깨끗한 물로 한 번 더 닦아내 기름기가 남지 않도록 해주세요.
4. **마무리 건조**: 수건이나 마른 걸레로 완전히 건조시켜 미끄러짐 방지!

⚠️ 물티슈는 알코올 성분 때문에 표면 손상이 갈 수 있으니 자주 사용하는 건 권장하지 않아요.
"""

응답이 위와 같이 나오게 됩니다.

해당 응답값의 타입 역시 RunResult로 Agent에서 사용하는 값 입니다.

이때 우리는 output_type에 원하는 데이터 타입을 주입해주고 내려받을 수 있습니다.

from pydantic import BaseModel

class CleanInfoData(BaseModel):
    result: str
    clean_tools: list

class CleaningProAgent(Agent):
    def __init__(self):
        super().__init__(
            name="청소 전문가 에이전트",
            instructions="너는 청소 전문가 에이전트야, 전문 청소지식으로 청소의 올바른 방법을 알려주고 잘못된 방법을 지적해야해",
            output_type=CleanInfoData
        )
    ...

user_message = "어제 고기 구워 먹고 나서 거실 바닥이 기름기로 미끌미끌해졌어요. 물티슈로는 안 닦이네요. 어떻게 해야 깨끗해질까요?"

clean_agent = CleaningProAgent()

prompt = test_cleaning_prompt(message=user_message)

result = clean_agent.excute(prompt=prompt)

print(result)

"""
RunResult:
- Last agent: Agent(name="청소 전문가 에이전트", ...)
- Final output (CleanInfoData):
    {
      "result": "AI가 적어준 화장실 청소법",
      "clean_tools": ['솔','락스','스펀지'],
    }
...
"""

여기서 RunResultfinal_output_as() 메서드를 사용하여 원하는 타입으로 output이 내려왔는지 검증과 함께 변환을 한차례 거쳐야 합니다.

result = clean_agent.excute(prompt=prompt)
result = result.final_output_as(CleanInfoData, raise_if_incorrect_type=True)
print(result)

raise_if_incorrect_typeTrue로 주게되면 타입이 우리가 전달해준 타입과 맞지 않을 경우 TypeError를 발생시킵니다.

def final_output_as(self, cls: type[T], raise_if_incorrect_type: bool = False) -> T:
    """
    A convenience method to cast the final output to a specific type. By default, the cast
    is only for the typechecker. If you set `raise_if_incorrect_type` to True, we'll raise a
    TypeError if the final output is not of the given type.

    Args:
        cls: The type to cast the final output to.
        raise_if_incorrect_type: If True, we'll raise a TypeError if the final output is not of
            the given type.

    Returns:
        The final output casted to the given type.
    """
    if raise_if_incorrect_type and not isinstance(self.final_output, cls):
        raise TypeError(f"Final output is not of type {cls.__name__}")

    return cast(T, self.final_output)

타입이 맞다면 최종적으로 전달받은 타입으로 변환시켜줍니다.

이렇게 하면 우리는 result = "AI가 적어준 화장실 청소법" clean_tools = ["솔","락스","스펀지"] 으로 바로 확인 해 볼 수 있습니다.

model_dump 메서드를 사용한다면 {'result': 'AI가 적어준 화장실 청소법','clean_tools': ['솔','락스','스펀지'] } 이렇게도 받을 수 있습니다.

반응형