Editing & Debugging Code Agent with VS Code#

This tutorial shows how to develop, test, and debug a Code Agent in Dataiku using Code Studio with VS Code. You will learn how to edit, sync, and safely manage Code Agents in your projects. For setting up a Code Studio template, see the documentation.

Prerequisites#

A Dataiku 11+ instance.
Administrator privileges for your user profile.
A Kubernetes cluster is configured. For details, visit Elastic AI Computation.
A base image is built. Typically, this is built using a command such as ./bin/dssadmin build-base-image --type container-exec. For details, visit Build the Base Image.
A project with some Code Agent and Code Tool. If you haven’t got one, you can follow this tutorial.
A Code Studio Template with VS Code and the Code Environment used by your GenAi component. Refer to this tutorial if you don’t know how to set up a Code Studio template.

Editing a Code Agent#

If you haven’t set up a Code Studio, create one. Select the Code Studio attached to your project, and click the Start button. Once the Code Studio is running, you can select your Code Agent in the code-agents folder within VS Code.

Note

You should not create a Code Agent directly in the code-agents folder. Creating a Code Agent requires more than just creating a simple file. If you create a new file in this folder, you will be able to perform some tests, but the ‘Sync files with DSS’ button won’t save it back to DSS.

To find your agent, you need to know its ID. This code snippet can help you find it. A Code Agent can have multiple versions. You also need to know which version you want to edit. Suppose your agent has the name Extract_topics, its ID is pO0H7cvh, and you want to edit the version named v1, you should edit the file code-agents/Extract_topics_pO0H7cvh_v1.py.

When making changes to a Code Agent, you may want to test it. If you use the code below, the agent will run, but you won’t be able to see your changes, as you are working on a local copy of the code.

import dataiku

client = dataiku.api_client()
project = client.get_default_project()


CODE_AGENT = "" # Fill with your Code Agent ID
llm  = project.get_llm(CODE_AGENT)

completion = llm.new_completion()
completion.with_message("How to edit an Code Agent in Dataiku?")
resp = completion.execute()
print(resp.success)
print(resp.trace)

To view your changes, synchronize your files with DSS and then run the code above. However, if you do that, the Code Agent will be updated for the entire project, and users may be affected by your changes. Usually, it is not the way that you want to go when you are coding. Hopefully, Dataiku provides a helper to test your modifications before publishing them.

Let’s take a real example. Imagine having a Code Agent (Extract_topics) that reformulates a question into smaller, more easily answerable questions, as shown in the code below. You also need to test/modify it, for example, to see if changing the base LLM would improve your results.

import dataiku
from dataiku.llm.python import BaseLLM

OPENAI_CONNECTION_NAME = "REPLACE_WITH_YOUR_CONNECTION_NAME"  # example: "openAI"

class MyLLM(BaseLLM):
    def __init__(self):
        pass

    def process(self, query, settings, trace):
        llm = dataiku.api_client().get_default_project().get_llm(f"openai:{OPENAI_CONNECTION_NAME}:gpt-4.1-mini")
        completion = llm.new_completion().with_message("Reformulates a question into smaller, more easily answerable questions", "system")
        for message in query["messages"]:
            if message.get("content"):
                completion = completion.with_message(message["content"], message.get("role", "user"))
        llm_resp = completion.execute()

        resp_text = "Here are the questions you need to answer."
        resp_text = resp_text + "\n" + llm_resp.text
        return {"text": resp_text}

If you run this code with the following input, you will obtain the following response:

{
   "messages": [
      {
         "role": "user",
         "content": "What is the capital of the country where the Eiffel Tower is built?"
      }
   ],
   "context": {}
}

Here are the questions you need to answer.
To answer your question, let's break it down into smaller parts:

1. Which country is the Eiffel Tower located in?
2. What is the capital city of that country?

Would you like me to provide the answers?

Using the Dataiku helper#

Now, you need to see if changing the LLM will impact the result. In VS Code, in the code-agents folder, create a file named test_agent.py with the following content:

import json
from Extract_topics_pO0H7cvh_v1 import MyLLM
from dataiku.llm.python.tests import run_completion_query

resp = run_completion_query(MyLLM, "What is the capital of the country where the Eiffel Tower is built?")
print(json.dumps(resp, indent=4))

The highlighted line shows how to use the helper for testing your Code Agent. Now that you know how to use your Code Agent locally, you can test your modifications. For example, if you change your LLM definition with llm = dataiku.api_client().get_default_project().get_llm(f"openai:{OPENAI_CONNECTION_NAME}:gpt-5-mini"), and run the test again, you will obtain:

{
    "text": "Here are the questions you need to answer.\nHere are smaller, easy-to-answer questions that lead to the original answer:\n\n1. In which city is the Eiffel Tower located?  \n2. Which country is that city in?  \n3. What is the capital of that country?",
    "trace": {
        "type": "span",
        "name": "DKU_TEST_AGENT_CALL",
        "children": [],
        "attributes": {},
        "inputs": {},
        "outputs": {},
        "begin": "2025-12-16T13:23:53.538000Z",
        "end": "2025-12-16T13:28:08.331000Z",
        "duration": 254793
    }
}

Making it more usable#

With the helper, you can easily test your modifications. However, if you synchronize your files, you will see that your test file (test_agent.py) has disappeared. This is due to the synchronization; as the test file is not a Code Agent, it is not saved, and as it is not present as an Agent, the synchronization deletes it. To be able to keep your test file from one synchronization to another, you need to move it from the code-agents folder to the code_studio-versioned folder. However, if you simply move the file, you will encounter an error, as Extract_topics_pO0H7cvh_v1.py is no longer in the same folder as the test file. You will need to modify the test file to include the code-agents folder in the path, as shown in the code below.

import sys
import os

current_file_path = os.path.abspath(__file__)
parent_dir = os.path.dirname(os.path.dirname(current_file_path))
sibling_dir = os.path.join(parent_dir, 'code-agents')
sys.path.insert(0, sibling_dir)

import json
from Extract_topics_pO0H7cvh_v1 import MyLLM
from dataiku.llm.python.tests import run_completion_query

resp = run_completion_query(MyLLM, "What is the capital of the country where the Eiffel Tower is built?")
print(json.dumps(resp, indent=4))

As you are able to run your Code Agent in Code Studio, you can use the Debug action of Code Studio, as you do for classical debugging.

If your agent is using a tool, you will also be able to debug it by putting a breakpoint in the tool.

Do not forget to synchronize your files with Dataiku once you have finished testing/debugging your Code Agent.

Wrapping up#

In this tutorial, you learned how to develop and test a Code Agent in Dataiku using Code Studio and VS Code. You discovered how to safely edit, synchronize, and debug your agents, as well as run local tests before publishing changes.

Here is the complete code of the tutorial: