diff --git a/src/codegen/agents/code_agent.py b/src/codegen/agents/code_agent.py
index 693c0cd44..99406ef40 100644
--- a/src/codegen/agents/code_agent.py
+++ b/src/codegen/agents/code_agent.py
@@ -106,12 +106,11 @@ def __init__(
             **metadata,
         }
 
-    def run(self, prompt: str, image_urls: Optional[list[str]] = None) -> str:
-        """Run the agent with a prompt and optional images.
+    def run(self, prompt: str) -> str:
+        """Run the agent with a prompt.
 
         Args:
             prompt: The prompt to run
-            image_urls: Optional list of base64-encoded image strings. Example: ["data:image/png;base64,<base64_str>"]
             thread_id: Optional thread ID for message history
 
         Returns:
@@ -125,15 +124,14 @@ def run(self, prompt: str, image_urls: Optional[list[str]] = None) -> str:
             "recursion_limit": 100,
         }
 
-        # Prepare content with prompt and images if provided
-        content = [{"type": "text", "text": prompt}]
-        if image_urls:
-            content += [{"type": "image_url", "image_url": {"url": image_url}} for image_url in image_urls]
+        # this message has a reducer which appends the current message to the existing history
+        # see more https://langchain-ai.github.io/langgraph/concepts/low_level/#reducers
+        input = {"query": prompt}
 
         config = RunnableConfig(configurable={"thread_id": self.thread_id}, tags=self.tags, metadata=self.metadata, recursion_limit=200)
         # we stream the steps instead of invoke because it allows us to access intermediate nodes
 
-        stream = self.agent.stream({"messages": [HumanMessage(content=content)]}, config=config, stream_mode="values")
+        stream = self.agent.stream(input, config=config, stream_mode="values")
 
         _tracer = MessageStreamTracer(logger=self.logger)
 
@@ -145,7 +143,7 @@ def run(self, prompt: str, image_urls: Optional[list[str]] = None) -> str:
 
         for s in traced_stream:
             if len(s["messages"]) == 0 or isinstance(s["messages"][-1], HumanMessage):
-                message = HumanMessage(content=content)
+                message = HumanMessage(content=prompt)
             else:
                 message = s["messages"][-1]
 
diff --git a/src/codegen/agents/scratch.ipynb b/src/codegen/agents/scratch.ipynb
index d39e4874f..1b50f8920 100644
--- a/src/codegen/agents/scratch.ipynb
+++ b/src/codegen/agents/scratch.ipynb
@@ -6,7 +6,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from codegen.agents.code_agent import CodeAgent"
+    "from codegen.agents.code_agent import CodeAgent\n",
+    "\n",
+    "\n",
+    "CodeAgent"
    ]
   },
   {
@@ -43,16 +46,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "image = \"\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "agent = CodeAgent(codebase)"
+    "agent = CodeAgent(codebase)\n",
+    "agent.run(\"What is the main character's name? also show the source code where you find the answer\", logger=ConsoleLogger())"
    ]
   },
   {
@@ -61,7 +56,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "agent.run(\"Tell me about the images you see.\", image_urls=[f\"data:image/png;base64,{image}\", f\"data:image/png;base64,{image}\"])"
+    "agent.run(\"What is the main character's name?\")"
    ]
   },
   {
@@ -69,9 +64,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "agent.run(\"What is the main character's name?\")"
-   ]
+   "source": []
   }
  ],
  "metadata": {
@@ -90,7 +83,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.1"
+   "version": "3.13.0"
   }
  },
  "nbformat": 4,