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

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

注記

Databricks では、AI エージェント スキーマ ( ChatModelSplitChatMessageRequestStringResponse) が非推奨になりました。 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_chunkpredict_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 つのノートブックで定義されていますが、別のドライバー ノートブックを作成する必要があります。 このドライバーノートブックは、エージェントを モデルレジストリに記録し、モデルサービングを使用してエージェントをデプロイします。

この分離は、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"}