Implement RAG summarization and context pipeline
- Add ContextPipeline for async RAG lookups - Implement RAG result summarization via LLMProcessor - Add CLI flag for PDF ingestion - Strip markdown code blocks from LLM responses - Update TUI context display to use ListItems
This commit is contained in:
@@ -0,0 +1,67 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Tuple
|
||||
|
||||
from src.llm.models import ContextUpdate
|
||||
from src.rag.manager import RAGManager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ContextPipeline:
|
||||
def __init__(self, rag_manager: RAGManager):
|
||||
self.rag_manager = rag_manager
|
||||
|
||||
async def process_message(
|
||||
self, speaker: str, text: str, context_queue: asyncio.Queue
|
||||
):
|
||||
"""
|
||||
Processes a single message and pushes summarized insights to the context queue.
|
||||
"""
|
||||
try:
|
||||
# Use RAGManager.retrieve with summarize=True to get concise insights
|
||||
# Run in a thread to avoid blocking the event loop
|
||||
insights = await asyncio.to_thread(
|
||||
self.rag_manager.retrieve, text, summarize=True
|
||||
)
|
||||
|
||||
if insights:
|
||||
logger.info(
|
||||
f"ContextPipeline: Found {len(insights)} insights for text: {text}"
|
||||
)
|
||||
for insight in insights:
|
||||
await context_queue.put(insight)
|
||||
else:
|
||||
logger.debug(f"ContextPipeline: No insights found for text: {text}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"ContextPipeline error processing message: {e}")
|
||||
|
||||
async def run(
|
||||
self,
|
||||
transcript_queue: asyncio.Queue,
|
||||
context_queue: asyncio.Queue,
|
||||
stop_event: asyncio.Event,
|
||||
):
|
||||
"""
|
||||
Main loop that listens to the transcript queue and triggers RAG lookups.
|
||||
"""
|
||||
logger.info("Context Pipeline started.")
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
# Get raw text from transcript queue (speaker, text)
|
||||
speaker, text = await transcript_queue.get()
|
||||
|
||||
# For now, implement the basic flow: every message triggers a lookup.
|
||||
# If performance becomes an issue, a filter can be added here.
|
||||
await self.process_message(speaker, text, context_queue)
|
||||
|
||||
# Mark the task as done
|
||||
transcript_queue.task_done()
|
||||
except Exception as e:
|
||||
logger.error(f"Context Pipeline loop error: {e}")
|
||||
|
||||
# Small sleep to avoid tight loop
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
logger.info("Context Pipeline stopped.")
|
||||
@@ -147,7 +147,9 @@ class PipelineOrchestrator:
|
||||
if filter_result.filtered_text:
|
||||
try:
|
||||
snippets = await asyncio.to_thread(
|
||||
self.rag_manager.retrieve, filter_result.filtered_text
|
||||
self.rag_manager.retrieve,
|
||||
filter_result.filtered_text,
|
||||
summarize=True,
|
||||
)
|
||||
rag_snippets = snippets
|
||||
except Exception as e:
|
||||
@@ -197,8 +199,8 @@ class PipelineOrchestrator:
|
||||
)
|
||||
)
|
||||
|
||||
# f. Also push the RAG snippets used for extraction to the context queue
|
||||
for snippet in rag_snippets:
|
||||
# f. Push the distilled RAG snippets from extraction to the context queue
|
||||
for snippet in extraction_result.context_updates:
|
||||
await self.context_queue.put(snippet)
|
||||
|
||||
except Exception as e:
|
||||
@@ -235,7 +237,9 @@ class PipelineOrchestrator:
|
||||
logger.info(f"RAG: Triggering query for: {query}")
|
||||
try:
|
||||
# Run retrieval in a thread to avoid blocking the event loop
|
||||
updates = await asyncio.to_thread(self.rag_manager.retrieve, query)
|
||||
updates = await asyncio.to_thread(
|
||||
self.rag_manager.retrieve, query, summarize=True
|
||||
)
|
||||
for update in updates:
|
||||
await self.context_queue.put(update)
|
||||
logger.info(
|
||||
@@ -295,10 +299,21 @@ class PipelineOrchestrator:
|
||||
self.is_running = True
|
||||
self.listener.start()
|
||||
|
||||
# Initialize Context Pipeline
|
||||
from src.pipeline.context_pipeline import ContextPipeline
|
||||
|
||||
self.context_pipeline = ContextPipeline(self.rag_manager)
|
||||
stop_event = asyncio.Event()
|
||||
|
||||
# Start workers as background tasks
|
||||
tasks = [
|
||||
asyncio.create_task(self.stt_worker()),
|
||||
asyncio.create_task(self.llm_worker()),
|
||||
asyncio.create_task(
|
||||
self.context_pipeline.run(
|
||||
self.transcript_queue, self.context_queue, stop_event
|
||||
)
|
||||
),
|
||||
asyncio.create_task(self.tui_worker()),
|
||||
]
|
||||
|
||||
@@ -310,6 +325,7 @@ class PipelineOrchestrator:
|
||||
pass
|
||||
finally:
|
||||
self.is_running = False
|
||||
stop_event.set()
|
||||
self.listener.stop()
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
|
||||
Reference in New Issue
Block a user