Add reportlab dependency for PDF generation and enhance document processing

- Added reportlab to pyproject.toml and requirements.txt for PDF generation capabilities.
- Updated document processing logic to include a new function for generating redacted PDFs from the processed document structure.
- Enhanced logging and error handling during PDF generation to improve user feedback and debugging.
- Refactored the document processor to return detailed processing results, including removed indices and cost metrics.

JUPYTER_USAGE.md

@@ -0,0 +1,143 @@
+# Jupyter Notebook Usage
+
+This document shows how to use the document processing function in Jupyter notebooks for integration into larger processing pipelines.
+
+## Simple Usage
+
+```python
+from processing.document_processor import process_document_with_redaction
+
+# Process a single document
+result = process_document_with_redaction(
+    file_path="path/to/your/document.pdf",
+    endpoint="your-azure-openai-endpoint",
+    api_key="your-azure-openai-key",
+    api_version="2024-02-15-preview",
+    deployment="o3-mini"  # or "o4-mini", "o3", "o4"
+)
+
+# Access the results
+original_md = result.original_document_md
+redacted_md = result.redacted_document_md
+input_tokens = result.input_tokens
+output_tokens = result.output_tokens
+cost = result.cost
+
+print(f"Processing complete!")
+print(f"Input tokens: {input_tokens:,}")
+print(f"Output tokens: {output_tokens:,}")
+print(f"Total cost: ${cost:.4f}")
+```
+
+## Batch Processing
+
+```python
+import os
+from processing.document_processor import process_document_with_redaction
+
+# Configuration
+AZURE_OPENAI_ENDPOINT = "your-azure-openai-endpoint"
+AZURE_OPENAI_KEY = "your-azure-openai-key"
+AZURE_OPENAI_VERSION = "2024-02-15-preview"
+AZURE_OPENAI_DEPLOYMENT = "o3-mini"
+
+# Process multiple documents
+pdf_directory = "path/to/pdf/files"
+results = []
+
+for filename in os.listdir(pdf_directory):
+    if filename.endswith('.pdf'):
+        file_path = os.path.join(pdf_directory, filename)
+        
+        print(f"Processing {filename}...")
+        
+        try:
+            result = process_document_with_redaction(
+                file_path=file_path,
+                endpoint=AZURE_OPENAI_ENDPOINT,
+                api_key=AZURE_OPENAI_KEY,
+                api_version=AZURE_OPENAI_VERSION,
+                deployment=AZURE_OPENAI_DEPLOYMENT
+            )
+            
+            results.append({
+                'filename': filename,
+                'original_md': result.original_document_md,
+                'redacted_md': result.redacted_document_md,
+                'input_tokens': result.input_tokens,
+                'output_tokens': result.output_tokens,
+                'cost': result.cost
+            })
+            
+            print(f"  ✓ Completed - Cost: ${result.cost:.4f}")
+            
+        except Exception as e:
+            print(f"  ✗ Error processing {filename}: {e}")
+
+# Summary
+total_cost = sum(r['cost'] for r in results)
+total_input_tokens = sum(r['input_tokens'] for r in results)
+total_output_tokens = sum(r['output_tokens'] for r in results)
+
+print(f"\nBatch processing complete!")
+print(f"Documents processed: {len(results)}")
+print(f"Total input tokens: {total_input_tokens:,}")
+print(f"Total output tokens: {total_output_tokens:,}")
+print(f"Total cost: ${total_cost:.4f}")
+```
+
+## Environment Variables
+
+You can also use environment variables for configuration:
+
+```python
+import os
+from dotenv import load_dotenv
+from processing.document_processor import process_document_with_redaction
+
+# Load environment variables
+load_dotenv()
+
+# Get configuration from environment
+AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
+AZURE_OPENAI_KEY = os.getenv("AZURE_OPENAI_KEY")
+AZURE_OPENAI_VERSION = os.getenv("AZURE_OPENAI_VERSION")
+AZURE_OPENAI_DEPLOYMENT = os.getenv("AZURE_OPENAI_DEPLOYMENT")
+
+# Process document
+result = process_document_with_redaction(
+    file_path="document.pdf",
+    endpoint=AZURE_OPENAI_ENDPOINT,
+    api_key=AZURE_OPENAI_KEY,
+    api_version=AZURE_OPENAI_VERSION,
+    deployment=AZURE_OPENAI_DEPLOYMENT
+)
+```
+
+## Return Value
+
+The function returns a `ProcessingResult` object with the following attributes:
+
+- `original_document_md`: Markdown version of the original document
+- `redacted_document_md`: Markdown version with medication sections removed
+- `input_tokens`: Number of input tokens used
+- `output_tokens`: Number of output tokens generated
+- `cost`: Total cost in USD
+
+## Supported Models
+
+The function supports the following Azure OpenAI deployment names:
+- `o3-mini` (GPT-4o Mini) - Cheapest option
+- `o4-mini` (GPT-4o Mini) - Same as o3-mini
+- `o3` (GPT-3.5 Turbo) - Medium cost
+- `o4` (GPT-4o) - Most expensive but most capable
+
+## Error Handling
+
+The function will raise exceptions for:
+- File not found
+- Invalid Azure OpenAI credentials
+- API rate limits
+- Network errors
+
+Make sure to handle these appropriately in your pipeline. 

pyproject.toml

@@ -10,4 +10,5 @@ dependencies = [
     "pyyaml>=6.0",
     "python-dotenv>=1.1.1",
     "openai>=1.91.0",
+    "reportlab>=4.4.2",
 ]

requirements.txt

@@ -4,4 +4,5 @@ docling
 streamlit
 pyyaml
 python-dotenv
-openai
+openai
+reportlab

src/processing/document_processor.py

@@ -3,11 +3,12 @@ import time
 import logging
 import json
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, Tuple
 
 # Don't import DocumentConverter at module level to prevent early initialization
 # from docling.document_converter import DocumentConverter
 from processing.sections import SectionExtractor
+from utils.cost_tracker import cost_tracker
 
 # Remove global converter initialization - will be done lazily
 # _docling_converter = DocumentConverter()
@@ -23,6 +24,92 @@ class DocumentResult:
     redacted_markdown: str
     redacted_json: dict
 
+@dataclass
+class ProcessingResult:
+    """Simple result for Jupyter notebook usage."""
+    original_document_md: str
+    redacted_document_md: str
+    original_document_json: dict
+    redacted_document_json: dict
+    removed_indices: list  # Add the actual indices that were removed
+    input_tokens: int
+    output_tokens: int
+    cost: float
+
+def process_document_with_redaction(
+    file_path: str,
+    endpoint: str,
+    api_key: str,
+    api_version: str,
+    deployment: str,
+    section_extractor: Optional[SectionExtractor] = None
+) -> ProcessingResult:
+    """
+    Process a document and return a simple tuple with results.
+    
+    Args:
+        file_path: Path to the PDF file to process
+        endpoint: Azure OpenAI endpoint
+        api_key: Azure OpenAI API key
+        api_version: Azure OpenAI API version
+        deployment: Azure OpenAI deployment name
+        section_extractor: Optional custom section extractor
+    
+    Returns:
+        ProcessingResult with (original_document_md, redacted_document_md, input_tokens, output_tokens, cost)
+    """
+    logger.info(f"Processing document: {file_path}")
+    
+    # Reset cost tracker for this processing session
+    cost_tracker.reset_session()
+    
+    # Create section extractor if not provided
+    if section_extractor is None:
+        from processing.sections import ReasoningSectionExtractor
+        section_extractor = ReasoningSectionExtractor(
+            endpoint=endpoint,
+            api_key=api_key,
+            api_version=api_version,
+            deployment=deployment,
+        )
+    
+    # Process the document
+    processor = DocumentProcessor(section_extractor=section_extractor)
+    result = processor.process(file_path)
+    
+    # Get the actual removed indices from the section extractor
+    removed_indices = []
+    if section_extractor:
+        # Extract the removed indices from the LLM response
+        extraction_result = section_extractor.llm_extractor.extract_medication_sections(result.structured_json)
+        removed_indices = extraction_result.get("indices_to_remove", [])
+    
+    # Get cost summary
+    cost_summary = cost_tracker.get_session_summary()
+    total_input_tokens = cost_summary.get("total_tokens", 0)
+    total_output_tokens = 0  # We'll calculate this from the breakdown
+    total_cost = cost_summary.get("total_cost", 0.0)
+    
+    # Calculate output tokens from model breakdown
+    for model_stats in cost_summary.get("model_breakdown", {}).values():
+        total_output_tokens += model_stats.get("output_tokens", 0)
+    
+    # Calculate input tokens (total - output)
+    total_input_tokens = total_input_tokens - total_output_tokens
+    
+    logger.info(f"Processing complete - Input: {total_input_tokens}, Output: {total_output_tokens}, Cost: ${total_cost:.4f}")
+    
+    return ProcessingResult(
+        original_document_md=result.structured_markdown,
+        redacted_document_md=result.redacted_markdown,
+        original_document_json=result.structured_json,
+        redacted_document_json=result.redacted_json,
+        removed_indices=removed_indices,
+        input_tokens=total_input_tokens,
+        output_tokens=total_output_tokens,
+        cost=total_cost
+    )
+
 class DocumentProcessor:
     """Handles parsing of documents with Docling and redacting specified sections."""
     def __init__(self, section_extractor: Optional[SectionExtractor] = None):
@@ -80,12 +167,12 @@ class DocumentProcessor:
         
         # Persist outputs to files (JSON and redacted text) for auditing
         base_name = os.path.splitext(os.path.basename(file_path))[0]
-        # Use temp directory for output files - try to use the same temp dir as the main app
-        temp_dir = "temp_files"
+        # Use the same temp directory as the main application
+        temp_dir = os.environ.get('TEMP_DIR', '/tmp/docling_temp')
         try:
             os.makedirs(temp_dir, exist_ok=True)
         except PermissionError:
