Celes Renata
6179382d1e
- Recommendation worker now resolves thesis-rewriter config from DB and passes ollama_config to generate_recommendation. Thesis rewriting is now active when the thesis-rewriter agent exists in ai_agents. Refreshes config every 50 jobs. - Event classifier now resolves its own config separately from the document extractor via 'event-classifier' slug. Uses a separate OllamaClient when the model differs from the extractor. Refreshes alongside the extractor every 100 jobs. - Document extractor was already wired (existing code). - Added 8 unit tests for AgentConfigResolver covering: DB resolution, variant override, not-found, DB errors, TTL caching, cache refresh, and invalidation.
118 lines
4.3 KiB
Python
118 lines
4.3 KiB
Python
"""Recommendation worker entrypoint - polls Redis for recommendation jobs."""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
|
|
import asyncpg
|
|
from minio import Minio
|
|
|
|
from services.recommendation.worker import generate_recommendation
|
|
from services.shared.agent_config import AgentConfigResolver
|
|
from services.shared.config import OllamaConfig, load_config
|
|
from services.shared.logging import setup_logging
|
|
from services.shared.redis_keys import QUEUE_RECOMMENDATION, queue_key
|
|
|
|
logger = logging.getLogger("recommendation_main")
|
|
|
|
|
|
async def main() -> None:
|
|
config = load_config()
|
|
setup_logging("recommendation", level=config.log_level, json_output=config.json_logs)
|
|
|
|
pool = await asyncpg.create_pool(dsn=config.postgres.dsn, min_size=2, max_size=8)
|
|
minio_client = Minio(
|
|
config.minio.endpoint,
|
|
access_key=config.minio.access_key,
|
|
secret_key=config.minio.secret_key,
|
|
secure=config.minio.secure,
|
|
)
|
|
|
|
import redis.asyncio as aioredis
|
|
|
|
redis_client = aioredis.from_url(config.redis.url)
|
|
queue = queue_key(QUEUE_RECOMMENDATION)
|
|
|
|
# Resolve thesis rewriter config from DB
|
|
resolver = AgentConfigResolver(pool, ttl_seconds=60)
|
|
ollama_config: OllamaConfig | None = None
|
|
try:
|
|
resolved = await resolver.resolve("thesis-rewriter")
|
|
if resolved is not None:
|
|
ollama_config = OllamaConfig(
|
|
base_url=config.ollama.base_url,
|
|
model=resolved.model_name,
|
|
timeout=resolved.timeout_seconds,
|
|
max_retries=resolved.max_retries,
|
|
max_tokens=resolved.max_tokens,
|
|
)
|
|
logger.info(
|
|
"Thesis rewriter enabled: model=%s variant=%s",
|
|
resolved.model_name, resolved.variant_id,
|
|
)
|
|
else:
|
|
logger.info("No DB config for thesis-rewriter — thesis rewriting disabled")
|
|
except Exception:
|
|
logger.warning("Failed to resolve thesis-rewriter config — thesis rewriting disabled", exc_info=True)
|
|
|
|
logger.info("Recommendation worker started, polling %s", queue)
|
|
|
|
refresh_counter = 0
|
|
|
|
try:
|
|
while True:
|
|
raw = await redis_client.lpop(queue)
|
|
if raw is None:
|
|
await asyncio.sleep(1)
|
|
continue
|
|
|
|
payload = raw
|
|
job = json.loads(payload)
|
|
ticker = job.get("ticker", "")
|
|
window = job.get("window", "7d")
|
|
|
|
logger.info("Processing recommendation job for %s/%s", ticker, window)
|
|
|
|
# Refresh resolver every 50 jobs to pick up config changes
|
|
refresh_counter += 1
|
|
if refresh_counter % 50 == 0:
|
|
try:
|
|
resolved = await resolver.resolve("thesis-rewriter")
|
|
if resolved is not None:
|
|
new_config = OllamaConfig(
|
|
base_url=config.ollama.base_url,
|
|
model=resolved.model_name,
|
|
timeout=resolved.timeout_seconds,
|
|
max_retries=resolved.max_retries,
|
|
max_tokens=resolved.max_tokens,
|
|
)
|
|
if ollama_config is None or new_config.model != ollama_config.model:
|
|
logger.info("Thesis rewriter config updated: model=%s", resolved.model_name)
|
|
ollama_config = new_config
|
|
except Exception:
|
|
logger.warning("Failed to refresh thesis-rewriter config", exc_info=True)
|
|
|
|
try:
|
|
rec = await generate_recommendation(
|
|
pool, ticker, window,
|
|
minio_client=minio_client,
|
|
ollama_config=ollama_config,
|
|
)
|
|
if rec:
|
|
logger.info(
|
|
"Recommendation generated for %s: %s %s",
|
|
ticker, rec.action.value, rec.mode.value,
|
|
)
|
|
else:
|
|
logger.info("No recommendation generated for %s (no trend data)", ticker)
|
|
except Exception:
|
|
logger.exception("Recommendation failed for %s", ticker)
|
|
finally:
|
|
await pool.close()
|
|
await redis_client.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|