Update app.py

app.py

@@ -1,155 +1,156 @@
-import streamlit as st
+import gradio as gr
 import requests
 import json
-import re
 import os
 
 FIREWORKS_API_KEY = os.getenv("FIREWORKS_API_KEY")
 API_URL = "https://api.fireworks.ai/inference/v1/chat/completions"
 MODEL_NAME = "accounts/fireworks/models/glm-4p6"
 
-st.set_page_config(page_title="Intelligent Documentation Agent", layout="wide")
 
-# ---- Sidebar ----
-st.sidebar.title("⚙️ Configuration")
-api_key_input = st.sidebar.text_input("🔑 Fireworks API Key", value=FIREWORKS_API_KEY or "", type="password")
-if api_key_input:
-    FIREWORKS_API_KEY = api_key_input
-
-# ---- Tabs ----
-tab1, tab2 = st.tabs(["📄 Generate Documentation", "💬 Chat with Code"])
-
-# ---------------- TAB 1: DOCUMENTATION GENERATOR ----------------
-with tab1:
-    st.title("Intelligent Documentation Generator")
-
-    st.markdown("Upload a Python file, paste code, or provide a GitHub link to generate smart documentation.")
+# ---------- Fireworks GLM Request ----------
+def query_fireworks(messages, max_tokens=2048, temperature=0.6):
+    headers = {
+        "Accept": "application/json",
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {FIREWORKS_API_KEY}",
+    }
+    payload = {
+        "model": MODEL_NAME,
+        "max_tokens": max_tokens,
+        "temperature": temperature,
+        "messages": messages,
+    }
+    response = requests.post(API_URL, headers=headers, data=json.dumps(payload))
+    if response.status_code == 200:
+        return response.json()["choices"][0]["message"]["content"]
+    else:
+        return f"Error: {response.text}"
 
-    code_input_type = st.radio(
-        "Choose input type:",
-        ["Paste Code", "Upload File", "GitHub Link"],
-        horizontal=True
-    )
 
-    code_content = ""
+# ---------- Documentation Generator ----------
+def generate_docs(code_input_type, code_content, file_obj, github_url, api_key):
+    global FIREWORKS_API_KEY
+    if api_key:
+        FIREWORKS_API_KEY = api_key
 
-    # --- Option 1: Paste code ---
+    code = ""
     if code_input_type == "Paste Code":
-        code_content = st.text_area("Paste your Python code here", height=300)
-
-    # --- Option 2: Upload file ---
-    elif code_input_type == "Upload File":
-        uploaded_file = st.file_uploader("Upload a Python file", type=["py"])
-        if uploaded_file is not None:
-            code_content = uploaded_file.read().decode("utf-8")
-
-    # --- Option 3: GitHub link ---
-    elif code_input_type == "GitHub Link":
-        github_url = st.text_input("Enter a GitHub raw file link (e.g. https://raw.githubusercontent.com/...)")
-        if github_url:
-            if "raw.githubusercontent.com" not in github_url:
-                github_url = github_url.replace("github.com", "raw.githubusercontent.com").replace("/blob/", "/")
-            try:
-                r = requests.get(github_url)
-                if r.status_code == 200:
-                    code_content = r.text
-                else:
-                    st.error(f"Unable to fetch file (HTTP {r.status_code})")
-            except Exception as e:
-                st.error(f"Error fetching file: {e}")
-
-    if st.button("🚀 Generate Documentation"):
-        if not code_content.strip():
-            st.warning("Please provide code first.")
+        code = code_content
+    elif code_input_type == "Upload File" and file_obj is not None:
+        code = file_obj.read().decode("utf-8")
+    elif code_input_type == "GitHub Link" and github_url:
+        if "raw.githubusercontent.com" not in github_url:
+            github_url = github_url.replace("github.com", "raw.githubusercontent.com").replace("/blob/", "/")
+        r = requests.get(github_url)
+        if r.status_code == 200:
+            code = r.text
         else:
