Skip to content

Add support for latest models #39

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
agamm merged 5 commits into from
Jan 24, 2026
Merged

Add support for latest models #39

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
9c8bb4b
Add support for latest models
agamm Jan 23, 2026
7e0c7fd
Add comprehensive GPT-5 model family support
agamm Jan 24, 2026
2ef4c4e
Fix syntax error in Gemini test file
agamm Jan 24, 2026
49f2bc4
Complete GPT-5 model family integration and syntax fixes
agamm Jan 24, 2026
8b1a15b
Fix syntax error in test_batch_results_with_real_google_format
agamm Jan 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ from batchata import Batch

# Simple batch processing
batch = Batch(results_dir="./output")
.set_default_params(model="claude-sonnet-4-20250514") # or "gpt-4.1-2025-04-14" or "gemini-2.5-flash"
.set_default_params(model="gpt-5.2-latest") # or "claude-sonnet-4-5-20250929" or "gemini-3.0-pro-latest"
.add_cost_limit(usd=5.0)

for file in files:
Expand Down Expand Up @@ -79,7 +79,7 @@ batch = Batch(
items_per_batch=3
)
.set_state(file="./invoice_state.json", reuse_state=False)
.set_default_params(model="claude-sonnet-4-20250514", temperature=0.0)
.set_default_params(model="gpt-5.2-pro-2025-12-11", temperature=0.0)
.add_cost_limit(usd=5.0)
.add_time_limit(minutes=10) # Time limit of 10 minutes
.set_verbosity("warn")
Expand Down
4 changes: 2 additions & 2 deletions batchata/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@

# Simple batch processing
batch = Batch(results_dir="./output")
.set_default_params(model="claude-sonnet-4-20250514")
.set_default_params(model="claude-sonnet-4-5-20250929")
.add_cost_limit(usd=5.0)

# Add jobs
Expand All @@ -45,7 +45,7 @@ class DocumentAnalysis(BaseModel):
key_points: list[str]

batch = Batch(results_dir="./results")
.set_default_params(model="claude-sonnet-4-20250514")
.set_default_params(model="claude-sonnet-4-5-20250929")

