org id not None (#2696)
This commit is contained in:
@@ -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 = (
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user