forkflux-api 0.1.0__py3-none-any.whl

forkflux_api/agents/dependencies.py

@@ -0,0 +1,21 @@
+from fastapi import Depends, Request
+from forkflux_api.agents.respositories import TargetRoleRepository
+from forkflux_api.agents.services import TargetRoleService
+from forkflux_api.database import get_session
+from sqlalchemy.ext.asyncio import AsyncSession
+
+
+def get_trace_id(request: Request) -> str:
+    return request.state.trace_id
+
+
+def get_target_role_repo(
+    session: AsyncSession = Depends(get_session), trace_id: str = Depends(get_trace_id)
+) -> TargetRoleRepository:
+    return TargetRoleRepository(session=session, trace_id=trace_id)
+
+
+def get_target_role_service(
+    repository: TargetRoleRepository = Depends(get_target_role_repo), trace_id: str = Depends(get_trace_id)
+) -> TargetRoleService:
+    return TargetRoleService(target_role_repo=repository, trace_id=trace_id)

forkflux_api/agents/dto.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+
+
+@dataclass(slots=True)
+class TargetRoleCreate:
+    role_key: str
+    role_label: str
+
+
+@dataclass(slots=True)
+class AgentIdentityCreate:
+    agent_label: str
+    role_id: int
+    tool_family: str | None
+
+
+@dataclass(slots=True)
+class AgentApiTokenCreate:
+    agent_id: int

forkflux_api/agents/exceptions.py

@@ -0,0 +1,28 @@
+class AgentApiTokenNotFoundError(Exception):
+    code = "agent_api_token.not_found"
+    msg = "Agent API token not found."
+
+
+class AgentApiTokenConflictError(Exception):
+    code = "agent_api_token.conflict"
+    msg = "Agent API token conflicts with existing data constraints."
+
+
+class AgentIdentityNotFoundError(Exception):
+    code = "agent_identity.not_found"
+    msg = "Agent identity not found."
+
+
+class AgentIdentityConflictError(Exception):
+    code = "agent_identity.conflict"
+    msg = "Agent identity conflicts with existing data constraints."
+
+
+class TargetRoleConflictError(Exception):
+    code = "target_role.conflict"
+    msg = "Target role already exists."
+
+
+class TargetRoleNotFoundError(Exception):
+    code = "target_role.not_found"
+    msg = "Target role not found."

forkflux_api/agents/handlers.py

@@ -0,0 +1,19 @@
+from fastapi import APIRouter, Depends
+from forkflux_api.agents.dependencies import get_target_role_service
+from forkflux_api.agents.models import AgentIdentity
+from forkflux_api.agents.schemas import GetMeResponse, ListRolesResponse
+from forkflux_api.agents.services import TargetRoleService
+from forkflux_api.dependencies import get_current_agent, verify_token
+
+router = APIRouter(prefix="/agents", tags=["agents"], dependencies=[Depends(verify_token)])
+
+
+@router.get("/roles", response_model=list[ListRolesResponse])
+async def list_roles(service: TargetRoleService = Depends(get_target_role_service)):
+    roles = await service.get_all_roles()
+    return roles
+
+
+@router.get("/me", response_model=GetMeResponse)
+async def get_me(current_agent: AgentIdentity = Depends(get_current_agent)):
+    return current_agent

forkflux_api/agents/models.py

