Timmy-time-dashboard/src/timmy/paperclip.py

"""Paperclip integration for Timmy.

This module provides a client for the Paperclip API, and a poller for
running research tasks.
"""

from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass

import httpx

from config import settings
from timmy.research_tools import get_llm_client, google_web_search
from timmy.research_triage import triage_research_report

logger = logging.getLogger(__name__)


@dataclass
class PaperclipTask:
    """A task from the Paperclip API."""

    id: str
    kind: str
    context: dict


class PaperclipClient:
    """A client for the Paperclip API."""

    def __init__(self) -> None:
        self.base_url = settings.paperclip_url
        self.api_key = settings.paperclip_api_key
        self.agent_id = settings.paperclip_agent_id
        self.company_id = settings.paperclip_company_id
        self.timeout = settings.paperclip_timeout

    async def get_tasks(self) -> list[PaperclipTask]:
        """Get a list of tasks from the Paperclip API."""
        async with httpx.AsyncClient(timeout=self.timeout) as client:
            resp = await client.get(
                f"{self.base_url}/api/tasks",
                headers={"Authorization": f"Bearer {self.api_key}"},
                params={
                    "agent_id": self.agent_id,
                    "company_id": self.company_id,
                    "status": "queued",
                },
            )
            resp.raise_for_status()
            tasks = resp.json()
            return [
                PaperclipTask(id=t["id"], kind=t["kind"], context=t["context"])
                for t in tasks
            ]

    async def update_task_status(
        self, task_id: str, status: str, result: str | None = None
    ) -> None:
        """Update the status of a task."""
        async with httpx.AsyncClient(timeout=self.timeout) as client:
            await client.patch(
                f"{self.base_url}/api/tasks/{task_id}",
                headers={"Authorization": f"Bearer {self.api_key}"},
                json={"status": status, "result": result},
            )


class ResearchOrchestrator:
    """Orchestrates research tasks."""

    async def get_gitea_issue(self, issue_number: int) -> dict:
        """Get a Gitea issue by its number."""
        owner, repo = settings.gitea_repo.split("/", 1)
        api_url = f"{settings.gitea_url}/api/v1/repos/{owner}/{repo}/issues/{issue_number}"
        async with httpx.AsyncClient(timeout=15) as client:
            resp = await client.get(
                api_url,
                headers={"Authorization": f"token {settings.gitea_token}"},
            )
            resp.raise_for_status()
            return resp.json()

    async def post_gitea_comment(self, issue_number: int, comment: str) -> None:
        """Post a comment to a Gitea issue."""
        owner, repo = settings.gitea_repo.split("/", 1)
        api_url = f"{settings.gitea_url}/api/v1/repos/{owner}/{repo}/issues/{issue_number}/comments"
        async with httpx.AsyncClient(timeout=15) as client:
            await client.post(
                api_url,
                headers={"Authorization": f"token {settings.gitea_token}"},
                json={"body": comment},
            )

    async def run_research_pipeline(self, issue_title: str) -> str:
        """Run the research pipeline."""
        search_results = await google_web_search(issue_title)
        
        llm_client = get_llm_client()
        response = await llm_client.completion(
            f"Summarize the following search results and generate a research report:\\n\\n{search_results}",
            max_tokens=2048,
        )
        return response.text

    async def run(self, context: dict) -> str:
        """Run a research task."""
        issue_number = context.get("issue_number")
        if not issue_number:
            return "Missing issue_number in task context"

        issue = await self.get_gitea_issue(issue_number)

        report = await self.run_research_pipeline(issue["title"])

        triage_results = await triage_research_report(report, source_issue=issue_number)

        comment = f"Research complete for issue #{issue_number}.\\n\\n"
        if triage_results:
            comment += "Created the following issues:\\n"
            for result in triage_results:
                if result["gitea_issue"]:
                    comment += f"- #{result['gitea_issue']['number']}: {result['action_item'].title}\\n"
        else:
            comment += "No new issues were created.\\n"

        await self.post_gitea_comment(issue_number, comment)

        return f"Research complete for issue #{issue_number}"


class PaperclipPoller:
    """Polls the Paperclip API for new tasks."""

    def __init__(self) -> None:
        self.client = PaperclipClient()
        self.orchestrator = ResearchOrchestrator()
        self.poll_interval = settings.paperclip_poll_interval

    async def poll(self) -> None:
        """Poll the Paperclip API for new tasks."""
        if self.poll_interval == 0:
            return

        while True:
            try:
                tasks = await self.client.get_tasks()
                for task in tasks:
                    if task.kind == "research":
                        await self.run_research_task(task)
            except httpx.HTTPError as exc:
                logger.warning("Error polling Paperclip: %s", exc)

            await asyncio.sleep(self.poll_interval)

    async def run_research_task(self, task: PaperclipTask) -> None:
        """Run a research task."""
        await self.client.update_task_status(task.id, "running")
        try:
            result = await self.orchestrator.run(task.context)
            await self.client.update_task_status(task.id, "completed", result)
        except Exception as exc:
            logger.error("Error running research task: %s", exc, exc_info=True)
            await self.client.update_task_status(task.id, "failed", str(exc))


async def start_paperclip_poller() -> None:
    """Start the Paperclip poller."""
    if settings.paperclip_enabled:
        poller = PaperclipPoller()
        asyncio.create_task(poller.poll())