Databricks에서 OpenAI Agents SDK로 멀티 에이전트 시스템 구축하기

openai-agents==0.0.7
unitycatalog-openai[databricks]==0.2.0
mlflow==2.21.0

-- 정책 번호로 청구 내역 조회
CREATE OR REPLACE FUNCTION catalog.schema.search_claims_details_by_policy_no (
    input_policy_no STRING COMMENT 'Policy number'
)
RETURNS TABLE
COMMENT 'Returns policy details about a customer given policy_no.'
RETURN
SELECT *
FROM catalog.schema.claims_table
WHERE policy_no = input_policy_no;

-- 정책 문서 벡터 검색
CREATE OR REPLACE FUNCTION catalog.schema.policy_docs_vector_search (
    query STRING
    COMMENT 'Query string for insurance policy documentation.'
)
RETURNS TABLE
COMMENT 'Searches policy documentation, retrieving most relevant text documents.'
RETURN
SELECT
    chunked_text as page_content,
    map('doc_path', path, 'chunk_id', chunk_id) as metadata
FROM
    vector_search(
        index => 'catalog.schema.policy_docs_chunked_files_vs_index',
        query => query,
        num_results => 3
    );

from unitycatalog.ai.core.databricks import (
    DatabricksFunctionClient,
    FunctionExecutionResult,
)
from agents import function_tool, RunContextWrapper
from pydantic import BaseModel

class UserInfo(BaseModel):
    cust_id: str | None = None
    policy_no: str | None = None

@function_tool
def search_claims_details_by_policy_no(
    wrapper: RunContextWrapper[UserInfo], policy_no: str
) -> FunctionExecutionResult:
    print("[DEBUG]: search_claims_details_by_policy_no tool called")
    wrapper.context.policy_no = policy_no
    client = DatabricksFunctionClient()
    return client.execute_function(
        function_name="ai.insurance_agent.search_claims_details_by_policy_no",
        parameters={"input_policy_no": wrapper.context.policy_no},
    )

@function_tool
def policy_docs_vector_search(query: str) -> FunctionExecutionResult:
    print("[DEBUG]: policy_docs_vector_search tool called")
    client = DatabricksFunctionClient()
    return client.execute_function(
        function_name="ai.insurance_agent.policy_docs_vector_search",
        parameters={"query": query},
    )

import os
from agents import Agent
from agents.extensions.handoff_prompt import RECOMMENDED_PROMPT_PREFIX

os.environ["OPENAI_API_KEY"] = dbutils.secrets.get(...)

claims_detail_retrieval_agent = Agent[UserInfo](
    name="Claims Details Retrieval Agent",
    instructions=(
        f"{RECOMMENDED_PROMPT_PREFIX}"
        "You are a claims details retrieval agent. "
        "If speaking to a customer, you were transferred from the triage agent. "
        "Use this routine:\n"
        "1. Identify the customer's last question.\n"
        "2. Use search tools to retrieve claim data.\n"
        "3. Transfer back if unable to answer.\n"
    ),
    tools=[search_claims_details_by_policy_no],
    model="gpt-4o",
)

policy_qa_agent = Agent[UserInfo](
    name="Policy Q&A Agent",
    instructions=(
        f"{RECOMMENDED_PROMPT_PREFIX}"
        "You are an insurance policy Q&A agent. "
        "If speaking to a customer, you were transferred from the triage agent. "
        "Use this routine:\n"
        "1. Identify the customer's last question.\n"
        "2. Use search tools to answer policy questions.\n"
        "3. Transfer back if unable to answer.\n"
    ),
    tools=[policy_docs_vector_search],
    model="gpt-4o",
)

from openai import AsyncOpenAI
from agents import OpenAIChatCompletionsModel

BASE_URL = "https://your-databricks-workspace.com/serving-endpoints"
API_KEY = dbutils.secrets.get(...)

client = AsyncOpenAI(base_url=BASE_URL, api_key=API_KEY)

custom_llm = OpenAIChatCompletionsModel(
    model="databricks-claude-3-7-sonnet",
    openai_client=client
)

triage_agent = Agent[UserInfo](
    name="Triage agent",
    instructions=(
        f"{RECOMMENDED_PROMPT_PREFIX}"
        "You are a helpful triaging agent. "
        "Delegate questions to appropriate agents. "
        "If customers have no further questions, wish them well."
    ),
    handoffs=[claims_detail_retrieval_agent, policy_qa_agent],
    model="gpt-4o",
)

import mlflow
from agents import Runner

user_info = UserInfo(cust_id="7852", policy_no="102070455")

