weave-python 0.58.0__py3-none-any.whl → 0.59.0__py3-none-any.whl

weave/weavesql/weavedb/agent.py

@@ -13,6 +13,30 @@ import sqlalchemy.ext.asyncio
 from weave.weavesql.weavedb import models
 
 
+CLAIM_AGENT_RUN = """-- name: claim_agent_run \\:one
+UPDATE weave.agent_runs
+SET
+    status = :p1,
+    error_code = '',
+    safe_error_message = '',
+    started_at = COALESCE(started_at, :p2\\:\\:timestamptz),
+    finished_at = NULL
+WHERE organization_id = :p3
+  AND id = :p4
+  AND status = :p5
+RETURNING id, organization_id, agent_definition_id, agent_version_id, status, created_by_user_id, chat_session_id, workflow_run_id, workflow_step_run_id, input, output, error_code, safe_error_message, created_at, started_at, finished_at
+"""
+
+
+@dataclasses.dataclass()
+class ClaimAgentRunParams:
+    status: str
+    started_at: datetime.datetime
+    organization_id: uuid.UUID
+    agent_run_id: uuid.UUID
+    expected_status: str
+
+
 CREATE_AGENT_DEFINITION = """-- name: create_agent_definition \\:one
 INSERT INTO weave.agent_definitions (
     id,
@@ -632,6 +656,40 @@ class Querier:
     def __init__(self, conn: sqlalchemy.engine.Connection):
         self._conn = conn
 
+    def claim_agent_run(
+        self, arg: ClaimAgentRunParams
+    ) -> Optional[models.WeaveAgentRun]:
+        row = self._conn.execute(
+            sqlalchemy.text(CLAIM_AGENT_RUN),
+            {
+                "p1": arg.status,
+                "p2": arg.started_at,
+                "p3": arg.organization_id,
+                "p4": arg.agent_run_id,
+                "p5": arg.expected_status,
+            },
+        ).first()
+        if row is None:
+            return None
+        return models.WeaveAgentRun(
+            id=row[0],
+            organization_id=row[1],
+            agent_definition_id=row[2],
+            agent_version_id=row[3],
+            status=row[4],
+            created_by_user_id=row[5],
+            chat_session_id=row[6],
+            workflow_run_id=row[7],
+            workflow_step_run_id=row[8],
+            input=row[9],
+            output=row[10],
+            error_code=row[11],
+            safe_error_message=row[12],
+            created_at=row[13],
+            started_at=row[14],
+            finished_at=row[15],
+        )
+
     def create_agent_definition(
         self, arg: CreateAgentDefinitionParams
     ) -> Optional[models.WeaveAgentDefinition]:
@@ -1360,6 +1418,42 @@ class AsyncQuerier:
     def __init__(self, conn: sqlalchemy.ext.asyncio.AsyncConnection):
         self._conn = conn
 
+    async def claim_agent_run(
+        self, arg: ClaimAgentRunParams
+    ) -> Optional[models.WeaveAgentRun]:
+        row = (
+            await self._conn.execute(
+                sqlalchemy.text(CLAIM_AGENT_RUN),
+                {
+                    "p1": arg.status,
+                    "p2": arg.started_at,
+                    "p3": arg.organization_id,
+                    "p4": arg.agent_run_id,
+                    "p5": arg.expected_status,
+                },
+            )
+        ).first()
+        if row is None:
+            return None
+        return models.WeaveAgentRun(
+            id=row[0],
+            organization_id=row[1],
+            agent_definition_id=row[2],
+            agent_version_id=row[3],
+            status=row[4],
+            created_by_user_id=row[5],
+            chat_session_id=row[6],
+            workflow_run_id=row[7],
+            workflow_step_run_id=row[8],
+            input=row[9],
+            output=row[10],
+            error_code=row[11],
+            safe_error_message=row[12],
+            created_at=row[13],
+            started_at=row[14],
+            finished_at=row[15],
+        )
+
     async def create_agent_definition(
         self, arg: CreateAgentDefinitionParams
     ) -> Optional[models.WeaveAgentDefinition]:

{weave_python-0.58.0.dist-info → weave_python-0.59.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weave-python
-Version: 0.58.0
+Version: 0.59.0
 Author: Weave Labs
 License: FCL-1.0-ALv2
 License-File: LICENSE

{weave_python-0.58.0.dist-info → weave_python-0.59.0.dist-info}/RECORD

@@ -129,7 +129,7 @@ weave/weavesql/atcdb/models.py,sha256=FHBXYp_VOQIMJhB_uviKjDo9QTz1c88DRjx0l6otrS
 weave/weavesql/atcdb/run_template.py,sha256=kIa3JASDW0GgcErLg8NfFSh31KmzgHoLU102D19JsL0,12328
 weave/weavesql/atcdb/snapshot.py,sha256=Rgdvmy1GHr9n7lenLOEoEkukHi9eCzQGaGVUQBcqzYc,6411
 weave/weavesql/atcdb/team.py,sha256=4Q9DlRi4pMOxxoHEBuya6SZ-VMWumQBpG8cYhx8RBvA,12148
-weave/weavesql/weavedb/agent.py,sha256=KG4Ma6gngqKUL72mgp7P3u-g_Lle_AbLYVm2CMgNysU,66863
+weave/weavesql/weavedb/agent.py,sha256=mmZ6hGpN0ifVP7GABytv4hfDrswLIECJDDJGhWQbTIQ,69789
 weave/weavesql/weavedb/auth_identity.py,sha256=qW1JNoT9tlnRP1lRSry1wodITvsBc_BzS4lHlGrVmnw,134030
 weave/weavesql/weavedb/chat.py,sha256=STtu97aOKidcizsSLHpLFeo8DdYnv3PgpA2m6cW2hpE,58177
 weave/weavesql/weavedb/consolidation.py,sha256=kzZX-o3JF4n0r7rankiWN4I_eywsZUTHr9ujDmTG5nY,30785
@@ -143,7 +143,7 @@ weave/weavesql/weavedb/provider.py,sha256=zfBtTFiPGy-gJxil_Om3-N0E3aJpHYTsC1P3Jp
 weave/weavesql/weavedb/retrieval.py,sha256=tRGNZpPglYc1IBFWrQHKSulXTsWVpNgs49cB8L-nCPc,61207
 weave/weavesql/weavedb/task.py,sha256=AS2wig-e4UnlPHJ6Uk53XEa3lF9RLKwGpN4m9vmFqzM,23966
 weave/weavesql/weavedb/workflow.py,sha256=ICEaiczoMcJoAmvXc-lAX_kSfAQqaCfSGXZD9V1y4ac,79625
-weave_python-0.58.0.dist-info/METADATA,sha256=8Dal9YTEnico5kGxW-Zqb8EIqhRGI2f8Xa7LQl-gjkE,1444
-weave_python-0.58.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
-weave_python-0.58.0.dist-info/licenses/LICENSE,sha256=oO00b58BS5uBxjk80lqVRUm_p8yMsGilVGCdO4ulm5Y,4216
-weave_python-0.58.0.dist-info/RECORD,,
+weave_python-0.59.0.dist-info/METADATA,sha256=I3znFozj-ujKViZqIZSjc1HF2Dd-BDbPGu63eMNrNzU,1444
+weave_python-0.59.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+weave_python-0.59.0.dist-info/licenses/LICENSE,sha256=oO00b58BS5uBxjk80lqVRUm_p8yMsGilVGCdO4ulm5Y,4216
+weave_python-0.59.0.dist-info/RECORD,,