mirror of
https://github.com/furyhawk/agent_alpha.git
synced 2026-07-21 10:35:34 +00:00
furyhawk
bb4b614711
- Added `vectorstore.py` for managing vector storage operations using Milvus. - Introduced `rag_document.py` for handling RAG document lifecycle, including ingestion and status updates. - Created `rag_status.py` for streaming RAG ingestion status via Redis pub/sub. - Developed `rag_sync.py` for managing synchronization operations and logs. - Implemented `sync_source.py` for managing sync source configurations and triggering syncs. - Established a lightweight task dispatcher in `dispatcher.py` for background task execution. - Registered task functions for RAG ingestion and sync operations in `rag_tasks.py`. - Configured task dispatcher settings in `arq_settings.py`. - Added necessary database models and schemas for RAG document and sync operations.
50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
"""RAG status streaming — Redis pub/sub fan-out for SSE clients."""
|
|
|
|
import asyncio
|
|
import logging
|
|
from collections.abc import AsyncIterator
|
|
|
|
import redis.asyncio as aioredis
|
|
from fastapi.sse import ServerSentEvent
|
|
|
|
from backend.core.config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class RAGStatusService:
|
|
"""Streams RAG ingestion status events from Redis pub/sub."""
|
|
|
|
CHANNEL = "rag_status"
|
|
|
|
async def stream_events(self) -> AsyncIterator[ServerSentEvent]:
|
|
"""Yield ``ServerSentEvent`` items as they arrive on the ``rag_status`` channel.
|
|
|
|
The Redis client is created per-stream (one connection per subscriber). Cleanup is
|
|
guaranteed via ``finally`` even if the consumer disconnects mid-stream.
|
|
"""
|
|
client = aioredis.from_url(settings.valkey_url) # type: ignore[no-untyped-call]
|
|
pubsub = client.pubsub()
|
|
await pubsub.subscribe(self.CHANNEL)
|
|
event_id = 0
|
|
|
|
try:
|
|
async for message in pubsub.listen():
|
|
if message["type"] != "message":
|
|
continue
|
|
payload = message["data"]
|
|
data = payload.decode() if isinstance(payload, bytes) else payload
|
|
event_id += 1
|
|
yield ServerSentEvent(raw_data=data, event="status", id=str(event_id))
|
|
except asyncio.CancelledError:
|
|
# Client disconnected — propagate cancellation cleanly
|
|
raise
|
|
except Exception as exc:
|
|
logger.warning("RAG SSE stream error: %s", exc)
|
|
finally:
|
|
try:
|
|
await pubsub.unsubscribe(self.CHANNEL)
|
|
await client.aclose()
|
|
except Exception as exc:
|
|
logger.debug("RAG SSE cleanup error: %s", exc)
|