Integração com o Databricks

O Model Serving do Databricks hospeda tanto modelos externos quanto agentes criados com a interface ResponsesAgent do MLflow. Registre o AI Gateway como um modelo externo para acessar mais de 500 modelos de dentro do Databricks e instrumente agentes com OpenTelemetry para capturar suas chamadas de LLM no Toqan Router.

AI Gateway

Sirva qualquer modelo do Toqan Router por meio de um serving endpoint do Databricks, com acompanhamento de custos e fallbacks.

Observabilidade

Exporte traces de agentes em execução no Model Serving, com prompts, tokens e custo.

Pré-requisitos

Uma conta no Toqan Router e uma chave de API
Um workspace do Databricks com o Model Serving habilitado
MLflow 3.11 ou posterior, além de databricks-sdk e databricks-agents instalados localmente

Armazenar a chave de API

As duas seções abaixo leem a chave de um secret do Databricks. Mantenha-a lá em vez de colocá-la na configuração do endpoint, onde qualquer pessoa com acesso de visualização ao endpoint pode lê-la. Ignore create_scope se o scope já existir:

from databricks.sdk import WorkspaceClient

workspace = WorkspaceClient()
workspace.secrets.create_scope(scope="toqan")
workspace.secrets.put_secret(
    scope="toqan", key="api_key", string_value="<TOQAN_API_KEY>"
)

AI Gateway

Os modelos externos do Databricks aceitam qualquer endpoint compatível com OpenAI por meio do provedor custom, então o AI Gateway pode servir de base para um serving endpoint do Databricks. Defina external_model.name como o modelo a ser servido e aponte custom_provider_url para o AI Gateway:

import mlflow.deployments

client = mlflow.deployments.get_deploy_client("databricks")

client.create_endpoint(
    config={
        "name": "toqan-gateway",
        "config": {
            "served_entities": [
                {
                    "name": "toqan-gateway",
                    "external_model": {
                        "name": "openai/gpt-5.6-sol",
                        "provider": "custom",
                        "task": "llm/v1/chat",
                        "custom_provider_config": {
                            "custom_provider_url": "https://router.toqan.ai/eu/v3/router/chat/completions",
                            "bearer_token_auth": {
                                "token": "{{secrets/toqan/api_key}}"
                            },
                        },
                    },
                }
            ]
        },
    },
)

Consulte-o como qualquer outro serving endpoint do Databricks. As chamadas aparecem em Traces com modelo, uso de tokens e custo:

response = client.predict(
    endpoint="toqan-gateway",
    inputs={"messages": [{"role": "user", "content": "Hello"}]},
)
print(response["choices"][0]["message"]["content"])

Qualquer modelo habilitado funciona. Defina external_model.name como o slug do modelo no Toqan Router e consulte os modelos compatíveis. Um endpoint serve um único modelo, então crie um endpoint por modelo.

Observabilidade

Adicionar o OpenTelemetry a um agente já em execução no Model Serving envia as chamadas de LLM dele ao Toqan Router, com o conteúdo completo de solicitações e respostas, uso de tokens e custo. São quatro alterações: adicionar um exporter no nível do módulo, fazer o flush ao final de predict, declarar três pacotes extras e passar o endpoint e a chave no momento do deploy.

Isto abrange agentes implantados a partir de código com agents.deploy(). Agentes criados no construtor no-code Agent Bricks são executados em um endpoint gerenciado pelo Databricks, que não expõe variáveis de ambiente nem dependências do modelo, portanto não podem ser instrumentados dessa forma.

Adicionar o exporter ao agente

Adicione isto no nível do módulo no arquivo do agente, acima da classe do agente. Isso abre a conexão com o Toqan Router e anexa o instrumentor que registra as chamadas de LLM. O OpenAIInstrumentor registra chamadas feitas pelo SDK da OpenAI, que é o que WorkspaceClient().serving_endpoints.get_open_ai_client() retorna. Agentes criados com outra biblioteca cliente precisam do instrumentor OpenInference correspondente no lugar dele; caso contrário, o agente é implantado sem erros, mas não produz spans:

import os

from openinference.instrumentation.openai import OpenAIInstrumentor
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor

tracer_provider = TracerProvider()
tracer_provider.add_span_processor(
    BatchSpanProcessor(
        OTLPSpanExporter(
            endpoint=os.environ["TOQAN_OTEL_ENDPOINT"],
            headers={"Authorization": f"Bearer {os.environ['TOQAN_API_KEY']}"},
        )
    )
)
OpenAIInstrumentor().instrument(tracer_provider=tracer_provider)

Em seguida, faça o flush ao final de predict, antes de retornar. O Model Serving mantém o contêiner ativo entre as solicitações, então é o flush que envia os spans prontamente em vez de deixá-los na fila:

    def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
        response = ...  # existing agent logic, unchanged
        tracer_provider.force_flush(15000)
        return response

Declarar as dependências

O ambiente de serving instala apenas os pacotes declarados no modelo, então adicione estes ao pip_requirements existente em log_model. Sem eles, o agente falha ao importar no momento da solicitação:

        pip_requirements=[
            # existing requirements
            "opentelemetry-sdk",
            "opentelemetry-exporter-otlp-proto-http",
            "openinference-instrumentation-openai",
        ],

Passar o endpoint e a chave