@@ -0,0 +1,38 @@
+from datetime import datetime
+
+from forkflux_api.database import Base, UTCDateTime
+from sqlalchemy import BigInteger, Boolean, ForeignKey, Integer, Text
+from sqlalchemy.orm import Mapped, mapped_column
+
+PK_TYPE = BigInteger().with_variant(Integer, "sqlite")
+
+
+class TargetRole(Base):
+    __tablename__ = "target_role"
+
+    id: Mapped[int] = mapped_column(PK_TYPE, primary_key=True, autoincrement=True)
+    role_key: Mapped[str] = mapped_column(Text, nullable=False, unique=True)
+    role_label: Mapped[str] = mapped_column(Text, nullable=False)
+    created_at: Mapped[datetime] = mapped_column(UTCDateTime(), nullable=False)
+
+
+class AgentIdentity(Base):
+    __tablename__ = "agent_identity"
+
+    id: Mapped[int] = mapped_column(PK_TYPE, primary_key=True, autoincrement=True)
+    agent_label: Mapped[str] = mapped_column(Text, nullable=False)
+    role_id: Mapped[int] = mapped_column(ForeignKey("target_role.id"), nullable=False)
+    tool_family: Mapped[str | None] = mapped_column(Text, nullable=True)
+    created_at: Mapped[datetime] = mapped_column(UTCDateTime(), nullable=False)
+
+
+class AgentApiToken(Base):
+    __tablename__ = "agent_api_token"
+
+    id: Mapped[int] = mapped_column(PK_TYPE, primary_key=True, autoincrement=True)
+    token_hash: Mapped[str] = mapped_column(Text, nullable=False, unique=True)
+    agent_id: Mapped[int] = mapped_column(ForeignKey("agent_identity.id"), nullable=False)
+    is_active: Mapped[bool] = mapped_column(Boolean, nullable=False)
+    created_at: Mapped[datetime] = mapped_column(UTCDateTime(), nullable=False)
+    last_used_at: Mapped[datetime | None] = mapped_column(UTCDateTime(), nullable=True)
+    revoked_at: Mapped[datetime | None] = mapped_column(UTCDateTime(), nullable=True)

forkflux_api/agents/respositories.py

