piatra-ai

Thin wrapper over the official openai SDK, preset for the Piatra AI Gateway. It does exactly three things: presets the gateway base URL, reads the key from PIATRA_AI_KEY, and exposes typed piatra.* model-id constants. Everything else — chat, streaming, retries — is the OpenAI SDK's, unchanged.

Install

pip install piatra-ai

Imported as piatra_ai.

Usage

from piatra_ai import PiatraAI, MODELS

client = PiatraAI()
# api_key  defaults to os.environ["PIATRA_AI_KEY"]
# base_url defaults to https://ai.piatra.center/v1

res = client.chat.completions.create(
    model=MODELS.GLM,  # "piatra.glm"
    messages=[{"role": "user", "content": "What is default logic?"}],
)

Streaming and multimodal input are the OpenAI SDK's, unchanged:

stream = client.chat.completions.create(
    model=MODELS.GLM, messages=messages, stream=True,
)
for chunk in stream:
    ...

client.chat.completions.create(
    model=MODELS.MEDGEMMA,  # the vision-capable model
    messages=[{
        "role": "user",
        "content": [
            {"type": "text", "text": "describe this radiograph"},
            {"type": "image_url", "image_url": {"url": data_url}},
        ],
    }],
)

Configuration

Any other OpenAI client keyword argument is accepted and passed through unchanged.