feat: add queue processor with model resolution
This commit is contained in:
@@ -0,0 +1,100 @@
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
from src.backend.config import settings
|
||||
from src.backend.db.database import (
|
||||
get_db,
|
||||
pick_next_pending,
|
||||
get_topic,
|
||||
complete_request,
|
||||
fail_request,
|
||||
append_topic_summary,
|
||||
)
|
||||
from src.backend.services.worker_client import WorkerClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def resolve_model(
|
||||
request_model: str | None, topic_model: str | None
|
||||
) -> str:
|
||||
if request_model:
|
||||
return request_model
|
||||
if topic_model:
|
||||
return topic_model
|
||||
return settings.claude_model
|
||||
|
||||
|
||||
async def process_one(worker: WorkerClient) -> bool:
|
||||
db = await get_db()
|
||||
try:
|
||||
req = await pick_next_pending(db)
|
||||
if not req:
|
||||
return False
|
||||
|
||||
topic = await get_topic(db, req["topic_id"])
|
||||
if not topic:
|
||||
await fail_request(db, req["id"], "Topic not found")
|
||||
return True
|
||||
|
||||
model = resolve_model(req["model"], topic["model"])
|
||||
finally:
|
||||
await db.close()
|
||||
|
||||
try:
|
||||
result = await worker.process(
|
||||
prompt=req["prompt"],
|
||||
system_prompt=topic["system_prompt"],
|
||||
model=model,
|
||||
session_id=req["session_id"],
|
||||
)
|
||||
except Exception as e:
|
||||
db = await get_db()
|
||||
try:
|
||||
await fail_request(db, req["id"], str(e))
|
||||
finally:
|
||||
await db.close()
|
||||
return True
|
||||
|
||||
db = await get_db()
|
||||
try:
|
||||
if result.get("success"):
|
||||
await complete_request(
|
||||
db,
|
||||
req["id"],
|
||||
result["result"],
|
||||
session_id=result.get("session_id"),
|
||||
)
|
||||
|
||||
if req["summarize"]:
|
||||
try:
|
||||
summary_result = await worker.summarize(
|
||||
exchanges=[
|
||||
{"prompt": req["prompt"], "result": result["result"]}
|
||||
],
|
||||
mode="incremental",
|
||||
)
|
||||
if summary_result.get("success"):
|
||||
await append_topic_summary(
|
||||
db, req["topic_id"], summary_result["summary"]
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Summary failed for request %s: %s", req["id"], e)
|
||||
else:
|
||||
error_msg = result.get("error", "Unknown worker error")
|
||||
await fail_request(db, req["id"], str(error_msg))
|
||||
finally:
|
||||
await db.close()
|
||||
|
||||
return True
|
||||
|
||||
|
||||
async def queue_loop(worker: WorkerClient):
|
||||
while True:
|
||||
try:
|
||||
processed = await process_one(worker)
|
||||
if not processed:
|
||||
await asyncio.sleep(2)
|
||||
except Exception as e:
|
||||
logger.error("Queue loop error: %s", e)
|
||||
await asyncio.sleep(5)
|
||||
@@ -187,3 +187,74 @@ async def test_worker_client_health():
|
||||
client = WorkerClient()
|
||||
is_healthy = await client.health()
|
||||
assert is_healthy is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_model_request_overrides():
|
||||
from src.backend.services.queue import resolve_model
|
||||
|
||||
model = resolve_model(
|
||||
request_model="claude-opus-4-7",
|
||||
topic_model="claude-sonnet-4-6",
|
||||
)
|
||||
assert model == "claude-opus-4-7"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_model_topic_fallback():
|
||||
from src.backend.services.queue import resolve_model
|
||||
|
||||
model = resolve_model(
|
||||
request_model=None,
|
||||
topic_model="claude-sonnet-4-6",
|
||||
)
|
||||
assert model == "claude-sonnet-4-6"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_model_global_fallback():
|
||||
from src.backend.services.queue import resolve_model
|
||||
|
||||
model = resolve_model(request_model=None, topic_model=None)
|
||||
assert model == "claude-sonnet-4-6"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_one_request():
|
||||
from src.backend.services.queue import process_one
|
||||
from src.backend.db.database import (
|
||||
init_db,
|
||||
get_db,
|
||||
create_topic,
|
||||
create_request,
|
||||
get_request,
|
||||
)
|
||||
|
||||
await init_db()
|
||||
db = await get_db()
|
||||
try:
|
||||
topic = await create_topic(db, "test-queue", "System prompt")
|
||||
req = await create_request(db, "Hello", topic["id"])
|
||||
finally:
|
||||
await db.close()
|
||||
|
||||
mock_worker = AsyncMock()
|
||||
mock_worker.process.return_value = {
|
||||
"success": True,
|
||||
"result": "Hi there!",
|
||||
"session_id": None,
|
||||
}
|
||||
mock_worker.summarize.return_value = {
|
||||
"success": True,
|
||||
"summary": "Greeted the user.",
|
||||
}
|
||||
|
||||
await process_one(mock_worker)
|
||||
|
||||
db = await get_db()
|
||||
try:
|
||||
updated = await get_request(db, req["id"])
|
||||
assert updated["status"] == "completed"
|
||||
assert updated["result"] == "Hi there!"
|
||||
finally:
|
||||
await db.close()
|
||||
|
||||
Reference in New Issue
Block a user