File size: 5,611 Bytes
b950512 |
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 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 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() |