-            # Fallback to system temp directory if we can't create in current directory
+            # Fallback to system temp directory if we can't create in the main temp dir
             import tempfile
             temp_dir = os.path.join(tempfile.gettempdir(), "docling_temp_files")
             os.makedirs(temp_dir, exist_ok=True)
@@ -139,10 +226,177 @@ class DocumentProcessor:
                 logger.warning(f"Could not create cache directory {cache_dir}: {e}")
 
     def _export_redacted_markdown(self, document, redacted_json):
-        """Export redacted markdown using the redacted JSON structure."""
-        # Simply convert the redacted JSON back to markdown
-        return self._json_to_markdown(redacted_json)
+        """Export redacted markdown using Docling's Document class for proper formatting."""
+        try:
+            # Try different possible import paths for Docling Document class
+            try:
+                from docling.document import Document
+            except ImportError:
+                try:
+                    from docling import Document
+                except ImportError:
+                    try:
+                        from docling.core import Document
+                    except ImportError:
+                        # If all imports fail, use the fallback method
+                        logger.warning("Could not import Docling Document class from any known location")
+                        raise ImportError("Docling Document class not found")
+            
+            # Create a new Document from the redacted JSON
+            redacted_document = Document.from_dict(redacted_json)
+            
+            # Use Docling's export method for proper markdown formatting
+            redacted_md = redacted_document.export_to_markdown()
+            logger.info("Successfully generated redacted markdown using Docling Document class")
+            return redacted_md
+            
+        except Exception as e:
+            logger.warning(f"Failed to create Docling Document from redacted JSON: {e}")
+            logger.info("Falling back to manual JSON-to-markdown conversion")
+            # Fallback to the old method if Docling Document creation fails
+            return self._json_to_markdown(redacted_json)
     
+    def generate_redacted_pdf(self, redacted_json: dict, output_path: str) -> bool:
+        """
+        Generate a redacted PDF from the redacted JSON structure.
+        
+        Args:
+            redacted_json: The redacted document JSON structure
+            output_path: Path where the PDF should be saved
+            
+        Returns:
+            bool: True if PDF generation was successful, False otherwise
+        """
+        try:
+            # Import required libraries
+            from reportlab.lib.pagesizes import letter, A4
+            from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle
+            from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
+            from reportlab.lib.units import inch
+            from reportlab.lib import colors
+            import io
+            
+            logger.info(f"Generating redacted PDF: {output_path}")
+            
+            # Create PDF document
+            doc = SimpleDocTemplate(output_path, pagesize=A4)
+            story = []
+            
+            # Get styles
+            styles = getSampleStyleSheet()
+            normal_style = styles['Normal']
+            heading_style = styles['Heading1']
+            
+            # Create custom styles for better formatting
+            table_style = ParagraphStyle(
+                'TableStyle',
+                parent=normal_style,
+                fontName='Courier',
+                fontSize=9,
+                spaceAfter=6
+            )
+            
+            # Process text elements from JSON
+            texts = redacted_json.get("texts", [])
+            
+            # Group consecutive table-like elements together
+            i = 0
+            while i < len(texts):
+                text_elem = texts[i]
+                text_content = text_elem.get("text", "").strip()
+                label = text_elem.get("label", "")
+                level = text_elem.get("level", 0)
+                
+                if not text_content:
+                    i += 1
+                    continue
+                
+                # Handle different content types
+                if label == "section_header":
+                    # Create header with appropriate level
+                    if level == 1:
+                        story.append(Paragraph(text_content, heading_style))
+                    else:
+                        # Create sub-heading style
+                        sub_heading_style = ParagraphStyle(
+                            f'Heading{min(level, 3)}',
+                            parent=normal_style,
+                            fontSize=14 - level,
+                            spaceAfter=12,
+                            spaceBefore=12,
+                            textColor=colors.darkblue
+                        )
+                        story.append(Paragraph(text_content, sub_heading_style))
+                
+                elif label == "list_item":
+                    # Handle list items
+                    marker = text_elem.get("marker", "•")
+                    list_text = f"{marker} {text_content}"
+                    story.append(Paragraph(list_text, normal_style))
+                
+                elif '|' in text_content and text_content.count('|') > 1:
+                    # Handle table-like content - collect consecutive table rows
+                    table_rows = []
+                    
+                    # Add the current row
+                    cells = [cell.strip() for cell in text_content.split('|') if cell.strip()]
+                    if cells:
+                        table_rows.append(cells)
+                    
+                    # Look ahead for consecutive table rows
+                    j = i + 1
+                    while j < len(texts):
+                        next_text = texts[j].get("text", "").strip()
+                        if '|' in next_text and next_text.count('|') > 1:
+                            next_cells = [cell.strip() for cell in next_text.split('|') if cell.strip()]
+                            if next_cells:
+                                table_rows.append(next_cells)
+                            j += 1
+                        else:
+                            break
+                    
+                    # Create table if we have rows
+                    if table_rows:
+                        table = Table(table_rows)
+                        table.setStyle(TableStyle([
+                            ('ALIGN', (0, 0), (-1, -1), 'LEFT'),
+                            ('FONTNAME', (0, 0), (-1, -1), 'Courier'),
+                            ('FONTSIZE', (0, 0), (-1, -1), 9),
+                            ('BOTTOMPADDING', (0, 0), (-1, -1), 3),
+                            ('TOPPADDING', (0, 0), (-1, -1), 3),
+                            ('GRID', (0, 0), (-1, -1), 0.5, colors.grey),
+                            ('BACKGROUND', (0, 0), (-1, 0), colors.lightgrey),  # Header row
+                        ]))
+                        story.append(table)
+                        story.append(Spacer(1, 6))
+                        
+                        # Skip the rows we've already processed
+                        i = j - 1
+                    else:
+                        # Single row or no valid cells
+                        story.append(Paragraph(text_content, table_style))
+                
+                else:
+                    # Regular text content
+                    story.append(Paragraph(text_content, normal_style))
+                
+                # Add small spacing between elements
+                story.append(Spacer(1, 3))
+                i += 1
+            
+            # Build PDF
+            doc.build(story)
+            logger.info(f"Successfully generated redacted PDF: {output_path}")
+            return True
+            
+        except ImportError as e:
+            logger.error(f"Required PDF generation libraries not available: {e}")
+            logger.info("Install reportlab with: pip install reportlab")
+            return False
+        except Exception as e:
+            logger.error(f"Error generating redacted PDF: {e}")
+            return False
+
     def _json_to_markdown(self, json_data: dict) -> str:
         """Convert JSON document structure back to markdown format using Docling's structure."""
         markdown_lines = []

src/processing/llm_extractor.py

@@ -6,6 +6,7 @@ import logging
 from typing import Dict, Any
 
 from openai import AzureOpenAI
+from utils.cost_tracker import cost_tracker
 
 logger = logging.getLogger(__name__)
 
@@ -39,107 +40,137 @@ class AzureO1MedicationExtractor:
             })
 
         prompt = f"""
-You are a medical document analysis expert specializing in discharge letters. Your task is to identify ONLY the two formal medication lists that should be redacted, while preserving all medication mentions in clinical discussion.
-
-**CRITICAL: You should ONLY remove the two formal medication lists:**
-1. **Current medication list** (usually at the beginning of the document)
-2. **Discharge medication list** (usually at the end of the document, often under headers like "Als verdere behandeling stellen wij voor" or "Thuismedicatie")
-
-**Typical discharge letter structure:**
-- Patient information and admission details
-- Clinical discussion and treatment narrative (KEEP medication mentions here)
-- Current medication list (REMOVE this formal list)
-- Discharge instructions and follow-up
-- Discharge medication list (REMOVE this formal list)
-
-**DO NOT remove:**
-- Medication mentions in clinical discussion (e.g., "patient was treated with Eliquis")
-- Medication adjustments mentioned in the narrative
-- Dosage information in clinical context
-- Any medication information that appears in the main clinical text
-- Treatment decisions and clinical reasoning
-
-**ONLY remove:**
-- Complete medication lists with multiple drugs
-- Formal medication sections with headers
-- Standalone medication lists that are clearly separated from clinical text
-- Lists that appear to be formal medication documentation
+You are a medical document analysis expert specializing in discharge letters. Your task is to identify ONLY the formal medication lists that should be redacted, while preserving ALL other content including medical history tables.
+
+**CRITICAL: You should ONLY remove formal medication lists with explicit medication names, dosages, and frequencies.**
+
+**What to REMOVE (medication lists only):**
+1. **Current medication list** - sections with headers like "Huidige thuismedicatie", "Current medications", "Medicatie"
+2. **Discharge medication list** - sections with headers like "Als verdere behandeling stellen wij voor", "Thuismedicatie", "Discharge medications"
+
+**What medication lists look like:**
+- Header: "Huidige thuismedicatie" or similar
+- Followed by multiple lines with medication names, dosages, frequencies
+- Example: "Pantomed 20mg Tablet Oral - 1 tablet - 2 maal daags"
+- Example: "Forlax 10g Zakje Oral - 2 zakje - 1 maal daags (zo nodig)"
+
+**What to ABSOLUTELY NEVER REMOVE:**
+1. **Medical history tables** - Tables with "Datum" and "Bespreking" columns containing dates and medical events
+2. **Treatment history** - Narrative descriptions of medical procedures, treatments, or events
+3. **Clinical discussions** - Any text discussing medical conditions, procedures, or clinical decisions
+4. **Tables with dates and procedures** - Any table format showing timeline of medical events
+5. **Individual medication mentions in clinical text** - References to medications within clinical narratives
+
+**EXAMPLES OF CONTENT TO NEVER REMOVE:**
+- Tables like: "| Datum | Bespreking |" followed by medical events
+- "| 07/07/2017 | Niertransplantatie met donornier..."
+- "| 15/8/2017 | Uitgestelde transplantfunctie..."
+- "| 26/03/2018 | plaatsing peritoneaal dialysekatheter..."
+- Any text describing medical procedures, surgeries, or treatments
+- Clinical narratives mentioning medications in context (e.g., "behandeling met Sotrovimab")
+
+**KEY DISTINGUISHING FEATURES:**
+- **Medication lists**: Standalone sections with drug names + dosages + frequencies
+- **Medical history**: Tables or narratives describing medical events, procedures, surgeries
+- **Clinical text**: Discussions of treatment decisions, medical events, or conditions
+
+**If you see a table with dates and medical procedures, it is MEDICAL HISTORY, not a medication list.**
+**If you see clinical text discussing treatments or procedures, it is CLINICAL DISCUSSION, not a medication list.**
 
 Document structure:
 {text_analysis}
 
 **Analysis Instructions:**
-1. Look for formal medication sections with clear headers (e.g., "Thuismedicatie", "Als verdere behandeling stellen wij voor")
-2. Identify complete medication lists that contain multiple drugs with dosages
-3. **IGNORE** any medication mentions that appear within clinical discussion or narrative text
-4. Focus on structural elements that represent formal medication documentation
-5. Be conservative - if in doubt, do NOT remove
-6. Consider the position in the document (beginning/end vs. middle)
-
-**Examples of what to REMOVE:**
-- Complete lists under "Thuismedicatie" header
-- Formal medication lists under "Als verdere behandeling stellen wij voor"
-- Standalone medication sections with multiple drugs
-- Lists that appear at the beginning or end of the document
-
-**Examples of what to KEEP:**
-- "Patient was treated with Eliquis 2x 2.5mg" (clinical discussion)
-- "Stop Clopidogrel bij opname" (clinical instruction)
-- "Jardiance 10mg & Burinex 5mg" (if mentioned in clinical context)
-- Any medication mentioned in the context of treatment discussion
-
-Return your analysis as JSON:
+1. Look ONLY for formal medication sections with clear headers (e.g., "Thuismedicatie", "Huidige thuismedicatie")
+2. Identify sections that contain LISTS of medications with dosages and frequencies
+3. **NEVER identify medical history tables as medication lists**
+4. **NEVER identify clinical discussions as medication lists**
+5. Be extremely conservative - if in doubt, do NOT remove
+6. Focus ONLY on standalone medication documentation sections
+
+Return your analysis as a JSON object with this exact structure:
 {{
     "indices_to_remove": [list of integer indices - ONLY formal medication lists],
     "reasoning": {{
-        "formal_medication_lists": [list of identified formal medication list indices with explanations],
-        "clinical_medication_mentions": [list of clinical mentions that were correctly preserved],
-        "justification": "explanation of why only formal lists were selected for removal",
+        "justification": "explanation of why only formal medication lists were selected for removal",
         "confidence": "high/medium/low"
     }}
 }}
 """
