stratifyai 0.1.0__py3-none-any.whl

stratifyai/summarization.py

@@ -0,0 +1,179 @@
+"""Progressive summarization utilities for large files."""
+
+from typing import List, Optional
+from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
+
+from .client import LLMClient
+from .models import ChatRequest, Message
+from .chunking import chunk_content, get_chunk_metadata
+
+
+def summarize_chunk(
+    chunk: str,
+    client: LLMClient,
+    model: str = "gpt-4o-mini",
+    max_tokens: int = 1000,
+    context: Optional[str] = None
+) -> str:
+    """
+    Summarize a single chunk of content.
+    
+    Uses a cheaper model for cost efficiency.
+    
+    Args:
+        chunk: The content chunk to summarize
+        client: LLMClient instance (already configured with provider)
+        model: Model to use for summarization (default: gpt-4o-mini for cost)
+        max_tokens: Maximum tokens for summary
+        context: Optional context about the overall document
+        
+    Returns:
+        Summary of the chunk
+    """
+    # Build prompt
+    if context:
+        prompt = f"""Summarize the following section from a larger document.
+
+Context: {context}
+
+Section to summarize:
+{chunk}
+
+Provide a concise summary that preserves key information."""
+    else:
+        prompt = f"""Summarize the following text concisely, preserving key information:
+
+{chunk}"""
+    
+    # Create request
+    request = ChatRequest(
+        model=model,
+        messages=[Message(role="user", content=prompt)],
+        max_tokens=max_tokens
+    )
+    
+    # Get summary
+    response = client.chat_completion(request)
+    return response.content
+
+
+def summarize_chunks_progressive(
+    chunks: List[str],
+    client: LLMClient,
+    model: str = "gpt-4o-mini",
+    context: Optional[str] = None,
+    show_progress: bool = True
+) -> str:
+    """
+    Progressively summarize multiple chunks.
+    
+    Each chunk is summarized individually, then all summaries are combined.
+    
+    Args:
+        chunks: List of content chunks
+        client: LLMClient instance
+        model: Model to use for summarization
+        context: Optional context about the overall document
+        show_progress: Whether to show progress bar
+        
+    Returns:
+        Combined summary of all chunks
+    """
+    if not chunks:
+        return ""
+    
+    if len(chunks) == 1:
+        return summarize_chunk(chunks[0], client, model, context=context)
+    
+    summaries = []
+    
+    if show_progress:
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            TaskProgressColumn(),
+        ) as progress:
+            task = progress.add_task(
+                f"[cyan]Summarizing {len(chunks)} chunks...",
+                total=len(chunks)
+            )
+            
+            for i, chunk in enumerate(chunks, 1):
+                summary = summarize_chunk(
+                    chunk,
+                    client,
+                    model,
+                    context=f"{context} (Part {i}/{len(chunks)})" if context else f"Part {i}/{len(chunks)}"
+                )
+                summaries.append(f"**Part {i}/{len(chunks)}:**\n{summary}")
+                progress.update(task, advance=1)
+    else:
+        for i, chunk in enumerate(chunks, 1):
+            summary = summarize_chunk(
+                chunk,
+                client,
+                model,
+                context=f"{context} (Part {i}/{len(chunks)})" if context else f"Part {i}/{len(chunks)}"
+            )
+            summaries.append(f"**Part {i}/{len(chunks)}:**\n{summary}")
+    
+    # Combine summaries
+    combined = "\n\n".join(summaries)
+    
+    # If combined summaries are still very long, summarize the summaries
+    if len(combined) > 10000:  # Arbitrary threshold
+        final_summary = summarize_chunk(
+            combined,
+            client,
+            model,
+            context="Combined summaries of document sections"
+        )
+        return f"**Overall Summary:**\n{final_summary}\n\n**Detailed Summaries:**\n{combined}"
+    
+    return combined
+
+
+def summarize_file(
+    content: str,
+    client: LLMClient,
+    chunk_size: int = 50000,
+    model: str = "gpt-4o-mini",
+    context: Optional[str] = None,
+    show_progress: bool = True
+) -> dict:
+    """
+    Summarize a large file using progressive chunking.
+    
+    Args:
+        content: Full file content
+        client: LLMClient instance
+        chunk_size: Size of chunks in characters
+        model: Model to use for summarization
+        context: Optional context about the document
+        show_progress: Whether to show progress
+        
+    Returns:
+        Dictionary with summary and metadata
+    """
+    # Chunk the content
+    chunks = chunk_content(content, chunk_size=chunk_size)
+    metadata = get_chunk_metadata(chunks)
+    
+    # Summarize chunks
+    summary = summarize_chunks_progressive(
+        chunks,
+        client,
+        model=model,
+        context=context,
+        show_progress=show_progress
+    )
+    
+    return {
+        "summary": summary,
+        "original_length": len(content),
+        "summary_length": len(summary),
+        "reduction_percentage": round((1 - len(summary) / len(content)) * 100, 1),
+        "num_chunks": metadata["num_chunks"],
+        "chunk_metadata": metadata
+    }

stratifyai/utils/__init__.py

@@ -0,0 +1,11 @@
+"""Utility modules for StratifyAI."""
+
+from .token_counter import estimate_tokens, count_tokens_for_messages
+from .file_analyzer import analyze_file, FileAnalysis
+
+__all__ = [
+    "estimate_tokens",
+    "count_tokens_for_messages",
+    "analyze_file",
+    "FileAnalysis",
+]

stratifyai/utils/bedrock_validator.py

@@ -0,0 +1,136 @@
+"""Bedrock model validation utility.
+
+Validates AWS Bedrock model availability using boto3.
+"""
+
+import time
+from typing import Dict, List, Any, Optional
+
+try:
+    import boto3
+    from botocore.exceptions import ClientError, NoCredentialsError, BotoCoreError
+    BOTO3_AVAILABLE = True
+except ImportError:
+    BOTO3_AVAILABLE = False
+
+
+def validate_bedrock_models(
+    model_ids: List[str],
+    region_name: Optional[str] = None,
+) -> Dict[str, Any]:
+    """
+    Validate which Bedrock models are available in the user's AWS account/region.
+    
+    Args:
+        model_ids: List of model IDs to validate
+        region_name: AWS region (defaults to AWS_DEFAULT_REGION or us-east-1)
+        
+    Returns:
+        Dict containing:
+            - valid_models: List of model IDs that are available
+            - invalid_models: List of model IDs that are NOT available
+            - validation_time_ms: Time taken to validate in milliseconds
+            - error: Error message if validation failed (None if successful)
+    """
+    import os
+    
+    result = {
+        "valid_models": [],
+        "invalid_models": [],
+        "validation_time_ms": 0,
+        "error": None,
+    }
+    
+    if not BOTO3_AVAILABLE:
+        result["error"] = "boto3 not installed"
+        result["valid_models"] = model_ids  # Assume all valid if can't check
+        return result
+    
+    start_time = time.time()
+    
+    try:
+        # Get region from env or default
+        region = region_name or os.getenv("AWS_DEFAULT_REGION", "us-east-1")
+        
+        # Create bedrock client (not bedrock-runtime - we need list_foundation_models)
+        bedrock_client = boto3.client(
+            service_name="bedrock",
+            region_name=region
+        )
+        
+        # Get available foundation models
+        response = bedrock_client.list_foundation_models()
+        available_model_ids = {model["modelId"] for model in response.get("modelSummaries", [])}
+        
+        # Check each requested model
+        for model_id in model_ids:
+            if model_id in available_model_ids:
+                result["valid_models"].append(model_id)
+            else:
+                result["invalid_models"].append(model_id)
+        
+    except NoCredentialsError:
+        result["error"] = "AWS credentials not configured"
+        result["valid_models"] = model_ids  # Show all, let runtime handle auth
+        
+    except ClientError as e:
+        error_code = e.response.get("Error", {}).get("Code", "Unknown")
+        error_msg = e.response.get("Error", {}).get("Message", str(e))
+        result["error"] = f"AWS API error ({error_code}): {error_msg}"
+        result["valid_models"] = model_ids  # Show all on error
+        
+    except BotoCoreError as e:
+        result["error"] = f"AWS connection error: {str(e)}"
+        result["valid_models"] = model_ids  # Show all on error
+        
+    except Exception as e:
+        result["error"] = f"Validation failed: {str(e)}"
+        result["valid_models"] = model_ids  # Show all on error
+    
+    finally:
+        result["validation_time_ms"] = int((time.time() - start_time) * 1000)
+    
+    return result
+
+
+def get_validated_interactive_models(
+    region_name: Optional[str] = None,
+) -> Dict[str, Any]:
+    """
+    Get validated interactive Bedrock models with metadata.
+    
+    This is a convenience function that validates the curated interactive models
+    and returns them with their display metadata.
+    
+    Args:
+        region_name: AWS region (defaults to AWS_DEFAULT_REGION or us-east-1)
+        
+    Returns:
+        Dict containing:
+            - models: Dict mapping model_id to metadata (display_name, description, category)
+            - validation_result: Full validation result dict
+    """
+    from ..config import INTERACTIVE_BEDROCK_MODELS, BEDROCK_MODELS
+    
+    # Get list of interactive model IDs
+    model_ids = list(INTERACTIVE_BEDROCK_MODELS.keys())
+    
+    # Validate
+    validation_result = validate_bedrock_models(model_ids, region_name)
+    
+    # Build validated models dict with full metadata
+    models = {}
+    for model_id in validation_result["valid_models"]:
+        # Merge interactive metadata with full model config
+        interactive_meta = INTERACTIVE_BEDROCK_MODELS.get(model_id, {})
+        full_config = BEDROCK_MODELS.get(model_id, {})
+        
+        models[model_id] = {
+            **full_config,
+            **interactive_meta,
+        }
+    
+    return {
+        "models": models,
+        "validation_result": validation_result,
+    }

stratifyai/utils/code_extractor.py

@@ -0,0 +1,327 @@
+"""Code structure extraction for intelligent file analysis.
+
+This module extracts structural information from code files using AST to reduce
+token usage by 80%+ while preserving essential code structure.
+"""
+
+import ast
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List, Optional, Any
+
+
+@dataclass
+class FunctionInfo:
+    """Information about a function."""
+    name: str
+    line_number: int
+    params: List[str]
+    returns: Optional[str] = None
+    docstring: Optional[str] = None
+    decorators: List[str] = field(default_factory=list)
+    is_async: bool = False
+
+
+@dataclass
+class ClassInfo:
+    """Information about a class."""
+    name: str
+    line_number: int
+    bases: List[str]
+    methods: List[FunctionInfo]
+    docstring: Optional[str] = None
+    decorators: List[str] = field(default_factory=list)
+
+
+@dataclass
+class CodeStructure:
+    """Complete code structure information."""
+    file_path: str
+    language: str
+    imports: List[str]
+    functions: List[FunctionInfo]
+    classes: List[ClassInfo]
+    total_lines: int
+    docstring: Optional[str] = None
+    
+    def to_text(self) -> str:
+        """Convert structure to human-readable text.
+        
+        Returns:
+            Formatted code structure
+        """
+        lines = [
+            f"Code File: {self.file_path}",
+            f"Language: {self.language}",
+            f"Total Lines: {self.total_lines:,}",
+        ]
+        
+        if self.docstring:
+            lines.append(f"\nModule Docstring:\n  {self.docstring[:200]}")
+        
+        # Imports
+        if self.imports:
+            lines.append(f"\nImports ({len(self.imports)}):")
+            for imp in self.imports[:20]:  # Show first 20
+                lines.append(f"  - {imp}")
+            if len(self.imports) > 20:
+                lines.append(f"  ... and {len(self.imports) - 20} more")
+        
+        # Functions
+        if self.functions:
+            lines.append(f"\nFunctions ({len(self.functions)}):")
+            for func in self.functions:
+                decorators = f"@{', @'.join(func.decorators)} " if func.decorators else ""
+                async_prefix = "async " if func.is_async else ""
+                params = ", ".join(func.params)
+                returns = f" -> {func.returns}" if func.returns else ""
+                lines.append(f"  [Line {func.line_number}] {decorators}{async_prefix}def {func.name}({params}){returns}")
+                if func.docstring:
+                    lines.append(f"      \"{func.docstring[:100]}\"")
+        
+        # Classes
+        if self.classes:
+            lines.append(f"\nClasses ({len(self.classes)}):")
+            for cls in self.classes:
+                decorators = f"@{', @'.join(cls.decorators)} " if cls.decorators else ""
+                bases = f"({', '.join(cls.bases)})" if cls.bases else ""
+                lines.append(f"  [Line {cls.line_number}] {decorators}class {cls.name}{bases}:")
+                if cls.docstring:
+                    lines.append(f"      \"{cls.docstring[:100]}\"")
+                if cls.methods:
+                    lines.append(f"      Methods ({len(cls.methods)}):")
+                    for method in cls.methods[:10]:  # Show first 10 methods
+                        params = ", ".join(method.params)
+                        lines.append(f"        - {method.name}({params})")
+                    if len(cls.methods) > 10:
+                        lines.append(f"        ... and {len(cls.methods) - 10} more")
+        
+        return "\n".join(lines)
+
+
+class PythonASTVisitor(ast.NodeVisitor):
+    """AST visitor to extract code structure from Python files."""
+    
+    def __init__(self):
+        self.imports: List[str] = []
+        self.functions: List[FunctionInfo] = []
+        self.classes: List[ClassInfo] = []
+        self.current_class: Optional[str] = None
+    
+    def visit_Import(self, node: ast.Import):
+        """Visit import statement."""
+        for alias in node.names:
+            import_str = alias.name
+            if alias.asname:
+                import_str += f" as {alias.asname}"
+            self.imports.append(f"import {import_str}")
+        self.generic_visit(node)
+    
+    def visit_ImportFrom(self, node: ast.ImportFrom):
+        """Visit from...import statement."""
+        module = node.module or ""
+        for alias in node.names:
+            import_str = alias.name
+            if alias.asname:
+                import_str += f" as {alias.asname}"
+            self.imports.append(f"from {module} import {import_str}")
+        self.generic_visit(node)
+    
+    def visit_FunctionDef(self, node: ast.FunctionDef):
+        """Visit function definition."""
+        self._process_function(node, is_async=False)
+        self.generic_visit(node)
+    
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
+        """Visit async function definition."""
+        self._process_function(node, is_async=True)
+        self.generic_visit(node)
+    
+    def _process_function(self, node, is_async: bool):
+        """Process function/method node."""
+        # Extract parameters
+        params = []
+        for arg in node.args.args:
+            param_name = arg.arg
+            if arg.annotation:
+                try:
+                    param_name += f": {ast.unparse(arg.annotation)}"
+                except:
+                    pass
+            params.append(param_name)
+        
+        # Extract return type
+        returns = None
+        if node.returns:
+            try:
+                returns = ast.unparse(node.returns)
+            except:
+                pass
+        
+        # Extract docstring
+        docstring = ast.get_docstring(node)
+        if docstring:
+            docstring = docstring.split('\n')[0]  # First line only
+        
+        # Extract decorators
+        decorators = []
+        for decorator in node.decorator_list:
+            try:
+                decorators.append(ast.unparse(decorator))
+            except:
+                decorators.append("@decorator")
+        
+        func_info = FunctionInfo(
+            name=node.name,
+            line_number=node.lineno,
+            params=params,
+            returns=returns,
+            docstring=docstring,
+            decorators=decorators,
+            is_async=is_async
+        )
+        
+        # Add to appropriate list
+        if self.current_class:
+            # Find the class and add method
+            for cls in self.classes:
+                if cls.name == self.current_class:
+                    cls.methods.append(func_info)
+                    break
+        else:
+            self.functions.append(func_info)
+    
+    def visit_ClassDef(self, node: ast.ClassDef):
+        """Visit class definition."""
+        # Extract base classes
+        bases = []
+        for base in node.bases:
+            try:
+                bases.append(ast.unparse(base))
+            except:
+                bases.append("BaseClass")
+        
+        # Extract docstring
+        docstring = ast.get_docstring(node)
+        if docstring:
+            docstring = docstring.split('\n')[0]  # First line only
+        
+        # Extract decorators
+        decorators = []
+        for decorator in node.decorator_list:
+            try:
+                decorators.append(ast.unparse(decorator))
+            except:
+                decorators.append("@decorator")
+        
+        class_info = ClassInfo(
+            name=node.name,
+            line_number=node.lineno,
+            bases=bases,
+            methods=[],
+            docstring=docstring,
+            decorators=decorators
+        )
+        
+        self.classes.append(class_info)
+        
+        # Visit methods
+        old_class = self.current_class
+        self.current_class = node.name
+        self.generic_visit(node)
+        self.current_class = old_class
+
+
+def extract_python_structure(file_path: Path) -> CodeStructure:
+    """Extract structure from Python file using AST.
+    
+    Args:
+        file_path: Path to Python file
+        
+    Returns:
+        CodeStructure object
+        
+    Raises:
+        FileNotFoundError: If file doesn't exist
+        SyntaxError: If Python code is malformed
+    """
+    if not file_path.exists():
+        raise FileNotFoundError(f"Python file not found: {file_path}")
+    
+    # Read file
+    with open(file_path, 'r', encoding='utf-8') as f:
+        source_code = f.read()
+    
+    # Parse AST
+    try:
+        tree = ast.parse(source_code, filename=str(file_path))
+    except SyntaxError as e:
+        raise SyntaxError(f"Failed to parse {file_path}: {e}")
+    
+    # Extract module docstring
+    docstring = ast.get_docstring(tree)
+    if docstring:
+        docstring = docstring.split('\n')[0]  # First line only
+    
+    # Visit AST
+    visitor = PythonASTVisitor()
+    visitor.visit(tree)
+    
+    # Count lines
+    total_lines = source_code.count('\n') + 1
+    
+    return CodeStructure(
+        file_path=str(file_path),
+        language="Python",
+        imports=visitor.imports,
+        functions=visitor.functions,
+        classes=visitor.classes,
+        total_lines=total_lines,
+        docstring=docstring
+    )
+
+
+def analyze_code_file(file_path: Path) -> Dict[str, Any]:
+    """Analyze a code file and return structure information.
+    
+    Args:
+        file_path: Path to code file
+        
+    Returns:
+        Dictionary with structure and metadata
+    """
+    # Detect language from extension
+    extension = file_path.suffix.lower()
+    
+    if extension == '.py':
+        structure = extract_python_structure(file_path)
+        structure_text = structure.to_text()
+        
+        # Calculate sizes
+        original_size = file_path.stat().st_size
+        structure_size = len(structure_text)
+        
+        return {
+            'structure': structure,
+            'structure_text': structure_text,
+            'original_size_bytes': original_size,
+            'structure_size_bytes': structure_size,
+            'token_reduction_pct': ((original_size - structure_size) / original_size * 100) if original_size > 0 else 0.0,
+            'recommended_action': 'Use structure for LLM analysis instead of full code'
+        }
+    else:
+        # For non-Python files, return basic info
+        original_size = file_path.stat().st_size
+        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
+            lines = f.readlines()
+        
+        structure_text = f"Code File: {file_path}\nLanguage: {extension[1:] if extension else 'unknown'}\nTotal Lines: {len(lines)}\n\nNote: AST extraction only available for Python files."
+        
+        return {
+            'structure': None,
+            'structure_text': structure_text,
+            'original_size_bytes': original_size,
+            'structure_size_bytes': len(structure_text),
+            'token_reduction_pct': 0.0,
+            'recommended_action': 'Full file analysis required (non-Python)'
+        }