neuralgeekroot/AutomatedCodeReview

1import os2from dotenv import load_dotenv3from langgraph.graph import StateGraph, START, END4from typing_extensions import TypedDict5from enum import Enum6from langchain_groq import ChatGroq7from langchain_core.prompts import PromptTemplate8import streamlit as st9import langsmith10 11# Load environment variables12load_dotenv()13os.environ['GROQ_API_KEY'] = os.getenv('GROQ_API_KEY')14os.environ['LANGCHAIN_API_KEY'] = os.getenv('LANGCHAIN_API_KEY')15os.environ['LANGSMITH_TRACING_V2'] = 'true'16os.environ['LANGCHAIN_PROJECT_NAME'] = os.getenv('LANGCHAIN_PROJECT_NAME')17 18class Step(Enum):19    INPUT = "input"20    REVIEW = "review"21    IMPROVISATION = "improvisation"22    APPROVAL = "approval"23    APPROVED = "approved"24 25# Define the state structure26class CodingState(TypedDict):27    code: str28    step: Step29 30# Define prompts in a configuration dictionary31PROMPTS = {32    "coder": "Create a code as per the {code} provided",33    "peer": "Review the code provided by the coder: {code}. "34            "If correction is needed, return 'improvisation'. "35            "If suggestions are needed, return 'approval'. "36            "If the code is correct, return 'approved'.",37    "manager": "Add necessary docstrings to the following code and approve it: {code}"38}39 40# Initialize the LLM41llm = ChatGroq(model="qwen-2.5-32b")42 43# Define the coder node44def coder(state):45    """Based on the input message the code is created by the coder."""46    print("Coder Node: Generating code...")47    try:48        prompt = PromptTemplate.from_template(PROMPTS["coder"])49        chain = prompt | llm50        result = chain.invoke({'code': state['code']})51        print(f"Coder Node: Generated code: {result.content}")52        return {'code': result.content, 'step': Step.REVIEW}53    except Exception as e:54        print(f"Coder Node: Error generating code - {e}")55        return {'code': state['code'], 'step': Step.IMPROVISATION} 56 57# Define the peer node58def peer(state):59    """Reviewing the code provided by the coder and determining the next step."""60    print("Peer Node: Reviewing code...")61    try:62        prompt = PromptTemplate.from_template(PROMPTS["peer"])63        chain = prompt | llm64        result = chain.invoke({'code': state['code']})65 66        # Extract the decision step from result67        decision = result.content.strip().lower()68 69        # Validate decision70        valid_decisions = [Step.IMPROVISATION.value, Step.APPROVAL.value, Step.APPROVED.value]71        if decision not in valid_decisions:72            print(f"Peer Node: Invalid decision '{decision}'. Defaulting to 'approval'.")73            decision = Step.APPROVAL.value  # Default fallback74 75        return {"code": state["code"], "step": Step(decision)}76    except Exception as e:77        print(f"Peer Node: Error reviewing code - {e}")78        return {"code": state["code"], "step": Step.IMPROVISATION} 79 80def manager(state):81    """Add docstrings to the code and approve it."""82    print("Manager Node: Adding docstrings and approving code...")83    try:84        prompt = PromptTemplate.from_template(PROMPTS["manager"])85        chain = prompt | llm86        result = chain.invoke({'code': state['code']})87        print(f"Manager Node: Approved code: {result.content}")88        return {'code': result.content, 'step': Step.APPROVED}89    except Exception as e:90        print(f"Manager Node: Error approving code - {e}")91        return {'code': state['code'], 'step': Step.APPROVAL} 92 93# Define the code validity function94def code_validity(state):95    """Determine the next step based on the current state."""96    print(f"Code Validity: Current step: {state['step'].value}")97    if state['step'] == Step.IMPROVISATION:98        return "coder"99    elif state['step'] == Step.APPROVAL:100        return "manager"101    elif state['step'] == Step.APPROVED:102        return END103 104# Build the workflow105builder = StateGraph(CodingState)106 107# Add nodes108builder.add_node("coder", coder)109builder.add_node("peer", peer)110builder.add_node("manager", manager)111 112# Add edges113builder.add_edge(START, "coder")114builder.add_edge("coder", "peer")115builder.add_conditional_edges("peer", code_validity, {"coder": "coder", "manager": "manager", END: END})116builder.add_edge("manager", END)117 118# Compile the workflow119workflow = builder.compile()120 121# Streamlit frontend122st.title("Automated Code Peer Review")123st.write("Submit your code for an automated peer review using an open-source LLM.")124 125# Text area for code input126code = st.text_area("Paste your code here:", height=300)127 128if st.button("Generate Code"):129    if code.strip() == "":130        st.error("Please paste some code to review.")131    else:132        with st.spinner("Generating review..."):133            try:134                # Invoke the workflow135                result = workflow.invoke({"code": code, 'step': Step.INPUT})136                st.success("Review Generated!")137                st.write("### Code Review Feedback")138                st.write(result['code'])139 140            except Exception as e:141                st.error(f"An error occurred: {e}")