@@ -0,0 +1,150 @@
+from datetime import datetime, timezone
+
+import structlog
+from forkflux_api.agents.dto import AgentApiTokenCreate, AgentIdentityCreate, TargetRoleCreate
+from forkflux_api.agents.exceptions import (
+    AgentApiTokenConflictError,
+    AgentApiTokenNotFoundError,
+    AgentIdentityConflictError,
+    AgentIdentityNotFoundError,
+    TargetRoleConflictError,
+    TargetRoleNotFoundError,
+)
+from forkflux_api.agents.models import AgentApiToken, AgentIdentity, TargetRole
+from sqlalchemy import exists, select, update
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.ext.asyncio import AsyncSession
+
+
+class TargetRoleRepository:
+    def __init__(self, session: AsyncSession, trace_id: str) -> None:
+        self._session = session
+        self._logger = structlog.get_logger().bind(cls=self.__class__.__name__, trace_id=trace_id)
+
+    async def list(self) -> list[TargetRole]:
+        result = await self._session.execute(select(TargetRole))
+        return list(result.scalars().all())
+
+    async def get_by_role_key(self, role_key: str) -> TargetRole:
+        result = await self._session.execute(select(TargetRole).where(TargetRole.role_key == role_key))
+        target_role = result.scalar_one_or_none()
+        if target_role is None:
+            raise TargetRoleNotFoundError
+
+        return target_role
+
+    async def exists(self, role_key: str) -> bool:
+        log = self._logger.bind(method="exists", role_key=role_key)
+        result = await self._session.execute(select(exists().where(TargetRole.role_key == role_key)))
+        role_exists = result.scalar_one()
+
+        if role_exists:
+            log.info("target_role_exists_hit")
+        else:
+            log.info("target_role_exists_miss")
+
+        return role_exists
+
+    async def create(self, dto: TargetRoleCreate) -> TargetRole:
+        target_role = TargetRole(
+            role_key=dto.role_key,
+            role_label=dto.role_label,
+            created_at=datetime.now(timezone.utc),
+        )
+
+        self._session.add(target_role)
+        try:
+            await self._session.flush()
+        except IntegrityError as err:
+            await self._session.rollback()
+            raise TargetRoleConflictError from err
+
+        return target_role
+
+
+class AgentApiTokenRepository:
+    def __init__(self, session: AsyncSession, trace_id: str) -> None:
+        self._session = session
+        self._logger = structlog.get_logger().bind(cls=self.__class__.__name__, trace_id=trace_id)
+
+    async def get(self, token_hash: str) -> AgentApiToken:
+        result = await self._session.execute(
+            select(AgentApiToken).where(
+                AgentApiToken.token_hash == token_hash,
+                AgentApiToken.is_active.is_(True),
+            )
+        )
+        token = result.scalar_one_or_none()
+        if token is None:
+            raise AgentApiTokenNotFoundError
+
+        return token
+
+    async def create(self, dto: AgentApiTokenCreate, token_hash: str) -> AgentApiToken:
+        agent_api_token = AgentApiToken(
+            token_hash=token_hash,
+            agent_id=dto.agent_id,
+            is_active=True,
+            created_at=datetime.now(timezone.utc),
+        )
+
+        self._session.add(agent_api_token)
+        try:
+            await self._session.flush()
+        except IntegrityError as err:
+            await self._session.rollback()
+            raise AgentApiTokenConflictError from err
+
+        return agent_api_token
+
+    async def revoke(self, agent_id: int) -> int:
+        revoked_at = datetime.now(timezone.utc)
+        result = await self._session.execute(
+            update(AgentApiToken)
+            .where(
+                AgentApiToken.agent_id == agent_id,
+                AgentApiToken.is_active.is_(True),
+            )
+            .values(
+                is_active=False,
+                revoked_at=revoked_at,
+            )
+        )
+        await self._session.flush()
+
+        return result.rowcount or 0  # type: ignore[attr-defined]
+
+
+class AgentIdentityRepository:
+    def __init__(self, session: AsyncSession, trace_id: str) -> None:
+        self._session = session
+        self._logger = structlog.get_logger().bind(cls=self.__class__.__name__, trace_id=trace_id)
+
+    async def list(self) -> list[AgentIdentity]:
+        result = await self._session.execute(select(AgentIdentity))
+        return list(result.scalars().all())
+
+    async def get_by_id(self, agent_identity_id: int) -> AgentIdentity:
+        result = await self._session.execute(select(AgentIdentity).where(AgentIdentity.id == agent_identity_id))
+        agent_identity = result.scalar_one_or_none()
+        if agent_identity is None:
+            raise AgentIdentityNotFoundError
+
+        return agent_identity
+
+    async def create(self, dto: AgentIdentityCreate) -> AgentIdentity:
+        agent_identity = AgentIdentity(
+            agent_label=dto.agent_label,
+            role_id=dto.role_id,
+            tool_family=dto.tool_family,
+            created_at=datetime.now(timezone.utc),
+        )
+
+        self._session.add(agent_identity)
+        try:
+            await self._session.flush()
+        except IntegrityError as err:
+            await self._session.rollback()
+            raise AgentIdentityConflictError from err
+
+        return agent_identity

forkflux_api/agents/schemas.py

