app.py

import os
import gradio as gr
import psycopg2
from psycopg2.extras import RealDictCursor
from dotenv import load_dotenv
from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span
import openai
from contextlib import contextmanager

from langtrace_python_sdk import langtrace

load_dotenv()  # Load environment variables
langtrace.init(api_key = os.getenv('LANGTRACE_API_KEY'))

GPT_4O_MINI = "gpt-4o-mini"
GPT_4 = "chatgpt-4o-latest" 
O1_MINI = "o1-mini"

# Global cache for prompts
cached_prompts = {}

@contextmanager
def openai_session():
    """Context manager to properly handle OpenAI API sessions"""
    try:
        # Initialize client
        client = openai.OpenAI()
        yield client
    finally:
        # Clean up client resources
        if hasattr(client, 'close'):
            client.close()

@with_langtrace_root_span()
def call_model(prompt, model):
    print(f"calling {model} with prompt: {prompt[:100]}")
    with openai_session() as client:
        try:
            # Call API
            response = client.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": prompt}]
            )
            # Extract response text
            result = response.choices[0].message.content
            return result
        
        except Exception as e:
            return f"Error generating output: {str(e)}"

def get_db_connection():
    """Establishes and returns a new database connection."""
    db_params = {
        'dbname': os.getenv('DB_NAME'),
        'user': os.getenv('DB_USER'),
        'password': os.getenv('DB_PASSWORD'),
        'host': os.getenv('DB_HOST'),
        'port': os.getenv('DB_PORT')
    }
    conn = psycopg2.connect(**db_params)
    return conn

def load_prompts():
    """Fetches prompts from the DB and caches them as a dict mapping prompt name to a dict with body and id."""
    global cached_prompts
    conn = get_db_connection()
    cursor = conn.cursor(cursor_factory=RealDictCursor)
    cursor.execute("""
        SELECT code_review_prompt_name, code_review_prompt_body, code_review_prompt_id 
        FROM code_review_prompt
        ORDER BY code_review_prompt_name;
    """)
    rows = cursor.fetchall()
    # Cache format: { prompt_name: {"body": prompt_body, "id": prompt_id} }
    cached_prompts = { row["code_review_prompt_name"]: {"body": row["code_review_prompt_body"], "id": row["code_review_prompt_id"]} for row in rows }
    cursor.close()
    conn.close()
    return cached_prompts

# Initially load prompts
load_prompts()

def update_prompt_text(selected_prompt):
    """Uses the global cache to return the prompt body for the selected prompt."""
    return cached_prompts.get(selected_prompt, {}).get("body", "")

def build_final_prompt(custom_prompt, input_code):
    """
    Builds the final prompt for code review.
    """
    return (
        f"Please review the following code:\n\n"
        f"{custom_prompt}\n\n"
        f"<CODE>\n"
        f"{input_code}\n"
    )

def store_submission(prompt_name, custom_prompt, input_code, review_output):
    """
    Stores the submission in the database.
    """
    prompt_data = cached_prompts.get(prompt_name)
    prompt_id = prompt_data["id"] if prompt_data else None

    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute("""
        INSERT INTO code_review_submission (code_review_prompt_id, custom_prompt, input_code, ai_review)
        VALUES (%s, %s, %s, %s);
    """, (prompt_id, custom_prompt, input_code, review_output))
    conn.commit()
    cursor.close()
    conn.close()

def generate_ai_review(prompt_name, custom_prompt, input_code, model):
    """
    Generates an AI review by calling the API, then stores the submission.
    Uses prompt_name to look up the prompt id.
    """
    final_prompt = build_final_prompt(custom_prompt, input_code)
    review_output = call_model(final_prompt, model)
    store_submission(prompt_name, custom_prompt, input_code, review_output)
    return review_output

def call_four(prompt_name, custom_prompt, input_code):
    return generate_ai_review(prompt_name, custom_prompt, input_code, GPT_4)

def call_o1(prompt_name, custom_prompt, input_code):
    return generate_ai_review(prompt_name, custom_prompt, input_code, O1_MINI)

def call_four_mini(prompt_name, custom_prompt, input_code):
    return generate_ai_review(prompt_name, custom_prompt, input_code, GPT_4O_MINI)

