|
from langgraph.graph import END, START, StateGraph |
|
|
|
from langchain_core.tracers.context import tracing_v2_enabled |
|
|
|
from prompts import * |
|
import gradio as gr |
|
|
|
|
|
from utils import ( |
|
load_github_codebase, |
|
router, |
|
get_plan_for_codebase, |
|
parse_plan, |
|
explore_file, |
|
final_mermaid_code_generation, |
|
extract_mermaid_and_generate_graph, |
|
GraphState, |
|
) |
|
|
|
|
|
def visualize_github_repo(repo_name, repo_branch): |
|
yield ("Looking at the Repo!", None) |
|
documents = load_github_codebase(repo_name, repo_branch) |
|
yield ("Repo loaded!", None) |
|
|
|
|
|
workflow = StateGraph(GraphState) |
|
|
|
|
|
workflow.add_node("planner", get_plan_for_codebase) |
|
workflow.add_node("parse_plan", parse_plan) |
|
|
|
workflow.add_node("explore_file", explore_file) |
|
workflow.add_node("router", router) |
|
|
|
workflow.add_node("generate_mermaid_code", final_mermaid_code_generation) |
|
workflow.add_node("render_mermaid", extract_mermaid_and_generate_graph) |
|
|
|
|
|
|
|
workflow.add_edge(START, "planner") |
|
workflow.add_edge("planner", "parse_plan") |
|
|
|
|
|
workflow.add_conditional_edges( |
|
|
|
|
|
"parse_plan", |
|
|
|
router, |
|
) |
|
workflow.add_conditional_edges( |
|
"explore_file", |
|
router, |
|
) |
|
workflow.add_edge("generate_mermaid_code", "render_mermaid") |
|
workflow.add_edge("render_mermaid", END) |
|
|
|
|
|
|
|
|
|
|
|
app = workflow.compile() |
|
|
|
with tracing_v2_enabled(): |
|
for s in app.stream( |
|
{"messages": [], "documents": documents}, |
|
{"recursion_limit": 100}, |
|
): |
|
if "__end__" not in s: |
|
print(s) |
|
print("----") |
|
if "planner" in s: |
|
yield ("Planning the Exploration !", None) |
|
if "parse_plan" in s: |
|
yield ("Planning done! Parsing the plan!", None) |
|
if "explore_file" in s: |
|
yield ( |
|
f"Exploration started! Exploring file: {s['explore_file']['explored_files'][-1]} !", |
|
None, |
|
) |
|
if "generate_mermaid_code" in s: |
|
yield ( |
|
"Exploration done! Gathering thoughts and generating a graph!", |
|
None, |
|
) |
|
if "render_mermaid" in s: |
|
|
|
yield ( |
|
s["render_mermaid"]["messages"][-1].content, |
|
s["render_mermaid"]["final_graph"], |
|
) |
|
|
|
|
|
demo = gr.Interface( |
|
fn=visualize_github_repo, |
|
inputs=[ |
|
gr.Textbox( |
|
label="Repo Name", |
|
value="abhishekkrthakur/autoxgb", |
|
placeholder="Name of the Repo in format author/repo", |
|
), |
|
gr.Textbox( |
|
label="Repo Branch", |
|
value="main", |
|
placeholder="Branch to explore", |
|
), |
|
], |
|
outputs=[ |
|
gr.Textbox( |
|
label="Mermaid Graph", |
|
placeholder="Visualization of the functionalities of the Repo", |
|
), |
|
gr.Image( |
|
label="Rendered Graph", |
|
placeholder="Rendered Graph", |
|
), |
|
], |
|
title="Repo Functionality Visualizer", |
|
) |
|
|
|
demo.launch() |
|
|