synth.py

import os
import json
import argparse
from pathlib import Path
import litellm
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
import datetime
import platform

from utils.llm import StepIdentifier, TokenUsageTracker 
from utils.file import setup_run_directory, compile_latex_to_pdf
from utils.document import structure_to_markdown, structure_to_latex
from utils.generation import DocumentGenerator


# Create console for rich output
console = Console()

def check_model_cost_availability(model_name: str) -> str:
    """Checks if litellm can calculate costs for a given model."""
    try:
        prompt_cost, completion_cost = litellm.cost_per_token(
            model=model_name, prompt_tokens=100, completion_tokens=100
        )
        if prompt_cost > 0 or completion_cost > 0:
            return "[green]Available[/green]"
        else:
            # litellm knows the model, but costs are 0.
            return "[yellow]Known (Zero Cost/Potentially Free)[/yellow]"
    except Exception as e:
        # litellm no pricing data for this model
        return "[red]Unavailable[/red]"

# Parse command line arguments
def parse_arguments():
    parser = argparse.ArgumentParser(description="Generate synthetic documentation")
    
    # Required arguments
    parser.add_argument("document_type", choices=["tender", "offer"], help="Type of the document to be created, either offer or tender")
    parser.add_argument("document_title", help="Title of the document")
    parser.add_argument("document_info", help="Document information/guidance")
    
    # Optional arguments with defaults
    parser.add_argument("--guidance-doc", help="Path to guidance document. For the creation of tender documents this is supposed to be a description of the document structure. For the creation of the offer document it needs to be the final tender output json which this script produces.")
    parser.add_argument("--output-dir", default="generated", help="Directory to store output files")
    parser.add_argument("--model", default="openrouter/mistralai/mistral-small-3.1-24b-instruct", help="LLM to use for generation")
    parser.add_argument("--structure-model", default=None, help="LLM to use for structure generation")
    parser.add_argument("--recap-model", default=None, help="LLM to use for generating recaps between sections")
    parser.add_argument("--temperature", type=float, default=0.7, help="Temperature for LLM generation")
    parser.add_argument("--output-format", default="latex", choices=["markdown", "latex"], help="Output format (markdown or latex)")
    parser.add_argument("--no-track-costs", action="store_false", dest="track_costs", help="Disable token usage and cost tracking")
    parser.add_argument("--num-chapters", type=int, default=None, help="Number of chapters to generate (optional, mainly for tender)")
    parser.add_argument("--language", default="de", choices=["de", "en"], help="Language to use for generation (default: de - German)")
    parser.add_argument("--detail-level", choices=["low", "medium", "high"],
                        help="Detail level for generation only affects offer documents.",
                        default="high")
    parser.add_argument("--content-context-window", type=int, default=20,
                        help="Number of previous full chapter contents to consider for next section generation (default: 20)")

    args = parser.parse_args()

    # if structure_model is not provided, default it to the main model
    if not args.structure_model:
        args.structure_model = args.model

    # if recap_model is not provided, default it to the main model
    if not args.recap_model:
        args.recap_model = args.model

    # convert to config dictionary
    config = {
        "document_type": args.document_type,
        "document_title": args.document_title,
        "document_info": args.document_info,
        "detail_level": args.detail_level,
        "guidance_doc": args.guidance_doc,
        "output_dir": args.output_dir, 
        "model": args.model,
        "structure_model": args.structure_model,
        "recap_model": args.recap_model,
        "temperature": args.temperature,
        "output_format": args.output_format,
        "track_costs": args.track_costs,
        "num_chapters": args.num_chapters,
        "language": args.language,
        "content_context_window": args.content_context_window,
    }

    console.print(config["track_costs"])
    return config