def add_new_prompt(prompt_name, prompt_body):
    """Inserts a new prompt into the database, reloads the cache, and returns updated dropdown choices."""
    if not prompt_name.strip() or not prompt_body.strip():
        return "Error: Prompt Name and Body cannot be empty!", None, None

    conn = None
    try:
        conn = get_db_connection()
        cursor = conn.cursor(cursor_factory=RealDictCursor)
        cursor.execute("""
            INSERT INTO code_review_prompt (code_review_prompt_name, code_review_prompt_body)
            VALUES (%s, %s)
            ON CONFLICT (code_review_prompt_name) DO NOTHING
            RETURNING code_review_prompt_id;
        """, (prompt_name.strip(), prompt_body.strip()))
        inserted = cursor.fetchone()
        if inserted:
            conn.commit()
            message = f"Prompt '{prompt_name}' added successfully!"
        else:
            message = f"Prompt '{prompt_name}' already exists!"
        cursor.close()
        conn.close()
        # Reload the prompt cache after insertion
        load_prompts()
        new_choices = list(cached_prompts.keys())
        new_dropdown = gr.Dropdown(choices=new_choices, label="Select Review Type", value=new_choices[0] if new_choices else None)
        return message, new_dropdown, ""
    except psycopg2.IntegrityError:
        conn.rollback()
        return "Prompt already exists!", None, ""
    except Exception as e:
        if conn:
            conn.rollback()
        return f"Error: {str(e)}", None, ""
    finally:
        if conn:
            conn.close()

def reload_prompts():
    """Reloads the prompt cache and returns a new Dropdown with updated choices."""
    load_prompts()
    new_choices = list(cached_prompts.keys())
    new_dropdown = gr.Dropdown(choices=new_choices, label="Select Review Type", value=new_choices[0] if new_choices else None)
    return new_dropdown

# ---------------------- New Code for Loading Previous Submissions ---------------------- #

def select_submission_by_id(selected_id, submissions):
    """
    Given a selected submission_id (as string) and the list of submissions,
    finds the matching submission and returns its details.
    """
    index = next(
        (i for i, s in enumerate(submissions) if str(s["code_review_submission_id"]) == selected_id),
        0
    )
    details, code, review, index = show_submission(index, submissions)
    return index, details, code, review

def load_submissions_from_db():
    """Fetch all submissions ordered by their id in ascending order."""
    conn = get_db_connection()
    cursor = conn.cursor(cursor_factory=RealDictCursor)
    cursor.execute("""
        SELECT code_review_submission_id, code_review_prompt_id, custom_prompt, input_code, ai_review
        FROM code_review_submission
        ORDER BY created_at DESC; """)
    submissions = cursor.fetchall()
    cursor.close()
    conn.close()
    return submissions

def show_submission(index, submissions):
    """Returns formatted outputs for the submission at the given index:
    - Submission details (Markdown)
    - Input code (for gr.Code)
    - AI review output (Markdown)
    """
    if not submissions:
        return "No submissions available", "", "No submissions available", index
    # Clamp index within bounds
    index = max(0, min(index, len(submissions) - 1))
    submission = submissions[index]
    # Left panel Markdown for submission details
    submission_details = (
        f"**Submission ID:** {submission['code_review_submission_id']}\n\n"
        f"**Prompt ID:** {submission['code_review_prompt_id']}\n\n"
        f"**Custom Prompt:**\n{submission['custom_prompt']}\n\n"
    )
    # Input code for gr.Code
    input_code = submission['input_code']
    # Right panel Markdown for AI review output
    ai_review_output = f"**AI Review Output:**\n{submission['ai_review']}\n"
    return submission_details, input_code, ai_review_output, index


# --- Updated load_submissions function ---
def load_submissions():
    """Loads submissions from the database, resets the index,
    and returns a Dropdown update for submission choices."""
    submissions = load_submissions_from_db()
    index = 0
    details, code, review, index = show_submission(index, submissions)
    submission_choices = [str(s["code_review_submission_id"]) for s in submissions]
    default_choice = submission_choices[0] if submission_choices else None
    dropdown_update = gr.Dropdown(choices=submission_choices, value=default_choice)
    return submissions, index, details, code, review, dropdown_update

def next_submission(current_index, submissions):
    """Increment the index and display the next submission."""
    if submissions and current_index < len(submissions) - 1:
        current_index += 1
    details, code, review, current_index = show_submission(current_index, submissions)
    return current_index, details, code, review

def previous_submission(current_index, submissions):
    """Decrement the index and display the previous submission."""
    if submissions and current_index > 0:
        current_index -= 1
    details, code, review, current_index = show_submission(current_index, submissions)
    return current_index, details, code, review

# ----------------------------------------------------------------------------------------- #

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# Code Review Assistant")

    with gr.Tab("Code Review"):
        with gr.Row():
            with gr.Column(scale=1):
                prompt_dropdown = gr.Dropdown(
                    choices=list(cached_prompts.keys()),
                    label="Select Review Type",
                    value=list(cached_prompts.keys())[0] if cached_prompts else None
                )
                with gr.Accordion("Selected Prompt", open=False):
                    prompt_textbox = gr.Textbox(label="", lines=15, interactive=True)
        with gr.Row():
            generate_4o_btn = gr.Button("Generate Code Review 4o")
            generate_4omini_btn = gr.Button("Generate Code Review 4o-mini")
            generate_o1_btn = gr.Button("Generate Code Review o1-mini")
        with gr.Row():
            input_code = gr.Code(language="python", label="Input Code", lines=15, interactive=True)
            output_review = gr.Markdown(label="AI Review Output", container=True, show_copy_button=True)

    with gr.Tab("Manage Prompts"):
        with gr.Row():
            new_prompt_name = gr.Textbox(label="New Prompt Name", placeholder="Enter a prompt name")
            new_prompt_body = gr.Textbox(label="New Prompt Body", lines=3, placeholder="Enter the prompt details")
        add_prompt_btn = gr.Button("Add Prompt")
        reload_prompt_btn = gr.Button("Reload Prompts")
        prompt_status = gr.Textbox(label="Status", interactive=False)

    with gr.Tab("Submissions"):
        with gr.Row():
            load_submissions_btn = gr.Button("Load Submissions")
        # Hidden states to store submissions and the current index
        submissions_state = gr.State([])
        submission_index_state = gr.State(0)
        # Add a new dropdown to select submission by ID
        submission_dropdown = gr.Dropdown(label="Select Submission by ID", choices=[], value=None)
        with gr.Row():
            back_btn = gr.Button("Back")
            next_btn = gr.Button("Next")
        with gr.Row():
            with gr.Column():
                with gr.Accordion("Prompt Details", open=False):
                    submission_details_md = gr.Markdown(label="Submission Details")
                input_code_component = gr.Code(label="Input Code", language="python")
            with gr.Column():
                ai_review_md = gr.Markdown(label="AI Review Output", container=True, show_copy_button=True)
        
        # Load submissions from DB, initialize state, and update dropdown choices
        load_submissions_btn.click(
            fn=load_submissions,
            inputs=[],
            outputs=[submissions_state, submission_index_state, submission_details_md, input_code_component, ai_review_md, submission_dropdown]
        )
        
        # When a submission is selected from the dropdown, update the submission display
        submission_dropdown.change(
            fn=select_submission_by_id,
            inputs=[submission_dropdown, submissions_state],
            outputs=[submission_index_state, submission_details_md, input_code_component, ai_review_md]
        )
        
        # Go to previous submission
        back_btn.click(
            fn=previous_submission,
            inputs=[submission_index_state, submissions_state],
            outputs=[submission_index_state, submission_details_md, input_code_component, ai_review_md]
        )
        
        # Go to next submission
        next_btn.click(
            fn=next_submission,
            inputs=[submission_index_state, submissions_state],
            outputs=[submission_index_state, submission_details_md, input_code_component, ai_review_md]
        )



    # When dropdown changes, update the prompt text
    prompt_dropdown.change(
        fn=update_prompt_text,
        inputs=[prompt_dropdown],
        outputs=[prompt_textbox]
    )

    # On code review generation, pass both the dropdown (prompt name) and the editable prompt text
    generate_4o_btn.click(
        fn=call_four,
        inputs=[prompt_dropdown, prompt_textbox, input_code],
        outputs=[output_review]
    )
    generate_4omini_btn.click(
        fn=call_four_mini,
        inputs=[prompt_dropdown, prompt_textbox, input_code],
        outputs=[output_review]
    )
    generate_o1_btn.click(
        fn=call_o1,
        inputs=[prompt_dropdown, prompt_textbox, input_code],
        outputs=[output_review]
    )

    # On adding a new prompt, update status, replace dropdown with a new one, and clear new prompt body textbox.
    add_prompt_btn.click(
        fn=add_new_prompt,
        inputs=[new_prompt_name, new_prompt_body],
        outputs=[prompt_status, prompt_dropdown, new_prompt_body]
    )

    # Reload prompts when the reload button is clicked.
    reload_prompt_btn.click(
        fn=reload_prompts,
        inputs=[],
        outputs=[prompt_dropdown]
    )


demo.launch()