エージェントの入力スキーマと出力スキーマを定義する
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", }, ] }
サポートされている出力スキーマ
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 のコード例と、カスタム入力を使用するための UI ベースのメソッドを参照してください。
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 を使用したカスタムスキーマの例を示しています。
Langchainカスタムスキーマ
次のノートブックは、LangChain を使用したカスタム スキーマの例を示しています。 ノートブックの wrap_output 関数を変更して、メッセージ・ストリームから情報を解析および抽出できます。
AI Playground とエージェント レビュー アプリでcustom_inputs
を提供する
custom_inputs
フィールドを使用して追加の入力を含むカスタムエージェントスキーマを定義すると、AI Playgroundとエージェントレビューアプリの両方でこれらの入力を手動で提供できます。カスタム入力が指定されていない場合、エージェントはスキーマで指定されたデフォルト値を使用します。
AI Playground または Agent Review App で、歯車アイコン を選択します。
custom_inputsを有効にします。
エージェントの定義済み入力スキーマに一致するJSONオブジェクトを提供します。
JSON オブジェクトは、エージェントの入力スキーマと一致する必要があります。 たとえば、次のように定義された custom_inputs
データクラスがあるとします。
@dataclass
class CustomInputs():
id: int = 0
user: str = "default"
次に、 custom_inputs フィールドに入力する JSON 文字列は、次の例に示すように、 id
と user
の値を指定する必要があります。
{
"id": 123
"user": "dev_test",
}