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")
artifacts = []
if metadata.prompt:
await app.ARTIFACT_MANAGER.create_llm_artifact(
data=metadata.prompt.encode("utf-8"),
artifact_type=ArtifactType.LLM_PROMPT,
screenshots=screenshots,
step=step,
artifacts.append(
await app.ARTIFACT_MANAGER.prepare_llm_artifact(
data=metadata.prompt.encode("utf-8"),
artifact_type=ArtifactType.LLM_PROMPT,
screenshots=screenshots,
step=step,
)
)
if metadata.llm_request_json:
await app.ARTIFACT_MANAGER.create_llm_artifact(
data=metadata.llm_request_json.encode("utf-8"),
artifact_type=ArtifactType.LLM_REQUEST,
step=step,
artifacts.append(
await app.ARTIFACT_MANAGER.prepare_llm_artifact(
data=metadata.llm_request_json.encode("utf-8"),
artifact_type=ArtifactType.LLM_REQUEST,
step=step,
)
)
if metadata.llm_response_json:
await app.ARTIFACT_MANAGER.create_llm_artifact(
data=metadata.llm_response_json.encode("utf-8"),
artifact_type=ArtifactType.LLM_RESPONSE,
step=step,
artifacts.append(
await app.ARTIFACT_MANAGER.prepare_llm_artifact(
data=metadata.llm_response_json.encode("utf-8"),
artifact_type=ArtifactType.LLM_RESPONSE,
step=step,
)
)
if metadata.parsed_response_json:
await app.ARTIFACT_MANAGER.create_llm_artifact(
data=metadata.parsed_response_json.encode("utf-8"),
artifact_type=ArtifactType.LLM_RESPONSE_PARSED,
step=step,
artifacts.append(
await app.ARTIFACT_MANAGER.prepare_llm_artifact(
data=metadata.parsed_response_json.encode("utf-8"),
artifact_type=ArtifactType.LLM_RESPONSE_PARSED,
step=step,
)
)
if metadata.rendered_response_json:
await app.ARTIFACT_MANAGER.create_llm_artifact(
data=metadata.rendered_response_json.encode("utf-8"),
artifact_type=ArtifactType.LLM_RESPONSE_RENDERED,
step=step,
artifacts.append(
await app.ARTIFACT_MANAGER.prepare_llm_artifact(
data=metadata.rendered_response_json.encode("utf-8"),
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_input_tokens = (
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,
)
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(
self,
request: BulkArtifactCreationRequest,
@@ -636,7 +654,7 @@ class ArtifactManager:
return BulkArtifactCreationRequest(artifacts=artifacts, primary_key=ai_suggestion.ai_suggestion_id)
async def create_llm_artifact(
async def prepare_llm_artifact(
self,
data: bytes,
artifact_type: ArtifactType,
@@ -645,54 +663,40 @@ class ArtifactManager:
thought: Thought | None = None,
task_v2: TaskV2 | None = None,
ai_suggestion: AISuggestion | None = None,
) -> 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
"""
) -> BulkArtifactCreationRequest | None:
if step:
request = self._prepare_step_artifacts(
return self._prepare_step_artifacts(
step=step,
artifact_type=artifact_type,
data=data,
screenshots=screenshots,
)
await self._bulk_create_artifacts(request)
elif task_v2:
request = self._prepare_task_v2_artifacts(
return self._prepare_task_v2_artifacts(
task_v2=task_v2,
artifact_type=artifact_type,
data=data,
screenshots=screenshots,
)
await self._bulk_create_artifacts(request)
elif thought:
request = self._prepare_thought_artifacts(
return self._prepare_thought_artifacts(
thought=thought,
artifact_type=artifact_type,
data=data,
screenshots=screenshots,
)
await self._bulk_create_artifacts(request)
elif ai_suggestion:
request = self._prepare_ai_suggestion_artifacts(
return self._prepare_ai_suggestion_artifacts(
ai_suggestion=ai_suggestion,
artifact_type=artifact_type,
data=data,
screenshots=screenshots,
)
await self._bulk_create_artifacts(request)
else:
return None
async def update_artifact_data(
self,