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

AI エージェントは、Databricks の他の機能と互換性を持つために、特定の入力および出力スキーマ要件に従う必要があります。 この記事では、AI エージェントがこれらの要件に準拠していることを確認する方法と、互換性を確保しながらエージェントの入力スキーマと出力スキーマをカスタマイズする方法について説明します。

Mosaic AI は、 MLflow モデル シグネチャを使用して 、エージェントの入力スキーマと出力スキーマの要件を定義します。 モデル署名は、内部コンポーネントと外部コンポーネントにエージェントとの対話方法を指示し、それらがスキーマに準拠していることを検証します。

OpenAI チャット完了スキーマ (推奨)

Databricks では、 OpenAI チャット完了スキーマ を使用してエージェントの入力と出力を定義することをお勧めします。 このスキーマは広く採用されており、Databricks のものを含む多くのエージェント フレームワークとアプリケーションと互換性があります。

チャット完了 の入力スキーマ と 出力スキーマの詳細については、OpenAI のドキュメントを参照してください。

注：

OpenAIチャット完了スキーマは、エージェントの入力と出力を構造化するための標準にすぎません。 このスキーマの実装には、OpenAI のリソースやモデルの使用は含まれません。

MLflowは、APIs LangChainおよび PyFunc フレーバーのエージェントに便利な を提供し、チャット完了スキーマと互換性のあるエージェントを作成するのに役立ちます。

LangChainでチャット補完を実装する

エージェントが LangChain を使用している場合は、MLflow の ChatCompletionOutputParser() を使用して、エージェントの最終出力をチャット完了スキーマと互換性を持つように書式設定します。 LangGraph を使用する場合は、「 LangGraph カスタム スキーマ」を参照してください。

  from mlflow.langchain.output_parsers import ChatCompletionOutputParser

  chain = (
      {
          "user_query": itemgetter("messages")
          | RunnableLambda(extract_user_query_string),
          "chat_history": itemgetter("messages") | RunnableLambda(extract_chat_history),
      }
      | RunnableLambda(DatabricksChat)
      | ChatCompletionOutputParser()
  )

PyFuncでチャット補完を実装する

PyFunc を使用する場合、Databricks ではエージェントを mlflow.pyfunc.ChatModelのサブクラスとして記述することをお勧めします。 この方法には、次の利点があります。

• 型付きPythonクラスを使用して、チャット完了スキーマと互換性のあるエージェントコードを記述できます。

• エージェントをログに記録すると、MLflow は、 input_exampleがなくても、チャット完了と互換性のある署名を自動的に推論します。 これにより、エージェントの登録とデプロイのプロセスが簡素化されます。 ログ記録中のモデルの署名の推論を参照してください。

  from dataclasses import dataclass
  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 last_user_question_text.split(" "):
              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)
  

カスタム入力と出力

Databricks では、ほとんどのエージェントのユースケースで OpenAI チャット完了スキーマに準拠することをお勧めします。

ただし、一部のシナリオでは、 client_type や session_idなどの追加の入力や、将来の対話のためにチャット履歴に含めるべきではない取得ソース リンクなどの出力が必要になる場合があります。

これらのシナリオでは、Mosaic AI Agent Framework は、カスタム入力と出力による OpenAI チャット完了リクエストとレスポンスの拡張をサポートしています。

PyFuncおよびLangGraphエージェントのカスタム入力と出力を作成する方法については、次の例を参照してください。

警告

Agent Evaluation レビュー アプリは、現在、追加の入力フィールドを持つエージェントのトレースのレンダリングをサポートしていません。

PyFuncカスタムスキーマ

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

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

ノートブックを新しいタブで開くインポート用のリンクをコピー

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

ノートブックを新しいタブで開くインポート用のリンクをコピー

LangGraph カスタムスキーマ

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

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

ノートブックを新しいタブで開くインポート用のリンクをコピー

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

ノートブックを新しいタブで開くインポート用のリンクをコピー

AI Playground とエージェント レビュー アプリでcustom_inputsを提供する

エージェントが custom_inputs フィールドを使用して追加の入力を受け入れる場合は、 AI Playground と エージェントレビューアプリの両方でこれらの入力を手動で提供できます。

1. AI Playground または Agent Review App で、歯車アイコン を選択します。

2. custom_inputsを有効にします。

3. エージェントの定義済み入力スキーマに一致するJSONオブジェクトを提供します。

   

従来の入力スキーマと出力スキーマ

SplitChatMessageRequest 入力スキーマと StringResponse 出力スキーマは非推奨になりました。 これらのレガシ スキーマのいずれかを使用している場合、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"}