Batch LLM artifacts creation (#4322)

This commit is contained in:
Stanislav Novosad
2025-12-17 20:15:26 -07:00
committed by GitHub
parent f594474b9e
commit 1eca20b78a
3 changed files with 908 additions and 828 deletions

View File

@@ -1880,42 +1880,56 @@ class ForgeAgent:
LOG.debug("Persisting speculative LLM metadata") LOG.debug("Persisting speculative LLM metadata")
artifacts = []
if metadata.prompt: if metadata.prompt:
await app.ARTIFACT_MANAGER.create_llm_artifact( artifacts.append(
data=metadata.prompt.encode("utf-8"), await app.ARTIFACT_MANAGER.prepare_llm_artifact(
artifact_type=ArtifactType.LLM_PROMPT, data=metadata.prompt.encode("utf-8"),
screenshots=screenshots, artifact_type=ArtifactType.LLM_PROMPT,
step=step, screenshots=screenshots,
step=step,
)
) )
if metadata.llm_request_json: if metadata.llm_request_json:
await app.ARTIFACT_MANAGER.create_llm_artifact( artifacts.append(
data=metadata.llm_request_json.encode("utf-8"), await app.ARTIFACT_MANAGER.prepare_llm_artifact(
artifact_type=ArtifactType.LLM_REQUEST, data=metadata.llm_request_json.encode("utf-8"),
step=step, artifact_type=ArtifactType.LLM_REQUEST,
step=step,
)
) )
if metadata.llm_response_json: if metadata.llm_response_json:
await app.ARTIFACT_MANAGER.create_llm_artifact( artifacts.append(
data=metadata.llm_response_json.encode("utf-8"), await app.ARTIFACT_MANAGER.prepare_llm_artifact(
artifact_type=ArtifactType.LLM_RESPONSE, data=metadata.llm_response_json.encode("utf-8"),
step=step, artifact_type=ArtifactType.LLM_RESPONSE,
step=step,
)
) )
if metadata.parsed_response_json: if metadata.parsed_response_json:
await app.ARTIFACT_MANAGER.create_llm_artifact( artifacts.append(
data=metadata.parsed_response_json.encode("utf-8"), await app.ARTIFACT_MANAGER.prepare_llm_artifact(
artifact_type=ArtifactType.LLM_RESPONSE_PARSED, data=metadata.parsed_response_json.encode("utf-8"),
step=step, artifact_type=ArtifactType.LLM_RESPONSE_PARSED,
step=step,
)
) )
if metadata.rendered_response_json: if metadata.rendered_response_json:
await app.ARTIFACT_MANAGER.create_llm_artifact( artifacts.append(
data=metadata.rendered_response_json.encode("utf-8"), await app.ARTIFACT_MANAGER.prepare_llm_artifact(
artifact_type=ArtifactType.LLM_RESPONSE_RENDERED, data=metadata.rendered_response_json.encode("utf-8"),
step=step, artifact_type=ArtifactType.LLM_RESPONSE_RENDERED,
step=step,
)
) )
if artifacts:
await app.ARTIFACT_MANAGER.bulk_create_artifacts(artifacts)
incremental_cost = metadata.llm_cost if metadata.llm_cost and metadata.llm_cost > 0 else None incremental_cost = metadata.llm_cost if metadata.llm_cost and metadata.llm_cost > 0 else None
incremental_input_tokens = ( incremental_input_tokens = (
metadata.input_tokens if metadata.input_tokens and metadata.input_tokens > 0 else None metadata.input_tokens if metadata.input_tokens and metadata.input_tokens > 0 else None

File diff suppressed because it is too large Load Diff

View File

@@ -367,6 +367,24 @@ class ArtifactManager:
data=data, data=data,
) )
async def bulk_create_artifacts(
self,
requests: list[BulkArtifactCreationRequest | None],
) -> list[str]:
artifacts: list[ArtifactBatchData] = []
primary_key: str | None = None
for request in requests:
if request:
artifacts.extend(request.artifacts)
primary_key = request.primary_key
if primary_key is None or not artifacts:
return []
return await self._bulk_create_artifacts(
BulkArtifactCreationRequest(artifacts=artifacts, primary_key=primary_key)
)
async def _bulk_create_artifacts( async def _bulk_create_artifacts(
self, self,
request: BulkArtifactCreationRequest, request: BulkArtifactCreationRequest,
@@ -636,7 +654,7 @@ class ArtifactManager:
return BulkArtifactCreationRequest(artifacts=artifacts, primary_key=ai_suggestion.ai_suggestion_id) return BulkArtifactCreationRequest(artifacts=artifacts, primary_key=ai_suggestion.ai_suggestion_id)
async def create_llm_artifact( async def prepare_llm_artifact(
self, self,
data: bytes, data: bytes,
artifact_type: ArtifactType, artifact_type: ArtifactType,
@@ -645,54 +663,40 @@ class ArtifactManager:
thought: Thought | None = None, thought: Thought | None = None,
task_v2: TaskV2 | None = None, task_v2: TaskV2 | None = None,
ai_suggestion: AISuggestion | None = None, ai_suggestion: AISuggestion | None = None,
) -> None: ) -> BulkArtifactCreationRequest | None:
"""
Create LLM artifact with optional screenshots using bulk insert.
Args:
data: Main artifact data
artifact_type: Type of the main artifact
screenshots: Optional list of screenshot data
step: Optional Step entity
thought: Optional Thought entity
task_v2: Optional TaskV2 entity
ai_suggestion: Optional AISuggestion entity
"""
if step: if step:
request = self._prepare_step_artifacts( return self._prepare_step_artifacts(
step=step, step=step,
artifact_type=artifact_type, artifact_type=artifact_type,
data=data, data=data,
screenshots=screenshots, screenshots=screenshots,
) )
await self._bulk_create_artifacts(request)
elif task_v2: elif task_v2:
request = self._prepare_task_v2_artifacts( return self._prepare_task_v2_artifacts(
task_v2=task_v2, task_v2=task_v2,
artifact_type=artifact_type, artifact_type=artifact_type,
data=data, data=data,
screenshots=screenshots, screenshots=screenshots,
) )
await self._bulk_create_artifacts(request)
elif thought: elif thought:
request = self._prepare_thought_artifacts( return self._prepare_thought_artifacts(
thought=thought, thought=thought,
artifact_type=artifact_type, artifact_type=artifact_type,
data=data, data=data,
screenshots=screenshots, screenshots=screenshots,
) )
await self._bulk_create_artifacts(request)
elif ai_suggestion: elif ai_suggestion:
request = self._prepare_ai_suggestion_artifacts( return self._prepare_ai_suggestion_artifacts(
ai_suggestion=ai_suggestion, ai_suggestion=ai_suggestion,
artifact_type=artifact_type, artifact_type=artifact_type,
data=data, data=data,
screenshots=screenshots, screenshots=screenshots,
) )
await self._bulk_create_artifacts(request) else:
return None
async def update_artifact_data( async def update_artifact_data(
self, self,