gongzhiyongandClaude Sonnet 4.6
be60d9742e
The app.storage and app.tasks packages were never committed to git, causing ModuleNotFoundError on Azure deployment. Also adds the document and sandbox tool modules with their config fields. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
111 lines
3.1 KiB
Python
111 lines
3.1 KiB
Python
"""Azure Service Bus client for async task dispatch and consumption.
|
|
|
|
Long-running tasks (document generation, deep search) are published
|
|
to a Service Bus queue so the main request can return immediately
|
|
with a "task accepted" response. A background consumer picks up
|
|
messages and processes them.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import uuid
|
|
from datetime import datetime
|
|
|
|
from azure.servicebus.aio import ServiceBusClient, ServiceBusSender, ServiceBusReceiver
|
|
from azure.servicebus import ServiceBusMessage
|
|
|
|
from app.config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_client: ServiceBusClient | None = None
|
|
QUEUE_NAME = "soc-tasks"
|
|
|
|
|
|
async def _get_client() -> ServiceBusClient:
|
|
global _client
|
|
if _client is None:
|
|
_client = ServiceBusClient.from_connection_string(
|
|
settings.azure_service_bus_connection_string
|
|
)
|
|
return _client
|
|
|
|
|
|
async def send_task(
|
|
task_type: str,
|
|
payload: dict,
|
|
conversation_id: str | None = None,
|
|
) -> str:
|
|
"""Publish an async task message to Service Bus.
|
|
|
|
Args:
|
|
task_type: Type of task (e.g. "document_generate", "deep_search").
|
|
payload: Task-specific data.
|
|
conversation_id: Optional conversation context.
|
|
|
|
Returns:
|
|
A unique task_id for tracking.
|
|
"""
|
|
client = await _get_client()
|
|
task_id = uuid.uuid4().hex
|
|
|
|
message_body = json.dumps({
|
|
"task_id": task_id,
|
|
"task_type": task_type,
|
|
"conversation_id": conversation_id,
|
|
"payload": payload,
|
|
"created_at": datetime.utcnow().isoformat(),
|
|
})
|
|
|
|
sender: ServiceBusSender
|
|
async with client.get_queue_sender(queue_name=QUEUE_NAME) as sender:
|
|
await sender.send_messages(ServiceBusMessage(message_body))
|
|
|
|
logger.info("Sent task %s (type=%s) to Service Bus", task_id, task_type)
|
|
return task_id
|
|
|
|
|
|
async def receive_tasks(max_messages: int = 10, max_wait_time: int = 5) -> list[dict]:
|
|
"""Receive pending task messages from Service Bus.
|
|
|
|
Args:
|
|
max_messages: Maximum number of messages to receive.
|
|
max_wait_time: Maximum wait time in seconds.
|
|
|
|
Returns:
|
|
List of task dicts. Each message is completed (removed from queue)
|
|
after being returned.
|
|
"""
|
|
client = await _get_client()
|
|
tasks = []
|
|
|
|
receiver: ServiceBusReceiver
|
|
async with client.get_queue_receiver(
|
|
queue_name=QUEUE_NAME,
|
|
max_wait_time=max_wait_time,
|
|
) as receiver:
|
|
messages = await receiver.receive_messages(
|
|
max_message_count=max_messages,
|
|
max_wait_time=max_wait_time,
|
|
)
|
|
for msg in messages:
|
|
try:
|
|
body = json.loads(str(msg))
|
|
tasks.append(body)
|
|
await receiver.complete_message(msg)
|
|
except Exception:
|
|
logger.exception("Failed to process Service Bus message")
|
|
await receiver.dead_letter_message(msg, reason="parse_error")
|
|
|
|
return tasks
|
|
|
|
|
|
async def close_service_bus() -> None:
|
|
"""Close the Service Bus client."""
|
|
global _client
|
|
if _client is not None:
|
|
await _client.close()
|
|
_client = None
|