added agent fixes

Signed-off-by: Arya Pratap Singh <[email protected]>

agent/dsa_agent/agent.py

@@ -2,7 +2,6 @@
 from dsa_agent.state import AgentState
 from dsa_agent.nodes import (
     retrieve_question,
-    update_question,
     code_generation_in_node,
     visualize_code,
     complexity_analysis,
@@ -21,7 +20,7 @@
 
 workflow = StateGraph(AgentState)
 
-workflow.add_node("update_question", update_question)
+# workflow.add_node("update_question", update_question)
 workflow.add_node("retrieve_question", retrieve_question)
 workflow.add_node("code_generation_in_node", code_generation_in_node)
 workflow.add_node("visualize_code", visualize_code)
@@ -32,67 +31,61 @@
 workflow.add_conditional_edges(
     START,
     new_question,
-    {
-        "update_question": "update_question",
-        "retrieve_question": "retrieve_question",
-    },
-)
+    {"retrieve_question": "retrieve_question", "no_context": "no_context"},
 
+)
 
 
 workflow.add_edge("retrieve_question", "code_generation_in_node")
 
 
-
 workflow.add_conditional_edges(
     "code_generation_in_node",
     decide_to_generate_code,
     {
         "useful": "visualize_code",
-        "code_generation_in_node": "code_generation_in_node",
+        "no_context": "no_context",
     },
 )
 
 
 
-workflow.add_edge("visualize_code", "retrieve_question")
+# workflow.add_edge("visualize_code", "retrieve_question")
 
 
 workflow.add_conditional_edges(
     "visualize_code",
     decide_to_generate_visualisation,
     {
         "useful": "complexity_analysis",
-        "visualize_code": "visualize_code",
+        "no_context": "no_context",
     },
 )
 
-workflow.add_edge("complexity_analysis", "retrieve_question")
+# workflow.add_edge("complexity_analysis", "retrieve_question")
 
 
 workflow.add_conditional_edges(
     "complexity_analysis",
     decide_to_generate_complexity,
     {
         "useful": "explaination_code",
-        "complexity_analysis": "complexity_analysis",
+        "no_context": "no_context",
     },
 )
 
-workflow.add_edge("explaination_code", "retrieve_question")
+# workflow.add_edge("explaination_code", "retrieve_question")
 
 
 workflow.add_conditional_edges(
     "explaination_code",
     decide_to_generate_explanation,
     {
         "useful": END,
-        "explaination_code": "explaination_code",
+        "no_context": "no_context",
     },
 )
 
-workflow.add_edge("update_question", END)
-
 workflow.add_edge("no_context", END)
 
 graph = workflow.compile(checkpointer=MemorySaver())

agent/dsa_agent/edges/__init__.py

@@ -6,14 +6,16 @@
 from dsa_agent.edges.explanation_verify import explaination_verify
 
 def new_question(state):
-    messages = state["messages"]
-    last_message = messages[-1]
+    question = state["question"]
+    # last_message = messages[-1]
 
     if (
-        isinstance(last_message, SystemMessage)
-        and "QUESTION UPDATED" in last_message.content
+        isinstance(question, SystemMessage)
+        and "QUESTION UPDATED" in question.content
     ):
         return "update_question"
+
+    print("init edges new question : ", question)
 
     return "retrieve_question"
 
@@ -32,6 +34,7 @@ def decide_to_generate_code(state):
     print("---ASSESS GENERATED CODE---")
     code = state["code"]
     question = state["question"]
+    # testCases = state["testCases"]
 
     print("---GRADE CODE---")
     score = code_verify.invoke({"question": question, "code": code})

agent/dsa_agent/edges/code_verify.py

@@ -16,12 +16,13 @@ class VerifyCode(BaseModel):
 code_verification = llm.with_structured_output(VerifyCode)
 
 system = """You a 4000 rated Competitive Programmer that writes optimised and fully detailed working code in python \n 
-     Given a question and testCases you are expected to write a working code in python that passes all the testCases and is generic and well optimised \n"""
+     Given a question you are expected to write a working code in python that is generic and well optimised \n"""
+
 
 code_prompt = ChatPromptTemplate.from_messages(
     [
         ("system", system),
-        ("human", "User question: \n\n {question} \n\n Code Generated: \n\n {code}"),
+        ("human", "User question: \n\n {question} \n\n Code Generated: \n\n {code} \n\n"),
     ]
 )
 

agent/dsa_agent/nodes/__init__.py

@@ -16,9 +16,8 @@ def retrieve_question(state):
         state (dict): New key added to state, documents, that contains retrieved documents
     """
     print("---GETTING QUESTION---")
