Update app.py

app.py

@@ -1,87 +1,18 @@
 import streamlit as st
-import requests
-from transformers import pipeline
+from code_reviewer import review_code
+from pr_fetcher import fetch_pr_diff
 
 # -----------------------------
-# Load Hugging Face Model
-# -----------------------------
-@st.cache_resource
-def load_model():
-    return pipeline("text2text-generation", model="google/flan-t5-small")  # faster
-
-reviewer = load_model()
-
-# -----------------------------
-# Fetch PR diff from GitHub
-# -----------------------------
-def fetch_pr_diff(owner, repo, pr_number, token=None):
-    url = f"https://api.github.com/repos/{owner}/{repo}/pulls/{pr_number}"
-    headers = {}
-    if token:
-        headers["Authorization"] = f"token {token}"
-    r = requests.get(url, headers=headers)
-    if r.status_code != 200:
-        return None, f"Error fetching PR: {r.text}"
-    
-    pr = r.json()
-    diff_url = pr.get("diff_url")
-    diff_r = requests.get(diff_url, headers=headers)
-    if diff_r.status_code != 200:
-        return None, f"Error fetching diff: {diff_r.text}"
-    return diff_r.text, None
-
-# -----------------------------
-# Review code with LLM
-# -----------------------------
-def review_code(code_text):
-    try:
-        if len(code_text) > 2000:
-            code_text = code_text[:2000] + "\n... [TRUNCATED] ..."
-
-        prompt = f"""
-        You are an AI code reviewer.
-        Review the following GitHub Pull Request diff:
-
-        {code_text}
-
-        Give two things:
-        1. FEEDBACK: clear review notes (issues + improvements)
-        2. CODE: improved/refactored version
-
-        Format:
-
-        FEEDBACK:
-        <feedback>
-
-        CODE:
-        <code>
-        """
-
-        output = reviewer(prompt, max_length=700, do_sample=False)
-        response = output[0]['generated_text'].strip()
-
-        # Split output safely
-        feedback, code = "", ""
-        if "CODE:" in response:
-            parts = response.split("CODE:", 1)
-            feedback = parts[0].replace("FEEDBACK:", "").strip()
-            code = parts[1].strip()
-        else:
-            feedback = response  # fallback
-
-        return feedback, code
-
-    except Exception as e:
-        return f"⚠️ Error during review: {str(e)}", ""
-
-# -----------------------------
-# Streamlit UI
+# Streamlit Page Config
 # -----------------------------
 st.set_page_config(page_title="Code Review Agent", page_icon="🤖", layout="wide")
 
 st.title("🤖 Code Review Agent (Hugging Face Edition)")
-st.write("Fetch a Pull Request from GitHub and get AI-powered review + improved code")
+st.write("Fetch a Pull Request from GitHub and get **AI-powered review + improved code**")
 
+# -----------------------------
+# Sidebar Inputs
+# -----------------------------
 with st.sidebar:
     st.header("📂 GitHub PR Details")
     owner = st.text_input("GitHub Repo Owner", "kumar1905")
@@ -90,11 +21,15 @@ with st.sidebar:
     token = st.text_input("GitHub Token (optional if public repo)", type="password")
     fetch_btn = st.button("🚀 Fetch & Review PR")
 
+# -----------------------------
+# Fetch + Review
+# -----------------------------
 if fetch_btn:
     with st.spinner("Fetching PR diff..."):
-        diff, err = fetch_pr_diff(owner, repo, pr_number, token)
-        if err:
-            st.error(err)
+        diff = fetch_pr_diff(owner, repo, pr_number, token)
+
+        if diff.startswith("Error"):
+            st.error(diff)
         else:
             st.success("✅ PR diff fetched successfully!")
             with st.expander("📄 View Raw PR Diff"):
@@ -103,14 +38,22 @@ if fetch_btn:
             with st.spinner("Analyzing code with Hugging Face model..."):
                 feedback, code = review_code(diff)
 
+            # Feedback Section
             st.subheader("💡 Review Feedback")
             if feedback:
                 st.write(feedback)
             else:
                 st.warning("No feedback generated.")
 
+            # Improved Code Section
             st.subheader("✨ Suggested Improved Code")
             if code:
                 st.code(code, language="python")
             else:
                 st.warning("No improved code generated.")
+
+# -----------------------------
+# Footer
+# -----------------------------
+st.markdown("---")
+st.markdown("Built with ❤️ using [Streamlit](https://streamlit.io/) & [Hugging Face Transformers](https://huggingface.co/)")