feat: wire all 3 agents to DB config resolver
- 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.
This commit is contained in:
@@ -9,7 +9,8 @@ import asyncpg
|
||||
from minio import Minio
|
||||
|
||||
from services.recommendation.worker import generate_recommendation
|
||||
from services.shared.config import load_config
|
||||
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
|
||||
|
||||
@@ -32,8 +33,33 @@ async def main() -> None:
|
||||
|
||||
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)
|
||||
@@ -48,10 +74,30 @@ async def main() -> None:
|
||||
|
||||
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(
|
||||
|
||||
Reference in New Issue
Block a user