従来の入力および出力エージェント スキーマ
Databricks では、AI エージェント スキーマ ( ChatModel
、 SplitChatMessageRequest
、 StringResponse
) が非推奨になりました。 Databricks では、オーサリング エージェントに推奨される ChatAgent
スキーマに移行することをお勧めします。 コードでの AI エージェントの作成を参照してください。
AI エージェントは、Databricks の他の機能と互換性を持つために、特定の入力および出力スキーマ要件に従う必要があります。この記事では、従来のエージェントオーサリング署名とインターフェイスの使用方法について説明します。
ChatModel
インターフェース、 SplitChatMessageRequest
入力スキーマ、および StringResponse
出力スキーマ。
レガシーChatModelエージェントの作成
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
) も含まれています。
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
では、現在のクエリと履歴をエージェント入力として別々に渡すことができます。
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"}