intentkit 0.7.5.dev3__py3-none-any.whl → 0.8.34.dev7__py3-none-any.whl

intentkit/models/agent_data.py

@@ -1,24 +1,19 @@
+from __future__ import annotations
+
 import logging
-from datetime import datetime, timezone
+from datetime import UTC, datetime
 from decimal import Decimal
-from typing import Annotated, Any, Dict, Optional
+from typing import Annotated, Any
 
-from fastapi import HTTPException
-from intentkit.models.base import Base
-from intentkit.models.db import get_session
 from pydantic import BaseModel, ConfigDict
 from pydantic import Field as PydanticField
-from sqlalchemy import (
-    BigInteger,
-    Boolean,
-    Column,
-    DateTime,
-    Numeric,
-    String,
-    func,
-    select,
-)
+from sqlalchemy import BigInteger, Boolean, DateTime, Numeric, String, func, select
 from sqlalchemy.dialects.postgresql import JSON, JSONB
+from sqlalchemy.orm import Mapped, mapped_column
+
+from intentkit.models.base import Base
+from intentkit.models.db import get_session
+from intentkit.utils.error import IntentKitAPIError
 
 logger = logging.getLogger(__name__)
 
@@ -28,56 +23,87 @@ class AgentDataTable(Base):
 
     __tablename__ = "agent_data"
 
