agents: track indexing tasks (#676)

* plumb through agent_id

* worker can access db to persist task agent rs

* decorator works

* does this actually work?

* fix missing args

* Squashed commit of the following:

commit 2ea591d
Author: Scott <[email protected]>
Date:   Thu Aug 15 17:34:28 2024 -0400

    db changes

* wire up agent task status with controller

* return serialized json from tasks

* .

* raise job errors correctly and read back exception snippet

* get task status in client

* thanks tailwind

* link to the indexing page

* copypasta galore

* fix type errors

* lint

* revert ui change

* lint with new tools

* will this fix tests?

* fix lint

* simplify restore

* .

* clean ups

* small fixes

* fix

* feedback

* lint

* fixes

* fixes, still need to check db changes

* fe changes

* fix all

---------

Co-authored-by: Scott <[email protected]>
Co-authored-by: giannis <[email protected]>

Author: 3 people

.gitignore

@@ -156,3 +156,6 @@ secrets.toml
 /src/backend/config/configuration.yaml
 /src/backend/config/secrets.yaml
 logs/
+credentials.json
+token.json
+src/interfaces/assistants_web/bun.lockb

Makefile

@@ -78,6 +78,10 @@ win-setup:
 	poetry install --with setup --verbose
 	poetry run python src/backend/cli/main.py
 
+.PHONY: typecheck
+typecheck:
+	poetry run pyright
+
 .PHONY: lint
 lint:
 	poetry run ruff check
@@ -123,3 +127,7 @@ test-db:
 .PHONY: dev-sync
 dev-sync:
 	@docker compose up --build sync_worker sync_publisher flower -d
+
+.PHONY: dev-sync-down
+dev-sync-down:
+	@docker compose down sync_worker sync_publisher flower

pyproject.toml

@@ -116,7 +116,10 @@ requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.pyright]
-include = ["src/backend/services/metrics.py"]
+include = [
+  "src/backend/services/metrics.py", 
+  "src/backend/tools/google_drive/sync/actions/",
+  ]
 defineConstant = { DEBUG = true }
 reportMissingImports = true
 reportMissingTypeStubs = false

src/backend/alembic/env.py

@@ -27,6 +27,13 @@
 target_metadata = Base.metadata
 
 
