chore: checkpoint current backend/runtime changes
This commit is contained in:
qzl
@@ -1,15 +1,15 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Protocol, cast
|
||||
from typing import Any, Protocol
|
||||
from uuid import UUID
|
||||
|
||||
from core.agent.application.resume_service import ResumeService
|
||||
from core.agent.application.run_service import RunService
|
||||
from core.agent.infrastructure.events.redis_stream import RedisStreamEventStore
|
||||
from core.celery.app import celery_app
|
||||
from core.config.settings import config
|
||||
from core.logging import get_logger
|
||||
from services.base.redis import redis_service
|
||||
from core.taskiq.app import bulk_broker, critical_broker, default_broker
|
||||
from services.base.redis import get_or_init_redis_client
|
||||
|
||||
logger = get_logger("core.agent.infrastructure.queue.tasks")
|
||||
|
||||
@@ -29,13 +29,12 @@ class ResumeServiceLike(Protocol):
|
||||
|
||||
|
||||
async def _build_redis_publisher() -> PublishEvent:
|
||||
settings = cast(Any, config)
|
||||
client = redis_service.get_client()
|
||||
client = await get_or_init_redis_client()
|
||||
event_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,
|
||||
stream_prefix=config.agent_runtime.redis_stream_prefix,
|
||||
read_count=config.agent_runtime.redis_stream_read_count,
|
||||
block_ms=config.agent_runtime.redis_stream_block_ms,
|
||||
)
|
||||
|
||||
async def _publish(event_type: str, payload: dict[str, object]) -> None:
|
||||
@@ -70,22 +69,27 @@ async def run_agent_task(
|
||||
raise ValueError("session_id is required")
|
||||
UUID(session_id)
|
||||
|
||||
tool_call_id = ""
|
||||
user_input = ""
|
||||
if command_type == "resume":
|
||||
tool_call_id = str(command.get("tool_call_id", ""))
|
||||
if not tool_call_id:
|
||||
raise ValueError("tool_call_id is required")
|
||||
else:
|
||||
user_input = str(command.get("user_input", ""))
|
||||
if not user_input:
|
||||
raise ValueError("user_input is required")
|
||||
|
||||
start_event = "RUN_RESUMED" if command_type == "resume" else "RUN_STARTED"
|
||||
await publisher(start_event, {"session_id": session_id})
|
||||
|
||||
try:
|
||||
if command_type == "resume":
|
||||
tool_call_id = str(command.get("tool_call_id", ""))
|
||||
if not tool_call_id:
|
||||
raise ValueError("tool_call_id is required")
|
||||
result = await service_resume.resume(
|
||||
session_id=session_id,
|
||||
tool_call_id=tool_call_id,
|
||||
)
|
||||
else:
|
||||
user_input = str(command.get("user_input", ""))
|
||||
if not user_input:
|
||||
raise ValueError("user_input is required")
|
||||
result = await service_run.run(
|
||||
session_id=session_id,
|
||||
user_input=user_input,
|
||||
@@ -125,6 +129,16 @@ async def run_agent_task(
|
||||
raise
|
||||
|
||||
|
||||
@celery_app.task(name="tasks.agent.run_command")
|
||||
@default_broker.task(task_name="tasks.agent.run_command")
|
||||
async def run_command_task(command: dict[str, Any]) -> dict[str, object]:
|
||||
return await run_agent_task(command)
|
||||
|
||||
|
||||
@critical_broker.task(task_name="tasks.agent.run_command.critical")
|
||||
async def run_command_task_critical(command: dict[str, Any]) -> dict[str, object]:
|
||||
return await run_agent_task(command)
|
||||
|
||||
|
||||
@bulk_broker.task(task_name="tasks.agent.run_command.bulk")
|
||||
async def run_command_task_bulk(command: dict[str, Any]) -> dict[str, object]:
|
||||
return await run_agent_task(command)
|
||||
|
||||
@@ -1,73 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from celery import Celery
|
||||
from celery import signals as celery_signals
|
||||
from kombu import Queue
|
||||
|
||||
from core.config.settings import config
|
||||
from core.logging import get_logger
|
||||
from core.logging.celery import configure_celery_app
|
||||
from services.base.redis import redis_service
|
||||
|
||||
logger = get_logger("core.celery")
|
||||
|
||||
|
||||
def _init_redis_on_worker_startup(**_: object) -> None:
|
||||
import asyncio
|
||||
|
||||
logger.info("Initializing Redis service for Celery worker")
|
||||
try:
|
||||
result = asyncio.run(redis_service.initialize())
|
||||
if result:
|
||||
logger.info("Redis service initialized for Celery worker")
|
||||
else:
|
||||
logger.warning("Redis service initialization returned False")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.error("Failed to initialize Redis for Celery worker", error=str(exc))
|
||||
|
||||
|
||||
def create_celery_app() -> Celery:
|
||||
"""Create and configure the Celery application."""
|
||||
celery_settings = config.celery
|
||||
|
||||
app = Celery(
|
||||
"social_app",
|
||||
broker=config.celery_broker_url,
|
||||
backend=config.celery_result_backend,
|
||||
include=["core.agent.infrastructure.queue.tasks"],
|
||||
)
|
||||
|
||||
app.conf.update(
|
||||
task_serializer=celery_settings.task_serializer,
|
||||
result_serializer=celery_settings.result_serializer,
|
||||
accept_content=celery_settings.accept_content,
|
||||
timezone=celery_settings.timezone,
|
||||
enable_utc=celery_settings.enable_utc,
|
||||
task_track_started=celery_settings.task_track_started,
|
||||
task_time_limit=celery_settings.task_time_limit,
|
||||
task_soft_time_limit=celery_settings.task_soft_time_limit,
|
||||
task_default_retry_delay=celery_settings.task_default_retry_delay,
|
||||
task_default_queue="default",
|
||||
task_create_missing_queues=False,
|
||||
task_queues=(
|
||||
Queue("default"),
|
||||
Queue("critical"),
|
||||
Queue("bulk"),
|
||||
),
|
||||
task_routes={
|
||||
"tasks.critical.*": {"queue": "critical"},
|
||||
"tasks.bulk.*": {"queue": "bulk"},
|
||||
},
|
||||
task_acks_late=True,
|
||||
task_reject_on_worker_lost=True,
|
||||
worker_prefetch_multiplier=1,
|
||||
)
|
||||
|
||||
configure_celery_app(app, settings=config)
|
||||
|
||||
return app
|
||||
|
||||
|
||||
celery_app = create_celery_app()
|
||||
|
||||
celery_signals.worker_process_init.connect(_init_redis_on_worker_startup)
|
||||
@@ -63,19 +63,9 @@ class RuntimeSettings(BaseModel):
|
||||
return self
|
||||
|
||||
|
||||
class CelerySettings(BaseModel):
|
||||
class TaskiqSettings(BaseModel):
|
||||
broker_url: str | None = None
|
||||
result_backend: str | None = None
|
||||
task_serializer: str = "json"
|
||||
result_serializer: str = "json"
|
||||
accept_content: list[str] = Field(default_factory=lambda: ["json"])
|
||||
timezone: str = "UTC"
|
||||
enable_utc: bool = True
|
||||
task_track_started: bool = True
|
||||
task_time_limit: int = 300
|
||||
task_soft_time_limit: int = 240
|
||||
task_default_retry_delay: int = 30
|
||||
task_max_retries: int = 3
|
||||
result_backend_url: str | None = None
|
||||
|
||||
|
||||
class CorsSettings(BaseModel):
|
||||
@@ -189,7 +179,7 @@ class Settings(BaseSettings):
|
||||
storage: StorageSettings = StorageSettings()
|
||||
llm: LlmSettings = LlmSettings()
|
||||
agent_runtime: AgentRuntimeSettings = AgentRuntimeSettings()
|
||||
celery: CelerySettings = CelerySettings()
|
||||
taskiq: TaskiqSettings = TaskiqSettings()
|
||||
database: DatabaseSettings = DatabaseSettings()
|
||||
|
||||
@computed_field
|
||||
@@ -199,13 +189,13 @@ class Settings(BaseSettings):
|
||||
|
||||
@computed_field
|
||||
@property
|
||||
def celery_broker_url(self) -> str:
|
||||
return self.celery.broker_url or self.redis.url
|
||||
def taskiq_broker_url(self) -> str:
|
||||
return self.taskiq.broker_url or self.redis.url
|
||||
|
||||
@computed_field
|
||||
@property
|
||||
def celery_result_backend(self) -> str:
|
||||
return self.celery.result_backend or self.redis.url
|
||||
def taskiq_result_backend_url(self) -> str:
|
||||
return self.taskiq.result_backend_url or self.redis.url
|
||||
|
||||
model_config: ClassVar[SettingsConfigDict] = SettingsConfigDict(
|
||||
env_file=_resolve_env_file(),
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from core.logging import celery
|
||||
from core.logging.banner import log_service_banner
|
||||
from core.logging.config import configure_logging
|
||||
from core.logging.context import bind_context, clear_context, get_context
|
||||
@@ -8,7 +7,6 @@ from core.logging.logger import get_logger
|
||||
|
||||
__all__ = [
|
||||
"bind_context",
|
||||
"celery",
|
||||
"clear_context",
|
||||
"configure_logging",
|
||||
"get_context",
|
||||
|
||||
@@ -1,64 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import cast
|
||||
|
||||
from celery import Celery, signals
|
||||
|
||||
from core.config.settings import Settings
|
||||
from core.logging.banner import log_service_banner
|
||||
from core.logging.config import configure_logging
|
||||
from core.logging.context import bind_context, clear_context
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CelerySignalHandlers:
|
||||
on_setup_logging: Callable[..., None]
|
||||
on_after_setup_task_logger: Callable[..., None]
|
||||
on_task_prerun: Callable[..., None]
|
||||
on_task_postrun: Callable[..., None]
|
||||
|
||||
|
||||
def build_celery_signal_handlers(
|
||||
settings: Settings | None = None,
|
||||
) -> CelerySignalHandlers:
|
||||
active_settings = settings or Settings()
|
||||
|
||||
def on_setup_logging(*_args: object, **_kwargs: object) -> None:
|
||||
configure_logging(settings)
|
||||
log_service_banner(
|
||||
service_name=active_settings.runtime.service_name,
|
||||
environment=active_settings.runtime.environment,
|
||||
)
|
||||
|
||||
def on_after_setup_task_logger(*_args: object, **_kwargs: object) -> None:
|
||||
configure_logging(settings)
|
||||
|
||||
def on_task_prerun(*_args: object, **kwargs: object) -> None:
|
||||
task_id = cast(str | None, kwargs.get("task_id"))
|
||||
task = kwargs.get("task")
|
||||
task_name = getattr(task, "name", None)
|
||||
bind_context(task_id=task_id, task_name=task_name)
|
||||
|
||||
def on_task_postrun(*_args: object, **_kwargs: object) -> None:
|
||||
clear_context()
|
||||
|
||||
return CelerySignalHandlers(
|
||||
on_setup_logging=on_setup_logging,
|
||||
on_after_setup_task_logger=on_after_setup_task_logger,
|
||||
on_task_prerun=on_task_prerun,
|
||||
on_task_postrun=on_task_postrun,
|
||||
)
|
||||
|
||||
|
||||
def configure_celery_app(app: Celery, settings: Settings | None = None) -> None:
|
||||
app.conf.worker_hijack_root_logger = False
|
||||
|
||||
handlers = build_celery_signal_handlers(settings)
|
||||
signals.setup_logging.connect(handlers.on_setup_logging, weak=False)
|
||||
signals.after_setup_task_logger.connect(
|
||||
handlers.on_after_setup_task_logger, weak=False
|
||||
)
|
||||
signals.task_prerun.connect(handlers.on_task_prerun, weak=False)
|
||||
signals.task_postrun.connect(handlers.on_task_postrun, weak=False)
|
||||
@@ -0,0 +1,3 @@
|
||||
from core.taskiq.app import broker, bulk_broker, critical_broker, default_broker
|
||||
|
||||
__all__ = ["broker", "default_broker", "critical_broker", "bulk_broker"]
|
||||
@@ -0,0 +1,27 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from taskiq_redis import ListQueueBroker, RedisAsyncResultBackend
|
||||
|
||||
from core.config.settings import config
|
||||
from core.logging import configure_logging
|
||||
|
||||
|
||||
configure_logging(config)
|
||||
|
||||
def _build_broker(queue_name: str) -> ListQueueBroker:
|
||||
return ListQueueBroker(
|
||||
url=config.taskiq_broker_url,
|
||||
queue_name=queue_name,
|
||||
).with_result_backend(
|
||||
RedisAsyncResultBackend(redis_url=config.taskiq_result_backend_url)
|
||||
)
|
||||
|
||||
|
||||
default_broker = _build_broker("default")
|
||||
critical_broker = _build_broker("critical")
|
||||
bulk_broker = _build_broker("bulk")
|
||||
|
||||
# Backward-compatible export name for existing imports/tests.
|
||||
broker = default_broker
|
||||
|
||||
__all__ = ["broker", "default_broker", "critical_broker", "bulk_broker"]
|
||||
Reference in New Issue
Block a user