Adicione estes a environment_vars em agents.deploy(). OTEL_SERVICE_NAME é opcional e define o nome usado posteriormente para filtrar traces:

    environment_vars={
        "TOQAN_OTEL_ENDPOINT": "https://router.toqan.ai/eu/v2/otel/v1/traces",
        "TOQAN_API_KEY": "{{secrets/toqan/api_key}}",
        "OTEL_SERVICE_NAME": "databricks-support-agent",
    },

Configurar o OTLP apenas pelas variáveis padrão OTEL_EXPORTER_OTLP_* não exporta traces do Model Serving. O runtime gerenciado configura seu próprio tracing do MLflow, que tem precedência, e os spans chegam ao experimento do Databricks em vez do Toqan Router. Em vez disso, adicione o exporter ao agente conforme mostrado acima.

Uma referência a secret só é substituída quando é o valor inteiro de uma variável. Escrever Authorization=Bearer {{secrets/toqan/api_key}} envia as chaves literalmente, e o Toqan Router rejeita os spans com um 401. Passe a chave em sua própria variável, como acima, e monte o header no agente.

Visualizar traces

Os traces aparecem no AI Studio na aba Traces. Cada chamada de LLM chega como um span contendo as mensagens de solicitação e resposta, os tokens de prompt e de completion, o modelo e o provedor resolvidos e o custo calculado. Filtre pelo valor definido em OTEL_SERVICE_NAME para isolar um agente.

Exemplo completo

Este exemplo chama um endpoint nativo do Databricks. Para rotear o agente pelo AI Gateway, defina LLM_ENDPOINT como o endpoint criado acima.

Módulo completo do agente e script de deploy

import os

from databricks.sdk import WorkspaceClient
from mlflow.models import set_model
from mlflow.pyfunc import ResponsesAgent
from mlflow.types.responses import (
    ResponsesAgentRequest,
    ResponsesAgentResponse,
    create_text_output_item,
)
from openinference.instrumentation.openai import OpenAIInstrumentor
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor

LLM_ENDPOINT = "databricks-gemini-3-5-flash"

tracer_provider = TracerProvider()
tracer_provider.add_span_processor(
    BatchSpanProcessor(
        OTLPSpanExporter(
            endpoint=os.environ["TOQAN_OTEL_ENDPOINT"],
            headers={"Authorization": f"Bearer {os.environ['TOQAN_API_KEY']}"},
        )
    )
)
OpenAIInstrumentor().instrument(tracer_provider=tracer_provider)


class SupportAgent(ResponsesAgent):
    def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
        client = WorkspaceClient().serving_endpoints.get_open_ai_client()
        completion = client.chat.completions.create(
            model=LLM_ENDPOINT,
            messages=[{"role": "user", "content": self._question(request)}],
            max_tokens=300,
        )
        tracer_provider.force_flush(15000)
        return ResponsesAgentResponse(
            output=[
                create_text_output_item(
                    text=completion.choices[0].message.content, id="answer"
                )
            ]
        )

    def _question(self, request: ResponsesAgentRequest) -> str:
        last_message = request.input[-1]
        as_dict = (
            last_message
            if isinstance(last_message, dict)
            else last_message.model_dump()
        )
        return as_dict.get("content", "")


set_model(SupportAgent())

import mlflow
from databricks import agents
from mlflow.models.resources import DatabricksServingEndpoint

mlflow.set_registry_uri("databricks-uc")

with mlflow.start_run():
    logged_agent = mlflow.pyfunc.log_model(
        name="agent",
        python_model="agent.py",
        resources=[
            DatabricksServingEndpoint(endpoint_name="databricks-gemini-3-5-flash")
        ],
        pip_requirements=[
            "mlflow>=3.11",
            "databricks-sdk",
            "openai",
            "httpx",
            "opentelemetry-sdk",
            "opentelemetry-exporter-otlp-proto-http",
            "openinference-instrumentation-openai",
        ],
    )

registered_agent = mlflow.register_model(
    logged_agent.model_uri, "catalog.schema.support_agent"
)

agents.deploy(
    "catalog.schema.support_agent",
    registered_agent.version,
    environment_vars={
        "TOQAN_OTEL_ENDPOINT": "https://router.toqan.ai/eu/v2/otel/v1/traces",
        "TOQAN_API_KEY": "{{secrets/toqan/api_key}}",
        "OTEL_SERVICE_NAME": "databricks-support-agent",
    },
)

Visão geral

Instrumente seu app

Explorar

Monitorar

Exportar

Revisar

Privacidade e segurança de dados

Integrações com frameworks

AI Gateway

Observabilidade

Pré-requisitos

Armazenar a chave de API

AI Gateway

Observabilidade

Adicionar o exporter ao agente

Declarar as dependências

Passar o endpoint e a chave

Visualizar traces

Exemplo completo

AI Gateway

Observabilidade

​Pré-requisitos

​Armazenar a chave de API

​AI Gateway

​Observabilidade

​Adicionar o exporter ao agente

​Declarar as dependências

​Passar o endpoint e a chave

​Visualizar traces

​Exemplo completo

Pré-requisitos

Armazenar a chave de API

AI Gateway

Observabilidade

Adicionar o exporter ao agente

Declarar as dependências

Passar o endpoint e a chave

Visualizar traces

Exemplo completo