feat: add batch job tracking and generate batch APIs (#1727)

Co-authored-by: Matt Zhou <mattzh1314@gmail.com>
This commit is contained in:
Sarah Wooders
2025-04-17 17:02:07 -07:00
committed by GitHub
parent cd94962882
commit 41d85c4114
22 changed files with 690 additions and 262 deletions

View File

@@ -102,7 +102,7 @@ async def poll_batch_updates(server: SyncServer, batch_jobs: List[LLMBatchJob],
results: List[BatchPollingResult] = await asyncio.gather(*coros)
# Update the server with batch status changes
server.batch_manager.bulk_update_batch_statuses(updates=results)
server.batch_manager.bulk_update_llm_batch_statuses(updates=results)
logger.info(f"[Poll BatchJob] Bulk-updated {len(results)} LLM batch(es) in the DB at job level.")
return results
@@ -176,7 +176,7 @@ async def poll_running_llm_batches(server: "SyncServer") -> None:
try:
# 1. Retrieve running batch jobs
batches = server.batch_manager.list_running_batches()
batches = server.batch_manager.list_running_llm_batches()
metrics.total_batches = len(batches)
# TODO: Expand to more providers
@@ -193,7 +193,7 @@ async def poll_running_llm_batches(server: "SyncServer") -> None:
# 6. Bulk update all items for newly completed batch(es)
if item_updates:
metrics.updated_items_count = len(item_updates)
server.batch_manager.bulk_update_batch_items_results_by_agent(item_updates)
server.batch_manager.bulk_update_batch_llm_items_results_by_agent(item_updates)
else:
logger.info("[Poll BatchJob] No item-level updates needed.")

View File

@@ -6,25 +6,25 @@ from letta.schemas.enums import AgentStepStatus, JobStatus
class BatchPollingResult(NamedTuple):
batch_id: str
llm_batch_id: str
request_status: JobStatus
batch_response: Optional[BetaMessageBatch]
class ItemUpdateInfo(NamedTuple):
batch_id: str
llm_batch_id: str
agent_id: str
request_status: JobStatus
batch_request_result: Optional[BetaMessageBatchIndividualResponse]
class StepStatusUpdateInfo(NamedTuple):
batch_id: str
llm_batch_id: str
agent_id: str
step_status: AgentStepStatus
class RequestStatusUpdateInfo(NamedTuple):
batch_id: str
llm_batch_id: str
agent_id: str
request_status: JobStatus