with mlflow.start_span(name="insurance_agent", span_type="AGENT") as span:
    print("[AGENT] Hello! How may I assist you?")
    while True:
        user_input = input("[USER]: ")
        if user_input.lower() == "exit":
            print("[AGENT]: Bye!")
            break
        if not user_input:
            continue
        try:
            result = await Runner.run(
                starting_agent=triage_agent,
                input=user_input,
                context=user_info
            )
            print("\n[AGENT]:", result.final_output)
        except Exception as e:
            print(f"\nError occurred: {str(e)}")

import os
os.environ["OPENAI_AGENTS_DISABLE_TRACING"] = "1"

from agents import set_tracing_disabled
set_tracing_disabled(disabled=True)

import mlflow
from mlflow.pyfunc import ChatAgent

mlflow.openai.autolog()

class InsuranceChatAgent(ChatAgent):
    def __init__(self, starting_agent: Agent):
        self.starting_agent = starting_agent

    def _convert_to_input_text(self, messages: List[ChatAgentMessage]) -> str:
        """가장 최근의 사용자 메시지를 입력 텍스트로 추출"""
        for message in reversed(messages):
            if message.role == "user":
                return message.content
        return ""

    def _create_user_context(
        self,
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[Dict[str, Any]] = None
    ) -> UserInfo:
        """MLflow 입력을 UserInfo 객체로 변환"""
        user_info = UserInfo()
        if context:
            conversation_id = getattr(context, "conversation_id", None)
            if conversation_id:
                user_info.conversation_id = conversation_id

            user_id = getattr(context, "user_id", None)
            if user_id:
                user_info.user_id = user_id
        return user_info

from mlflow.entities import SpanType

class InsuranceChatAgent(ChatAgent):
    ...

    @mlflow.trace(name="insurance_chat_agent", span_type=SpanType.AGENT)
    def predict(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[Dict[str, Any]] = None
    ) -> ChatAgentResponse:
        input_text = self._convert_to_input_text(messages)
        user_info = self._create_user_context(context, custom_inputs)
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            result = loop.run_until_complete(
                Runner.run(
                    starting_agent=self.starting_agent,
                    input=input_text,
                    context=user_info,
                )
            )
        finally:
            loop.close()
        return ChatAgentResponse(
            messages=[
                ChatAgentMessage(
                    role="assistant",
                    content=result.final_output,
                    id=str(uuid4())
                )
            ]
        )

    @mlflow.trace(name="insurance_chat_agent_stream", span_type=SpanType.AGENT)
    def predict_stream(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[Dict[str, Any]] = None
    ) -> Generator[ChatAgentResponse, None, None]:
        response = self.predict(messages, context, custom_inputs)
        for message in response.messages:
            yield ChatAgentChunk(delta=message)

AGENT = InsuranceChatAgent(starting_agent=triage_agent)
mlflow.models.set_model(AGENT)

import os
from mlflow.models.resources import (
    DatabricksFunction,
    DatabricksServingEndpoint,
    DatabricksVectorSearchIndex
)

resources = [
    DatabricksVectorSearchIndex("catalog.schema.policy_docs_chunked_files_vs_index"),
    DatabricksServingEndpoint("databricks-bge-large-en"),
    DatabricksFunction("catalog.schema.search_claims_details_by_policy_no"),
    DatabricksFunction("catalog.schema.policy_docs_vector_search")
]

mlflow.set_registry_uri("databricks-uc")
mlflow.set_experiment("/Users/user@company.com/insurance_agent_experiment")
mlflow.openai.autolog()

with mlflow.start_run():
    logged_model_info = mlflow.pyfunc.log_model(
        artifact_path="insurance_chat_agent",
        python_model=os.path.join(os.getcwd(), "insurance_chat_agent.py"),
        input_example={
            "messages": [
                {
                    "role": "user",
                    "content": "hi, id like to check on my existing claims?",
                }
            ],
            "context": {"conversation_id": "123", "user_id": "123"},
        },
        pip_requirements=[
            "mlflow",
            "openai-agents",
            "unitycatalog-openai[databricks]==0.2.0",
            "pydantic",
        ],
        resources=resources
    )

model_name = "insurance_chat_agent"
UC_MODEL_NAME = f"catalog.schema.{model_name}"

uc_registered_model_info = mlflow.register_model(
    model_uri=logged_model_info.model_uri,
    name=UC_MODEL_NAME
)

from databricks import agents

agents.deploy(
    UC_MODEL_NAME,
    uc_registered_model_info.version,
    environment_vars={
        "OPENAI_API_KEY": "{{secrets/my_scope/my_api_key}}",
        "DATABRICKS_TOKEN": "{{secrets/my_scope/my_token}}",
    },
    tags={"endpoint_desc": "insurance_chat_agent"},
)