143 lines
4.8 KiB
Python
143 lines
4.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from contextlib import suppress
|
|
from socket import gethostname
|
|
from typing import Any
|
|
|
|
from app.core.logging import get_logger
|
|
from app.core.websocket.broadcaster import broadcaster
|
|
from app.db.session import async_session_factory
|
|
from app.services.ai_client import AIProviderClient
|
|
from app.ai_tasks.prompts import get_effective_prompt
|
|
from app.services.earth_news import (
|
|
NEWS_ENRICH_PROMPT_KEY,
|
|
_infer_news_enrichment,
|
|
build_target_location_patch,
|
|
parsed_news_item_from_job_payload,
|
|
)
|
|
from app.services.earth_news_queue import (
|
|
NewsTargetLocationMessage,
|
|
get_news_target_location_queue,
|
|
save_target_location_patch,
|
|
)
|
|
from app.services.earth_news_store import update_earth_news_item_enrichment as update_earth_news_item_location
|
|
|
|
|
|
logger = get_logger(__name__, service="earth_news")
|
|
|
|
WORKER_BATCH_SIZE = 4
|
|
WORKER_BLOCK_MS = 5000
|
|
WORKER_BACKOFF_SECONDS = 5.0
|
|
|
|
_worker_task: asyncio.Task | None = None
|
|
|
|
|
|
async def _build_provider_client() -> AIProviderClient | None:
|
|
try:
|
|
from app.api.v1.settings import get_runtime_ai_provider_config
|
|
|
|
async with async_session_factory() as session:
|
|
runtime_config = await get_runtime_ai_provider_config(session)
|
|
return AIProviderClient(
|
|
service_url=runtime_config["service_url"],
|
|
service_token=runtime_config["service_token"],
|
|
timeout=runtime_config["timeout_seconds"],
|
|
retry_attempts=runtime_config["retry_attempts"],
|
|
llm_config=runtime_config.get("llm_config") or {},
|
|
)
|
|
except Exception as exc:
|
|
logger.warning_event(
|
|
"Failed to build Earth news AI provider client",
|
|
event="earth_news.target_location.provider_unavailable",
|
|
context={"error": str(exc)},
|
|
)
|
|
return None
|
|
|
|
|
|
async def process_target_location_message(
|
|
message: NewsTargetLocationMessage,
|
|
*,
|
|
provider_client: AIProviderClient | None,
|
|
) -> dict[str, Any]:
|
|
item = parsed_news_item_from_job_payload(message.payload)
|
|
async with async_session_factory() as session:
|
|
prompt = await get_effective_prompt(session, NEWS_ENRICH_PROMPT_KEY)
|
|
target, localizations = await _infer_news_enrichment(
|
|
item,
|
|
provider_client=provider_client,
|
|
prompt=prompt,
|
|
)
|
|
item.target_location = target
|
|
item.localizations = localizations or item.localizations
|
|
patch = build_target_location_patch(item, target)
|
|
await save_target_location_patch(item.id, patch)
|
|
async with async_session_factory() as session:
|
|
await update_earth_news_item_location(session, item_id=item.id, patch=patch)
|
|
await session.commit()
|
|
await broadcaster.broadcast_custom(
|
|
"earth_news",
|
|
{
|
|
"item_id": item.id,
|
|
"patch": patch,
|
|
},
|
|
)
|
|
return patch
|
|
|
|
|
|
async def _run_target_location_worker() -> None:
|
|
consumer_name = f"{gethostname()}:{id(asyncio.current_task())}"
|
|
queue = get_news_target_location_queue()
|
|
while True:
|
|
try:
|
|
messages = await queue.consume_batch(
|
|
consumer_name=consumer_name,
|
|
count=WORKER_BATCH_SIZE,
|
|
block_ms=WORKER_BLOCK_MS,
|
|
)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc:
|
|
logger.warning_event(
|
|
"Earth news target location worker queue read failed",
|
|
event="earth_news.target_location.worker_read_failed",
|
|
context={"error": str(exc)},
|
|
)
|
|
await asyncio.sleep(WORKER_BACKOFF_SECONDS)
|
|
continue
|
|
|
|
if not messages:
|
|
continue
|
|
provider_client = await _build_provider_client()
|
|
for message in messages:
|
|
try:
|
|
await process_target_location_message(message, provider_client=provider_client)
|
|
await queue.ack(message.message_id)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc:
|
|
logger.warning_event(
|
|
"Earth news target location worker job failed",
|
|
event="earth_news.target_location.worker_job_failed",
|
|
context={"item_id": message.item_id, "error": str(exc)},
|
|
)
|
|
with suppress(Exception):
|
|
await queue.retry_or_dead_letter(message, error=str(exc))
|
|
|
|
|
|
def start_earth_news_target_worker() -> None:
|
|
global _worker_task
|
|
if _worker_task is None or _worker_task.done():
|
|
_worker_task = asyncio.create_task(_run_target_location_worker())
|
|
|
|
|
|
async def stop_earth_news_target_worker() -> None:
|
|
global _worker_task
|
|
task = _worker_task
|
|
if task is None:
|
|
return
|
|
task.cancel()
|
|
with suppress(asyncio.CancelledError):
|
|
await task
|
|
_worker_task = None
|