Karthix1/ai-code-explainer

1import streamlit as st
2import os
3import json
4from openai import OpenAI
5from dotenv import load_dotenv
6import re
7
8load_dotenv()
9
10client = OpenAI(
11  base_url="https://openrouter.ai/api/v1",
12  api_key=os.getenv("OPENROUTER_API_KEY"),
13)
14
15def generate_code_analysis_with_retry(code_snippet: str, model_name: str,language: str, max_retries: int = 3):
16    """
17    Generates code analysis with a self-correction loop.
18    It tries to get valid JSON, and if it fails, it tells the AI its mistake and retries.
19    Returns a parsed dictionary on success, or None on failure.
20    """
21    
22    # Define the initial user request
23    initial_prompt = f"""
24    You are an expert {language} programmer. Analyze the following code snippet and provide a plain-English explanation and a Google-style docstring.
25    
26    Code:
27    ```
28    {code_snippet}
29    ```
30    
31    Respond with ONLY a single, valid JSON object with two keys: "explanation" and "docstring". Do not include any markdown formatting, comments, or other text outside of the JSON.
32    """
33    
34    # Initialize the conversation history for the AI
35    messages = [{"role": "user", "content": initial_prompt}]
36    
37    # Start the self-correction loop
38    for attempt in range(max_retries):
39        st.write(f" Attempt {attempt + 1} of {max_retries}...") 
40        
41        try:
42            # === ACT: Call the AI ===
43            response = client.chat.completions.create(
44                model=model_name,
45                messages=messages
46            )
47            raw_output = response.choices[0].message.content
48            
49            
50            # Use regex to find the JSON object within the potentially messy string
51            match = re.search(r"\{.*\}", raw_output, re.DOTALL)
52            if not match:
53                raise ValueError("No JSON object found in the response.")
54            
55            cleaned_json_str = match.group(0)
56            
57            parsed_json = json.loads(cleaned_json_str)
58            
59            if "explanation" not in parsed_json or "docstring" not in parsed_json:
60                raise ValueError("JSON is missing required keys ('explanation', 'docstring').")
61
62            st.success(f"Analysis successful on attempt {attempt + 1}!")
63            return parsed_json 
64
65        except (json.JSONDecodeError, ValueError, IndexError) as e:
66            # === REASON & REACT: If an error occurred, start the correction process ===
67            st.warning(f"Attempt {attempt + 1} failed: {e}. Trying to self-correct...")
68            
69            # Add the AI's failed response to the conversation history
70            messages.append({"role": "assistant", "content": raw_output})
71            
72            # Create the corrective prompt, showing the AI its own mistake
73            corrective_prompt = f"""
74            Your previous response could not be parsed.
75            Error: "{e}"
76            Your full response was:
77            ---
78            {raw_output}
79            ---
80            Please correct your mistake. Look at the error and your previous response. 
81            Provide the response again as a single, valid JSON object with the keys "explanation" and "docstring". 
82            DO NOT wrap it in markdown or add any other text.
83            """
84            
85            # Add  corrective instruction to the conversation
86            messages.append({"role": "user", "content": corrective_prompt})
87            
88    st.error(f"Failed to get a valid response after {max_retries} attempts.")
89    return None
90
91
92
93st.set_page_config(layout="wide")
94st.title("AI Code Explainer & Docstring Generator")
95st.write("Powered by OpenRouter.ai with a Self-Correction Loop")
96
97code_input = st.text_area(
98    "Paste your Python function or code block here:", 
99    height=250, 
100    placeholder="def my_function(arg1, arg2):\n    # Your code here\n    return result"
101)
102
103model_choice = st.selectbox(
104    "Choose your AI model:",
105    (
106        "Google: Gemma 3n",
107        "MoonshotAI: Kimi Dev ",
108        "NVIDIA: Nemotron Nano 9B",
109        "Mistral: Mistral 7B Instruct",
110    ),
111    help="Free models from OpenRouter. Different models have different strengths."
112)
113
114MODEL_MAPPING = {
115    "Google: Gemma 3n": "google/gemma-3n-e2b-it:free",
116    "MoonshotAI: Kimi Dev ": "moonshotai/kimi-dev-72b:free",
117    "NVIDIA: Nemotron Nano 9B": "nvidia/nemotron-nano-9b-v2:free",
118    "Mistral: Mistral 7B Instruct": "mistralai/mistral-7b-instruct:free",
119}
120selected_model_id = MODEL_MAPPING[model_choice]
121
122language = st.selectbox("Select Language", ["Python", "JavaScript", "Java", "Go"])
123if st.button("Analyze Code", type="primary"):
124    if code_input:
125        analysis_dict = generate_code_analysis_with_retry(code_input, selected_model_id,language)
126        
127        if analysis_dict:
128            st.subheader("Final Analysis Results")
129            col1, col2 = st.columns(2)
130            
131            with col1:
132                st.info("💬 Plain English Explanation")
133                st.write(analysis_dict.get("explanation", "No explanation was generated."))
134            
135            with col2:
136                st.success("📝 Generated Docstring")
137                st.code(analysis_dict.get("docstring", "No docstring was generated."), language="python")
138    else:
139        st.warning("Please paste some code into the text area above.")