feat(agent): complete closed-loop runtime and pricing fallback
This commit is contained in:
@@ -0,0 +1,75 @@
|
||||
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(),
|
||||
)
|
||||
Reference in New Issue
Block a user