Bagikan melalui

Menulis agen AI dalam kode

Halaman ini menunjukkan cara menulis agen AI di Python menggunakan Mosaic AI Agent Framework dan pustaka penulisan agen populer seperti LangGraph, PyFunc, dan OpenAI.

Persyaratan

Petunjuk / Saran

Databricks merekomendasikan untuk menginstal versi terbaru klien MLflow Python saat mengembangkan agen.

Untuk menulis dan menyebarkan agen menggunakan pendekatan di halaman ini, instal perangkat berikut:

  • databricks-agents 0.16.0 atau lebih tinggi
  • mlflow 2.20.2 atau lebih tinggi
  • Python 3.10 atau lebih tinggi.
    • Untuk memenuhi persyaratan ini, Anda dapat menggunakan komputasi tanpa server atau Databricks Runtime 13.3 LTS atau lebih tinggi.
%pip install -U -qqqq databricks-agents>=0.16.0 mlflow>=2.20.2

Databricks juga merekomendasikan penginstalan paket integrasi Databricks AI Bridge saat membuat agen. Paket integrasi ini (seperti databricks-langchain, databricks-openai) menyediakan lapisan API bersama untuk berinteraksi dengan fitur AI Databricks, seperti Databricks AI/BI Genie dan Vector Search, di seluruh kerangka kerja penulisan agen dan SDK.

LangChain/LangGraph

%pip install -U -qqqq databricks-langchain

OpenAI

%pip install -U -qqqq databricks-openai

Agen Python Asli

%pip install -U -qqqq databricks-ai-bridge

Menggunakan ChatAgent untuk agen pembuat

Databricks merekomendasikan antarmuka MLflow ChatAgent untuk mengembangkan agen tingkat produksi. Spesifikasi skema obrolan ini mirip dengan, tetapi tidak kompatibel secara ketat dengan, skema OpenAI ChatCompletion .

ChatAgent dengan mudah membungkus agen yang ada untuk kompatibilitas Databricks.

ChatAgent memberikan manfaat berikut:

  • Kemampuan agen tingkat lanjut

    • Dukungan multi-agen
    • Streaming output: Aktifkan pengalaman pengguna interaktif dengan streaming output dalam bagian-bagian yang lebih kecil.
    • Sejarah pesan pemanggilan alat yang komprehensif: Mengembalikan beberapa pesan, termasuk pesan pemanggilan alat yang bersifat perantara, untuk meningkatkan kualitas dan pengelolaan percakapan.
    • Dukungan untuk konfirmasi pemanggilan alat

  • Pengembangan, penyebaran, dan pemantauan yang disederhanakan

    • Tulis agen Anda menggunakan kerangka kerja apa pun: Bungkus agen yang ada menggunakan ChatAgent antarmuka untuk mendapatkan kompatibilitas langsung dengan AI Playground, Evaluasi Agen, dan Pemantauan Agen.
    • Antarmuka penulisan bertipe: Menulis kode agen dengan menggunakan kelas Python bertipe, mendapatkan manfaat dari fitur pelengkapan otomatis di IDE dan notebook.
    • Inferensi tanda tangan otomatis: MLflow secara otomatis menyimpulkan tanda tangan ChatAgent saat mencatat agen, menyederhanakan pendaftaran dan penyebaran. Lihat Infer Model Signature selama pengelogan.
    • tabel inferensi Gateway AI yang disempurnakan: Tabel inferensi Gateway AI diaktifkan secara otomatis untuk agen yang telah disebarkan, menyediakan akses ke metadata log permintaan terperinci.

Untuk mempelajari cara membuat ChatAgent, lihat contoh di bagian berikut dan dokumentasi MLflow - Apa itu antarmuka ChatAgent.

Bagaimana jika saya sudah memiliki agen?

Jika Anda sudah memiliki agen yang dibangun dengan LangChain, LangGraph, atau kerangka kerja serupa, Anda tidak perlu menulis ulang agen Anda untuk menggunakannya di Databricks. Sebagai gantinya, cukup bungkus agen Anda yang ada dengan antarmuka MLflow ChatAgent :

  1. Tulis kelas pembungkus Python yang mewarisi dari mlflow.pyfunc.ChatAgent.

    Di dalam kelas pembungkus, pertahankan agen Anda yang sudah ada sebagai atribut self.agent = your_existing_agent.

  2. KelasChatAgent mengharuskan Anda untuk menerapkan metodepredict yang menangani permintaan non-streaming.

    predict harus menerima:

    • messages: list[ChatAgentMessage], yang merupakan daftar ChatAgentMessage masing-masing dengan peran (seperti "pengguna" atau "asisten"), perintah, dan ID.

    • (Opsional) context: Optional[ChatContext] dan custom_inputs: Optional[dict] untuk data tambahan.

    import uuid