@@ -0,0 +1,17 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class ListRolesResponse(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+
+    role_key: str
+    role_label: str
+
+
+class GetMeResponse(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+
+    id: int
+    agent_label: str
+    role_id: int
+    tool_family: str | None

forkflux_api/agents/services.py

@@ -0,0 +1,117 @@
+import hashlib
+import secrets
+
+import structlog
+from forkflux_api.agents.dto import AgentApiTokenCreate, AgentIdentityCreate, TargetRoleCreate
+from forkflux_api.agents.models import AgentApiToken, AgentIdentity, TargetRole
+from forkflux_api.agents.respositories import AgentApiTokenRepository, AgentIdentityRepository, TargetRoleRepository
+
+
+class TargetRoleService:
+    def __init__(self, target_role_repo: TargetRoleRepository, trace_id: str) -> None:
+        self._logger = structlog.get_logger().bind(cls=self.__class__.__name__, trace_id=trace_id)
+        self._target_role_repo = target_role_repo
+
+    async def get_all_roles(self) -> list[TargetRole]:
+        log = self._logger.bind(method="get_all_roles")
+        log.info("operation_started")
+
+        roles = await self._target_role_repo.list()
+
+        log.info("operation_completed", roles_count=len(roles))
+        return roles
+
+    async def get_by_role_key(self, role_key: str) -> TargetRole:
+        log = self._logger.bind(method="get_by_role_key", role_key=role_key)
+        log.info("operation_started")
+
+        role = await self._target_role_repo.get_by_role_key(role_key)
+
+        log.info("operation_completed")
+        return role
+
+    async def is_role_exists(self, role_key: str) -> bool:
+        log = self._logger.bind(method="is_role_exists", role_key=role_key)
+        log.info("operation_started")
+
+        exists = await self._target_role_repo.exists(role_key)
+
+        log.info("operation_completed", role_exists=exists)
+        return exists
+
+    async def create_role(self, dto: TargetRoleCreate) -> TargetRole:
+        log = self._logger.bind(method="create_role", role_key=dto.role_key)
+        log.info("operation_started")
+
+        role = await self._target_role_repo.create(dto)
+
+        log.info("operation_completed")
+        return role
+
+
+class AgentApiTokenService:
+    def __init__(self, agent_api_token_repo: AgentApiTokenRepository, trace_id: str) -> None:
+        self._logger = structlog.get_logger().bind(cls=self.__class__.__name__, trace_id=trace_id)
+        self._agent_api_token_repo = agent_api_token_repo
+
+    async def get_token(self, token_hash: str) -> AgentApiToken:
+        log = self._logger.bind(method="get_token", token_hash=token_hash)
+        log.info("operation_started")
+
+        token = await self._agent_api_token_repo.get(token_hash)
+
+        log.info("operation_completed", token_id=token.id, agent_id=token.agent_id)
+        return token
+
+    async def create_token(self, dto: AgentApiTokenCreate) -> str:
+        log = self._logger.bind(method="create_token", agent_id=dto.agent_id)
+        log.info("operation_started")
+
+        raw_token = secrets.token_urlsafe(32)
+        token_hash = hashlib.sha256(raw_token.encode()).hexdigest()
+        await self._agent_api_token_repo.create(dto=dto, token_hash=token_hash)
+
+        log.info("operation_completed")
+        return raw_token
+
+    async def revoke_token(self, agent_id: int) -> int:
+        log = self._logger.bind(method="revoke_token", agent_id=agent_id)
+        log.info("operation_started")
+
+        revoked_count = await self._agent_api_token_repo.revoke(agent_id)
+
+        log.info("operation_completed", revoked_count=revoked_count)
+        return revoked_count
+
+
+class AgentIdentityService:
+    def __init__(self, agent_identity_repo: AgentIdentityRepository, trace_id: str) -> None:
+        self._logger = structlog.get_logger().bind(cls=self.__class__.__name__, trace_id=trace_id)
+        self._agent_identity_repo = agent_identity_repo
+
+    async def get_all_agents(self) -> list[AgentIdentity]:
+        log = self._logger.bind(method="get_all_agents")
+        log.info("operation_started")
+
+        agents = await self._agent_identity_repo.list()
+
+        log.info("operation_completed", agents_count=len(agents))
+        return agents
+
+    async def get_by_id(self, agent_identity_id: int) -> AgentIdentity:
+        log = self._logger.bind(method="get_by_id", agent_identity_id=agent_identity_id)
+        log.info("operation_started")
+
+        agent = await self._agent_identity_repo.get_by_id(agent_identity_id)
+
+        log.info("operation_completed")
+        return agent
+
+    async def create_agent(self, dto: AgentIdentityCreate) -> AgentIdentity:
+        log = self._logger.bind(method="create_agent", agent_label=dto.agent_label, role_id=dto.role_id)
+        log.info("operation_started")
+
+        agent = await self._agent_identity_repo.create(dto)
+
+        log.info("operation_completed", agent_identity_id=agent.id)
+        return agent

forkflux_api/cli.py

@@ -0,0 +1,208 @@
+import asyncio
+import logging
+import os
+import pathlib
+import sys
+from functools import wraps
+from uuid import uuid4
+
+import structlog
+import typer
+import uvicorn
+from alembic import command
+from alembic.config import Config
+from rich.console import Console
+from rich.table import Table
+
+sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))
+
+from forkflux_api.agents.dto import AgentApiTokenCreate, AgentIdentityCreate, TargetRoleCreate
+from forkflux_api.agents.exceptions import (
+    AgentApiTokenConflictError,
+    AgentIdentityConflictError,
+    TargetRoleConflictError,
+    TargetRoleNotFoundError,
+)
+from forkflux_api.agents.respositories import AgentApiTokenRepository, AgentIdentityRepository, TargetRoleRepository
+from forkflux_api.agents.services import AgentApiTokenService, AgentIdentityService, TargetRoleService
+from forkflux_api.database import session_manager
+
+app = typer.Typer(help="ForkFlux Management CLI")
+console = Console()
+
+_CLI_LOGGING_CONFIGURED = False
+
+agents_role_app = typer.Typer(help="Agents role management")
+agent_app = typer.Typer(help="Agents management")
+
+app.add_typer(agents_role_app, name="agents-role")
+app.add_typer(agent_app, name="agent")
+
+
+def _configure_cli_logging() -> None:
+    """Suppress INFO logs for CLI-invoked services, keep WARNING/ERROR visible."""
+    global _CLI_LOGGING_CONFIGURED
+
+    if _CLI_LOGGING_CONFIGURED:
+        return
+
+    logging.basicConfig(level=logging.WARNING, force=True)
+    structlog.configure(wrapper_class=structlog.make_filtering_bound_logger(logging.WARNING))
+
+    _CLI_LOGGING_CONFIGURED = True
+
+
+def apply_migrations() -> None:
+    console.print("Apply database migrations")
+    current_dir = os.path.dirname(__file__)
+    alembic_cfg = Config(toml_file="../pyproject.toml")
+    alembic_cfg.set_main_option("script_location", os.path.join(current_dir, "migrations"))
+    command.upgrade(alembic_cfg, "head")
+
+
+@app.command(help="Run the server")
+def serve(host: str = "0.0.0.0", port: int = 8080) -> None:  # noqa: S104
+    apply_migrations()
+
+    console.print("Starting server...", style="bold green")
+    uvicorn.run(
+        "forkflux_api.main:app",
+        host=host,
+        port=port,
+        forwarded_allow_ips="*",
+        workers=2,
+        loop="none" if sys.platform == "win32" else "auto",
+    )
+
+
+@app.command(help="Initialize the database and add some example data")
+def init() -> None:
+    apply_migrations()
+
+    asyncio.run(_init_async())
+
+
+async def _init_async() -> None:
+    _configure_cli_logging()
+
+    console.print("Lets add 2 roles - developer and QA")
+    await add_role.__wrapped__(role_key="developer", role_label="Developer")
+    await add_role.__wrapped__(role_key="qa", role_label="QA")
+
+    console.print("Lets add 2 agents - agent-1 and agent-2")
+    await add_agent.__wrapped__(agent_label="agent-1", role_key="developer")
+    await add_agent.__wrapped__(agent_label="agent-2", role_key="qa")
+
+
+@agents_role_app.command("list")
+@lambda f: wraps(f)(lambda *a, **kw: asyncio.run(f(*a, **kw)))
+async def list_roles() -> None:
+    _configure_cli_logging()
+    trace_id = str(uuid4())
+
+    async with session_manager() as session:
+        repo = TargetRoleRepository(session=session, trace_id=trace_id)
+        roles = await TargetRoleService(target_role_repo=repo, trace_id=trace_id).get_all_roles()
+
+    table = Table("Key", "Label")
+    for role in roles:
+        table.add_row(role.role_key, role.role_label)
+    console.print(table)
+
+
+@agents_role_app.command("add")
+@lambda f: wraps(f)(lambda *a, **kw: asyncio.run(f(*a, **kw)))
+async def add_role(role_key: str, role_label: str) -> None:
+    """
+    Adds a new role with the specified key and label.
+    """
+    _configure_cli_logging()
+    trace_id = str(uuid4())
+
+    async with session_manager() as session:
+        try:
+            repo = TargetRoleRepository(session=session, trace_id=trace_id)
+            dto = TargetRoleCreate(role_key=role_key, role_label=role_label)
+            new_role = await TargetRoleService(target_role_repo=repo, trace_id=trace_id).create_role(dto=dto)
+            console.print(f"Role {new_role.role_key} created successfully")
+        except TargetRoleConflictError:
+            console.print(f"Role with key {role_key} already exists", style="bold red")
+
+
+@agent_app.command("list")
+@lambda f: wraps(f)(lambda *a, **kw: asyncio.run(f(*a, **kw)))
+async def list_agents() -> None:
+    _configure_cli_logging()
+    trace_id = str(uuid4())
+
+    async with session_manager() as session:
+        agent_repo = AgentIdentityRepository(session=session, trace_id=trace_id)
+        agents = await AgentIdentityService(agent_identity_repo=agent_repo, trace_id=trace_id).get_all_agents()
+        role_repo = TargetRoleRepository(session=session, trace_id=trace_id)
+        roles = await TargetRoleService(target_role_repo=role_repo, trace_id=trace_id).get_all_roles()
+
+    roles_mapping = {role.id: role.role_key for role in roles}
+
+    table = Table("ID", "Label", "Role key")
+    for agent in agents:
+        table.add_row(str(agent.id), agent.agent_label, roles_mapping[agent.role_id])
+    console.print(table)
+
+
+@agent_app.command("add")
+@lambda f: wraps(f)(lambda *a, **kw: asyncio.run(f(*a, **kw)))
+async def add_agent(agent_label: str, role_key: str, tool_family: str | None = None) -> None:
+    """
+    Adds a new agent with the specified label, role, and tool family (optional).
+    """
+    _configure_cli_logging()
+    trace_id = str(uuid4())
+
+    async with session_manager() as session:
+        try:
+            role_repo = TargetRoleRepository(session=session, trace_id=trace_id)
+            role = await TargetRoleService(target_role_repo=role_repo, trace_id=trace_id).get_by_role_key(role_key)
+        except TargetRoleNotFoundError:
+            console.print(f"Role with key {role_key} not found", style="bold red")
+            return
+
+        try:
+            agent_repo = AgentIdentityRepository(session=session, trace_id=trace_id)
+            agent_dto = AgentIdentityCreate(agent_label=agent_label, role_id=role.id, tool_family=tool_family)
+            new_agent = await AgentIdentityService(agent_identity_repo=agent_repo, trace_id=trace_id).create_agent(
+                dto=agent_dto
+            )
+            console.print(f"Agent {new_agent.agent_label} created successfully")
+        except AgentIdentityConflictError:
+            console.print("Can't create new agent", style="bold red")
+            return
+
+        try:
+            token_repo = AgentApiTokenRepository(session=session, trace_id=trace_id)
+            token_dto = AgentApiTokenCreate(agent_id=new_agent.id)
+            new_token = await AgentApiTokenService(agent_api_token_repo=token_repo, trace_id=trace_id).create_token(
+                dto=token_dto
+            )
+            console.print(f"Token {new_token} for agent {new_agent.agent_label} created successfully")
+        except AgentApiTokenConflictError:
+            console.print("Can't create new token", style="bold red")
+            return
+
+
+@agent_app.command("revoke-token")
+@lambda f: wraps(f)(lambda *a, **kw: asyncio.run(f(*a, **kw)))
+async def agent_revoke_token(agent_id: int) -> None:
+    """
+    Revokes the token associated with a specified agent.
+    """
+    _configure_cli_logging()
+    trace_id = str(uuid4())
+
+    async with session_manager() as session:
+        token_repo = AgentApiTokenRepository(session=session, trace_id=trace_id)
+        await AgentApiTokenService(agent_api_token_repo=token_repo, trace_id=trace_id).revoke_token(agent_id=agent_id)
+        console.print(f"Token for agent {agent_id} revoked successfully")
+
+
+if __name__ == "__main__":
+    app()