import json
import os

from axllm import OpenAICompatibleClient, agent
from axllm.runtime_quickjs import AxQuickJsCodeRuntime

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},
)

# The handbook can be arbitrarily large. Listing it in `contextFields` keeps it
# in the agent's runtime so it never inflates the model prompt -- the agent reads
# it through code, not through tokens. That is the whole point of an Ax agent
# over a plain gen() call: the source material stays out of the context window.
handbook = """
# Acme Cloud -- Support Handbook

## Billing
- Invoices are issued on the 1st of each month and are due net-15.
- Plan downgrades take effect at the END of the current billing cycle, not immediately.
- Refunds are issued to the original payment method within 5 business days.

## Access
- Seats can be added by any workspace Owner under Settings -> Members.
- SSO (SAML) is available on Enterprise; SCIM provisioning is Owner-only.

## Incidents
- Status and uptime are published at status.acme.example.
- Sev-1 incidents page the on-call within 5 minutes; updates post every 30 minutes.

## Data
- Exports are available in CSV and JSON from Settings -> Data.
- Deleted workspaces are recoverable for 30 days, then permanently purged.
""".strip()

assistant = agent(
    'question:string, handbook:string -> answer:string, citations:string[] "Handbook sections the answer relies on"',
    # Keep the handbook in the runtime, out of the prompt.
    {"contextFields": ["handbook"], "runtime": {"language": "JavaScript"}},
)

result = assistant.forward(
    client,
    {
        "question": "A customer downgraded their plan today. When does it take effect, and can they get a refund for the current cycle?",
        "handbook": handbook,
    },
    {"runtime": AxQuickJsCodeRuntime(), "max_actor_steps": 12},
)

print(json.dumps(result, indent=2, sort_keys=True))

import json
import os

from axllm import OpenAICompatibleClient, agent
from axllm.runtime_quickjs import AxQuickJsCodeRuntime

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},
)

# A raw, noisy incident report. It lives in `contextFields`, so the agent works
# it inside the runtime; `contextPolicy: lean` keeps the prompt compact by
# preferring live runtime state and summaries over replaying the raw text.
report = """
[2026-03-02 14:01:22Z] INFO  gateway       deploy svc-checkout-edge v812 -> prod (channel: canary 10%)
[2026-03-02 14:03:10Z] WARN  checkout-api  p95 latency 1180ms (baseline 240ms) region=eu-west-1
[2026-03-02 14:04:55Z] ERROR checkout-api  502 from svc-payments-gw: upstream timeout (10s) tenant_tier=enterprise
[2026-03-02 14:05:01Z] ERROR checkout-api  502 from svc-payments-gw: upstream timeout (10s) tenant_tier=enterprise
[2026-03-02 14:05:40Z] WARN  payments-gw   circuit half-open, 3 retries exhausted for order=ord_99214
[2026-03-02 14:06:12Z] INFO  gateway       canary widened 10% -> 50% for svc-checkout-edge v812
[2026-03-02 14:07:33Z] ERROR checkout-api  502 from svc-payments-gw: upstream timeout (10s) tenant_tier=enterprise
[2026-03-02 14:08:02Z] ERROR checkout-api  user-visible: "Payment could not be processed" shown to 1,284 sessions
[2026-03-02 14:09:48Z] WARN  payments-gw   connection pool exhausted (max=64) waiting=210
[2026-03-02 14:11:20Z] INFO  on-call       paged: SEV-2 opened (eu-west-1 checkout error rate 38%)
[2026-03-02 14:14:05Z] INFO  gateway       rollback svc-checkout-edge v812 -> v811 (channel: prod 100%)
[2026-03-02 14:17:41Z] INFO  checkout-api  p95 latency 260ms, error rate 0.4% region=eu-west-1
[2026-03-02 14:19:10Z] INFO  on-call       SEV-2 mitigated, monitoring for 30m
""".strip()

triage = agent(
    'report:string, question:string -> severity:class "low, medium, high, critical", rootCause:string, nextSteps:string[], evidence:string[] "Quoted log lines that support the assessment"',
    {
        "contextFields": ["report"],
        "contextPolicy": {"preset": "lean", "budget": "balanced"},
        "runtime": {"language": "JavaScript"},
    },
)

