f5c4ab50f4
* auto fixes * auto fix pt2 and transitive deps and undefined var checking locals() * manual fixes (ignored or letta-code fixed) * fix circular import * remove all ignores, add FastAPI rules and Ruff rules * add ty and precommit * ruff stuff * ty check fixes * ty check fixes pt 2 * error on invalid
56 lines
2.5 KiB
Python
56 lines
2.5 KiB
Python
import uuid
|
|
from datetime import datetime
|
|
from typing import TYPE_CHECKING, List, Optional, Union
|
|
|
|
if TYPE_CHECKING:
|
|
from letta.orm.llm_batch_items import LLMBatchItem
|
|
from letta.orm.organization import Organization
|
|
|
|
from anthropic.types.beta.messages import BetaMessageBatch
|
|
from sqlalchemy import DateTime, ForeignKey, Index, String
|
|
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
|
|
|
from letta.orm.custom_columns import CreateBatchResponseColumn, PollBatchResponseColumn
|
|
from letta.orm.mixins import OrganizationMixin
|
|
from letta.orm.sqlalchemy_base import SqlalchemyBase
|
|
from letta.schemas.enums import JobStatus, ProviderType
|
|
from letta.schemas.llm_batch_job import LLMBatchJob as PydanticLLMBatchJob
|
|
|
|
|
|
class LLMBatchJob(SqlalchemyBase, OrganizationMixin):
|
|
"""Represents a single LLM batch request made to a provider like Anthropic"""
|
|
|
|
__tablename__ = "llm_batch_job"
|
|
__table_args__ = (
|
|
Index("ix_llm_batch_job_created_at", "created_at"),
|
|
Index("ix_llm_batch_job_status", "status"),
|
|
)
|
|
|
|
__pydantic_model__ = PydanticLLMBatchJob
|
|
|
|
# TODO: We want to migrate all the ORM models to do this, so we will need to move this to the SqlalchemyBase
|
|
# TODO: Some still rely on the Pydantic object to do this
|
|
id: Mapped[str] = mapped_column(String, primary_key=True, default=lambda: f"batch_req-{uuid.uuid4()}")
|
|
|
|
status: Mapped[JobStatus] = mapped_column(String, default=JobStatus.created, doc="The current status of the batch.")
|
|
|
|
llm_provider: Mapped[ProviderType] = mapped_column(String, doc="LLM provider used (e.g., 'Anthropic')")
|
|
|
|
create_batch_response: Mapped[Union[BetaMessageBatch]] = mapped_column(
|
|
CreateBatchResponseColumn, doc="Full JSON response from initial batch creation"
|
|
)
|
|
latest_polling_response: Mapped[Union[BetaMessageBatch]] = mapped_column(
|
|
PollBatchResponseColumn, nullable=True, doc="Last known polling result from LLM provider"
|
|
)
|
|
|
|
last_polled_at: Mapped[Optional[datetime]] = mapped_column(
|
|
DateTime(timezone=True), nullable=True, doc="Last time we polled the provider for status"
|
|
)
|
|
|
|
letta_batch_job_id: Mapped[str] = mapped_column(
|
|
String, ForeignKey("jobs.id", ondelete="CASCADE"), nullable=False, doc="ID of the Letta batch job"
|
|
)
|
|
|
|
organization: Mapped["Organization"] = relationship("Organization", back_populates="llm_batch_jobs")
|
|
items: Mapped[List["LLMBatchItem"]] = relationship("LLMBatchItem", back_populates="batch", lazy="selectin")
|