batch.add_job(
file="document.pdf",
Expand Down
2 changes: 1 addition & 1 deletion batchata/core/batch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ class Batch:
```python
batch = Batch("./results", max_parallel_batches=10, items_per_batch=10)
.set_state(file="./state.json", reuse_state=True)
.set_default_params(model="claude-sonnet-4-20250514", temperature=0.7)
.set_default_params(model="claude-sonnet-4-5-20250929", temperature=0.7)
.add_cost_limit(usd=15.0)
.add_job(messages=[{"role": "user", "content": "Hello"}])
.add_job(file="./path/to/file.pdf", prompt="Generate summary of file")
Expand Down
22 changes: 22 additions & 0 deletions batchata/providers/anthropic/models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,28 @@
max_output_tokens=4096,
batch_discount=0.5,
supports_images=True,
supports_files=False,
supports_citations=True,
supports_structured_output=True,
file_types=[]
),
"claude-opus-4-5-20250929": ModelConfig(
name="claude-opus-4-5-20250929",
max_input_tokens=200000,
max_output_tokens=8192,
batch_discount=0.5,
supports_images=True,
supports_files=True,
supports_citations=True,
supports_structured_output=True,
file_types=[".pdf", ".txt", ".docx", ".jpg", ".png", ".gif", ".webp"]
),
"claude-sonnet-4-5-20250929": ModelConfig(
name="claude-sonnet-4-5-20250929",
max_input_tokens=200000,
max_output_tokens=8192,
batch_discount=0.5,
supports_images=True,
supports_files=True,
supports_citations=True,
supports_structured_output=True,
Expand Down
18 changes: 11 additions & 7 deletions batchata/providers/anthropic/parse_results.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -310,10 +310,14 @@ def _save_raw_response(result: Any, job_id: str, raw_files_dir: str) -> None:

def _calculate_cost(input_tokens: int, output_tokens: int, model: str, batch_discount: float = 0.5) -> float:
"""Calculate cost for tokens using tokencost."""
from tokencost import calculate_cost_by_tokens

# Calculate costs using tokencost
input_cost = float(calculate_cost_by_tokens(input_tokens, model, token_type="input"))
output_cost = float(calculate_cost_by_tokens(output_tokens, model, token_type="output"))

return (input_cost + output_cost) * batch_discount
try:
from tokencost import calculate_cost_by_tokens

# Calculate costs using tokencost
input_cost = float(calculate_cost_by_tokens(input_tokens, model, token_type="input"))
output_cost = float(calculate_cost_by_tokens(output_tokens, model, token_type="output"))

return (input_cost + output_cost) * batch_discount
except Exception as e:
logger.warning(f"Failed to calculate cost for model {model} using tokencost: {e}. Returning 0 cost.")
return 0.0
45 changes: 39 additions & 6 deletions batchata/providers/gemini/models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@
# Google Gemini models with batch processing support
# Batch mode provides 50% discount on standard API pricing
GEMINI_MODELS = {
"gemini-2.5-pro": ModelConfig(
name="gemini-2.5-pro",
"gemini-3.0-pro-latest": ModelConfig(
name="gemini-3.0-pro-latest",
max_input_tokens=2097152, # 2M context
max_output_tokens=8192,
batch_discount=0.5, # 50% discount confirmed in docs
Expand All @@ -17,8 +17,41 @@
supports_structured_output=True,
file_types=[".pdf", ".txt", ".jpg", ".png", ".gif", ".webp"]
),
"gemini-2.5-flash": ModelConfig(
name="gemini-2.5-flash",
"gemini-3.0-pro": ModelConfig(
name="gemini-3.0-pro",
max_input_tokens=2097152, # 2M context
max_output_tokens=8192,
batch_discount=0.5, # 50% discount confirmed in docs
supports_images=True,
supports_files=True,
supports_citations=False,
supports_structured_output=True,
file_types=[".pdf", ".txt", ".jpg", ".png", ".gif", ".webp"]
),
"gemini-3.0-flash-latest": ModelConfig(
name="gemini-3.0-flash-latest",
max_input_tokens=1048576, # 1M context
max_output_tokens=8192,
batch_discount=0.5,
supports_images=True,
supports_files=True,
supports_citations=False,
supports_structured_output=True,
file_types=[".pdf", ".txt", ".jpg", ".png", ".gif", ".webp"]
),
"gemini-3.0-flash": ModelConfig(
name="gemini-3.0-flash",
max_input_tokens=1048576, # 1M context
max_output_tokens=8192,
batch_discount=0.5,
supports_images=True,
supports_files=True,
supports_citations=False,
supports_structured_output=True,
file_types=[".pdf", ".txt", ".jpg", ".png", ".gif", ".webp"]
),
"gemini-3.0-flash-lite-latest": ModelConfig(
name="gemini-3.0-flash-lite-latest",
max_input_tokens=1048576, # 1M context
max_output_tokens=8192,
batch_discount=0.5,
Expand All @@ -28,8 +61,8 @@
supports_structured_output=True,
file_types=[".pdf", ".txt", ".jpg", ".png", ".gif", ".webp"]
),
"gemini-2.5-flash-lite": ModelConfig(
name="gemini-2.5-flash-lite",
"gemini-3.0-flash-lite": ModelConfig(
name="gemini-3.0-flash-lite",
max_input_tokens=1048576, # 1M context
max_output_tokens=8192,
batch_discount=0.5,
Expand Down
9 changes: 6 additions & 3 deletions batchata/providers/gemini/parse_results.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,10 @@
from typing import Dict, List, Optional

from ...core.job_result import JobResult
from ...utils import to_dict
from ...utils import to_dict, get_logger


logger = get_logger(__name__)


def parse_results(results: List[Dict], job_mapping: Dict[str, 'Job'], raw_files_dir: Optional[str] = None, batch_discount: float = 0.5, batch_id: Optional[str] = None) -> List[JobResult]:
Expand Down Expand Up @@ -129,8 +132,8 @@ def _calculate_cost(model: str, input_tokens: int, output_tokens: int, batch_dis
total_cost = float(input_cost + output_cost)
return total_cost * (1 - batch_discount)

except (ImportError, ModuleNotFoundError, AttributeError, ValueError):
# Return zero cost if tokencost library unavailable or calculation fails
except Exception as e:
logger.warning(f"Failed to calculate cost for model {model} using tokencost: {e}. Returning 0 cost.")
return 0.0


Expand Down
4 changes: 2 additions & 2 deletions batchata/providers/model_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ class ModelConfig:
"""Configuration for a specific model.

Attributes:
name: Model identifier (e.g., "claude-sonnet-4-20250514")
name: Model identifier (e.g., "claude-sonnet-4-5-20250929")
max_input_tokens: Maximum input context length
max_output_tokens: Maximum tokens that can be generated
batch_discount: Discount factor for batch processing (e.g., 0.5 for 50% off)
Expand All @@ -20,7 +20,7 @@ class ModelConfig:
file_types: List of supported file extensions
"""

name: str # e.g., "claude-sonnet-4-20250514"
name: str # e.g., "claude-sonnet-4-5-20250929"
max_input_tokens: int
max_output_tokens: int
batch_discount: float # e.g., 0.5 for 50% off
Expand Down
Loading
Loading