diff --git a/.dockerignore b/.dockerignore
index 4093c0a..3d422e8 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -147,11 +147,4 @@ data/
 reports/
 
 # Synthetic data conversations
-src/agents/utils/example_inputs/
-src/agents/utils/synthetic_conversations/
-src/agents/utils/synthetic_conversation_generation.py
-src/agents/utils/testbench_prompts.py
-src/agents/utils/langgraph_viz.py
-
-# development agents
-src/agents/student_agent/
\ No newline at end of file
+src/agents/utils/example_inputs/
\ No newline at end of file
diff --git a/.github/workflows/dev.yml b/.github/workflows/dev.yml
index 1191a86..1236ba3 100644
--- a/.github/workflows/dev.yml
+++ b/.github/workflows/dev.yml
@@ -50,6 +50,7 @@ jobs:
         if: always()
         run: |
           source .venv/bin/activate
+          export PYTHONPATH=$PYTHONPATH:.
           pytest --junit-xml=./reports/pytest.xml --tb=auto -v
 
       - name: Upload test results
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 8400ca3..1da0493 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -50,6 +50,7 @@ jobs:
         if: always()
         run: |
           source .venv/bin/activate
+          export PYTHONPATH=$PYTHONPATH:.
           pytest --junit-xml=./reports/pytest.xml --tb=auto -v
 
       - name: Upload test results
diff --git a/.gitignore b/.gitignore
index 4b52234..11f861b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -50,6 +50,7 @@ coverage.xml
 *.py,cover
 .hypothesis/
 .pytest_cache/
+reports/
 
 # Translations
 *.mo
diff --git a/Dockerfile b/Dockerfile
index 9150687..38276cc 100755
--- a/Dockerfile
+++ b/Dockerfile
@@ -25,7 +25,7 @@ COPY src ./src
 
 COPY index.py .
 
-COPY index_test.py .
+COPY tests ./tests
 
 # Set the Lambda function handler
 CMD ["index.handler"]
\ No newline at end of file
diff --git a/README.md b/README.md
index bf84b60..ef662b0 100755
--- a/README.md
+++ b/README.md
@@ -43,11 +43,11 @@ In GitHub, choose Use this template > Create a new repository in the repository
 
 Choose the owner, and pick a name for the new repository.
 
-> [!IMPORTANT] If you want to deploy the evaluation function to Lambda Feedback, make sure to choose the Lambda Feedback organization as the owner.
+> [!IMPORTANT] If you want to deploy the chat function to Lambda Feedback, make sure to choose the `Lambda Feedback` organization as the owner.
 
-Set the visibility to Public or Private.
+Set the visibility to `Public` or `Private`.
 
-> [!IMPORTANT] If you want to use GitHub deployment protection rules, make sure to set the visibility to Public.
+> [!IMPORTANT] If you want to use GitHub deployment protection rules, make sure to set the visibility to `Public`.
 
 Click on Create repository.
 
@@ -78,9 +78,9 @@ Also, don't forget to update or delete the Quickstart chapter from the `README.m
 
 ## Development
 
-You can create your own invocation to your own agents hosted anywhere. Copy or update the `base_agent` from `src/agents/` and edit it to match your LLM agent requirements. Import the new invocation in the `module.py` file.
+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/agents/llm_factory.py`.
+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`.
 
 ### Prerequisites
 
@@ -90,23 +90,37 @@ You agent can be based on an LLM hosted anywhere, you have available currently O
 ### Repository Structure
 
 ```bash
-.github/workflows/
-    dev.yml                           # deploys the DEV function to Lambda Feedback
-    main.yml                          # deploys the STAGING function to Lambda Feedback
-    test-report.yml                   # gathers Pytest Report of function tests
-
-docs/               # docs for devs and users
-
-src/module.py       # chat_module function implementation
-src/module_test.py  # chat_module function tests
-src/agents/         # find all agents developed for the chat functionality
-src/agents/utils/test_prompts.py      # allows testing of any LLM agent on a couple of example inputs containing Lambda Feedback Questions and synthetic student conversations
+.
+├── .github/workflows/
+│   ├── dev.yml                           # deploys the DEV function to Lambda Feedback
+│   ├── main.yml                          # deploys the STAGING and PROD functions to Lambda Feedback
+│   └── test-report.yml                   # gathers Pytest Report of function tests
+├── 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                         
+└── tests/                                # contains all tests for the chat function
+    ├── 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
+    ├── test_index.py                     # pytests
+    └── test_module.py                    # pytests
 ```
 
 
 ## Testing the Chat Function
 
-To test your function, you can either call the code directly through a python script. Or you can build the respective chat function docker container locally and call it through an API request. Below you can find details on those processes.
+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.
+
+### Run Unit Tests
+
+You can run the unit tests using `pytest`.
+
+```bash
+pytest
+```
 
 ### Run the Chat Script
 
@@ -116,9 +130,9 @@ You can run the Python function itself. Make sure to have a main function in eit
 python src/module.py
 ```
 
-You can also use the `testbench_agents.py` script to test the agents with example inputs from Lambda Feedback questions and synthetic conversations.
+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 src/agents/utils/testbench_agents.py
+python tests/manual_agent_run.py
 ```
 
 ### Calling the Docker Image Locally