def main():
    # Parse command line arguments
    config = parse_arguments()

    run_error = None

    # offer creation always needs a json tender structure
    if config["document_type"] == "offer":
        if not config["guidance_doc"] or not config["guidance_doc"].endswith(".json"):
            console.print("[bold red]Error:[/bold red] For document type 'offer', --guidance-doc must be provided and point to a tender JSON file.", style="red")
            return

        # check if the doc exists early
        if not os.path.exists(config["guidance_doc"]):
             console.print(f"[bold red]Error:[/bold red] Guidance document '{config['guidance_doc']}' not found.", style="red")
             return
        
    console.print(
        Panel.fit(
            "[bold cyan]Document Generation Process[/bold cyan]",
            border_style="blue",
            padding=(1, 10),
        )
    )

    # display configuration
    config_table = Table(show_header=True, header_style="bold magenta")
    config_table.add_column("Setting", style="dim")
    config_table.add_column("Value")

    for key, value in config.items():
        config_table.add_row(key, str(value) if value is not None else "[dim]Not Set[/dim]")

    console.print(
        Panel(config_table, title="[bold]Configuration", border_style="green")
    )
    console.print()

    if config.get("track_costs", True): # Only check if tracking is enabled
        console.print("[bold]Model cost calculation availability (via litellm):[/bold]")
        cost_check_table = Table(show_header=False, box=None, padding=(0, 1, 0, 0))
        cost_check_table.add_column("Model Type", style="dim")
        cost_check_table.add_column("Model Name")
        cost_check_table.add_column("Cost Tracking Status")

        structure_model_name = config["structure_model"]
        main_model_name = config["model"]

        structure_cost_status = check_model_cost_availability(structure_model_name)
        main_cost_status = check_model_cost_availability(main_model_name)

        cost_check_table.add_row("Structure Model", structure_model_name, structure_cost_status)
        cost_check_table.add_row("Main Model", main_model_name, main_cost_status)

        console.print(cost_check_table)
    else:
        console.print("[dim]Cost tracking is disabled via --track-costs=False.[/dim]\n")
    console.print()

    # initialize token tracker
    token_tracker = TokenUsageTracker()

    # create run directory at the start
    run_dir_path = setup_run_directory(config)
    
    # --- instantiate the DocumentGenerator ---
    try:
        generator = DocumentGenerator(config, token_tracker, run_dir_path)
    except Exception as e:
         console.print(f"[bold red]Error initializing DocumentGenerator:[/bold red] {e}")
         import traceback
         console.print(f"[red]{traceback.format_exc()}[/red]")
         return

    try:
        # Step 1: Generate structure
        generator.generate_structure()

        # confirmation step
        structure = generator.structure
        console.print("\n[bold yellow]Generated Document Structure:[/bold yellow]")
        for i, chapter in enumerate(structure["chapters"]):
            console.print(f"\n[bold cyan]Chapter {i+1}:[/bold cyan] {chapter.get('title', 'Untitled')}")
            console.print(f"[dim]Keywords:[/dim] {chapter.get('keywords', 'N/A')}")
            console.print(f"[dim]Brief Summary:[/dim] {chapter.get('brief_summary', 'N/A')}")
            console.print(f"[dim]Relevant Tender Sections:[/dim] {chapter.get('relevant_tender_chapters', 'N/A')}")


        # Ask for confirmation
        if (
            not console.input(
                "\n[bold yellow]Continue with this structure? (y/n):[/bold yellow] "
            )
            .lower()
            .startswith("y")
        ):
            console.print("[red]Process aborted by user[/red]")
            if config.get("track_costs", True) and token_tracker.total_cost > 0:
                 console.print("\n[bold blue]Token Usage and Cost Summary (Aborted Run):[/bold blue]")
                 console.print(token_tracker.get_summary_table())
            return

        console.print("[bold green]✓[/bold green] Basic structure confirmed.\n")

        # Step 2: Expand summaries
        generator.expand_summaries()
        console.print("[bold green]✓[/bold green] Chapter summaries expanded.\n")


        # Step 3: Generate full content
        generator.generate_sections() 
        console.print("[bold green]✓[/bold green] Full content generated.\n")

        # --- final Output Generation ---
        final_structure = generator.structure 
        output_format = config.get("output_format", "markdown")
        final_dir = run_dir_path / "final" 

        # Step 4 & 5: Generate final documents based on output format
        if output_format == "markdown":
            console.print(
                "[bold blue]Step 4/5:[/bold blue] [yellow]Creating final markdown document...[/yellow]"
            )
            with console.status("[bold green]Writing markdown file...", spinner="dots"):
                markdown_content = structure_to_markdown(final_structure)
                output_path = final_dir / "final_document.md"
                try:
                    with open(output_path, "w", encoding="utf-8") as f:
                        f.write(markdown_content)
                    console.print(
                        f"[bold green]✓[/bold green] Final markdown document saved to: [bold cyan]{output_path}[/bold cyan]\n"
                    )
                except IOError as e:
                     console.print(f"[bold red]Error writing markdown file {output_path}: {e}[/bold red]\n")


        elif output_format == "latex":
            console.print(
                "[bold blue]Step 4/5:[/bold blue] [yellow]Creating LaTeX document...[/yellow]"
            )
            with console.status("[bold green]Generating LaTeX document...", spinner="dots"):
                latex_content = structure_to_latex(config, final_structure, str(config["document_type"]))
                latex_path = final_dir / "final_document.tex"
                try:
                    with open(latex_path, "w", encoding="utf-8") as f:
                        f.write(latex_content)
                    console.print(
                        f"[bold green]✓[/bold green] LaTeX document saved to: [bold cyan]{latex_path}[/bold cyan]"
                    )
                except IOError as e:
                     console.print(f"[bold red]Error writing LaTeX file {latex_path}: {e}[/bold red]")


            console.print(
                "[bold blue]Step 5/5:[/bold blue] [yellow]Compiling PDF document...[/yellow]"
            )
            with console.status("[bold green]Compiling PDF...", spinner="dots"):
                compile_result = compile_latex_to_pdf(latex_path) 
                pdf_path = final_dir / "final_document.pdf"

                # Check if PDF exists after compilation attempt
                if pdf_path.exists():
                    console.print(
                        f"[bold green]✓[/bold green] PDF document potentially compiled: [bold cyan]{pdf_path}[/bold cyan]"
                    )
                    if compile_result:
                         console.print("[dim]Compilation reported success.[/dim]")
                    else:
                         console.print("[yellow]Warning:[/yellow] [dim]Compilation might have had issues. Please check the PDF and .log file.[/dim]")

                else:
                    console.print(
                        f"[bold red]✗[/bold red] Failed to compile PDF. Check LaTeX errors in '{latex_path.with_suffix('.log')}'."
                    )
                    run_error = {
                        "message": "PDF compilation failed"
                    }
            console.print() 


        # --- Final Report ---
        if config.get("track_costs", True):
            console.print("\n[bold blue]Token Usage and Cost Summary:[/bold blue]")

            if token_tracker.calls:
                 console.print(token_tracker.get_summary_table())

                 # also save to file
                 usage_summary = {
                     "total_prompt_tokens": token_tracker.total_prompt_tokens,
                     "total_completion_tokens": token_tracker.total_completion_tokens,
                     "total_cost_usd": token_tracker.total_cost, 
                     "calls": token_tracker.calls,
                 }
                 usage_path = final_dir / "token_usage_summary.json"
                 try:
                     with open(usage_path, "w", encoding="utf-8") as f:
                         json.dump(usage_summary, f, ensure_ascii=False, indent=2)
                     console.print(f"[dim]Token usage summary saved to: {usage_path}[/dim]")
                 except IOError as e:
                      console.print(f"[bold red]Error saving token usage summary: {e}[/bold red]")
            else:
                 console.print("[dim]No token usage tracked.[/dim]")

        console.print(
            Panel.fit(
                "[bold green]Document Generation Process Finished![/bold green]",
                border_style="green",
                padding=(1, 10),
            )
        )

    except FileNotFoundError as e_fnf:
         # Catch template loading errors during generation steps
         console.print("[bold red]\n--- Error: Required File Not Found ---[/bold red]")
         console.print(f"[red]{e}[/red]")
         console.print("[bold red]--- Generation aborted ---[/bold red]")
         run_error = {"message": str(e_fnf), "traceback": traceback.format_exc()}
    except ValueError as e_val:
         # Catch configuration or validation errors
         console.print("[bold red]\n--- Configuration Error ---[/bold red]")
         console.print(f"[red]{e}[/red]")
         console.print("[bold red]--- Generation aborted ---[/bold red]")
         run_error = {"message": str(e_val), "traceback": traceback.format_exc()}
    except Exception as e_gen:
        # catch any other unexpected errors
        console.print("[bold red]\n--- An Unexpected Error Occurred During Generation ---[/bold red]")
        import traceback
        console.print(f"[red]{traceback.format_exc()}[/red]")
        console.print("[bold red]--- Generation aborted ---[/bold red]")
        run_error = {"message": str(e_gen), "traceback": traceback.format_exc()}

        # attempt to print cost on failure
        if config.get("track_costs", True) and token_tracker.total_cost > 0:
            try:
                console.print("\n[bold blue]Token Usage and Cost Summary (Failed Run):[/bold blue]")
                console.print(token_tracker.get_summary_table())
            except Exception as cost_e:
                 console.print(f"[red]Could not display cost summary: {cost_e}[/red]")
    finally:
         # save run details 
         try:
            run_details_path = run_dir_path / "run_details.json"
            run_details = {
                "run_directory": str(run_dir_path),
                "config": config,
                "status": "completed" if pdf_path and pdf_path.exists() else f"failed: {run_error} " if run_error else "unknown",
                "output_files": {
                    "markdown": str(final_dir / "final_document.md") if config["output_format"] == "markdown" else None,
                    "latex": str(final_dir / "final_document.tex") if config["output_format"] == "latex" else None,
                    "pdf": str(final_dir / "final_document.pdf") if config["output_format"] == "latex" and pdf_path.exists() else None,
                    "token_usage": str(final_dir / "token_usage_summary.json") if config.get("track_costs", True) else None
                },
                "generation_stats": {
                    "total_chapters": len(generator.structure["chapters"]) if hasattr(generator, "structure") else 0,
                    "tokens": {
                        "total_tokens": token_tracker.get_total_tokens(),
                        "average_tokens_per_chapter": token_tracker.get_average_tokens_for_step(StepIdentifier.CONTENT_GENERATION),
                        "token_usage": {
                            "prompt_tokens": token_tracker.total_prompt_tokens,
                            "completion_tokens": token_tracker.total_completion_tokens,
                            "total_cost_usd": token_tracker.total_cost,
                        },
                    } if config.get("track_costs", True) else "Cost tracking disabled via --track-costs=False",
                    "times": {
                        "structure_generation": f"{generator.timing_stats.get('structure_generation', 0):.2f} seconds",
                        "summary_expansion": f"{generator.timing_stats.get('summary_expansion', 0):.2f} seconds",
                        "content_generation": f"{generator.timing_stats.get('content_generation', 0):.2f} seconds",
                        "total_generation_time": f"{sum(generator.timing_stats.values()):.2f} seconds" if hasattr(generator, "timing_stats") else None
                    }
                },
                "system_info": {
                    "python_version": platform.python_version(),
                    "os": platform.system(),
                    "machine": platform.machine()
                }
            }
            
            with open(run_details_path, "w", encoding="utf-8") as f:
                 json.dump(run_details, f, ensure_ascii=False, indent=2)
         except Exception as config_e:
             console.print(f"[red]Could not save configuration: {config_e}[/red]")


if __name__ == "__main__":
    main()