メインコンテンツまでスキップ

従来の入力および出力エージェント スキーマ

注記

Databricks では、AI エージェント スキーマ ( ChatModel、 SplitChatMessageRequest、 StringResponse) が非推奨になりました。 Databricks では、オーサリング エージェントに推奨される ChatAgent スキーマに移行することをお勧めします。 「コードでの AI エージェントの作成」を参照してください。

AI エージェントは、Databricks の他の機能と互換性を持つために、特定の入力および出力スキーマ要件に従う必要があります。この記事では、従来のエージェントオーサリング署名とインターフェイスの使用方法について説明します。 ChatModel インターフェース、 SplitChatMessageRequest 入力スキーマ、および StringResponse 出力スキーマ。

レガシーChatModelエージェントの作成​

important

Databricks では、エージェントや AI アプリの生成に ChatAgent インターフェイスを使用することをお勧めします。 ChatModel から ChatAgent に移行するには、 MLflow のドキュメント - ChatModel から ChatAgent への移行を参照してください。

ChatModel は、OpenAI の ChatCompletion スキーマを拡張する MLflow の従来のエージェント作成インターフェイスであり、互換性を維持できます ChatCompletion標準をサポートするプラットフォームと、カスタム機能を追加します。 詳細については、「 MLflow: ChatModel の概要 」を参照してください。

エージェントを mlflow.pyfunc.ChatModel のサブクラスとしてオーサリングすると、次の利点があります。

  • 配信されたエージェントを呼び出すときにストリーミング エージェントの出力を有効にします (要求本文の {stream: true} をバイパスします)。
  • エージェントのサービスを受けると、AIゲートウェイ推論テーブルが自動的に有効になり、リクエスタ名などの拡張リクエストログメタデータにアクセスできるようになります。
  • 型指定された Python クラスを使用して、ChatCompletion スキーマと互換性のあるエージェントコードを記述できます。
  • MLflow は、エージェントのログを記録するときに、 input_exampleがなくても、チャット完了と互換性のある署名を自動的に推論します。 これにより、エージェントの登録とデプロイのプロセスが簡素化されます。 ログ記録中のモデルの署名の推論を参照してください。

次のコードは、Databricks ノートブックで実行するのが最適です。 ノートブックは、エージェントの開発、テスト、および反復処理に便利な環境を提供します。

MyAgent クラスは mlflow.pyfunc.ChatModelを拡張し、必要な predict メソッドを実装します。これにより、Mosaic AI Agent Framework との互換性が確保されます。

このクラスには、ストリーミング出力を処理するためのオプションのメソッド ( _create_chat_completion_chunk と predict_stream ) も含まれています。

Python
import re
from typing import Optional, Dict, List, Generator
from mlflow.pyfunc import ChatModel
from mlflow.types.llm import (
    # Non-streaming helper classes
    ChatCompletionRequest,
    ChatCompletionResponse,
    ChatCompletionChunk,
    ChatMessage,
    ChatChoice,
    ChatParams,
    # Helper classes for streaming agent output
    ChatChoiceDelta,
    ChatChunkChoice,
)

class MyAgent(ChatModel):
    """
    Defines a custom agent that processes ChatCompletionRequests
    and returns ChatCompletionResponses.
    """
    def predict(self, context, messages: list[ChatMessage], params: ChatParams) -> ChatCompletionResponse:
        last_user_question_text = messages[-1].content
        response_message = ChatMessage(
            role="assistant",
            content=(
                f"I will always echo back your last question. Your last question was: {last_user_question_text}. "
            )
        )
        return ChatCompletionResponse(
            choices=[ChatChoice(message=response_message)]
        )

    def _create_chat_completion_chunk(self, content) -> ChatCompletionChunk:
        """Helper for constructing a ChatCompletionChunk instance for wrapping streaming agent output"""
        return ChatCompletionChunk(
                choices=[ChatChunkChoice(
                    delta=ChatChoiceDelta(
                        role="assistant",
                        content=content
                    )
                )]
            )

    def predict_stream(
        self, context, messages: List[ChatMessage], params: ChatParams
    ) -> Generator[ChatCompletionChunk, None, None]:
        last_user_question_text = messages[-1].content
        yield self._create_chat_completion_chunk(f"Echoing back your last question, word by word.")
        for word in re.findall(r"\S+\s*", last_user_question_text):
            yield self._create_chat_completion_chunk(word)

agent = MyAgent()
model_input = ChatCompletionRequest(
    messages=[ChatMessage(role="user", content="What is Databricks?")]
)
response = agent.predict(context=None, model_input=model_input)
print(response)

エージェント クラス MyAgent は 1 つのノートブックで定義されていますが、別のドライバー ノートブックを作成する必要があります。 driver ノートブックは、エージェントを Model Registry にログに記録し、モデルサービングを使用してエージェントをデプロイします。

この分離は、MLflow の Models from Code 手法を使用してモデルをログ記録するために Databricks が推奨するワークフローに従います。

SplitChatMessageRequest 入力スキーマ (非推奨)​

SplitChatMessagesRequest では、現在のクエリと履歴をエージェント入力として別々に渡すことができます。

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

StringResponse 出力スキーマ (非推奨)​

StringResponse エージェントの応答を、単一の文字列 content フィールドを持つオブジェクトとして返すことができます。

{"content": "This is an example string response"}
に最終更新