wehub-resource-sync
97e91a83f3
Ruff / Ruff (push) Has been cancelled
Test / Core Tests (push) Has been cancelled
Test / Offline Coverage Tests (Python 3.10) (push) Has been cancelled
Test / Offline Coverage Tests (Python 3.11) (push) Has been cancelled
Test / Offline Coverage Tests (Python 3.12) (push) Has been cancelled
Test / Offline Coverage Tests (Python 3.13) (push) Has been cancelled
Test / Offline Coverage Tests (Python 3.9) (push) Has been cancelled
Test / Full Coverage (Python 3.11) (push) Has been cancelled
Test / Core Provider Tests (OpenAI) (push) Has been cancelled
Test / Core Provider Tests (Anthropic) (push) Has been cancelled
Test / Core Provider Tests (Google) (push) Has been cancelled
Test / Core Provider Tests (Other) (push) Has been cancelled
Test / Anthropic Tests (push) Has been cancelled
Test / Gemini Tests (push) Has been cancelled
Test / Google GenAI Tests (push) Has been cancelled
Test / Vertex AI Tests (push) Has been cancelled
Test / OpenAI Tests (push) Has been cancelled
Test / Writer Tests (push) Has been cancelled
Test / Auto Client Tests (push) Has been cancelled
ty / type-check (push) Has been cancelled
561 lines
19 KiB
Python
561 lines
19 KiB
Python
import os
|
|
from rich.console import Console
|
|
from rich.table import Table
|
|
from rich.live import Live
|
|
import typer
|
|
import time
|
|
import json
|
|
import warnings
|
|
from instructor.batch import BatchProcessor, BatchJobInfo
|
|
|
|
from tqdm import tqdm
|
|
|
|
app = typer.Typer()
|
|
|
|
console = Console()
|
|
|
|
|
|
def generate_table(
|
|
batch_jobs: list[BatchJobInfo], provider: str, full_id: bool = False
|
|
):
|
|
"""Generate enhanced table for batch jobs using unified BatchJobInfo objects
|
|
|
|
Args:
|
|
batch_jobs: List of batch job info objects
|
|
provider: Provider name (openai, anthropic)
|
|
full_id: If True, show full batch IDs without truncation
|
|
"""
|
|
table = Table(title=f"{provider.title()} Batch Jobs")
|
|
|
|
# Adjust column width based on full_id flag
|
|
id_max_width = None if full_id else 20
|
|
table.add_column("Batch ID", style="dim", max_width=id_max_width, no_wrap=True)
|
|
table.add_column("Status", min_width=10)
|
|
table.add_column("Created", style="dim", min_width=10)
|
|
table.add_column("Started", style="dim", min_width=10)
|
|
table.add_column("Duration", style="dim", min_width=7)
|
|
|
|
# Add provider-specific columns for request counts
|
|
if provider == "openai":
|
|
table.add_column("Completed", justify="right", min_width=8)
|
|
table.add_column("Failed", justify="right", min_width=6)
|
|
table.add_column("Total", justify="right", min_width=6)
|
|
elif provider == "anthropic":
|
|
table.add_column("Succeeded", justify="right", min_width=8)
|
|
table.add_column("Errored", justify="right", min_width=7)
|
|
table.add_column("Processing", justify="right", min_width=9)
|
|
|
|
for batch_job in batch_jobs:
|
|
# Color code status
|
|
status_color = {
|
|
"pending": "yellow",
|
|
"processing": "blue",
|
|
"completed": "green",
|
|
"failed": "red",
|
|
"cancelled": "red",
|
|
"expired": "red",
|
|
}.get(batch_job.status.value, "white")
|
|
|
|
colored_status = f"[{status_color}]{batch_job.status.value}[/{status_color}]"
|
|
|
|
# Format timestamps
|
|
created_str = (
|
|
batch_job.timestamps.created_at.strftime("%m/%d %H:%M")
|
|
if batch_job.timestamps.created_at
|
|
else "N/A"
|
|
)
|
|
started_str = (
|
|
batch_job.timestamps.started_at.strftime("%m/%d %H:%M")
|
|
if batch_job.timestamps.started_at
|
|
else "N/A"
|
|
)
|
|
|
|
# Calculate duration
|
|
duration_str = "N/A"
|
|
if batch_job.timestamps.started_at and batch_job.timestamps.completed_at:
|
|
duration = (
|
|
batch_job.timestamps.completed_at - batch_job.timestamps.started_at
|
|
)
|
|
total_minutes = duration.total_seconds() / 60
|
|
if total_minutes < 60:
|
|
duration_str = f"{int(total_minutes)}m"
|
|
else:
|
|
hours = total_minutes / 60
|
|
duration_str = f"{hours:.1f}h"
|
|
elif batch_job.timestamps.started_at and batch_job.status.value == "processing":
|
|
from datetime import datetime, timezone
|
|
|
|
duration = datetime.now(timezone.utc) - batch_job.timestamps.started_at
|
|
total_minutes = duration.total_seconds() / 60
|
|
if total_minutes < 60:
|
|
duration_str = f"{int(total_minutes)}m"
|
|
else:
|
|
hours = total_minutes / 60
|
|
duration_str = f"{hours:.1f}h"
|
|
|
|
# Truncate batch ID for display only if full_id is False
|
|
batch_id_display = str(batch_job.id)
|
|
if not full_id and len(batch_id_display) > 18:
|
|
batch_id_display = batch_id_display[:15] + "..."
|
|
|
|
if provider == "openai":
|
|
table.add_row(
|
|
batch_id_display,
|
|
colored_status,
|
|
created_str,
|
|
started_str,
|
|
duration_str,
|
|
str(batch_job.request_counts.completed or 0),
|
|
str(batch_job.request_counts.failed or 0),
|
|
str(batch_job.request_counts.total or 0),
|
|
)
|
|
elif provider == "anthropic":
|
|
table.add_row(
|
|
str(batch_job.id),
|
|
colored_status,
|
|
created_str,
|
|
started_str,
|
|
duration_str,
|
|
str(batch_job.request_counts.succeeded or 0),
|
|
str(batch_job.request_counts.errored or 0),
|
|
str(batch_job.request_counts.processing or 0),
|
|
)
|
|
|
|
return table
|
|
|
|
|
|
def get_jobs(limit: int = 10, provider: str = "openai") -> list[BatchJobInfo]:
|
|
"""Get batch jobs for the specified provider using BatchProcessor"""
|
|
|
|
# Create a dummy model string for the provider
|
|
# We just need the provider part for listing batches
|
|
model_map = {
|
|
"openai": "openai/gpt-4o-mini",
|
|
"anthropic": "anthropic/claude-3-sonnet",
|
|
}
|
|
|
|
if provider not in model_map:
|
|
raise ValueError(f"Unsupported provider: {provider}")
|
|
|
|
# Create a dummy response model (not used for listing)
|
|
from pydantic import BaseModel
|
|
|
|
class DummyModel(BaseModel):
|
|
dummy: str = "dummy"
|
|
|
|
try:
|
|
# Create BatchProcessor instance
|
|
processor = BatchProcessor(model_map[provider], DummyModel)
|
|
# Get batch jobs
|
|
return processor.list_batches(limit=limit)
|
|
except Exception as e:
|
|
console.print(f"[red]Error listing {provider} batch jobs: {e}[/red]")
|
|
return []
|
|
|
|
|
|
@app.command(name="list", help="See all existing batch jobs")
|
|
def watch(
|
|
limit: int = typer.Option(10, help="Total number of batch jobs to show"),
|
|
poll: int = typer.Option(
|
|
10, help="Time in seconds to wait for the batch job to complete"
|
|
),
|
|
screen: bool = typer.Option(False, help="Enable or disable screen output"),
|
|
live: bool = typer.Option(
|
|
False, help="Enable live polling to continuously update the table"
|
|
),
|
|
provider: str = typer.Option(
|
|
"openai",
|
|
help="Provider to use (e.g., 'openai', 'anthropic')",
|
|
),
|
|
# Deprecated flag for backward compatibility
|
|
use_anthropic: bool = typer.Option(
|
|
None,
|
|
help="[DEPRECATED] Use --model instead. Use Anthropic API instead of OpenAI",
|
|
),
|
|
full_id: bool = typer.Option(
|
|
False,
|
|
"--full-id",
|
|
help="Show full batch IDs without truncation",
|
|
),
|
|
):
|
|
"""
|
|
Monitor the status of the most recent batch jobs
|
|
"""
|
|
# Handle deprecated flag
|
|
if use_anthropic is not None:
|
|
warnings.warn(
|
|
"--use-anthropic is deprecated. Use --provider 'anthropic' instead.",
|
|
DeprecationWarning,
|
|
stacklevel=2,
|
|
)
|
|
if use_anthropic:
|
|
provider = "anthropic"
|
|
|
|
# Check if required API key is available for the provider
|
|
required_keys = {
|
|
"anthropic": "ANTHROPIC_API_KEY",
|
|
"openai": "OPENAI_API_KEY",
|
|
}
|
|
|
|
if provider in required_keys and not os.getenv(required_keys[provider]):
|
|
console.print(
|
|
f"[red]Error: {required_keys[provider]} environment variable not set for {provider}[/red]"
|
|
)
|
|
return
|
|
|
|
batch_jobs = get_jobs(limit, provider)
|
|
table = generate_table(batch_jobs, provider, full_id=full_id)
|
|
|
|
if not live:
|
|
# Show table once and exit
|
|
console.print(table)
|
|
return
|
|
|
|
# Live polling mode
|
|
with Live(table, refresh_per_second=2, screen=screen) as live_table:
|
|
while True:
|
|
batch_jobs = get_jobs(limit, provider)
|
|
table = generate_table(batch_jobs, provider, full_id=full_id)
|
|
live_table.update(table)
|
|
time.sleep(poll)
|
|
|
|
|
|
@app.command(
|
|
help="Create a batch job from a file",
|
|
)
|
|
def create_from_file(
|
|
file_path: str = typer.Option(help="File containing the batch job requests"),
|
|
model: str = typer.Option(
|
|
"openai/gpt-4o-mini",
|
|
help="Model in format 'provider/model-name' (e.g., 'openai/gpt-4', 'anthropic/claude-3-sonnet')",
|
|
),
|
|
description: str = typer.Option(
|
|
"Instructor batch job",
|
|
help="Description/metadata for the batch job",
|
|
),
|
|
completion_window: str = typer.Option(
|
|
"24h",
|
|
help="Completion window for the batch job (OpenAI only)",
|
|
),
|
|
# Deprecated flag for backward compatibility
|
|
use_anthropic: bool = typer.Option(
|
|
None,
|
|
help="[DEPRECATED] Use --model instead. Use Anthropic API instead of OpenAI",
|
|
),
|
|
):
|
|
"""Create a batch job from a file using the unified BatchProcessor"""
|
|
# Handle deprecated flag
|
|
if use_anthropic is not None:
|
|
warnings.warn(
|
|
"--use-anthropic is deprecated. Use --model 'anthropic/claude-3-sonnet' instead.",
|
|
DeprecationWarning,
|
|
stacklevel=2,
|
|
)
|
|
if use_anthropic:
|
|
model = "anthropic/claude-3-sonnet"
|
|
|
|
try:
|
|
# Create a dummy response model (not used for direct file submission)
|
|
from pydantic import BaseModel
|
|
|
|
class DummyModel(BaseModel):
|
|
dummy: str = "dummy"
|
|
|
|
# Create BatchProcessor instance
|
|
processor = BatchProcessor(model, DummyModel)
|
|
|
|
# Prepare metadata
|
|
metadata = {
|
|
"description": description,
|
|
}
|
|
|
|
with console.status(f"[bold green]Submitting batch job...", spinner="dots"):
|
|
batch_id = processor.submit_batch(
|
|
file_path, metadata=metadata, completion_window=completion_window
|
|
)
|
|
|
|
console.print(f"[bold green]Batch job created with ID: {batch_id}[/bold green]")
|
|
|
|
# Show updated batch list
|
|
provider_name = model.split("/", 1)[0]
|
|
watch(limit=5, poll=2, screen=False, live=False, provider=provider_name)
|
|
|
|
except Exception as e:
|
|
console.print(f"[bold red]Error creating batch job: {e}[/bold red]")
|
|
|
|
|
|
@app.command(help="Cancel a batch job")
|
|
def cancel(
|
|
batch_id: str = typer.Option(help="Batch job ID to cancel"),
|
|
provider: str = typer.Option(
|
|
"openai",
|
|
help="Provider to use (e.g., 'openai', 'anthropic')",
|
|
),
|
|
# Deprecated flag for backward compatibility
|
|
use_anthropic: bool = typer.Option(
|
|
None,
|
|
help="[DEPRECATED] Use --provider 'anthropic' instead. Use Anthropic API instead of OpenAI",
|
|
),
|
|
):
|
|
"""Cancel a batch job using the unified BatchProcessor"""
|
|
# Handle deprecated flag
|
|
if use_anthropic is not None:
|
|
warnings.warn(
|
|
"--use-anthropic is deprecated. Use --provider 'anthropic' instead.",
|
|
DeprecationWarning,
|
|
stacklevel=2,
|
|
)
|
|
if use_anthropic:
|
|
provider = "anthropic"
|
|
|
|
try:
|
|
# Create a dummy response model (not used for cancellation)
|
|
from pydantic import BaseModel
|
|
|
|
class DummyModel(BaseModel):
|
|
dummy: str = "dummy"
|
|
|
|
# Create a dummy model string for the provider
|
|
model_map = {
|
|
"openai": "openai/gpt-4o-mini",
|
|
"anthropic": "anthropic/claude-3-sonnet",
|
|
}
|
|
|
|
if provider not in model_map:
|
|
console.print(f"[red]Unsupported provider: {provider}[/red]")
|
|
return
|
|
|
|
# Create BatchProcessor instance
|
|
processor = BatchProcessor(model_map[provider], DummyModel)
|
|
|
|
with console.status(
|
|
f"[bold yellow]Cancelling {provider} batch job...", spinner="dots"
|
|
):
|
|
processor.cancel_batch(batch_id)
|
|
|
|
console.print(
|
|
f"[bold green]Batch {batch_id} cancelled successfully![/bold green]"
|
|
)
|
|
|
|
# Show updated status
|
|
watch(limit=5, poll=2, screen=False, live=False, provider=provider)
|
|
|
|
except NotImplementedError as e:
|
|
console.print(f"[yellow]Note: {e}[/yellow]")
|
|
except Exception as e:
|
|
console.print(f"[bold red]Error cancelling batch {batch_id}: {e}[/bold red]")
|
|
|
|
|
|
@app.command(help="Delete a completed batch job")
|
|
def delete(
|
|
batch_id: str = typer.Option(help="Batch job ID to delete"),
|
|
provider: str = typer.Option(
|
|
"openai",
|
|
help="Provider to use (e.g., 'openai', 'anthropic')",
|
|
),
|
|
):
|
|
"""Delete a batch job using the unified BatchProcessor"""
|
|
try:
|
|
# Create a dummy response model (not used for deletion)
|
|
from pydantic import BaseModel
|
|
|
|
class DummyModel(BaseModel):
|
|
dummy: str = "dummy"
|
|
|
|
# Create a dummy model string for the provider
|
|
model_map = {
|
|
"openai": "openai/gpt-4o-mini",
|
|
"anthropic": "anthropic/claude-3-sonnet",
|
|
}
|
|
|
|
if provider not in model_map:
|
|
console.print(f"[red]Unsupported provider: {provider}[/red]")
|
|
return
|
|
|
|
# Create BatchProcessor instance
|
|
processor = BatchProcessor(model_map[provider], DummyModel)
|
|
|
|
with console.status(
|
|
f"[bold yellow]Deleting {provider} batch job...", spinner="dots"
|
|
):
|
|
processor.delete_batch(batch_id)
|
|
|
|
console.print(
|
|
f"[bold green]Batch {batch_id} deleted successfully![/bold green]"
|
|
)
|
|
|
|
# Show updated status
|
|
watch(limit=5, poll=2, screen=False, live=False, provider=provider)
|
|
|
|
except NotImplementedError as e:
|
|
console.print(f"[yellow]Note: {e}[/yellow]")
|
|
except Exception as e:
|
|
console.print(f"[bold red]Error deleting batch {batch_id}: {e}[/bold red]")
|
|
|
|
|
|
@app.command(help="Download the file associated with a batch job")
|
|
def download_file(
|
|
batch_id: str = typer.Option(help="Batch job ID to download"),
|
|
download_file_path: str = typer.Option(help="Path to download file to"),
|
|
provider: str = typer.Option(
|
|
"openai",
|
|
help="Provider to use (e.g., 'openai', 'anthropic')",
|
|
),
|
|
):
|
|
try:
|
|
if provider == "anthropic":
|
|
from anthropic import Anthropic
|
|
|
|
client = Anthropic()
|
|
# TODO: Remove beta fallback when stable API is available
|
|
try:
|
|
batches_client = client.messages.batches
|
|
except AttributeError:
|
|
batches_client = client.beta.messages.batches
|
|
batch = batches_client.retrieve(batch_id)
|
|
if batch.processing_status != "ended":
|
|
raise ValueError("Only completed Jobs can be downloaded")
|
|
|
|
results_url = batch.results_url
|
|
if not results_url:
|
|
raise ValueError("Results URL not available")
|
|
|
|
with open(download_file_path, "w") as file:
|
|
for result in tqdm(client.messages.batches.results(batch_id)):
|
|
file.write(json.dumps(result.model_dump()) + "\n")
|
|
else:
|
|
from openai import OpenAI
|
|
|
|
client = OpenAI()
|
|
batch = client.batches.retrieve(batch_id=batch_id)
|
|
status = batch.status
|
|
|
|
if status != "completed":
|
|
raise ValueError("Only completed Jobs can be downloaded")
|
|
|
|
file_id = batch.output_file_id
|
|
|
|
assert file_id, f"Equivalent Output File not found for {batch_id}"
|
|
file_response = client.files.content(file_id)
|
|
|
|
with open(download_file_path, "w") as file:
|
|
file.write(file_response.text)
|
|
|
|
except Exception as e:
|
|
console.log(f"[bold red]Error downloading file for {batch_id}: {e}")
|
|
|
|
|
|
@app.command(help="Retrieve results from a batch job")
|
|
def results(
|
|
batch_id: str = typer.Option(help="Batch job ID to get results from"),
|
|
output_file: str = typer.Option(help="File to save the results to"),
|
|
model: str = typer.Option(
|
|
"openai/gpt-4o-mini",
|
|
help="Model in format 'provider/model-name' (e.g., 'openai/gpt-4', 'anthropic/claude-3-sonnet')",
|
|
),
|
|
):
|
|
"""Retrieve and save batch job results"""
|
|
provider, _ = model.split("/", 1)
|
|
|
|
try:
|
|
if provider == "openai":
|
|
from openai import OpenAI
|
|
|
|
client = OpenAI()
|
|
batch = client.batches.retrieve(batch_id=batch_id)
|
|
|
|
if batch.status != "completed":
|
|
console.print(
|
|
f"[yellow]Batch status is '{batch.status}', not completed[/yellow]"
|
|
)
|
|
return
|
|
|
|
file_id = batch.output_file_id
|
|
if not file_id:
|
|
console.print("[red]No output file available[/red]")
|
|
return
|
|
|
|
file_response = client.files.content(file_id)
|
|
with open(output_file, "w") as f:
|
|
f.write(file_response.text)
|
|
console.print(f"[bold green]Results saved to: {output_file}[/bold green]")
|
|
|
|
elif provider == "anthropic":
|
|
from anthropic import Anthropic
|
|
|
|
client = Anthropic()
|
|
batch = client.beta.messages.batches.retrieve(batch_id)
|
|
|
|
if batch.processing_status != "ended":
|
|
console.print(
|
|
f"[yellow]Batch status is '{batch.processing_status}', not ended[/yellow]"
|
|
)
|
|
return
|
|
|
|
# Get results from Anthropic batch API
|
|
results_iter = client.beta.messages.batches.results(batch_id)
|
|
|
|
with open(output_file, "w") as f:
|
|
for result in results_iter:
|
|
f.write(json.dumps(result.model_dump()) + "\n")
|
|
console.print(f"[bold green]Results saved to: {output_file}[/bold green]")
|
|
|
|
else:
|
|
console.print(f"[red]Unsupported provider: {provider}[/red]")
|
|
|
|
except Exception as e:
|
|
console.log(f"[bold red]Error retrieving results for {batch_id}: {e}")
|
|
|
|
|
|
@app.command(help="Create batch job using BatchProcessor")
|
|
def create(
|
|
messages_file: str = typer.Option(help="JSONL file with message conversations"),
|
|
model: str = typer.Option(
|
|
"openai/gpt-4o-mini",
|
|
help="Model in format 'provider/model-name' (e.g., 'openai/gpt-4', 'anthropic/claude-3-sonnet')",
|
|
),
|
|
response_model: str = typer.Option(
|
|
help="Python class path for response model (e.g., 'examples.User')"
|
|
),
|
|
output_file: str = typer.Option(
|
|
"batch_requests.jsonl", help="Output file for batch requests"
|
|
),
|
|
max_tokens: int = typer.Option(1000, help="Maximum tokens per request"),
|
|
temperature: float = typer.Option(0.1, help="Temperature for generation"),
|
|
):
|
|
"""Create a batch job using the unified BatchProcessor"""
|
|
try:
|
|
# Import the response model dynamically
|
|
module_path, class_name = response_model.rsplit(".", 1)
|
|
import importlib
|
|
|
|
module = importlib.import_module(module_path)
|
|
response_class = getattr(module, class_name)
|
|
|
|
# Load messages from file
|
|
messages_list = []
|
|
with open(messages_file) as f:
|
|
for line in f:
|
|
if line.strip():
|
|
messages_list.append(json.loads(line))
|
|
|
|
# Create batch processor
|
|
processor = BatchProcessor(model, response_class)
|
|
|
|
# Create batch file
|
|
with console.status(
|
|
f"[bold green]Creating batch file with {len(messages_list)} requests...",
|
|
spinner="dots",
|
|
):
|
|
processor.create_batch_from_messages(
|
|
messages_list, output_file, max_tokens, temperature
|
|
)
|
|
|
|
console.print(f"[bold green]Batch file created: {output_file}[/bold green]")
|
|
console.print(
|
|
f"[yellow]Use 'instructor batch create-from-file --file-path {output_file}' to submit the batch[/yellow]"
|
|
)
|
|
|
|
except Exception as e:
|
|
console.log(f"[bold red]Error creating batch: {e}")
|