vitaai 0.3.0

Python SDK for VitaAI self-hosted LLM endpoints — drop-in Gemini API replacement

vitaai · Python SDK

Drop-in Python SDK for VitaAI self-hosted LLM endpoints.
Designed to mirror the Google Gemini SDK interface exactly — so migrating existing Gemini projects takes changing one import.

---

Install

pip install vitaai

---

Quick start

from vitaai import VitaAI

client = VitaAI(
    api_key="your-api-key",
    base_url="http://your-server-ip:8000"
)

response = client.models.generate_content(
    model="gemma-2b-it",
    contents="What is AI?",
    max_output_tokens=100
)
print(response.text)
print("Tokens used:", response.usage_metadata.total_token_count)

---

Usage

Generate content

response = client.models.generate_content(
    model="gemma-2b-it",
    contents="What is the capital of France?",
)
print(response.text)
print(response.usage_metadata.total_token_count)

Control output length

response = client.models.generate_content(
    model="gemma-2b-it",
    contents="Explain machine learning",
    max_output_tokens=200
)
print(response.text)

Generation config

response = client.models.generate_content(
    model="gemma-2b-it",
    contents="Write a product description",
    generation_config={
        "temperature": 0.8,
        "max_output_tokens": 256,
        "top_p": 0.95,
    },
)
print(response.text)

System instruction

response = client.models.generate_content(
    model="gemma-2b-it",
    contents="Tell me a joke.",
    system_instruction="You are a dry, deadpan comedian. Keep it under 2 sentences.",
)
print(response.text)

Multi-turn conversation

response = client.models.generate_content(
    model="gemma-2b-it",
    contents=[
        {"role": "user",  "parts": [{"text": "My name is Alex."}]},
        {"role": "model", "parts": [{"text": "Nice to meet you, Alex!"}]},
        {"role": "user",  "parts": [{"text": "What's my name?"}]},
    ],
)
print(response.text)

Streaming

for chunk in client.models.generate_content_stream(
    model="gemma-2b-it",
    contents="Write a short story",
    max_output_tokens=300
):
    print(chunk.text, end="", flush=True)
print()

Embeddings

response = client.models.embed_content(
    contents="The quick brown fox jumps over the lazy dog"
)
print(len(response.embedding.values))   # number of dimensions e.g. 384
print(response.embedding.values)        # list of floats

List available models

for model in client.models.list():
    print(model["name"], model.get("displayName", ""))

---

Error handling

from vitaai.transport import AuthenticationError, RateLimitError, ServerError

try:
    response = client.models.generate_content(
        model="gemma-2b-it",
        contents="Hello"
    )
    print(response.text)

except AuthenticationError:
    print("Invalid API key")
except RateLimitError:
    print("Too many requests")
except ServerError:
    print("Server error")
except Exception as e:
    print(f"Error: {e}")

---

Migration from Google Gemini

Before (google.genai)	After (vitaai)
from google import genai	from vitaai import VitaAI
client = genai.Client()	client = VitaAI()
client.models.generate_content(...)	client.models.generate_content(...)
response.text	response.text ✅ same
response.usage_metadata	response.usage_metadata ✅ same

---

Endpoint contract

Your self-hosted backend must expose these routes:

Method	Path	Purpose
POST	/v1/models/generate	Text generation
POST	/v1/models/generate?stream=true	Streaming
POST	/v1/models/embed	Embeddings
GET	/v1/models/list	List models

---

Environment variables

Variable	Description
VITAAI_API_KEY	API key (alternative to passing api_key=)
VITAAI_BASE_URL	Override base URL