import json
import os

from axllm import OpenAICompatibleClient, ax


api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not api_key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")

client = OpenAICompatibleClient(
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    model_config={"temperature": 0},
)
program = ax('question:string -> answer:string')
out = program.forward(client, {"question": "In one sentence, explain Ax as a language-agnostic LLM programming library."})
print(json.dumps(out, indent=2, sort_keys=True))

import json
import os

from axllm import OpenAICompatibleClient, ax


api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not api_key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")

client = OpenAICompatibleClient(
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    model_config={"temperature": 0},
)
program = ax('ticket:string -> priority:class "high, normal, low", summary:string, labels:string[]')
out = program.forward(client, {"ticket": "Checkout has failed for enterprise customers since 09:00. Support wants a concise summary and tags."})
print(json.dumps(out, indent=2, sort_keys=True))

import json
import os

from axllm import OpenAICompatibleClient, ax, f


api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not api_key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")

client = OpenAICompatibleClient(
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    model_config={"temperature": 0},
)
signature = (
    f()
    .input("requestText", f.string("Booking request").min(10).max(500))
    .input("contactEmail", f.string("Contact email").email())
    .output("partySize", f.number("Guests").min(1).max(12))
    .output(
        "bookingCode",
        f.string("Three letters, a dash, and four digits").regex(
            r"^[A-Z]{3}-\d{4}$", "Must look like ABC-1234"
        ),
    )
    .output(
        "guestProfile",
        f.object(
            {
                "fullName": f.string("Primary guest").min(2),
                "dietaryNotes": f.string("Dietary requirements").optional(),
            }
        ),
    )
    .build()
)
output = ax(signature).forward(
    client,
    {
        "requestText": "Book dinner for four people under the name Ada Lovelace.",
        "contactEmail": "ada@example.com",
    },
)
print(json.dumps(output, indent=2, sort_keys=True))

import json
import os
import uuid

from axllm import OpenAICompatibleClient, set_usage_observer


api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not api_key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")

events = []
set_usage_observer(events.append)
client = OpenAICompatibleClient(
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    usage_context={
        "tenantId": "tenant-42",
        "feature": "support-chat",
        "attributes": {"environment": "example"},
    },
)

try:
    client.chat(
        {"chat_prompt": [{"role": "user", "content": "Reply with one short greeting."}]},
        {"usageContext": {"userId": "user-7", "requestId": str(uuid.uuid4())}},
    )
    print(json.dumps(events, indent=2, sort_keys=True))
finally:
    set_usage_observer(None)

import json
import os

from axllm import OpenAICompatibleClient, ax


api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not api_key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")

client = OpenAICompatibleClient(
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    model_config={"temperature": 0},
)
program = ax('context:string, question:string -> answer:string, citations:string[]')
out = program.forward(client, {"context": "Ax uses signatures, ai(), ax(), agent(), flow(), and optimize() for production LLM programs.", "question": "How should a new developer think about Ax?"})
print(json.dumps(out, indent=2, sort_keys=True))

import os

from axllm import (
    AxBalancer,
    AxBalancerAdaptiveStrategy,
    AxBalancerOptions,
    AxInMemoryBalancerStatsStore,
    OpenAICompatibleClient,
)


api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not api_key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")

model = os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini")
services = [
    OpenAICompatibleClient(
        api_key=api_key,
        model=model,
        base_url=os.getenv("OPENAI_PRIMARY_BASE_URL", "https://api.openai.com/v1"),
    ),
    OpenAICompatibleClient(
        api_key=os.getenv("OPENAI_BACKUP_API_KEY", api_key),
        model=model,
        base_url=os.getenv("OPENAI_BACKUP_BASE_URL", "https://api.openai.com/v1"),
    ),
]

# Reuse this store across balancers in one process. A Redis/database adapter can
# implement the same atomic get/observe contract for multi-process state.
stats_store = AxInMemoryBalancerStatsStore()
route_keys = ["openai-primary", "openai-backup"]
events = []
strategy = AxBalancerAdaptiveStrategy(
    deadline_ms=6_000,
    bad_outcome_cost=0.02,
    expected_tokens={"promptTokens": 1_200, "completionTokens": 300},
    namespace="support-summary-v1",
    route_key=lambda _service, index: route_keys[index],
    slice=lambda context: "streaming" if context["options"].get("stream") else "interactive",
    stats_store=stats_store,
    on_routing_event=lambda event: events.append(event),
)
balancer = AxBalancer(services, AxBalancerOptions(strategy=strategy))
response = balancer.chat(
    {"model": model, "chat_prompt": [{"role": "user", "content": "Summarize why shared routing state matters."}]}
)
print(response["results"][0]["content"])
print([event["type"] for event in events])