How to structure a Python Code Tool for agent use?

Hi there, a bit confused about how exactly to use the Code Tool with python.

Given a Code Tool with the following:

def get_jurisdiction(query_location_ids: str):
    id_name_map = {
        384: "USA National",
        42442: "China National"
    }
    if not isinstance(query_location_ids, str) or not query_location_ids:
        return {}
    id_jurisdiction_map = {}
    for query_location_id in query_location_ids.split(","):
        int_id = int(query_location_id.strip())
        id_jurisdiction_map[int_id] = id_name_map.get(int_id, "Unknown")
    return id_jurisdiction_map

return get_jurisdiction(_query.query_location_ids)

I specify the input schema using Generate From JSON Example:

{
	"query_location_ids": "12,345,6789"
}

But get this error:

[
  {"response": "There was an error: \"name '_query' is not defined\""}
]

I’m probably just missing something simple, but couldn’t find it after searching.
Any help appreciated

Hi @sac.mm.xlv !
You’re very close. make sure you define the query name correctly and use the same name in your function, the parameter name in your schema must match the parameter name in your function signature,

1 Like

The _query object isn’t a thing in the Code Tool, that’s why it’s undefined. Your input variables are available directly so just do return get_jurisdiction(query_location_ids) on that last line and it should work.

1 Like