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()
|
client = WorkerClient()
|
||||||
is_healthy = await client.health()
|
is_healthy = await client.health()
|
||||||
assert is_healthy is True
|
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