+
         logger.info(f"Prompt length: {len(prompt)}")
         logger.info(f"Number of text elements: {len(text_analysis)}")
+        
         try:
             response = self.client.chat.completions.create(
-            messages=[
-                {
-                    "role": "system",
-                    "content": "You are a helpful assistant.",
-                },
-                {
-                    "role": "user",
-                    "content": prompt,
-                }
-            ],
-            max_completion_tokens=100000,  # adjust as needed
-            model=self.deployment
+                messages=[
+                    {
+                        "role": "system",
+                        "content": "You are a helpful assistant that analyzes medical documents and identifies formal medication lists for redaction.",
+                    },
+                    {
+                        "role": "user",
+                        "content": prompt,
+                    }
+                ],
+                max_completion_tokens=100000,
+                model=self.deployment,
+                response_format={"type": "json_object"}
             )
+            
+            # Record token usage and cost
+            if hasattr(response, 'usage') and response.usage:
+                cost_tracker.record_usage(
+                    prompt_tokens=response.usage.prompt_tokens,
+                    completion_tokens=response.usage.completion_tokens,
+                    model=self.model_name
+                )
+                logger.info(f"API call completed - Input: {response.usage.prompt_tokens}, "
+                           f"Output: {response.usage.completion_tokens}, "
+                           f"Total: {response.usage.total_tokens} tokens")
+            
         except Exception as e:
             logger.error(f"Exception during LLM call: {e}", exc_info=True)
             return {"indices_to_remove": [], "reasoning": {"confidence": "low"}}
 
         try:
-            logger.error(f"Raw LLM response: {response.choices[0].message.content!r}")
+            logger.info(f"Raw LLM response: {response.choices[0].message.content!r}")
+            
+            # Parse the structured JSON response
             result = json.loads(response.choices[0].message.content)
             
-            # Validate and limit the number of elements to remove
+            # Get the indices to remove
             indices_to_remove = result.get("indices_to_remove", [])
             
-            # Be conservative - limit to maximum 10 elements to prevent over-removal
-            if len(indices_to_remove) > 10:
-                logger.warning(f"LLM suggested removing {len(indices_to_remove)} elements, limiting to 10 most likely formal medication lists")
-                # Keep only the first 10 (assuming they're ordered by importance)
-                indices_to_remove = indices_to_remove[:10]
-                result["indices_to_remove"] = indices_to_remove
-                result["reasoning"]["justification"] += " [LIMITED: Only top 10 elements selected to prevent over-removal]"
+            # Log what the LLM suggested
+            logger.info(f"LLM suggested removing {len(indices_to_remove)} elements: {indices_to_remove}")
             
-            # Log the reasoning for transparency
-            reasoning = result.get("reasoning", {})
-            logger.info(f"LLM reasoning: {reasoning}")
+            # Log detailed information about what's being removed
+            if indices_to_remove:
+                logger.info("DETAILED ANALYSIS OF LLM SUGGESTIONS:")
+                logger.info("=" * 60)
+                
+                for idx in indices_to_remove:
+                    if idx < len(text_analysis):
+                        text_content = text_analysis[idx].get("text", "")
+                        text_label = text_analysis[idx].get("label", "")
+                        logger.info(f"Index {idx} ({text_label}): '{text_content}'")
+                    else:
+                        logger.error(f"Index {idx} is out of bounds (max: {len(text_analysis)-1})")
+                
+                logger.info("=" * 60)
+                
+                # Log the reasoning if provided
+                reasoning = result.get("reasoning", {})
+                if reasoning:
+                    logger.info(f"LLM reasoning: {reasoning}")
+                
+                logger.info(f"Final removal list: {len(indices_to_remove)} elements will be removed")
+            else:
+                logger.info("No elements will be removed")
             
             return result
+            
         except Exception as e:
             logger.error(f"Failed to parse LLM response: {e}")
             return {"indices_to_remove": [], "reasoning": {"confidence": "low"}}

src/processing/sections.py

@@ -21,18 +21,57 @@ class ReasoningSectionExtractor:
         reasoning = extraction_result.get("reasoning", {})
         
         # Log detailed reasoning for transparency
-        logger.info(f"O1-mini reasoning: {reasoning}")
+        logger.info(f"LLM reasoning summary: {reasoning}")
+        
+        # Get the texts for detailed logging
+        texts = doc_json.get("texts", [])
         
         # Provide specific feedback about what was removed
         if indices_to_remove:
             logger.info(f"Removing {len(indices_to_remove)} text elements: {indices_to_remove}")
             
-            # Show what specific content is being removed
-            texts = doc_json.get("texts", [])
+            # Categorize and show what specific content is being removed
+            medication_headers = []
+            medication_items = []
+            other_content = []
+            
             for idx in indices_to_remove:
                 if idx < len(texts):
-                    text_content = texts[idx].get("text", "")[:100]
-                    logger.info(f"  Removing text {idx}: '{text_content}...'")
+                    text_content = texts[idx].get("text", "")
+                    text_label = texts[idx].get("label", "")
+                    
+                    # Categorize the content
+                    if any(keyword in text_content.lower() for keyword in ['medicatie', 'thuismedicatie', 'medication', 'drugs']):
+                        medication_headers.append((idx, text_content))
+                    elif any(keyword in text_content.lower() for keyword in ['tablet', 'capsule', 'mg', 'ml', 'zakje', 'oral', 'maal daags']):
+                        medication_items.append((idx, text_content))
+                    else:
+                        other_content.append((idx, text_content))
+                    
+                    # Log with more detail
+                    logger.info(f"  → Removing index {idx} ({text_label}): '{text_content[:150]}{'...' if len(text_content) > 150 else ''}'")
+                else:
+                    logger.warning(f"  → Invalid index {idx}: exceeds document length ({len(texts)})")
+            
+            # Summary of what was categorized
+            if medication_headers:
+                logger.info(f"Medication headers removed: {len(medication_headers)} items")
+                for idx, content in medication_headers:
+                    logger.info(f"    Header {idx}: {content}")
+            
+            if medication_items:
+                logger.info(f"Medication items removed: {len(medication_items)} items")
+                for idx, content in medication_items[:5]:  # Show first 5 to avoid spam
+                    logger.info(f"    Item {idx}: {content[:100]}...")
+                if len(medication_items) > 5:
+                    logger.info(f"    ... and {len(medication_items) - 5} more medication items")
+            
+            if other_content:
+                logger.warning(f"⚠️  NON-MEDICATION content removed: {len(other_content)} items")
+                for idx, content in other_content:
+                    logger.warning(f"    ⚠️  Index {idx}: {content[:200]}...")
+                logger.warning("⚠️  Please review: non-medication content was removed - this may indicate an issue with the LLM detection")
+                
         else:
             logger.info("No formal medication lists identified for removal")
         
@@ -46,6 +85,7 @@ class ReasoningSectionExtractor:
         # Log the result
         removed_count = len(texts) - len(redacted_texts)
         logger.info(f"Successfully removed {removed_count} text elements from document structure")
+        logger.info(f"Document structure: {len(texts)} → {len(redacted_texts)} text elements")
         
         return redacted_json
     

src/streamlit_app.py

@@ -3,6 +3,8 @@
 
 import os
 import tempfile
+import json
+from datetime import datetime
 
 # Get a writable temp directory first
 try:
@@ -91,26 +93,12 @@ directories_to_create = [
     os.environ['ACCELERATE_CACHE'],
 ]
 