@@ -156,7 +170,7 @@ curl --location 'http://localhost:8080/2015-03-31/functions/function/invocations
 #### Call Docker Container
 ##### A. Call Docker with Python Requests
 
-In the `src/agents/utils` folder you can find the `requests_testscript.py` script that calls the POST URL of the running docker container. It reads any kind of input files with the expected schema. You can use this to test your curl calls of the chatbot.
+In the `tests/` folder you can find the `manual_agent_requests.py` script that calls the POST URL of the running docker container. It reads any kind of input files with the expected schema. You can use this to test your curl calls of the chatbot.
 
 ##### B. Call Docker Container through API request
 
@@ -183,7 +197,6 @@ Body with optional Params:
         "conversational_style":" ",
         "question_response_details": "",
         "include_test_data": true,
-        "agent_type": {agent_name}
     }
 }
 ```
diff --git a/docs/dev.md b/docs/dev.md
index a528199..81d1433 100644
--- a/docs/dev.md
+++ b/docs/dev.md
@@ -12,7 +12,15 @@
 
 ## Testing the Chat Function
 
-To test your function, you can either call the code directly through a python script. Or you can build the respective chat function docker container locally and call it through an API request. Below you can find details on those processes.
+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.
+
+### Run Unit Tests
+
+You can run the unit tests using `pytest`.
+
+```bash
+pytest
+```
 
 ### Run the Chat Script
 
@@ -22,9 +30,9 @@ You can run the Python function itself. Make sure to have a main function in eit
 python src/module.py
 ```
 
-You can also use the `testbench_agents.py` script to test the agents with example inputs from Lambda Feedback questions and synthetic conversations.
+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 src/agents/utils/testbench_agents.py
+python tests/manual_agent_run.py
 ```
 
 ### Calling the Docker Image Locally
@@ -62,7 +70,7 @@ curl --location 'http://localhost:8080/2015-03-31/functions/function/invocations
 #### Call Docker Container
 ##### A. Call Docker with Python Requests
 
-In the `src/agents/utils` folder you can find the `requests_testscript.py` script that calls the POST URL of the running docker container. It reads any kind of input files with the expected schema. You can use this to test your curl calls of the chatbot.
+In the `tests/` folder you can find the `manual_agent_requests.py` script that calls the POST URL of the running docker container. It reads any kind of input files with the expected schema. You can use this to test your curl calls of the chatbot.
 
 ##### B. Call Docker Container through API request
 
@@ -89,7 +97,6 @@ Body with optional Params:
         "conversational_style":" ",
         "question_response_details": "",
         "include_test_data": true,
-        "agent_type": {agent_name}
     }
 }
 ```
diff --git a/index.py b/index.py
index 62f4c47..0ad738d 100644
--- a/index.py
+++ b/index.py
@@ -1,10 +1,6 @@
 import json
