File size: 4,138 Bytes
bf25b85
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f0f5a18
bf25b85
f0f5a18
bf25b85
 
 
f0f5a18
bf25b85
 
 
 
f0f5a18
bf25b85
 
 
 
 
f0f5a18
bf25b85
 
 
 
 
 
 
 
 
f0f5a18
bf25b85
 
 
 
 
 
 
 
 
 
 
 
f0f5a18
 
 
 
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 Public 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()