Skip to main content

A trivial set of API bindings for AI models, because I'd like them to be easy to use.

Project description

TrivialAI

(A set of httpx / boto3-based, trivial bindings for AI models — now with optional streaming)

Install

pip install trivialai
# Optional: HTTP/2 for OpenAI/Anthropic
# pip install "trivialai[http2]"
# Optional: AWS Bedrock support (via boto3)
# pip install "trivialai[bedrock]"
  • Requires Python ≥ 3.9.
  • Uses httpx for HTTP-based providers and boto3 for Bedrock.

Quick start

>>> from trivialai import claude, gcp, ollama, chatgpt, bedrock

Synchronous usage (unchanged ergonomics)

Ollama

>>> client = ollama.Ollama("gemma2:2b", "http://localhost:11434/")
# or ollama.Ollama("deepseek-coder-v2:latest", "http://localhost:11434/")
# or ollama.Ollama("mannix/llama3.1-8b-abliterated:latest", "http://localhost:11434/")
>>> client.generate("sys msg", "Say hi with 'platypus'.").content
"Hi there—platypus!"
>>> client.generate_json("sys msg", "Return {'name': 'Platypus'} as JSON").content
{'name': 'Platypus'}

Claude (Anthropic API)

>>> client = claude.Claude("claude-3-5-sonnet-20240620", os.environ["ANTHROPIC_API_KEY"])
>>> client.generate("sys msg", "Say hi with 'platypus'.").content
"Hello, platypus!"

GCP (Vertex AI)

>>> client = gcp.GCP("gemini-1.5-flash-001", "/path/to/gcp_creds.json", "us-central1")
>>> client.generate("sys msg", "Say hi with 'platypus'.").content
"Hello, platypus!"

ChatGPT (OpenAI API)

>>> client = chatgpt.ChatGPT("gpt-4o-mini", os.environ["OPENAI_API_KEY"])
>>> client.generate("sys msg", "Say hi with 'platypus'.").content
"Hello, platypus!"

AWS Bedrock (Claude / Llama / Nova / etc)

Bedrock support is provided via the Bedrock client, which implements the same LLMMixin interface as the others.

1) One-time AWS setup

  1. Enable Bedrock and model access

    • In the AWS console, pick a Bedrock-supported region (e.g. us-east-1).
    • Go to Amazon Bedrock → Model access and enable access for the models you want (e.g. Claude 3.5 Sonnet, Llama, Nova, etc).
  2. IAM permissions

    Grant your user/role permission to call Bedrock runtime APIs, for example:

    {
      "Version": "2012-10-17",
      "Statement": [
        {
          "Effect": "Allow",
          "Action": [
            "bedrock:Converse",
            "bedrock:ConverseStream",
            "bedrock:InvokeModel",
            "bedrock:InvokeModelWithResponseStream"
          ],
          "Resource": "*"
        }
      ]
    }
    

    You can restrict Resource to specific model ARNs later.

  3. Credentials

    TrivialAI can use either:

    • the normal AWS credential chain (aws configure, env vars, instance role), or
    • explicit credentials passed into the Bedrock constructor.

2) Choosing the right model_id

Bedrock distinguishes between:

  • Foundation model IDs, like: anthropic.claude-3-5-sonnet-20241022-v2:0
  • Inference profile IDs, which are region-prefixed, like: us.anthropic.claude-3-5-sonnet-20241022-v2:0

Some newer models (like Claude 3.5 Sonnet v2) must be called via the inference profile ID from certain regions. If you see a ValidationException complaining about “Invocation of model ID ... with on-demand throughput isn’t supported; retry with an inference profile”, swap to the us.-prefixed ID.

3) Minimal Bedrock demo

from trivialai import bedrock

# Using an inference profile ID for Claude 3.5 Sonnet v2 from us-east-1:
client = bedrock.Bedrock(
    model_id="us.anthropic.claude-3-5-sonnet-20241022-v2:0",
    region="us-east-1",
    # Either rely on normal AWS creds...
    # aws_profile="my-dev-profile",
    # ...or pass explicit keys (for testing):
    aws_access_key_id="AKIA...",
    aws_secret_access_key="SECRET...",
)

res = client.generate(
    "This is a test message. Make sure your reply contains the word 'margarine'",
    "Hello there! Can you hear me?"
)
print(res.content)
# -> "Yes, I can hear you! ... margarine ..."

# With JSON parsing:
res_json = client.generate_json(
    "You are a JSON-only assistant.",
    "Return {'name':'Platypus'} as JSON."
)
print(res_json.content)
# -> {'name': 'Platypus'}

