org id not None (#2696)

This commit is contained in:
Asher Foa
2025-06-13 10:44:04 -04:00
committed by GitHub
parent 296d2f903b
commit 687f3ec547
3 changed files with 13 additions and 15 deletions

View File

@@ -359,8 +359,9 @@ class AgentDB:
async def get_total_unique_step_order_count_by_task_ids(
self,
*,
task_ids: list[str],
organization_id: str | None = None,
organization_id: str,
) -> int:
"""
Get the total count of unique (step.task_id, step.order) pairs of StepModel for the given task ids
@@ -1009,6 +1010,8 @@ class AgentDB:
async def get_artifacts_by_entity_id(
self,
*,
organization_id: str,
artifact_type: ArtifactType | None = None,
task_id: str | None = None,
step_id: str | None = None,
@@ -1016,7 +1019,6 @@ class AgentDB:
workflow_run_block_id: str | None = None,
thought_id: str | None = None,
task_v2_id: str | None = None,
organization_id: str | None = None,
) -> list[Artifact]:
try:
async with self.Session() as session:
@@ -1053,16 +1055,18 @@ class AgentDB:
async def get_artifact_by_entity_id(
self,
*,
artifact_type: ArtifactType,
organization_id: str,
task_id: str | None = None,
step_id: str | None = None,
workflow_run_id: str | None = None,
workflow_run_block_id: str | None = None,
thought_id: str | None = None,
task_v2_id: str | None = None,
organization_id: str | None = None,
) -> Artifact | None:
artifacts = await self.get_artifacts_by_entity_id(
organization_id=organization_id,
artifact_type=artifact_type,
task_id=task_id,
step_id=step_id,
@@ -1070,7 +1074,6 @@ class AgentDB:
workflow_run_block_id=workflow_run_block_id,
thought_id=thought_id,
task_v2_id=task_v2_id,
organization_id=organization_id,
)
return artifacts[0] if artifacts else None
@@ -2401,9 +2404,10 @@ class AgentDB:
async def get_thoughts(
self,
*,
task_v2_id: str,
thought_types: list[ThoughtType] | None = None,
organization_id: str | None = None,
thought_types: list[ThoughtType],
organization_id: str,
) -> list[Thought]:
async with self.Session() as session:
query = (

View File

@@ -1135,13 +1135,10 @@ async def get_artifacts(
)
analytics.capture("skyvern-oss-agent-entity-artifacts-get")
params = {
"organization_id": current_org.organization_id,
entity_type_to_param[entity_type]: entity_id,
}
artifacts = await app.DATABASE.get_artifacts_by_entity_id(**params) # type: ignore
artifacts = await app.DATABASE.get_artifacts_by_entity_id(organization_id=current_org.organization_id, **params) # type: ignore
if settings.ENV != "local" or settings.GENERATE_PRESIGNED_URLS:
signed_urls = await app.ARTIFACT_MANAGER.get_share_links(artifacts)

View File

@@ -1351,12 +1351,9 @@ def _generate_random_string(length: int = 5) -> str:
return "".join(random.choices(RANDOM_STRING_POOL, k=length))
async def get_thought_timelines(
task_v2_id: str,
organization_id: str | None = None,
) -> list[WorkflowRunTimeline]:
async def get_thought_timelines(*, task_v2_id: str, organization_id: str) -> list[WorkflowRunTimeline]:
thoughts = await app.DATABASE.get_thoughts(
task_v2_id,
task_v2_id=task_v2_id,
organization_id=organization_id,
thought_types=[
ThoughtType.plan,