utils.py

from langchain.agents import tool
from typing import Literal
import json
from PIL import Image

from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, ToolMessage
from langgraph.graph import END, MessagesState

from render_mermaid import render_mermaid
from langchain_community.document_loaders import GithubFileLoader

# from langchain_ollama import ChatOllama
from prompts import *
from constants import file_extensions
from __init__ import llm, llm_structured


class GraphState(MessagesState):
    working_knowledge: str
    all_files: list[str]
    remaining_files: list[str]
    explored_files: list[str]
    explored_summaries: str
    document_summaries_store: dict
    documents: list
    final_graph: Image


def load_github_codebase(repo: str, branch: str):
    loader = GithubFileLoader(
        repo=repo,  # the repo name
        branch=branch,  # the branch name
        github_api_url="https://api.github.com",
        file_filter=lambda file_path: file_path.endswith(tuple(file_extensions)),
        # file_filter=lambda filepath: True,
        encoding="utf-8",
    )
    documents = loader.load()
    return documents


def get_file_content_summary(file_path: str, state: GraphState):
    """Returns the functional summary of a file. Please note that the file_path should not be null.

    Args:
        file_path: The path of the file for which the summary is required."""

    summary = check_summary_in_store(file_path, state)
    if summary:
        return summary
    for document in state["documents"]:
        if document.metadata["path"] == file_path:
            doc_content = document.page_content
            break
    # print(content)
    summary = llm.invoke(
        [SystemMessage(content=summarizer_prompt), HumanMessage(content=doc_content)]
    ).content
    summary = json.dumps({"FilePath": file_path, "Summary": summary})
    save_summary_in_store(file_path, summary, state)
    return summary


def explore_file(state: GraphState):
    file_path = state["remaining_files"].pop()

    summary_dict = json.loads(get_file_content_summary(file_path, state))
    if summary_dict["FilePath"] in state["explored_files"]:
        return state
    knowledge_str = f"""* File Path: {summary_dict['FilePath']}\n\tSummary: {summary_dict['Summary']}\n\n"""
    state["explored_summaries"] += knowledge_str
    state["explored_files"].append(file_path)
    return state


@tool
def generate_final_mermaid_code():
    """Generate the final mermaid code for the codebase once all the files are explored and the working knowledge is complete."""
    return "generate_mermaid_code"


def check_summary_in_store(file_path: str, state: GraphState):
    if file_path in state["document_summaries_store"]:
        return state["document_summaries_store"][file_path]
    return None


def save_summary_in_store(file_path: str, summary: str, state: GraphState):
    state["document_summaries_store"][file_path] = summary


def get_all_filesnames_in_codebase(state: GraphState):
    """Get a list of all files (as filepaths) in the codebase."""
    filenames = []
    for document in state["documents"]:
        filenames.append(document.metadata["path"])

    return {
        "all_files": filenames,
        "explored_files": [],
        "remaining_files": filenames,
        "explored_summaries": "",
        "document_summaries_store": {},
    }


def parse_plan(state: GraphState):
    """Parse the plan and return the next action."""
    if "File Exploration Plan" in state["working_knowledge"]:
        plan_working = state["working_knowledge"].split("File Exploration Plan")[1]
    else:
        plan_working = state["working_knowledge"]
    response = llm_structured.invoke(plan_parser.format(plan_list=plan_working))[
        "plan_list"
    ]
    if len(response) > 25:
        response = response[:25]
    # response = eval(llm.invoke(plan_parser.format(plan_list=plan_working)).content)
    return {"remaining_files": response}


def router(state: GraphState):
    """Route the conversation to the appropriate node based on the current state of the conversation."""
    if state["remaining_files"] != []:
        return "explore_file"
    else:
        return "generate_mermaid_code"


def get_plan_for_codebase(state: GraphState):
    new_state = get_all_filesnames_in_codebase(state)
    planner_content = "# File Structure\n" + str(new_state["all_files"])
    plan = llm.invoke(
        [SystemMessage(content=planner_prompt), HumanMessage(content=planner_content)]
    )

    knowledge_str = f"""# Plan\n{plan.content}"""
    new_state["working_knowledge"] = knowledge_str
    # print(new_state)
    return new_state


def final_mermaid_code_generation(state: GraphState):
    final_graph_content = (
        "# Disjoint Codebase Understanding\n"
        + state["working_knowledge"]
        + "\n\n# Completed Explorations\n"
        + state["explored_summaries"]
    )
    response = llm.invoke(
        [
            SystemMessage(content=final_graph_prompt),
            HumanMessage(content=final_graph_content),
        ]
    )
    return {"messages": [response]}


import time


def extract_mermaid_and_generate_graph(state: GraphState):
    mermaid_code = state["messages"][-1].content
    if "mermaid" in mermaid_code:
        mermaid_code = mermaid_code.split("mermaid")[-1]
    response = llm.invoke(
        [SystemMessage(content=mermaid_extracter), HumanMessage(content=mermaid_code)]
    ).content
    response = response.split("```mermaid")[-1].split("```")[0]
    # Save the mermaid code in a file with the current timestamp
    # print(response)
    file_name = f"mermaid/{int(time.time())}.png"
    render_mermaid(response, file_name)

    # Read image to return as output
    img = Image.open(file_name)
    return {"messages": [AIMessage(response)], "final_graph": img}


def need_to_update_working_knowledge(state: GraphState):
    messages = state["messages"]
    last_message = messages[-1]
    # prev_to_last_message = messages[-2]
    # If the last call is a tool message, we need to update the working knowledge
    if last_message.content == "generate_mermaid_code":
        return "generate_mermaid_code"
    if isinstance(last_message, ToolMessage):
        return "tools_knowledge_update"
    # Otherwise, we continue with the agent
    return "agent"