cli.py

import argparse
from tomlToChain.Builder import Builder
from tomlToChain.flowgrapher import generate_dot_file
from tomlToChain.linter import lint_toml
from dotenv import load_dotenv

load_dotenv()

def run_chain(toml_file):
    """Runs the LLM chain based on the provided TOML file."""
    print(f"🚀 Running LLM Chain from {toml_file}...")
    builder = Builder(toml_file=toml_file)
    chain = builder.get_chain()
    chain.run_chain()


def generate_flow(toml_file, output_file="flow_graph.dot"):
    """Generates a DOT flow graph for the TOML chain."""
    print(f"📊 Generating DOT flow graph for {toml_file}...")
    builder = Builder(toml_file=toml_file)
    chain = builder.get_chain()
    generate_dot_file(builder, chain, filename=output_file)

def lint(toml_file):
    """Lint the TOML file for validation."""
    print(f"🔍 Linting TOML file: {toml_file}")
    lint_toml(toml_file)

def main():
    parser = argparse.ArgumentParser(
        description="LLM Chain Execution Framework - Run, Lint, or Generate Flow Graph from TOML files."
    )

    parser.add_argument(
        "-t", "--toml", type=str, required=True, help="Path to the TOML file."
    )
    parser.add_argument(
        "-f", "--flow", action="store_true", help="Generate a DOT flow graph from the TOML file."
    )
    parser.add_argument(
        "-o", "--output", type=str, default="flow_graph.dot", help="Output file for the DOT flow graph (default: flow_graph.dot)."
    )
    parser.add_argument(
        "-l", "--lint", action="store_true", help="Lint the TOML file for validation."
    )

    args = parser.parse_args()

    if args.lint:
        lint(args.toml)
    elif args.flow:
        generate_flow(args.toml, args.output)
    else:
        run_chain(args.toml)


if __name__ == "__main__":
    main()