-    id = Column(String, primary_key=True, comment="Same as Agent.id")
-    evm_wallet_address = Column(String, nullable=True, comment="EVM wallet address")
-    solana_wallet_address = Column(
+    id: Mapped[str] = mapped_column(
+        String, primary_key=True, comment="Same as Agent.id"
+    )
+    evm_wallet_address: Mapped[str | None] = mapped_column(
+        String, nullable=True, comment="EVM wallet address"
+    )
+    solana_wallet_address: Mapped[str | None] = mapped_column(
         String, nullable=True, comment="Solana wallet address"
     )
-    cdp_wallet_data = Column(String, nullable=True, comment="CDP wallet data")
-    twitter_id = Column(String, nullable=True, comment="Twitter user ID")
-    twitter_username = Column(String, nullable=True, comment="Twitter username")
-    twitter_name = Column(String, nullable=True, comment="Twitter display name")
-    twitter_access_token = Column(String, nullable=True, comment="Twitter access token")
-    twitter_access_token_expires_at = Column(
+    cdp_wallet_data: Mapped[str | None] = mapped_column(
+        String, nullable=True, comment="CDP wallet data"
+    )
+    twitter_id: Mapped[str | None] = mapped_column(
+        String, nullable=True, comment="Twitter user ID"
+    )
+    twitter_username: Mapped[str | None] = mapped_column(
+        String, nullable=True, comment="Twitter username"
+    )
+    twitter_name: Mapped[str | None] = mapped_column(
+        String, nullable=True, comment="Twitter display name"
+    )
+    twitter_access_token: Mapped[str | None] = mapped_column(
+        String, nullable=True, comment="Twitter access token"
+    )
+    twitter_access_token_expires_at: Mapped[datetime | None] = mapped_column(
         DateTime(timezone=True),
         nullable=True,
         comment="Twitter access token expiration time",
     )
-    twitter_refresh_token = Column(
+    twitter_refresh_token: Mapped[str | None] = mapped_column(
         String, nullable=True, comment="Twitter refresh token"
     )
-    twitter_self_key_refreshed_at = Column(
+    twitter_self_key_refreshed_at: Mapped[datetime | None] = mapped_column(
         DateTime(timezone=True),
         nullable=True,
         comment="Twitter self-key userinfo last refresh time",
     )
-    twitter_is_verified = Column(
+    twitter_is_verified: Mapped[bool] = mapped_column(
         Boolean,
         nullable=False,
         default=False,
         comment="Whether the Twitter account is verified",
     )
-    telegram_id = Column(String, nullable=True, comment="Telegram user ID")
-    telegram_username = Column(String, nullable=True, comment="Telegram username")
-    telegram_name = Column(String, nullable=True, comment="Telegram display name")
-    error_message = Column(String, nullable=True, comment="Last error message")
-    api_key = Column(
+    telegram_id: Mapped[str | None] = mapped_column(
+        String, nullable=True, comment="Telegram user ID"
+    )
+    telegram_username: Mapped[str | None] = mapped_column(
+        String, nullable=True, comment="Telegram username"
+    )
+    telegram_name: Mapped[str | None] = mapped_column(
+        String, nullable=True, comment="Telegram display name"
+    )
+    discord_id: Mapped[str | None] = mapped_column(
+        String, nullable=True, comment="Discord user ID"
+    )
+    discord_username: Mapped[str | None] = mapped_column(
+        String, nullable=True, comment="Discord username"
+    )
+    discord_name: Mapped[str | None] = mapped_column(
+        String, nullable=True, comment="Discord display name"
+    )
+    error_message: Mapped[str | None] = mapped_column(
+        String, nullable=True, comment="Last error message"
+    )
+    api_key: Mapped[str | None] = mapped_column(
         String, nullable=True, unique=True, comment="API key for the agent"
     )
-    api_key_public = Column(
+    api_key_public: Mapped[str | None] = mapped_column(
         String, nullable=True, unique=True, comment="Public API key for the agent"
     )
-    created_at = Column(
+    created_at: Mapped[datetime] = mapped_column(
         DateTime(timezone=True),
         nullable=False,
         server_default=func.now(),
         comment="Timestamp when the agent data was created",
     )
-    updated_at = Column(
+    updated_at: Mapped[datetime] = mapped_column(
         DateTime(timezone=True),
         nullable=False,
         server_default=func.now(),
-        onupdate=lambda: datetime.now(timezone.utc),
+        onupdate=lambda: datetime.now(UTC),
         comment="Timestamp when the agent data was last updated",
     )
 
@@ -94,70 +120,70 @@ class AgentData(BaseModel):
         ),
     ]
     evm_wallet_address: Annotated[
-        Optional[str],
+        str | None,
         PydanticField(
             default=None,
             description="EVM wallet address",
         ),
     ] = None
     solana_wallet_address: Annotated[
-        Optional[str],
+        str | None,
         PydanticField(
             default=None,
             description="Solana wallet address",
         ),
     ] = None
     cdp_wallet_data: Annotated[
-        Optional[str],
+        str | None,
         PydanticField(
             default=None,
             description="CDP wallet data",
         ),
     ] = None
     twitter_id: Annotated[
-        Optional[str],
+        str | None,
         PydanticField(
             default=None,
             description="Twitter user ID",
         ),
     ] = None
     twitter_username: Annotated[
-        Optional[str],
+        str | None,
         PydanticField(
             default=None,
             description="Twitter username",
         ),
     ] = None
     twitter_name: Annotated[
-        Optional[str],
+        str | None,
         PydanticField(
             default=None,
             description="Twitter display name",
         ),
     ] = None
     twitter_access_token: Annotated[
-        Optional[str],
+        str | None,
         PydanticField(
             default=None,
             description="Twitter access token",
         ),
     ] = None
     twitter_access_token_expires_at: Annotated[
-        Optional[datetime],
+        datetime | None,
         PydanticField(
             default=None,
             description="Twitter access token expiration time",
         ),
     ] = None
     twitter_refresh_token: Annotated[
-        Optional[str],
+        str | None,
         PydanticField(
             default=None,
             description="Twitter refresh token",
         ),
     ] = None
     twitter_self_key_refreshed_at: Annotated[
-        Optional[datetime],
+        datetime | None,
         PydanticField(
             default=None,
             description="Twitter self-key userinfo last refresh time",
@@ -171,42 +197,63 @@ class AgentData(BaseModel):
         ),
     ] = None
     telegram_id: Annotated[
-        Optional[str],
+        str | None,
         PydanticField(
             default=None,
             description="Telegram user ID",
         ),
     ] = None
     telegram_username: Annotated[
-        Optional[str],
+        str | None,
         PydanticField(
             default=None,
             description="Telegram username",
         ),
     ] = None
     telegram_name: Annotated[
-        Optional[str],
+        str | None,
         PydanticField(
             default=None,
             description="Telegram display name",
         ),
     ] = None
+    discord_id: Annotated[
+        str | None,
+        PydanticField(
+            default=None,
+            description="Discord user ID",
+        ),
+    ] = None
+    discord_username: Annotated[
+        str | None,
+        PydanticField(
+            default=None,
+            description="Discord username",
+        ),
+    ] = None
+    discord_name: Annotated[
+        str | None,
+        PydanticField(
+            default=None,
+            description="Discord display name",
+        ),
+    ] = None
     error_message: Annotated[
-        Optional[str],
+        str | None,
         PydanticField(
             default=None,
             description="Last error message",
         ),
     ] = None
     api_key: Annotated[
-        Optional[str],
+        str | None,
         PydanticField(
             default=None,
             description="API key for the agent",
         ),
     ] = None
     api_key_public: Annotated[
-        Optional[str],
+        str | None,
         PydanticField(
             default=None,
             description="Public API key for the agent",
@@ -215,14 +262,14 @@ class AgentData(BaseModel):
     created_at: Annotated[
         datetime,
         PydanticField(
-            default_factory=lambda: datetime.now(timezone.utc),
+            default_factory=lambda: datetime.now(UTC),
             description="Timestamp when the agent data was created",
         ),
     ]
     updated_at: Annotated[
         datetime,
         PydanticField(
-            default_factory=lambda: datetime.now(timezone.utc),
+            default_factory=lambda: datetime.now(UTC),
             description="Timestamp when the agent data was last updated",
         ),
     ]
@@ -247,7 +294,7 @@ class AgentData(BaseModel):
             return cls.model_construct(id=agent_id)
 
     @classmethod
-    async def get_by_api_key(cls, api_key: str) -> Optional["AgentData"]:
+    async def get_by_api_key(cls, api_key: str) -> AgentData | None:
         """Get agent data by API key.
 
         Args:
@@ -343,20 +390,22 @@ class AgentPluginDataTable(Base):
 
     __tablename__ = "agent_plugin_data"
 
-    agent_id = Column(String, primary_key=True)
-    plugin = Column(String, primary_key=True)
-    key = Column(String, primary_key=True)
-    data = Column(JSON().with_variant(JSONB(), "postgresql"), nullable=True)
-    created_at = Column(
+    agent_id: Mapped[str] = mapped_column(String, primary_key=True)
+    plugin: Mapped[str] = mapped_column(String, primary_key=True)
+    key: Mapped[str] = mapped_column(String, primary_key=True)
+    data: Mapped[dict[str, Any] | None] = mapped_column(
+        JSON().with_variant(JSONB(), "postgresql"), nullable=True
+    )
+    created_at: Mapped[datetime] = mapped_column(
         DateTime(timezone=True),
         nullable=False,
         server_default=func.now(),
     )
-    updated_at = Column(
+    updated_at: Mapped[datetime] = mapped_column(
         DateTime(timezone=True),
         nullable=False,
         server_default=func.now(),
-        onupdate=lambda: datetime.now(timezone.utc),
+        onupdate=lambda: datetime.now(UTC),
     )
 
 
@@ -388,28 +437,28 @@ class AgentPluginData(BaseModel):
         PydanticField(description="Key for this specific piece of data"),
     ]
     data: Annotated[
-        Dict[str, Any],
+        dict[str, Any],
         PydanticField(default=None, description="JSON data stored for this key"),
     ]
     created_at: Annotated[
         datetime,
         PydanticField(
             description="Timestamp when this data was created",
-            default_factory=lambda: datetime.now(timezone.utc),
+            default_factory=lambda: datetime.now(UTC),
         ),
     ]
     updated_at: Annotated[
         datetime,
         PydanticField(
             description="Timestamp when this data was last updated",
-            default_factory=lambda: datetime.now(timezone.utc),
+            default_factory=lambda: datetime.now(UTC),
         ),
     ]
 
     @classmethod
     async def get(
         cls, agent_id: str, plugin: str, key: str
-    ) -> Optional["AgentPluginData"]:
+    ) -> "AgentPluginData" | None:
         """Get plugin data for an agent.
 
         Args:
@@ -476,44 +525,50 @@ class AgentQuotaTable(Base):
 
     __tablename__ = "agent_quotas"
 
-    id = Column(String, primary_key=True)
-    plan = Column(String, default="self-hosted")
-    message_count_total = Column(BigInteger, default=0)
-    message_limit_total = Column(BigInteger, default=99999999)
-    message_count_monthly = Column(BigInteger, default=0)
-    message_limit_monthly = Column(BigInteger, default=99999999)
-    message_count_daily = Column(BigInteger, default=0)
-    message_limit_daily = Column(BigInteger, default=99999999)
-    last_message_time = Column(DateTime(timezone=True), default=None, nullable=True)
-    autonomous_count_total = Column(BigInteger, default=0)
-    autonomous_limit_total = Column(BigInteger, default=99999999)
-    autonomous_count_monthly = Column(BigInteger, default=0)
-    autonomous_limit_monthly = Column(BigInteger, default=99999999)
-    last_autonomous_time = Column(DateTime(timezone=True), default=None, nullable=True)
-    twitter_count_total = Column(BigInteger, default=0)
-    twitter_limit_total = Column(BigInteger, default=99999999)
-    twitter_count_monthly = Column(BigInteger, default=0)
-    twitter_limit_monthly = Column(BigInteger, default=99999999)
-    twitter_count_daily = Column(BigInteger, default=0)
-    twitter_limit_daily = Column(BigInteger, default=99999999)
-    last_twitter_time = Column(DateTime(timezone=True), default=None, nullable=True)
-    free_income_daily = Column(Numeric(22, 4), default=0)
-    avg_action_cost = Column(Numeric(22, 4), default=0)
-    min_action_cost = Column(Numeric(22, 4), default=0)
-    max_action_cost = Column(Numeric(22, 4), default=0)
-    low_action_cost = Column(Numeric(22, 4), default=0)
-    medium_action_cost = Column(Numeric(22, 4), default=0)
-    high_action_cost = Column(Numeric(22, 4), default=0)
-    created_at = Column(
+    id: Mapped[str] = mapped_column(String, primary_key=True)
+    plan: Mapped[str] = mapped_column(String, default="self-hosted")
+    message_count_total: Mapped[int] = mapped_column(BigInteger, default=0)
+    message_limit_total: Mapped[int] = mapped_column(BigInteger, default=99999999)
+    message_count_monthly: Mapped[int] = mapped_column(BigInteger, default=0)
+    message_limit_monthly: Mapped[int] = mapped_column(BigInteger, default=99999999)
+    message_count_daily: Mapped[int] = mapped_column(BigInteger, default=0)
+    message_limit_daily: Mapped[int] = mapped_column(BigInteger, default=99999999)
+    last_message_time: Mapped[datetime | None] = mapped_column(
+        DateTime(timezone=True), default=None, nullable=True
+    )
+    autonomous_count_total: Mapped[int] = mapped_column(BigInteger, default=0)
+    autonomous_limit_total: Mapped[int] = mapped_column(BigInteger, default=99999999)
+    autonomous_count_monthly: Mapped[int] = mapped_column(BigInteger, default=0)
+    autonomous_limit_monthly: Mapped[int] = mapped_column(BigInteger, default=99999999)
+    last_autonomous_time: Mapped[datetime | None] = mapped_column(
+        DateTime(timezone=True), default=None, nullable=True
+    )
+    twitter_count_total: Mapped[int] = mapped_column(BigInteger, default=0)
+    twitter_limit_total: Mapped[int] = mapped_column(BigInteger, default=99999999)
+    twitter_count_monthly: Mapped[int] = mapped_column(BigInteger, default=0)
+    twitter_limit_monthly: Mapped[int] = mapped_column(BigInteger, default=99999999)
+    twitter_count_daily: Mapped[int] = mapped_column(BigInteger, default=0)
+    twitter_limit_daily: Mapped[int] = mapped_column(BigInteger, default=99999999)
+    last_twitter_time: Mapped[datetime | None] = mapped_column(
+        DateTime(timezone=True), default=None, nullable=True
+    )
+    free_income_daily: Mapped[Decimal] = mapped_column(Numeric(22, 4), default=0)
+    avg_action_cost: Mapped[Decimal] = mapped_column(Numeric(22, 4), default=0)
+    min_action_cost: Mapped[Decimal] = mapped_column(Numeric(22, 4), default=0)
+    max_action_cost: Mapped[Decimal] = mapped_column(Numeric(22, 4), default=0)
+    low_action_cost: Mapped[Decimal] = mapped_column(Numeric(22, 4), default=0)
+    medium_action_cost: Mapped[Decimal] = mapped_column(Numeric(22, 4), default=0)
+    high_action_cost: Mapped[Decimal] = mapped_column(Numeric(22, 4), default=0)
+    created_at: Mapped[datetime] = mapped_column(
         DateTime(timezone=True),
         nullable=False,
         server_default=func.now(),
     )
-    updated_at = Column(
+    updated_at: Mapped[datetime] = mapped_column(
         DateTime(timezone=True),
         nullable=False,
         server_default=func.now(),
-        onupdate=lambda: datetime.now(timezone.utc),
+        onupdate=lambda: datetime.now(UTC),
     )
 
 
@@ -547,7 +602,7 @@ class AgentQuota(BaseModel):
         int, PydanticField(default=99999999, description="Daily message limit")
     ]
     last_message_time: Annotated[
-        Optional[datetime],
+        datetime | None,
         PydanticField(default=None, description="Last message timestamp"),
     ]
     autonomous_count_total: Annotated[
@@ -578,7 +633,7 @@ class AgentQuota(BaseModel):
         ),
     ]
     last_autonomous_time: Annotated[
-        Optional[datetime],
+        datetime | None,
         PydanticField(default=None, description="Last autonomous operation timestamp"),
     ]
     twitter_count_total: Annotated[
@@ -603,7 +658,7 @@ class AgentQuota(BaseModel):
         PydanticField(default=99999999, description="Daily Twitter operations limit"),
     ]
     last_twitter_time: Annotated[
-        Optional[datetime],
+        datetime | None,
         PydanticField(default=None, description="Last Twitter operation timestamp"),
     ]
     free_income_daily: Annotated[
@@ -638,14 +693,14 @@ class AgentQuota(BaseModel):
         datetime,
         PydanticField(
             description="Timestamp when this quota was created",
-            default_factory=lambda: datetime.now(timezone.utc),
+            default_factory=lambda: datetime.now(UTC),
         ),
     ]
     updated_at: Annotated[
         datetime,
         PydanticField(
             description="Timestamp when this quota was last updated",
-            default_factory=lambda: datetime.now(timezone.utc),
+            default_factory=lambda: datetime.now(UTC),
         ),
     ]
 
@@ -754,7 +809,11 @@ class AgentQuota(BaseModel):
                 await session.execute(stmt)
         except Exception as e:
             logger.error(f"Error adding free income: {str(e)}")
-            raise HTTPException(status_code=500, detail=f"Database error: {str(e)}")
+            raise IntentKitAPIError(
+                status_code=500,
+                key="DatabaseError",
+                message=f"Database error: {str(e)}",
+            )
 
     async def add_message(self) -> None:
         """Add a message to the agent's message count."""
@@ -766,7 +825,7 @@ class AgentQuota(BaseModel):
                 quota_record.message_count_total += 1
                 quota_record.message_count_monthly += 1
                 quota_record.message_count_daily += 1
-                quota_record.last_message_time = datetime.now(timezone.utc)
+                quota_record.last_message_time = datetime.now(UTC)
                 db.add(quota_record)
                 await db.commit()
 
@@ -786,7 +845,7 @@ class AgentQuota(BaseModel):
                 # Update record
                 quota_record.autonomous_count_total += 1
                 quota_record.autonomous_count_monthly += 1
-                quota_record.last_autonomous_time = datetime.now(timezone.utc)
+                quota_record.last_autonomous_time = datetime.now(UTC)
                 db.add(quota_record)
                 await db.commit()
 
@@ -807,7 +866,7 @@ class AgentQuota(BaseModel):
                 # Update record
                 quota_record.twitter_count_total += 1
                 quota_record.twitter_count_daily += 1
-                quota_record.last_twitter_time = datetime.now(timezone.utc)
+                quota_record.last_twitter_time = datetime.now(UTC)
                 db.add(quota_record)
                 await db.commit()