File size: 4,113 Bytes
bf25b85 |
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 |
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)
# Define a new graph
workflow = StateGraph(GraphState)
# Define the two nodes we will cycle
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)
# Set the entrypoint as `agent`
# This means that this node is the first one called
workflow.add_edge(START, "planner")
workflow.add_edge("planner", "parse_plan")
# We now add a conditional edge
workflow.add_conditional_edges(
# First, we define the start node. We use `agent`.
# This means these are the edges taken after the `agent` node is called.
"parse_plan",
# Next, we pass in the function that will determine which node is called next.
router,
)
workflow.add_conditional_edges(
"explore_file",
router,
)
workflow.add_edge("generate_mermaid_code", "render_mermaid")
workflow.add_edge("render_mermaid", END)
# Finally, we compile it!
# This compiles it into a LangChain Runnable,
# meaning you can use it as you would any other runnable.
# Note that we're (optionally) passing the memory when compiling the graph
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", # Input component label
value="abhishekkrthakur/autoxgb", # Input component placeholder
placeholder="Name of the Repo in format author/repo", # Input component description
),
gr.Textbox(
label="Repo Branch", # Input component label
value="main", # Input component placeholder
placeholder="Branch to explore", # Input component description
),
],
outputs=[
gr.Textbox(
label="Mermaid Graph", # Output component label
placeholder="Visualization of the functionalities of the Repo", # Output component description
),
gr.Image(
label="Rendered Graph", # Output component label
placeholder="Rendered Graph", # Output component description
),
],
title="Repo Functionality Visualizer",
)
demo.launch()
|