-try:
-    from .src.module import chat_module
-    from .src.agents.utils.types import JsonType
-except ImportError:
-    from src.module import chat_module
-    from src.agents.utils.types import JsonType
+from src.module import chat_module
+from src.agent.utils.types import JsonType
 
 def handler(event: JsonType, context):
     """
diff --git a/src/__init__.py b/src/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/src/agents/base_agent/base_agent.py b/src/agent/agent.py
similarity index 91%
rename from src/agents/base_agent/base_agent.py
rename to src/agent/agent.py
index 0fb199d..31e23cb 100644
--- a/src/agents/base_agent/base_agent.py
+++ b/src/agent/agent.py
@@ -1,13 +1,7 @@
-try:
-    from ..llm_factory import OpenAILLMs, GoogleAILLMs
-    from .base_prompts import \
-        role_prompt, conv_pref_prompt, update_conv_pref_prompt, summary_prompt, update_summary_prompt, summary_system_prompt
-    from ..utils.types import InvokeAgentResponseType
-except ImportError:
-    from src.agents.llm_factory import OpenAILLMs, GoogleAILLMs
-    from src.agents.base_agent.base_prompts import \
-        role_prompt, conv_pref_prompt, update_conv_pref_prompt, summary_prompt, update_summary_prompt, summary_system_prompt
-    from src.agents.utils.types import InvokeAgentResponseType
+from src.agent.utils.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
@@ -62,7 +56,7 @@ def call_model(self, state: State, config: RunnableConfig) -> str:
         system_message = self.role_prompt
 
         # Adding external student progress and question context details from data queries
-        question_response_details = config["configurable"].get("question_response_details", "")
+        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"
 
@@ -98,8 +92,8 @@ def summarize_conversation(self, state: State, config: RunnableConfig) -> dict:
         """Summarize the conversation."""
 
         summary = state.get("summary", "")
-        previous_summary = config["configurable"].get("summary", "")
-        previous_conversationalStyle = config["configurable"].get("conversational_style", "")
+        previous_summary = config.get("configurable", {}).get("summary", "")
+        previous_conversationalStyle = config.get("configurable", {}).get("conversational_style", "")
         if previous_summary:
             summary = previous_summary
         
diff --git a/src/agents/base_agent/base_prompts.py b/src/agent/prompts.py
similarity index 89%
rename from src/agents/base_agent/base_prompts.py
rename to src/agent/prompts.py
index 683caab..a3577c8 100644
--- a/src/agents/base_agent/base_prompts.py
+++ b/src/agent/prompts.py
@@ -1,8 +1,43 @@
-# NOTE:
-# PROMPTS generated with the help of ChatGPT GPT-4o Nov 2024
-
+# 
+# NOTE: Default prompts generated with the help of ChatGPT GPT-4o Nov 2024
+# 
+# Description of the prompts:
+# 
+# 1. role_prompt: Sets the overall role and behaviour of the chatbot.
+# 
+# 2. summary_prompt: Used to generate a summary of the conversation.
+# 2. update_summary_prompt: Used to update the conversation summary with new messages.
+# 2. summary_system_prompt: Provides context for the chatbot based on the existing summary.
+# 
+# 3. conv_pref_prompt: Used to analyze and extract the student's conversational style and learning preferences.
+# 3. update_conv_pref_prompt: Used to update the conversational style based on new interactions.
+# 
+
+# 1. Role Prompt
 role_prompt = "You are an excellent tutor that aims to provide clear and concise explanations to students. I am the student. Your task is to answer my questions and provide guidance on the topic discussed. Ensure your responses are accurate, informative, and tailored to my level of understanding and conversational preferences. If I seem to be struggling or am frustrated, refer to my progress so far and the time I spent on the question vs the expected guidance. If I 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.\n\n"
 
+# 2. Summary Prompts
+summary_guidelines = """Ensure the summary is:
+
+Concise: Keep the summary brief while including all essential information.
+Structured: Organize the summary into sections such as 'Topics Discussed' and 'Top 3 Key Detailed Ideas'.
+Neutral and Accurate: Avoid adding interpretations or opinions; focus only on the content shared.
+When summarizing: If the conversation is technical, highlight significant concepts, solutions, and terminology. If context involves problem-solving, detail the problem and the steps or solutions provided. If the user asks for creative input, briefly describe the ideas presented.
+Last messages: Include the most recent 5 messages to provide context for the summary.
+
+Provide the summary in a bulleted format for clarity. Avoid redundant details while preserving the core intent of the discussion."""
+
+summary_prompt = f"""Summarize the conversation between a student and a tutor. Your summary should highlight the major topics discussed during the session, followed by a detailed recollection of the last five significant points or ideas. Ensure the summary flows smoothly to maintain the continuity of the discussion.
+
+{summary_guidelines}"""
+
+update_summary_prompt = f"""Update the summary by taking into account the new messages above.
+
+{summary_guidelines}"""
+
+summary_system_prompt = "You are continuing a tutoring session with the student. Background context: {summary}. Use this context to inform your understanding but do not explicitly restate, refer to, or incorporate the details directly in your responses unless the user brings them up. Respond naturally to the user's current input, assuming prior knowledge from the summary."
+
+# 3. Conversational Preference Prompt
 pref_guidelines = """**Guidelines:**
 - Use concise, objective language.
 - Note the student's educational goals, such as understanding foundational concepts, passing an exam, getting top marks, code implementation, hands-on practice, etc.
@@ -57,23 +92,3 @@
 
 {pref_guidelines}
 """
-
-summary_guidelines = """Ensure the summary is:
-
-Concise: Keep the summary brief while including all essential information.
-Structured: Organize the summary into sections such as 'Topics Discussed' and 'Top 3 Key Detailed Ideas'.
-Neutral and Accurate: Avoid adding interpretations or opinions; focus only on the content shared.
-When summarizing: If the conversation is technical, highlight significant concepts, solutions, and terminology. If context involves problem-solving, detail the problem and the steps or solutions provided. If the user asks for creative input, briefly describe the ideas presented.
-Last messages: Include the most recent 5 messages to provide context for the summary.
-
-Provide the summary in a bulleted format for clarity. Avoid redundant details while preserving the core intent of the discussion."""
-
-summary_prompt = f"""Summarize the conversation between a student and a tutor. Your summary should highlight the major topics discussed during the session, followed by a detailed recollection of the last five significant points or ideas. Ensure the summary flows smoothly to maintain the continuity of the discussion.
-
-{summary_guidelines}"""
-
-update_summary_prompt = f"""Update the summary by taking into account the new messages above.
-
-{summary_guidelines}"""
-
-summary_system_prompt = "You are continuing a tutoring session with the student. Background context: {summary}. Use this context to inform your understanding but do not explicitly restate, refer to, or incorporate the details directly in your responses unless the user brings them up. Respond naturally to the user's current input, assuming prior knowledge from the summary."
\ No newline at end of file
diff --git a/src/agents/utils/example_inputs/example_input_1.json b/src/agent/utils/example_inputs/example_input_1.json
similarity index 100%
rename from src/agents/utils/example_inputs/example_input_1.json
rename to src/agent/utils/example_inputs/example_input_1.json
diff --git a/src/agents/utils/example_inputs/example_input_2.json b/src/agent/utils/example_inputs/example_input_2.json
similarity index 100%
rename from src/agents/utils/example_inputs/example_input_2.json
rename to src/agent/utils/example_inputs/example_input_2.json
diff --git a/src/agents/utils/example_inputs/example_input_3.json b/src/agent/utils/example_inputs/example_input_3.json
similarity index 100%
rename from src/agents/utils/example_inputs/example_input_3.json
rename to src/agent/utils/example_inputs/example_input_3.json
diff --git a/src/agents/llm_factory.py b/src/agent/utils/llm_factory.py
similarity index 100%
rename from src/agents/llm_factory.py
rename to src/agent/utils/llm_factory.py
diff --git a/src/agents/utils/parse_json_context_to_prompt.py b/src/agent/utils/parse_json_context_to_prompt.py
similarity index 99%
rename from src/agents/utils/parse_json_context_to_prompt.py
rename to src/agent/utils/parse_json_context_to_prompt.py
index e9f9dcd..a233126 100644
--- a/src/agents/utils/parse_json_context_to_prompt.py
+++ b/src/agent/utils/parse_json_context_to_prompt.py
@@ -3,7 +3,7 @@
 """
 
 from typing import List, Optional, Dict, Any, Union
