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