feat(router): add field for agents

webisopen · Jan 24, 2025 · bd0d1ed · bd0d1ed
1 parent 211d3d9
commit bd0d1ed
Show file tree

Hide file tree

Showing 6 changed files with 30 additions and 65 deletions.
diff --git a/migrations/versions/215e5a803b40_create_initial_tables.py b/migrations/versions/215e5a803b40_create_initial_tables.py
@@ -120,6 +120,12 @@ def upgrade() -> None:
             postgresql.ENUM("inactive", "active", name="agent_status"),
             nullable=False,
         ),
+        sa.Column(
+            "type",
+            postgresql.ENUM("IP", "DeFi", "DeSci", "Others", name="agent_type"),
+            nullable=False,
+            index=True,
+        ),
         sa.Column(
             "created_at",
             postgresql.TIMESTAMP(timezone=True),

diff --git a/openagent/database/migrations/env.py b/openagent/database/migrations/env.py
diff --git a/openagent/database/models/agent.py b/openagent/database/models/agent.py
@@ -14,6 +14,16 @@ def __str__(self):
         return self.value
 
 
+class AgentType(str, enum.Enum):
+    IP = "IP"
+    DEFI = "DeFi"
+    DESCI = "DeSci"
+    OTHERS = "Others"
+
+    def __str__(self):
+        return self.value
+
+
 class Agent(Base):
     __tablename__ = "agents"
 
@@ -35,6 +45,10 @@ class Agent(Base):
         Enum(AgentStatus, values_callable=lambda x: [e.value for e in x]),
         nullable=False,
     )
+    type = Column(
+        Enum(AgentType, values_callable=lambda x: [e.value for e in x]),
+        nullable=False,
+    )
     created_at = Column(DateTime, default=lambda: datetime.now(UTC), nullable=False)
     updated_at = Column(
         DateTime,
@@ -44,6 +58,8 @@ class Agent(Base):
     )
 
     def __init__(self, *args, **kwargs):
+        if "type" in kwargs and isinstance(kwargs["type"], AgentType):
+            kwargs["type"] = kwargs["type"].value
         if "status" in kwargs and isinstance(kwargs["status"], AgentStatus):
             kwargs["status"] = kwargs["status"].value
         super().__init__(*args, **kwargs)

diff --git a/openagent/router/routes/agent.py b/openagent/router/routes/agent.py
@@ -104,6 +104,7 @@ def create_agent(
             telegram=request.telegram,
             website=request.website,
             tool_configs=request.get_tool_configs_data(),
+            type=request.type,
             status=AgentStatus.INACTIVE,  # default status
         )
 

diff --git a/openagent/router/routes/models/request.py b/openagent/router/routes/models/request.py
@@ -1,3 +1,4 @@
+from openagent.database.models.agent import AgentType
 from pydantic import BaseModel
 
 from openagent.tools import ToolConfig
@@ -16,6 +17,7 @@ class CreateAgentRequest(BaseModel):
     twitter: str | None = None
     telegram: str | None = None
     website: str | None = None
+    type: AgentType
     tool_configs: list[ToolConfig] | None = None
 
     def get_tool_configs_data(self) -> list[dict]:

diff --git a/openagent/router/routes/models/response.py b/openagent/router/routes/models/response.py
@@ -1,8 +1,9 @@
+from datetime import datetime
 from typing import Generic, TypeVar
 
 from pydantic import BaseModel, ConfigDict
 
-from openagent.database.models.agent import AgentStatus
+from openagent.database.models.agent import AgentStatus, AgentType
 from openagent.database.models.tool import ToolType
 from openagent.tools import ToolConfig
 
@@ -32,7 +33,10 @@ class AgentResponse(BaseModel):
     telegram: str | None = None
     website: str | None = None
     tool_configs: list[ToolConfig] | None = None
+    type: AgentType
     status: AgentStatus
+    created_at: datetime
+    updated_at: datetime
 
 
 class AgentListResponse(BaseModel):