-from .prompt_context_templates import PromptFormatter
+from src.agent.utils.prompt_context_templates import PromptFormatter
 
 # Definitions questionSubmissionSummary type
 class StudentLatestSubmission:
@@ -150,7 +150,7 @@ def parse_json_to_structured_prompt(
     question_submission_summary: Optional[List[StudentWorkResponseArea]],
     question_information: Optional[QuestionDetails],
     question_access_information: Optional[QuestionAccessInformation]
-) -> Optional[str]:
+) -> str:
     """
     Parse JSON data into a well-structured, LLM-friendly prompt.
     
@@ -322,7 +322,7 @@ def parse_json_to_prompt(
     questionSubmissionSummary: Optional[List[StudentWorkResponseArea]],
     questionInformation: Optional[QuestionDetails],
     questionAccessInformation: Optional[QuestionAccessInformation]
-) -> Optional[str]:
+) -> str:
     """
     Legacy wrapper for backward compatibility.
     Recommended to use parse_json_to_structured_prompt for new code.
diff --git a/src/agents/utils/prompt_context_templates.py b/src/agent/utils/prompt_context_templates.py
similarity index 100%
rename from src/agents/utils/prompt_context_templates.py
rename to src/agent/utils/prompt_context_templates.py
diff --git a/src/agents/utils/types.py b/src/agent/utils/types.py
similarity index 100%
rename from src/agents/utils/types.py
rename to src/agent/utils/types.py
diff --git a/src/agents/__init__.py b/src/agents/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/src/agents/student_agent/student_agent.py b/src/agents/student_agent/student_agent.py
deleted file mode 100644
index 4c5f54e..0000000
--- a/src/agents/student_agent/student_agent.py
+++ /dev/null
@@ -1,145 +0,0 @@
-try:
-    from ..llm_factory import OpenAILLMs
-    from .student_prompts import \
-        base_student_persona, curious_student_persona, contradicting_student_persona, reliant_student_persona, confused_student_persona, unrelated_student_persona, \
-        process_prompt
-    from ..utils.types import InvokeAgentResponseType
-except ImportError:
-    from src.agents.llm_factory import OpenAILLMs
-    from src.agents.student_agent.student_prompts import \
-        base_student_persona, curious_student_persona, contradicting_student_persona, reliant_student_persona, confused_student_persona, unrelated_student_persona, \
-        process_prompt
-    from src.agents.utils.types import InvokeAgentResponseType
-
-from langgraph.graph import StateGraph, START, END
-from langchain_core.messages import SystemMessage, RemoveMessage, HumanMessage, AIMessage
-from langchain_core.runnables.config import RunnableConfig
-from langgraph.graph.message import add_messages
-from typing import Annotated, TypeAlias
-from typing_extensions import TypedDict
-
-"""
-Student agent for synthetic evaluation of the other LLM tutors. This agent is designed to be a student that requires help in the conversation. 
-[LLM workflow with a summarisation, and chat agent that receives an external conversation history].
-
-This agent is designed to:
-- [role_prompt]             role of a student to ask questions on the topic  
-- [student_type]            student's learning profile and comprehension level [many profiles can be chosen from the student_prompts.py]
-"""
-
-ValidMessageTypes: TypeAlias = SystemMessage | HumanMessage | AIMessage
-AllMessageTypes: TypeAlias = ValidMessageTypes | RemoveMessage
-
-class State(TypedDict):
-    messages: Annotated[list[AllMessageTypes], add_messages]
-    summary: str
-
-class StudentAgent:
-    def __init__(self, student_type: str):
-        llm = OpenAILLMs(temperature=0.75)
-        self.llm = llm.get_llm()
-        self.summary = ""
-        self.conversationalStyle = ""
-        self.type = student_type
-
-        # Define Agent's specific Personas
-        self.role_prompt = process_prompt
-        if self.type == "base":
-            self.role_prompt += base_student_persona
-        elif self.type == "curious":
-            self.role_prompt += curious_student_persona
-        elif self.type == "contradicting":
-            self.role_prompt += contradicting_student_persona
-        elif self.type == "reliant":
-            self.role_prompt += reliant_student_persona
-        elif self.type == "confused":
-            self.role_prompt += confused_student_persona
-        elif self.type == "unrelated":
-            self.role_prompt += unrelated_student_persona
-        else:
-            raise Exception("Unknown Student Agent Type")
-        # Define a new graph for the conversation & compile it
-        self.workflow = StateGraph(State)
-        self.workflow_definition()
-        self.app = self.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
-        system_message = self.role_prompt
-
-        # Adding external student progress and question context details from data queries
-        question_response_details = config["configurable"].get("question_response_details", "")
-        if question_response_details:
-            # convert "my" to "your" in the question_response_details to preserve the student agent as the user
-            question_response_details = question_response_details.replace("My", "Your")
-            question_response_details = question_response_details.replace("my", "your")
-            question_response_details = question_response_details.replace("I am", "you are")
-            system_message += f"\n\n## Known Learning Materials: {question_response_details} \n\n"
-
-        # Adding summary and conversational style to the system message
-        summary = state.get("summary", "")
-        previous_summary = config["configurable"].get("summary", "")
-        if previous_summary:
-            summary = previous_summary
-        if summary:
-            system_message += f"## Summary of conversation earlier: {summary} \n\n"
-
-        messages = [SystemMessage(content=system_message)] + state['messages']
-
-        valid_messages = self.check_for_valid_messages(messages)
-        response = self.llm.invoke(valid_messages)
-
-        # Save summary for fetching outside the class
-        self.summary = summary
-
-        return {"summary": summary, "messages": [response]}
-    
-    def check_for_valid_messages(self, messages: list[AllMessageTypes]) -> list[ValidMessageTypes]:
-        """ Removing the RemoveMessage() from the list of messages """
-
-        valid_messages: list[ValidMessageTypes] = []
-        for message in messages:
-            if message.type != 'remove':
-                valid_messages.append(message)
-        return valid_messages  
-
-    def workflow_definition(self) -> None:
-        self.workflow.add_node("call_llm", self.call_model)
-
-        self.workflow.add_edge(START, "call_llm")
-        self.workflow.add_edge("call_llm", END)
-
-    def get_summary(self) -> str:
-        return self.summary
-
-    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
-    
-def invoke_student_agent(query: str, conversation_history: list, summary: str, student_type:str, question_response_details: str, session_id: str) -> InvokeAgentResponseType:
-    """
-    Call a base student agents that forms a basic conversation with the tutor agent.
-    """
-    print(f'in invoke_student_agent(), student_type: {student_type}')
-    agent = StudentAgent(student_type=student_type)
-
-    config = {"configurable": {"thread_id": session_id, "summary": summary, "question_response_details": question_response_details}}
-    response_events = agent.app.invoke({"messages": conversation_history + [AIMessage(content=query)]}, config=config, stream_mode="values") #updates
-    pretty_printed_response = agent.pretty_response_value(response_events) # get last event/ai answer in the response
-
-    # Gather Metadata from the agent
-    summary = agent.get_summary()
-
-    return {
-        "input": query,
-        "output": pretty_printed_response,
-        "intermediate_steps": [str(summary), conversation_history]
-    }
\ No newline at end of file
diff --git a/src/agents/student_agent/student_prompts.py b/src/agents/student_agent/student_prompts.py
deleted file mode 100644
index e7a8e8e..0000000
--- a/src/agents/student_agent/student_prompts.py
+++ /dev/null
@@ -1,10 +0,0 @@
-# PROMPTS generated with the help of ChatGPT GPT-4o Nov 2024
-
-process_prompt = "Maintain the flow of the conversation by responding directly to the latest message in one sentence. Stay in character as "
-
-base_student_persona = "a student who seeks assistance. Ask questions from a first-person perspective, requesting clarification on how to solve the promblem from the known materials."
-curious_student_persona = "a curious and inquisitive student. Ask thoughtful, detailed questions from a first-person perspective to clarify concepts, explore real-life applications, and uncover complexities. Don’t hesitate to challenge assumptions and ask for clarification when needed."
-contradicting_student_persona = "a skeptical student. Ask questions from a first-person perspective, questioning my reasoning, identifying potential flaws, and challenging explanations. Request clarification whenever something seems unclear or incorrect."
-reliant_student_persona = "a student who relies heavily on your help. Ask questions from a first-person perspective, seeking help for even small problems, and requesting clarification or further assistance to ensure understanding."
-confused_student_persona = "a student who feels confused and uncertain about the topic. Ask questions from a first-person perspective, expressing uncertainty about the material and requesting clarification on both the topic and the tutor’s reasoning."
-unrelated_student_persona = "a student who engages in casual conversation. Ask lighthearted or unrelated questions from a first-person perspective, discussing personal interests or unrelated topics rather than focusing on the material."
\ No newline at end of file
diff --git a/src/agents/utils/langgraph_viz.py b/src/agents/utils/langgraph_viz.py
deleted file mode 100644
index 79e15db..0000000
--- a/src/agents/utils/langgraph_viz.py
+++ /dev/null
@@ -1,13 +0,0 @@
-"""
-Helper script to visualise the agent graph using pygraphviz.
-Setup on mac [see more here https://github.com/pygraphviz/pygraphviz/blob/main/INSTALL.txt]:
-# $ brew install graphviz
-# $ pip install pygraphviz
-"""
-
-agent = ...
-
-
-graph = agent.app.get_graph()
-print(graph)
-graph.draw_png("./graph.png")
\ No newline at end of file
diff --git a/src/agents/utils/synthetic_conversation_generation.py b/src/agents/utils/synthetic_conversation_generation.py
deleted file mode 100644
index 9235756..0000000
--- a/src/agents/utils/synthetic_conversation_generation.py
+++ /dev/null
@@ -1,132 +0,0 @@
-"""
-## Synthetic Dataset Generator ##
--> GOAL: Generate a synthetic dataset of conversations between a tutor and a student [both LLMs].
-
-For each question/scenario example in the example_inputs folder, a pipeline of two agents will be invoked. 
-The agents will play the role of a tutor and a student conversing about the question/scenario.
-
-The conversations will be 20 turns long, with the tutor and student taking turns to send a message.
-
-The tutor can be one of the following types:
-- Informational Agent (base)
-The tutor agent can be selected by changing the "agent_type" field in this script.
-
-The student can have multiple skill levels and conversational styles. Those are defined by the prompts used by the LLM.
-
-Any of the models accessible through the API calls defined in the 'llm_factory.py' can be used for either the tutor and the agent LLM.
-"""
-
-import csv
-import json
-try:
-  from ..student_agent.student_agent import invoke_student_agent
-  from .parse_json_context_to_prompt import parse_json_to_prompt
-  from ..base_agent.base_agent import invoke_base_agent
-except ImportError:
-  from src.agents.student_agent.student_agent import invoke_student_agent
-  from src.agents.utils.parse_json_context_to_prompt import parse_json_to_prompt
-  from src.agents.base_agent.base_agent import invoke_base_agent
-import os
-
-
-def generate_synthetic_conversations(raw_text: str, num_turns: int, student_agent_type: str, tutor_agent_type: str):
-  """
-  Generate a synthetic dataset of conversations between a tutor and a student [both LLMs].
-  """
-  if tutor_agent_type == "base": 
-    invoke_tutor_agent = invoke_base_agent
-  else:
-    raise ValueError("Invalid tutor agent type")
-    
-  parsed_json = json.loads(raw_text)
-  params = parsed_json["params"]
-  conversation_id = params["conversation_id"]
-  include_test_data = params["include_test_data"]
-  summary = ""
-  conversational_style = ""
-  question_response_details = params["question_response_details"]
-  question_submission_summary = question_response_details["questionSubmissionSummary"] if "questionSubmissionSummary" in question_response_details else []
-  question_information = question_response_details["questionInformation"] if "questionInformation" in question_response_details else {}
-  question_access_information = question_response_details["questionAccessInformation"] if "questionAccessInformation" in question_response_details else {}
-  question_response_details_prompt = parse_json_to_prompt(
-    question_submission_summary,
-    question_information,
-    question_access_information
-  )
-  
-  #  Generate Conversation
-  conversation_history = []
-  message = "Ask a question."
-  for i in range(0,num_turns):
-    print(f"Turn {i+1} of {num_turns}")
-    if len(conversation_history) == 0:
-      message = "Ask me a question regarding your thoughts on the learning materials that you are currently woking on."
-    else:
-      message = conversation_history[-1]["content"]
-
-    if i % 2 == 0:
-      # Student starts
-      student_response = invoke_student_agent(message, conversation_history[:-1], summary, student_agent_type, question_response_details_prompt, conversation_id)
-      conversation_history.append({
-        "role": "user",
-        "content": student_response["output"]
-      })
-    else:
-      tutor_response = invoke_tutor_agent(message, conversation_history, summary, conversational_style, question_response_details_prompt, conversation_id)
-      conversation_history.append({
-        "role": "assistant",
-        "content": tutor_response["output"]
-      })
-
-      if "summary" in tutor_response:
-        summary = tutor_response["summary"]
-      if "conversationalStyle" in tutor_response:
-        conversational_style = tutor_response["conversationalStyle"]
-  
-  #  Save Conversation
-  conversation_output = {
-    "conversation_id": conversation_id+"_"+student_agent_type+"_"+tutor_agent_type+"_synthetic",
-    "student_agent_type": student_agent_type,
-    "tutor_agent_type": tutor_agent_type,
-    "conversation": conversation_history
-  }
-  return conversation_output
-
-
-if __name__ == "__main__":
-  num_turns = 6
-  tutor_agent_types   = ["base"]                           
-  # Students can be "base", "curious", "contradicting", "reliant", "confused", "unrelated"
-  student_agent_types = ["base", "curious", "contradicting", "reliant", "confused", "unrelated"]  
-
-  #  Read all question files
-  questions = []
-  example_inputs_folder = "src/agents/utils/example_inputs/"
-  output_folder = "src/agents/utils/synthetic_conversations/"
-  for filename in os.listdir(example_inputs_folder):
-    if filename.endswith("1.json"):
-      questions.append(os.path.join(example_inputs_folder, filename))
-
-    for tutor_agent_type in tutor_agent_types:
-      # Open CSV file for writing
-      csv_filename = os.path.join(output_folder, "all_conversations_"+tutor_agent_type+".csv")
-      with open(csv_filename, "w", newline='') as csvfile:
-        csv_writer = csv.writer(csvfile)
-        # Write the header
-        csv_writer.writerow(["tutor", "student", "conversation", "conversation_id"])
-      
-        for student_agent_type in student_agent_types:
-          for question in questions:
-            print(f"Generating synthetic conversation for {question} with tutor: {tutor_agent_type} and student: {student_agent_type}")
-            with open(question, "r") as file:
-              raw_text = file.read()
-
-              conversation = generate_synthetic_conversations(raw_text, num_turns, student_agent_type, tutor_agent_type)
-
-              conversation_output_filename = output_folder + question.split('/')[-1].replace(".json", "_"+student_agent_type+"_"+tutor_agent_type+"_conversation.json")
-              with open(conversation_output_filename, "w") as file:
-                json.dump(conversation, file, indent=2)
-
-              # Write to CSV
-              conversation_id = conversation["conversation_id"]
-              csv_writer.writerow([tutor_agent_type, student_agent_type, conversation["conversation"], conversation_id])
diff --git a/src/agents/utils/synthetic_conversations/NOTE.md b/src/agents/utils/synthetic_conversations/NOTE.md
deleted file mode 100644
index 4bb113d..0000000
--- a/src/agents/utils/synthetic_conversations/NOTE.md
+++ /dev/null
@@ -1,4 +0,0 @@
-For evaluation purposes of the developed agent, you can use `synthetic_conversation_generation.py` to review the performance of your LLM tutor by running a multi-agent communication with a student agent (available in `src/agents/`).
-
-This folder contains all the synthetic conversations generated by an LLM student discussing with an LLM tutor.
-The files are generated by running the `synthetic_conversation_generation.py`.
\ No newline at end of file
diff --git a/src/agents/utils/testbench_agents.py b/src/agents/utils/testbench_agents.py
deleted file mode 100644
index 46d0d36..0000000
--- a/src/agents/utils/testbench_agents.py
+++ /dev/null
@@ -1,82 +0,0 @@
-"""
- 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
-try:
-    from .parse_json_context_to_prompt import parse_json_to_prompt
-    from ..base_agent.base_agent import invoke_base_agent
-except ImportError:
-    from src.agents.utils.parse_json_context_to_prompt import parse_json_to_prompt
-    from src.agents.base_agent.base_agent import invoke_base_agent
-
-# File path for the input text
-path = "src/agents/utils/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
-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?" 
-    # NOTE: ########################################
-
-    # replace "mock" in the message and conversation history with the actual message
-    parsed_json["message"] = message
-    parsed_json["params"]["conversation_history"][-1]["content"] = message
-
-    params = parsed_json["params"]
-
-    if "include_test_data" in params:
-        include_test_data = params["include_test_data"]
-    if "conversation_history" in params:
-        conversation_history = params["conversation_history"]
-    if "summary" in params:
-        summary = params["summary"]
-    if "conversational_style" in params:
-        conversationalStyle = params["conversational_style"]
-    if "question_response_details" in params:
-        question_response_details = params["question_response_details"]
-        question_submission_summary = question_response_details["questionSubmissionSummary"] if "questionSubmissionSummary" in question_response_details else []
-        question_information = question_response_details["questionInformation"] if "questionInformation" in question_response_details else {}
-        question_access_information = question_response_details["questionAccessInformation"] if "questionAccessInformation" in question_response_details else {}
-        question_response_details_prompt = parse_json_to_prompt(
-            question_submission_summary,
-            question_information,
-            question_access_information
-        )
-        print("Question Response Details Prompt:", question_response_details_prompt, "\n\n")
-
-    if "conversation_id" in params:
-        conversation_id = params["conversation_id"]
-    else:
-        raise Exception("Internal Error: The conversation id is required in the parameters of the chat module.")
-
-    """
-      STEP 3: Call the LLM agent to get a response to the user's message
-    """
-    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)
-    
-    print(response)
-    print("AI Response:", response['output'])
-    
-
-except json.JSONDecodeError as e:
-    print("Error decoding JSON:", e)
-
-
-
diff --git a/src/module.py b/src/module.py
index 7cbbdd8..922979f 100755
--- a/src/module.py
+++ b/src/module.py
@@ -3,14 +3,9 @@
 from lf_toolkit.chat.result import ChatResult as Result
 from lf_toolkit.chat.params import ChatParams as Params
 
-try:
-    from .agents.utils.parse_json_context_to_prompt import parse_json_to_prompt
-    from .agents.base_agent.base_agent import invoke_base_agent
-    from .agents.utils.types import JsonType
-except ImportError:
-    from src.agents.utils.parse_json_context_to_prompt import parse_json_to_prompt
-    from src.agents.base_agent.base_agent import invoke_base_agent
-    from src.agents.utils.types import JsonType
+from src.agent.utils.parse_json_context_to_prompt 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:
     """
