add cached steps flag to db (is_script_cached in Step) (#4202)
This commit is contained in:
@@ -257,6 +257,7 @@ class AgentDB:
|
||||
retry_index: int,
|
||||
organization_id: str | None = None,
|
||||
status: StepStatus = StepStatus.created,
|
||||
created_by: str | None = None,
|
||||
) -> Step:
|
||||
try:
|
||||
async with self.Session() as session:
|
||||
@@ -266,6 +267,7 @@ class AgentDB:
|
||||
retry_index=retry_index,
|
||||
status=status,
|
||||
organization_id=organization_id,
|
||||
created_by=created_by,
|
||||
)
|
||||
session.add(new_step)
|
||||
await session.commit()
|
||||
@@ -595,6 +597,7 @@ class AgentDB:
|
||||
incremental_output_tokens: int | None = None,
|
||||
incremental_reasoning_tokens: int | None = None,
|
||||
incremental_cached_tokens: int | None = None,
|
||||
created_by: str | None = None,
|
||||
) -> Step:
|
||||
try:
|
||||
async with self.Session() as session:
|
||||
@@ -627,6 +630,8 @@ class AgentDB:
|
||||
step.reasoning_token_count = incremental_reasoning_tokens + (step.reasoning_token_count or 0)
|
||||
if incremental_cached_tokens is not None:
|
||||
step.cached_token_count = incremental_cached_tokens + (step.cached_token_count or 0)
|
||||
if created_by is not None:
|
||||
step.created_by = created_by
|
||||
|
||||
await session.commit()
|
||||
updated_step = await self.get_step(step_id, organization_id)
|
||||
|
||||
@@ -142,6 +142,7 @@ class StepModel(Base):
|
||||
cached_token_count = Column(Integer, default=0)
|
||||
step_cost = Column(Numeric, default=0)
|
||||
finished_at = Column(DateTime, nullable=True)
|
||||
created_by = Column(String, nullable=True)
|
||||
|
||||
|
||||
class OrganizationModel(Base):
|
||||
|
||||
@@ -234,6 +234,7 @@ def convert_to_step(step_model: StepModel, debug_enabled: bool = False) -> Step:
|
||||
reasoning_token_count=step_model.reasoning_token_count,
|
||||
cached_token_count=step_model.cached_token_count,
|
||||
step_cost=step_model.step_cost,
|
||||
created_by=step_model.created_by,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -71,6 +71,7 @@ class Step(BaseModel):
|
||||
reasoning_token_count: int | None = None
|
||||
cached_token_count: int | None = None
|
||||
step_cost: float = 0
|
||||
created_by: str | None = None
|
||||
is_speculative: bool = False
|
||||
speculative_original_status: StepStatus | None = None
|
||||
speculative_llm_metadata: SpeculativeLLMMetadata | None = None
|
||||
|
||||
Reference in New Issue
Block a user