diff --git a/AGENTS.md b/AGENTS.md
new file mode 100644
index 0000000..30c7082
--- /dev/null
+++ b/AGENTS.md
@@ -0,0 +1,88 @@
+# AGENTS.md
+
+This file provides guidance to AI agents when working with code in this repository.
+
+## Project Overview
+
+This is a chat function connecting students to an AI educational chatbot that is integrated with the **Lambda-Feedback** educational platform. It deploys as an AWS Lambda function (containerized via Docker) that receives student chat messages with educational context and returns LLM-powered chatbot responses.
+
+## Commands
+
+**Testing:**
+```bash
+pytest                              # Run all unit tests
+python tests/manual_agent_run.py   # Test agent locally with example inputs
+python tests/manual_agent_requests.py  # Test running Docker container
+```
+
+**Docker:**
+```bash
+docker build -t llm_chat .
+docker run --env-file .env -p 8080:8080 llm_chat
+```
+
+**Manual API test (while Docker is running):**
+```bash
+curl -X POST http://localhost:8080/2015-03-31/functions/function/invocations \
+  -H 'Content-Type: application/json' \
+  -d '{"body":"{\"conversationId\": \"12345Test\", \"messages\": [{\"role\": \"USER\", \"content\": \"hi\"}], \"user\": {\"type\": \"LEARNER\"}}"}'
+```
+
+**Run a single test:**
+```bash
+pytest tests/test_module.py        # Run specific test file
+pytest tests/test_index.py::test_function_name  # Run specific test
+```
+
+## Architecture
+
+### Request Flow
+
+```
+Lambda event → index.py (handler)
+  → validates via lf_toolkit ChatRequest schema
+  → src/module.py (chat_module)
+    → extracts muEd API context (messages, conversationId, question context, user type)
+    → parses educational context to prompt text via src/agent/context.py
+    → src/agent/agent.py (BaseAgent / LangGraph)
+      → routes to call_llm or summarize_conversation node
+      → calls LLM provider (OpenAI / Google / Azure / Ollama)
+  → returns ChatResponse (output, summary, conversationalStyle, processingTime)
+```
+
+### Key Files
+
+| File | Role |
+|------|------|
+| `index.py` | AWS Lambda entry point; parses event body, validates schema |
+| `src/module.py` | Transforms muEd API request → invokes agent → builds ChatResponse |
+| `src/agent/agent.py` | LangGraph stateful graph; manages message history and summarization |
+| `src/agent/prompts.py` | System prompts for tutor behavior, summarization, style detection |
+| `src/agent/llm_factory.py` | Factory classes for each LLM provider (OpenAI, Google, Azure, Ollama) |
+| `src/agent/context.py` | Converts muEd question/submission context dicts to LLM prompt text |
+| `tests/utils.py` | Shared test helpers: `assert_valid_chat_request`, `assert_valid_chat_response` |
+| `tests/example_inputs/` | Real muEd payloads used for end-to-end tests |
+
+### Agent Logic (LangGraph)
+
+`BaseAgent` maintains a state graph with two nodes:
+- **`call_llm`**: Invokes the LLM with system prompt + conversation summary + conversational style preference
+- **`summarize_conversation`**: Triggered when message count exceeds ~11; summarizes history and also extracts the student's preferred conversational style
+
+Messages are trimmed after summarization to keep context window manageable. The `summary` and `conversationalStyle` fields persist across calls via the `ChatRequest` metadata.
+
+### muEd API Format
+
+`src/module.py` handles the muEd request format (https://mued.org/). The `context` field in `ChatRequest` contains nested educational data (question parts, student submissions, task info) that gets parsed into a tutoring prompt via `src/agent/context.py`.
+
+### LLM Configuration
+
+LLM provider and model are set via environment variables (see `.env.example`). The `llm_factory.py` selects the provider at runtime. The Lambda function name/identity is set in `config.json`.
+
+The agent uses **two separate LLM instances** — `self.llm` for chat responses and `self.summarisation_llm` for conversation summarisation and style analysis. By default both use the same provider, but you can point them at different models (e.g. a cheaper model for summarisation) by changing the class in `agent.py`.
+
+## Deployment
+
+- Pushing to `dev` branch triggers the dev deployment GitHub Actions workflow
+- Pushing to `main` triggers staging deployment, with manual approval required for production
+- All environment variables (API keys, model names) are injected via GitHub Actions secrets/variables — do not hardcode them
diff --git a/CLAUDE.md b/CLAUDE.md
new file mode 100644
index 0000000..b2f1d63
--- /dev/null
+++ b/CLAUDE.md
@@ -0,0 +1,88 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Project Overview
+
+This is a chat function connecting students to an AI educational chatbot that is integrated with the **Lambda-Feedback** educational platform. It deploys as an AWS Lambda function (containerized via Docker) that receives student chat messages with educational context and returns LLM-powered chatbot responses.
+
+## Commands
+
+**Testing:**
+```bash
+pytest                              # Run all unit tests
+python tests/manual_agent_run.py   # Test agent locally with example inputs
+python tests/manual_agent_requests.py  # Test running Docker container
+```
+
+**Docker:**
+```bash
+docker build -t llm_chat .
+docker run --env-file .env -p 8080:8080 llm_chat
+```
+
+**Manual API test (while Docker is running):**
+```bash
+curl -X POST http://localhost:8080/2015-03-31/functions/function/invocations \
+  -H 'Content-Type: application/json' \
+  -d '{"body":"{\"conversationId\": \"12345Test\", \"messages\": [{\"role\": \"USER\", \"content\": \"hi\"}], \"user\": {\"type\": \"LEARNER\"}}"}'
+```
+
+**Run a single test:**
+```bash
+pytest tests/test_module.py        # Run specific test file
+pytest tests/test_index.py::test_function_name  # Run specific test
+```
+
+## Architecture
+
+### Request Flow
+
+```
+Lambda event → index.py (handler)
+  → validates via lf_toolkit ChatRequest schema
+  → src/module.py (chat_module)
+    → extracts muEd API context (messages, conversationId, question context, user type)
+    → parses educational context to prompt text via src/agent/context.py
+    → src/agent/agent.py (BaseAgent / LangGraph)
+      → routes to call_llm or summarize_conversation node
+      → calls LLM provider (OpenAI / Google / Azure / Ollama)
+  → returns ChatResponse (output, summary, conversationalStyle, processingTime)
+```
+
+### Key Files
+
+| File | Role |
+|------|------|
+| `index.py` | AWS Lambda entry point; parses event body, validates schema |
+| `src/module.py` | Transforms muEd API request → invokes agent → builds ChatResponse |
+| `src/agent/agent.py` | LangGraph stateful graph; manages message history and summarization |
+| `src/agent/prompts.py` | System prompts for tutor behavior, summarization, style detection |
+| `src/agent/llm_factory.py` | Factory classes for each LLM provider (OpenAI, Google, Azure, Ollama) |
+| `src/agent/context.py` | Converts muEd question/submission context dicts to LLM prompt text |
+| `tests/utils.py` | Shared test helpers: `assert_valid_chat_request`, `assert_valid_chat_response` |
+| `tests/example_inputs/` | Real muEd payloads used for end-to-end tests |
+
+### Agent Logic (LangGraph)
+
+`BaseAgent` maintains a state graph with two nodes:
+- **`call_llm`**: Invokes the LLM with system prompt + conversation summary + conversational style preference
+- **`summarize_conversation`**: Triggered when message count exceeds ~11; summarizes history and also extracts the student's preferred conversational style
+
+Messages are trimmed after summarization to keep context window manageable. The `summary` and `conversationalStyle` fields persist across calls via the `ChatRequest` metadata.
+
+### muEd API Format
+
+`src/module.py` handles the muEd request format (https://mued.org/). The `context` field in `ChatRequest` contains nested educational data (question parts, student submissions, task info) and the `user` field contains user-specific information (e.g., user type, preferences, task progress) that gets parsed into a tutoring prompt via `src/agent/context.py`.
+
+### LLM Configuration
+
+LLM provider and model are set via environment variables (see `.env.example`). The `llm_factory.py` selects the provider at runtime. The Lambda function name/identity is set in `config.json`.
+
+The agent uses **two separate LLM instances** — `self.llm` for chat responses and `self.summarisation_llm` for conversation summarisation and style analysis. By default both use the same provider, but you can point them at different models (e.g. a cheaper model for summarisation) by changing the class in `agent.py`.
+
+## Deployment
+
+- Pushing to `dev` branch triggers the dev deployment GitHub Actions workflow
+- Pushing to `main` triggers staging deployment, with manual approval required for production
+- All environment variables (API keys, model names) are injected via GitHub Actions secrets/variables — do not hardcode them
diff --git a/README.md b/README.md
index 1bf3503..fb03792 100755
--- a/README.md
+++ b/README.md
@@ -81,7 +81,9 @@ Also, don't forget to update or delete the Quickstart chapter from the `README.m
 
 You can create your own invocation to your own agents hosted anywhere. Copy or update the `agent.py` from `src/agent/` and edit it to match your LLM agent requirements. Import the new invocation in the `module.py` file.
 
-You agent can be based on an LLM hosted anywhere, you have available currently OpenAI, AzureOpenAI, and Ollama models but you can introduce your own API call in the `src/agent/utils/llm_factory.py`.
+Your agent can be based on an LLM hosted anywhere. OpenAI, Google AI, Azure OpenAI, and Ollama are available out of the box via `src/agent/llm_factory.py`, and you can add your own provider there too.
+
+The agent uses **two separate LLM instances** — `self.llm` for chat responses and `self.summarisation_llm` for conversation summarisation and style analysis. By default both use the same provider, but you can point them at different models (e.g. a cheaper or faster model for summarisation) by changing the class in `agent.py`.
 
 ### Prerequisites
 
@@ -99,13 +101,17 @@ You agent can be based on an LLM hosted anywhere, you have available currently O
 ├── docs/                                 # docs for devs and users
 ├── src/
 │   ├── agent/
-│   │   ├── utils/                        # utils for the agent, including the llm_factory
-│   │   ├── agent.py                      # the agent logic
-│   │   └── prompts.py                    # the system prompts defining the behaviour of the chatbot
-│   └── module.py                         
+│   │   ├── agent.py                      # LangGraph stateful agent logic
+│   │   ├── context.py                    # converts muEd context dicts to LLM prompt text
+│   │   ├── llm_factory.py                # factory classes for each LLM provider
+│   │   └── prompts.py                    # system prompts defining the behaviour of the chatbot
+│   └── module.py
 └── tests/                                # contains all tests for the chat function
+    ├── example_inputs/                   # muEd example payloads for end-to-end tests
     ├── manual_agent_requests.py          # allows testing of the docker container through API requests
     ├── manual_agent_run.py               # allows testing of any LLM agent on a couple of example inputs
+    ├── utils.py                          # shared test helpers
+    ├── test_example_inputs.py            # pytests for the example input files
     ├── test_index.py                     # pytests
     └── test_module.py                    # pytests
 ```
@@ -165,7 +171,7 @@ This will start the chat function and expose it on port `8080` and it will be op
 ```bash
 curl --location 'http://localhost:8080/2015-03-31/functions/function/invocations' \
 --header 'Content-Type: application/json' \
---data '{"body":"{\"message\": \"hi\", \"params\": {\"conversation_id\": \"12345Test\", \"conversation_history\": [{\"type\": \"user\", 
+--data '{"body":"{\"conversationId\": \"12345Test\", \"messages\": [{\"role\": \"USER\", \"content\": \"hi\"}], \"user\": {\"type\": \"LEARNER\"}}"}'
 ```
 
 #### Call Docker Container
@@ -184,21 +190,98 @@ http://localhost:8080/2015-03-31/functions/function/invocations
 Body (stringified within body for API request):
 
 ```JSON
-{"body":"{\"message\": \"hi\", \"params\": {\"conversation_id\": \"12345Test\", \"conversation_history\": [{\"type\": \"user\", \"content\": \"hi\"}]}}"}
+{"body":"{\"conversationId\": \"12345Test\", \"messages\": [{\"role\": \"USER\", \"content\": \"hi\"}], \"user\": {\"type\": \"LEARNER\"}}"}
 ```
 
-Body with optional Params:
-```JSON
+Body with optional fields:
+```json
 {
-    "message":"hi",
-    "params":{
-        "conversation_id":"12345Test",
-        "conversation_history":[{"type":"user","content":"hi"}],
-        "summary":" ",
-        "conversational_style":" ",
-        "question_response_details": "",
-        "include_test_data": true,
+  "conversationId": "<uuid>",
+  "messages": [
+    { "role": "USER", "content": "<previous user message>" },
+    { "role": "ASSISTANT", "content": "<previous assistant reply>" },
+    { "role": "USER", "content": "<current message>" }
+  ],
+  "user": {
+    "type": "LEARNER",
+    "preference": {
+      "conversationalStyle": "<stored style string>"
+    },
+    "taskProgress": {
+      "timeSpentOnQuestion": "30 minutes",
+      "accessStatus": "a good amount of time spent on this question today.",
+      "markedDone": "This question is still being worked on.",
+      "currentPart": {
+        "position": 0,
+        "timeSpentOnPart": "10 minutes",
+        "markedDone": "This part is not marked done.",
+        "responseAreas": [
+          {
+            "responseType": "EXPRESSION",
+            "totalSubmissions": 3,
+            "wrongSubmissions": 2,
+            "latestSubmission": {
+              "submission": "<student's last answer>",
+              "feedback": "<feedback text from evaluator>",
+              "answer": "<reference answer used for evaluation>"
+            }
+          }
+        ]
+      }
     }
+  },
+  "context": {
+    "summary": "<compressed conversation history>",
+    "set": {
+      "title": "Fundamentals",
+      "number": 2,
+      "description": "<set description>"
+    },
+    "question": {
+      "title": "Understanding Polymorphism",
+      "number": 3,
+      "guidance": "<teacher guidance>",
+      "content": "<master question content>",
+      "estimatedTime": "15-25 minutes",
+      "parts": [
+        {
+          "position": 0,
+          "content": "<part prompt>",
+          "answerContent": "<part answer>",
+          "workedSolutionSections": [
+            { "position": 0, "title": "Step 1", "content": "..." }
+          ],
+          "structuredTutorialSections": [
+            { "position": 0, "title": "Hint", "content": "..." }
+          ],
+          "responseAreas": [
+            {
+              "position": 0,
+              "responseType": "EXPRESSION",
+              "answer": "<reference answer>",
+              "preResponseText": "<label shown before input>"
+            }
+          ]
+        }
+      ]
+    }
+  }
+}
+```
+
+Response:
+
+```json
+{
+  "output": {
+    "role": "ASSISTANT",
+    "content": "<assistant reply text>"
+  },
+  "metadata": {
+    "summary": "<updated conversation summary>",
+    "conversationalStyle": "<updated style string>",
+    "processingTimeMs": 1234
+  }
 }
 ```
 
diff --git a/docs/dev.md b/docs/dev.md
index e0f56b1..ef0fec9 100644
--- a/docs/dev.md
+++ b/docs/dev.md
@@ -10,32 +10,6 @@ Helping approach: always responds with a follow-up question
 </code>
 </pre>
 
-## Inputs
-Body:
-```JSON
-{
-    "message":"hi",
-    "params":{
-        "conversation_id":"12345Test",
-        "conversation_history":[{"type":"user","content":"hi"}],
-        "include_test_data": true,
-    }
-}
-```
-
-## Outputs
-```JSON
-{
-    "chatbotResponse":"hi back",
-    "metadata": {
-        "summary": "",
-        "conversational_style": "",
-        "conversation_history": [],
-    },
-    "processing_time": 0
-}
-```
-
 ## Testing the Chat Function
 
 To test your function, you can run the unit tests, call the code directly through a python script, or build the respective chat function docker container locally and call it through an API request. Below you can find details on those processes.
@@ -50,12 +24,6 @@ pytest
 
 ### Run the Chat Script
 
-You can run the Python function itself. Make sure to have a main function in either `src/module.py` or `index.py`.
-
-```bash
-python src/module.py
-```
-
 You can also use the `manual_agent_run.py` script to test the agents with example inputs from Lambda Feedback questions and synthetic conversations.
 ```bash
 python tests/manual_agent_run.py
@@ -90,7 +58,7 @@ This will start the chat function and expose it on port `8080` and it will be op
 ```bash
 curl --location 'http://localhost:8080/2015-03-31/functions/function/invocations' \
 --header 'Content-Type: application/json' \
---data '{"body":"{\"message\": \"hi\", \"params\": {\"conversation_id\": \"12345Test\", \"conversation_history\": [{\"type\": \"user\", 
+--data '{"body":"{\"conversationId\": \"12345Test\", \"messages\": [{\"role\": \"USER\", \"content\": \"hi\"}], \"user\": {\"type\": \"LEARNER\"}}"}'
 ```
 
 #### Call Docker Container
@@ -106,24 +74,100 @@ POST URL:
 http://localhost:8080/2015-03-31/functions/function/invocations
 ```
 
-Body (stringified within body for API request):
+Input body (stringified within body for API request):
 
 ```JSON
-{"body":"{\"message\": \"hi\", \"params\": {\"conversation_id\": \"12345Test\", \"conversation_history\": [{\"type\": \"user\", \"content\": \"hi\"}]}}"}
+{"body":"{\"conversationId\": \"12345Test\", \"messages\": [{\"role\": \"USER\", \"content\": \"hi\"}], \"user\": {\"type\": \"LEARNER\"}}"}
 ```
 
-Body with optional Params:
-```JSON
+Input body with optional fields:
+```json
 {
-    "message":"hi",
-    "params":{
-        "conversation_id":"12345Test",
-        "conversation_history":[{"type":"user","content":"hi"}],
-        "summary":" ",
-        "conversational_style":" ",
-        "question_response_details": "",
-        "include_test_data": true,
+  "conversationId": "<uuid>",
+  "messages": [
+    { "role": "USER", "content": "<previous user message>" },
+    { "role": "ASSISTANT", "content": "<previous assistant reply>" },
+    { "role": "USER", "content": "<current message>" }
+  ],
+  "user": {
+    "type": "LEARNER",
+    "preference": {
+      "conversationalStyle": "<stored style string>"
+    },
+    "taskProgress": {
+      "timeSpentOnQuestion": "30 minutes",
+      "accessStatus": "a good amount of time spent on this question today.",
+      "markedDone": "This question is still being worked on.",
+      "currentPart": {
+        "position": 0,
+        "timeSpentOnPart": "10 minutes",
+        "markedDone": "This part is not marked done.",
+        "responseAreas": [
+          {
+            "responseType": "EXPRESSION",
+            "totalSubmissions": 3,
+            "wrongSubmissions": 2,
+            "latestSubmission": {
+              "submission": "<student's last answer>",
+              "feedback": "<feedback text from evaluator>",
+              "answer": "<reference answer used for evaluation>"
+            }
+          }
+        ]
+      }
     }
+  },
+  "context": {
+    "summary": "<compressed conversation history>",
+    "set": {
+      "title": "Fundamentals",
+      "number": 2,
+      "description": "<set description>"
+    },
+    "question": {
+      "title": "Understanding Polymorphism",
+      "number": 3,
+      "guidance": "<teacher guidance>",
+      "content": "<master question content>",
+      "estimatedTime": "15-25 minutes",
+      "parts": [
+        {
+          "position": 0,
+          "content": "<part prompt>",
+          "answerContent": "<part answer>",
+          "workedSolutionSections": [
+            { "position": 0, "title": "Step 1", "content": "..." }
+          ],
+          "structuredTutorialSections": [
+            { "position": 0, "title": "Hint", "content": "..." }
+          ],
+          "responseAreas": [
+            {
+              "position": 0,
+              "responseType": "EXPRESSION",
+              "answer": "<reference answer>",
+              "preResponseText": "<label shown before input>"
+            }
+          ]
+        }
+      ]
+    }
+  }
 }
 ```
 
+Output Response:
+
+```json
+{
+  "output": {
+    "role": "ASSISTANT",
+    "content": "<assistant reply text>"
+  },
+  "metadata": {
+    "summary": "<updated conversation summary>",
+    "conversationalStyle": "<updated style string>",
+    "processingTimeMs": 1234
+  }
+}
+```
diff --git a/index.py b/index.py
index 0ad738d..d0610fe 100644
--- a/index.py
+++ b/index.py
@@ -1,52 +1,36 @@
 import json
+from pydantic import ValidationError
+
+from lf_toolkit.chat import ChatRequest
 from src.module import chat_module
-from src.agent.utils.types import JsonType
 
-def handler(event: JsonType, context):
+
+def handler(event, context):
     """
     Lambda handler function
     """
-    # Log the input event for debugging purposes
-    # print("Received event:", " ".join(json.dumps(event, indent=2).splitlines()))
-
     if "body" in event:
         try:
             event = json.loads(event["body"])
         except json.JSONDecodeError:
             return {
                 "statusCode": 400,
-                "body": "Invalid JSON format in the body or body not found. Please check the input."
+                "body": "Invalid JSON format in the body. Please check the input.",
             }
 
-    if "message" not in event:
-        return {
-            "statusCode": 400,
-            "body": "Missing 'message' key in event. Please confirm the key in the json body."
-        }
-    if "params" not in event:
-        return {
-            "statusCode": 400,
-            "body": "Missing 'params' key in event. Please confirm the key in the json body. Make sure it contains the necessary conversation_id."
-        }
-    
-    message = event.get("message")
-    params = event.get("params")
+    try:
+        request = ChatRequest.model_validate(event)
+    except ValidationError as e:
+        return {"statusCode": 400, "body": e.json()}
 
     try:
-        chatbot_response = chat_module(message, params)
+        result = chat_module(request)
     except Exception as e:
         return {
             "statusCode": 500,
-            "body": f"An error occurred within the chat_module(): {str(e)}"
+            "body": f"An error occurred within the chat_module(): {str(e)}",
         }
 
-    # Create a response
-    response = {
-        "statusCode": 200,
-        "body": json.dumps(chatbot_response)
-    }
-
-    # Log the response for debugging purposes
+    response = {"statusCode": 200, "body": result.model_dump_json()}
     print("Returning response:", " ".join(json.dumps(response, indent=2).splitlines()))
-
-    return response
\ No newline at end of file
+    return response
diff --git a/src/agent/agent.py b/src/agent/agent.py
index 31e23cb..5a43018 100644
--- a/src/agent/agent.py
+++ b/src/agent/agent.py
@@ -1,7 +1,6 @@
-from src.agent.utils.llm_factory import OpenAILLMs, GoogleAILLMs
+from src.agent.llm_factory import OpenAILLMs, GoogleAILLMs
 from src.agent.prompts import \
     role_prompt, conv_pref_prompt, update_conv_pref_prompt, summary_prompt, update_summary_prompt, summary_system_prompt
-from src.agent.utils.types import InvokeAgentResponseType
 
 from langgraph.graph import StateGraph, START, END
 from langchain_core.messages import SystemMessage, RemoveMessage, HumanMessage, AIMessage
@@ -11,12 +10,12 @@
 from typing_extensions import TypedDict
 
 """
-Base agent for development [LLM workflow with a summarisation, profiling, and chat agent that receives an external conversation history].
+Base agent: LLM chat with automatic conversation summarisation and style analysis.
 
-This agent is designed to:
-- [summarise_prompt]        summarise the conversation after 'max_messages_to_summarize' number of messages is reached in the conversation
-- [conv_pref_prompt]        analyse the conversation style of the student 
-- [role_prompt]             role of a tutor to answer student's questions on the topic  
+Nodes:
+- call_llm:               responds to the student using role_prompt + optional question context + conversation summary
+- summarize_conversation: triggered after max_messages_to_summarize messages; compresses history and
+                          analyses the student's conversational style (returned in output metadata only)
 """
 
 ValidMessageTypes: TypeAlias = SystemMessage | HumanMessage | AIMessage
@@ -29,14 +28,11 @@ class State(TypedDict):
 
 class BaseAgent:
     def __init__(self):
-        llm = OpenAILLMs()                  # OpenAILLMs() or GoogleAILLMs()
-        self.llm = llm.get_llm()
-        summarisation_llm = OpenAILLMs()    # OpenAILLMs() or GoogleAILLMs()
-        self.summarisation_llm = summarisation_llm.get_llm()
-        self.summary = ""
-        self.conversationalStyle = ""
-
-        # Define Agent's specific Parameters
+        # Main chat LLM — change to GoogleAILLMs(), AzureLLMs(), or OllamaLLMs() if preferred
+        self.llm = OpenAILLMs().get_llm()
+        # Summarisation LLM — can be set to a different/cheaper model than the chat LLM
+        self.summarisation_llm = OpenAILLMs().get_llm()
+
         self.max_messages_to_summarize = 11
         self.role_prompt = role_prompt
         self.summary_prompt = summary_prompt
@@ -44,23 +40,22 @@ def __init__(self):
         self.conversation_preference_prompt = conv_pref_prompt
         self.update_conversation_preference_prompt = update_conv_pref_prompt
 
-        # Define a new graph for the conversation & compile it
-        self.workflow = StateGraph(State)
-        self.workflow_definition()
-        self.app = self.workflow.compile()
+        workflow = StateGraph(State)
+        workflow.add_node("call_llm", self.call_model)
+        workflow.add_node("summarize_conversation", self.summarize_conversation)
+        workflow.add_conditional_edges(source=START, path=self.should_summarize)
+        workflow.add_edge("summarize_conversation", "call_llm")
+        workflow.add_edge("call_llm", END)
+        self.app = workflow.compile()
 
-    def call_model(self, state: State, config: RunnableConfig) -> str:
-        """Call the LLM model knowing the role system prompt, the summary and the conversational style."""
-        
-        # Default AI tutor role prompt
+    def call_model(self, state: State, config: RunnableConfig) -> dict:
+        """Invoke the chat LLM with role prompt, optional question context, and conversation summary."""
         system_message = self.role_prompt
 
-        # Adding external student progress and question context details from data queries
-        question_response_details = config.get("configurable", {}).get("question_response_details", "")
-        if question_response_details:
-            system_message += f"## Known Question Materials: {question_response_details} \n\n"
+        context_prompt = config.get("configurable", {}).get("context_prompt", "")
+        if context_prompt:
+            system_message += f"## Known Question Materials: {context_prompt} \n\n"
 
-        # Adding summary and conversational style to the system message
         summary = state.get("summary", "")
         conversationalStyle = state.get("conversationalStyle", "")
         if summary:
@@ -68,130 +63,56 @@ def call_model(self, state: State, config: RunnableConfig) -> str:
         if conversationalStyle:
             system_message += f"## Known conversational style and preferences of the student for this conversation: {conversationalStyle}. \n\nYour answer must be in line with this conversational style."
 
-        messages = [SystemMessage(content=system_message)] + state['messages']
+        messages = [SystemMessage(content=system_message)] + state["messages"]
+        response = self.llm.invoke(self._valid(messages))
+        return {"messages": [response]}
 
-        valid_messages = self.check_for_valid_messages(messages)
-        response = self.llm.invoke(valid_messages)
+    def summarize_conversation(self, state: State) -> dict:
+        """Summarise history and analyse conversational style when message count exceeds the threshold."""
+        summary = state.get("summary", "")
+        conversationalStyle = state.get("conversationalStyle", "")
 
-        # Save summary for fetching outside the class
-        self.summary = summary
-        self.conversationalStyle = conversationalStyle
+        summary_message = (
+            f"This is summary of the conversation to date: {summary}\n\n" + self.update_summary_prompt
+        ) if summary else self.summary_prompt
 
-        return {"summary": summary, "messages": [response]}
-    
-    def check_for_valid_messages(self, messages: list[AllMessageTypes]) -> list[ValidMessageTypes]:
-        """ Removing the RemoveMessage() from the list of messages """
+        style_message = (
+            f"This is the previous conversational style of the student for this conversation: {conversationalStyle}\n\n" + self.update_conversation_preference_prompt
+        ) if conversationalStyle else self.conversation_preference_prompt
 
-        valid_messages: list[ValidMessageTypes] = []
-        for message in messages:
-            if message.type != 'remove':
-                valid_messages.append(message)
-        return valid_messages
-    
-    def summarize_conversation(self, state: State, config: RunnableConfig) -> dict:
-        """Summarize the conversation."""
+        summary_response = self.summarisation_llm.invoke(self._valid(state["messages"][:-1] + [HumanMessage(content=summary_message)]))
+        conversational_style_response = self.summarisation_llm.invoke(self._valid(state["messages"][:-1] + [HumanMessage(content=style_message)]))
 
-        summary = state.get("summary", "")
-        previous_summary = config.get("configurable", {}).get("summary", "")
-        previous_conversationalStyle = config.get("configurable", {}).get("conversational_style", "")
-        if previous_summary:
-            summary = previous_summary
-        
-        if summary:
-            summary_message = (
-                f"This is summary of the conversation to date: {summary}\n\n" +
-                self.update_summary_prompt
-            )
-        else:
-            summary_message = self.summary_prompt
-        
-        if previous_conversationalStyle:
-            conversationalStyle_message = (
-                f"This is the previous conversational style of the student for this conversation: {previous_conversationalStyle}\n\n" +
-                self.update_conversation_preference_prompt
-            )
-        else:
-            conversationalStyle_message = self.conversation_preference_prompt
-
-        # STEP 1: Summarize the conversation
-        messages = state["messages"][:-1] + [HumanMessage(content=summary_message)] 
-        valid_messages = self.check_for_valid_messages(messages)
-        summary_response = self.summarisation_llm.invoke(valid_messages)
-
-        # STEP 2: Analyze the conversational style
-        messages = state["messages"][:-1] + [HumanMessage(content=conversationalStyle_message)]
-        valid_messages = self.check_for_valid_messages(messages)
-        conversationalStyle_response = self.summarisation_llm.invoke(valid_messages)
-
-        # Delete messages that are no longer wanted, except the last ones
         delete_messages: list[AllMessageTypes] = [RemoveMessage(id=m.id) for m in state["messages"][:-3]]
+        return {"summary": summary_response.content, "conversationalStyle": conversational_style_response.content, "messages": delete_messages}
 
-        return {"summary": summary_response.content, "conversationalStyle": conversationalStyle_response.content, "messages": delete_messages}
-    
     def should_summarize(self, state: State) -> str:
-        """
-        Return the next node to execute. 
-        If there are more than X messages, then we summarize the conversation.
-        Otherwise, we call the LLM.
-        """
-
-        messages = state["messages"]
-        valid_messages = self.check_for_valid_messages(messages)
-        nr_messages = len(valid_messages)
-        if len(valid_messages) == 0:
-            raise Exception("Internal Error: No valid messages found in the conversation history. Conversation history might be empty.")
-        if "system" in valid_messages[-1].type:
-            nr_messages -= 1
-
-        # always pairs of (sent, response) + 1 latest message
-        if nr_messages > self.max_messages_to_summarize:
-            return "summarize_conversation"
-        return "call_llm"    
-
-    def workflow_definition(self) -> None:
-        self.workflow.add_node("call_llm", self.call_model)
-        self.workflow.add_node("summarize_conversation", self.summarize_conversation)
-
-        self.workflow.add_conditional_edges(source=START, path=self.should_summarize)
-        self.workflow.add_edge("summarize_conversation", "call_llm")
-        self.workflow.add_edge("call_llm", END)
-
-    def get_summary(self) -> str:
-        return self.summary
-    
-    def get_conversational_style(self) -> str:
-        return self.conversationalStyle
-
-    def print_update(self, update: dict) -> None:
-        for k, v in update.items():
-            for m in v["messages"]:
-                m.pretty_print()
-            if "summary" in v:
-                print(v["summary"])
-
-    def pretty_response_value(self, event: dict) -> str:
-        return event["messages"][-1].content
-    
-agent = BaseAgent()
-def invoke_base_agent(query: str, conversation_history: list, summary: str, conversationalStyle: str, question_response_details: str, session_id: str) -> InvokeAgentResponseType:
-    """
-    Call an agent that has no conversation memory and expects to receive all past messages in the params and the latest human request in the query.
-    If conversation history longer than X, the agent will summarize the conversation and will provide a conversational style analysis.
-    """
-    print(f'in invoke_base_agent(), thread_id = {session_id}')
+        valid = self._valid(state["messages"])
+        if not valid:
+            raise Exception("Internal Error: No valid messages found in the conversation history.")
+        nr_messages = len(valid) - (1 if "system" in valid[-1].type else 0)
+        return "summarize_conversation" if nr_messages > self.max_messages_to_summarize else "call_llm"
 
-    config = {"configurable": {"thread_id": session_id, "summary": summary, "conversational_style": conversationalStyle, "question_response_details": question_response_details}}
-    response_events = agent.app.invoke({"messages": conversation_history, "summary": summary, "conversational_style": conversationalStyle}, config=config, stream_mode="values") #updates
-    pretty_printed_response = agent.pretty_response_value(response_events) # get last event/ai answer in the response
+    def _valid(self, messages: list[AllMessageTypes]) -> list[ValidMessageTypes]:
+        return [m for m in messages if m.type != "remove"]
 
-    # Gather Metadata from the agent
-    summary = agent.get_summary()
-    conversationalStyle = agent.get_conversational_style()
 
-    print(f'in invoke_base_agent(), response generated by chatbot')
+agent = BaseAgent()
 
+def invoke_base_agent(messages: list, summary: str, conversationalStyle: str, context_prompt: str, session_id: str | None = None) -> dict:
+    """
+    Invoke the agent with the full conversation history and return the chatbot reply plus updated metadata.
+    Summarisation and style analysis trigger automatically once message count exceeds max_messages_to_summarize.
+    """
+    print(f"in invoke_base_agent(), thread_id = {session_id}")
+    config: RunnableConfig = {"configurable": {"context_prompt": context_prompt}}
+    state = agent.app.invoke(
+        State(messages=messages, summary=summary, conversationalStyle=conversationalStyle),
+        config=config,
+    )
+    print("in invoke_base_agent(), response generated by chatbot")
     return {
-        "input": query,
-        "output": pretty_printed_response,
-        "intermediate_steps": [str(summary), conversationalStyle, conversation_history]
-    }
\ No newline at end of file
+        "output": state["messages"][-1].content,
+        "summary": state.get("summary", ""),
+        "conversationalStyle": state.get("conversationalStyle", ""),
+    }
diff --git a/src/agent/context.py b/src/agent/context.py
new file mode 100644
index 0000000..a038747
--- /dev/null
+++ b/src/agent/context.py
@@ -0,0 +1,155 @@
+from typing import Optional, Dict, Any, List
+
+from src.agent.prompts import response_format_prompt
+
+
+def parse_json_to_prompt(context: dict, task_progress: dict) -> str:
+    """Convert muEd context and task progress directly into an LLM-friendly prompt string."""
+
+    question = context.get("question")
+    if not question:
+        return "# ERROR: Question details unavailable\n\nPlease describe the question you're working on so I can assist you effectively."
+
+    set_data = context.get("set", {})
+    current_part = task_progress.get("currentPart", {}) if task_progress else {}
+    current_part_position = current_part.get("position")
+    submissions = current_part.get("responseAreas", [])
+
+    sections = []
+
+    # 1. Question header
+    title_parts = []
+    set_number = set_data.get("number")
+    if set_number is not None and set_data.get("title"):
+        title_parts.append(f"## Set {set_number + 1}: {set_data['title']}")
+
+    question_num = ""
+    q_number = question.get("number")
+    if set_number is not None and q_number is not None:
+        question_num = f"{set_number + 1}.{q_number + 1}"
+    title_parts.append(f"### Question {question_num}: {question.get('title', '')}")
+
+    current_part_letter = _part_letter(current_part_position) if current_part_position is not None else ""
+    progress_indicator = f"Currently working on: Part ({current_part_letter})" if current_part_letter else ""
+
+    estimated_time = question.get("estimatedTime", "Not specified")
+    sections.append(f"""
+# Question Context
+
+{chr(10).join(title_parts)}
+
+{progress_indicator}
+
+### Question Details
+- Guidance: {question.get('guidance') or 'None provided'}
+- Description: {question.get('content') or 'No description available'}
+- Expected Duration: {estimated_time}
+
+---
+""")
+
+    # 2. Progress summary
+    if task_progress:
+        progress_items = []
+        if task_progress.get("timeSpentOnQuestion"):
+            progress_items.append(f"- Time spent today: {task_progress['timeSpentOnQuestion']}")
+        if task_progress.get("accessStatus"):
+            progress_items.append(f"- Status: {task_progress['accessStatus']}")
+        if task_progress.get("markedDone"):
+            progress_items.append(f"- Completion: {task_progress['markedDone']}")
+        if progress_items:
+            sections.append(f"\n## Progress Summary\n\n{chr(10).join(progress_items)}\n\n---\n")
+
+    # 3. Parts
+    for i, part in enumerate(question.get("parts", [])):
+        part_position = part.get("position", i)
+        is_current = current_part_position == part_position
+        time_on_part = current_part.get("timeSpentOnPart") if is_current else None
+        sections.append(_format_part(part, part_position, is_current, time_on_part, submissions))
+
+    # Combine
+    intro = (
+        "\n# Personalized Learning Assistant\n\n"
+        "I have detailed information about your current question, including your progress, responses, "
+        "and any feedback you've received. This context helps me provide targeted assistance based on "
+        "your specific situation.\n\n"
+    )
+    valid_sections = [s.strip() for s in sections if s and s.strip()]
+    response_format = (
+        "# Response Formatting\n" + response_format_prompt
+        if response_format_prompt
+        else ""
+    )
+    content = intro + "\n".join(valid_sections) + "\n" + response_format
+    content = content.replace("&#x20;&#x20;", " ").replace("&#x20", " ")
+    return "\n".join(line for line in content.split("\n") if line.strip() or not line).strip()
+
+
+def _part_letter(position: int) -> str:
+    return chr(96 + (position + 1))
+
+def _format_part(part: dict, part_position: int, is_current: bool, time_on_part: Optional[str], submissions: list) -> str:
+    letter = _part_letter(part_position)
+    status_text = " [CURRENTLY WORKING ON]" if is_current else ""
+    header = f"## Part ({letter}){status_text}"
+    if is_current and time_on_part and time_on_part != "No recorded duration":
+        header += f"\n\n*Time spent on this part: {time_on_part}*"
+
+    content_text = part.get("content", "").strip()
+    content = f"### Part Content\n\n{content_text}" if content_text else "### Part Content\n\nNo content provided"
+
+    response_areas = []
+    for j, ra in enumerate(part.get("responseAreas", [])):
+        ra_position = ra.get("position", j)
+        student_work = _get_student_work(ra_position, submissions)
+        response_areas.append(_format_response_area(ra_position, ra.get("preResponseText"), ra.get("answer"), student_work))
+    ra_block = f"\n### Response Areas\n\n{''.join(response_areas)}" if response_areas else ""
+
+    answer = part.get("answerContent")
+    answer_block = f"### Final Answer\n\n{answer}" if answer else "### Final Answer\n\nNo direct answer specified for this part"
+
+    solutions = [
+        f"{ws.get('title', f'#### Solution {i+1}')}\n\n{ws.get('content', '').strip() or 'No content available'}"
+        for i, ws in enumerate(part.get("workedSolutionSections", []))
+    ]
+    solutions_block = "### Worked Solutions\n\n" + "\n".join(solutions) if solutions else "### Worked Solutions\n\nNone available"
+
+    tutorials = [
+        f"{ts.get('title', f'#### Tutorial {i+1}')}\n\n{ts.get('content', '').strip() or 'No content available'}"
+        for i, ts in enumerate(part.get("structuredTutorialSections", []))
+    ]
+    tutorials_block = "### Structured Tutorials\n\n" + "\n".join(tutorials) if tutorials else "### Structured Tutorials\n\nNone available"
+
+    return "\n".join([header, content, ra_block, answer_block, solutions_block, tutorials_block]) + "\n---\n"
+
+def _get_student_work(ra_position: int, submissions: list) -> Dict[str, Any]:
+    if ra_position < len(submissions):
+        s = submissions[ra_position]
+        latest = s.get("latestSubmission") or {}
+        if latest:
+            return {
+                "has_submissions": True,
+                "latest_response": latest.get("submission"),
+                "latest_feedback": latest.get("feedback"),
+                "total_submissions": s.get("totalSubmissions"),
+                "total_wrong": s.get("wrongSubmissions"),
+            }
+    return {"has_submissions": False}
+
+def _format_response_area(position: int, task_description: Optional[str], expected_answer: Any, student_work: Dict[str, Any]) -> str:
+    task_text = f"- Task: {task_description}" if task_description else "- Task: Not specified"
+    if not student_work.get("has_submissions"):
+        submission_text = "- Your Work on this response area: No response submitted yet"
+    else:
+        submission_text = (
+            f"- Your Work on this response area:\n"
+            f"  - Latest response: {student_work.get('latest_response', 'None')}\n"
+            f"  - Latest feedback: {student_work.get('latest_feedback', 'None')}\n"
+            f"  - Total attempts: {student_work.get('total_submissions', 0)} out of which {student_work.get('total_wrong', 0)} were incorrect"
+        )
+    return (
+        f"\n#### Response Area {position + 1}\n\n"
+        f"{task_text}\n"
+        f"- Expected Answer (confidential): {expected_answer}\n"
+        f"{submission_text}\n"
+    )
diff --git a/src/agent/utils/llm_factory.py b/src/agent/llm_factory.py
similarity index 91%
rename from src/agent/utils/llm_factory.py
rename to src/agent/llm_factory.py
index 5b6b1d5..07e887b 100644
--- a/src/agent/utils/llm_factory.py
+++ b/src/agent/llm_factory.py
@@ -18,11 +18,11 @@ def __init__(self, temperature: int = 0):
                         temperature=temperature,
                         max_tokens=None,
                     )
-        self._azure_embedding = AzureOpenAIEmbeddings(azure_deployment=os.environ['AZURE_OPENAI_EMBEDDING_1536_DEPLOYMENT'], 
+        self._azure_embedding = AzureOpenAIEmbeddings(azure_deployment=os.environ['AZURE_OPENAI_EMBEDDING_1536_DEPLOYMENT'],
                                         openai_api_version=os.environ["AZURE_OPENAI_API_VERSION"],
                                         model=os.environ["AZURE_OPENAI_EMBEDDING_1536_MODEL"])
 
-    def get_llm(self):    
+    def get_llm(self):
         return self._azure_llm
 
     def get_embedding(self):
@@ -31,7 +31,7 @@ def get_embedding(self):
 class OllamaLLMs:
     def __init__(self):
         self._ollama_llm = Ollama(
-            model=os.environ['OLLAMA_MODEL'], # Any of the available models listed in the API docs
+            model=os.environ['OLLAMA_MODEL'],
             base_url=os.environ['OLLAMA_BASE_URL'],
             headers={
                 'X-API-Key': os.environ['OLLAMA_API_KEY'],
@@ -52,7 +52,7 @@ def get_llm(self):
 
     def get_embedding(self):
         return self._ollama_embedding
-    
+
 class OpenAILLMs:
     def __init__(self, temperature: int = 0):
         self._openai_llm = ChatOpenAI(
@@ -65,21 +65,20 @@ def __init__(self, temperature: int = 0):
             model='text-embedding-ada-002',
             api_key=os.environ['OPENAI_API_KEY'],
         )
-    
+
     def get_llm(self):
         return self._openai_llm
-    
+
     def get_embedding(self):
         return self._openai_embedding
 
 class GoogleAILLMs:
     def __init__(self, temperature: int = 0):
-
         self._google_llm = ChatGoogleGenerativeAI(
             model=os.environ['GOOGLE_AI_MODEL'],
             temperature=temperature,
             google_api_key=os.environ['GOOGLE_AI_API_KEY'],
         )
-    
+
     def get_llm(self):
-        return self._google_llm
\ No newline at end of file
+        return self._google_llm
diff --git a/src/agent/prompts.py b/src/agent/prompts.py
index b8e6808..f7e6fcf 100644
--- a/src/agent/prompts.py
+++ b/src/agent/prompts.py
@@ -4,6 +4,7 @@
 # Description of the prompts:
 # 
 # 1. role_prompt: Sets the overall role and behaviour of the chatbot.
+# 1b. response_format_prompt: Provides specific instructions on how the chatbot should format its responses, for example regarding mathematical equations formatting and language.
 # 
 # 2. summary_prompt: Used to generate a summary of the conversation.
 # 2. update_summary_prompt: Used to update the conversation summary with new messages.
@@ -44,6 +45,12 @@
 
 Remember: Your role is to be the question-asker, not the answer-giver. Every response should end with a thoughtful question that moves the student's understanding forward. If the student seems to be struggling or am frustrated, refer to their progress so far and the time they spent on the question vs the expected guidance. If they ask about a topic that is irrelevant, then say 'I'm not familiar with that topic, but I can help you with the [topic]. You do not need to end your messages with a concluding statement.
 """
+
+# 1b. Response Format Prompt
+response_format_prompt = """Mathematical equations are in KaTeX format, preserve them the same. Ensure mathematical equations are surrounded by one '$' for in-line equations and '$$' for block equations.
+Example: '$E=mc^2$' or '$$E=mc^2$$'.
+Use British English spellings."""
+
 # 2. Summary Prompts
 summary_guidelines = """Ensure the summary is:
 
diff --git a/src/agent/utils/example_inputs/example_input_1.json b/src/agent/utils/example_inputs/example_input_1.json
deleted file mode 100644
index 5bd8071..0000000
--- a/src/agent/utils/example_inputs/example_input_1.json
+++ /dev/null
@@ -1,168 +0,0 @@
-{
-  "message": "i dont remember anything",
-  "params": {
-    "include_test_data": true,
-    "conversation_history": [
-      { "type": "user", "content": "what should I do?" },
-      {
-        "type": "assistant",
-        "content": "It seems like you're currently working on Part (a) of the dot product question. Since you haven't submitted an answer yet, let's take a moment to break it down together. \n\nWhat do you remember about how to calculate the dot product of two vectors? Can you describe the steps you would take?"
-      },
-      { "type": "user", "content": "i dont remember anything" }
-    ],
-    "summary": "",
-    "conversational_style": "",
-    "question_response_details": {
-      "questionSubmissionSummary": [
-        {
-          "publishedPartId": "0e0432e4-90a2-47a1-b597-55f76596b7d5",
-          "publishedPartPosition": 0,
-          "publishedResponseAreaId": "ce28b12b-f583-4d83-a0e4-36b17127746a",
-          "publishedResponseAreaPosition": 0,
-          "responseAreaUniversalId": "f8857876-482d-411c-8ffc-734fb07712cf",
-          "publishedResponseAreaPreResponseText": "$\\vec{a} \\ \\cdot \\ \\vec{b} \\ =$",
-          "publishedResponseType": "NUMBER",
-          "publishedResponseConfig": null,
-          "totalSubmissions": 0,
-          "totalWrongSubmissions": 0
-        },
-        {
-          "publishedPartId": "a74a6fef-8c94-474c-b381-8d97a4b54725",
-          "publishedPartPosition": 1,
-          "publishedResponseAreaId": "f3be55ed-af9b-4483-ba45-636ccbad7a24",
-          "publishedResponseAreaPosition": 0,
-          "responseAreaUniversalId": "054bbcce-facb-4aaa-a1d2-7bec7627a6ef",
-          "publishedResponseAreaPreResponseText": "$\\left(\\vec{a}  - \\vec{b}\\right)\\cdot \\vec{c}\\  =$",
-          "publishedResponseType": "NUMBER",
-          "publishedResponseConfig": null,
-          "totalSubmissions": 0,
-          "totalWrongSubmissions": 0
-        },
-        {
-          "publishedPartId": "66160c45-0554-470a-88ae-82f36e755ef6",
-          "publishedPartPosition": 2,
-          "publishedResponseAreaId": "79260525-1edb-4f19-bcdd-d827b5b692f3",
-          "publishedResponseAreaPosition": 0,
-          "responseAreaUniversalId": "bba9308c-55c7-4733-9315-37bc26fc8ec1",
-          "publishedResponseAreaPreResponseText": "$\\left( \\ \\vec{a}  \\cdot \\vec{c} \\ \\right) \\vec{b}\\  =$",
-          "publishedResponseType": "MATRIX",
-          "publishedResponseConfig": { "cols": 1, "rows": 3 },
-          "totalSubmissions": 0,
-          "totalWrongSubmissions": 0
-        }
-      ],
-      "questionInformation": {
-        "questionTitle": "Dot Product",
-        "questionGuidance": "",
-        "questionContent": "$$\n\\vec{a}=\\begin{bmatrix}1 \\\\ 3\\\\ -2\\end{bmatrix} \\quad \\vec{b}=\\begin{bmatrix}0 \\\\ 3\\\\ 1\\end{bmatrix} \\quad \\vec{c}=\\begin{bmatrix} 1 \\\\\\ -1\\\\ -3\\end{bmatrix}\n$$",
-        "durationLowerBound": 1,
-        "durationUpperBound": 4,
-        "parts": [
-          {
-            "publishedPartId": "0e0432e4-90a2-47a1-b597-55f76596b7d5",
-            "publishedPartPosition": 0,
-            "publishedPartContent": "",
-            "publishedPartAnswerContent": "",
-            "publishedWorkedSolutionSections": [
-              {
-                "id": "a7b1150a-05b7-4337-b902-5c6b1835cc74",
-                "position": 0,
-                "title": "",
-                "content": "Assuming the given basis is orthonormal, the dot product between two vectors, $\\vec{a}$ and $\\vec{b}$, can be calculated simply by multiplying their corresponding components and summing them up:\n\n&#x20;      &#x20;\n\n$$\n\\begin{array}{rl}\\vec{a} \\cdot \\vec{b} &=\\begin{bmatrix}1 \\\\ 3\\\\ -2\\end{bmatrix}•\\begin{bmatrix}0 \\\\ 3\\\\ 1\\end{bmatrix} \\\\\\\\ &= (1 \\cdot 0) + (3 \\cdot 3) + (-2 \\cdot 1) \\\\\\\\ &= 0 + 9 + (-2) \\\\\\\\ &= 7\\end{array}\n$$\n\n&#x20;        &#x20;\n\nTherefore, the dot product between vectors $\\vec{a}$ and $\\vec{b}$ is \\$7\\$"
-              }
-            ],
-            "publishedResponseAreas": [
-              {
-                "id": "ce28b12b-f583-4d83-a0e4-36b17127746a",
-                "position": 0,
-                "universalResponseAreaId": "f8857876-482d-411c-8ffc-734fb07712cf",
-                "preResponseText": "$\\vec{a} \\ \\cdot \\ \\vec{b} \\ =$",
-                "Response": {
-                  "id": "4d0dfa85-fc34-4649-bb6e-0f194e8a5a04",
-                  "responseType": "NUMBER",
-                  "config": null,
-                  "answer": 7
-                },
-                "responseType": "NUMBER",
-                "answer": 7
-              }
-            ]
-          },
-          {
-            "publishedPartId": "a74a6fef-8c94-474c-b381-8d97a4b54725",
-            "publishedPartPosition": 1,
-            "publishedPartContent": "",
-            "publishedPartAnswerContent": "",
-            "publishedWorkedSolutionSections": [
-              {
-                "id": "0e91a91d-c536-4120-a70f-5cfea57b4235",
-                "position": 0,
-                "title": "",
-                "content": "To calculate $\\left(\\vec{a} - \\vec{b}\\right) \\cdot \\vec{c}$ , we first need to find the vector resulting from the subtraction of $\\vec{b}$ from $\\vec{a}$:\n\n&#x20;    &#x20;\n\n$$\n\\begin{array}{rl}\\vec{a} - \\vec{b} &= \\begin{bmatrix}1 \\\\ 3 \\\\ -2\\end{bmatrix} - \\begin{bmatrix}0 \\\\ 3 \\\\ 1\\end{bmatrix} \\\\\\\\&= \\begin{bmatrix}1 - 0 \\\\ 3 - 3 \\\\ -2 - 1\\end{bmatrix} \\\\\\\\&= \\begin{bmatrix}1 \\\\ 0 \\\\ -3\\end{bmatrix}\\end{array}\n$$\n\n\n\nNext, we can compute the dot product of $\\left(\\vec{a} - \\vec{b}\\right)$ and $\\vec{c}$\n\n&#x20;     &#x20;\n\n$$\n\\begin{array}{rl}\\left(\\vec{a} - \\vec{b}\\right) \\cdot \\vec{c} &= \\begin{bmatrix}1 \\\\ 0 \\\\ -3\\end{bmatrix} \\cdot \\begin{bmatrix}1 \\\\ -1 \\\\ -3\\end{bmatrix} \\\\\\\\ &= (1 \\cdot 1) + (0 \\cdot -1) + (-3 \\cdot -3) \\\\\\\\&= 1 + 0 + 9 \\\\\\\\&= 10\\end{array}\n$$\n\n\n\nTherefore, $\\left(\\vec{a} - \\vec{b}\\right) \\cdot \\vec{c}$ is equal to \\$10"
-              }
-            ],
-            "publishedResponseAreas": [
-              {
-                "id": "f3be55ed-af9b-4483-ba45-636ccbad7a24",
-                "position": 0,
-                "universalResponseAreaId": "054bbcce-facb-4aaa-a1d2-7bec7627a6ef",
-                "preResponseText": "$\\left(\\vec{a}  - \\vec{b}\\right)\\cdot \\vec{c}\\  =$",
-                "Response": {
-                  "id": "01edd054-b1fd-4049-8551-a8551d7a63ea",
-                  "responseType": "NUMBER",
-                  "config": null,
-                  "answer": 10
-                },
-                "responseType": "NUMBER",
-                "answer": 10
-              }
-            ]
-          },
-          {
-            "publishedPartId": "66160c45-0554-470a-88ae-82f36e755ef6",
-            "publishedPartPosition": 2,
-            "publishedPartContent": "",
-            "publishedPartAnswerContent": "",
-            "publishedWorkedSolutionSections": [
-              {
-                "id": "c3225276-7bf8-4962-a459-d9583442c26e",
-                "position": 0,
-                "title": "",
-                "content": "To calculate $\\left( \\ \\vec{a} \\cdot \\vec{c} \\ \\right) \\vec{b}$ , we first need to find the dot product of vectors $\\vec{a}$ and $\\vec{c}$\n\n&#x20;       &#x20;\n\n$$\n\\begin{array}{rl}\\vec{a} \\cdot \\vec{c} &= \\begin{bmatrix}1 \\\\ 3 \\\\ -2\\end{bmatrix} \\cdot \\begin{bmatrix}1 \\\\ -1 \\\\ -3\\end{bmatrix} \\\\\\\\ &= (1 \\cdot 1) + (3 \\cdot -1) + (-2 \\cdot -3) \\\\\\\\&= 1 - 3 + 6 \\\\\\\\&= 4\\end{array}\n$$\n\n&#x20;          &#x20;\n\nNext, we can scale $\\vec{b}$ by $\\vec{a} \\cdot \\vec{c}$\n\n&#x20;        &#x20;\n\n$$\n\\begin{array}{rl}\\left( \\ \\vec{a} \\cdot \\vec{c} \\ \\right) \\vec{b} &= 4\\begin{bmatrix}0 \\\\ 3 \\\\ 1\\end{bmatrix}  \\\\\\\\ &= \\begin{bmatrix}4\\cdot0 \\\\ 4\\cdot3 \\\\ 4\\cdot1\\end{bmatrix}  \\\\\\\\&= \\begin{bmatrix}0 \\\\ 12 \\\\ 4\\end{bmatrix} \\end{array}\n$$\n\n&#x20;        &#x20;\n"
-              }
-            ],
-            "publishedResponseAreas": [
-              {
-                "id": "79260525-1edb-4f19-bcdd-d827b5b692f3",
-                "position": 0,
-                "universalResponseAreaId": "bba9308c-55c7-4733-9315-37bc26fc8ec1",
-                "preResponseText": "$\\left( \\ \\vec{a}  \\cdot \\vec{c} \\ \\right) \\vec{b}\\  =$",
-                "Response": {
-                  "id": "cc759374-175f-40be-8481-2c81e34795b6",
-                  "responseType": "MATRIX",
-                  "config": { "cols": 1, "rows": 3 },
-                  "answer": [["0"], ["12"], ["4"]]
-                },
-                "responseType": "MATRIX",
-                "answer": [["0"], ["12"], ["4"]]
-              }
-            ]
-          }
-        ]
-      },
-      "questionAccessInformation": {
-        "estimatedMinimumTime": "1 minute",
-        "estimaredMaximumTime": "4 minutes",
-        "timeTaken": "20 minutes",
-        "accessStatus": "too much time spent on this question.",
-        "markedDone": "",
-        "currentPart": {
-          "id": "0e0432e4-90a2-47a1-b597-55f76596b7d5",
-          "position": 0
-        }
-      }
-    },
-    "conversation_id": "7a65b6ed-85d1-4621-8efb-4fc8e9c5a8de",
-    "agent_type": "base"
-  }
-}
diff --git a/src/agent/utils/example_inputs/example_input_2.json b/src/agent/utils/example_inputs/example_input_2.json
deleted file mode 100644
index 503ece0..0000000
--- a/src/agent/utils/example_inputs/example_input_2.json
+++ /dev/null
@@ -1,143 +0,0 @@
-{
-  "message": "what is the function value when x is between -1 0",
-  "params": {
-    "include_test_data": true,
-    "conversation_history": [
-      {
-        "type": "user",
-        "content": "what is the function value when x is between -1 0"
-      }
-    ],
-    "summary": "",
-    "conversational_style": "",
-    "question_response_details": {
-      "questionSubmissionSummary": [
-        {
-          "publishedPartId": "08e6f713-def7-4f97-83d4-34b0e67f5222",
-          "publishedPartPosition": 0,
-          "publishedResponseAreaId": "b8fc25d6-7cf4-4afc-984a-fa75fb0af7e7",
-          "publishedResponseAreaPosition": 0,
-          "responseAreaUniversalId": "e849584a-f58c-4330-858b-b507b9d8c56c",
-          "publishedResponseAreaPreResponseText": "$a_0=$",
-          "publishedResponseType": "EXPRESSION",
-          "publishedResponseConfig": {
-            "allowPhoto": true,
-            "allowHandwrite": true
-          },
-          "totalSubmissions": 0,
-          "totalWrongSubmissions": 0
-        },
-        {
-          "publishedPartId": "08e6f713-def7-4f97-83d4-34b0e67f5222",
-          "publishedPartPosition": 0,
-          "publishedResponseAreaId": "2afd4463-ffbd-4b6b-8244-19c9c27e00f1",
-          "publishedResponseAreaPosition": 1,
-          "responseAreaUniversalId": "f20abf75-15ee-46ee-bee5-2897c71b2493",
-          "publishedResponseAreaPreResponseText": "$a_n=$",
-          "publishedResponseType": "EXPRESSION",
-          "publishedResponseConfig": {
-            "allowPhoto": true,
-            "allowHandwrite": true
-          },
-          "totalSubmissions": 0,
-          "totalWrongSubmissions": 0
-        },
-        {
-          "publishedPartId": "08e6f713-def7-4f97-83d4-34b0e67f5222",
-          "publishedPartPosition": 0,
-          "publishedResponseAreaId": "1477280a-64c7-43f3-a31d-091b95fd0900",
-          "publishedResponseAreaPosition": 2,
-          "responseAreaUniversalId": "f3cb33ee-419b-4b60-ab2a-2b7425710b5a",
-          "publishedResponseAreaPreResponseText": "$b_n=$",
-          "publishedResponseType": "EXPRESSION",
-          "publishedResponseConfig": {
-            "allowPhoto": true,
-            "allowHandwrite": true
-          },
-          "totalSubmissions": 0,
-          "totalWrongSubmissions": 0
-        }
-      ],
-      "questionInformation": {
-        "questionTitle": "Piecewise function Fourier series",
-        "questionGuidance": "",
-        "questionContent": "Find $a_0$, $a_n$ and $b_n$ for the Fourier series of $f(x)$, which is assumed to have period $4$.\n\n&#x20;&#x20;\n\n$$\nf(x)= \\begin{cases}0, & -2 \\leq x<-1 \\\\\\ \\frac{2 k}{3}, & -1 \\leq x<1 \\\\\\ -\\frac{k}{2}, & 1 \\leq x<2.\\end{cases}\n$$\n",
-        "durationLowerBound": 2,
-        "durationUpperBound": 10,
-        "parts": [
-          {
-            "publishedPartId": "08e6f713-def7-4f97-83d4-34b0e67f5222",
-            "publishedPartPosition": 0,
-            "publishedPartContent": "",
-            "publishedPartAnswerContent": "$$\na_0=\\frac{5k}{12}\n$$\n\n&#x20;&#x20;\n\n&#x20;&#x20;\n\n$$\n\\begin{align*}\na_n &= \\dfrac{11k}{6n \\pi} \\sin \\left( \\dfrac{n \\pi}{2} \\right) \\\\[1em]\n &= \\dfrac{11k}{6n \\pi} \\frac{1-(-1)^n}{2}  (-1)^{^{\\frac{n+3}{2}}} \\\\[1em]\n&= (-1)^{n+1}\\frac{11k}{6(2n-1)\\pi}\n\\end{align*}\n$$\n\n&#x20;&#x20;\n\n&#x20;&#x20;\n\n$$\n\\begin{align*}\nb_n &= \\frac{k}{2n \\pi} \\left( \\cos \\left( \\frac{n \\pi}{2} \\right) - \\cos(n \\pi) \\right) \\\\[1em]\nb_n &= \\frac{k}{2n \\pi} \\left( \\frac{1+(-1)^n}{2}(-1)^{\\frac{n}{2}} - (-1)^n \\right)\n\\end{align*}\n\n$$\n",
-            "publishedWorkedSolutionSections": [
-              {
-                "id": "d41e73da-331d-4897-b6f6-15031469502c",
-                "position": 0,
-                "title": "",
-                "content": "Recall the Fourier series equations for period $2L$:\n\n&#x20;&#x20;\n\n$$\na_0 = \\frac{1}{L} \\int_{-L}^L {f(x)} \\, \\text{d}x\n\n$$\n\n&#x20;&#x20;\n\n$$\na_n = \\frac{1}{L} \\int_{-L}^L {f(x)} \\cos\\left(\\frac{n \\pi x}{L}\\right) \\, \\text{d}x\n$$\n\n&#x20;&#x20;\n\n$$\nb_n = \\frac{1}{L} \\int_{-L}^L {f(x)} \\sin\\left(\\frac{n \\pi x}{L}\\right) \\, \\text{d}x\n$$\n\n***\n\n$$\n2L=4\n$$\n\n$$\nL=2\n$$\n\n***\n\n### **Finding $a_0$:**\n\n$$\na_0 = \\frac{1}{2} \\int_{-2}^2 {f(x)} \\, \\text{d}x\n$$\n\n&#x20;&#x20;\n\n$$\na_0 = \\frac{1}{2} \\left( \\int_{-2}^{-1} 0 \\, \\text{d}x + \\int_{-1}^1 \\frac{2k}{3} \\, \\text{d}x + \\int_1^2 -\\frac{k}{2} \\, \\text{d}x \\right)\n\n$$\n\n&#x20;&#x20;\n\n$$\na_0=\\frac{5k}{12}\n$$\n\n***\n\n### **Finding $a_n$:**\n\n$$\na_n = \\frac{1}{2} \\int_{-2}^2 {f(x)} \\cos \\left( \\frac{n \\pi x}{L} \\right) \\, \\text{d}x\n\n$$\n\n&#x20;&#x20;\n\n$$\na_n = \\frac{1}{2} \\left( \\int_{-2}^{-1} 0 \\, \\text{d}x + \\int_{-1}^1 \\frac{2k}{3} \\cos \\left( \\frac{n \\pi x}{2} \\right) \\, \\text{d}x + \\int_1^2 -\\frac{k}{2} \\cos \\left( \\frac{n \\pi x}{2} \\right) \\, \\text{d}x \\right)\n\n$$\n\n&#x20;&#x20;\n\nAfter evaluating the integrals, the following is obtained:\n\n&#x20;&#x20;\n\n$$\na_n=\\frac{k}{3}\\left(\\frac{2}{n \\pi} \\sin \\left(\\frac{n \\pi}{2}\\right)-\\frac{2}{n \\pi} \\sin \\left(-\\frac{n \\pi}{2}\\right)\\right)-\\frac{k}{4}\\left(\\frac{2}{n \\pi} \\sin (n \\pi)-\\frac{2}{n \\pi} \\sin \\left(\\frac{n \\pi}{2}\\right)\\right)\n$$\n\n&#x20;&#x20;\n\n*   The second $\\sin$ term can be written as: $-\\frac{2}{n \\pi} \\sin \\left(-\\frac{n \\pi}{2}\\right)=\\frac{2}{n \\pi} \\sin \\left(\\frac{n \\pi}{2}\\right)$\n*   The third $\\sin$ term is always zero.\n\n&#x20;&#x20;\n\n$$\na_n=\\frac{k}{3}\\left(\\frac{4}{n \\pi} \\sin \\left(\\frac{n \\pi}{2}\\right)\\right)-\\frac{k}{4}\\left(-\\frac{2}{n \\pi} \\sin \\left(\\frac{n \\pi}{2}\\right)\\right)\n$$\n\n&#x20;&#x20;\n\nSome manipulation results in:\n\n&#x20;&#x20;\n\n$$\na_n = \\dfrac{11k}{6n \\pi} \\sin \\left( \\dfrac{n \\pi}{2} \\right)\n$$\n\n**Further simplification of $a_n$:**\n\nThe $\\sin \\left( \\dfrac{n \\pi}{2} \\right)$ term can be simplified by considering the pattern with increasing $n$:\n\n&#x20;&#x20;\n\n$$\n\\sin\\left(\\frac{n\\pi}{2}\\right) = \\begin{cases}\n0 & \\text{if $n=0$} \\\\\n1 & \\text{if $n=1$} \\\\\n0 & \\text{if $n=2$} \\\\\n-1 & \\text{if $n=3$} \\\\\n0 & \\text{if $n=4$} \\\\\n\\vdots & \\vdots\n\\end{cases}\n$$\n\nThis can be achieved as follows:\n\n$$\n\\sin \\left( \\dfrac{n \\pi}{2} \\right)=\\frac{1-(-1)^n}{2} (-1)^{^{\\frac{n+3}{2}}}\n$$\n\n&#x20;&#x20;\n\n$$\na_n = \\dfrac{11k}{6n \\pi} \\frac{1-(-1)^n}{2}  (-1)^{^{\\frac{n+3}{2}}}\n$$\n\n$$\n\n\n$$\n\n&#x20;&#x20;\n\nHowever, since $\\sin \\left( \\dfrac{n \\pi}{2} \\right)$ is zero for all even $n$, the expression can alternatively be written as:\n\n&#x20;&#x20;\n\n$$\na_n=(-1)^{n+1}\\frac{11k}{6(2n-1)\\pi} \n$$\n\n***\n\n### **Finding $b_n$:**\n\n$$\nb_n = \\frac{1}{2} \\int_{-2}^2 {f(x)} \\sin\\left(\\frac{n \\pi x}{L}\\right) \\, \\text{d}x\n$$\n\n&#x20;&#x20;\n\n$$\nb_n = \\frac{1}{2} \\left( \\int_{-2}^{-1} 0 \\, \\text{d}x + \\int_{-1}^1 \\frac{2k}{3} \\sin \\left( \\frac{n \\pi x}{2} \\right) \\, \\text{d}x + \\int_1^2 -\\frac{k}{2} \\sin \\left( \\frac{n \\pi x}{2} \\right) \\, \\text{d}x \\right)\n$$\n\n&#x20;&#x20;\n\nAfter evaluating the integrals, the following is obtained:\n\n&#x20;  &#x20;\n\n$$\nb_n=\\frac{k}{3}\\left(-\\frac{2}{n \\pi} \\cos \\left(\\frac{n \\pi}{2}\\right)+\\frac{2}{n \\pi} \\cos \\left(\\frac{n \\pi}{2}\\right)\\right)-\\frac{k}{4}\\left(-\\frac{2}{n \\pi} \\cos (n \\pi)+\\frac{2}{n \\pi} \\cos \\left(\\frac{n \\pi}{2}\\right)\\right)\n$$\n\n(note that the second $\\cos$ term has positive argument, because $\\cos$ is an even function.)\n\n&#x20;&#x20;\n\nSimplifying this expression yields the answer. Note that $\\cos(n\\pi)$ has been replaced with $(-1)^n$.\n\n&#x20;&#x20;\n\n$$\nb_n = \\frac{k}{2n \\pi} \\left( \\cos \\left( \\frac{n \\pi}{2} \\right) - \\cos(n \\pi) \\right)\n$$\n\n&#x20;&#x20;\n\n**Further simplification of** $b_n$**:**\n\n*   The $\\cos(n\\pi)$ term can be replaced by $(-1)^n$.\n*   The $\\cos\\left(\\frac{n\\pi}{2}\\right)$ term can be simplified by considering the pattern with increasing $n$:\n\n    &#x20;&#x20;\n\n$$\n\\cos\\left(\\frac{n\\pi}{2}\\right) = \\begin{cases}\n1 & \\text{if $n=0$} \\\\\n0 & \\text{if $n=1$} \\\\\n-1 & \\text{if $n=2$} \\\\\n0 & \\text{if $n=3$} \\\\\n1 & \\text{if $n=4$} \\\\\n\\vdots & \\vdots\n\\end{cases}\n$$\n\nThis can be achieved as follows:\n\n$$\n\\cos\\left(\\frac{n\\pi}{2}\\right) =\\frac{1+(-1)^n}{2}(-1)^{\\frac{n}{2}}\n$$\n\nFinally, this yields:\n\n&#x20;&#x20;\n\n$$\nb_n = \\frac{k}{2n \\pi} \\left( \\frac{1+(-1)^n}{2}(-1)^{\\frac{n}{2}} - (-1)^n \\right)\n$$\n\n&#x20;&#x20;\n\n(Note that this expression may seem less concise that simply including the cos(n\\*pi/2) but is much more desirable and efficient in a numerical algorithm).\n"
-              }
-            ],
-            "publishedResponseAreas": [
-              {
-                "id": "b8fc25d6-7cf4-4afc-984a-fa75fb0af7e7",
-                "position": 0,
-                "universalResponseAreaId": "e849584a-f58c-4330-858b-b507b9d8c56c",
-                "preResponseText": "$a_0=$",
-                "Response": {
-                  "id": "d5d9f641-0724-4fa9-b3a7-2139c551c0a5",
-                  "responseType": "EXPRESSION",
-                  "config": { "allowPhoto": true, "allowHandwrite": true },
-                  "answer": "5k/12"
-                },
-                "responseType": "EXPRESSION",
-                "answer": "5k/12"
-              },
-              {
-                "id": "2afd4463-ffbd-4b6b-8244-19c9c27e00f1",
-                "position": 1,
-                "universalResponseAreaId": "f20abf75-15ee-46ee-bee5-2897c71b2493",
-                "preResponseText": "$a_n=$",
-                "Response": {
-                  "id": "c7cb9bae-d67e-42db-9d71-3b4d2c3c8b44",
-                  "responseType": "EXPRESSION",
-                  "config": { "allowPhoto": true, "allowHandwrite": true },
-                  "answer": "(-1)^(n+1) 11k/(6(2n-1)pi)"
-                },
-                "responseType": "EXPRESSION",
-                "answer": "(-1)^(n+1) 11k/(6(2n-1)pi)"
-              },
-              {
-                "id": "1477280a-64c7-43f3-a31d-091b95fd0900",
-                "position": 2,
-                "universalResponseAreaId": "f3cb33ee-419b-4b60-ab2a-2b7425710b5a",
-                "preResponseText": "$b_n=$",
-                "Response": {
-                  "id": "9b868669-feb4-4c33-977e-2297bf4f7948",
-                  "responseType": "EXPRESSION",
-                  "config": { "allowPhoto": true, "allowHandwrite": true },
-                  "answer": "k/(2n pi) ( (1+(-1)^n)/(2) (-1)^(n/2) - (-1)^n)"
-                },
-                "responseType": "EXPRESSION",
-                "answer": "k/(2n pi) ( (1+(-1)^n)/(2) (-1)^(n/2) - (-1)^n)"
-              }
-            ]
-          }
-        ]
-      },
-      "questionAccessInformation": {
-        "estimatedMinimumTime": "2 minutes",
-        "estimaredMaximumTime": "10 minutes",
-        "timeTaken": "less than one minute",
-        "accessStatus": "too little time spent on this question.",
-        "markedDone": "",
-        "currentPart": {
-          "id": "08e6f713-def7-4f97-83d4-34b0e67f5222",
-          "position": 0
-        }
-      }
-    },
-    "conversation_id": "6779b184-41b5-4384-ade9-1c06d2cd91a5",
-    "agent_type": "base"
-  }
-}
diff --git a/src/agent/utils/example_inputs/example_input_3.json b/src/agent/utils/example_inputs/example_input_3.json
deleted file mode 100644
index d93ac92..0000000
--- a/src/agent/utils/example_inputs/example_input_3.json
+++ /dev/null
@@ -1,460 +0,0 @@
-{
-  "message": "hi",
-  "params": {
-    "include_test_data": true,
-    "conversation_history": [
-      { "type": "user", "content": "hi" },
-      { "type": "ai", "content": "Hello! How can I help you today?" }
-    ],
-    "summary": "",
-    "conversational_style": "",
-    "question_response_details": {
-      "questionSubmissionSummary": [
-        {
-          "publishedPartId": "04d2cab2-eeca-4c1f-bcef-6401fe4ec635",
-          "publishedPartPosition": 0,
-          "publishedResponseAreaId": "1194c5b3-1831-433e-b78e-c96c94f99117",
-          "publishedResponseAreaPosition": 0,
-          "responseAreaUniversalId": "ee1df8dd-bccb-46d2-b816-1a707115495e",
-          "publishedResponseAreaPreResponseText": "$\\vec{a}+\\vec{b}=$",
-          "publishedResponseType": "MATRIX",
-          "publishedResponseConfig": { "cols": 1, "rows": 3 },
-          "totalSubmissions": 1,
-          "totalWrongSubmissions": 0,
-          "latestSubmission": {
-            "universalResponseAreaId": "ee1df8dd-bccb-46d2-b816-1a707115495e",
-            "answer": "[[\"-1\"],[\"5\"],[\"-1\"]]",
-            "submission": [["-1"], ["5"], ["-1"]],
-            "feedback": "Correct",
-            "rawResponse": {
-              "result": {
-                "feedback": "",
-                "is_correct": true,
-                "detailed_feedback": [
-                  [
-                    {
-                      "feedback": "",
-                      "is_correct": true,
-                      "response_latex": "-1",
-                      "response_simplified": "-1"
-                    }
-                  ],
-                  [
-                    {
-                      "feedback": "",
-                      "is_correct": true,
-                      "response_latex": "5",
-                      "response_simplified": "5"
-                    }
-                  ],
-                  [
-                    {
-                      "feedback": "",
-                      "is_correct": true,
-                      "response_latex": "-1",
-                      "response_simplified": "-1"
-                    }
-                  ]
-                ]
-              },
-              "command": "eval"
-            }
-          }
-        },
-        {
-          "publishedPartId": "04d2cab2-eeca-4c1f-bcef-6401fe4ec635",
-          "publishedPartPosition": 0,
-          "publishedResponseAreaId": "4e354eda-fb78-4c26-a3ae-f8d1c06297a2",
-          "publishedResponseAreaPosition": 1,
-          "responseAreaUniversalId": "7f7ad90a-0cfb-44e0-b5d7-6be585206508",
-          "publishedResponseAreaPreResponseText": "$\\vec{b}+\\vec{a}=$",
-          "publishedResponseType": "MATRIX",
-          "publishedResponseConfig": { "cols": 1, "rows": 3 },
-          "totalSubmissions": 1,
-          "totalWrongSubmissions": 0,
-          "latestSubmission": {
-            "universalResponseAreaId": "7f7ad90a-0cfb-44e0-b5d7-6be585206508",
-            "answer": "[[\"-1\"],[\"5\"],[\"-1\"]]",
-            "submission": [["-1"], ["5"], ["-1"]],
-            "feedback": "Correct",
-            "rawResponse": {
-              "result": {
-                "feedback": "",
-                "is_correct": true,
-                "detailed_feedback": [
-                  [
-                    {
-                      "feedback": "",
-                      "is_correct": true,
-                      "response_latex": "-1",
-                      "response_simplified": "-1"
-                    }
-                  ],
-                  [
-                    {
-                      "feedback": "",
-                      "is_correct": true,
-                      "response_latex": "5",
-                      "response_simplified": "5"
-                    }
-                  ],
-                  [
-                    {
-                      "feedback": "",
-                      "is_correct": true,
-                      "response_latex": "-1",
-                      "response_simplified": "-1"
-                    }
-                  ]
-                ]
-              },
-              "command": "eval"
-            }
-          }
-        },
-        {
-          "publishedPartId": "dbe25bc9-6c42-4e58-aac7-d15a783ef337",
-          "publishedPartPosition": 1,
-          "publishedResponseAreaId": "85c33011-0e12-435c-944e-399043e92779",
-          "publishedResponseAreaPosition": 0,
-          "responseAreaUniversalId": "bcf1b636-f188-4c68-a59a-c153fb4419d2",
-          "publishedResponseAreaPreResponseText": "$3\\vec{c}=$",
-          "publishedResponseType": "MATRIX",
-          "publishedResponseConfig": { "cols": 1, "rows": 3 },
-          "totalSubmissions": 0,
-          "totalWrongSubmissions": 0
-        },
-        {
-          "publishedPartId": "dbe25bc9-6c42-4e58-aac7-d15a783ef337",
-          "publishedPartPosition": 1,
-          "publishedResponseAreaId": "ae4a9db7-ff54-489d-9a21-143621e6e7aa",
-          "publishedResponseAreaPosition": 1,
-          "responseAreaUniversalId": "bfab6be0-11bd-4799-ba10-1e9a0182da36",
-          "publishedResponseAreaPreResponseText": "$-\\vec{a}=$",
-          "publishedResponseType": "MATRIX",
-          "publishedResponseConfig": { "cols": 1, "rows": 3 },
-          "totalSubmissions": 0,
-          "totalWrongSubmissions": 0
-        },
-        {
-          "publishedPartId": "dbe25bc9-6c42-4e58-aac7-d15a783ef337",
-          "publishedPartPosition": 1,
-          "publishedResponseAreaId": "3ae38184-9dbc-4c5f-872b-75ddfd691b20",
-          "publishedResponseAreaPosition": 2,
-          "responseAreaUniversalId": "170e2cf2-9d0d-438e-a406-cafe7ff1017b",
-          "publishedResponseAreaPreResponseText": "$\\frac{\\vec{b}}{2}=$",
-          "publishedResponseType": "MATRIX",
-          "publishedResponseConfig": { "cols": 1, "rows": 3 },
-          "totalSubmissions": 0,
-          "totalWrongSubmissions": 0
-        },
-        {
-          "publishedPartId": "2cb9d565-8f1c-4901-bf2b-0b89f12ec9e2",
-          "publishedPartPosition": 2,
-          "publishedResponseAreaId": "6c64c869-358d-477d-b835-7c5e268b08b2",
-          "publishedResponseAreaPosition": 0,
-          "responseAreaUniversalId": "dfa4be9e-c9e2-442d-95a4-c2c61d25fe51",
-          "publishedResponseAreaPreResponseText": "$3\\vec{a}-3\\vec{c}=$",
-          "publishedResponseType": "MATRIX",
-          "publishedResponseConfig": { "cols": 1, "rows": 3 },
-          "totalSubmissions": 0,
-          "totalWrongSubmissions": 0
-        },
-        {
-          "publishedPartId": "2cb9d565-8f1c-4901-bf2b-0b89f12ec9e2",
-          "publishedPartPosition": 2,
-          "publishedResponseAreaId": "0b5d2dab-7c6d-4694-8c3b-56df1afaba96",
-          "publishedResponseAreaPosition": 1,
-          "responseAreaUniversalId": "e695440a-65b8-4ebd-b055-6173d9d84148",
-          "publishedResponseAreaPreResponseText": "$3(\\vec{a}-\\vec{c})=$",
-          "publishedResponseType": "MATRIX",
-          "publishedResponseConfig": { "cols": 1, "rows": 3 },
-          "totalSubmissions": 0,
-          "totalWrongSubmissions": 0
-        },
-        {
-          "publishedPartId": "3ee15335-1bfc-466a-8fff-66abd742dc33",
-          "publishedPartPosition": 3,
-          "publishedResponseAreaId": "8a4d3fed-e450-475d-a09b-1e58f8dbbbe2",
-          "publishedResponseAreaPosition": 0,
-          "responseAreaUniversalId": "4fd2eb4f-dec2-4695-9148-09104935d431",
-          "publishedResponseAreaPreResponseText": "$-5(\\vec{a}+\\vec{c})+\\vec{b}=$",
-          "publishedResponseType": "MATRIX",
-          "publishedResponseConfig": { "cols": 1, "rows": 3 },
-          "totalSubmissions": 0,
-          "totalWrongSubmissions": 0
-        },
-        {
-          "publishedPartId": "b2620eb0-9de6-404e-a82c-b4ad8e015d40",
-          "publishedPartPosition": 4,
-          "publishedResponseAreaId": "1a5fb3fb-001d-4ed5-9606-2a73fcae036c",
-          "publishedResponseAreaPosition": 0,
-          "responseAreaUniversalId": "c98844ba-3fd2-4154-8713-1d76bfaf610c",
-          "publishedResponseAreaPreResponseText": "$3\\vec{a}-2\\vec{c}+3\\vec{b}=$",
-          "publishedResponseType": "MATRIX",
-          "publishedResponseConfig": { "cols": 1, "rows": 3 },
-          "totalSubmissions": 0,
-          "totalWrongSubmissions": 0
-        }
-      ],
-      "questionInformation": {
-        "questionTitle": "Vector Arithmetics",
-        "questionGuidance": null,
-        "questionContent": "$$\n\\vec{a}=\\begin{bmatrix}1 \\\\ 2\\\\ 3\\end{bmatrix} \\quad \\vec{b}=\\begin{bmatrix}-2 \\\\ 3\\\\ -4\\end{bmatrix} \\quad \\vec{c}=\\begin{bmatrix} 0 \\\\\\ 4\\\\ -1\\end{bmatrix}\n$$",
-        "durationLowerBound": 1,
-        "durationUpperBound": 3,
-        "parts": [
-          {
-            "publishedPartId": "04d2cab2-eeca-4c1f-bcef-6401fe4ec635",
-            "publishedPartPosition": 0,
-            "publishedPartContent": "",
-            "publishedPartAnswerContent": "",
-            "publishedWorkedSolutionSections": [
-              {
-                "id": "9cb3a347-c05f-4101-b522-a0c6d2fbc6bf",
-                "position": 0,
-                "title": "",
-                "content": "As we made no mention that the vectors are represented in different basis, we can assume that they all share the same basis set.\n\n&#x20; &#x20;\n\n$$\n\\begin{array}{rl}\n\\vec{a}+\\vec{b} &=\\begin{bmatrix}1 \\\\ 2\\\\ 3\\end{bmatrix}+\\begin{bmatrix}-2 \\\\ 3\\\\ -4\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}1-2 \\\\ 2+3\\\\ 3-4\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}-1 \\\\ 5\\\\ -1\\end{bmatrix}\n\\end{array}\n$$\n\n&#x20; &#x20;\n\n***\n\n&#x20; &#x20;\n\nSince vector addition is commutative, $\\vec{b}+\\vec{a}=\\vec{a}+\\vec{b}$\n"
-              }
-            ],
-            "publishedResponseAreas": [
-              {
-                "id": "1194c5b3-1831-433e-b78e-c96c94f99117",
-                "position": 0,
-                "universalResponseAreaId": "ee1df8dd-bccb-46d2-b816-1a707115495e",
-                "preResponseText": "$\\vec{a}+\\vec{b}=$",
-                "Response": {
-                  "id": "0c40d7de-15a0-4916-af34-5ee2b330f3dd",
-                  "responseType": "MATRIX",
-                  "config": { "cols": 1, "rows": 3 },
-                  "answer": [["-1"], ["5"], ["-1"]]
-                },
-                "responseType": "MATRIX",
-                "answer": [["-1"], ["5"], ["-1"]]
-              },
-              {
-                "id": "4e354eda-fb78-4c26-a3ae-f8d1c06297a2",
-                "position": 1,
-                "universalResponseAreaId": "7f7ad90a-0cfb-44e0-b5d7-6be585206508",
-                "preResponseText": "$\\vec{b}+\\vec{a}=$",
-                "Response": {
-                  "id": "a01003b5-f11d-4d06-9ec2-35efc11b7b64",
-                  "responseType": "MATRIX",
-                  "config": { "cols": 1, "rows": 3 },
-                  "answer": [["-1"], ["5"], ["-1"]]
-                },
-                "responseType": "MATRIX",
-                "answer": [["-1"], ["5"], ["-1"]]
-              }
-            ]
-          },
-          {
-            "publishedPartId": "dbe25bc9-6c42-4e58-aac7-d15a783ef337",
-            "publishedPartPosition": 1,
-            "publishedPartContent": "",
-            "publishedPartAnswerContent": "",
-            "publishedWorkedSolutionSections": [
-              {
-                "id": "7d696718-c6f5-4e11-a5fb-bd122c1274cf",
-                "position": 0,
-                "title": "$3\\vec{c}$",
-                "content": "$$\n\\begin{array}{rl}\n3\\vec{c}&=3\\begin{bmatrix}0 \\\\ 4\\\\ -1\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}3\\cdot0 \\\\ 3\\cdot4\\\\ 3\\cdot-1\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}0 \\\\ 12\\\\ -3\\end{bmatrix}\n\\end{array}\n$$"
-              },
-              {
-                "id": "8134095d-f49e-4834-b8de-12baa9625d7c",
-                "position": 0,
-                "title": "",
-                "content": ""
-              },
-              {
-                "id": "010b312d-2699-439f-bcab-9f8ac4b96114",
-                "position": 1,
-                "title": "$-\\vec{a}$",
-                "content": "$$\n\\begin{array}{rl}\n-\\vec{a}&=3\\begin{bmatrix}1 \\\\ 2\\\\ 3\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}-1\\cdot1 \\\\ -1\\cdot2\\\\ -1\\cdot3\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}-1 \\\\ -2\\\\ -3\\end{bmatrix}\n\\end{array}\n$$"
-              },
-              {
-                "id": "20602e93-de57-4e69-afcd-e05f454155de",
-                "position": 2,
-                "title": "$\\frac{\\vec{b}}{2}$",
-                "content": "$$\n\\begin{array}{rl}\n\\frac{\\vec{b}}{2}&=\\frac{1}{2}\\begin{bmatrix}-2 \\\\ 3\\\\ -4\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}\\frac{1}{2}\\cdot-2 \\\\ \\frac{1}{2}\\cdot3\\\\ \\frac{1}{2}\\cdot-4\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}-1 \\\\ 1.5\\\\ -2\\end{bmatrix}\n\\end{array}\n$$"
-              }
-            ],
-            "publishedResponseAreas": [
-              {
-                "id": "85c33011-0e12-435c-944e-399043e92779",
-                "position": 0,
-                "universalResponseAreaId": "bcf1b636-f188-4c68-a59a-c153fb4419d2",
-                "preResponseText": "$3\\vec{c}=$",
-                "Response": {
-                  "id": "8623c885-b49d-476e-b316-f45a5052cacd",
-                  "responseType": "MATRIX",
-                  "config": { "cols": 1, "rows": 3 },
-                  "answer": [["0"], ["12"], ["-3"]]
-                },
-                "responseType": "MATRIX",
-                "answer": [["0"], ["12"], ["-3"]]
-              },
-              {
-                "id": "ae4a9db7-ff54-489d-9a21-143621e6e7aa",
-                "position": 1,
-                "universalResponseAreaId": "bfab6be0-11bd-4799-ba10-1e9a0182da36",
-                "preResponseText": "$-\\vec{a}=$",
-                "Response": {
-                  "id": "150494b9-1ed6-44f3-8d52-8bdbd74f7c75",
-                  "responseType": "MATRIX",
-                  "config": { "cols": 1, "rows": 3 },
-                  "answer": [["-1"], ["-2"], ["-3"]]
-                },
-                "responseType": "MATRIX",
-                "answer": [["-1"], ["-2"], ["-3"]]
-              },
-              {
-                "id": "3ae38184-9dbc-4c5f-872b-75ddfd691b20",
-                "position": 2,
-                "universalResponseAreaId": "170e2cf2-9d0d-438e-a406-cafe7ff1017b",
-                "preResponseText": "$\\frac{\\vec{b}}{2}=$",
-                "Response": {
-                  "id": "ed65b833-5cb6-4198-aa2f-a5b74d86b3f1",
-                  "responseType": "MATRIX",
-                  "config": { "cols": 1, "rows": 3 },
-                  "answer": [["-1"], ["1.5"], ["-2"]]
-                },
-                "responseType": "MATRIX",
-                "answer": [["-1"], ["1.5"], ["-2"]]
-              }
-            ]
-          },
-          {
-            "publishedPartId": "2cb9d565-8f1c-4901-bf2b-0b89f12ec9e2",
-            "publishedPartPosition": 2,
-            "publishedPartContent": "",
-            "publishedPartAnswerContent": "",
-            "publishedWorkedSolutionSections": [
-              {
-                "id": "8162a17d-bed7-48a9-b487-0bd6280a3988",
-                "position": 0,
-                "title": "$3\\vec{a}-3\\vec{c}$",
-                "content": "$$\n\\begin{array}{rl}\n3\\vec{a}-3\\vec{c} &=3\\begin{bmatrix}1 \\\\ 2\\\\ 3\\end{bmatrix}-3\\begin{bmatrix}0 \\\\ 4\\\\ -1\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}3 \\\\ 6\\\\ 9\\end{bmatrix}-\\begin{bmatrix}0 \\\\ 12\\\\ -3\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}3-0 \\\\ 6-12\\\\ 9-(-3)\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}3 \\\\ -6\\\\ 12\\end{bmatrix}\n\\end{array}\n$$"
-              },
-              {
-                "id": "7526f2d6-baf1-4aea-bb9f-6486f7233f97",
-                "position": 0,
-                "title": "",
-                "content": "As we made no mention that the vectors are represented in different basis, we can assume that they all share the same basis set.\n\n&#x20; &#x20;\n\n$$\n\\begin{array}{rl}\n3\\vec{a}-2\\vec{c}+3\\vec{b} &=3\\begin{bmatrix}1 \\\\ 2\\\\ 3\\end{bmatrix}-2\\begin{bmatrix}0 \\\\ 4\\\\ -1\\end{bmatrix}+3\\begin{bmatrix}-2 \\\\ 3\\\\ -4\\end{bmatrix}\n\\\\\\\\\n&= 3\\left(\\begin{bmatrix}1 \\\\ 2\\\\ 3\\end{bmatrix}+\\begin{bmatrix}-2 \\\\ 3\\\\ -4\\end{bmatrix}\\right)-2\\begin{bmatrix}0 \\\\ 4\\\\ -1\\end{bmatrix}\n\\\\\\\\\n&= 3\\begin{bmatrix}1-2 \\\\ 2+3\\\\ 3-4\\end{bmatrix}-2\\begin{bmatrix}0 \\\\ 4\\\\ -1\\end{bmatrix}\n\\\\\\\\\n&= 3\\begin{bmatrix}-1 \\\\ 5\\\\ -1\\end{bmatrix}-2\\begin{bmatrix}0 \\\\ 4\\\\ -1\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}-3 \\\\ 15\\\\ -3\\end{bmatrix}+\\begin{bmatrix}0 \\\\ -8\\\\ 2\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}-3+0 \\\\ 15-8\\\\ -3+2\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}-3 \\\\ 7\\\\ -1\\end{bmatrix}\n\\end{array}\n$$"
-              },
-              {
-                "id": "9d88d02a-e63d-40d0-8200-5d81718e2259",
-                "position": 1,
-                "title": "$3(\\vec{a}-\\vec{c})$",
-                "content": "$$\n\\begin{array}{rl}\n3(\\vec{a}-\\vec{c}) &=3\\left(\\begin{bmatrix}1 \\\\ 2\\\\ 3\\end{bmatrix}-\\begin{bmatrix}0 \\\\ 4\\\\ -1\\end{bmatrix}\\right)\n\\\\\\\\\n&= 3\\begin{bmatrix}1-0 \\\\ 2-4\\\\ 3-(-1)\\end{bmatrix}\n\\\\\\\\\n&= 3\\begin{bmatrix}1 \\\\ -2\\\\ 4\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}3 \\\\ -6\\\\ 12\\end{bmatrix}\n\\end{array}\n$$"
-              }
-            ],
-            "publishedResponseAreas": [
-              {
-                "id": "6c64c869-358d-477d-b835-7c5e268b08b2",
-                "position": 0,
-                "universalResponseAreaId": "dfa4be9e-c9e2-442d-95a4-c2c61d25fe51",
-                "preResponseText": "$3\\vec{a}-3\\vec{c}=$",
-                "Response": {
-                  "id": "53199f12-d660-4ffe-95ae-9d28c9f260ba",
-                  "responseType": "MATRIX",
-                  "config": { "cols": 1, "rows": 3 },
-                  "answer": [["3"], ["-6"], ["12"]]
-                },
-                "responseType": "MATRIX",
-                "answer": [["3"], ["-6"], ["12"]]
-              },
-              {
-                "id": "0b5d2dab-7c6d-4694-8c3b-56df1afaba96",
-                "position": 1,
-                "universalResponseAreaId": "e695440a-65b8-4ebd-b055-6173d9d84148",
-                "preResponseText": "$3(\\vec{a}-\\vec{c})=$",
-                "Response": {
-                  "id": "c189d436-d5a4-4c61-b26e-95ac30c281b5",
-                  "responseType": "MATRIX",
-                  "config": { "cols": 1, "rows": 3 },
-                  "answer": [["3"], ["-6"], ["12"]]
-                },
-                "responseType": "MATRIX",
-                "answer": [["3"], ["-6"], ["12"]]
-              }
-            ]
-          },
-          {
-            "publishedPartId": "3ee15335-1bfc-466a-8fff-66abd742dc33",
-            "publishedPartPosition": 3,
-            "publishedPartContent": "",
-            "publishedPartAnswerContent": "",
-            "publishedWorkedSolutionSections": [
-              {
-                "id": "916cf686-7a41-4169-8fa2-b2e9d6924375",
-                "position": 0,
-                "title": "",
-                "content": "As we made no mention that the vectors are represented in different basis, we can assume that they all share the same basis set.\n\n&#x20; &#x20;\n\n$$\n\\begin{array}{rl}\n-5(\\vec{a}+\\vec{c})+\\vec{b} &=-5\\left(\\begin{bmatrix}1 \\\\ 2\\\\ 3\\end{bmatrix}+\\begin{bmatrix}0 \\\\ 4\\\\ -1\\end{bmatrix}\\right)+\\begin{bmatrix}-2 \\\\ 3\\\\ -4\\end{bmatrix}\n\\\\\\\\\n&= -5\\begin{bmatrix}1 \\\\ 2\\\\ 3\\end{bmatrix}-5\\begin{bmatrix}0 \\\\ 4\\\\ -1\\end{bmatrix}+\\begin{bmatrix}-2 \\\\ 3\\\\ -4\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}-5 \\\\ -10\\\\ -15\\end{bmatrix}+\\begin{bmatrix}0 \\\\ -20\\\\ 5\\end{bmatrix}+\\begin{bmatrix}-2 \\\\ 3\\\\ -4\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}-5+0-2 \\\\ -10-20+3\\\\ -15+5-4\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}-7 \\\\ -27\\\\ -14\\end{bmatrix}\n\\end{array}\n$$"
-              }
-            ],
-            "publishedResponseAreas": [
-              {
-                "id": "8a4d3fed-e450-475d-a09b-1e58f8dbbbe2",
-                "position": 0,
-                "universalResponseAreaId": "4fd2eb4f-dec2-4695-9148-09104935d431",
-                "preResponseText": "$-5(\\vec{a}+\\vec{c})+\\vec{b}=$",
-                "Response": {
-                  "id": "eeb448a9-81b9-4f4d-836f-d04029aa449c",
-                  "responseType": "MATRIX",
-                  "config": { "cols": 1, "rows": 3 },
-                  "answer": [["-7"], ["-27"], ["-14"]]
-                },
-                "responseType": "MATRIX",
-                "answer": [["-7"], ["-27"], ["-14"]]
-              }
-            ]
-          },
-          {
-            "publishedPartId": "b2620eb0-9de6-404e-a82c-b4ad8e015d40",
-            "publishedPartPosition": 4,
-            "publishedPartContent": "",
-            "publishedPartAnswerContent": "",
-            "publishedWorkedSolutionSections": [
-              {
-                "id": "3d15fcbc-5173-4d89-a13a-3296c25680ec",
-                "position": 0,
-                "title": "",
-                "content": "As we made no mention that the vectors are represented in different basis, we can assume that they all share the same basis set.\n\n&#x20; &#x20;\n\n$$\n\\begin{array}{rl}\n3\\vec{a}-2\\vec{c}+3\\vec{b} &=3\\begin{bmatrix}1 \\\\ 2\\\\ 3\\end{bmatrix}-2\\begin{bmatrix}0 \\\\ 4\\\\ -1\\end{bmatrix}+3\\begin{bmatrix}-2 \\\\ 3\\\\ -4\\end{bmatrix}\n\\\\\\\\\n&= 3\\left(\\begin{bmatrix}1 \\\\ 2\\\\ 3\\end{bmatrix}+\\begin{bmatrix}-2 \\\\ 3\\\\ -4\\end{bmatrix}\\right)-2\\begin{bmatrix}0 \\\\ 4\\\\ -1\\end{bmatrix}\n\\\\\\\\\n&= 3\\begin{bmatrix}1-2 \\\\ 2+3\\\\ 3-4\\end{bmatrix}-2\\begin{bmatrix}0 \\\\ 4\\\\ -1\\end{bmatrix}\n\\\\\\\\\n&= 3\\begin{bmatrix}-1 \\\\ 5\\\\ -1\\end{bmatrix}-2\\begin{bmatrix}0 \\\\ 4\\\\ -1\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}-3 \\\\ 15\\\\ -3\\end{bmatrix}+\\begin{bmatrix}0 \\\\ -8\\\\ 2\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}-3+0 \\\\ 15-8\\\\ -3+2\\end{bmatrix}\n\\\\\\\\\n&= \\begin{bmatrix}-3 \\\\ 7\\\\ -1\\end{bmatrix}\n\\end{array}\n$$"
-              }
-            ],
-            "publishedResponseAreas": [
-              {
-                "id": "1a5fb3fb-001d-4ed5-9606-2a73fcae036c",
-                "position": 0,
-                "universalResponseAreaId": "c98844ba-3fd2-4154-8713-1d76bfaf610c",
-                "preResponseText": "$3\\vec{a}-2\\vec{c}+3\\vec{b}=$",
-                "Response": {
-                  "id": "c5f1c9a9-eb13-4b2e-83ad-f9a3bd87a056",
-                  "responseType": "MATRIX",
-                  "config": { "cols": 1, "rows": 3 },
-                  "answer": [["-3"], ["7"], ["-1"]]
-                },
-                "responseType": "MATRIX",
-                "answer": [["-3"], ["7"], ["-1"]]
-              }
-            ]
-          }
-        ]
-      },
-      "questionAccessInformation": {
-        "estimatedMinimumTime": "1 minute",
-        "estimaredMaximumTime": "3 minutes",
-        "timeTaken": "13 minutes",
-        "accessStatus": "too much time spent on this question.",
-        "markedDone": "Part (a) was marked done",
-        "currentPart": {
-          "id": "04d2cab2-eeca-4c1f-bcef-6401fe4ec635",
-          "position": 0
-        }
-      }
-    },
-    "conversation_id": "2f95c229-af43-4c1f-918a-0fd428178acf",
-    "agent_type": "base"
-  }
-}
diff --git a/src/agent/utils/parse_json_context_to_prompt.py b/src/agent/utils/parse_json_context_to_prompt.py
deleted file mode 100644
index a233126..0000000
--- a/src/agent/utils/parse_json_context_to_prompt.py
+++ /dev/null
@@ -1,334 +0,0 @@
-"""
-Refactored JSON to prompt parser using improved, clearer structure.
-"""
-
-from typing import List, Optional, Dict, Any, Union
-from src.agent.utils.prompt_context_templates import PromptFormatter
-
-# Definitions questionSubmissionSummary type
-class StudentLatestSubmission:
-    def __init__(
-        self,
-        universalResponseAreaId: Optional[str] = None,
-        answer: Optional[str] = None,
-        submission: Optional[str] = None,
-        feedback: Optional[str] = None,
-        rawResponse: Optional[dict] = None,
-    ):
-        self.universalResponseAreaId = universalResponseAreaId
-        self.answer = answer
-        self.submission = submission
-        self.feedback = feedback
-        self.rawResponse = rawResponse
-
-class StudentWorkResponseArea:
-    def __init__(
-        self,
-        publishedPartId: Optional[str] = None,
-        publishedPartPosition: Optional[int] = None,
-        publishedResponseAreaId: Optional[str] = None,
-        publishedResponseAreaPosition: Optional[int] = None,
-        responseAreaUniversalId: Optional[str] = None,
-        publishedResponseAreaPreResponseText: Optional[str] = None,
-        publishedResponseType: Optional[str] = None,
-        publishedResponseConfig: Optional[dict] = None,
-        totalSubmissions: Optional[int] = None,
-        totalWrongSubmissions: Optional[int] = None,
-        latestSubmission: Optional[StudentLatestSubmission] = None,
-    ):
-        self.publishedPartId = publishedPartId
-        self.publishedPartPosition = publishedPartPosition
-        self.publishedResponseAreaId = publishedResponseAreaId
-        self.publishedResponseAreaPosition = publishedResponseAreaPosition
-        self.responseAreaUniversalId = responseAreaUniversalId
-        self.publishedResponseAreaPreResponseText = publishedResponseAreaPreResponseText
-        self.publishedResponseType = publishedResponseType
-        self.publishedResponseConfig = publishedResponseConfig
-        self.latestSubmission = StudentLatestSubmission(**latestSubmission) if latestSubmission else None
-        self.totalSubmissions = totalSubmissions
-        self.totalWrongSubmissions = totalWrongSubmissions
-
-# questionInformation type
-class ResponseAreaDetails:
-    def __init__(
-        self,
-        id: Optional[str] = None,
-        position: Optional[int] = None,
-        universalResponseAreaId: Optional[str] = None,
-        preResponseText: Optional[str] = None,
-        responseType: Optional[str] = None,
-        answer: Optional[dict] = None,
-        Response: Optional[dict] = None,
-    ):
-        self.id = id
-        self.position = position
-        self.universalResponseAreaId = universalResponseAreaId
-        self.preResponseText = preResponseText
-        self.responseType = responseType
-        self.answer = answer
-        self.Response = Response
-
-class PartDetails:
-    def __init__(
-        self,
-        publishedPartId: Optional[str] = None,
-        publishedPartPosition: Optional[int] = None,
-        publishedPartContent: Optional[str] = None,
-        publishedPartAnswerContent: Optional[str] = None,
-        publishedWorkedSolutionSections: Optional[List[dict]] = [],
-        publishedStructuredTutorialSections: Optional[List[dict]] = [],
-        publishedResponseAreas: Optional[List[Optional[ResponseAreaDetails]]] = [],
-    ):
-        self.publishedPartId = publishedPartId
-        self.publishedPartPosition = publishedPartPosition
-        self.publishedPartContent = publishedPartContent
-        self.publishedPartAnswerContent = publishedPartAnswerContent
-        self.publishedWorkedSolutionSections = publishedWorkedSolutionSections
-        self.publishedStructuredTutorialSections = publishedStructuredTutorialSections
-        self.publishedResponseAreas = [ResponseAreaDetails(**publishedResponseArea) for publishedResponseArea in publishedResponseAreas]
-
-class QuestionDetails:
-    def __init__(
-        self,
-        setNumber: Optional[int] = None,
-        setName: Optional[str] = None,
-        setDescription: Optional[str] = None,
-        questionNumber: Optional[int] = None,
-        questionTitle: Optional[str] = None,
-        questionGuidance: Optional[str] = None,
-        questionContent: Optional[str] = None,
-        durationLowerBound: Optional[int] = None,
-        durationUpperBound: Optional[int] = None,
-        parts: Optional[List[PartDetails]] = [],
-    ):
-        self.setNumber = setNumber
-        self.setName = setName
-        self.setDescription = setDescription
-        self.questionNumber = questionNumber
-        self.questionTitle = questionTitle
-        self.questionGuidance = questionGuidance
-        self.questionContent = questionContent
-        self.durationLowerBound = durationLowerBound
-        self.durationUpperBound = durationUpperBound
-        self.parts = [PartDetails(**part) for part in parts] 
-
-# questionAccessInformation type
-class CurrentPart:
-    def __init__(
-        self, 
-        id: str = None, 
-        position: int = None, 
-        universalPartId: Optional[str] = None,
-        timeTakenPart: Optional[str] = None, 
-        markedDonePart: Optional[str] = None
-    ):
-        self.id = id
-        self.position = position
-        self.universalPartId = universalPartId
-        self.timeTakenPart = timeTakenPart
-        self.markedDonePart = markedDonePart
-
-class QuestionAccessInformation:
-    def __init__(
-        self,
-        estimatedMinimumTime: Optional[str] = None,
-        estimaredMaximumTime: Optional[str] = None,
-        timeTaken: Optional[str] = None,
-        accessStatus: Optional[str] = None,
-        markedDone: Optional[str] = None,
-        currentPart: Optional[Dict[str, Union[str, int]]] = {},
-    ):
-        self.estimatedMinimumTime = estimatedMinimumTime
-        self.estimaredMaximumTime = estimaredMaximumTime
-        self.timeTaken = timeTaken
-        self.accessStatus = accessStatus
-        self.markedDone = markedDone
-        self.currentPart = CurrentPart(**currentPart)
-
-
-def parse_json_to_structured_prompt(
-    question_submission_summary: Optional[List[StudentWorkResponseArea]],
-    question_information: Optional[QuestionDetails],
-    question_access_information: Optional[QuestionAccessInformation]
-) -> str:
-    """
-    Parse JSON data into a well-structured, LLM-friendly prompt.
-    
-    Args:
-        question_submission_summary: Student's work and submissions
-        question_information: Question details and structure
-        question_access_information: Current progress and timing info
-        
-    Returns:
-        Formatted prompt string or error message
-    """
-    
-    if not question_information:
-        return PromptFormatter.format_error_message()
-    
-    # Convert to proper objects
-    submission_summary = [StudentWorkResponseArea(**summary) for summary in question_submission_summary]
-    question_info = QuestionDetails(**question_information)
-    access_info = QuestionAccessInformation(**question_access_information) if question_access_information else None
-    
-    # TODO: EXPERIMENTAL - Remove later
-    # if question_info.setNumber is not None:
-    #     if (question_info.setNumber + 1) % 2 != 0:
-    #         return PromptFormatter.format_no_context_message()
-    
-    # Build prompt sections
-    sections = []
-    
-    # 1. Question Header
-    current_part_letter = None
-    if access_info and access_info.currentPart:
-        current_part_letter = PromptFormatter.get_part_letter(access_info.currentPart.position)
-    
-    set_info = {
-        'number': question_info.setNumber,
-        'name': question_info.setName
-    }
-    
-    question_data = {
-        'number': question_info.questionNumber,
-        'title': question_info.questionTitle,
-        'guidance': question_info.questionGuidance,
-        'content': question_info.questionContent,
-        'duration_lower': question_info.durationLowerBound,
-        'duration_upper': question_info.durationUpperBound
-    }
-    
-    sections.append(PromptFormatter.format_question_header(
-        set_info, 
-        question_data, 
-        current_part_letter
-    ))
-    
-    # 2. Progress Summary (if available)
-    if access_info:
-        progress_section = PromptFormatter.format_progress_summary(
-            access_info.timeTaken,
-            access_info.accessStatus,
-            access_info.markedDone
-        )
-        if progress_section:
-            sections.append(progress_section)
-    
-    # 3. Parts Details
-    for part in question_info.parts:
-        sections.append(_format_single_part(
-            part, 
-            access_info.currentPart if access_info else None,
-            submission_summary
-        ))
-    
-    # 4. Combine into final prompt
-    return PromptFormatter.format_complete_prompt(sections)
-
-
-def _format_single_part(
-    part: PartDetails, 
-    current_part: Optional[CurrentPart],
-    submissions: List[StudentWorkResponseArea]
-) -> str:
-    """Format a single part with all its components."""
-    
-    if not part:
-        return ""
-    
-    part_sections = []
-    part_letter = PromptFormatter.get_part_letter(part.publishedPartPosition)
-    
-    # Determine if this is the current part
-    is_current = current_part and current_part.id == part.publishedPartId
-    time_on_part = current_part.timeTakenPart if is_current and current_part else None
-    
-    # 1. Part Header
-    part_sections.append(PromptFormatter.format_part_header(
-        part_letter, 
-        is_current, 
-        time_on_part
-    ))
-    
-    # 2. Part Content
-    part_sections.append(PromptFormatter.format_part_content(part.publishedPartContent))
-    
-    # 3. Response Areas
-    response_areas = []
-    for response_area in part.publishedResponseAreas:
-        student_work = _extract_student_work_for_area(response_area, submissions)
-        response_areas.append(PromptFormatter.format_single_response_area(
-            response_area.position,
-            response_area.preResponseText,
-            response_area.answer,
-            student_work
-        ))
-    
-    if response_areas:
-        part_sections.append(PromptFormatter.format_response_areas(response_areas))
-    
-    # 4. Final Part Answer
-    part_sections.append(PromptFormatter.format_part_answer(part.publishedPartAnswerContent))
-    
-    # 5. Worked Solutions
-    solutions_data = []
-    if part.publishedWorkedSolutionSections:
-        for ws in part.publishedWorkedSolutionSections:
-            solutions_data.append({
-                'title': ws.get('title', ''),
-                'content': ws.get('content', ''),
-                'position': ws.get('position', 0)
-            })
-
-    #  6. Structured Tutorial Sections
-    tutorial_data = []
-    if part.publishedStructuredTutorialSections:
-        for ts in part.publishedStructuredTutorialSections:
-            tutorial_data.append({
-                'title': ts.get('title', ''),
-                'content': ts.get('content', ''),
-                'position': ts.get('position', 0)
-            })
-    
-    part_sections.append(PromptFormatter.format_worked_solutions(solutions_data))
-    part_sections.append(PromptFormatter.format_structured_tutorials(tutorial_data))
-
-    return "\n".join(part_sections) + "\n---\n"
-
-
-def _extract_student_work_for_area(
-    response_area: ResponseAreaDetails, 
-    submissions: List[StudentWorkResponseArea]
-) -> Dict[str, Any]:
-    """Extract student work data for a specific response area."""
-    
-    for submission in submissions:
-        if (submission.publishedResponseAreaId == response_area.id and 
-            submission.latestSubmission):
-            
-            return {
-                'has_submissions': True,
-                'latest_response': submission.latestSubmission.submission,
-                'latest_feedback': submission.latestSubmission.feedback,
-                'total_submissions': submission.totalSubmissions,
-                'total_wrong': submission.totalWrongSubmissions
-            }
-    
-    return {'has_submissions': False}
-
-
-# Convenience function that maintains the original interface
-def parse_json_to_prompt(
-    questionSubmissionSummary: Optional[List[StudentWorkResponseArea]],
-    questionInformation: Optional[QuestionDetails],
-    questionAccessInformation: Optional[QuestionAccessInformation]
-) -> str:
-    """
-    Legacy wrapper for backward compatibility.
-    Recommended to use parse_json_to_structured_prompt for new code.
-    """
-    return parse_json_to_structured_prompt(
-        questionSubmissionSummary,
-        questionInformation,
-        questionAccessInformation
-    )
diff --git a/src/agent/utils/prompt_context_templates.py b/src/agent/utils/prompt_context_templates.py
deleted file mode 100644
index e77a208..0000000
--- a/src/agent/utils/prompt_context_templates.py
+++ /dev/null
@@ -1,263 +0,0 @@
-"""
-Improved prompt templates with clearer structure for LLM consumption.
-Uses hierarchical organization and consistent formatting.
-"""
-
-from typing import Optional, List, Dict, Any
-
-class PromptFormatter:
-    """Centralized prompt formatting with clear structure."""
-    
-    @staticmethod
-    def format_error_message() -> str:
-        """Return structured error message."""
-        return """
-# ERROR: Question details unavailable
-
-Please describe the question you're working on so I can assist you effectively.
-"""
-
-    @staticmethod
-    def format_no_context_message() -> str:
-        """Return structured no-context message."""
-        return """
-# NOTICE: Question details not provided
-
-Please tell me about the question you're working on. I'll use British English spellings.
-"""
-
-    @staticmethod
-    def format_question_header(
-        set_info: Dict[str, Any],
-        question_info: Dict[str, Any],
-        current_part: Optional[str] = None
-    ) -> str:
-        """Format the main question header with metadata."""
-        
-        # Build title components
-        title_parts = []
-        if set_info.get('number') is not None and set_info.get('name'):
-            title_parts.append(f"## Set {set_info['number'] + 1}: {set_info['name']}")
-        
-        question_num = ""
-        if set_info.get('number') is not None and question_info.get('number') is not None:
-            question_num = f"{set_info['number'] + 1}.{question_info['number'] + 1}"
-        
-        title_parts.append(f"## Question {question_num}: {question_info['title']}")
-        
-        # Current progress indicator
-        progress_indicator = ""
-        if current_part:
-            progress_indicator = f"### Currently working on: Part ({current_part})"
-        
-        # Question metadata
-        guidance = question_info.get('guidance', 'None provided')
-        content = question_info.get('content', 'No description available')
-        
-        # Duration formatting
-        duration_text = "- Expected Duration: "
-        if question_info.get('duration_lower') and question_info.get('duration_upper'):
-            duration_text += f"{question_info['duration_lower']}-{question_info['duration_upper']} minutes"
-        else:
-            duration_text += "Not specified"
-        
-        return f"""
-# Question Context
-
-{"\n".join(title_parts)}
-
-{progress_indicator}
-
-### Question Details
-- Guidance: {guidance}
-- Description: {content}
-{duration_text}
-
-> Note: Mathematical equations are in KaTeX format, preserve them the same. Ensure mathematical equations are surrounded by one '$' for in-line equations and '$$' for block equations. 
-Example: '$E=mc^2$' or '$$E=mc^2$$'.
-Use British English spellings.
-
----
-"""
-
-    @staticmethod
-    def format_progress_summary(
-        time_taken: Optional[str] = None,
-        access_status: Optional[str] = None,
-        marked_done: Optional[str] = None
-    ) -> str:
-        """Format progress and timing information."""
-        if not any([time_taken, access_status, marked_done]):
-            return ""
-        
-        progress_items = []
-        if time_taken:
-            progress_items.append(f"- Time spent today: {time_taken}")
-        if access_status:
-            progress_items.append(f"- Status: {access_status}")
-        if marked_done:
-            progress_items.append(f"- Completion: {marked_done}")
-        
-        return f"""
-# Progress Summary
-
-{"\n".join(progress_items)}
-
----
-"""
-
-    @staticmethod
-    def format_part_header(
-        part_letter: str,
-        is_current: bool = False,
-        time_on_part: Optional[str] = None
-    ) -> str:
-        """Format part header with clear indicators."""
-        
-        status_text = " [CURRENTLY WORKING ON]" if is_current else ""
-        
-        header = f"## Part ({part_letter}){status_text}"
-        
-        if is_current and time_on_part:
-            time_display = time_on_part if time_on_part != 'No recorded duration' else 'No time recorded'
-            header += f"\n\n*Time spent on this part: {time_display}*"
-        
-        return header
-
-    @staticmethod
-    def format_part_content(content: Optional[str]) -> str:
-        """Format part content with clear labeling."""
-        if not content or not content.strip():
-            return "### Part Content\n\nNo content provided"
-        
-        return f"### Part Content\n\n{content.strip()}"
-
-    @staticmethod
-    def format_response_areas(response_areas: List[str]) -> str:
-        """Format multiple response areas with clear separation."""
-        if not response_areas:
-            return "### Response Areas\n\nNone defined"
-        
-        return f"""
-### Response Areas
-
-{"\n".join(response_areas)}
-"""
-
-    @staticmethod
-    def format_single_response_area(
-        position: int,
-        task_description: Optional[str],
-        expected_answer: Any,
-        student_work: Dict[str, Any]
-    ) -> str:
-        """Format a single response area with student work."""
-        
-        # Format task description
-        task_text = f"- Task: {task_description}" if task_description else "- Task: Not specified"
-        
-        # Format expected answer (keep secret)
-        answer_text = f"- Expected Answer (confidential): {expected_answer}"
-        
-        # Format student submissions
-        submission_text = PromptFormatter._format_student_submissions(student_work)
-        
-        return f"""
-#### Response Area {position + 1}
-
-{task_text}
-{answer_text}
-{submission_text}
-"""
-
-    @staticmethod
-    def _format_student_submissions(student_work: Dict[str, Any]) -> str:
-        """Format student submission history."""
-        if not student_work.get('has_submissions'):
-            return "- Your Work: No responses submitted yet"
-        
-        latest = student_work.get('latest_response', 'None')
-        feedback = student_work.get('latest_feedback', 'None')
-        total = student_work.get('total_submissions', 0)
-        wrong = student_work.get('total_wrong', 0)
-        
-        return f"""- Your Work:
-  - Latest response: {latest}
-  - Latest feedback: {feedback}
-  - Total attempts: {total}
-  - Incorrect attempts: {wrong}"""
-
-    @staticmethod
-    def format_part_answer(answer_content: Optional[str]) -> str:
-        """Format the final part answer."""
-        if not answer_content:
-            return "### Final Answer\n\nNo direct answer specified for this part"
-        
-        return f"### Final Answer\n\n{answer_content}"
-
-    @staticmethod
-    def format_worked_solutions(solutions: List[Dict[str, Any]]) -> str:
-        """Format worked solutions section."""
-        if not solutions:
-            return "### Worked Solutions\n\nNone available"
-        
-        solution_texts = []
-        for i, solution in enumerate(solutions):
-            title = solution.get('title', f'Solution {i + 1}')
-            content = solution.get('content', '').strip() or 'No content available'
-            solution_texts.append(f"#### {title}\n\n{content}")
-        
-        return f"""### Worked Solutions
-
-{"\n".join(solution_texts)}"""
-
-    @staticmethod
-    def format_structured_tutorials(tutorials: List[Dict[str, Any]]) -> str:
-        """Format structured tutorials section."""
-        if not tutorials:
-            return "### Structured Tutorials\n\nNone available"
-        
-        tutorial_texts = []
-        for i, tutorial in enumerate(tutorials):
-            title = tutorial.get('title', f'Tutorial {i + 1}')
-            content = tutorial.get('content', '').strip() or 'No content available'
-            tutorial_texts.append(f"#### {title}\n\n{content}")
-        
-        return f"""### Structured Tutorials
-
-{"\n".join(tutorial_texts)}"""
-
-    @staticmethod
-    def format_complete_prompt(sections: List[str]) -> str:
-        """Combine all sections into a complete, well-structured prompt."""
-        
-        intro = """
-# Personalized Learning Assistant
-
-I have detailed information about your current question, including your progress, responses, and any feedback you've received. This context helps me provide targeted assistance based on your specific situation.
-
-"""
-        
-        # Filter out empty sections and join
-        valid_sections = [section.strip() for section in sections if section and section.strip()]
-        
-        content = intro + "\n".join(valid_sections)
-        
-        # Clean up formatting
-        content = content.replace("&#x20;&#x20;", " ").replace("&#x20", " ")
-        content = "\n".join(line for line in content.split("\n") if line.strip() or not line)
-        
-        return content.strip()
-
-    @staticmethod
-    def get_part_letter(position: int) -> str:
-        """Convert position to lowercase letter (1-indexed)."""
-        return chr(96 + (position + 1))
-
-
-# Legacy function wrappers for backward compatibility
-def get_error_prompt() -> str:
-    return PromptFormatter.format_error_message()
-
-def get_no_context_prompt() -> str:
-    return PromptFormatter.format_no_context_message()
diff --git a/src/agent/utils/types.py b/src/agent/utils/types.py
deleted file mode 100644
index d038dbc..0000000
--- a/src/agent/utils/types.py
+++ /dev/null
@@ -1,4 +0,0 @@
-from typing import Any, Dict, TypeAlias
-
-JsonType = Dict[str, Any]
-InvokeAgentResponseType: TypeAlias = Dict[str, Any]
diff --git a/src/module.py b/src/module.py
index 922979f..133fdc2 100755
--- a/src/module.py
+++ b/src/module.py
@@ -1,85 +1,74 @@
 import time
-from typing import Any
-from lf_toolkit.chat.result import ChatResult as Result
-from lf_toolkit.chat.params import ChatParams as Params
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
 
-from src.agent.utils.parse_json_context_to_prompt import parse_json_to_prompt
+from lf_toolkit.chat import ChatRequest, ChatResponse, Message
+from lf_toolkit.shared.mued_api_v0_1_0 import Role
+
+from src.agent.context import parse_json_to_prompt
 from src.agent.agent import invoke_base_agent
-from src.agent.utils.types import JsonType
 
-def chat_module(message: Any, params: Params) -> JsonType:
+
+def chat_module(request: ChatRequest) -> ChatResponse:
     """
-    Function used by student to converse with a chatbot.
-    ---
-    The handler function passes three arguments to module():
+    Main entry point — receives a ChatRequest and returns a ChatResponse.
 
-    - `message` which is the message sent by the student.
-    - `params` which are any extra parameters that may be useful,
-        e.g., conversation history and summary, conversational style of user, conversation id.
+    The ChatRequest contains:
+    - `messages`: conversation history; the last message is the current student input.
+    - `conversationId`: unique session identifier.
+    - `user`: optional user info (preference, task progress).
+    - `context`: optional educational context (question details, prior conversation summary).
 
-    The output of this function is what is returned as the API response
-    and therefore must be JSON-encodable. It must also conform to the
-    response schema.
+    Edit src/agent/prompts.py to change the chatbot's behaviour.
+    Edit src/agent/agent.py to change the agent logic (summarisation threshold, LLM provider, etc.).
+    """
 
-    Any standard python library may be used, as well as any package
-    available on pip (provided it is added to requirements.txt).
+    conversation_id = request.conversationId
 
-    The way you wish to structure you code (all in this function, or
-    split into many) is entirely up to you. All that matters are the
-    return types and that module() is the main function used
-    to output the Chatbot response.
-    """
+    context = request.context or {}
+    task_progress = (request.user.taskProgress or {}) if request.user else {}
+
+    summary = context.get("summary", "") or ""
+
+    conversationalStyle = ""
+    if request.user and request.user.preference:
+        conversationalStyle = request.user.preference.model_dump().get("conversationalStyle", "") or ""
+
+    messages = _to_langchain_messages(request.messages)
 
-    result = Result()
-
-    # EXTRACT PARAMETERS
-    conversation_id = params.get("conversation_id", None)
-    if conversation_id is None:
-        raise Exception("Internal Error: The conversation id is required in the parameters of the chat module.")
-
-    include_test_data = params.get("include_test_data", False) or False
-    conversation_history = params.get("conversation_history", []) or []
-    summary = params.get("summary", "") or ""
-    conversationalStyle = params.get("conversational_style", "") or ""
-
-    question_response_details = params.get("question_response_details", {})
-    if isinstance(question_response_details, dict):
-        question_submission_summary = question_response_details.get("questionSubmissionSummary", [])
-        question_information = question_response_details.get("questionInformation", {})
-        question_access_information = question_response_details.get("questionAccessInformation", {})
-    else:
-        print("ERROR:: question_response_details is not a dict")
-        raise Exception("Internal Error: The question response details parameter is malformed.")
-    
-    # PARSE QUESTION RESPONSE DETAILS TO PROMPT
     try:
-        question_response_details_prompt = parse_json_to_prompt(
-            question_submission_summary,
-            question_information,
-            question_access_information
-        )
+        context_prompt = parse_json_to_prompt(context, task_progress)
     except Exception as e:
         print("ERROR:: ", e)
-        raise Exception("Internal Error: The question response details could not be parsed.")
-    
+        raise Exception("Internal Error: The context prompt could not be parsed.")
 
-    # RUN THE AGENT AND MEASURE PROCESSING TIME
     start_time = time.time()
-
-    chatbot_response = invoke_base_agent(query=message, \
-                            conversation_history=conversation_history, \
-                            summary=summary, \
-                            conversationalStyle=conversationalStyle, \
-                            question_response_details=question_response_details_prompt, \
-                            session_id=conversation_id)
-
+    chatbot_response = invoke_base_agent(
+        messages=messages,
+        summary=summary,
+        conversationalStyle=conversationalStyle,
+        context_prompt=context_prompt,
+        session_id=conversation_id,
+    )
     end_time = time.time()
 
-    result._processing_time = end_time - start_time
-    result.add_response("chatbot_response", chatbot_response["output"])
-    result.add_metadata("summary", chatbot_response["intermediate_steps"][0])
-    result.add_metadata("conversational_style", chatbot_response["intermediate_steps"][1])
-    result.add_metadata("conversation_history", chatbot_response["intermediate_steps"][2])
-    result.add_processing_time(end_time - start_time)
+    return ChatResponse(
+        output=Message(role=Role.ASSISTANT, content=chatbot_response["output"]),
+        metadata={
+            "summary": chatbot_response["summary"],
+            "conversationalStyle": chatbot_response["conversationalStyle"],
+            "processingTimeMs": round((end_time - start_time) * 1000),
+        },
+    )
+
 
-    return result.to_dict(include_test_data=include_test_data)
\ No newline at end of file
+def _to_langchain_messages(messages):
+    result = []
+    for m in messages:
+        role = str(m.role).upper()
+        if role == "USER":
+            result.append(HumanMessage(content=m.content))
+        elif role == "ASSISTANT":
+            result.append(AIMessage(content=m.content))
+        elif role == "SYSTEM":
+            result.append(SystemMessage(content=m.content))
+    return result
diff --git a/tests/example_inputs/example_input_0.json b/tests/example_inputs/example_input_0.json
new file mode 100644
index 0000000..3154cfc
--- /dev/null
+++ b/tests/example_inputs/example_input_0.json
@@ -0,0 +1,7 @@
+{
+  "conversationId": "example-conversation-id-0",
+  "messages": [
+    { "role": "USER", "content": "Hi, how do I solve this problem?" }
+  ],
+  "user": { "type": "LEARNER" }
+}
diff --git a/tests/example_inputs/example_input_1.json b/tests/example_inputs/example_input_1.json
new file mode 100644
index 0000000..ada3502
--- /dev/null
+++ b/tests/example_inputs/example_input_1.json
@@ -0,0 +1,123 @@
+{
+  "conversationId": "example-conversation-id-1",
+  "messages": [
+    { "role": "USER", "content": "How can I start solving this question?" }
+  ],
+  "user": {
+    "type": "LEARNER",
+    "preference": { "conversationalStyle": "" },
+    "taskProgress": {
+      "timeSpentOnQuestion": "58 minutes",
+      "accessStatus": "too much time spent on this question today.",
+      "markedDone": "This question is still being worked on.",
+      "currentPart": {
+        "position": 0,
+        "timeSpentOnPart": "37 minutes",
+        "markedDone": "This part is not marked done.",
+        "responseAreas": [
+          {
+            "responseType": "NUMBER",
+            "totalSubmissions": 1,
+            "wrongSubmissions": 1,
+            "latestSubmission": {
+              "submission": 1,
+              "feedback": "Incorrect",
+              "answer": "26"
+            }
+          },
+          {
+            "responseType": "NUMBER",
+            "totalSubmissions": 1,
+            "wrongSubmissions": 1,
+            "latestSubmission": {
+              "submission": 20,
+              "feedback": "Incorrect",
+              "answer": "53"
+            }
+          }
+        ]
+      }
+    }
+  },
+  "context": {
+    "summary": "",
+    "set": {
+      "title": "Computing",
+      "number": 8,
+      "description": "Exercises on number systems and binary arithmetic."
+    },
+    "question": {
+      "title": "Binary Numbers",
+      "number": 0,
+      "guidance": "",
+      "content": "",
+      "estimatedTime": "5-10 minutes",
+      "parts": [
+        {
+          "position": 0,
+          "content": "Convert the following numbers to decimal: $11010_2$, $110101_2$.",
+          "answerContent": "$11010_2 = 26_{10}$\\n\\n$110101_2 = 53_{10}$&#x20;\\n\\n&#x20;&#x20;",
+          "workedSolutionSections": [
+            {
+              "position": 0,
+              "title": "",
+              "content": "$11010_2 = (0\\\\times1) + (1\\\\times2) + (0\\\\times4) + (1\\\\times8) + (1\\\\times16) = \\\\boxed{26_{10}}$&#x20;\\n\\n***\\n\\n$110101_2 = (1\\\\times1) + (0\\\\times2) + (1\\\\times4) + (0\\\\times8) + (1\\\\times16) + (1\\\\times32) = \\\\boxed{53_{10}}$&#x20;\\n\\n&#x20;"
+            }
+          ],
+          "structuredTutorialSections": [],
+          "responseAreas": [
+            {
+              "position": 0,
+              "responseType": "NUMBER",
+              "answer": 26,
+              "preResponseText": "$11010_2 =$"
+            },
+            {
+              "position": 1,
+              "responseType": "NUMBER",
+              "answer": 53,
+              "preResponseText": "$110101_2$"
+            }
+          ]
+        },
+        {
+          "position": 1,
+          "content": "Convert the following numbers to binary: $101_{10}$, $16_{10}$.",
+          "answerContent": "$101_{10} = 1100101_2$&#x20;\\n\\n$16_{10} = 10000_2$&#x20;\\n\\n&#x20;&#x20;",
+          "workedSolutionSections": [
+            {
+              "position": 0,
+              "title": "",
+              "content": ""
+            },
+            {
+              "position": 0,
+              "title": "By repeated division",
+              "content": "$101_{10}$ can be divided by 2 to give a quotient of 50 and a remainder of 1. The remainder is the first binary digit: $1_2$.\\n\\n***\\n\\n$50_{10}$ divided by 2 gives a quotient of 25 and a remainder of zero, so the second binary digit is 0: $01_2$.\\n\\n***\\n\\n$25_{10}$ divided by 2 gives a quotient of 12 and a remainder of 1, so the third binary digit is 1: $101_2$\\n\\n***\\n\\n$12_{10}$ divided by 2 gives a quotient of 6 and a remainder of 0, so the fourth binary digit is 0: $0101_2$\\n\\n***\\n\\n$6_{10}$ divided by 2 gives a quotient of 3 and a remainder of 0, so the fifth binary digit is 0: $00101_2$\\n\\n***\\n\\n$3_{10}$ divided by 2 gives a quotient of 1 and a remainder of 1, so the sixth binary digit is 1: $100101_2$\\n\\n***\\n\\n$1_{10}$ divided by 2 gives a quotient of zero and a remainder of 1, so the seventh binary digit is 1: $\\\\boxed{1100101_2}$.&#x20;\\n\\n&#x20;&#x20;"
+            },
+            {
+              "position": 1,
+              "title": "By inspection",
+              "content": "$101_{10}$ can be written in terms of powers of 2 as $101 = 64 + 32 + 4 + 1 = 2^6 + 2^5 + 2^2 + 2^0$, so its binary representation is $\\\\boxed{1100101_2}$.\\n\\n***\\n\\n$16_{10}$ is already a power of 2 ($2^4$), so its binary representation will have 1 bit set: $\\\\boxed{10000_2}$.&#x20;\\n\\n&#x20; &#x20;"
+            }
+          ],
+          "structuredTutorialSections": [],
+          "responseAreas": [
+            {
+              "position": 0,
+              "responseType": "NUMBER",
+              "answer": 1100101,
+              "preResponseText": "$101_{10} = $"
+            },
+            {
+              "position": 1,
+              "responseType": "NUMBER",
+              "answer": 10000,
+              "preResponseText": "$16_{10} = $"
+            }
+          ]
+        }
+      ]
+    }
+  }
+}
diff --git a/tests/example_inputs/example_input_2.json b/tests/example_inputs/example_input_2.json
new file mode 100644
index 0000000..a0dd5db
--- /dev/null
+++ b/tests/example_inputs/example_input_2.json
@@ -0,0 +1,97 @@
+{
+  "conversationId": "example-conversation-id-2",
+  "messages": [{ "role": "USER", "content": "help" }],
+  "user": {
+    "type": "LEARNER",
+    "preference": {
+      "conversationalStyle": "student prefers direct and short responses"
+    },
+    "taskProgress": {
+      "timeSpentOnQuestion": "less than one minute",
+      "accessStatus": "",
+      "markedDone": "This question is still being worked on.",
+      "currentPart": {
+        "position": 0,
+        "timeSpentOnPart": "less than one minute",
+        "markedDone": "This part is marked done.",
+        "responseAreas": [
+          {
+            "responseType": "NUMERIC_UNITS",
+            "totalSubmissions": 3,
+            "wrongSubmissions": 2,
+            "latestSubmission": {
+              "submission": "939.7 J",
+              "feedback": "Correct",
+              "answer": "\"939.7 J\""
+            }
+          }
+        ]
+      }
+    }
+  },
+  "context": {
+    "summary": "noting was discussed yet",
+    "set": { "title": "Physics", "number": 7, "description": "" },
+    "question": {
+      "title": "Work done",
+      "number": 0,
+      "guidance": "See section 1.11.",
+      "content": "**(L3)** A father pulls a child on a sled with a rope that has a constant tension of $F =100\\\\text{N}$ and makes an angle of $20^\\\\circ$ to the horizontal. Calculate the work he performs after pulling the child for $10\\\\text{m}$.&#x20;\\n\\n![](image)\\n",
+      "estimatedTime": "10 minutes",
+      "parts": [
+        {
+          "position": 0,
+          "content": "Calculate the work he performs after pulling the child for $10\\\\text{m}$.&#x20;\\n",
+          "answerContent": "$W = 939.7\\\\text{J}$\\n",
+          "workedSolutionSections": [
+            {
+              "position": 0,
+              "title": "",
+              "content": "We require the force component parallel to the direction of travel, and so $W = \\\\vec{F}\\\\cdot\\\\vec{d}$ (the projection of $\\\\vec{F}$ onto $\\\\vec{d}$) (**section 1.11**). Hence applying the dot product definition to evaluate this (**section 1.8**):\\n\\n***\\n\\n$$\\n\\\\begin{aligned}\\nW & = \\\\vec{F}\\\\cdot \\\\vec{d} \\\\\\\\\\n& = 100\\\\cdot10\\\\cos\\\\theta \\\\\\\\\\n& = 939.7\\\\text{J} \\n\\\\end{aligned}\\n$$\\n"
+            }
+          ],
+          "structuredTutorialSections": [
+            {
+              "position": 0,
+              "title": null,
+              "content": "Can you express the work done as a dot product (**section 1.11**)?\\n\\n\\n***\\n\\nEvaluate the dot product using the dot product definition (**section 1.8**).\\n"
+            }
+          ],
+          "responseAreas": [
+            {
+              "position": 0,
+              "responseType": "NUMERIC_UNITS",
+              "answer": "939.7 J",
+              "preResponseText": "$W=$"
+            }
+          ]
+        },
+        {
+          "position": 1,
+          "content": "Why does he not pull at a shallower angle (closer to the horizontal)?\\n",
+          "answerContent": "",
+          "workedSolutionSections": [
+            {
+              "position": 0,
+              "title": "",
+              "content": "Pulling at a shallower angle (closer to horizontal) would require less work or allow him to pull less strongly, but it reduces the reaction force that he feels too, which reduces his maximum pulling force, and also allow him to gain more leverage as he pulls (he can lean forward and use his weight to pull whilst pivoted at his feet).&#x20;\\n"
+            }
+          ],
+          "structuredTutorialSections": [],
+          "responseAreas": [
+            {
+              "position": 0,
+              "responseType": "MULTIPLE_CHOICE",
+              "answer": [
+                "[true] ... reduces the reaction force acting on him, hence reducing the maximum pulling force.",
+                "[false] ... is a more awkward angle for the man to maintain, causing him to strain his muscles.",
+                "[false] ... increases the friction between the sledge and the ground, making it more difficult to move."
+              ],
+              "preResponseText": "Pulling at a shallower angle..."
+            }
+          ]
+        }
+      ]
+    }
+  }
+}
diff --git a/tests/example_inputs/example_input_3.json b/tests/example_inputs/example_input_3.json
new file mode 100644
index 0000000..d62d794
--- /dev/null
+++ b/tests/example_inputs/example_input_3.json
@@ -0,0 +1,244 @@
+{
+  "conversationId": "example-conversation-id-3",
+  "messages": [
+    { "role": "USER", "content": "hi, I need help with this question" },
+    {
+      "role": "ASSISTANT",
+      "content": "Of course! I can see you're working on vector arithmetic. What would you like help with?"
+    },
+    { "role": "USER", "content": "I'm confused about how to add vectors" },
+    {
+      "role": "ASSISTANT",
+      "content": "Vector addition is straightforward — you add the corresponding components. For vectors a and b, you add the first elements together, then the second, then the third."
+    },
+    { "role": "USER", "content": "ok so for a+b I get [-1, 5, -1]?" },
+    {
+      "role": "ASSISTANT",
+      "content": "Exactly right! a+b = [-1, 5, -1]. Well done."
+    },
+    {
+      "role": "USER",
+      "content": "and b+a should be the same because of commutativity?"
+    },
+    {
+      "role": "ASSISTANT",
+      "content": "Correct — vector addition is commutative, so b+a = a+b = [-1, 5, -1]."
+    },
+    {
+      "role": "USER",
+      "content": "I submitted both and got the first one right but the second one wrong, I put [-1, 5, 1] by mistake"
+    },
+    {
+      "role": "ASSISTANT",
+      "content": "That's just a sign error on the last component — easy to do. The last element of a is 3 and of b is -4, so their sum is -1, not 1."
+    },
+    {
+      "role": "USER",
+      "content": "got it. now I'm stuck on part b, what does 3c mean?"
+    },
+    {
+      "role": "ASSISTANT",
+      "content": "Scalar multiplication — you multiply every component of c by 3. So 3c = [3×0, 3×4, 3×(-1)] = [0, 12, -3]."
+    },
+    { "role": "USER", "content": "I am still confused" }
+  ],
+  "user": {
+    "type": "LEARNER",
+    "preference": {
+      "conversationalStyle": "The student prefers concise answers with worked examples."
+    },
+    "taskProgress": {
+      "timeSpentOnQuestion": "13 minutes",
+      "accessStatus": "too much time spent on this question.",
+      "markedDone": "Part (a) was marked done.",
+      "currentPart": {
+        "position": 0,
+        "timeSpentOnPart": "10 minutes",
+        "markedDone": "This part is not marked done.",
+        "responseAreas": [
+          {
+            "responseType": "MATRIX",
+            "totalSubmissions": 1,
+            "wrongSubmissions": 0,
+            "latestSubmission": {
+              "submission": [["-1"], ["5"], ["-1"]],
+              "feedback": "Correct",
+              "answer": [["-1"], ["5"], ["-1"]]
+            }
+          },
+          {
+            "responseType": "MATRIX",
+            "totalSubmissions": 1,
+            "wrongSubmissions": 1,
+            "latestSubmission": {
+              "submission": [["-1"], ["5"], ["1"]],
+              "feedback": "Missing negative sign in the last element.",
+              "answer": [["-1"], ["5"], ["-1"]]
+            }
+          }
+        ]
+      }
+    }
+  },
+  "context": {
+    "summary": "",
+    "set": {
+      "title": "Linear Algebra",
+      "number": 1,
+      "description": "Vectors, matrices, and linear transformations."
+    },
+    "question": {
+      "title": "Vector Arithmetics",
+      "number": 1,
+      "guidance": null,
+      "content": "$$\n\\vec{a}=\\begin{bmatrix}1 \\\\ 2\\\\ 3\\end{bmatrix} \\quad \\vec{b}=\\begin{bmatrix}-2 \\\\ 3\\\\ -4\\end{bmatrix} \\quad \\vec{c}=\\begin{bmatrix} 0 \\\\\\ 4\\\\ -1\\end{bmatrix}\n$$",
+      "estimatedTime": "1-3 minutes",
+      "parts": [
+        {
+          "position": 0,
+          "content": "",
+          "answerContent": "",
+          "workedSolutionSections": [
+            {
+              "position": 1,
+              "title": "",
+              "content": "$$\n\\vec{a}+\\vec{b} = \\begin{bmatrix}-1 \\\\ 5\\\\ -1\\end{bmatrix}\n$$\n\nSince vector addition is commutative, $\\vec{b}+\\vec{a}=\\vec{a}+\\vec{b}$.\n"
+            }
+          ],
+          "structuredTutorialSections": [],
+          "responseAreas": [
+            {
+              "position": 0,
+              "responseType": "MATRIX",
+              "answer": [["-1"], ["5"], ["-1"]],
+              "preResponseText": "$\\vec{a}+\\vec{b}=$"
+            },
+            {
+              "position": 1,
+              "responseType": "MATRIX",
+              "answer": [["-1"], ["5"], ["-1"]],
+              "preResponseText": "$\\vec{b}+\\vec{a}=$"
+            }
+          ]
+        },
+        {
+          "position": 1,
+          "content": "",
+          "answerContent": "",
+          "workedSolutionSections": [
+            {
+              "position": 1,
+              "title": "$3\\vec{c}$",
+              "content": "$$\n3\\vec{c} = \\begin{bmatrix}0 \\\\ 12\\\\ -3\\end{bmatrix}\n$$"
+            },
+            {
+              "position": 2,
+              "title": "$-\\vec{a}$",
+              "content": "$$\n-\\vec{a} = \\begin{bmatrix}-1 \\\\ -2\\\\ -3\\end{bmatrix}\n$$"
+            },
+            {
+              "position": 3,
+              "title": "$\\frac{\\vec{b}}{2}$",
+              "content": "$$\n\\frac{\\vec{b}}{2} = \\begin{bmatrix}-1 \\\\ 1.5\\\\ -2\\end{bmatrix}\n$$"
+            }
+          ],
+          "structuredTutorialSections": [],
+          "responseAreas": [
+            {
+              "position": 0,
+              "responseType": "MATRIX",
+              "answer": [["0"], ["12"], ["-3"]],
+              "preResponseText": "$3\\vec{c}=$"
+            },
+            {
+              "position": 1,
+              "responseType": "MATRIX",
+              "answer": [["-1"], ["-2"], ["-3"]],
+              "preResponseText": "$-\\vec{a}=$"
+            },
+            {
+              "position": 2,
+              "responseType": "MATRIX",
+              "answer": [["-1"], ["1.5"], ["-2"]],
+              "preResponseText": "$\\frac{\\vec{b}}{2}=$"
+            }
+          ]
+        },
+        {
+          "position": 2,
+          "content": "",
+          "answerContent": "",
+          "workedSolutionSections": [
+            {
+              "position": 1,
+              "title": "$3\\vec{a}-3\\vec{c}$",
+              "content": "$$\n3\\vec{a}-3\\vec{c} = \\begin{bmatrix}3 \\\\ -6\\\\ 12\\end{bmatrix}\n$$"
+            },
+            {
+              "position": 2,
+              "title": "$3(\\vec{a}-\\vec{c})$",
+              "content": "$$\n3(\\vec{a}-\\vec{c}) = \\begin{bmatrix}3 \\\\ -6\\\\ 12\\end{bmatrix}\n$$"
+            }
+          ],
+          "structuredTutorialSections": [],
+          "responseAreas": [
+            {
+              "position": 0,
+              "responseType": "MATRIX",
+              "answer": [["3"], ["-6"], ["12"]],
+              "preResponseText": "$3\\vec{a}-3\\vec{c}=$"
+            },
+            {
+              "position": 1,
+              "responseType": "MATRIX",
+              "answer": [["3"], ["-6"], ["12"]],
+              "preResponseText": "$3(\\vec{a}-\\vec{c})=$"
+            }
+          ]
+        },
+        {
+          "position": 3,
+          "content": "",
+          "answerContent": "",
+          "workedSolutionSections": [
+            {
+              "position": 1,
+              "title": "",
+              "content": "$$\n-5(\\vec{a}+\\vec{c})+\\vec{b} = \\begin{bmatrix}-7 \\\\ -27\\\\ -14\\end{bmatrix}\n$$"
+            }
+          ],
+          "structuredTutorialSections": [],
+          "responseAreas": [
+            {
+              "position": 0,
+              "responseType": "MATRIX",
+              "answer": [["-7"], ["-27"], ["-14"]],
+              "preResponseText": "$-5(\\vec{a}+\\vec{c})+\\vec{b}=$"
+            }
+          ]
+        },
+        {
+          "position": 4,
+          "content": "",
+          "answerContent": "",
+          "workedSolutionSections": [
+            {
+              "position": 1,
+              "title": "",
+              "content": "$$\n3\\vec{a}-2\\vec{c}+3\\vec{b} = \\begin{bmatrix}-3 \\\\ 7\\\\ -1\\end{bmatrix}\n$$"
+            }
+          ],
+          "structuredTutorialSections": [],
+          "responseAreas": [
+            {
+              "position": 0,
+              "responseType": "MATRIX",
+              "answer": [["-3"], ["7"], ["-1"]],
+              "preResponseText": "$3\\vec{a}-2\\vec{c}+3\\vec{b}=$"
+            }
+          ]
+        }
+      ]
+    }
+  }
+}
diff --git a/tests/manual_agent_requests.py b/tests/manual_agent_requests.py
index 9d4c729..57023b5 100644
--- a/tests/manual_agent_requests.py
+++ b/tests/manual_agent_requests.py
@@ -9,7 +9,7 @@
 url = "http://localhost:8080/2015-03-31/functions/function/invocations"
 
 # File path for the input text
-path = "src/agent/utils/example_inputs/"
+path = "tests/example_inputs/"
 input_file = path + "example_input_1.json"
 
 # Step 1: Read the input file
diff --git a/tests/manual_agent_run.py b/tests/manual_agent_run.py
index b69f609..43a500b 100644
--- a/tests/manual_agent_run.py
+++ b/tests/manual_agent_run.py
@@ -1,46 +1,38 @@
 """
- Conversation turn-based Testbench of the agent's performance.
- Select an example input file and write your query. Then run the agent to get the response.
+Conversation turn-based Testbench of the agent's performance.
+Select an example input file and write your query. Then run the agent to get the response.
 """
 
 import json
+from lf_toolkit.chat import ChatRequest
 from src.module import chat_module
 
 # File path for the input text
-path = "src/agent/utils/example_inputs/"
+path = "tests/example_inputs/"
 input_file = path + "example_input_1.json"
 
 # Step 1: Read the input file
 with open(input_file, "r") as file:
     raw_text = file.read()
-    
-# Step 5: Parse into JSON
+
+# Step 2: Parse into JSON
 try:
     parsed_json = json.loads(raw_text)
 
-    """
-      STEP 2: Extract the parameters from the JSON
-    """
     # NOTE: #### This is the testing message #####
-    message = "Hi, how do I solve this problem?" 
+    message = "Hi, how do I solve this problem?"
     # NOTE: ########################################
 
-    # In the JSON, replace "mock" in the message and conversation history with the testing message
-    parsed_json["message"] = message
-    parsed_json["params"]["conversation_history"][-1]["content"] = message
+    # Replace the last message content with the testing message
+    parsed_json["messages"][-1]["content"] = message
+
+    # Step 3: Validate into a ChatRequest
+    request = ChatRequest.model_validate(parsed_json)
 
-    params = parsed_json["params"]
+    # Step 4: Call the chat module to get a response
+    response = chat_module(request)
 
-    """
-      STEP 3: Call the chat module to get a response to the user's message
-    """
-    response = chat_module(message, params)
-    
-    print(json.dumps(response, indent=4))
-    
+    print(json.dumps(json.loads(response.model_dump_json()), indent=4))
 
 except json.JSONDecodeError as e:
     print("Error decoding JSON:", e)
-
-
-
diff --git a/tests/test_example_inputs.py b/tests/test_example_inputs.py
new file mode 100644
index 0000000..3373fd1
--- /dev/null
+++ b/tests/test_example_inputs.py
@@ -0,0 +1,41 @@
+import unittest
+import json
+import os
+from index import handler
+from tests.utils import assert_valid_chat_request, assert_valid_chat_response
+
+EXAMPLE_INPUTS_DIR = "tests/example_inputs"
+
+
+class TestExampleInputs(unittest.TestCase):
+    """
+    End-to-end tests for each example input in tests/example_inputs/.
+    Each file must be a valid muEd ChatRequest and produce a valid muEd ChatResponse.
+    """
+
+    def _test(self, filename: str):
+        with open(os.path.join(EXAMPLE_INPUTS_DIR, filename)) as f:
+            payload = json.load(f)
+        assert_valid_chat_request(self, payload)
+        result = handler({"body": json.dumps(payload)}, None)
+        assert_valid_chat_response(self, result)
+        return payload, json.loads(result["body"])
+
+    def test_example_input_0_simple(self):
+        self._test("example_input_0.json")
+
+    def test_example_input_1(self):
+        self._test("example_input_1.json")
+
+    def test_example_input_2_metadata_persists(self):
+        """Example input 2 has an existing summary and conversationalStyle — verifies they persist in the response."""
+        payload, result_body = self._test("example_input_2.json")
+        self.assertEqual(result_body["metadata"]["summary"], payload["context"]["summary"])
+        self.assertEqual(result_body["metadata"]["conversationalStyle"], payload["user"]["preference"]["conversationalStyle"])
+
+    def test_example_input_3_metadata_updates(self):
+        """Example input 3 has 13 messages — enough to trigger summarisation.
+        Verifies that summary and conversationalStyle are populated in the response."""
+        _, result_body = self._test("example_input_3.json")
+        self.assertGreater(len(result_body["metadata"]["summary"]), 0, "summary should be populated after summarisation")
+        self.assertGreater(len(result_body["metadata"]["conversationalStyle"]), 0, "conversationalStyle should be populated after summarisation")
diff --git a/tests/test_index.py b/tests/test_index.py
index 822049d..b6047ec 100644
--- a/tests/test_index.py
+++ b/tests/test_index.py
@@ -1,62 +1,30 @@
 import unittest
 import json
 from index import handler
+from tests.utils import assert_valid_chat_request, assert_valid_chat_response
+
+
+def make_event(body: dict) -> dict:
+    return {"body": json.dumps(body)}
+
+
+BASE_BODY = {
+    "messages": [{"role": "USER", "content": "Hello, World"}],
+    "conversationId": "1234Test",
+}
+
 
 class TestChatIndexFunction(unittest.TestCase):
-    """
-    TestCase Class used to test the algorithm.
-    ---
-    Tests are used here to check that the algorithm written
-    is working as it should.
-
-    It's best practise to write these tests first to get a
-    kind of 'specification' for how your algorithm should
-    work, and you should run these tests before committing
-    your code to AWS.
-
-    Read the docs on how to use unittest here:
-    https://docs.python.org/3/library/unittest.html
-
-    Use module() to check your algorithm works
-    as it should.
-
-    The expected input of the hander is a JsonType.
-    """
-
-    def test_missing_argument(self):
-        arguments = ["message", "params"]
-
-        for arg in arguments:
-            event = {
-                "message": "Hello, World",
-                "params": {"conversation_id": "1234Test", "conversation_history": [{"type": "user", "content": "Hello, World"}]}
-            }
-            event.pop(arg)
-            event = {"body":json.dumps(event)}
-
-            result = handler(event, None)
-
-            self.assertEqual(result.get("statusCode"), 400)
-    
-    def test_correct_arguments(self):
-        event = {
-            "message": "Hello, World",
-            "params": {"conversation_id": "1234Test", "conversation_history": [{"type": "user", "content": "Hello, World"}]}
-        }
-        event = {"body":json.dumps(event)}
-
-        result = handler(event, None)
-
-        self.assertEqual(result.get("statusCode"), 200)
-
-    def test_correct_response(self):
-        event = {
-            "message": "Hello, World",
-            "params": {"conversation_id": "1234Test", "conversation_history": [{"type": "user", "content": "Hello, World"}]}
-        }
-        event = {"body":json.dumps(event)}
-
-        result = handler(event, None)
-
-        self.assertEqual(result.get("statusCode"), 200)
-        
\ No newline at end of file
+
+    def test_missing_messages(self):
+        body = {k: v for k, v in BASE_BODY.items() if k != "messages"}
+        result = handler(make_event(body), None)
+        self.assertEqual(result.get("statusCode"), 400)
+
+    def test_invalid_json_body(self):
+        result = handler({"body": "not valid json"}, None)
+        self.assertEqual(result.get("statusCode"), 400)
+
+    def test_response_format(self):
+        assert_valid_chat_request(self, BASE_BODY)
+        assert_valid_chat_response(self, handler(make_event(BASE_BODY), None))
diff --git a/tests/test_module.py b/tests/test_module.py
index 26c96de..43be647 100755
--- a/tests/test_module.py
+++ b/tests/test_module.py
@@ -1,84 +1,19 @@
 import unittest
-from lf_toolkit.chat.result import ChatResult as Result
-from lf_toolkit.chat.params import ChatParams as Params
+from lf_toolkit.chat import ChatRequest
 from src.module import chat_module
+from tests.utils import assert_valid_chat_response
 
-class TestChatModuleFunction(unittest.TestCase):
-    """
-    TestCase Class used to test the algorithm.
-    ---
-    Tests are used here to check that the algorithm written
-    is working as it should.
-
-    It's best practise to write these tests first to get a
-    kind of 'specification' for how your algorithm should
-    work, and you should run these tests before committing
-    your code to AWS.
-
-    Read the docs on how to use unittest here:
-    https://docs.python.org/3/library/unittest.html
-
-    Use module() to check your algorithm works
-    as it should.
-    """
-
-    def test_missing_parameters(self):
-        # Checking state for missing parameters on default agent
-        response = "Hello, World"
-        expected_params = Params(include_test_data=True, conversation_history=['{ "type": "user", "content": response }'], \
-                                    summary="", conversational_style="", \
-                                    question_response_details={}, conversation_id="1234Test")
-
-        for p in expected_params:
-            params = expected_params.copy()
-            # except for the special parameters
-            if p not in ["include_test_data", "conversation_id", "conversation_history"]:
-                params.pop(p)
-
-                result = chat_module(response, params)
-
-                self.assertIsNotNone(result)
-                self.assertEqual("error" in result, False)
-            elif p == "include_test_data":
-                params.pop(p)
 
-                result = chat_module(response, params)
+def make_request(**kwargs):
+    defaults = {
+        "messages": [{"role": "USER", "content": "Hello, World"}],
+        "conversationId": "1234Test",
+    }
+    defaults.update(kwargs)
+    return ChatRequest.model_validate(defaults)
 
-                # check if result has nothing except for the chatbot_response
-                self.assertIsNotNone(result.get("chatbot_response"))
-                self.assertEqual(len(result), 1)
-            elif p == "conversation_id":
-                params.pop(p)
-                
-                with self.assertRaises(Exception) as cm:
-                    chat_module(response, params)
 
-                self.assertTrue("Internal Error" in str(cm.exception))
-                self.assertTrue("conversation id" in str(cm.exception))
-            elif p == "conversation_history":
-                params.pop(p)
-
-                with self.assertRaises(Exception) as cm:
-                    chat_module(response, params)
-
-                self.assertTrue("Internal Error" in str(cm.exception))
-                self.assertTrue("conversation history" in str(cm.exception))
-
-    def test_agent_output(self):
-        # Checking the output of the agent
-        response = "Hello, World"
-        params = Params(conversation_id="1234Test", conversation_history=['{ "type": "user", "content": response }'])
-
-        result = chat_module(response, params)
-
-        self.assertIsNotNone(result.get("chatbot_response"))
-
-    def test_processing_time_calc(self):
-        # Checking the processing time calculation
-        response = "Hello, World"
-        params = Params(include_test_data=True, conversation_id="1234Test", conversation_history=['{ "type": "user", "content": response }'])
-
-        result = chat_module(response, params)
+class TestChatModuleFunction(unittest.TestCase):
 
-        self.assertIsNotNone(result.get("processing_time"))
-        self.assertGreaterEqual(result.get("processing_time"), 0)
\ No newline at end of file
+    def test_response_format(self):
+        assert_valid_chat_response(self, chat_module(make_request()))
diff --git a/tests/utils.py b/tests/utils.py
new file mode 100644
index 0000000..113c36e
--- /dev/null
+++ b/tests/utils.py
@@ -0,0 +1,35 @@
+import json
+import unittest
+from lf_toolkit.chat import ChatRequest, ChatResponse
+
+
+def assert_valid_chat_request(test: unittest.TestCase, payload: dict):
+    """Assert a dict is a valid muEd ChatRequest (at least one message)."""
+    try:
+        request = ChatRequest.model_validate(payload)
+    except Exception as e:
+        test.fail(f"Payload is not a valid ChatRequest: {e}")
+    test.assertGreater(len(request.messages), 0, "messages must not be empty")
+
+
+def assert_valid_chat_response(test: unittest.TestCase, result):
+    """
+    Assert a result matches the expected muEd ChatResponse format.
+    Accepts either a ChatResponse object or a Lambda handler result dict.
+    """
+    if isinstance(result, ChatResponse):
+        body = json.loads(result.model_dump_json())
+    else:
+        test.assertEqual(result.get("statusCode"), 200)
+        body = json.loads(result["body"])
+
+    output = body.get("output", {})
+    test.assertEqual(output.get("role"), "ASSISTANT")
+    test.assertIsInstance(output.get("content"), str)
+    test.assertGreater(len(output["content"]), 0)
+
+    metadata = body.get("metadata", {})
+    test.assertIsInstance(metadata.get("summary"), str)
+    test.assertIsInstance(metadata.get("conversationalStyle"), str)
+    test.assertIsInstance(metadata.get("processingTimeMs"), int)
+    test.assertGreaterEqual(metadata["processingTimeMs"], 0)