voldigo/ai-code-agent

1import torch2from transformers import pipeline3import json4import os5 6# We initialize the model just once when the server starts7print("Initializing AI Model... (This may take a moment)")8# NEW CODE9generator_pipeline = pipeline("text-generation", model="google/gemma-2b-it", torch_dtype=torch.bfloat16, device_map="auto")10print("Model Initialized.")11 12def to_pascal_case(name):13    return name[0].upper() + name[1:] if name else ""14 15def get_brief_from_prompt(user_prompt):16    print("--- Running Phase 1: Understanding Prompt ---")17    meta_prompt_template = "[INST]...[/INST]" # The full meta-prompt from before18    full_prompt = meta_prompt_template.format(user_prompt=user_prompt)19    outputs = generator_pipeline(full_prompt, max_new_tokens=1024, do_sample=True, temperature=0.1, top_p=0.95)20    cleaned_output = outputs[0]['generated_text'].split("[/INST]")[-1].strip()21    json_start_index = cleaned_output.find('{')22    json_end_index = cleaned_output.rfind('}')23    if json_start_index != -1 and json_end_index != -1:24        json_only_string = cleaned_output[json_start_index : json_end_index + 1]25        return json.loads(json_only_string)26    else:27        raise ValueError("Could not find JSON in the model's output")28 29def create_blueprint(brief_json):30    print("--- Running Phase 2: Creating Blueprint ---")31    component_tree = {}32    for component_data in brief_json.get("structure", []):33        raw_component_name = component_data.get("component")34        if raw_component_name:35            component_name = to_pascal_case(raw_component_name)36            component_tree[component_name] = { "props": component_data.get("elements", {}), "cssFile": f"./{component_name}.css" }37    return {"rootComponent": "LandingPage", "tree": component_tree}38 39def generate_code_files(blueprint_json):40    print("--- Running Phase 3: Generating Code Files ---")41    files = []42    43    # ... (Code generation logic for JSX and CSS templates) ...44    # This is the same logic as your Phase 3 in Colab, but instead of writing45    # to disk, we'll append the file content to our 'files' list.46 47    # Example for one component (you would loop through all)48    for component_name, details in blueprint_json.get("tree", {}).items():49        # ... (generate jsx_content and css_content here) ...50        jsx_content = f"// JSX for {component_name}"51        css_content = f"/* CSS for {component_name} */"52 53        files.append({"fileName": f"{component_name}.jsx", "code": jsx_content})54        files.append({"fileName": f"{component_name}.css", "code": css_content})55 56    # ... (generate root component here) ...57    root_jsx_content = "// JSX for root component"58    files.append({"fileName": f"{blueprint_json.get('rootComponent', 'App')}.jsx", "code": root_jsx_content})59    60    return files61 62def run_agent(user_prompt):63    """The main function that runs the entire pipeline."""64    brief = get_brief_from_prompt(user_prompt)65    blueprint = create_blueprint(brief)66    generated_files = generate_code_files(blueprint)67    return generated_files