result = triage.forward(
    client,
    {
        "report": report,
        "question": "What happened, how bad was it, and what should the on-call do next? Cite the lines you relied on.",
    },
    {"runtime": AxQuickJsCodeRuntime(), "max_actor_steps": 12},
)

print(json.dumps(result, indent=2, sort_keys=True))

import json
import os
import threading
import time
from axllm import ai, agent, fn, run_control
from axllm.runtime_quickjs import AxQuickJsCodeRuntime

key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY.")
client = ai("openai", api_key=key, model="gpt-6-astra",
            model_config={"thinkingTokenBudget": "low", "max_tokens": 4096})
pending, finished, overlap = threading.Event(), threading.Event(), threading.Event()
control = run_control()
applied = []

def slow_reference(_):
    pending.set()
    time.sleep(6)
    finished.set()
    return "REF-42"

def local_label(_):
    if pending.wait(3) and not finished.is_set():
        overlap.set()
    return "LAUNCH"

def on_event(event):
    if event["type"] == "tool.started" and not applied:
        applied.append("queued")
        control.steer("Include the word VERIFIED in the final answer.")
        control.set_thinking_token_budget("medium")
    if event["type"] == "applied":
        applied.append(event["timing"])

control.on_event(on_event)
program = agent("question -> answer", {"directResponse": "off", "runtime": {"language": "JavaScript"}, "functions": [
    fn("slow_reference").description("Look up a reference; takes a few seconds.")
        .execution("background").handler(slow_reference).build(),
    fn("local_label").description("Read an independent local label immediately.")
        .execution("background").handler(local_label).build(),
]})
result = program.forward(client, {"question": "Use the native tools tools_slow_reference and tools_local_label. First call tools_slow_reference. While it is pending, call tools_local_label. Call each tool only once; do not call a tool again while its result is pending. In the executor, call the native tools directly rather than invoking them from actor code; then use final(...) in the code runtime to pass their results to the responder. Return both exact results in one sentence."},
                         {"control": control, "runtime": AxQuickJsCodeRuntime(), "serviceTier": "standard", "maxSteps": 6, "max_actor_steps": 12})
assert overlap.is_set(), "The model did not perform independent work while the background tool was pending"
assert all(word in result["answer"] for word in ("REF-42", "LAUNCH", "VERIFIED")), result
assert "next-response" in applied, applied
print(json.dumps({"result": result, "background_overlap": True, "control_timing": applied[1:]}, indent=2))

import json
import os
from axllm import ai, agent, run_control
from axllm.runtime_quickjs import AxQuickJsCodeRuntime

key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY.")
client = ai("openai", api_key=key, model="gpt-6-astra",
            model_config={"thinkingTokenBudget": "low", "max_tokens": 4096})
control = run_control()
events = []
control.on_event(events.append)
control.steer("Include VERIFIED in the final answer.")
control.steer("Include CHILD-CHECK in the final answer.", target="root/team.researcher")
control.set_thinking_token_budget("medium", target="root/team.researcher/executor")
child = agent("question -> answer", {"directResponse": "off", "runtime": AxQuickJsCodeRuntime()})
parent = agent("question -> answer", {"directResponse": "off", "runtime": AxQuickJsCodeRuntime()})
parent.add_child_agent("team", "researcher", child)
result = parent.forward(client, {"question": "Delegate to team.researcher exactly once by calling await team.researcher({question: \"Compute 37 + 5 and return the exact sum.\"}) in actor code. Pass the complete child answer as evidence to final(...), then report it in your final answer."},
                        {"control": control, "serviceTier": "standard", "max_actor_steps": 8})
assert all(word in result["answer"] for word in ("42", "VERIFIED", "CHILD-CHECK")), result
assert "team.researcher" in parent.get_usage()["children"]
assert any(event["type"] == "applied" and event["path"] == "root/team.researcher/executor" for event in events)
print(json.dumps({"result": result, "child_usage": parent.get_usage()["children"]}, indent=2))

import json
import os

from axllm import OpenAICompatibleClient, agent
from axllm.runtime_quickjs import AxQuickJsCodeRuntime

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},
)

