76 lines
2.5 KiB
Python
76 lines
2.5 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from typing import Any, cast
|
||
|
|
from uuid import UUID
|
||
|
|
|
||
|
|
from fastapi import Depends
|
||
|
|
import redis.asyncio as redis
|
||
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
|
|
||
|
|
from core.agent.infrastructure.events.redis_stream import RedisStreamEventStore
|
||
|
|
from core.agent.infrastructure.queue.tasks import run_command_task
|
||
|
|
from core.config.settings import config
|
||
|
|
from core.db import get_db
|
||
|
|
from v1.agent.repository import AgentRepository
|
||
|
|
from v1.agent.service import AgentService
|
||
|
|
|
||
|
|
|
||
|
|
class CeleryQueueClient:
|
||
|
|
def __init__(self) -> None:
|
||
|
|
settings = cast(Any, config)
|
||
|
|
self._redis = redis.from_url(settings.redis.url, decode_responses=True)
|
||
|
|
|
||
|
|
async def enqueue(
|
||
|
|
self, *, command: dict[str, object], dedup_key: str | None
|
||
|
|
) -> str:
|
||
|
|
redis_key = None
|
||
|
|
if dedup_key:
|
||
|
|
redis_key = f"agent:dedup:{dedup_key}"
|
||
|
|
locked = await self._redis.set(redis_key, "__inflight__", nx=True, ex=300)
|
||
|
|
if not locked:
|
||
|
|
existing = await self._redis.get(redis_key)
|
||
|
|
if existing and existing != "__inflight__":
|
||
|
|
return existing
|
||
|
|
|
||
|
|
payload = dict(command)
|
||
|
|
if dedup_key:
|
||
|
|
payload["dedup_key"] = dedup_key
|
||
|
|
delay = getattr(run_command_task, "delay")
|
||
|
|
result = delay(payload)
|
||
|
|
task_id = str(result.id)
|
||
|
|
if redis_key is not None:
|
||
|
|
await self._redis.set(redis_key, task_id, ex=300)
|
||
|
|
return task_id
|
||
|
|
|
||
|
|
|
||
|
|
class RedisEventStream:
|
||
|
|
def __init__(self) -> None:
|
||
|
|
settings = cast(Any, config)
|
||
|
|
client = redis.from_url(settings.redis.url, decode_responses=True)
|
||
|
|
self._store = RedisStreamEventStore(
|
||
|
|
client=client,
|
||
|
|
stream_prefix=settings.agent_runtime.redis_stream_prefix,
|
||
|
|
read_count=settings.agent_runtime.redis_stream_read_count,
|
||
|
|
block_ms=settings.agent_runtime.redis_stream_block_ms,
|
||
|
|
)
|
||
|
|
|
||
|
|
async def read(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
session_id: str,
|
||
|
|
last_event_id: str | None,
|
||
|
|
) -> list[dict[str, Any]]:
|
||
|
|
rows = await self._store.read_events(
|
||
|
|
session_id=UUID(session_id),
|
||
|
|
last_event_id=last_event_id,
|
||
|
|
)
|
||
|
|
return [{**row, "cursor": last_event_id} for row in rows]
|
||
|
|
|
||
|
|
|
||
|
|
def get_agent_service(session: AsyncSession = Depends(get_db)) -> AgentService:
|
||
|
|
return AgentService(
|
||
|
|
repository=AgentRepository(session),
|
||
|
|
queue=CeleryQueueClient(),
|
||
|
|
stream=RedisEventStream(),
|
||
|
|
)
|