AI Gateway

# 동일한 OpenAI 호환 API로 다양한 모델 호출
from openai import OpenAI

client = OpenAI(
    api_key=os.environ["DATABRICKS_TOKEN"],
    base_url=f"{os.environ['DATABRICKS_HOST']}/serving-endpoints"
)

# 엔드포인트 이름만 바꾸면 모델 전환
response = client.chat.completions.create(
    model="my-gateway-endpoint",  # AI Gateway 엔드포인트
    messages=[{"role": "user", "content": "매출 리포트를 요약해 주세요"}],
    max_tokens=500
)

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import (
    EndpointCoreConfigInput,
    ServedEntityInput,
    ExternalModel,
    OpenAiConfig
)

w = WorkspaceClient()

# OpenAI GPT-4o를 AI Gateway 뒤에 배치
w.serving_endpoints.create(
    name="gpt4o-gateway",
    config=EndpointCoreConfigInput(
        served_entities=[
            ServedEntityInput(
                external_model=ExternalModel(
                    name="gpt-4o",
                    provider="openai",
                    task="llm/v1/chat",
                    openai_config=OpenAiConfig(
                        openai_api_key="{{secrets/my-scope/openai-key}}"
                    )
                )
            )
        ]
    )
)

from databricks.sdk.service.serving import (
    AiGatewayConfig,
    AiGatewayGuardrails,
    AiGatewayGuardrailParameters,
    AiGatewayGuardrailPiiParameters
)

# 가드레일이 적용된 엔드포인트 설정
w.serving_endpoints.update_ai_gateway(
    name="customer-support-agent",
    guardrails=AiGatewayGuardrails(
        input=AiGatewayGuardrailParameters(
            # 입력에서 부적절한 콘텐츠 차단
            safety=True,
            # 허용되지 않은 주제 차단
            invalid_keywords=["경쟁사", "해킹", "불법"]
        ),
        output=AiGatewayGuardrailParameters(
            # 출력에서 부적절한 콘텐츠 차단
            safety=True,
            # PII 자동 마스킹
            pii=AiGatewayGuardrailPiiParameters(
                behavior="BLOCK"  # BLOCK 또는 MASK
            ),
            # 금칙어 차단
            invalid_keywords=["내부 전용", "기밀"]
        )
    )
)

from databricks.sdk.service.serving import (
    AiGatewayRateLimit,
    AiGatewayRateLimitRenewalPeriod,
    AiGatewayRateLimitKey
)

w.serving_endpoints.update_ai_gateway(
    name="customer-support-agent",
    rate_limits=[
        # 사용자별 분당 10회 제한
        AiGatewayRateLimit(
            calls=10,
            renewal_period=AiGatewayRateLimitRenewalPeriod.MINUTE,
            key=AiGatewayRateLimitKey.USER
        ),
        # 엔드포인트 전체 분당 100회 제한
        AiGatewayRateLimit(
            calls=100,
            renewal_period=AiGatewayRateLimitRenewalPeriod.MINUTE,
            key=AiGatewayRateLimitKey.ENDPOINT
        )
    ]
)

-- 엔드포인트별 일일 사용량 및 비용 추적
SELECT
    DATE(timestamp) AS date,
    COUNT(*) AS total_requests,
    SUM(request:usage:prompt_tokens) AS total_input_tokens,
    SUM(request:usage:completion_tokens) AS total_output_tokens,
    SUM(request:usage:total_tokens) AS total_tokens,
    -- 대략적 비용 추정 (모델별 단가는 다를 수 있음)
    ROUND(SUM(request:usage:total_tokens) / 1000000 * 2.5, 2) AS estimated_cost_usd
FROM catalog.schema.customer_support_agent_inference_log
WHERE DATE(timestamp) >= CURRENT_DATE() - INTERVAL 30 DAYS
GROUP BY DATE(timestamp)
ORDER BY date DESC;

-- Model Serving 엔드포인트별 비용 추적
SELECT
    usage_date,
    usage_metadata.endpoint_name,
    SUM(usage_quantity) AS total_dbus,
    ROUND(SUM(usage_quantity) * 0.07, 2) AS estimated_cost_usd
FROM system.billing.usage
WHERE billing_origin_product = 'MODEL_SERVING'
    AND usage_date >= CURRENT_DATE() - INTERVAL 30 DAYS
GROUP BY 1, 2
ORDER BY total_dbus DESC;

# GPT-4o 80%, Claude 20%으로 트래픽 분할
w.serving_endpoints.update_config(
    name="my-gateway",
    served_entities=[
        ServedEntityInput(
            external_model=ExternalModel(name="gpt-4o", provider="openai", ...),
            traffic_percentage=80
        ),
        ServedEntityInput(
            external_model=ExternalModel(name="claude-3-5-sonnet", provider="anthropic", ...),
            traffic_percentage=20
        )
    ]
)

# 추론 테이블 자동 캡처 설정
w.serving_endpoints.update_config(
    name="my-gateway",
    auto_capture_config=AutoCaptureConfigInput(
        catalog_name="ml_catalog",
        schema_name="serving_logs",
        enabled=True
    )
)

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

# 1. 엔드포인트 생성 (Foundation Model 사용)
endpoint = w.serving_endpoints.create(
    name="cs-chatbot-prod",
    config=EndpointCoreConfigInput(
        served_entities=[
            ServedEntityInput(
                entity_name="databricks-meta-llama-3-1-70b-instruct",
                workload_size="Small",
                scale_to_zero_enabled=False  # 프로덕션은 콜드 스타트 방지
            )
        ],
        auto_capture_config=AutoCaptureConfigInput(
            catalog_name="prod_catalog",
            schema_name="serving",
            enabled=True
        )
    )
)

# 2. 가드레일 설정
w.serving_endpoints.update_ai_gateway(
    name="cs-chatbot-prod",
    guardrails=AiGatewayGuardrails(
        input=AiGatewayGuardrailParameters(
            safety=True,
            invalid_keywords=["해킹", "불법", "경쟁사"]
        ),
        output=AiGatewayGuardrailParameters(
            safety=True,
            pii=AiGatewayGuardrailPiiParameters(behavior="MASK")
        )
    ),
    rate_limits=[
        AiGatewayRateLimit(
            calls=20,
            renewal_period=AiGatewayRateLimitRenewalPeriod.MINUTE,
            key=AiGatewayRateLimitKey.USER
        )
    ]
)