-            st.info("Generating intelligent documentation with GLM-4.6...")
-
-            messages = [
-                {
-                    "role": "system",
-                    "content": (
-                        "You are an AI documentation assistant that analyzes Python code and generates "
-                        "multi-layer documentation. Include: 1) Overview, 2) Key classes/functions, "
-                        "3) Usage examples, 4) Dependencies, 5) Suggested improvements."
-                    ),
-                },
-                {"role": "user", "content": code_content},
-            ]
-
-            payload = {
-                "model": MODEL_NAME,
-                "max_tokens": 4096,
-                "temperature": 0.6,
-                "messages": messages,
-            }
-
-            headers = {
-                "Accept": "application/json",
-                "Content-Type": "application/json",
-                "Authorization": f"Bearer {FIREWORKS_API_KEY}",
-            }
-
-            response = requests.post(API_URL, headers=headers, data=json.dumps(payload))
-            if response.status_code == 200:
-                output = response.json()["choices"][0]["message"]["content"]
-                st.subheader("📘 Generated Documentation")
-                st.markdown(output)
-                st.session_state["doc_output"] = output
-                st.session_state["code_context"] = code_content
-            else:
-                st.error(f"Error: {response.text}")
-
-# ---------------- TAB 2: CHATBOT ----------------
-with tab2:
-    st.title("💬 Chat with Your Code")
-
-    if "code_context" not in st.session_state:
-        st.info("First, generate or upload a code file in the previous tab to start chatting.")
-    else:
-        code_context = st.session_state["code_context"]
-        if "chat_history" not in st.session_state:
-            st.session_state["chat_history"] = []
-
-        user_query = st.text_input("Ask something about the code:")
-        if st.button("Ask"):
-            if not user_query.strip():
-                st.warning("Please enter a question.")
-            else:
-                messages = [
-                    {
-                        "role": "system",
-                        "content": (
-                            "You are a code analysis assistant. You answer questions about the following Python code. "
-                            "Explain clearly and refer to specific functions or logic when relevant."
-                        ),
-                    },
-                    {"role": "user", "content": f"Code:\n{code_context}\n\nQuestion:\n{user_query}"},
-                ]
-
-                payload = {
-                    "model": MODEL_NAME,
-                    "max_tokens": 2048,
-                    "temperature": 0.5,
-                    "messages": messages,
-                }
-
-                headers = {
-                    "Accept": "application/json",
-                    "Content-Type": "application/json",
-                    "Authorization": f"Bearer {FIREWORKS_API_KEY}",
-                }
-
-                response = requests.post(API_URL, headers=headers, data=json.dumps(payload))
-                if response.status_code == 200:
-                    reply = response.json()["choices"][0]["message"]["content"]
-                    st.session_state["chat_history"].append(("user", user_query))
-                    st.session_state["chat_history"].append(("assistant", reply))
-                else:
-                    st.error(f"Error: {response.text}")
-
-        # Display chat
-        for role, msg in st.session_state["chat_history"]:
-            if role == "user":
-                st.markdown(f"**You:** {msg}")
-            else:
-                st.markdown(f"**GLM-4.6:** {msg}")
+            return f"Unable to fetch file (HTTP {r.status_code})", ""
+
+    if not code.strip():
+        return "⚠️ Please provide valid Python code.", ""
+
+    messages = [
+        {
+            "role": "system",
+            "content": (
+                "You are an AI documentation assistant that analyzes Python code and generates "
+                "multi-layer documentation. Include: 1) Overview, 2) Key classes/functions, "
+                "3) Usage examples, 4) Dependencies, 5) Suggested improvements."
+            ),
+        },
+        {"role": "user", "content": code},
+    ]
+
+    output = query_fireworks(messages, max_tokens=4096)
+    return output, code
+
+
+# ---------- Code Chat ----------
+def chat_with_code(user_input, chat_history, code_context, api_key):
+    global FIREWORKS_API_KEY
+    if api_key:
+        FIREWORKS_API_KEY = api_key
+
+    if not code_context.strip():
+        return "⚠️ No code context found. Please generate documentation first.", chat_history
+
+    messages = [
+        {
+            "role": "system",
+            "content": (
+                "You are a code analysis assistant. You answer questions about the following Python code. "
+                "Explain clearly and refer to specific functions or logic when relevant."
+            ),
+        },
+        {"role": "user", "content": f"Code:\n{code_context}\n\nQuestion:\n{user_input}"},
+    ]
+
+    reply = query_fireworks(messages)
+    chat_history.append(("You", user_input))
+    chat_history.append(("GLM-4.6", reply))
+    return "", chat_history
+
+
+# ---------- Gradio UI ----------
+
+with gr.Blocks(title="Intelligent Documentation Generator Agent") as demo:
+    gr.Markdown("# 🧠 Intelligent Documentation Generator Agent (GLM-4.6 on Fireworks AI)")
+    gr.Markdown(
+        "Generate intelligent documentation or chat with your Python code. "
+        "Provide your Fireworks API key to get started."
+    )
+
+    with gr.Row():
+        api_key_box = gr.Textbox(
+            label="🔑 Fireworks API Key",
+            placeholder="Enter your Fireworks API Key",
+            type="password",
+        )
+
+    with gr.Tabs():
+        # ---- Documentation Tab ----
+        with gr.TabItem("📄 Generate Documentation"):
+            code_input_type = gr.Radio(
+                ["Paste Code", "Upload File", "GitHub Link"], label="Choose input type", value="Paste Code"
+            )
+
+            code_box = gr.Textbox(label="Paste your Python code", lines=15, visible=True)
+            file_upload = gr.File(label="Upload Python file", visible=False)
+            github_box = gr.Textbox(label="GitHub raw file link", visible=False)
+
+            def toggle_inputs(input_type):
+                return (
+                    gr.update(visible=input_type == "Paste Code"),
+                    gr.update(visible=input_type == "Upload File"),
+                    gr.update(visible=input_type == "GitHub Link"),
+                )
+
+            code_input_type.change(toggle_inputs, [code_input_type], [code_box, file_upload, github_box])
+
+            generate_btn = gr.Button("🚀 Generate Documentation")
+
+            output_box = gr.Markdown(label="📘 Generated Documentation")
+            hidden_code_context = gr.Textbox(visible=False)
+
+            generate_btn.click(
+                generate_docs,
+                [code_input_type, code_box, file_upload, github_box, api_key_box],
+                [output_box, hidden_code_context],
+            )
+
+        # ---- Chat Tab ----
+        with gr.TabItem("💬 Chat with Code"):
+            gr.Markdown("Chat about the Python file you uploaded or analyzed in the previous tab.")
+            chatbot = gr.Chatbot(label="GLM-4.6 Code Assistant")
+            msg = gr.Textbox(label="Ask a question about the code")
+            send_btn = gr.Button("Ask")
+
+            send_btn.click(
+                chat_with_code,
+                [msg, chatbot, hidden_code_context, api_key_box],
+                [msg, chatbot],
+            )
+
+demo.launch()