@@ -36,40 +31,39 @@ def chat_module(message: Any, params: Params) -> JsonType:
     """
 
     result = Result()
-    include_test_data = False
-    conversation_history = []
-    summary = ""
-    conversationalStyle = ""
-    question_response_details_prompt = ""
 
-    if "include_test_data" in params:
-        include_test_data = params["include_test_data"]
-    if "conversation_history" in params:
-        conversation_history = params["conversation_history"]
-    if "summary" in params:
-        summary = params["summary"]
-    if "conversational_style" in params:
-        conversationalStyle = params["conversational_style"]
-    if "question_response_details" in params:
-        question_response_details = params["question_response_details"]
-        question_submission_summary = question_response_details["questionSubmissionSummary"] if "questionSubmissionSummary" in question_response_details else []
-        question_information = question_response_details["questionInformation"] if "questionInformation" in question_response_details else {}
-        question_access_information = question_response_details["questionAccessInformation"] if "questionAccessInformation" in question_response_details else {}
-        try:
-            question_response_details_prompt = parse_json_to_prompt(
-                question_submission_summary,
-                question_information,
-                question_access_information
-            )
-            print("INFO:: ", question_response_details_prompt)
-        except Exception as e:
-            print("ERROR:: ", e)
-            raise Exception("Internal Error: The question response details could not be parsed.")
-    if "conversation_id" in params:
-        conversation_id = params["conversation_id"]
-    else:
+    # 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
+        )
+    except Exception as e:
+        print("ERROR:: ", e)
+        raise Exception("Internal Error: The question response details could not be parsed.")
     
+
+    # RUN THE AGENT AND MEASURE PROCESSING TIME
     start_time = time.time()
 
     chatbot_response = invoke_base_agent(query=message, \
diff --git a/src/agents/utils/requests_testscript.py b/tests/manual_agent_requests.py
similarity index 93%
rename from src/agents/utils/requests_testscript.py
rename to tests/manual_agent_requests.py
index 6b8e433..9d4c729 100644
--- a/src/agents/utils/requests_testscript.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/agents/utils/example_inputs/"
+path = "src/agent/utils/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
new file mode 100644
index 0000000..b69f609
--- /dev/null
+++ b/tests/manual_agent_run.py
@@ -0,0 +1,46 @@
+"""
+ 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 src.module import chat_module
+
+# File path for the input text
+path = "src/agent/utils/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
+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?" 
+    # 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
+
+    params = parsed_json["params"]
+
+    """
+      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))
+    
+
+except json.JSONDecodeError as e:
+    print("Error decoding JSON:", e)
+
+
+
diff --git a/index_test.py b/tests/test_index.py
similarity index 95%
rename from index_test.py
rename to tests/test_index.py
index fb8e7bf..822049d 100644
--- a/index_test.py
+++ b/tests/test_index.py
@@ -1,10 +1,6 @@
 import unittest
 import json
-
-try:
-    from .index import handler
-except ImportError:
-    from index import handler
+from index import handler
 
 class TestChatIndexFunction(unittest.TestCase):
     """
