Update app1.py

app1.py

@@ -1,214 +1,133 @@
-# app.py
-
 import os
 import streamlit as st
 from dotenv import load_dotenv
-from langchain_groq import ChatGroq
-from langgraph.graph import StateGraph, END
+from langgraph.graph import StateGraph, END, START
 from langgraph.checkpoint.memory import MemorySaver
-from langchain_core.messages import AIMessage, HumanMessage
-from typing import Annotated
-from typing_extensions import TypedDict
-from langchain_together import Together
-from tools import execute_python_code
-import io
-import contextlib
+from langchain_core.messages import HumanMessage, AIMessage
+from langchain_core.runnables import Runnable
+from langchain_together import ChatTogether
+from typing import TypedDict, List, Optional
 import traceback
-import time
+import io
+import sys
 
 # Load environment variables
 load_dotenv()
-os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")
-together_api_key = os.getenv("TOGETHER_API_KEY")
-
-# LangGraph State definition
-class State(TypedDict):
-    messages: Annotated[list, ...]
-    name: str
-    birthday: str
+llm = ChatTogether(model="meta-llama/Meta-Llama-3-8B-Instruct")
+
+# Define the graph state
+class GraphState(TypedDict):
     input: str
-    code: str
-    explanation: str
-    execution_result: str
-
-# LLM
-code_generator = Together(
-    model="deepseek-ai/DeepSeek-R1-Distill-Llama-70B-free",
-    temperature=0.2,
-    max_tokens=1500,
-    api_key=together_api_key,
-)
-
-# Memory
-memory = MemorySaver()
-
-# Define LangGraph Nodes
-def generate_code(state: State):
-    user_prompt = state["input"]
-    system_prompt = """You are an expert Python coding assistant specializing in LangGraph applications.
-Generate clean, working Python code for the user's request with these requirements:
-1. The code MUST use the LangGraph framework (langgraph library).
-2. Implement a proper flow graph using StateGraph.
-3. Include all necessary imports and make sure the code is complete.
-4. Include code to visualize the flow graph.
-5. Output ONLY the final Python code.
-
-User request:"""
-    full_prompt = system_prompt + user_prompt
-
-    for attempt in range(3):
-        try:
-            response = code_generator.invoke(full_prompt)
-            return {**state, "code": str(response)}
-        except Exception as e:
-            if "503" in str(e):
-                time.sleep(2)
-            else:
-                raise e
-    raise Exception("Code generation failed after retries.")
-
-def execute_code(state: State):
+    messages: List[HumanMessage | AIMessage]
+    code: Optional[str]
+    execution_result: Optional[str]
+    explanation: Optional[str]
+
+# === Node 1: Generate Code ===
+def generate_code(state: GraphState) -> GraphState:
+    prompt = f"""You are a senior Python developer.
+Generate Python code for the following user request. Just return the code only, no explanation.
+
+Request: {state['input']}
+"""
+    messages = state["messages"] + [HumanMessage(content=prompt)]
+    response = llm.invoke(messages)
+    return {
+        **state,
+        "messages": messages + [AIMessage(content=response.content)],
+        "code": response.content
+    }
+
+# === Node 2: Execute Code ===
+def execute_code(state: GraphState) -> GraphState:
     code = state.get("code", "")
-    buffer = io.StringIO()
     try:
-        with contextlib.redirect_stdout(buffer):
+        buffer = io.StringIO()
+        with io.StringIO() as buf, io.StringIO() as err_buf:
+            sys.stdout = buf
+            sys.stderr = err_buf
             exec(code, {})
-        output = buffer.getvalue() or "✅ Code executed successfully with no output."
+            output = buf.getvalue()
+            error = err_buf.getvalue()
+        sys.stdout = sys.__stdout__
+        sys.stderr = sys.__stderr__
+        result = output if output else "✅ Code executed without output."
+        if error:
+            result += "\n⚠️ Error:\n" + error
     except Exception:
-        output = "❌ Execution Error:\n" + traceback.format_exc()
-    return {**state, "execution_result": output}
-
-def explain_code(state: State):
-    code = state["code"]
-    user_prompt = state["input"]
-    system_prompt = """You are a LangGraph expert who explains code clearly. Provide a detailed explanation of the code in three parts:
-1. LANGGRAPH FLOW: Describe nodes, edges, and how the graph flows.
-2. CODE FLOW: High-level architecture and logic.
-3. STEP-BY-STEP: Explain each part of the code so a beginner can understand it.
-4. VISUALIZATION: Instructions on how to run and see the graph output.
+        result = "❌ Exception:\n" + traceback.format_exc()
+    return {
+        **state,
+        "execution_result": result
+    }
+
+# === Node 3: Explain Code ===
+def explain_code(state: GraphState) -> GraphState:
+    prompt = f"""You are a code explainer. Please explain the following Python code:
+{state['code']}
 """
-
-    prompt = f"User Prompt: {user_prompt}\n\nCode:\n```python\n{code}\n```"
-    full_prompt = system_prompt + prompt
-
-    explanation = code_generator.invoke(full_prompt)
-    return {**state, "explanation": explanation}
-
-# LangGraph setup
-builder = StateGraph(State)
+    messages = state["messages"] + [HumanMessage(content=prompt)]
+    response = llm.invoke(messages)
+    return {
+        **state,
+        "messages": messages + [AIMessage(content=response.content)],
+        "explanation": response.content
+    }
+
+# === Build LangGraph ===
+builder = StateGraph(GraphState)
 builder.add_node("Generate_Code", generate_code)
 builder.add_node("Execute_Code", execute_code)
-builder.add_node("Code_Explainer", explain_code)
+builder.add_node("Explain_Code", explain_code)
+
 builder.set_entry_point("Generate_Code")
-graph = builder.compile(checkpointer=memory)
-
-# Streamlit UI Setup
-st.set_page_config(page_title="MitraVerse", layout="wide")
-
-st.markdown("""
-<style>
-.stChatMessage {
-    padding: 12px;
-    margin-bottom: 12px;
-    border-radius: 12px;
-    max-width: 90%;
-}
-.user {
-    background-color: #dcf8c6;
-    align-self: flex-end;
-}
-.bot {
-    background-color: #f1f0f0;
-    align-self: flex-start;
-}
-.input-box {
-    display: flex;
-    align-items: center;
-    gap: 0.5rem;
-}
-#floating-container {
-    display: flex;
-    align-items: center;
-    justify-content: space-between;
-    padding: 0.25rem 0.75rem;
-    background-color: #f9f9f9;
-    border-radius: 0.75rem;
-    margin-top: 1rem;
-    border: 1px solid #ccc;
-}
-.floating-popup {
-    margin-top: 0.5rem;
-    padding: 0.5rem;
-    border-radius: 0.5rem;
-    border: 1px solid #ccc;
-    background-color: white;
-}
-</style>
-""", unsafe_allow_html=True)
-
-st.title("🧠 MitraVerse")
+builder.add_edge("Generate_Code", "Execute_Code")
+builder.add_edge("Execute_Code", "Explain_Code")
+builder.set_finish_point("Explain_Code")
+
+graph = builder.compile()
+
+# === Streamlit App ===
+st.set_page_config(page_title="🧠 MitraVerse", page_icon="🧠")
+st.title("🧠 MitraVerse - LangGraph Code Assistant")
 
 # Initialize session state
-if "chat_history" not in st.session_state:
-    st.session_state.chat_history = []
-if "latest_code" not in st.session_state:
-    st.session_state.latest_code = ""
-if "latest_explanation" not in st.session_state:
-    st.session_state.latest_explanation = ""
-if "latest_input" not in st.session_state:
-    st.session_state.latest_input = ""
-
-# Display chat history
-for msg in st.session_state.chat_history:
-    role = "user" if isinstance(msg, HumanMessage) else "bot"
-    st.markdown(f"<div class='stChatMessage {role}'>{msg.content}</div>", unsafe_allow_html=True)
-
-# Input form
-with st.form("chat_form", clear_on_submit=True):
-    st.markdown('<div id="floating-container">', unsafe_allow_html=True)
-    st.markdown('</div>', unsafe_allow_html=True)
-    user_input = st.text_input("Ask me", label_visibility="collapsed", placeholder="Ask me Anything")
-    submitted = st.form_submit_button(label="Send")
-
-    if submitted and user_input:
-        st.session_state.chat_history.append(HumanMessage(content=user_input))
-        st.session_state.latest_input = user_input
-
-# Buttons to run each tool manually
-if st.session_state.latest_input:
-    if st.button("🔨 Generate Code"):
-        state_input = {
-            "messages": st.session_state.chat_history,
-            "input": st.session_state.latest_input
-        }
-        result = graph.invoke(state_input, node="Generate_Code")
-        st.session_state.latest_code = result["code"]
-        st.session_state.chat_history.append(
-            AIMessage(content="**💻 Generated Code:**\n\n```python\n" + result["code"] + "\n```")
-        )
-        st.code(result["code"], language="python")
-
-    if st.button("⚙️ Execute Code") and st.session_state.latest_code:
-        state_input = {
-            "code": st.session_state.latest_code,
-            "input": st.session_state.latest_input
-        }
-        result = graph.invoke(state_input, node="Execute_Code")
-        st.session_state.chat_history.append(
-            AIMessage(content="**🧪 Execution Result:**\n\n" + result["execution_result"])
-        )
-        st.text("🧪 Execution Result:")
-        st.text(result["execution_result"])
-
-    if st.button("🧠 Explain Code") and st.session_state.latest_code:
-        state_input = {
-            "code": st.session_state.latest_code,
-            "input": st.session_state.latest_input
-        }
-        result = graph.invoke(state_input, node="Code_Explainer")
-        st.session_state.latest_explanation = result["explanation"]
-        st.session_state.chat_history.append(
-            AIMessage(content="**🔍 Code Explanation:**\n\n" + result["explanation"])
-        )
-        with st.expander("🔍 Code Explanation"):
-            st.markdown(result["explanation"])
+for key in ["chat_history", "latest_code", "latest_explanation", "execution_result"]:
+    if key not in st.session_state:
+        st.session_state[key] = [] if key == "chat_history" else ""
+
+# User input form
+with st.container():
+    with st.form("chat_form", clear_on_submit=True):
+        user_input = st.text_input("Ask me anything", placeholder="e.g., Write a bubble sort in Python")
+        submitted = st.form_submit_button("🚀 Run End-to-End")
+
+        if submitted and user_input:
+            st.session_state.chat_history.append(HumanMessage(content=user_input))
+            state_input = {
+                "messages": st.session_state.chat_history,
+                "input": user_input,
+                "code": "",
+                "execution_result": "",
+                "explanation": ""
+            }
+            result = graph.invoke(state_input)
+            st.session_state.latest_code = result["code"]
+            st.session_state.execution_result = result["execution_result"]
+            st.session_state.latest_explanation = result["explanation"]
+
+# Show generated code
+if st.session_state.latest_code:
+    st.subheader("🧾 Generated Code")
+    st.code(st.session_state.latest_code, language="python")
+    st.download_button("📥 Download Code", st.session_state.latest_code, file_name="generated_code.py")
+
+# Show execution result
+if st.session_state.execution_result:
+    st.subheader("🧪 Execution Result")
+    st.text(st.session_state.execution_result)
+
+# Show code explanation
+if st.session_state.latest_explanation:
+    st.subheader("💡 Code Explanation")
+    st.markdown(st.session_state.latest_explanation)