Python AI SDK

3 months ago 3

A pure Python re-implementation of Vercel's popular AI SDK for TypeScript. Zero-configuration functions that work consistently across providers with first-class streaming, tool-calling, and structured output support.

Python is the defacto language for AI. However, to actually get started with AI, you'll need to 1. use a bloated external framework and install a bunch of dependencies, or 2. use an incredibly confusing API client (to simply call an LLM, you need client.chat.completions.create(**kwargs).result.choices[0].message.content).

Zero-configuration functions that work consistently across providers
First-class streaming & tool-calling support
Strong Pydantic types throughout - you know exactly what you're getting
Strict structured-output generation and streaming via Pydantic models
Provider-agnostic embeddings with built-in batching & retry logic
Tiny dependency footprint - no bloated external frameworks

Install via UV (Python package manager):

Or with pip:

pip install ai-sdk-python

That's it - no extra build steps or config files.

Get started in just a few lines of code.

from ai_sdk import generate_text, openai model = openai("gpt-4o-mini") res = generate_text(model=model, prompt="Tell me a haiku about Python") print(res.text)

import asyncio from ai_sdk import stream_text, openai async def main(): model = openai("gpt-4o-mini") stream_res = stream_text(model=model, prompt="Write a short story") async for chunk in stream_res.text_stream: print(chunk, end="", flush=True) asyncio.run(main())

from ai_sdk import generate_object, openai from pydantic import BaseModel class Person(BaseModel): name: str age: int model = openai("gpt-4o-mini") res = generate_object( model=model, schema=Person, prompt="Create a person named Alice, age 30" ) print(res.object) # Person(name='Alice', age=30)

from ai_sdk import embed_many, cosine_similarity, openai model = openai.embedding("text-embedding-3-small") texts = ["The cat sat on the mat.", "A dog was lying on the rug."] result = embed_many(model=model, values=texts) similarity = cosine_similarity(result.embeddings[0], result.embeddings[1]) print(f"Similarity: {similarity:.3f}")

from ai_sdk import tool, generate_text, openai from pydantic import BaseModel, Field # Using Pydantic models (recommended) class AddParams(BaseModel): a: float = Field(description="First number") b: float = Field(description="Second number") @tool( name="add", description="Add two numbers.", parameters=AddParams ) def add(a: float, b: float) -> float: return a + b model = openai("gpt-4o-mini") res = generate_text( model=model, prompt="What is 21 + 21?", tools=[add], ) print(res.text) # "The result is 42."

generate_text - Synchronous text generation with rich metadata
stream_text - Asynchronous streaming with real-time callbacks

generate_object - Structured output with Pydantic validation
stream_object - Streaming structured output with partial updates

embed - Single-value embedding helper
embed_many - Batch embedding with automatic batching
cosine_similarity - Semantic similarity calculation

tool - Define LLM-callable functions with Pydantic models or JSON schema

from ai_sdk import generate_text, openai from ai_sdk.types import CoreSystemMessage, CoreUserMessage, TextPart model = openai("gpt-4o-mini") messages = [ CoreSystemMessage(content="You are a helpful assistant."), CoreUserMessage(content=[TextPart(text="Respond with 'ack'.")]), ] res = generate_text(model=model, messages=messages) print(res.text)

Streaming Structured Output

import asyncio from ai_sdk import stream_object, openai from pydantic import BaseModel class Recipe(BaseModel): title: str ingredients: List[str] instructions: List[str] async def main(): model = openai("gpt-4o-mini") result = stream_object( model=model, schema=Recipe, prompt="Create a recipe for chocolate chip cookies" ) async for chunk in result.object_stream: print(chunk, end="", flush=True) recipe = await result.object() print(f"\n\nComplete recipe: {recipe}") asyncio.run(main())

from ai_sdk import embed_many, cosine_similarity, openai model = openai.embedding("text-embedding-3-small") # Knowledge base documents = [ "Python is a programming language.", "Machine learning involves training models on data.", "Databases store and retrieve information." ] # Search query query = "How do I learn to code?" # Embed everything all_texts = [query] + documents result = embed_many(model=model, values=all_texts) query_embedding = result.embeddings[0] doc_embeddings = result.embeddings[1:] # Find most similar document similarities = [] for i, doc_embedding in enumerate(doc_embeddings): sim = cosine_similarity(query_embedding, doc_embedding) similarities.append((sim, documents[i])) # Get top result top_result = max(similarities, key=lambda x: x[0]) print(f"Most relevant: {top_result[1]}")

from ai_sdk import tool, generate_text, openai import requests def get_weather(city: str) -> str: """Get current weather for a city.""" weather_data = { "New York": "72°F, Sunny", "London": "55°F, Rainy", "Tokyo": "68°F, Cloudy" } return weather_data.get(city, "Weather data not available") weather_tool = tool( name="get_weather", description="Get current weather information for a city.", parameters={ "type": "object", "properties": { "city": { "type": "string", "description": "The city name to get weather for" } }, "required": ["city"] }, execute=get_weather ) model = openai("gpt-4o-mini") res = generate_text( model=model, prompt="What's the weather like in New York?", tools=[weather_tool], ) print(res.text)

The SDK is provider-agnostic. Currently supported:

OpenAI - GPT models, embeddings, function calling
Anthropic - Claude models

from ai_sdk import generate_text, openai, anthropic # OpenAI openai_model = openai("gpt-4o-mini") res1 = generate_text(model=openai_model, prompt="Hello") # Anthropic anthropic_model = anthropic("claude-3-sonnet-20240229") res2 = generate_text(model=anthropic_model, prompt="Hello")

No complex setup - just import and use:

from ai_sdk import generate_text, openai res = generate_text(model=openai("gpt-4o-mini"), prompt="Hello!")

Swap providers without changing code:

# Works with any provider model = openai("gpt-4o-mini") # or anthropic("claude-3-sonnet-20240229") res = generate_text(model=model, prompt="Hello!")

Full Pydantic integration for type safety:

from pydantic import BaseModel from ai_sdk import generate_object class User(BaseModel): name: str age: int res = generate_object(model=model, schema=User, prompt="Create a user") user = res.object

Real-time text generation:

async for chunk in stream_text(model=model, prompt="Tell a story").text_stream: print(chunk, end="", flush=True)

5. Automatic Tool Calling

Define tools once, use everywhere:

add = tool(name="add", description="Add numbers", parameters={...}, execute=lambda x, y: x + y) res = generate_text(model=model, prompt="What's 2+2?", tools=[add])

Check out the examples directory for complete working examples: