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

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

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

{weave_python-0.57.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]:

{weave_python-0.57.0 → weave_python-0.59.0}/weave/weavesql/weavedb/auth_identity.py

@@ -276,6 +276,33 @@ class CreateUserRow:
     updated_at: datetime.datetime
 
 
+CREATE_USER_WITH_ID = """-- name: create_user_with_id \\:one
+INSERT INTO weave.users (
+    id,
+    status
+) VALUES (
+    :p1,
+    :p2
+)
+ON CONFLICT (id) DO NOTHING
+RETURNING id, status, created_at, updated_at
+"""
+
+
+@dataclasses.dataclass()
+class CreateUserWithIDParams:
+    id: uuid.UUID
+    status: str
+
+
+@dataclasses.dataclass()
+class CreateUserWithIDRow:
+    id: uuid.UUID
+    status: str
+    created_at: datetime.datetime
+    updated_at: datetime.datetime
+
+
 DELETE_EXPIRED_APPLICATION_SESSIONS = """-- name: delete_expired_application_sessions \\:execrows
 DELETE FROM weave.application_sessions
 WHERE expires_at <= now()
@@ -2075,6 +2102,21 @@ class Querier:
             updated_at=row[3],
         )
 
+    def create_user_with_id(
+        self, arg: CreateUserWithIDParams
+    ) -> Optional[CreateUserWithIDRow]:
+        row = self._conn.execute(
+            sqlalchemy.text(CREATE_USER_WITH_ID), {"p1": arg.id, "p2": arg.status}
+        ).first()
+        if row is None:
+            return None
+        return CreateUserWithIDRow(
+            id=row[0],
+            status=row[1],
+            created_at=row[2],
+            updated_at=row[3],
+        )
+
     def delete_expired_application_sessions(
         self, arg: DeleteExpiredApplicationSessionsParams
     ) -> int:
@@ -3358,6 +3400,23 @@ class AsyncQuerier:
             updated_at=row[3],
         )
 
+    async def create_user_with_id(
+        self, arg: CreateUserWithIDParams
+    ) -> Optional[CreateUserWithIDRow]:
+        row = (
+            await self._conn.execute(
+                sqlalchemy.text(CREATE_USER_WITH_ID), {"p1": arg.id, "p2": arg.status}
+            )
+        ).first()
+        if row is None:
+            return None
+        return CreateUserWithIDRow(
+            id=row[0],
+            status=row[1],
+            created_at=row[2],
+            updated_at=row[3],
+        )
+
     async def delete_expired_application_sessions(
         self, arg: DeleteExpiredApplicationSessionsParams
     ) -> int:

{weave_python-0.57.0 → weave_python-0.59.0}/weave/weavesql/weavedb/organization.py

@@ -84,6 +84,70 @@ class CreateOrganizationRow:
     updated_at: datetime.datetime
 
 
+CREATE_ORGANIZATION_WITH_ID = """-- name: create_organization_with_id \\:one
+INSERT INTO weave.organizations (
+    id,
+    slug,
+    name,
+    description,
+    primary_color,
+    logo,
+    logo_content_type,
+    timezone,
+    location
+) VALUES (
+    :p1,
+    :p2,
+    :p3,
+    :p4,
+    :p5,
+    :p6,
+    :p7,
+    :p8,
+    :p9
+)
+ON CONFLICT (slug) DO NOTHING
+RETURNING
+    id,
+    slug,
+    name,
+    description,
+    primary_color,
+    (logo IS NOT NULL) AS has_logo,
+    timezone,
+    location,
+    created_at,
+    updated_at
+"""
+
+
+@dataclasses.dataclass()
+class CreateOrganizationWithIDParams:
+    id: uuid.UUID
+    slug: str
+    name: str
+    description: str
+    primary_color: str
+    logo: Optional[memoryview]
+    logo_content_type: Optional[str]
+    timezone: str
+    location: str
+
+
+@dataclasses.dataclass()
+class CreateOrganizationWithIDRow:
+    id: uuid.UUID
+    slug: str
+    name: str
+    description: str
+    primary_color: str
+    has_logo: Optional[Any]
+    timezone: str
+    location: str
+    created_at: datetime.datetime
+    updated_at: datetime.datetime
+
+
 DELETE_ORGANIZATION = """-- name: delete_organization \\:execrows
 DELETE FROM weave.organizations
 WHERE id = :p1
@@ -325,6 +389,38 @@ class Querier:
             updated_at=row[9],
         )
 
+    def create_organization_with_id(
+        self, arg: CreateOrganizationWithIDParams
+    ) -> Optional[CreateOrganizationWithIDRow]:
+        row = self._conn.execute(
+            sqlalchemy.text(CREATE_ORGANIZATION_WITH_ID),
+            {
+                "p1": arg.id,
+                "p2": arg.slug,
+                "p3": arg.name,
+                "p4": arg.description,
+                "p5": arg.primary_color,
+                "p6": arg.logo,
+                "p7": arg.logo_content_type,
+                "p8": arg.timezone,
+                "p9": arg.location,
+            },
+        ).first()
+        if row is None:
+            return None
+        return CreateOrganizationWithIDRow(
+            id=row[0],
+            slug=row[1],
+            name=row[2],
+            description=row[3],
+            primary_color=row[4],
+            has_logo=row[5],
+            timezone=row[6],
+            location=row[7],
+            created_at=row[8],
+            updated_at=row[9],
+        )
+
     def delete_organization(self, arg: DeleteOrganizationParams) -> int:
         result = self._conn.execute(
             sqlalchemy.text(DELETE_ORGANIZATION), {"p1": arg.id}
@@ -485,6 +581,40 @@ class AsyncQuerier:
             updated_at=row[9],
         )
 
+    async def create_organization_with_id(
+        self, arg: CreateOrganizationWithIDParams
+    ) -> Optional[CreateOrganizationWithIDRow]:
+        row = (
+            await self._conn.execute(
+                sqlalchemy.text(CREATE_ORGANIZATION_WITH_ID),
+                {
+                    "p1": arg.id,
+                    "p2": arg.slug,
+                    "p3": arg.name,
+                    "p4": arg.description,
+                    "p5": arg.primary_color,
+                    "p6": arg.logo,
+                    "p7": arg.logo_content_type,
+                    "p8": arg.timezone,
+                    "p9": arg.location,
+                },
+            )
+        ).first()
+        if row is None:
+            return None
+        return CreateOrganizationWithIDRow(
+            id=row[0],
+            slug=row[1],
+            name=row[2],
+            description=row[3],
+            primary_color=row[4],
+            has_logo=row[5],
+            timezone=row[6],
+            location=row[7],
+            created_at=row[8],
+            updated_at=row[9],
+        )
+
     async def delete_organization(self, arg: DeleteOrganizationParams) -> int:
         result = await self._conn.execute(
             sqlalchemy.text(DELETE_ORGANIZATION), {"p1": arg.id}