# input example
[
  ChatAgentMessage(
    id=str(uuid.uuid4()),  # Generate a unique ID for each message
    role="user",
    content="What's the weather in Paris?"
  )
]

    predict harus mengembalikan ChatAgentResponse.

    import uuid

# output example
ChatAgentResponse(
  messages=[
    ChatAgentMessage(
      id=str(uuid.uuid4()),  # Generate a unique ID for each message
      role="assistant",
      content="It's sunny in Paris."
    )
  ]
)

  3. Mengonversi antar format

    Di predict, konversikan pesan masuk dari list[ChatAgentMessage] ke dalam format input yang diharapkan agen Anda.

    Setelah agen Anda menghasilkan respons, konversikan outputnya menjadi satu atau beberapa ChatAgentMessage objek dan bungkus dalam ChatAgentResponse.

Petunjuk / Saran

Mengonversi output LangChain secara otomatis

Jika Anda membungkus agen LangChain, Anda dapat menggunakan mlflow.langchain.output_parsers.ChatAgentOutputParser untuk secara otomatis mengonversi output LangChain ke dalam skema MLflow seperti ChatAgentMessage dan ChatAgentResponse.

Berikut ini adalah templat yang disederhanakan untuk mengonversi agen Anda:

from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import ChatAgentMessage, ChatAgentResponse, ChatAgentChunk
import uuid


class MyWrappedAgent(ChatAgent):
  def __init__(self, agent):
    self.agent = agent

  def predict(self, messages, context=None, custom_inputs=None):
    # Convert messages to your agent's format
    agent_input = ... # build from messages
    agent_output = self.agent.invoke(agent_input)
    # Convert output to ChatAgentMessage
    return ChatAgentResponse(
      messages=[ChatAgentMessage(role="assistant", content=agent_output, id=str(uuid.uuid4()),)]
    )

  def predict_stream(self, messages, context=None, custom_inputs=None):
    # If your agent supports streaming
    for chunk in self.agent.stream(...):
      yield ChatAgentChunk(delta=ChatAgentMessage(role="assistant", content=chunk, id=str(uuid.uuid4())))

Untuk contoh lengkapnya, lihat buku catatan di bagian berikut ini.

ChatAgent contoh

Buku catatan berikut menunjukkan cara menulis streaming dan non-streaming ChatAgents menggunakan pustaka populer OpenAI, LangGraph, dan AutoGen.

LangGraph

Jika Anda membungkus agen LangChain, Anda dapat menggunakan mlflow.langchain.output_parsers.ChatAgentOutputParser untuk secara otomatis mengonversi output LangChain ke dalam skema MLflow seperti ChatAgentMessage dan ChatAgentResponse.

Agen pemanggil aplikasi LangGraph

Ambil buku catatan

OpenAI

Agen pemanggil alat OpenAI

Ambil buku catatan

Agen pemanggil alat API OpenAI Responses

Ambil buku catatan

Agen khusus untuk obrolan OpenAI

Ambil buku catatan

AutoGen

Agen pemanggil alat AutoGen

Ambil buku catatan

DSPy

Agen DSPy untuk obrolan saja

Ambil buku catatan

Untuk mempelajari cara memperluas kemampuan agen ini dengan menambahkan alat, lihat alat agen AI.

Contoh multi-agen

Untuk mempelajari cara membuat sistem multi-agen menggunakan Genie, lihat Menggunakan Genie dalam sistem multi-agen.

Agen Streaming Output

Agen pemancar mengirimkan respons dalam aliran berkelanjutan dari potongan tambahan yang lebih kecil. Streaming mengurangi latensi yang dirasakan dan meningkatkan pengalaman pengguna untuk agen percakapan.

Untuk menulis streaming ChatAgent, tentukan predict_stream metode yang mengembalikan generator yang menghasilkan ChatAgentChunk objek - masing-masing ChatAgentChunk berisi sebagian respons. Baca selengkapnya tentang perilaku streaming ChatAgent ideal di dokumen MLflow.

Kode berikut menunjukkan contoh fungsi predict_stream, untuk contoh lengkap agen streaming, lihat contoh ChatAgent:

def predict_stream(
  self,
  messages: list[ChatAgentMessage],
  context: Optional[ChatContext] = None,
  custom_inputs: Optional[dict[str, Any]] = None,
) -> Generator[ChatAgentChunk, None, None]:
  # Convert messages to a format suitable for your agent
  request = {"messages": self._convert_messages_to_dict(messages)}

  # Stream the response from your agent
  for event in self.agent.stream(request, stream_mode="updates"):
    for node_data in event.values():
      # Yield each chunk of the response
      yield from (
        ChatAgentChunk(**{"delta": msg}) for msg in node_data["messages"]
      )

Persiapkan ChatAgentuntuk penyebaran model di Databricks Model Serving.

Databricks menyebarkan ChatAgentdi lingkungan terdistribusi pada Databricks Model Serving, yang berarti bahwa selama percakapan multi-giliran, replika server yang sama mungkin tidak dapat menangani semua permintaan. Perhatikan implikasi yang harus diperhatikan dalam mengelola status agen berikut:

  • Hindari penyimpanan sementara lokal: Saat menyebarkan ChatAgent, jangan berasumsi replika yang sama akan menangani semua permintaan dalam percakapan berkelanjutan. Mengkonstruksi kembali status internal menggunakan skema kamus ChatAgentRequest untuk setiap langkah.

  • keadaan yang aman dari thread: Rancang keadaan agen agar thread-safe, mencegah konflik dalam lingkungan multi-thread.

  • Menginisialisasi status dalam fungsi predict: Menginisialisasi status setiap kali fungsi predict dipanggil, bukan selama inisialisasi ChatAgent. Menyimpan status di tingkat ChatAgent dapat membocorkan informasi antara percakapan dan menyebabkan konflik karena satu replika ChatAgent dapat menangani permintaan dari beberapa percakapan.

Input dan output kustom

Beberapa skenario mungkin memerlukan input agen tambahan, seperti client_type dan session_id, atau output seperti tautan sumber pengambilan yang tidak boleh disertakan dalam riwayat obrolan untuk interaksi di masa mendatang.

Untuk skenario ini, MLflow ChatAgent secara asli mendukung bidang custom_inputs dan custom_outputs.

Peringatan

Aplikasi ulasan Agent Evaluation saat ini tidak mendukung penelusuran rendering untuk agen dengan bidang input tambahan.

Lihat buku catatan berikut untuk mempelajari cara mengatur input dan output kustom.

Buku catatan agen skema kustom OpenAI + PyFunc

Ambil buku catatan

Buku catatan agen skema kustom LangGraph

Ambil buku catatan

Berikan custom_inputs di AI Playground dan aplikasi ulasan agen

Jika agen Anda menerima input tambahan menggunakan bidang custom_inputs, Anda dapat memberikan input ini secara manual di AI Playground dan aplikasi peninjauan agen.

  1. Di AI Playground atau Aplikasi Tinjauan Agen, pilih ikon roda gigi..

  2. Aktifkan masukan_kustom.

  3. Berikan objek JSON yang cocok dengan skema input yang ditentukan agen Anda.

    Sediakan masukan_kustom pada lingkungan AI.

Tentukan skema pengalih kustom

Agen AI biasanya menggunakan pengambil data untuk menemukan dan melakukan pencarian pada data yang tidak terstruktur di indeks pencarian vektor. Misalnya alat retriever, lihat Membuat dan melacak alat retriever untuk data yang tidak terstruktur.

Lacak pengambil data ini di dalam agen Anda dengan rentang MLflow RETRIEVER untuk mengaktifkan fitur produk Databricks, termasuk:

  • Secara otomatis menampilkan tautan ke dokumen sumber yang diambil di antarmuka pengguna AI Playground
  • Menjalankan penilaian otomatis untuk keterkaitan dan relevansi pengambilan data dalam Evaluasi Agen

Nota

Databricks merekomendasikan penggunaan alat retriever yang disediakan oleh paket Databricks AI Bridge seperti databricks_langchain.VectorSearchRetrieverTool dan databricks_openai.VectorSearchRetrieverTool karena sudah sesuai dengan skema MLflow retriever. Lihat Kembangkan secara lokal alat pengambilan Pencarian Vektor dengan AI Bridge.

Jika agen Anda menyertakan rentang penjemput dengan skema kustom, panggil mlflow.models.set_retriever_schema saat Anda mendefinisikan agen Anda dalam kode. Ini mencocokkan kolom output retriever Anda ke bidang yang diharapkan oleh MLflow (primary_key, text_column, doc_uri).

import mlflow
# Define the retriever's schema by providing your column names
# For example, the following call specifies the schema of a retriever that returns a list of objects like
# [
#     {
#         'document_id': '9a8292da3a9d4005a988bf0bfdd0024c',
#         'chunk_text': 'MLflow is an open-source platform, purpose-built to assist machine learning practitioners...',
#         'doc_uri': 'https://mlflow.org/docs/latest/index.html',
#         'title': 'MLflow: A Tool for Managing the Machine Learning Lifecycle'
#     },
#     {
#         'document_id': '7537fe93c97f4fdb9867412e9c1f9e5b',
#         'chunk_text': 'A great way to get started with MLflow is to use the autologging feature. Autologging automatically logs your model...',
#         'doc_uri': 'https://mlflow.org/docs/latest/getting-started/',
#         'title': 'Getting Started with MLflow'
#     },
# ...
# ]
mlflow.models.set_retriever_schema(
    # Specify the name of your retriever span
    name="mlflow_docs_vector_search",
    # Specify the output column name to treat as the primary key (ID) of each retrieved document
    primary_key="document_id",
    # Specify the output column name to treat as the text content (page content) of each retrieved document
    text_column="chunk_text",
    # Specify the output column name to treat as the document URI of each retrieved document
    doc_uri="doc_uri",
    # Specify any other columns returned by the retriever
    other_columns=["title"],
)

Nota

Kolom doc_uri sangat penting saat mengevaluasi performa retriever. doc_uri adalah pengidentifikasi utama untuk dokumen yang dikembalikan oleh retriever, memungkinkan Anda membandingkannya dengan kumpulan evaluasi kebenaran dasar. Lihat Set evaluasi (MLflow 2).

Memparameterkan kode agen untuk penyebaran di seluruh lingkungan

Anda dapat memparametrisasi kode agen untuk menggunakan kembali kode agen yang sama di berbagai lingkungan.

Parameter adalah pasangan kunci-nilai yang Anda tentukan dalam kamus Python atau file .yaml.

Untuk mengonfigurasi kode, buat ModelConfig menggunakan kamus Python atau file .yaml . ModelConfig adalah sekumpulan parameter kunci-nilai yang memungkinkan manajemen konfigurasi fleksibel. Misalnya, Anda dapat menggunakan kamus selama pengembangan lalu mengonversinya ke file .yaml untuk penyebaran produksi dan CI/CD.

Untuk detail tentang ModelConfig, lihat dokumentasi MLflow.

Contoh ModelConfig ditunjukkan di bawah ini:

llm_parameters:
  max_tokens: 500
  temperature: 0.01
model_serving_endpoint: databricks-meta-llama-3-3-70b-instruct
vector_search_index: ml.docs.databricks_docs_index
prompt_template: 'You are a hello world bot. Respond with a reply to the user''s
  question that indicates your prompt template came from a YAML file. Your response
  must use the word "YAML" somewhere. User''s question: {question}'
prompt_template_input_vars:
  - question

Dalam kode agensi Anda, Anda dapat mengacu pada konfigurasi bawaan (pengembangan) dari file atau kamus .yaml.

import mlflow
# Example for loading from a .yml file
config_file = "configs/hello_world_config.yml"
model_config = mlflow.models.ModelConfig(development_config=config_file)

# Example of using a dictionary
config_dict = {
    "prompt_template": "You are a hello world bot. Respond with a reply to the user's question that is fun and interesting to the user. User's question: {question}",
    "prompt_template_input_vars": ["question"],
    "model_serving_endpoint": "databricks-meta-llama-3-3-70b-instruct",
    "llm_parameters": {"temperature": 0.01, "max_tokens": 500},
}

model_config = mlflow.models.ModelConfig(development_config=config_dict)

# Use model_config.get() to retrieve a parameter value
# You can also use model_config.to_dict() to convert the loaded config object
# into a dictionary
value = model_config.get('sample_param')

Kemudian, saat mencatat agen Anda, tentukan parameter model_config ke log_model untuk menentukan set parameter kustom yang akan digunakan saat memuat agen yang dicatat. Lihat dokumentasi MLflow - ModelConfig.

Penyebaran kesalahan pada streaming

Penyebaran kesalahan yang ditemui saat streaming oleh Mosaic AI dengan token terakhir di bawah databricks_output.error. Menjadi tanggung jawab klien yang melakukan panggilan untuk menangani dan menampilkan kesalahan ini dengan benar.

{
  "delta": …,
  "databricks_output": {
    "trace": {...},
    "error": {
      "error_code": BAD_REQUEST,
      "message": "TimeoutException: Tool XYZ failed to execute."
    }
  }
}

Langkah berikutnya