The Bedrock client fully participates in the same higher-level helpers:

  • generate_checked(...)
  • generate_json(...)
  • generate_tool_call(...)
  • generate_many_tool_calls(...)
  • stream_checked(...) / stream_json(...)

No special-casing required in downstream code.


Streaming (NDJSON-style events)

All providers expose a common streaming shape via stream(...) (sync iterator) and astream(...) (async):

Event schema

  • {"type":"start", "provider": "<ollama|openai|anthropic|gcp|bedrock>", "model": "..."}

  • {"type":"delta", "text":"...", "scratchpad":"..."}

    • For Ollama, scratchpad contains model “thinking” extracted from <think>…</think>.
    • For ChatGPT, Claude API, GCP, and Bedrock, scratchpad is "" (empty) in deltas.
  • {"type":"end", "content":"...", "scratchpad": <str|None>, "tokens": <int>}

  • {"type":"error", "message":"..."}

Example: streaming Ollama (sync)

>>> client = ollama.Ollama("gemma2:2b", "http://localhost:11434/")
>>> for ev in client.stream("sys", "Explain, think step-by-step."):
...     if ev["type"] == "delta":
...         # show model output live
...         print(ev["text"], end="")
...     elif ev["type"] == "end":
...         print("\n-- scratchpad --")
...         print(ev["scratchpad"])

Example: streaming Bedrock (sync)

>>> from trivialai import bedrock
>>> client = bedrock.Bedrock(
...     model_id="us.anthropic.claude-3-5-sonnet-20241022-v2:0",
...     region="us-east-1",
... )
>>> events = list(client.stream(
...     "This is a test message. Make sure your reply contains the word 'margarine'",
...     "Hello there! Can you hear me?"
... ))
>>> events[0]
{'type': 'start', 'provider': 'bedrock', 'model': 'us.anthropic.claude-3-5-sonnet-20241022-v2:0'}
>>> events[-1]
{'type': 'end', 'content': 'Yes, I can hear you! ... margarine ...', 'scratchpad': None, 'tokens': 36}

Example: parse-at-end streaming

If you want incremental updates and a structured parse at the end:

from trivialai.util import stream_checked, loadch

for ev in client.stream("sys", "Return a JSON object gradually."):
    # pass-through for UI
    if ev["type"] in {"start","delta"}:
        print(ev)
    elif ev["type"] == "end":
        # now emit the final parsed event
        for final_ev in stream_checked(iter([ev]), loadch):
            print(final_ev)  # {"type":"final","ok":True,"parsed":{...}}

Shortcut: stream_json(system, prompt) yields the same stream and a final parsed event using loadch.

Async flavor

async for ev in client.astream("sys", "Stream something."):
    ...

(For Bedrock, stream(...) is the native streaming interface; astream(...) currently falls back to the default LLMMixin behavior unless you wrap it yourself.)


Tool Calls

Use Tools to register Python functions, describe them to the model, and safely execute the model’s chosen call.

1) Define tools

You can register functions directly or with a decorator. Docstring = description. Type hints become the argument schema.

from typing import Optional, List
from trivialai.tools import Tools

tools = Tools()  # or Tools(extras={"api_key": "..."}), see below

@tools.define()
def screenshot(url: str, selectors: Optional[List[str]] = None) -> None:
    """Take a screenshot of a page; optionally highlight CSS selectors."""
    print("shot", url, selectors)

# Or:
def search(query: str, top_k: int = 5) -> List[str]:
    """Search and return top results."""
    return [f"res{i}" for i in range(top_k)]
tools.define(search)

2) Show tools to the model

tools.list() returns LLM-friendly metadata:

>>> tools.list()
[{
  "name": "screenshot",
  "description": "Take a screenshot of a page; optionally highlight CSS selectors.",
  "type": {"url": <class 'str'>, "selectors": typing.Optional[typing.List[str]]},
  "args": {
    "url": {"type": "string"},
    "selectors": {"type": "array", "items": {"type": "string"}, "nullable": True}
  }
},
{
  "name": "search",
  "description": "Search and return top results.",
  "type": {"query": <class 'str'>, "top_k": <class 'int'>},
  "args": {
    "query": {"type": "string"},
    "top_k": {"type": "int"}
  }
}]

3) Ask the model to choose a tool

All LLM clients (Ollama, Claude API, ChatGPT, GCP, Bedrock) support a helper that prompts for a tool call and validates it:

