diff --git a/alembic/versions/d169f2e66919_add_auth_header_and_timeout_to_creative_.py b/alembic/versions/d169f2e66919_add_auth_header_and_timeout_to_creative_.py new file mode 100644 index 000000000..24320f682 --- /dev/null +++ b/alembic/versions/d169f2e66919_add_auth_header_and_timeout_to_creative_.py @@ -0,0 +1,34 @@ +"""add_auth_header_and_timeout_to_creative_agents + +Revision ID: d169f2e66919 +Revises: c34b078f4e8a +Create Date: 2025-11-07 06:46:45.116343 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "d169f2e66919" +down_revision: str | Sequence[str] | None = "c34b078f4e8a" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + """Upgrade schema: Add auth_header and timeout columns to creative_agents table.""" + # Add auth_header column (nullable, matching SignalsAgent) + op.add_column("creative_agents", sa.Column("auth_header", sa.String(length=100), nullable=True)) + + # Add timeout column (not nullable, with default value) + op.add_column("creative_agents", sa.Column("timeout", sa.Integer(), nullable=False, server_default="30")) + + +def downgrade() -> None: + """Downgrade schema: Remove auth_header and timeout columns from creative_agents table.""" + op.drop_column("creative_agents", "timeout") + op.drop_column("creative_agents", "auth_header") diff --git a/src/core/creative_agent_registry.py b/src/core/creative_agent_registry.py index 5928abe38..8f3afd370 100644 --- a/src/core/creative_agent_registry.py +++ b/src/core/creative_agent_registry.py @@ -31,6 +31,7 @@ class CreativeAgent: enabled: bool = True priority: int = 1 # Lower = higher priority in search results auth: dict[str, Any] | None = None # Optional auth config for private agents + timeout: int = 30 # Request timeout in seconds @dataclass @@ -108,6 +109,9 @@ def _get_tenant_agents(self, tenant_id: str | None) -> list[CreativeAgent]: "type": db_agent.auth_type, "credentials": db_agent.auth_credentials, } + # Add auth_header if present (e.g., "Authorization", "x-api-key") + if db_agent.auth_header: + auth["header"] = db_agent.auth_header agents.append( CreativeAgent( @@ -116,6 +120,7 @@ def _get_tenant_agents(self, tenant_id: str | None) -> list[CreativeAgent]: enabled=db_agent.enabled, priority=db_agent.priority, auth=auth, + timeout=db_agent.timeout, ) ) @@ -160,7 +165,7 @@ async def _fetch_formats_from_agent( async with create_mcp_client( agent_url=agent.agent_url, auth=agent.auth, - timeout=30, + timeout=agent.timeout, max_retries=3, ) as client: # Build parameters for list_creative_formats diff --git a/src/core/database/models.py b/src/core/database/models.py index df25902dd..334c72478 100644 --- a/src/core/database/models.py +++ b/src/core/database/models.py @@ -717,7 +717,9 @@ class CreativeAgent(Base): enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True) priority: Mapped[int] = mapped_column(Integer, nullable=False, default=10) auth_type: Mapped[str | None] = mapped_column(String(50), nullable=True) + auth_header: Mapped[str | None] = mapped_column(String(100), nullable=True) auth_credentials: Mapped[str | None] = mapped_column(Text, nullable=True) + timeout: Mapped[int] = mapped_column(Integer, nullable=False, default=30) created_at: Mapped[DateTime] = mapped_column(DateTime, server_default=func.now()) updated_at: Mapped[DateTime] = mapped_column(DateTime, server_default=func.now(), onupdate=func.now())