t-youtube/backend/agent_worker.py

"""
Agent Queue Worker — polls agent_queue for pending items and
processes them against the local Qwen27B API (:8080).
"""
import asyncio
import json
import logging
import os
import aiosqlite
import httpx

import db

log = logging.getLogger(__name__)

QWEN_BASE_URL = os.environ.get("QWEN_BASE_URL", "http://localhost:8080")
QWEN_MODEL = os.environ.get("QWEN_MODEL", "current")
POLL_INTERVAL = int(os.environ.get("AGENT_POLL_INTERVAL", "30"))
MAX_CONCURRENT = int(os.environ.get("AGENT_MAX_CONCURRENT", "3"))
SYSTEM_PROMPT = os.environ.get(
    "AGENT_SYSTEM_PROMPT",
    "You are a helpful AI assistant. The user has watched a YouTube video and "
    "left you instructions. Read the video info below and respond to their "
    "prompt. Be concise and actionable."
)

semaphore = asyncio.Semaphore(MAX_CONCURRENT)


async def _call_qwen(prompt: str, video_title: str, video_context: str) -> str:
    """Call the local Qwen27B API with video context + user prompt."""
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {
            "role": "user",
            "content": (
                f"## Video Title\n{video_title}\n\n"
                f"## Video Context\n{video_context}\n\n"
                f"## User's Request\n{prompt}"
            ),
        },
    ]

    http_proxy = os.environ.get("HTTP_PROXY") or os.environ.get("http_proxy")
    https_proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("https_proxy")
    proxies = http_proxy or https_proxy

    async with httpx.AsyncClient(proxy=proxies, timeout=180) as client:
        resp = await client.post(
            f"{QWEN_BASE_URL}/v1/chat/completions",
            json={
                "model": QWEN_MODEL,
                "messages": messages,
                "max_tokens": 2048,
                "temperature": 0.3,
            },
        )
        resp.raise_for_status()
        data = resp.json()
        return data["choices"][0]["message"]["content"]


async def _process_item(item: dict) -> None:
    """Process a single agent_queue item."""
    video_id = item["video_id"]
    user_prompt = item["user_prompt"]

    async with aiosqlite.connect(db.DB_PATH) as conn:
        conn.row_factory = aiosqlite.Row

        # Mark as processing
        await conn.execute(
            "UPDATE agent_queue SET status = 'processing' WHERE id = ?",
            (item["id"],),
        )
        await conn.commit()

        # Get video info
        cursor = await conn.execute(
            "SELECT title, channel_name, summary, transcript FROM videos WHERE video_id = ?",
            (video_id,),
        )
        video = await cursor.fetchone()

    if not video:
        async with aiosqlite.connect(db.DB_PATH) as conn:
            await conn.execute(
                "UPDATE agent_queue SET status = 'failed', result = 'Video not found in DB' WHERE id = ?",
                (item["id"],),
            )
            await conn.commit()
        return

    # Build video context from whatever we have
    context_parts = []
    if video["channel_name"]:
        context_parts.append(f"Channel: {video['channel_name']}")
    if video["summary"] and not video["summary"].startswith("📺 Short clip"):
        context_parts.append(f"Summary:\n{video['summary']}")
    if video["transcript"] and video["transcript"] != "No transcript available":
        context_parts.append(f"Transcript (first 4000 chars):\n{video['transcript'][:4000]}")
    video_context = "\n\n".join(context_parts) if context_parts else "No additional context available."

    try:
        result = await _call_qwen(user_prompt, video["title"], video_context)
        status = "done"
    except Exception as e:
        log.warning(f"Qwen API call failed for item {item['id']}: {e}")
        result = f"Error: {e}"
        status = "failed"

    async with aiosqlite.connect(db.DB_PATH) as conn:
        await conn.execute(
            "UPDATE agent_queue SET status = ?, result = ?, processed_at = CURRENT_TIMESTAMP WHERE id = ?",
            (status, result, item["id"]),
        )
        await conn.commit()


async def process_queue_loop() -> None:
    """Main polling loop — runs forever."""
    log.info(
        f"Agent worker started: Qwen at {QWEN_BASE_URL}, "
        f"poll every {POLL_INTERVAL}s, max {MAX_CONCURRENT} concurrent"
    )

    while True:
        try:
            async with aiosqlite.connect(db.DB_PATH) as conn:
                conn.row_factory = aiosqlite.Row
                cursor = await conn.execute(
                    "SELECT id, video_id, user_prompt FROM agent_queue WHERE status = 'pending' ORDER BY created_at ASC LIMIT ?",
                    (MAX_CONCURRENT,),
                )
                items = await cursor.fetchall()

            if items:
                log.info(f"Agent queue: processing {len(items)} item(s)")
                tasks = [_process_item(dict(item)) for item in items]
                await asyncio.gather(*tasks)

        except Exception as e:
            log.error(f"Agent worker loop error: {e}")

        await asyncio.sleep(POLL_INTERVAL)


async def start_worker():
    """Entry point — called from main.py startup."""
    task = asyncio.create_task(process_queue_loop())
    log.info("Agent queue worker scheduled")
    return task


if __name__ == "__main__":
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
    )

    async def main():
        from db import init_db
        await init_db()
        log.info("Agent worker running in foreground. Ctrl+C to stop.")
        await process_queue_loop()

    asyncio.run(main())