+def include_object(object, name, type_, reflected, compare_to):
+    if type_ == "table" and reflected and compare_to is None:
+        return False
+    else:
+        return True
+
+
 def run_migrations_offline() -> None:
     """Run migrations in 'offline' mode.
 
@@ -45,6 +52,7 @@ def run_migrations_offline() -> None:
         target_metadata=target_metadata,
         literal_binds=True,
         dialect_opts={"paramstyle": "named"},
+        include_object=include_object,
     )
 
     with context.begin_transaction():
@@ -77,6 +85,7 @@ def process_revision_directives(context, revision, directives):
             connection=connection,
             target_metadata=target_metadata,
             process_revision_directives=process_revision_directives,
+            include_object=include_object,
         )
 
         with context.begin_transaction():

src/backend/alembic/versions/2024_08_22_ac3933258035_.py

@@ -0,0 +1,38 @@
+"""
+
+Revision ID: ac3933258035
+Revises: 08bcb9a24d9b
+Create Date: 2024-08-22 20:34:37.547325
+
+"""
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = 'ac3933258035'
+down_revision: Union[str, None] = '08bcb9a24d9b'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('agent_tasks',
+    sa.Column('agent_id', sa.String(), nullable=False),
+    sa.Column('task_id', sa.String(), nullable=False),
+    sa.Column('id', sa.String(), nullable=False),
+    sa.Column('created_at', sa.DateTime(), nullable=True),
+    sa.Column('updated_at', sa.DateTime(), nullable=True),
+    sa.ForeignKeyConstraint(['agent_id'], ['agents.id'], ondelete='CASCADE'),
+    sa.PrimaryKeyConstraint('id'),
+    sa.UniqueConstraint('agent_id', 'task_id', name='unique_agent_task')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('agent_tasks')
+    # ### end Alembic commands ###

src/backend/crud/agent_task.py

@@ -0,0 +1,28 @@
+from typing import List
+
+from sqlalchemy.orm import Session
+
+from backend.database_models.agent_task import AgentTask, SyncCeleryTaskMeta
+from backend.services.logger.utils import LoggerFactory
+from backend.services.transaction import validate_transaction
+
+logger = LoggerFactory().get_logger()
+
+
+@validate_transaction
+def create_agent_task(db: Session, agent_id: str, task_id: str) -> AgentTask:
+    agent_task = AgentTask(agent_id=agent_id, task_id=task_id)
+    db.add(agent_task)
+    db.commit()
+    db.refresh(agent_task)
+    return agent_task
+
+
+@validate_transaction
+def get_agent_tasks_by_agent_id(db: Session, agent_id: str) -> List[SyncCeleryTaskMeta]:
+    return (
+        db.query(SyncCeleryTaskMeta)
+        .join(AgentTask, AgentTask.task_id == SyncCeleryTaskMeta.task_id)
+        .filter(AgentTask.agent_id == agent_id)
+        .all()
+    )

src/backend/database_models/__init__.py

@@ -1,5 +1,6 @@
 # ruff: noqa
 from backend.database_models.agent import *
+from backend.database_models.agent_task import *
 from backend.database_models.agent_tool_metadata import *
 from backend.database_models.base import *
 from backend.database_models.blacklist import *

src/backend/database_models/agent_task.py

@@ -0,0 +1,43 @@
+from sqlalchemy import (
+    DateTime,
+    ForeignKey,
+    Integer,
+    LargeBinary,
+    String,
+    Text,
+    UniqueConstraint,
+)
+from sqlalchemy.orm import Mapped, mapped_column
+
+from backend.database_models.base import Base, MinimalBase
+
+
+class SyncCeleryTaskMeta(MinimalBase):
+    __tablename__ = "sync_celery_taskmeta"
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(155), unique=True)
+    status: Mapped[str] = mapped_column(String(50))
+    result: Mapped[bytes] = mapped_column(LargeBinary)
+    date_done: Mapped[DateTime] = mapped_column(DateTime)
+    traceback: Mapped[str] = mapped_column(Text)
+    name: Mapped[str] = mapped_column(String(155))
+    args: Mapped[bytes] = mapped_column(LargeBinary)
+    kwargs: Mapped[bytes] = mapped_column(LargeBinary)
+    worker: Mapped[str] = mapped_column(String(155))
+    retries: Mapped[int] = mapped_column(Integer)
+    queue: Mapped[str] = mapped_column(String(155))
+
+
+class AgentTask(Base):
+    __tablename__ = "agent_tasks"
+
+    agent_id: Mapped[str] = mapped_column(
+        ForeignKey("agents.id", ondelete="CASCADE"), nullable=False
+    )
+
+    task_id: Mapped[str] = mapped_column(nullable=False)
+
+    __table_args__ = (
+        UniqueConstraint("agent_id", "task_id", name="unique_agent_task"),
+    )

src/backend/database_models/base.py

@@ -42,6 +42,10 @@ def __new__(cls, *args, **kwargs):
         return object.__new__(cls)
 
 
+class MinimalBase(DeclarativeBase):
+    pass
+
+
 class Base(DeclarativeBase):
     id = mapped_column(
         String,

src/backend/routers/agent.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Optional
+from typing import List, Optional
 
 from fastapi import APIRouter, Depends, HTTPException
 from fastapi import File as RequestFile
@@ -11,6 +11,7 @@
 from backend.crud import agent as agent_crud
 from backend.crud import agent_tool_metadata as agent_tool_metadata_crud
 from backend.crud import snapshot as snapshot_crud
+from backend.crud.agent_task import get_agent_tasks_by_agent_id
 from backend.database_models.agent import Agent as AgentModel
 from backend.database_models.agent_tool_metadata import (
     AgentToolMetadata as AgentToolMetadataModel,
@@ -20,6 +21,7 @@
 from backend.schemas.agent import (
     Agent,
     AgentPublic,
+    AgentTaskResponse,
     AgentToolMetadata,
     AgentToolMetadataPublic,
     AgentVisibility,
@@ -40,6 +42,7 @@
     agent_to_metrics_agent,
 )
 from backend.services.agent import (
+    parse_task,
     raise_db_error,
     validate_agent_exists,
     validate_agent_tool_metadata_exists,
@@ -246,7 +249,6 @@ async def get_agent_by_id(
     agent_schema = Agent.model_validate(agent)
     ctx.with_agent(agent_schema)
     ctx.with_metrics_agent(agent_to_metrics_agent(agent))
-
     return agent
 
 
@@ -278,6 +280,23 @@ async def get_agent_deployments(
     ]
 
 
+@router.get(
+    "/{agent_id}/tasks",
+    response_model=List[AgentTaskResponse],
+    dependencies=[
+        Depends(validate_user_header),
+    ],
+)
+async def get_agent_tasks(
+    agent_id: str,
+    session: DBSessionDep,
+    ctx: Context = Depends(get_context),
+) -> List[AgentTaskResponse]:
+    raw_tasks = get_agent_tasks_by_agent_id(session, agent_id)
+    tasks = [parse_task(t) for t in raw_tasks]
+    return tasks
+
+
 @router.put(
     "/{agent_id}",
     response_model=AgentPublic,