エージェントの入力スキーマと出力スキーマを定義する

MLflow モデル シグネチャは 、AI エージェントの入力スキーマと出力スキーマの要件を定義します。 モデルシグネチャは、内部コンポーネントと外部コンポーネントにエージェントとの対話方法を指示します。 モデル署名は、入力がスキーマ要件に準拠していることを確認するための検証チェックとして機能します。

たとえば、Agent Evaluation レビューアプリを使用するには、エージェントは Agent Evaluation 入力スキーマに準拠している必要があります。

サポートされている入力スキーマ

Mosaic AI Agent フレームワークは、次の入力スキーマをサポートしています。

OpenAI チャット完了スキーマ

注：

Databricks では、OpenAI チャット完了スキーマが広く使用されており、多くのエージェント フレームワークやアプリケーションと相互運用できるため、OpenAI チャット完了スキーマを推奨しています。 OpenAI チャット完了スキーマがニーズを満たさない場合は、独自のスキーマを定義できます。 「カスタム・エージェント・スキーマ」を参照してください。

  • (推奨)Databricks では、 OpenAI チャット完了スキーマを使用することをお勧めします。 OpenAI チャット完了スキーマには、 messages パラメーターとしてオブジェクトの配列が必要です。 この形式は、RAGアプリケーションに最適です。

    question = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-Augmented Generation?",
        },
        {
            "role": "assistant",
            "content": "RAG, or Retrieval Augmented Generation, is a generative AI design pattern that combines a large language model (LLM) with external knowledge retrieval. This approach allows for real-time data connection to generative AI applications, improving their accuracy and quality by providing context from your data to the LLM during inference. Databricks offers integrated tools that support various RAG scenarios, such as unstructured data, structured data, tools & function calling, and agents.",
        },
        {
            "role": "user",
            "content": "How to build RAG for unstructured data",
        },
    ]
}

SplitChatMessageRequest (英語)

SplitChatMessagesRequest は、マルチターン チャット アプリケーション、特に現在のクエリと履歴を個別に管理する場合に推奨されます。

  question = {
      "query": "What is MLflow",
      "history": [
          {
              "role": "user",
              "content": "What is Retrieval-augmented Generation?"
          },
          {
              "role": "assistant",
              "content": "RAG is"
          }
      ]
  }

Langchain 式言語

エージェントがLangChainを使用している場合は、LangChain Expression Languageでチェーンを記述できます。 チェーン定義コードでは、 itemgetter を使用して、入力形式に応じてメッセージを取得したり、オブジェクト query または history したりできます。

サポートされている出力スキーマ

Mosaic AI Agent Framework は、次の出力スキーマをサポートしています。

チャット完了レスポンス

(推奨)ChatCompletionResponse は、OpenAI 応答形式の相互運用性を持つ顧客に推奨されます。

LangChain - ChatCompletionsOutputParser (英語)

エージェントが LangChain を使用している場合は、最後のチェーン ステップとして MLflow の ChatCompletionsOutputParser() を使用します。 これにより、LangChain AIメッセージがエージェント互換の形式にフォーマットされます。


  from mlflow.langchain.output_parsers import ChatCompletionsOutputParser

  chain = (
      {
          "user_query": itemgetter("messages")
          | RunnableLambda(extract_user_query_string),
          "chat_history": itemgetter("messages") | RunnableLambda(extract_chat_history),
      }
      | RunnableLambda(DatabricksChat)
      | ChatCompletionsOutputParser()
  )

PyFunc - 入力クラスと出力クラスに注釈を付ける

PyFuncを使用している場合、Databricksでは、タイプヒントを使用して、mlflow.models.rag_signaturesで定義されたクラスのサブクラスである入力および出力データクラスでpredict()関数に注釈を付けることをお勧めします。

形式が確実に守られるように、predict()内のデータクラスから出力オブジェクトを構築できます。返されたオブジェクトは、シリアル化できるように辞書表現に変換する必要があります。


  from mlflow.models.rag_signatures import ChatCompletionRequest, ChatCompletionResponse, ChainCompletionChoice, Message

  class RAGModel(PythonModel):
    ...
      def predict(self, context, model_input: ChatCompletionRequest) -> ChatCompletionResponse:
        ...
        return asdict(ChatCompletionResponse(
            choices=[ChainCompletionChoice(message=Message(content=text))]
        ))

明示的署名と推論署名

MLflow 、実行時にエージェントの入力スキーマと出力スキーマを推測し、署名を自動的に作成できます。 サポートされている入力スキーマと出力スキーマを使用する場合、推論された署名はエージェントフレームワークと互換性があります。 サポートされているスキーマに関する情報については、「 サポートされている入力スキーマ」を参照してください。

ただし、カスタム・エージェント・スキーマを使用する場合は、 カスタム・エージェント・スキーマの手順に従って、モデル署名を明示的に定義する必要があります。

カスタム・エージェント・スキーマ

エージェントのスキーマをカスタマイズして、エージェントとの間で追加のフィールドを送受信するには、 サポートされている入力/出力スキーマのサブクラスを作成します。 次に、追加のフィールドを格納するために、追加のキー ( custom_inputs と custom_outputs ) を追加します。 以下の Pyfunc と Langchain のコード例を参照してください。

databricks-agents SDK、AI Playground や Review App などの Databricks クライアント UI、およびその他の Mosaic AI Agent Framework 機能を使用するには、エージェントのスキーマが次の要件を満たしている必要があります。

  • エージェントは mlflow バージョン2.17.1以降を使用する必要があります。

  • エージェントノートブックで、サブクラスに追加された追加のフィールドを Optional としてマークし、デフォルト値を割り当てます。

  • ドライバー ノートブックで、サブクラスのインスタンスを使用してinfer_signatureを使用してModelSignatureを構築します。

  • ドライバー ノートブックで、サブクラスで asdict を呼び出して入力例を作成します。

PyFuncカスタムスキーマ

上記の要件に加えて、PyFunc ベースのエージェントは、Mosaic AI エージェント機能と対話するために、次の要件も満たす必要があります。

PyFuncカスタムスキーマの要件

エージェントノートブックでは、predict 関数と predict ストリーム関数は、次の要件を満たしている必要があります。

  • 入力サブクラスの型ヒントを用意します。

  • ドット表記を使用してデータクラスフィールドにアクセスします(たとえば、model_input["custom_inputs"]の代わりにmodel_input.custom_input.idを使用します)。

  • dictionaryを返します。サブクラスのインスタンスで asdict を呼び出して、戻り値をディクショナリとして書式設定できます。

次のノートブックは、PyFunc を使用したカスタムスキーマの例を示しています。

PyFunc カスタムスキーマエージェントノートブック

ノートブックを新しいタブで開く

PyFuncカスタムスキーマドライバーノートブック

ノートブックを新しいタブで開く

Langchainカスタムスキーマ

次のノートブックは、LangChain を使用したカスタム スキーマの例を示しています。 ノートブックの wrap_output 関数を変更して、メッセージ・ストリームから情報を解析および抽出できます。

Langchain カスタム スキーマ エージェント ノートブック

ノートブックを新しいタブで開く

Langchain カスタム スキーマ ドライバー ノートブック

ノートブックを新しいタブで開く