diff --git a/src/module_test.py b/tests/test_module.py
similarity index 89%
rename from src/module_test.py
rename to tests/test_module.py
index 222eef2..26c96de 100755
--- a/src/module_test.py
+++ b/tests/test_module.py
@@ -1,9 +1,7 @@
 import unittest
-
-try:
-    from .module import Params, chat_module
-except ImportError:
-    from module import Params, chat_module
+from lf_toolkit.chat.result import ChatResult as Result
+from lf_toolkit.chat.params import ChatParams as Params
+from src.module import chat_module
 
 class TestChatModuleFunction(unittest.TestCase):
     """
@@ -27,7 +25,7 @@ class TestChatModuleFunction(unittest.TestCase):
     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 }], \
+        expected_params = Params(include_test_data=True, conversation_history=['{ "type": "user", "content": response }'], \
                                     summary="", conversational_style="", \
                                     question_response_details={}, conversation_id="1234Test")
 
@@ -69,7 +67,7 @@ def test_missing_parameters(self):
     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 }])
+        params = Params(conversation_id="1234Test", conversation_history=['{ "type": "user", "content": response }'])
 
         result = chat_module(response, params)
 
@@ -78,7 +76,7 @@ def test_agent_output(self):
     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 }])
+        params = Params(include_test_data=True, conversation_id="1234Test", conversation_history=['{ "type": "user", "content": response }'])
 
         result = chat_module(response, params)