diff --git a/src/llm/models.py b/src/llm/models.py index caae1d5..bf53b26 100644 --- a/src/llm/models.py +++ b/src/llm/models.py @@ -54,6 +54,16 @@ class ContextUpdate(BaseModel): ) +class FilterResult(BaseModel): + contextual_info: str = Field( + ..., + description="Information interesting to the user but not useful for structured extraction", + ) + filtered_text: str = Field( + ..., description="Cleaned transcript used for structured data extraction" + ) + + class ExtractionResult(BaseModel): lore_updates: List[LoreUpdate] = Field( default_factory=list, description="List of discovered lore facts", alias="lore" @@ -68,6 +78,11 @@ class ExtractionResult(BaseModel): description="List of significant plot points or events", alias="events", ) + context_updates: List[ContextUpdate] = Field( + default_factory=list, + description="List of context updates", + alias="context", + ) class Config: populate_by_name = True diff --git a/src/llm/processor.py b/src/llm/processor.py index 99fc04c..fdc8cda 100644 --- a/src/llm/processor.py +++ b/src/llm/processor.py @@ -1,11 +1,13 @@ import logging import os +from posix import system +from this import s from typing import Any, Dict, Optional from openai import OpenAI from pydantic import ValidationError -from .models import ExtractionResult +from .models import ExtractionResult, FilterResult from .prompts import EXTRACTION_SYSTEM_PROMPT, NOISE_FILTER_SYSTEM_PROMPT logger = logging.getLogger(__name__) @@ -90,13 +92,28 @@ class LLMProcessor: logger.error(f"LLM Error: {e}") return "" - def filter_transcript(self, text: str, context: Optional[str] = None) -> str: + def filter_transcript( + self, text: str, context: Optional[str] = None + ) -> FilterResult: """ Stage 1: Raw Transcript -> Filtered Text. """ - result = self._call_llm(NOISE_FILTER_SYSTEM_PROMPT, text, context=context) + result = self._call_llm( + NOISE_FILTER_SYSTEM_PROMPT, + text, + context=context, + response_format={"type": "json_object"}, + ) logger.info(f"LLM Processor (Filter): {text} -> {result}") - return result + + import json + + try: + data = json.loads(result) + return FilterResult(**data) + except (json.JSONDecodeError, ValidationError) as e: + logger.error(f"Filter Parsing Error: {e}") + return FilterResult(contextual_info="", filtered_text=result) def extract_structured_data( self, filtered_text: str, context: Optional[str] = None @@ -109,18 +126,18 @@ class LLMProcessor: # Using standard chat.completions.create with JSON mode for better compatibility with vLLM logger.info("LLM Processor (Extract): Sending request to backend...") - messages = [ - {"role": "system", "content": EXTRACTION_SYSTEM_PROMPT}, - ] + system_prompt = EXTRACTION_SYSTEM_PROMPT if context: - messages.append( - { - "role": "system", - "content": f"Context from previous conversation:\n{context}", - } - ) + system_prompt += f"\n{context}" + + messages = [ + {"role": "system", "content": system_prompt}, + ] messages.append({"role": "user", "content": filtered_text}) + for message in messages: + logger.info(f"LLM Processor (Extract): Message: {message}") + response = self.client.chat.completions.create( model=self.model, messages=messages, @@ -142,15 +159,3 @@ class LLMProcessor: logger.error(f"Extraction Error: {e}") # Return an empty ExtractionResult if parsing fails return ExtractionResult() - - def process_pipeline( - self, raw_text: str, context: Optional[str] = None - ) -> ExtractionResult: - """ - Executes the two-stage pipeline: Raw Transcript -> Filtered Text -> Structured Data. - """ - filtered_text = self.filter_transcript(raw_text, context=context) - if not filtered_text: - return ExtractionResult() - - return self.extract_structured_data(filtered_text, context=context) diff --git a/src/llm/prompts.py b/src/llm/prompts.py index 957beba..c4cbe61 100644 --- a/src/llm/prompts.py +++ b/src/llm/prompts.py @@ -2,7 +2,11 @@ NOISE_FILTER_SYSTEM_PROMPT = """ You are a D&D Game Master's assistant. Given a transcript, remove all out-of-character (OOC) chatter, logistical discussions (e.g., 'Where is my d20?'), and non-relevant noise. -Output only the in-character dialogue and game-relevant events. + +You must output your response as a JSON object with the following keys: +- "contextual_info": Information that is interesting or relevant to the story/session but doesn't fit into lore, character state, or significant events (e.g., flavor text, atmospheric descriptions, player questions, or player commentary that adds context). +- "filtered_text": The cleaned transcript used for structured data extraction. + Keep the original speakers' names if they are present in the transcript. Do not add any commentary or summaries. Just filter the text. """ @@ -26,6 +30,7 @@ Return a JSON object with exactly these keys: - "category": (string) 'NPC', 'Location', 'WorldBuilding', or 'Plot' - "entity_name": (string) The name of the NPC, Location, or entity - "content": (string) The actual lore fact or description + - "context": (string, optional) Helpful information for the DM (e.g., descriptions of characters, spell details, game mechanics) discovered via the knowledge base or the transcript. 2. "character_state": A list of objects. Each object MUST have: - "character_name": (string) Name of the character - "hp_change": (integer, optional) Change in HP @@ -33,6 +38,9 @@ Return a JSON object with exactly these keys: - "status_effects_removed": (list of strings) - "inventory_changes": (list of objects with "item", "quantity", "action") 3. "events": A list of strings. Each string should be a concise description of a significant plot development. +4. "context": A list of objects. Each object MUST have: + - "contextual_info": (string) The contextual information for the event + - "source": (string) The source of the context (e.g., "players handbook, page 68") Example Output: { @@ -40,7 +48,7 @@ Example Output: { "category": "NPC", "entity_name": "Thorne", - "content": "A gruff dwarf who runs the local tavern." + "content": "A gruff dwarf who runs the local tavern.", } ], "character_state": [ @@ -54,6 +62,12 @@ Example Output: ], "events": [ "The party discovered the secret entrance to the crypt." + ], + "context": [ + { + "contextual_info": "fireball does 1d6 damage, with an area of effect of 10 feet circle", + "source": "players handbook, page 68" + } ] } diff --git a/src/pipeline/orchestrator.py b/src/pipeline/orchestrator.py index 0bd8153..5ddf09c 100644 --- a/src/pipeline/orchestrator.py +++ b/src/pipeline/orchestrator.py @@ -131,32 +131,76 @@ class PipelineOrchestrator: context_text = full_history_text # 2. Prepare Context (Wiki / Database of Knowledge) - wiki_context = self._get_wiki_context() + # wiki_context = self._get_wiki_context() # Combine both - combined_context = f"Conversation History:\n{context_text}\n\nWiki Knowledge:\n{wiki_context}" + combined_context = f"Conversation History:\n{context_text}\n\n" - # Process via LLM (Filter -> Extract) - # Run in a separate thread to avoid blocking the event loop - result = await asyncio.to_thread( - self.processor.process_pipeline, raw_text, context=combined_context + # --- New RAG Flow --- + # a. Filter transcript first to get cleaned text + filter_result = await asyncio.to_thread( + self.processor.filter_transcript, raw_text, context=combined_context + ) + + # b. Use filtered text to retrieve relevant snippets from RAG + rag_snippets = [] + if filter_result.filtered_text: + try: + snippets = await asyncio.to_thread( + self.rag_manager.retrieve, filter_result.filtered_text + ) + rag_snippets = snippets + except Exception as e: + logger.error(f"RAG Retrieval Error in llm_worker: {e}") + + # c. Combine RAG snippets with existing combined_context + logger.info(f"LLM Processor (Extract): rag_snippets: {rag_snippets}") + rag_context_text = "\n".join([s.snippet for s in rag_snippets]) + augmented_context = combined_context + if rag_context_text: + augmented_context += ( + f"\n\nRelevant RAG Context:\n{rag_context_text}" + ) + + # d. Extract structured data using the augmented context + extraction_result = await asyncio.to_thread( + self.processor.extract_structured_data, + filter_result.filtered_text if filter_result.filtered_text else "", + context=augmented_context, ) if ( - result.lore_updates - or result.character_updates - or result.significant_events + extraction_result.lore_updates + or extraction_result.character_updates + or extraction_result.significant_events ): logger.info( - f"LLM Worker: Proposal generated. Putting into proposal queue. (Lore: {len(result.lore_updates)}, Char: {len(result.character_updates)})" + f"LLM Worker: Proposal generated. Putting into proposal queue. (Lore: {len(extraction_result.lore_updates)}, Char: {len(extraction_result.character_updates)})" ) - await self.proposal_queue.put(result) + await self.proposal_queue.put(extraction_result) - # Trigger RAG query based on extracted entities - await self._trigger_rag_queries(result) + # Trigger RAG query based on extracted entities (for TUI updates) + await self._trigger_rag_queries(extraction_result) else: logger.info("LLM Worker: No relevant game data extracted.") + # e. If the filter found contextual info, push it to the context queue + if filter_result.contextual_info: + logger.info( + f"LLM Worker: Contextual info found: {filter_result.contextual_info}" + ) + await self.context_queue.put( + ContextUpdate( + query="Filter", + snippet=filter_result.contextual_info, + source="Transcript", + ) + ) + + # f. Also push the RAG snippets used for extraction to the context queue + for snippet in rag_snippets: + await self.context_queue.put(snippet) + except Exception as e: logger.error(f"LLM Worker error: {e}") diff --git a/src/ui/tui.py b/src/ui/tui.py index dcee0fb..0f2cd58 100644 --- a/src/ui/tui.py +++ b/src/ui/tui.py @@ -3,7 +3,7 @@ from typing import List, Optional, Union from textual.app import App, ComposeResult from textual.containers import Container, Horizontal, Vertical -from textual.widgets import Button, DataTable, Footer, Input, Label, Static +from textual.widgets import Button, DataTable, Footer, Input, Label, ListView, Static from src.llm.models import CharacterStateUpdate, ExtractionResult, LoreUpdate from src.persistence.characters import update_character_state @@ -12,18 +12,16 @@ from src.persistence.lore import update_lore class ConfirmationApp(App): CSS = """ - Screen { - layout: horizontal; + #main-container { + height: 100%; + } + + #top-pane { + height: 70%; } #left-pane { - width: 30%; - border: solid; - padding: 1; - } - - #middle-pane { - width: 30%; + width: 60%; border: solid; padding: 1; } @@ -35,6 +33,12 @@ class ConfirmationApp(App): layout: vertical; } + #middle-pane { + height: 30%; + border: solid; + padding: 1; + } + #details-container { height: auto; margin-bottom: 1; @@ -84,35 +88,38 @@ class ConfirmationApp(App): def compose(self) -> ComposeResult: yield Container( - Horizontal( - Vertical( - DataTable(id="update-table"), - id="left-pane", + Vertical( + Horizontal( + Vertical( + DataTable(id="update-table"), + id="left-pane", + ), + Vertical( + Vertical( + Label("Details:", id="details-label"), + Static("No update selected", id="details-text"), + id="details-container", + ), + Vertical( + Label("Edit Value:"), + Input(id="edit-input"), + Button("Save Edit", id="save-edit"), + id="edit-container", + ), + Horizontal( + Button("Accept", id="btn-accept"), + Button("Reject", id="btn-reject"), + Button("Edit", id="btn-edit"), + id="actions-container", + ), + id="right-pane", + ), ), - Vertical( - Static("No context available", id="context-pane"), + Horizontal( + ListView(id="context-pane"), id="middle-pane", ), - Vertical( - Vertical( - Label("Details:", id="details-label"), - Static("No update selected", id="details-text"), - id="details-container", - ), - Vertical( - Label("Edit Value:"), - Input(id="edit-input"), - Button("Save Edit", id="save-edit"), - id="edit-container", - ), - Horizontal( - Button("Accept", id="btn-accept"), - Button("Reject", id="btn-reject"), - Button("Edit", id="btn-edit"), - id="actions-container", - ), - id="right-pane", - ), + id="main-container", ), Footer(), ) @@ -169,11 +176,13 @@ class ConfirmationApp(App): while True: try: update = await self.context_queue.get() - context_pane = self.query_one("#context-pane", Static) + context_pane = self.query_one("#context-pane", ListView) # Format the update for display display_text = f"Query: {update.query}\nSource: {update.source}\n\n{update.snippet}" - context_pane.update(display_text) + + # Add a new Static widget to the ListView + context_pane.mount(Static(display_text)) if hasattr(self.context_queue, "task_done"): self.context_queue.task_done()