-    messages = state["messages"]
-    question = messages[-1].content
-
+    question = state["question"]
+    print(f"---nodes new QUESTION DETECTED: {question}---")
     return {
         **state,
         "question": question,
@@ -49,9 +48,12 @@ def code_generation_in_node(state):
 
     print("---TRANSFORM QUERY---")
     question = state["question"]
-    testCases = state["testCases"]
+    # testCases = state["testCases"]
+    # print(testCases)
+    # if testCases == []:
+        # testCases = ["X RETURNS Y"]
 
-    code_generated = code_writer.invoke({"question": question , "testCases": testCases})
+    code_generated = code_writer.invoke({"question": question })
     return {**state, "code": code_generated}
 
 def visualize_code(state):
@@ -85,7 +87,7 @@ def complexity_analysis(state):
     code = state["code"]
 
     complexity = complexity_generated.invoke({"code": code})
-    return {**state, "time_complexity": list(complexity)[0], "space_complexity": list(complexity)[1]}
+    return {**state, "time_complexity": complexity, "space_complexity": complexity}
 
 def explaination_code(state):
     '''
@@ -107,7 +109,7 @@ def explaination_code(state):
 def no_context(state):
     print("---NO CONTEXT---")
 
-    messages = state["messages"]
+    messages = state["question"]
     messages.append(HumanMessage("I'm sorry, I can't find any relevant information."))
 
     return state

agent/dsa_agent/nodes/coding.py

@@ -8,14 +8,14 @@
 llm = ChatGroq(model="llama-3.3-70b-versatile", temperature=0)
 
 system = """You a 4000 rated Competitive Programmer that writes optimised and fully detailed working code in python \n 
-     Given a question and testCases you are expected to write a working code in python that passes all the testCases and is generic and well optimised \n"""
+     Given a question you are expected to write a working code in python that is generic and well optimised \n"""
 
 re_write_prompt = ChatPromptTemplate.from_messages(
     [
         ("system", system),
         (
             "human",
-            "Here is the question: \n\n {question} \n and here are the test cases : \n\n {testCases}. Formulate an optimise answer in python.",
+            "Here is the question: \n\n {question} \n. Formulate an optimised answer in python.",
         ),
     ]
 )

agent/dsa_agent/state.py

@@ -9,18 +9,18 @@ class AgentState(TypedDict):
 
     Attributes:
         question: str
-        testCases: List[str]
+        # testCases: List[str]
         code: Optional[str]
         explanation: Optional[str]
         time_complexity: Optional[str]
         space_complexity: Optional[str]
         visualization: Optional[str]
-        messages: list of messages
+        # messages: list of messages
     """
 
     question: str
-    testCases: List[str]
-    messages: Annotated[Sequence[BaseMessage], add_messages]
+    # testCases: Optional[List[str]] # type: ignore
+    # messages: Annotated[Sequence[BaseMessage], add_messages]
     code: Optional[str] 
     explanation: Optional[str]
     time_complexity: Optional[str] 

backend/tools/visualisation_generator.py

@@ -8,9 +8,9 @@ def generate_mermaid_visualization(self, code: str) -> str:
         
         Create a flowchart using Mermaid syntax. Focus on key steps and decision points.
         The diagram should start with 'flowchart TD' and use proper Mermaid syntax.
-        Include clear labels and connections between nodes.
+        Include clear labels and connections between nodes. Ensure the diagram is easy to understand and follow.
 
-        DO NOT INCLUDE ANYTHING ELSE JUST THE MERMAID CODE FOR THE FLOWCHART AND MAKE IT WELL STRUCTURED AND LABELLED.
+        DO NOT INCLUDE ANYTHING ELSE JUST THE MERMAID CODE FOR THE FLOWCHART AND MAKE IT WELL STRUCTURED AND LABELLED AND ONLY AND ONLY RETURN THE MERMAID CODE.
         """
 
         response = self.model.invoke([HumanMessage(content=prompt)]) # type: ignore

frontend-interface/app/layout.tsx

@@ -20,7 +20,7 @@ export default function RootLayout({
 }) {
   return (
     <html lang="en">
-      <body className={comfortaa.className}>
+      <body suppressHydrationWarning className={comfortaa.className}>
         <CopilotKit runtimeUrl="/api/copilotkit" agent="dsa_agent">
           {children}
         </CopilotKit>

frontend-interface/app/page.tsx

@@ -150,21 +150,8 @@ const DSASolutionInterface = () => {
 
   const handleRetry = () => {
     setShowChat(true);
-    const initialMessage = `I need an alternative solution for the following problem:
-
-    Problem: ${question}
-
-    Test Cases: ${testCases
-      .map((tc) => `Input: ${tc.input}, Output: ${tc.output}`)
-      .join("\n")}
-
-    Expected Time Complexity: ${timeComplexity}
-    Expected Space Complexity: ${spaceComplexity}
-
-    Previous Solution:
-    ${solution?.code}
-
-    Please provide a different approach or optimization.`;
+    const initialMessage = `
+    question: ${question}`;
 
     run(() => new TextMessage({ role: Role.System, content: initialMessage }));
   };
@@ -173,8 +160,7 @@ const DSASolutionInterface = () => {
   const { state, setState, run } = useCoAgent<AgentState>({
     name: "dsa_agent",
     initialState: {
-      question: "",
-      testCases: [],
+      question: "find greatest of 3 numbers",
       timeComplexity: "",
       spaceComplexity: "",
       code: "",
@@ -186,7 +172,7 @@ const DSASolutionInterface = () => {
   const { isLoading, appendMessage, visibleMessages } = useCopilotChat();
 
   return (
-    <div className="flex min-h-screen bg-gradient-to-br from-gray-50 to-gray-100 dark:from-gray-900 dark:to-gray-800">
+    <div  className="flex min-h-screen bg-gradient-to-br from-gray-50 to-gray-100 dark:from-gray-900 dark:to-gray-800">
       <ChatInterface
         isLoading={isLoading}
         appendMessage={appendMessage}