Create app.py

app.py

@@ -0,0 +1,156 @@
+import gradio as gr
+import requests
+import json
+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"
+
+
+# ---------- 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}"
+
+
+# ---------- 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
+
+    code = ""
+    if code_input_type == "Paste Code":
+        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:
+            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()