#!/usr/bin/env python3
"""AOCore embeddings via the OpenAI Python SDK.

The /v1/embeddings endpoint accepts a single string OR a list of strings.
The response is the standard OpenAI shape: data[].embedding is the vector.

Run:
    pip install openai
    export AOSENTRY_KEY="sk-..."
    python embeddings.py
"""

import os
import sys

from openai import OpenAI


def main() -> int:
    api_key = os.environ.get("AOSENTRY_KEY")
    if not api_key:
        print("Set AOSENTRY_KEY to your minted API key (sk-...).", file=sys.stderr)
        return 1

    base_url = os.environ.get("AOSENTRY_BASE_URL", "https://gateway.core.aocyber.ai/v1")

    client = OpenAI(base_url=base_url, api_key=api_key)

    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=[
            "AOCore is an OpenAI-compatible LLM gateway.",
            "It adds audit, guardrails, and spend control.",
        ],
    )

    # response.data is a list with one entry per input. response.usage carries
    # the gateway-reported prompt/total tokens (same shape as OpenAI's).
    for i, item in enumerate(response.data):
        vector = item.embedding
        print(f"embedding[{i}]: dim={len(vector)} first3={vector[:3]}")

    print(f"\nusage: prompt_tokens={response.usage.prompt_tokens} "
          f"total_tokens={response.usage.total_tokens}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())