Spaces:
Sleeping
Sleeping
| import os | |
| from prompts import * | |
| import gradio as gr | |
| def visualize_github_repo(repo_name, repo_branch, mistral_api_key): | |
| os.environ['MISTRAL_API_KEY'] = mistral_api_key | |
| from langgraph.graph import END, START, StateGraph | |
| from langchain_core.tracers.context import tracing_v2_enabled | |
| 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, | |
| ) | |
| yield "Looking at the Repo!" | |
| documents = load_github_codebase(repo_name, repo_branch) | |
| yield "Repo loaded!" | |
| # 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]} !" | |
| if "generate_mermaid_code" in s: | |
| yield "Exploration done! Gathering thoughts and generating a graph!", | |
| if "render_mermaid" in s: | |
| yield s["render_mermaid"]["messages"][-1].content | |
| 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 | |
| ), | |
| gr.Textbox( | |
| label="Mistral API Key", # Input component label | |
| # value="main", # Input component placeholder | |
| placeholder="Mistral API Key from mistral.ai", # Input component description | |
| ), | |
| ], | |
| outputs = gr.Textbox(label="Mermaid Graph", placeholder="Visualization of the functionalities of the Repo",), # Output component description | |
| title="Repo Functionality Visualizer", | |
| ) | |
| demo.launch() | |