# A long, messy brief -- exactly the kind of input you do not want replayed into
# the prompt on every turn. `contextFields` holds it in the runtime, the
# `checkpointed` policy compacts older turns once the prompt grows, and
# `maxRuntimeChars` caps how much runtime output is echoed back.
brief = """
# Migration brief: monolith -> services (draft, unordered notes)

Current: single Rails monolith, Postgres primary + 1 replica, Sidekiq for jobs.
Pain: deploys take 40m, one bad migration locks the orders table, on-call burnout.
Constraints: no downtime windows > 5m, PCI scope must shrink, team of 6, 2 quarters.
Hot paths: checkout (writes orders, payments), search (read-heavy), notifications (async).
Known landmines: payments code has no tests; search shares the orders DB; a nightly
cron rebuilds the catalog and pins CPU for ~20m; the replica lags up to 90s under load.
Org wants: independent deploys for checkout, smaller blast radius, an audit trail.
Nice to have: event log for orders, read-model for search, feature flags.
Hard no: a big-bang rewrite; introducing Kubernetes this year.
""".strip()

specialist = agent(
    'brief:string, goal:string -> plan:string[] "Ordered, concrete steps", answer:string, risks:string[]',
    {
        "contextFields": ["brief"],
        "contextPolicy": {"preset": "checkpointed", "budget": "balanced"},
        "maxRuntimeChars": 3000,
        "runtime": {"language": "JavaScript"},
    },
)

result = specialist.forward(
    client,
    {
        "brief": brief,
        "goal": "Propose a safe, incremental 2-quarter plan to split checkout out first, respecting the hard constraints.",
    },
    {"runtime": AxQuickJsCodeRuntime(), "max_actor_steps": 12},
)

print(json.dumps(result, indent=2, sort_keys=True))

import json
import os

from axllm import AnthropicClient, GoogleGeminiClient, OpenAICompatibleClient, ax

openai_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
google_key = os.getenv("GOOGLE_APIKEY") or os.getenv("GOOGLE_API_KEY")
anthropic_key = os.getenv("ANTHROPIC_APIKEY") or os.getenv("ANTHROPIC_API_KEY")
if not (openai_key and google_key and anthropic_key):
    raise SystemExit("Set OPENAI_APIKEY, GOOGLE_APIKEY, and ANTHROPIC_APIKEY to run this multi-provider panel.")

# A panel of three different providers, each answering the same question
# independently. Plain ax() composition (no agent runtime): fan out to the
# panel, judge the candidates, then synthesize one grounded answer.
panel = [
    ("openai/gpt-5.4-mini", OpenAICompatibleClient(api_key=openai_key, model="gpt-5.4-mini", model_config={"temperature": 0})),
    ("google/gemini-3.5-flash", GoogleGeminiClient(api_key=google_key, model="gemini-3.5-flash")),
    ("anthropic/claude-haiku-4.5", AnthropicClient(api_key=anthropic_key, model="claude-haiku-4-5")),
]

researcher = ax("question:string -> answer:string, keyFindings:string[], citations:string[], confidence:number")
researcher.set_instruction("Answer independently. Use evidence. Call out uncertainty. Do not optimize for consensus.")

judge = ax("question:string, candidates:json -> consensus:string[], contradictions:string[], uniqueInsights:string[], blindSpots:string[]")
judge.set_instruction("Compare the candidates. Find agreement, conflicts, missing coverage, and unique useful points.")

synthesizer = ax("question:string, candidates:json, review:json -> answer:string, citations:string[], caveats:string[]")
synthesizer.set_instruction("Write one final answer grounded in the candidates and review. Resolve conflicts explicitly.")

question = "What are the strongest arguments for and against a national carbon tax?"

candidates = []
for model, client in panel:
    response = researcher.forward(client, {"question": question})
    candidates.append({"model": model, **response})

# The judge + synthesizer run on one of the panel clients (OpenAI here).
orchestrator = panel[0][1]
review = judge.forward(orchestrator, {"question": question, "candidates": candidates})
final = synthesizer.forward(orchestrator, {"question": question, "candidates": candidates, "review": review})

print(json.dumps(final, indent=2, sort_keys=True))