次の方法で共有

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

  • [アーティクル]

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

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

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

Mosaic AI Agent Framework では、次の入力スキーマがサポートされています。

OpenAI チャット完了スキーマ

Note

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

  • (Recommended) 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 式言語でチェーンを記述できます。 チェーン定義コードでは、 itemgetter を使用して、入力形式に応じてメッセージまたは query または history オブジェクトを取得できます。

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

Mosaic AI Agent Framework では、次の出力スキーマがサポートされています。

ChatCompletionResponse

(Recommended) OpenAI 応答形式の相互運用性をお持ちのお客様には、ChatCompletionResponse をお勧めします。

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(fake_model)
      | ChatCompletionsOutputParser() # or use this for ChatCompletionResponse
  )

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 は、実行時にエージェントの入力スキーマと出力スキーマを推論し、署名を自動的に作成できます。 サポートされている入力スキーマと出力スキーマを使用する場合、推論された署名は Agent Framework と互換性があります。 サポートされているスキーマの詳細については、「 サポートされる入力スキーマ」を参照してください。

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

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

エージェントのスキーマをカスタマイズして、エージェントとの間で追加のフィールドを渡したり返したりするには、 サポートされている入力/出力スキーマのサブクラスを作成します。 次に、追加のキー custom_inputs 追加し、追加のフィールドを含む custom_outputs を追加します。 以下の Pyfunc および Langchain のコード例 参照してください。

databricks-agents SDK、AI Playground やレビュー アプリなどの Databricks クライアント UI、およびその他のモザイク AI エージェント フレームワーク機能を使用するには、エージェントのスキーマが次の要件を満たしている必要があります。

  • エージェントはバージョン 2.17.1 以降 mlflow 使用する必要があります。
  • エージェント ノートブックで、サブクラスに追加された追加のフィールドを Optional としてマークし、既定値を割り当てます。
  • ドライバー ノートブックで、サブクラスのインスタンスでinfer_signatureを使用してModelSignatureを構築します。
  • ドライバー ノートブックで、サブクラスで asdict を呼び出して入力例を作成します。

PyFunc カスタム スキーマ

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

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

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

  • 入力サブクラスの型ヒントがあります。
  • ドット表記を使用してデータクラス フィールドにアクセスします (たとえば、model_input["custom_inputs"]ではなくmodel_input.custom_input.idを使用します)。
  • dictionary を返します。 サブクラスのインスタンスで asdict を呼び出して、リターンをディクショナリとして書式設定できます。

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

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

ノートブックを入手

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

ノートブックを入手

Langchain カスタム スキーマ

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

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

ノートブックを入手

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

ノートブックを入手