Update app.py

app.py

@@ -25,9 +25,12 @@ class State(TypedDict):
     messages: Annotated[list, add_messages]
     name: str
     birthday: str
+    input : str
+    code = str
  
 # LLM
 llm = ChatGroq(model="deepseek-r1-distill-llama-70b")
+code_generator = 
  
 # Memory
 memory = MemorySaver()
@@ -36,13 +39,56 @@ memory = MemorySaver()
 def ai_assistance(state: State):
     result = llm.invoke(state["messages"])
     return {"messages": [result]}
- 
+
+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 (using builder.show() or similar methods).
+    5. Structure the code with proper node functions, state definitions, and graph compilation.
+
+    IMPORTANT: Output ONLY the final Python code.  
+    DO NOT include any explanations, comments, or text before, inside, or after the code.  
+    Start the output with the necessary import statements (e.g., "from langgraph import StateGraph, State, Transition").  
+    No additional text, no markdown fences, just the pure code.
+
+    User request:"""    
+    
+    full_prompt = system_prompt + user_prompt
+    code = llm.invoke(full_prompt)
+    
+    return {"input": user_prompt, "code": code.content}
+
+def router(state: State):
+    last_msg = state["messages"][-1].content.lower()
+    if any(keyword in last_msg for keyword in ["generate", "code", "script", "langgraph", "flow"]):
+        return "Generate_Code"
+    return "AI_Assistance"
+
 # Build LangGraph
+# builder = StateGraph(State)
+# builder.add_node("AI_Assistance", ai_assistance)
+# builder.add_edge(START, "AI_Assistance")
+# builder.add_edge("AI_Assistance", END)
+# graph = builder.compile(checkpointer=memory)
 builder = StateGraph(State)
 builder.add_node("AI_Assistance", ai_assistance)
-builder.add_edge(START, "AI_Assistance")
+builder.add_node("Generate_Code", generate_code)
+builder.add_node("Router", router)
+
+builder.set_entry_point("Router")
+builder.add_edge("Router", "AI_Assistance")
+builder.add_edge("Router", "Generate_Code")
 builder.add_edge("AI_Assistance", END)
+builder.add_edge("Generate_Code", END)
+
 graph = builder.compile(checkpointer=memory)
+
  
 # Streamlit UI setup
 st.set_page_config(page_title="LangGraph Chatbot", layout="wide")
@@ -111,9 +157,22 @@ with st.container():
         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))
+        
             config = {"configurable": {"thread_id": st.session_state.thread_id}}
-            result = graph.invoke({"messages": st.session_state.chat_history}, config=config)
-            st.session_state.chat_history.append(result["messages"][-1])
-            st.rerun()
+            state_input = {
+                "messages": st.session_state.chat_history,
+                "input": user_input,
+            }
+        
+            result = graph.invoke(state_input, config=config)
+        
+            # If it’s a code response, append it as-is
+            if result.get("code"):
+                st.session_state.chat_history.append(AIMessage(content=result["code"]))
+            else:
+                st.session_state.chat_history.append(result["messages"][-1])
+        
+            st.rerun()