Update app.py

app.py

@@ -15,6 +15,12 @@ from tools import execute_python_code, web_search, deep_think
 import io
 import contextlib
 import traceback
+from langchain_core.messages import AIMessage
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain.prompts import PromptTemplate
+from langchain.chains import LLMChain
+from langchain_core.messages import AIMessage
+from typing import List
 # Load environment
 load_dotenv()
 # os.environ["TAVILY_API_KEY"] = os.getenv("TAVILY_API_KEY")
@@ -50,18 +56,81 @@ def ai_assistance(state: State):
     result =code_generator.invoke(state["messages"])
     return {"messages": state['messages']+[result]}
 
+# def agent_node(state: State):
+#     # Use your LLM here (e.g., Together, OpenAI, etc.)
+#     model = ChatGoogleGenerativeAI(model = "gemini-2.0-flash-001").bind_tools(tools)
+#     follow_up_prompt = "Break down this task into subtasks and ask follow-up questions if needed:\n\n"
+#     last_user_msg = state["messages"][-1].content
+#     full_prompt = follow_up_prompt + last_user_msg
+
+#     response = model.invoke(full_prompt)
+
+#     return {"messages": state["messages"] + [AIMessage(content=response)]}
+
+llm = ChatGroq(temperature=0)
+
+# Template to extract subtasks from the user's input
+subtask_prompt = PromptTemplate.from_template(
+    """You are an expert AI agent designer.
+
+Given the user's goal:
+"{user_goal}"
+
+1. Break this goal into a clear list of subtasks (in bullet points).
+2. If any clarification is needed, ask relevant follow-up questions.
+
+Respond in this format:
+---
+Subtasks:
+- ...
+- ...
+Follow-Up Questions (if any):
+- ...
+---"""
+)
+
+subtask_chain = LLMChain(llm=llm, prompt=subtask_prompt)
+
 def agent_node(state: State):
-    # Use your LLM here (e.g., Together, OpenAI, etc.)
-    model = ChatGoogleGenerativeAI(model = "gemini-2.0-flash-001").bind_tools(tools)
-    follow_up_prompt = "Break down this task into subtasks and ask follow-up questions if needed:\n\n"
-    last_user_msg = state["messages"][-1].content
-    full_prompt = follow_up_prompt + last_user_msg
+    user_input = state["messages"][-1].content
 
-    response = model.invoke(full_prompt)
+    # Get subtasks and possible questions
+    result = subtask_chain.invoke({"user_goal": user_input})
+    response_text = result["text"]
 
-    return {"messages": state["messages"] + [AIMessage(content=response)]}
+    # Parse subtasks and follow-up questions
+    subtasks, questions = parse_subtasks_and_questions(response_text)
+
+    # Append AI response to messages
+    state["messages"].append(AIMessage(content=response_text))
+
+    # Save subtasks and questions into state
+    return {
+        "messages": state["messages"],
+        "subtasks": subtasks,
+        "follow_up_questions": questions
+    }
+
+# ✂️ Helper function to parse bullet points
+def parse_subtasks_and_questions(text: str):
+    subtasks = []
+    questions = []
+
+    collecting = None
+    for line in text.strip().splitlines():
+        line = line.strip()
+        if line.lower().startswith("subtasks:"):
+            collecting = "subtasks"
+        elif line.lower().startswith("follow-up questions"):
+            collecting = "questions"
+        elif line.startswith("-"):
+            if collecting == "subtasks":
+                subtasks.append(line[1:].strip())
+            elif collecting == "questions":
+                questions.append(line[1:].strip())
+
+    return subtasks, questions
 
-    
 import time
 def generate_code(state: State):
     user_prompt = state["input"]
@@ -203,18 +272,22 @@ def router(state):
     else:
         return "AI_Assistance"
 builder = StateGraph(State)
+
+# Define Nodes
 builder.add_node("AI_Assistance", ai_assistance)
+builder.add_node("LLM_Agent", agent_node)
 builder.add_node("Generate_Code", generate_code)
-# builder.add_node("LLM_Agent", agent_node)
 builder.add_node("Code_Explainer", explain_code)
 
+# Entry Point
 builder.set_entry_point("AI_Assistance")
 
-builder.add_edge("AI_Assistance", "Generate_Code")
-# builder.add_edge("Generate_Code", "LLM_Agent")
-# builder.add_edge("LLM_Agent", "Code_Explainer")
-builder.add_edge("Generate_Code","Code_Explainer")
+# Define Flow
+builder.add_edge("AI_Assistance", "LLM_Agent")
+builder.add_edge("LLM_Agent", "Generate_Code")
+builder.add_edge("Generate_Code", "Code_Explainer")
 builder.add_edge("Code_Explainer", END)
+
 graph = builder.compile(checkpointer=memory) 
 # Streamlit UI setup
 st.set_page_config(page_title="MitraVerse", layout="wide")