Spaces:
Running
Running
import os | |
from dotenv import load_dotenv | |
from typing import TypedDict | |
from langchain_groq import ChatGroq | |
from pydantic import BaseModel, Field | |
from langsmith import traceable | |
from langgraph.graph import StateGraph, START, END | |
from IPython.display import Image, display | |
load_dotenv() | |
os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY") | |
os.environ["LANGCHAIN_TRACING_V2"] = "true" | |
os.environ["LANGCHAIN_API_KEY"] = os.getenv("LANGCHAIN_API_KEY") | |
llm = ChatGroq(model="qwen-2.5-32b") | |
# Graph state | |
class State(TypedDict): | |
code_snippet: str # input | |
readability_feedback: str # intermediate | |
security_feedback: str # intermediate | |
best_practices_feedback: str # intermediate | |
feedback_aggregator: str # output | |
# Nodes | |
def get_readability_feedback(state: State): | |
"""First LLM call to check code readability""" | |
msg = llm.invoke( | |
f"Provide readability feedback for the following code:\n\n {state['code_snippet']}" | |
) | |
return {"readability_feedback": msg.content} | |
def get_security_feedback(state: State): | |
"""Second LLM call to check for security vulnerabilities in code""" | |
msg = llm.invoke( | |
f"Check for potential security vulnerabilities in the following code and provide feedback:\n\n {state['code_snippet']}" | |
) | |
return {"security_feedback": msg.content} | |
def get_best_practices_feedback(state: State): | |
"""Third LLM call to check for adherence to coding best practices""" | |
msg = llm.invoke( | |
f"Evaluate the adherence to coding best practices in the following code and provide feedback:\n\n {state['code_snippet']}" | |
) | |
return {"best_practices_feedback": msg.content} | |
def aggregate_feedback(state: State): | |
"""Combine all the feedback from the three LLM calls into a single output""" | |
combined = f"Here's the overall feedback for the code:\n\n" | |
combined += f"READABILITY FEEDBACK:\n{state['readability_feedback']}\n\n" | |
combined += f"SECURITY FEEDBACK:\n{state['security_feedback']}\n\n" | |
combined += f"BEST PRACTICES FEEDBACK:\n{state['best_practices_feedback']}" | |
return {"feedback_aggregator": combined} | |
# Build workflow | |
parallel_builder = StateGraph(State) | |
# Add nodes - Corrected node names | |
parallel_builder.add_node("get_readability_feedback", get_readability_feedback) | |
parallel_builder.add_node("get_security_feedback", get_security_feedback) | |
parallel_builder.add_node("get_best_practices_feedback", get_best_practices_feedback) | |
parallel_builder.add_node("aggregate_feedback", aggregate_feedback) | |
# Add edges to connect nodes | |
parallel_builder.add_edge(START, "get_readability_feedback") | |
parallel_builder.add_edge(START, "get_security_feedback") | |
parallel_builder.add_edge(START, "get_best_practices_feedback") | |
parallel_builder.add_edge("get_readability_feedback", "aggregate_feedback") | |
parallel_builder.add_edge("get_security_feedback", "aggregate_feedback") | |
parallel_builder.add_edge("get_best_practices_feedback", "aggregate_feedback") | |
parallel_builder.add_edge("aggregate_feedback", END) | |
parallel_workflow = parallel_builder.compile() | |
# Show workflow - Try and except to handle the timeout | |
try: | |
display(Image(parallel_workflow.get_graph().draw_mermaid_png())) | |
except Exception as e: | |
print(f"Error generating Mermaid diagram: {e}") | |
# Invoke | |
# Here is an example of a program, you can change it for any python code. | |
full_program = """ | |
import os | |
from dotenv import load_dotenv | |
load_dotenv() | |
print(os.getenv("LANGCHAIN_API_KEY")) | |
""" | |
state = parallel_workflow.invoke({"code_snippet": full_program}) | |
print(state["feedback_aggregator"]) | |