-# Monkey patch os.makedirs to prevent root directory access
-original_makedirs = os.makedirs
-
-def safe_makedirs(name, mode=0o777, exist_ok=False):
-    """Safe version of makedirs that prevents root directory access."""
-    # Check if trying to create directory in root filesystem
-    if name.startswith('/') and not name.startswith('/tmp') and not name.startswith('/app'):
-        # Redirect to temp directory
-        basename = os.path.basename(name)
-        safe_name = os.path.join(TEMP_DIR, basename)
-        print(f"Redirecting root directory creation from {name} to {safe_name}")
-        return original_makedirs(safe_name, mode, exist_ok)
-    return original_makedirs(name, mode, exist_ok)
-
-# Apply the monkey patch
-os.makedirs = safe_makedirs
-
 for directory in directories_to_create:
     try:
-        os.makedirs(directory, exist_ok=True)
+        # Create directory and all parent directories
+        os.makedirs(directory, mode=0o777, exist_ok=True)
+        # Ensure the directory has write permissions
+        os.chmod(directory, 0o777)
     except Exception as e:
         print(f"Warning: Could not create directory {directory}: {e}")
 
@@ -121,11 +109,10 @@ import shutil
 from processing.document_processor import DocumentProcessor
 from processing.sections import ReasoningSectionExtractor
 from utils.logging_utils import get_log_handler
+from utils.cost_tracker import cost_tracker
 from dotenv import load_dotenv
 import sys
-import html
 import difflib
-import re
 import time
 
 # Configure logging early to avoid issues
@@ -212,16 +199,40 @@ def get_temp_files_info():
         
         files = os.listdir(TEMP_DIR)
         total_size = 0
+        file_details = []
         
         for filename in files:
             try:
                 file_path = os.path.join(TEMP_DIR, filename)
                 if os.path.isfile(file_path):
-                    total_size += os.path.getsize(file_path)
+                    file_size = os.path.getsize(file_path)
+                    total_size += file_size
+                    file_details.append({
+                        'name': filename,
+                        'size': file_size,
+                        'type': 'file'
+                    })
+                elif os.path.isdir(file_path):
+                    file_details.append({
+                        'name': filename,
+                        'size': 0,
+                        'type': 'directory'
+                    })
             except (PermissionError, OSError) as e:
                 logging.warning(f"Error accessing file {filename}: {e}")
+                file_details.append({
+                    'name': filename,
+                    'size': 0,
+                    'type': 'error'
+                })
                 continue
         
+        # Log detailed information for debugging
+        if file_details:
+            logging.info(f"Temp directory contents ({TEMP_DIR}):")
+            for detail in file_details:
+                logging.info(f"  - {detail['name']} ({detail['type']}): {detail['size']} bytes")
+        
         return len(files), total_size
     except PermissionError as e:
         logging.warning(f"Permission error accessing temp directory {TEMP_DIR}: {e}")
@@ -298,215 +309,7 @@ st.markdown("""
     border-radius: 10px;
     border: 1px solid #e9ecef;
 }
-
-/* Synchronized scrolling styles */
-.sync-scroll-container {
-    display: flex;
-    gap: 20px;
-    height: 600px;
-    font-family: 'Courier New', monospace;
-    font-size: 12px;
-}
-
-.sync-scroll-panel {
-    flex: 1;
-    border: 1px solid #ddd;
-    border-radius: 5px;
-    overflow: hidden;
-    display: flex;
-    flex-direction: column;
-}
-
-.sync-scroll-header {
-    background-color: #f8f9fa;
-    padding: 10px;
-    border-bottom: 1px solid #ddd;
-    font-weight: bold;
-}
-
-.sync-scroll-content {
-    flex: 1;
-    overflow-y: auto;
-    padding: 10px;
-    background-color: #fff;
-    scroll-behavior: smooth;
-    transition: scroll-top 0.1s ease-out;
-}
-
-/* Prevent scroll chaining */
-.sync-scroll-content::-webkit-scrollbar {
-    width: 8px;
-}
-
-.sync-scroll-content::-webkit-scrollbar-track {
-    background: #f1f1f1;
-}
-
-.sync-scroll-content::-webkit-scrollbar-thumb {
-    background: #888;
-    border-radius: 4px;
-}
-
-.sync-scroll-content::-webkit-scrollbar-thumb:hover {
-    background: #555;
-}
 </style>
-
-<script>
-// Improved synchronized scrolling implementation with better debugging
-console.log('Starting sync scroll setup...');
-
-function setupSyncScroll() {
-    console.log('setupSyncScroll called');
-    
-    // Wait for elements to be available
-    setTimeout(function() {
-        console.log('Looking for scroll elements...');
-        const originalContent = document.getElementById('original-content');
-        const redactedContent = document.getElementById('redacted-content');
-        
-        console.log('Original content element:', originalContent);
-        console.log('Redacted content element:', redactedContent);
-        
-        if (originalContent && redactedContent) {
-            console.log('Both elements found, setting up sync...');
-            
-            let isScrolling = false;
-            let scrollTimeout;
-            
-            function syncScroll(source, target) {
-                if (!isScrolling) {
-                    isScrolling = true;
-                    console.log('Syncing scroll from', source.id, 'to', target.id, 'scrollTop:', source.scrollTop);
-                    target.scrollTop = source.scrollTop;
-                    
-                    // Clear existing timeout
-                    if (scrollTimeout) {
-                        clearTimeout(scrollTimeout);
-                    }
-                    
-                    // Reset flag after a short delay
-                    scrollTimeout = setTimeout(() => { 
-                        isScrolling = false; 
-                        console.log('Scroll sync completed');
-                    }, 100);
-                }
-            }
-            
-            // Remove existing listeners to prevent duplicates
-            if (originalContent._syncScrollHandler) {
-                originalContent.removeEventListener('scroll', originalContent._syncScrollHandler);
-            }
-            if (redactedContent._syncScrollHandler) {
-                redactedContent.removeEventListener('scroll', redactedContent._syncScrollHandler);
-            }
-            
-            // Create new handlers
-            originalContent._syncScrollHandler = function(e) {
-                console.log('Original content scrolled:', e.target.scrollTop);
-                syncScroll(originalContent, redactedContent);
-            };
-            
-            redactedContent._syncScrollHandler = function(e) {
-                console.log('Redacted content scrolled:', e.target.scrollTop);
-                syncScroll(redactedContent, originalContent);
-            };
-            
-            // Add event listeners
-            originalContent.addEventListener('scroll', originalContent._syncScrollHandler, { passive: true });
-            redactedContent.addEventListener('scroll', redactedContent._syncScrollHandler, { passive: true });
-            
-            console.log('Event listeners added successfully');
-            
-            // Show status indicator
-            const statusElement = document.getElementById('sync-status');
-            if (statusElement) {
-                statusElement.style.display = 'block';
-                console.log('Status indicator shown');
-            }
-            
-            // Test the synchronization
-            setTimeout(() => {
-                console.log('Testing scroll sync...');
-                console.log('Original scrollTop:', originalContent.scrollTop);
-                console.log('Redacted scrollTop:', redactedContent.scrollTop);
-                
-                // Try a small scroll to test
-                originalContent.scrollTop = 10;
-                setTimeout(() => {
-                    console.log('After test scroll - Original:', originalContent.scrollTop, 'Redacted:', redactedContent.scrollTop);
-                }, 50);
-            }, 200);
-            
-        } else {
-            console.log('Elements not found, will retry...');
-            // Retry with exponential backoff
-            setTimeout(setupSyncScroll, 300);
-        }
-    }, 200);
-}
-
-// Multiple initialization strategies
-function initializeSyncScroll() {
-    console.log('Initializing sync scroll...');
-    
-    // Strategy 1: Immediate setup
-    setupSyncScroll();
-    
-    // Strategy 2: Setup after DOM ready
-    if (document.readyState === 'loading') {
-        document.addEventListener('DOMContentLoaded', function() {
-            console.log('DOM loaded, setting up sync scroll...');
-            setupSyncScroll();
-        });
-    }
-    
-    // Strategy 3: Setup after window load
-    window.addEventListener('load', function() {
-        console.log('Window loaded, setting up sync scroll...');
-        setupSyncScroll();
-    });
-    
-    // Strategy 4: Periodic retry for first 10 seconds
-    let attempts = 0;
-    const maxAttempts = 20;
-    const retryInterval = setInterval(function() {
-        attempts++;
-        console.log('Retry attempt', attempts);
-        
-        const originalContent = document.getElementById('original-content');
-        const redactedContent = document.getElementById('redacted-content');
-        
-        if (originalContent && redactedContent) {
-            console.log('Elements found on retry, setting up...');
-            setupSyncScroll();
-            clearInterval(retryInterval);
-        } else if (attempts >= maxAttempts) {
-            console.log('Max retry attempts reached, giving up');
-            clearInterval(retryInterval);
-        }
-    }, 500);
-}
-
-// Start initialization
-initializeSyncScroll();
-
-// Listen for Streamlit-specific events
-if (window.parent && window.parent.postMessage) {
-    console.log('Streamlit environment detected');
-    
-    // Listen for any messages that might indicate a rerun
-    window.addEventListener('message', function(event) {
-        console.log('Received message:', event.data);
-        if (event.data && (event.data.type === 'streamlit:rerun' || event.data.type === 'streamlit:setComponentValue')) {
-            console.log('Streamlit rerun detected, reinitializing sync scroll...');
-            setTimeout(setupSyncScroll, 1000);
-        }
-    });
-}
-
-console.log('Sync scroll script loaded');
-</script>
 """, unsafe_allow_html=True)
 
 # Configure root logger only once (avoid duplicate handlers on reruns)
@@ -528,6 +331,7 @@ Use the buttons below to view the original structure or process with redaction.
 if st.button("🧹 Clear All Data", type="secondary", help="Remove all temporary files and reset the application"):
     if clear_all_data():
         st.success("✅ All data cleared successfully! The application has been reset.")
+        cost_tracker.reset_session()  # Reset cost tracking when clearing data
         st.rerun()
     else:
         st.error("❌ Error clearing data. Please try again.")
@@ -578,6 +382,32 @@ with col1:
         # Show warning if total size is large
         if total_size > 50 * 1024 * 1024:  # 50MB
             st.warning("⚠️ Large temporary files detected. Consider clearing data to free up space.")
+        
+        # Debug: Show temp files (expandable)
+        with st.expander("🔍 Debug: View temporary files"):
+            try:
+                if os.path.exists(TEMP_DIR):
+                    files = os.listdir(TEMP_DIR)
+                    if files:
+                        st.write("**Temporary files in directory:**")
+                        for filename in files:
+                            file_path = os.path.join(TEMP_DIR, filename)
+                            try:
+                                if os.path.isfile(file_path):
+                                    size = os.path.getsize(file_path)
+                                    st.write(f"📄 {filename} ({format_file_size(size)})")
+                                elif os.path.isdir(file_path):
+                                    st.write(f"📁 {filename} (directory)")
+                                else:
+                                    st.write(f"❓ {filename} (unknown)")
+                            except Exception as e:
+                                st.write(f"❌ {filename} (error: {e})")
+                    else:
+                        st.write("No files found in temp directory")
+                else:
+                    st.write("Temp directory does not exist")
+            except Exception as e:
+                st.write(f"Error accessing temp directory: {e}")
     else:
         st.caption("📁 No temporary files")
 
@@ -593,72 +423,6 @@ with col2:
     else:
         st.caption("No files to delete")
 
-def create_diff_content(original_text: str, redacted_text: str, view_type: str) -> str:
-    """Create HTML content for diff view with highlighting."""
-    import difflib
-    import re
-    
-    # Normalize the text to reduce formatting differences
-    def normalize_text(text):
-        # Remove extra whitespace and normalize line endings
-        lines = text.split('\n')
-        normalized_lines = []
-        for line in lines:
-            # Strip whitespace but preserve content
-            stripped = line.strip()
-            if stripped:
-                # Normalize header formatting differences
-                # Convert ## to # for level 1 headers
-                if re.match(r'^##\s+', stripped):
-                    stripped = re.sub(r'^##\s+', '# ', stripped)
-                # Normalize quote formatting
-                if stripped.startswith('&gt; '):
-                    stripped = stripped.replace('&gt; ', '> ')
-                elif stripped.startswith('+ > '):
-                    stripped = stripped.replace('+ > ', '> ')
-                
-                normalized_lines.append(stripped)
-        return normalized_lines
-    
-    original_lines = normalize_text(original_text)
-    redacted_lines = normalize_text(redacted_text)
-    
-    # Use difflib to get a more sophisticated diff
-    differ = difflib.Differ()
-    diff = list(differ.compare(original_lines, redacted_lines))
-    
-    html_lines = []
-    
-    if view_type == 'original':
-        # Show original with removed content highlighted
-        for line in diff:
-            if line.startswith('  '):  # Unchanged line
-                escaped_line = html.escape(line[2:])
-                html_lines.append(f'<div style="padding: 2px 4px; margin: 1px 0; color: #333;">{escaped_line}</div>')
-            elif line.startswith('- '):  # Removed line
-                escaped_line = html.escape(line[2:])
-                html_lines.append(f'<div style="background-color: #ffebee; color: #c62828; padding: 2px 4px; margin: 1px 0; border-left: 3px solid #f44336; font-weight: bold;">- {escaped_line}</div>')
-            elif line.startswith('+ '):  # Added line (show as empty space in original view)
-                html_lines.append(f'<div style="background-color: #e8f5e8; color: #2e7d32; padding: 2px 4px; margin: 1px 0; border-left: 3px solid #4caf50; font-style: italic; opacity: 0.7;">+ (added in redacted version)</div>')
-            elif line.startswith('? '):  # Ignore difflib hints
-                continue
-    
-    elif view_type == 'redacted':
-        # Show redacted content with added content highlighted
-        for line in diff:
-            if line.startswith('  '):  # Unchanged line
-                escaped_line = html.escape(line[2:])
-                html_lines.append(f'<div style="padding: 2px 4px; margin: 1px 0; color: #333;">{escaped_line}</div>')
-            elif line.startswith('- '):  # Removed line (show as empty space in redacted view)
-                html_lines.append(f'<div style="background-color: #ffebee; color: #c62828; padding: 2px 4px; margin: 1px 0; border-left: 3px solid #f44336; font-style: italic; opacity: 0.7;">- (removed from original)</div>')
-            elif line.startswith('+ '):  # Added line
-                escaped_line = html.escape(line[2:])
-                html_lines.append(f'<div style="background-color: #e8f5e8; color: #2e7d32; padding: 2px 4px; margin: 1px 0; border-left: 3px solid #4caf50; font-weight: bold;">+ {escaped_line}</div>')
-            elif line.startswith('? '):  # Ignore difflib hints
-                continue
-    
-    return '\n'.join(html_lines)
-
 if uploaded_files:
     # UI to select which file to work with (if multiple files uploaded)
     file_names = [f.name for f in uploaded_files]
@@ -698,32 +462,42 @@ if uploaded_files:
                     # Save uploaded file to a temporary location
                     temp_path = save_uploaded_file(uploaded_file, selected_file)
                     
-                    # Create a DocumentProcessor with a SectionExtractor for our target sections
-                    section_extractor = ReasoningSectionExtractor(
-                        endpoint=AZURE_OPENAI_ENDPOINT,
-                        api_key=AZURE_OPENAI_KEY,
-                        api_version=AZURE_OPENAI_VERSION,
-                        deployment=AZURE_OPENAI_DEPLOYMENT,
-                    )
-                    processor = DocumentProcessor(section_extractor=section_extractor)
+                    # Ensure the deployment name is in the cost tracker
+                    if AZURE_OPENAI_DEPLOYMENT and AZURE_OPENAI_DEPLOYMENT not in cost_tracker.get_available_models():
+                        model_type = cost_tracker.guess_model_type(AZURE_OPENAI_DEPLOYMENT)
+                        cost_tracker.add_deployment_pricing(AZURE_OPENAI_DEPLOYMENT, model_type)
+                    
+                    # Use the new processing function
+                    from processing.document_processor import process_document_with_redaction
                     
                     # Attach an in-memory log handler to capture logs for this file
                     log_handler, log_buffer = get_log_handler()
                     root_logger = logging.getLogger()
                     root_logger.addHandler(log_handler)
+                    
                     try:
-                        # Process the document (Docling parse + section redaction)
-                        result = processor.process(temp_path)
+                        # Process the document using the new function
+                        processing_result = process_document_with_redaction(
+                            file_path=temp_path,
+                            endpoint=AZURE_OPENAI_ENDPOINT,
+                            api_key=AZURE_OPENAI_KEY,
+                            api_version=AZURE_OPENAI_VERSION,
+                            deployment=AZURE_OPENAI_DEPLOYMENT,
+                        )
+                        
+                        # Save results in session state (maintaining compatibility with existing UI)
+                        st.session_state.processed_results[selected_file] = {
+                            "structured_json": processing_result.original_document_json,
+                            "redacted_md": processing_result.redacted_document_md,
+                            "redacted_json": processing_result.redacted_document_json,  # Now this is actually redacted!
+                            "original_markdown": processing_result.original_document_md,
+                            "processing_result": processing_result  # Store the new result
+                        }
+                        
                     finally:
                         # Remove handler and stop capturing logs
                         root_logger.removeHandler(log_handler)
                     
-                    # Save results in session state
-                    st.session_state.processed_results[selected_file] = {
-                        "structured_json": result.structured_json,
-                        "redacted_md": result.redacted_markdown,
-                        "redacted_json": result.redacted_json
-                    }
                     # Combine log records into a single text
                     log_text = "\n".join(log_buffer)
                     st.session_state.logs[selected_file] = log_text
@@ -807,25 +581,7 @@ if uploaded_files:
             structured_json = data["structured_json"]
             redacted_md = data["redacted_md"]
             redacted_json = data["redacted_json"]
-            
-            # Get the original markdown from the structured JSON
-            # We need to reconstruct the original markdown from the structured JSON
-            # For now, we'll use the structured_markdown from the DocumentResult
-            # But we need to store this in the session state
-            
-            # Create a DocumentProcessor to get the original markdown
-            if "original_markdown" not in st.session_state.processed_results[selected_file]:
-                # Save uploaded file to a temporary location
-                temp_path = save_uploaded_file(uploaded_file, selected_file)
-                
-                # Create a DocumentProcessor without section extraction to get original markdown
-                processor = DocumentProcessor(section_extractor=None)
-                result = processor.process(temp_path)
-                
-                # Store the original markdown
-                st.session_state.processed_results[selected_file]["original_markdown"] = result.structured_markdown
-            
-            original_md = st.session_state.processed_results[selected_file]["original_markdown"]
+            original_md = data["original_markdown"]
             
             # Show processing summary
             original_texts = structured_json.get("texts", [])
@@ -844,67 +600,230 @@ if uploaded_files:
                 st.subheader("Original vs Redacted Content")
                 st.caption("Compare the original document content with the redacted version")
                 
-                # Add status indicator
-                st.markdown("""
-                <div id="sync-status" style="padding: 8px; background-color: #e8f5e8; border: 1px solid #4caf50; border-radius: 4px; margin-bottom: 10px; display: none;">
-                    ✅ <strong>Synchronized scrolling is active</strong> - Scroll either panel to sync both views
-                </div>
-                """, unsafe_allow_html=True)
+                # Get the actual removed indices from the processing result
+                actual_removed_indices = []
+                if "processing_result" in st.session_state.processed_results[selected_file]:
+                    processing_result = st.session_state.processed_results[selected_file]["processing_result"]
+                    actual_removed_indices = processing_result.removed_indices
                 
-                # Create a diff-like interface with synchronized scrolling and highlighting
-                diff_html = f"""
-                <div class="sync-scroll-container">
-                    <div class="sync-scroll-panel">
-                        <div class="sync-scroll-header">
-                            📋 Original Document
-                        </div>
-                        <div id="original-content" class="sync-scroll-content">
-                            {create_diff_content(original_md, redacted_md, 'original')}
-                        </div>
-                    </div>
-                    <div class="sync-scroll-panel">
-                        <div class="sync-scroll-header">
-                            🔒 Redacted Document
-                        </div>
-                        <div id="redacted-content" class="sync-scroll-content">
-                            {create_diff_content(original_md, redacted_md, 'redacted')}
-                        </div>
-                    </div>
-                </div>
-                """
+                # Create a more intelligent side-by-side comparison based on JSON structure
+                col1, col2 = st.columns(2)
                 
-                st.markdown(diff_html, unsafe_allow_html=True)
+                with col1:
+                    st.markdown("**📋 Original Document**")
+                    
+                    # Display original content with removed sections highlighted
+                    for i, text_elem in enumerate(original_texts):
+                        text_content = text_elem.get("text", "")
+                        label = text_elem.get("label", "")
+                        
+                        # Check if this element was removed
+                        is_removed = i in actual_removed_indices
+                        
+                        if is_removed:
+                            # Highlight removed content in red
+                            st.markdown(f"""
+                            <div style="background-color: #ffebee; color: #c62828; padding: 8px; margin: 4px 0; border-left: 4px solid #f44336; border-radius: 4px;">
+                                <strong>Text {i} ({label}) - REMOVED:</strong><br>
+                                {text_content}
+                            </div>
+                            """, unsafe_allow_html=True)
+                        else:
+                            # Show normal content
+                            content_preview = text_content[:150] + "..." if len(text_content) > 150 else text_content
+                            st.markdown(f"""
+                            <div style="padding: 4px; margin: 2px 0; border-radius: 4px;">
+                                <strong>Text {i} ({label}) - {len(text_content)} chars:</strong><br>
+                                <code style="background-color: #f5f5f5; padding: 2px; border-radius: 2px;">{content_preview}</code>
+                            </div>
+                            """, unsafe_allow_html=True)
                 
-                # Add a hidden component to trigger JavaScript setup after Streamlit reruns
-                st.markdown("""
-                <script>
-                // Trigger setup after Streamlit rerun
-                if (window.parent && window.parent.postMessage) {
-                    // Wait for Streamlit to finish rendering
-                    setTimeout(function() {
-                        setupSyncScroll();
-                    }, 500);
-                }
-                </script>
-                """, unsafe_allow_html=True)
+                with col2:
+                    st.markdown("**🔒 Redacted Document**")
+                    
+                    # Display redacted content (only non-removed elements)
+                    redacted_index = 0
+                    for i, text_elem in enumerate(original_texts):
+                        text_content = text_elem.get("text", "")
+                        label = text_elem.get("label", "")
+                        
+                        # Check if this element was removed
+                        is_removed = i in actual_removed_indices
+                        
+                        if is_removed:
+                            # Show placeholder for removed content
+                            st.markdown(f"""
+                            <div style="background-color: #ffebee; color: #c62828; padding: 8px; margin: 4px 0; border-left: 4px solid #f44336; border-radius: 4px; font-style: italic; opacity: 0.7;">
+                                <strong>Text {i} ({label}) - REMOVED</strong><br>
+                                [Content removed by redaction]
+                            </div>
+                            """, unsafe_allow_html=True)
+                        else:
+                            # Show the actual content from redacted texts
+                            if redacted_index < len(redacted_texts):
+                                redacted_content = redacted_texts[redacted_index].get("text", "")
+                                content_preview = redacted_content[:150] + "..." if len(redacted_content) > 150 else redacted_content
+                                st.markdown(f"""
+                                <div style="padding: 4px; margin: 2px 0; border-radius: 4px;">
+                                    <strong>Text {i} ({label}) - {len(redacted_content)} chars:</strong><br>
+                                    <code style="background-color: #f5f5f5; padding: 2px; border-radius: 2px;">{content_preview}</code>
+                                </div>
+                                """, unsafe_allow_html=True)
+                                redacted_index += 1
+                            else:
+                                st.markdown(f"""
+                                <div style="padding: 4px; margin: 2px 0; border-radius: 4px; background-color: #f5f5f5;">
+                                    <strong>Text {i} ({label}):</strong><br>
+                                    [Content preserved]
+                                </div>
+                                """, unsafe_allow_html=True)
                 
-                  
-                # Add legend for the diff highlighting
+                # Add legend
                 st.markdown("---")
                 col1, col2 = st.columns(2)
                 with col1:
-                    st.markdown("**🎨 Diff Legend:**")
+                    st.markdown("**🎨 Comparison Legend:**")
                     st.markdown("🔴 **Red background** = Removed content")
-                    st.markdown("🟢 **Green background** = Added content")
-                    st.markdown("⚪ **White background** = Unchanged content")                
+                    st.markdown("⚪ **White background** = Preserved content")
+                    st.markdown("📝 **Italic text** = Placeholder for removed content")
                 
                 with col2:
-                    st.markdown("**💡 Tips:**")
-                    st.markdown("Look for red-highlighted sections")
-                    st.markdown("These show what was redacted")
-                    st.markdown("Use scroll to navigate long documents")
+                    st.markdown("**💡 How to read:**")
+                    st.markdown("Left panel shows original with removed sections highlighted")
+                    st.markdown("Right panel shows redacted version with placeholders")
+                    st.markdown("Compare corresponding text indices to see changes")
                 
-             
+                # Add debug information to help identify missing content
+                with st.expander("🔍 Debug: Content Analysis"):
+                    st.write("**Searching for table content...**")
+                    
+                    # Search for table-related content in original texts
+                    table_elements = []
+                    for i, text_elem in enumerate(original_texts):
+                        text_content = text_elem.get("text", "")
+                        label = text_elem.get("label", "")
+                        
+                        if "Bespreking" in text_content or "|" in text_content or "table" in label.lower():
+                            table_elements.append({
+                                "index": i,
+                                "label": label,
+                                "content": text_content[:200] + "..." if len(text_content) > 200 else text_content,
+                                "is_removed": i in actual_removed_indices
+                            })
+                    
+                    if table_elements:
+                        st.write(f"**Found {len(table_elements)} table-related elements:**")
+                        for elem in table_elements:
+                            status = "🔴 REMOVED" if elem["is_removed"] else "✅ PRESERVED"
+                            st.write(f"**Text {elem['index']} ({elem['label']}) - {status}:**")
+                            st.write(f"`{elem['content']}`")
+                            st.write("---")
+                    else:
+                        st.write("**No table-related content found in original texts**")
+                    
+                    # Also check redacted texts
+                    st.write("**Table content in redacted texts:**")
+                    table_elements_redacted = []
+                    for i, text_elem in enumerate(redacted_texts):
+                        text_content = text_elem.get("text", "")
+                        label = text_elem.get("label", "")
+                        
+                        if "Bespreking" in text_content or "|" in text_content or "table" in label.lower():
+                            table_elements_redacted.append({
+                                "index": i,
+                                "label": label,
+                                "content": text_content[:200] + "..." if len(text_content) > 200 else text_content
+                            })
+                    
+                    if table_elements_redacted:
+                        st.write(f"**Found {len(table_elements_redacted)} table-related elements in redacted content:**")
+                        for elem in table_elements_redacted:
+                            st.write(f"**Text {elem['index']} ({elem['label']}):**")
+                            st.write(f"`{elem['content']}`")
+                            st.write("---")
+                    else:
+                        st.write("**No table-related content found in redacted texts**")
+                
+                # Add download buttons for redacted content
+                st.markdown("---")
+                st.subheader("📥 Download Redacted Content")
+                
+                col1, col2, col3 = st.columns(3)
+                
+                with col1:
+                    # Download redacted markdown
+                    st.download_button(
+                        label="📄 Download Redacted Markdown",
+                        data=redacted_md,
+                        file_name=f"{selected_file}_redacted.md",
+                        mime="text/markdown",
+                        help="Download the redacted document as Markdown format"
+                    )
+                
+                with col2:
+                    # Generate and download redacted PDF
+                    pdf_generated = False
+                    pdf_bytes = None
+                    
+                    if st.button("📋 Generate Redacted PDF", help="Generate a PDF version of the redacted document"):
+                        with st.spinner("Generating redacted PDF..."):
+                            try:
+                                # Create a DocumentProcessor to access PDF generation
+                                temp_path = save_uploaded_file(uploaded_file, selected_file)
+                                processor = DocumentProcessor(section_extractor=None)
+                                
+                                # Generate PDF path
+                                base_name = os.path.splitext(selected_file)[0]
+                                pdf_path = os.path.join(TEMP_DIR, f"{base_name}_redacted.pdf")
+                                
+                                # Generate the PDF
+                                success = processor.generate_redacted_pdf(redacted_json, pdf_path)
+                                
+                                if success:
+                                    # Read the generated PDF and store for download
+                                    with open(pdf_path, "rb") as pdf_file:
+                                        pdf_bytes = pdf_file.read()
+                                    pdf_generated = True
+                                    st.success("✅ PDF generated successfully!")
+                                else:
+                                    st.error("❌ Failed to generate PDF. Check logs for details.")
+                                    
+                            except Exception as e:
+                                st.error(f"❌ Error generating PDF: {e}")
+                                st.info("💡 Make sure reportlab is installed: `pip install reportlab`")
+                    
+                    # Show download button if PDF was generated
+                    if pdf_generated and pdf_bytes:
+                        st.download_button(
+                            label="📥 Download Redacted PDF",
+                            data=pdf_bytes,
+                            file_name=f"{os.path.splitext(selected_file)[0]}_redacted.pdf",
+                            mime="application/pdf",
+                            help="Download the redacted document as PDF"
+                        )
+                        
+                        # Show debug information about what's in the PDF
+                        with st.expander("🔍 Debug: PDF Content Analysis"):
+                            st.write("**Content that will be included in the PDF:**")
+                            texts_in_pdf = redacted_json.get("texts", [])
+                            st.write(f"Total text elements: {len(texts_in_pdf)}")
+                            
+                            for i, text_elem in enumerate(texts_in_pdf):
+                                text_content = text_elem.get("text", "")[:100] + "..." if len(text_elem.get("text", "")) > 100 else text_elem.get("text", "")
+                                label = text_elem.get("label", "")
+                                st.write(f"**Text {i} ({label}):** {text_content}")
+                    elif not pdf_generated:
+                        st.info("💡 Click 'Generate Redacted PDF' to create a PDF version")
+                
+                with col3:
+                    # Download redacted JSON structure
+                    st.download_button(
+                        label="🔧 Download Redacted JSON",
+                        data=json.dumps(redacted_json, indent=2, ensure_ascii=False),
+                        file_name=f"{selected_file}_redacted.json",
+                        mime="application/json",
+                        help="Download the redacted document structure as JSON"
+                    )
             
             with tab2:
                 st.subheader("Document Structure Analysis")                
@@ -922,19 +841,139 @@ if uploaded_files:
             with tab3:
                 st.subheader("Processing Details")
                 
+                # Show cost analysis for this processing session
+                st.subheader("💰 Cost Analysis")
+                
+                # Get cost data from the processing result
+                if "processing_result" in st.session_state.processed_results[selected_file]:
+                    processing_result = st.session_state.processed_results[selected_file]["processing_result"]
+                    
+                    col1, col2, col3 = st.columns(3)
+                    with col1:
+                        st.metric("Total Cost", f"${processing_result.cost:.4f}")
+                    with col2:
+                        st.metric("Input Tokens", f"{processing_result.input_tokens:,}")
+                    with col3:
+                        st.metric("Output Tokens", f"{processing_result.output_tokens:,}")
+                    
+                    # Add download button for cost report
+                    cost_report = {
+                        "timestamp": datetime.now().isoformat(),
+                        "total_cost": processing_result.cost,
+                        "input_tokens": processing_result.input_tokens,
+                        "output_tokens": processing_result.output_tokens,
+                        "total_tokens": processing_result.input_tokens + processing_result.output_tokens,
+                        "document_processed": selected_file,
+                        "model_used": AZURE_OPENAI_DEPLOYMENT
+                    }
+                    
+                    st.download_button(
+                        label="📥 Download Cost Report (JSON)",
+                        data=json.dumps(cost_report, indent=2),
+                        file_name=f"cost_report_{selected_file}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
+                        mime="application/json"
+                    )
+                    
+                    # Show model information
+                    model_info = cost_tracker.get_model_info(AZURE_OPENAI_DEPLOYMENT)
+                    if model_info:
+                        st.subheader("Model Information")
+                        st.write(f"**Model:** {model_info.description}")
+                        st.write(f"**Input cost:** ${model_info.input_cost_per_1k_tokens:.4f}/1K tokens")
+                        st.write(f"**Output cost:** ${model_info.output_cost_per_1k_tokens:.4f}/1K tokens")
+                        
+                        # Calculate cost breakdown
+                        input_cost = (processing_result.input_tokens / 1000) * model_info.input_cost_per_1k_tokens
+                        output_cost = (processing_result.output_tokens / 1000) * model_info.output_cost_per_1k_tokens
+                        st.write(f"**Cost breakdown:** Input: ${input_cost:.4f}, Output: ${output_cost:.4f}")
+                else:
+                    # Fallback to old cost summary method
+                    cost_summary = cost_tracker.get_session_summary()
+                    
+                    if cost_summary["usage_count"] > 0:
+                        col1, col2, col3 = st.columns(3)
+                        with col1:
+                            st.metric("Total Cost", f"${cost_summary['total_cost']:.4f}")
+                        with col2:
+                            st.metric("Total Tokens", f"{cost_summary['total_tokens']:,}")
+                        with col3:
+                            st.metric("API Calls", cost_summary["usage_count"])
+                        
+                        # Add download button for cost report
+                        cost_report = {
+                            "timestamp": datetime.now().isoformat(),
+                            "total_cost": cost_summary["total_cost"],
+                            "total_tokens": cost_summary["total_tokens"],
+                            "api_calls": cost_summary["usage_count"],
+                            "model_breakdown": cost_summary["model_breakdown"],
+                            "document_processed": selected_file
+                        }
+                        
+                        st.download_button(
+                            label="📥 Download Cost Report (JSON)",
+                            data=json.dumps(cost_report, indent=2),
+                            file_name=f"cost_report_{selected_file}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
+                            mime="application/json"
+                        )
+                        
+                        # Show detailed model breakdown
+                        if cost_summary["model_breakdown"]:
+                            st.subheader("Model Usage Breakdown")
+                            for model, stats in cost_summary["model_breakdown"].items():
+                                model_info = cost_tracker.get_model_info(model)
+                                model_display_name = model_info.description if model_info else model
+                                
+                                with st.expander(f"{model_display_name} - ${stats['cost']:.4f}"):
+                                    col1, col2 = st.columns(2)
+                                    with col1:
+                                        st.write(f"**Input tokens:** {stats['input_tokens']:,}")
+                                        st.write(f"**Output tokens:** {stats['output_tokens']:,}")
+                                    with col2:
+                                        st.write(f"**Total tokens:** {stats['total_tokens']:,}")
+                                        st.write(f"**API calls:** {stats['usage_count']}")
+                                    
+                                    # Show cost breakdown
+                                    if model_info:
+                                        input_cost = (stats['input_tokens'] / 1000) * model_info.input_cost_per_1k_tokens
+                                        output_cost = (stats['output_tokens'] / 1000) * model_info.output_cost_per_1k_tokens
+                                        st.write(f"**Cost breakdown:** Input: ${input_cost:.4f}, Output: ${output_cost:.4f}")
+                    else:
+                        st.info("No API calls recorded for this session")
+                
                 # Show what was removed
                 if removed_count > 0:
                     st.info(f"**Removed {removed_count} text elements from the document structure.**")
                     
-                    # Show the removed text elements
+                    # Show the removed text elements - use the actual indices from the processing result
                     st.subheader("Removed Text Elements:")
-                    removed_texts = []
-                    for i, text_elem in enumerate(original_texts):
-                        if i >= len(redacted_texts) or text_elem.get("text", "") != redacted_texts[i].get("text", ""):
-                            removed_texts.append((i, text_elem.get("text", "")[:100] + "..." if len(text_elem.get("text", "")) > 100 else text_elem.get("text", "")))
                     
-                    for idx, text in removed_texts:
-                        st.text(f"Text {idx}: {text}")
+                    # Get the actual indices that were removed from the processing result
+                    if "processing_result" in st.session_state.processed_results[selected_file]:
+                        # Get the actual removed indices from the LLM response
+                        processing_result = st.session_state.processed_results[selected_file]["processing_result"]
+                        actual_removed_indices = processing_result.removed_indices
+                        
+                        if actual_removed_indices:
+                            st.info(f"**Elements removed by LLM analysis ({len(actual_removed_indices)} elements):**")
+                            
+                            for idx in actual_removed_indices:
+                                if idx < len(original_texts):
+                                    text_content = original_texts[idx].get("text", "")
+                                    st.text(f"Text {idx}: {text_content[:100]}{'...' if len(text_content) > 100 else ''}")
+                                else:
+                                    st.text(f"Text {idx}: [Index out of bounds]")
+                        else:
+                            st.info("**No elements were identified for removal by the LLM.**")
+                    else:
+                        # Fallback to the old method if processing result not available
+                        st.warning("**Note: Using fallback calculation method**")
+                        removed_texts = []
+                        for i, text_elem in enumerate(original_texts):
+                            if i >= len(redacted_texts) or text_elem.get("text", "") != redacted_texts[i].get("text", ""):
+                                removed_texts.append((i, text_elem.get("text", "")[:100] + "..." if len(text_elem.get("text", "")) > 100 else text_elem.get("text", "")))
+                        
+                        for idx, text in removed_texts:
+                            st.text(f"Text {idx}: {text}")
                 else:
                     st.info("No text elements were removed during processing.")
                 

src/utils/cost_tracker.py

@@ -0,0 +1,241 @@
+import logging
+from dataclasses import dataclass
+from typing import Dict, Optional
+from datetime import datetime
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class ModelPricing:
+    """Pricing information for Azure OpenAI models."""
+    model_name: str
+    input_cost_per_1k_tokens: float  # Cost per 1000 input tokens
+    output_cost_per_1k_tokens: float  # Cost per 1000 output tokens
+    description: str
+
+@dataclass
+class TokenUsage:
+    """Token usage statistics for a single API call."""
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    model: str
+    timestamp: datetime
+
+@dataclass
+class CostAnalysis:
+    """Cost analysis for document processing."""
+    total_input_tokens: int
+    total_output_tokens: int
+    total_cost: float
+    model_breakdown: Dict[str, Dict[str, float]]  # {model: {"input_cost": x, "output_cost": y, "total_cost": z}}
+    processing_time: float
+    timestamp: datetime
+
+class CostTracker:
+    """Tracks token usage and calculates costs for Azure OpenAI API calls."""
+    
+    # Hardcoded pricing for Azure OpenAI models (current as of 2024)
+    # Source: https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/
+    MODEL_PRICING = {
+        # Standard model names
+        "gpt-4o-mini": ModelPricing(
+            model_name="gpt-4o-mini",
+            input_cost_per_1k_tokens=0.00015,  # $0.00015 per 1K input tokens
+            output_cost_per_1k_tokens=0.0006,  # $0.0006 per 1K output tokens
+            description="GPT-4o Mini (O3 Mini)"
+        ),
+        "gpt-4o": ModelPricing(
+            model_name="gpt-4o",
+            input_cost_per_1k_tokens=0.0025,  # $0.0025 per 1K input tokens
+            output_cost_per_1k_tokens=0.01,   # $0.01 per 1K output tokens
+            description="GPT-4o (O4)"
+        ),
+        "gpt-35-turbo": ModelPricing(
+            model_name="gpt-35-turbo",
+            input_cost_per_1k_tokens=0.0005,  # $0.0005 per 1K input tokens
+            output_cost_per_1k_tokens=0.0015, # $0.0015 per 1K output tokens
+            description="GPT-3.5 Turbo (O3)"
+        ),
+        # Azure deployment names (custom names set in Azure)
+        "o3-mini": ModelPricing(
+            model_name="o3-mini",
+            input_cost_per_1k_tokens=0.00015,  # $0.00015 per 1K input tokens
+            output_cost_per_1k_tokens=0.0006,  # $0.0006 per 1K output tokens
+            description="O3 Mini (GPT-4o Mini)"
+        ),
+        "o4-mini": ModelPricing(
+            model_name="o4-mini",
+            input_cost_per_1k_tokens=0.00015,  # $0.00015 per 1K input tokens
+            output_cost_per_1k_tokens=0.0006,  # $0.0006 per 1K output tokens
+            description="O4 Mini (GPT-4o Mini)"
+        ),
+        "o3": ModelPricing(
+            model_name="o3",
+            input_cost_per_1k_tokens=0.0005,  # $0.0005 per 1K input tokens
+            output_cost_per_1k_tokens=0.0015, # $0.0015 per 1K output tokens
+            description="O3 (GPT-3.5 Turbo)"
+        ),
+        "o4": ModelPricing(
+            model_name="o4",
+            input_cost_per_1k_tokens=0.0025,  # $0.0025 per 1K input tokens
+            output_cost_per_1k_tokens=0.01,   # $0.01 per 1K output tokens
+            description="O4 (GPT-4o)"
+        ),
+        # Alternative model names that might be used in Azure deployments
+        "gpt-4o-mini-2024-07-18": ModelPricing(
+            model_name="gpt-4o-mini-2024-07-18",
+            input_cost_per_1k_tokens=0.00015,  # $0.00015 per 1K input tokens
+            output_cost_per_1k_tokens=0.0006,  # $0.0006 per 1K output tokens
+            description="GPT-4o Mini (O3 Mini) - Latest"
+        ),
+        "gpt-4o-2024-05-13": ModelPricing(
+            model_name="gpt-4o-2024-05-13",
+            input_cost_per_1k_tokens=0.0025,  # $0.0025 per 1K input tokens
+            output_cost_per_1k_tokens=0.01,   # $0.01 per 1K output tokens
+            description="GPT-4o (O4) - Latest"
+        ),
+        "gpt-35-turbo-0125": ModelPricing(
+            model_name="gpt-35-turbo-0125",
+            input_cost_per_1k_tokens=0.0005,  # $0.0005 per 1K input tokens
+            output_cost_per_1k_tokens=0.0015, # $0.0015 per 1K output tokens
+            description="GPT-3.5 Turbo (O3) - Latest"
+        ),
+    }
+    
+    def __init__(self):
+        self.usage_history: list[TokenUsage] = []
+        self.current_session_tokens = 0
+        self.current_session_cost = 0.0
+    
+    def record_usage(self, prompt_tokens: int, completion_tokens: int, model: str) -> TokenUsage:
+        """Record token usage from an API call."""
+        total_tokens = prompt_tokens + completion_tokens
+        usage = TokenUsage(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=total_tokens,
+            model=model,
+            timestamp=datetime.now()
+        )
+        
+        self.usage_history.append(usage)
+        self.current_session_tokens += total_tokens
+        
+        # Calculate cost for this usage
+        cost = self._calculate_cost(prompt_tokens, completion_tokens, model)
+        self.current_session_cost += cost
+        
+        logger.info(f"Recorded usage: {prompt_tokens} input + {completion_tokens} output = {total_tokens} total tokens "
+                   f"for model {model}, cost: ${cost:.6f}")
+        
+        return usage
+    
+    def _calculate_cost(self, input_tokens: int, output_tokens: int, model: str) -> float:
+        """Calculate cost for given token usage and model."""
+        if model not in self.MODEL_PRICING:
+            logger.warning(f"Unknown model pricing for {model}, using default pricing")
+            # Try to guess the model type based on the name
+            if "mini" in model.lower():
+                # Assume it's a mini model (cheapest)
+                model = "o3-mini"
+            elif "o4" in model.lower():
+                # Assume it's O4 (most expensive)
+                model = "o4"
+            elif "o3" in model.lower():
+                # Assume it's O3 (medium)
+                model = "o3"
+            else:
+                # Default to cheapest option
+                model = "o3-mini"
+        
+        pricing = self.MODEL_PRICING[model]
+        
+        input_cost = (input_tokens / 1000) * pricing.input_cost_per_1k_tokens
+        output_cost = (output_tokens / 1000) * pricing.output_cost_per_1k_tokens
+        
+        return input_cost + output_cost
+    
+    def get_session_summary(self) -> Dict[str, any]:
+        """Get summary of current session usage."""
+        if not self.usage_history:
+            return {
+                "total_tokens": 0,
+                "total_cost": 0.0,
+                "model_breakdown": {},
+                "usage_count": 0
+            }
+        
+        model_breakdown = {}
+        for usage in self.usage_history:
+            if usage.model not in model_breakdown:
+                model_breakdown[usage.model] = {
+                    "input_tokens": 0,
+                    "output_tokens": 0,
+                    "total_tokens": 0,
+                    "cost": 0.0,
+                    "usage_count": 0
+                }
+            
+            model_breakdown[usage.model]["input_tokens"] += usage.prompt_tokens
+            model_breakdown[usage.model]["output_tokens"] += usage.completion_tokens
+            model_breakdown[usage.model]["total_tokens"] += usage.total_tokens
+            model_breakdown[usage.model]["usage_count"] += 1
+            model_breakdown[usage.model]["cost"] += self._calculate_cost(
+                usage.prompt_tokens, usage.completion_tokens, usage.model
+            )
+        
+        return {
+            "total_tokens": self.current_session_tokens,
+            "total_cost": self.current_session_cost,
+            "model_breakdown": model_breakdown,
+            "usage_count": len(self.usage_history)
+        }
+    
+    def reset_session(self):
+        """Reset current session statistics."""
+        self.usage_history = []
+        self.current_session_tokens = 0
+        self.current_session_cost = 0.0
+        logger.info("Cost tracker session reset")
+    
+    def get_available_models(self) -> list[str]:
+        """Get list of available models with pricing."""
+        return list(self.MODEL_PRICING.keys())
+    
+    def get_model_info(self, model: str) -> Optional[ModelPricing]:
+        """Get pricing information for a specific model."""
+        return self.MODEL_PRICING.get(model)
+    
+    def add_deployment_pricing(self, deployment_name: str, model_type: str = "o3-mini"):
+        """Add pricing for a custom deployment name by mapping it to an existing model type."""
+        if deployment_name in self.MODEL_PRICING:
+            return  # Already exists
+        
+        # Map deployment name to existing model pricing
+        if model_type in self.MODEL_PRICING:
+            base_pricing = self.MODEL_PRICING[model_type]
+            self.MODEL_PRICING[deployment_name] = ModelPricing(
+                model_name=deployment_name,
+                input_cost_per_1k_tokens=base_pricing.input_cost_per_1k_tokens,
+                output_cost_per_1k_tokens=base_pricing.output_cost_per_1k_tokens,
+                description=f"{deployment_name} ({base_pricing.description})"
+            )
+            logger.info(f"Added pricing for deployment {deployment_name} based on {model_type}")
+        else:
+            logger.warning(f"Unknown model type {model_type} for deployment {deployment_name}")
+    
+    def guess_model_type(self, deployment_name: str) -> str:
+        """Guess the model type based on deployment name."""
+        deployment_lower = deployment_name.lower()
+        if "mini" in deployment_lower:
+            return "o3-mini"
+        elif "o4" in deployment_lower:
+            return "o4"
+        elif "o3" in deployment_lower:
+            return "o3"
+        else:
+            return "o3-mini"  # Default to cheapest
+
+# Global cost tracker instance
+cost_tracker = CostTracker() 

uv.lock

@@ -200,6 +200,7 @@ dependencies = [
     { name = "openai" },
     { name = "python-dotenv" },
     { name = "pyyaml" },
+    { name = "reportlab" },
     { name = "streamlit" },
 ]
 
@@ -209,6 +210,7 @@ requires-dist = [
     { name = "openai", specifier = ">=1.91.0" },
     { name = "python-dotenv", specifier = ">=1.1.1" },
     { name = "pyyaml", specifier = ">=6.0" },
+    { name = "reportlab", specifier = ">=4.4.2" },
     { name = "streamlit", specifier = ">=1.46.0" },
 ]
 
@@ -1333,6 +1335,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/45/94/bc295babb3062a731f52621cdc992d123111282e291abaf23faa413443ea/regex-2024.11.6-cp313-cp313-win_amd64.whl", hash = "sha256:2b3361af3198667e99927da8b84c1b010752fa4b1115ee30beaa332cabc3ef1a", size = 273545, upload-time = "2024-11-06T20:11:15Z" },
 ]
 
+[[package]]
+name = "reportlab"
+version = "4.4.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "charset-normalizer" },
+    { name = "pillow" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ec/9b/3483c7e4ad33d15f22d528872439e5bc92485814d7e7d10dbc3130368a83/reportlab-4.4.2.tar.gz", hash = "sha256:fc6283048ddd0781a9db1d671715990e6aa059c8d40ec9baf34294c4bd583a36", size = 3509063, upload-time = "2025-06-18T12:20:19.526Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/9f/74/ed990bc9586605d4e46f6b0e0b978a5b8e757aa599e39664bee26d6dc666/reportlab-4.4.2-py3-none-any.whl", hash = "sha256:58e11be387457928707c12153b7e41e52533a5da3f587b15ba8f8fd0805c6ee2", size = 1953624, upload-time = "2025-06-18T12:20:16.152Z" },
+]
+
 [[package]]
 name = "requests"
 version = "2.32.4"