from trivialai import ollama
client = ollama.Ollama("gemma2:2b", "http://localhost:11434/")

res = client.generate_tool_call(
    tools,
    system="You are a tool-use router.",
    prompt="Take a screenshot of https://example.com and highlight the search box."
)

# Validated, parsed dict:
>>> res.content
{'functionName': 'screenshot', 'args': {'url': 'https://example.com', 'selectors': ['#search']}}

Multiple calls? Use generate_many_tool_calls(...):

multi = client.generate_many_tool_calls(
    tools,
    prompt="Search for 'platypus', then screenshot the first result."
)
# -> [{'functionName': 'search', ...}, {'functionName': 'screenshot', ...}]

4) Validate/execute (with robust errors)

  • Validation rules: all required params present; optional params may be omitted; unknown params are rejected.
  • On invalid input, methods raise TransformError (no None returns).
from trivialai.util import TransformError

tool_call = res.content  # {'functionName': 'screenshot', 'args': {...}}

# Validate explicitly (optional; call() validates too)
assert tools.validate(tool_call)

# Execute
try:
    tools.call(tool_call)
except TransformError as e:
    print("Tool call failed:", e.message, e.raw)

If you already have a raw JSON string from a model and want to validate+parse:

parsed = tools.transform('{"functionName":"search","args":{"query":"platypus"}}')
# or for a list of calls:
calls = tools.transform_multi('[{"functionName":"search","args":{"query":"platypus"}}]')

5) Extras / environment defaults

Attach fixed kwargs (e.g., tokens, org IDs) that merge into every call:

tools = Tools(extras={"api_key": "SECRET"})  # extras override user args by default
tools.call(tool_call)

# Per-call control:
tools.call_with_extras({"api_key": "OTHER"}, tool_call, override=True)   # extras win
tools.call_with_extras({"api_key": "OTHER"}, tool_call, override=False)  # user args win

Notes

  • Return values are whatever your function returns—side effects are on you. Keep tools small and deterministic when possible.
  • tools.list() keeps the original type hints for backward compatibility and adds a normalized args schema that’s friendlier for prompts.
  • Safety: only register functions you actually want the model to invoke.

Embeddings

The embeddings module uses httpx and supports Ollama embeddings:

from trivialai.embedding import OllamaEmbedder
embed = OllamaEmbedder(model="nomic-embed-text", server="http://localhost:11434")
vec = embed("hello world")

Notes & compatibility

  • Dependencies: httpx replaces requests. Use httpx[http2] if you want HTTP/2 for OpenAI/Anthropic. Use boto3 for AWS Bedrock.
  • Python: ≥ 3.9 (we use asyncio.to_thread).
  • Scratchpad: only Ollama surfaces <think> content; others emit scratchpad as "" in deltas and None in the final event.
  • GCP/Vertex AI: primarily for setup/auth. No native provider streaming; astream falls back to a single final chunk unless you override.
  • Bedrock: stream(...) uses converse_stream(); token counts (when available) are surfaced as tokens in the final end event.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

trivialai-0.3.8.tar.gz (35.7 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

trivialai-0.3.8-py3-none-any.whl (28.8 kB view details)

Uploaded Python 3

File details

Details for the file trivialai-0.3.8.tar.gz.

File metadata

  • Download URL: trivialai-0.3.8.tar.gz
  • Upload date:
  • Size: 35.7 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.9.25

File hashes

Hashes for trivialai-0.3.8.tar.gz
Algorithm Hash digest
SHA256 f624139569e31798722b87df15a83f447494b15cdd471f4da1aca0b89a2a40c3
MD5 fcf05080a6a9475169bb11520166a357
BLAKE2b-256 d54dc9364a9351e11bd750da1db8aa478a900f25e762fa464eba09f84aedfba1

See more details on using hashes here.

File details

Details for the file trivialai-0.3.8-py3-none-any.whl.

File metadata

  • Download URL: trivialai-0.3.8-py3-none-any.whl
  • Upload date:
  • Size: 28.8 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.9.25

File hashes

Hashes for trivialai-0.3.8-py3-none-any.whl
Algorithm Hash digest
SHA256 e4adfa8be6d45d67927a91fb8dcda61af1f19bd2d843f1ddc53ee1d8e0bb8402
MD5 6712ee4046abb5327806bedc749052d7
BLAKE2b-256 e02283eb705a409e3cc2a5ac83c24d8ab580d15cd0e54f95aa82b8e34906e5e5

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page