analyze_documents.py

#!/usr/bin/env python3
"""
Document analysis script using LLM to generate summaries and key insights.
Groups pages into documents (like .eleventy.js) and analyzes each one.
"""

import os
import json
import re
from pathlib import Path
from typing import Dict, List, Optional
from collections import defaultdict
from openai import OpenAI
from tqdm import tqdm
from dotenv import load_dotenv


class DocumentAnalyzer:
    """Analyze grouped documents using LLM"""

    def __init__(self, api_url: str, api_key: str, model: str = "gpt-4o"):
        self.client = OpenAI(api_key=api_key, base_url=api_url)
        self.model = model
        self.results_dir = Path("./results")
        self.analyses_file = Path("./analyses.json")

    def normalize_doc_num(self, doc_num: Optional[str]) -> Optional[str]:
        """Normalize document number to handle LLM variations"""
        if not doc_num:
            return None
        return str(doc_num).lower().replace(r'[^a-z0-9-]', '-').replace(r'-+', '-').strip('-')

    def load_and_group_documents(self) -> List[Dict]:
        """Load all JSON files and group into documents (matching .eleventy.js logic)"""
        pages = []

        # Recursively read all JSON files
        for json_file in self.results_dir.glob("**/*.json"):
            try:
                with open(json_file, 'r', encoding='utf-8') as f:
                    content = json.load(f)

                    relative_path = json_file.relative_to(self.results_dir)
                    pages.append({
                        'path': str(relative_path),
                        'filename': json_file.stem,
                        'folder': str(relative_path.parent) if relative_path.parent != Path('.') else 'root',
                        **content
                    })
            except Exception as e:
                print(f"Warning: Could not load {json_file}: {e}")

        print(f"Loaded {len(pages)} pages")

        # Group by normalized document number
        document_map = defaultdict(list)

        for page in pages:
            doc_num = page.get('document_metadata', {}).get('document_number')

            if not doc_num:
                # Use filename as fallback
                normalized = self.normalize_doc_num(page['filename']) or page['filename']
            else:
                normalized = self.normalize_doc_num(doc_num)

            document_map[normalized].append(page)

        # Helper to extract numeric page number
        def get_page_num(page):
            page_num = page.get('document_metadata', {}).get('page_number', 0) or 0
            if isinstance(page_num, int):
                return page_num
            # Handle formats like "24 of 66" or "24/66"
            if isinstance(page_num, str):
                # Extract first number
                match = re.search(r'(\d+)', page_num)
                if match:
                    return int(match.group(1))
            return 0

        # Convert to sorted documents
        documents = []
        for normalized_num, doc_pages in document_map.items():
            # Sort pages by page number
            doc_pages.sort(key=get_page_num)

            # Get metadata
            first_page = doc_pages[0]
            raw_doc_nums = list(set(
                p.get('document_metadata', {}).get('document_number')
                for p in doc_pages
                if p.get('document_metadata', {}).get('document_number')
            ))

            # Combine full text from all pages
            full_text = '\n\n--- PAGE BREAK ---\n\n'.join(
                p.get('full_text', '') for p in doc_pages
            )

            # Collect all entities
            all_entities = {
                'people': set(),
                'organizations': set(),
                'locations': set(),
                'dates': set(),
                'reference_numbers': set()
            }

            for page in doc_pages:
                if 'entities' in page and page['entities']:
                    for key in all_entities.keys():
                        if key in page['entities'] and page['entities'][key]:
                            all_entities[key].update(page['entities'][key])

            documents.append({
                'unique_id': normalized_num,
                'document_number': raw_doc_nums[0] if len(raw_doc_nums) == 1 else normalized_num,
                'page_count': len(doc_pages),
                'full_text': full_text,
                'document_metadata': first_page.get('document_metadata', {}),
                'entities': {k: sorted(list(v)) for k, v in all_entities.items()}
            })

        print(f"Grouped into {len(documents)} documents")
        return sorted(documents, key=lambda d: d['document_number'])

    def get_analysis_prompt(self) -> str:
        """Get the system prompt for document analysis"""
        return """You are an expert legal document analyst specializing in court documents, depositions, and legal filings.

Analyze the provided document and return a concise summary with key insights.

Your analysis should include:
1. **Document Type**: What kind of document is this? (deposition, court filing, letter, email, affidavit, etc.)
2. **Key Topics**: What are the main subjects/topics discussed? (2-3 bullet points)
3. **Key People**: Who are the most important people mentioned and their roles?
4. **Significance**: Why is this document potentially important? What does it reveal or establish?
5. **Summary**: A 2-3 sentence summary of the document's content

Be factual, concise, and focus on what makes this document notable or significant.

Return ONLY valid JSON in this format:
{
  "document_type": "string",
  "key_topics": ["topic1", "topic2", "topic3"],
  "key_people": [
    {"name": "person name", "role": "their role or significance in this doc"}
  ],
  "significance": "Why this document matters (1-2 sentences)",
  "summary": "Brief summary (2-3 sentences)"
}"""

    def analyze_document(self, document: Dict) -> Optional[Dict]:
        """Analyze a single document using LLM"""
        try:
            # Limit text length for API (keep first ~8000 chars if too long)
            full_text = document['full_text']
            if len(full_text) > 8000:
                full_text = full_text[:8000] + "\n\n[... document continues ...]"

            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {
                        "role": "system",
                        "content": self.get_analysis_prompt()
                    },
                    {
                        "role": "user",
                        "content": f"Analyze this document:\n\n{full_text}"
                    }
                ],
                temperature=0.2,
                max_tokens=1000
            )

            content = response.choices[0].message.content.strip()

            # Extract JSON
            json_match = re.search(r'```(?:json)?\s*\n(.*?)\n```', content, re.DOTALL)
            if json_match:
                content = json_match.group(1).strip()
            else:
                json_match = re.search(r'\{.*\}', content, re.DOTALL)
                if json_match:
                    content = json_match.group(0).strip()

            analysis = json.loads(content)

            return {
                'document_id': document['unique_id'],
                'document_number': document['document_number'],
                'page_count': document['page_count'],
                'analysis': analysis
            }

        except Exception as e:
            print(f"Error analyzing document {document['document_number']}: {e}")
            return None

    def analyze_all(self, limit: Optional[int] = None) -> List[Dict]:
        """Analyze all documents"""
        print("=" * 60)
        print("DOCUMENT ANALYSIS")
        print("=" * 60)

        # Load existing analyses to resume
        existing_analyses = {}
        if self.analyses_file.exists():
            try:
                with open(self.analyses_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    existing_analyses = {a['document_id']: a for a in data.get('analyses', [])}
                print(f"Found {len(existing_analyses)} existing analyses")
            except Exception as e:
                print(f"Could not load existing analyses: {e}")

        documents = self.load_and_group_documents()

        if limit:
            documents = documents[:limit]
            print(f"Limited to {limit} documents for this run")

        analyses = []
        skipped = 0

        for doc in tqdm(documents, desc="Analyzing documents"):
            # Skip if already analyzed
            if doc['unique_id'] in existing_analyses:
                analyses.append(existing_analyses[doc['unique_id']])
                skipped += 1
                continue

            analysis = self.analyze_document(doc)
            if analysis:
                analyses.append(analysis)
                # Save incrementally
                self.save_analyses(analyses)

        print(f"\n✅ Analyzed {len(analyses) - skipped} new documents")
        print(f"   Skipped {skipped} already-analyzed documents")
        print(f"   Total analyses: {len(analyses)}")

        return analyses

    def save_analyses(self, analyses: List[Dict]):
        """Save analyses to JSON file"""
        output = {
            'total': len(analyses),
            'analyses': analyses
        }

        with open(self.analyses_file, 'w', encoding='utf-8') as f:
            json.dump(output, f, indent=2, ensure_ascii=False)


def main():
    load_dotenv()

    import argparse
    parser = argparse.ArgumentParser(description="Analyze documents using LLM")
    parser.add_argument("--api-url", help="OpenAI-compatible API base URL")
    parser.add_argument("--api-key", help="API key")
    parser.add_argument("--model", help="Model name")
    parser.add_argument("--limit", type=int, help="Limit number of documents to analyze")
    parser.add_argument("--force", action="store_true", help="Re-analyze all documents (ignore existing)")

    args = parser.parse_args()

    api_url = args.api_url or os.getenv("OPENAI_API_URL")
    api_key = args.api_key or os.getenv("OPENAI_API_KEY")
    model = args.model or os.getenv("OPENAI_MODEL", "gpt-4o")

    analyzer = DocumentAnalyzer(api_url, api_key, model)

    # Clear existing if force flag
    if args.force and analyzer.analyses_file.exists():
        analyzer.analyses_file.unlink()
        print("Removed existing analyses (--force mode)")

    analyses = analyzer.analyze_all(limit=args.limit)
    analyzer.save_analyses(analyses)

    print(f"\n✅ Saved analyses to {analyzer.analyses_file}")


if __name__ == "__main__":
    main()