weave-python 0.58.0__tar.gz → 0.59.0__tar.gz

{weave_python-0.58.0 → weave_python-0.59.0}/PKG-INFO

@@ -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 → weave_python-0.59.0}/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]: