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"\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()