potiuk commented on code in PR #70703:
URL: https://github.com/apache/airflow/pull/70703#discussion_r4070366233
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py:
##########
@@ -215,6 +267,114 @@ def run_agent(
response_type="dict",
)
+ def create_agent(
+ self,
+ *,
+ database: str,
+ schema: str,
+ agent_name: str,
+ comment: str | None = None,
+ profile: dict[str, Any] | None = None,
+ models: dict[str, Any] | None = None,
+ instructions: dict[str, Any] | None = None,
+ orchestration: dict[str, Any] | None = None,
+ tools: list[dict[str, Any]] | None = None,
+ tool_resources: dict[str, Any] | None = None,
+ create_mode: CreateMode = CreateMode.ERROR_IF_EXISTS,
+ timeout: int | None = 600,
+ ) -> JsonDict:
+ """
+ Create a Snowflake Cortex Agent.
+
+ :param database: Database in which to create the agent.
+ :param schema: Schema in which to create the agent.
+ :param agent_name: Name of the Cortex Agent.
+ :param comment: Optional comment. Optional. Defaults to ``None``.
+ :param profile: Agent profile configuration. Optional. Defaults to
``None``.
+ :param models: Model configuration. Optional. Defaults to ``None``.
+ :param instructions: Agent instructions. Optional. Defaults to
``None``.
+ :param orchestration: Orchestration configuration. Optional. Defaults
to ``None``.
+ :param tools: Agent tools. Optional. Defaults to ``None``.
+ :param tool_resources: Tool resource configuration. Optional. Defaults
to ``None``.
+ :param create_mode: Resource creation mode. One of ``errorIfExists``,
``orReplace``
+ or ``ifNotExists``. Optional. Defaults to ``errorIfExists``.
+ :param timeout: Maximum time in seconds to wait for the Cortex Agent
request
+ to complete. Defaults to ``600``.
+ :return: JSON response confirming creation.
+ """
+ payload = {
+ "name": agent_name,
+ **self._build_agent_payload(
+ comment=comment,
+ profile=profile,
+ models=models,
+ instructions=instructions,
+ orchestration=orchestration,
+ tools=tools,
+ tool_resources=tool_resources,
+ ),
+ }
+
+ endpoint = f"/api/v2/databases/{quote(database,
safe='')}/schemas/{quote(schema, safe='')}/agents"
+
+ return self._request(
+ method="POST",
+ endpoint=endpoint,
+ payload=payload,
+ params={"createMode": create_mode.value},
+ timeout=timeout,
+ response_type="dict",
+ )
+
+ def update_agent(
+ self,
+ *,
+ database: str,
+ schema: str,
+ agent_name: str,
+ comment: str | None = None,
+ profile: dict[str, Any] | None = None,
+ models: dict[str, Any] | None = None,
+ instructions: dict[str, Any] | None = None,
+ orchestration: dict[str, Any] | None = None,
+ tools: list[dict[str, Any]] | None = None,
+ tool_resources: dict[str, Any] | None = None,
+ timeout: int | None = 600,
+ ) -> JsonDict:
+ """
+ Update a Snowflake Cortex Agent.
+
+ :param database: Database containing the agent.
+ :param schema: Schema containing the agent.
+ :param agent_name: Name of the Cortex Agent.
Review Comment:
`update_agent` accepts `comment`, `profile`, `models`, `instructions`,
`orchestration`, `tools` and `tool_resources`, but the docstring goes straight
from `agent_name` to `timeout` — seven of the eleven parameters are
undocumented. `create_agent` a few lines up documents all of them, and these
render in the published provider API reference, so `update_agent` would ship
with seven undocumented arguments next to a sibling that documents them.
```suggestion
:param agent_name: Name of the Cortex Agent.
:param comment: Optional comment. Optional. Defaults to ``None``.
:param profile: Agent profile configuration. Optional. Defaults to
``None``.
:param models: Model configuration. Optional. Defaults to ``None``.
:param instructions: Agent instructions. Optional. Defaults to
``None``.
:param orchestration: Orchestration configuration. Optional.
Defaults to ``None``.
:param tools: Agent tools. Optional. Defaults to ``None``.
:param tool_resources: Tool resource configuration. Optional.
Defaults to ``None``.
```
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py:
##########
@@ -215,6 +267,114 @@ def run_agent(
response_type="dict",
)
+ def create_agent(
+ self,
+ *,
+ database: str,
+ schema: str,
+ agent_name: str,
+ comment: str | None = None,
+ profile: dict[str, Any] | None = None,
+ models: dict[str, Any] | None = None,
+ instructions: dict[str, Any] | None = None,
+ orchestration: dict[str, Any] | None = None,
+ tools: list[dict[str, Any]] | None = None,
+ tool_resources: dict[str, Any] | None = None,
+ create_mode: CreateMode = CreateMode.ERROR_IF_EXISTS,
+ timeout: int | None = 600,
+ ) -> JsonDict:
+ """
+ Create a Snowflake Cortex Agent.
+
+ :param database: Database in which to create the agent.
+ :param schema: Schema in which to create the agent.
+ :param agent_name: Name of the Cortex Agent.
+ :param comment: Optional comment. Optional. Defaults to ``None``.
+ :param profile: Agent profile configuration. Optional. Defaults to
``None``.
+ :param models: Model configuration. Optional. Defaults to ``None``.
+ :param instructions: Agent instructions. Optional. Defaults to
``None``.
+ :param orchestration: Orchestration configuration. Optional. Defaults
to ``None``.
+ :param tools: Agent tools. Optional. Defaults to ``None``.
+ :param tool_resources: Tool resource configuration. Optional. Defaults
to ``None``.
+ :param create_mode: Resource creation mode. One of ``errorIfExists``,
``orReplace``
+ or ``ifNotExists``. Optional. Defaults to ``errorIfExists``.
+ :param timeout: Maximum time in seconds to wait for the Cortex Agent
request
+ to complete. Defaults to ``600``.
+ :return: JSON response confirming creation.
+ """
+ payload = {
+ "name": agent_name,
+ **self._build_agent_payload(
+ comment=comment,
+ profile=profile,
+ models=models,
+ instructions=instructions,
+ orchestration=orchestration,
+ tools=tools,
+ tool_resources=tool_resources,
+ ),
+ }
+
+ endpoint = f"/api/v2/databases/{quote(database,
safe='')}/schemas/{quote(schema, safe='')}/agents"
+
+ return self._request(
+ method="POST",
+ endpoint=endpoint,
+ payload=payload,
+ params={"createMode": create_mode.value},
+ timeout=timeout,
+ response_type="dict",
+ )
+
+ def update_agent(
+ self,
+ *,
+ database: str,
+ schema: str,
+ agent_name: str,
+ comment: str | None = None,
+ profile: dict[str, Any] | None = None,
+ models: dict[str, Any] | None = None,
+ instructions: dict[str, Any] | None = None,
+ orchestration: dict[str, Any] | None = None,
+ tools: list[dict[str, Any]] | None = None,
+ tool_resources: dict[str, Any] | None = None,
+ timeout: int | None = 600,
+ ) -> JsonDict:
+ """
+ Update a Snowflake Cortex Agent.
+
+ :param database: Database containing the agent.
+ :param schema: Schema containing the agent.
+ :param agent_name: Name of the Cortex Agent.
+ :param timeout: Maximum time in seconds to wait for the Cortex Agent
request
+ to complete. Defaults to ``600``.
+ :return: JSON response confirming the update, or an empty dictionary
when
+ Snowflake returns a successful response without a body.
+ """
+ endpoint = (
+ f"/api/v2/databases/{quote(database, safe='')}"
+ f"/schemas/{quote(schema, safe='')}"
+ f"/agents/{quote(agent_name, safe='')}"
+ )
+
+ return self._request(
+ method="PUT",
+ endpoint=endpoint,
+ payload=self._build_agent_payload(
Review Comment:
This is the one I'd like confirmed before the PR lands.
`_build_agent_payload` omits every unset field, which makes
`update_agent(comment="new comment")` *look* like a partial patch. But `PUT
/agents/{name}` is REST-conventionally a full replace — if Snowflake implements
it that way, that call silently clears the agent's `tools`, `models`,
`instructions` and `orchestration`. I couldn't find anything in [Snowflake's
Cortex Agents REST API
docs](https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-agents-rest-api)
that states which semantics apply.
That's the difference between "change the comment" and "delete the agent's
tool configuration", and the current signature actively encourages the
dangerous reading. Could you check against a live account and write the answer
into the docstring? If it turns out to be a replace, this probably wants an
explicit warning in the docstring, or a `describe_agent`-then-merge path so
callers can't lose configuration by omission.
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py:
##########
@@ -124,6 +139,43 @@ def _request(
return data
+ def _build_agent_payload(
Review Comment:
Nit, take it or leave it: this never touches `self`, so it could be a
`@staticmethod`.
##########
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_cortex_agent.py:
##########
@@ -381,6 +384,157 @@ def test_get_text_response(
):
assert SnowflakeCortexAgentHook.get_text_response(response) == expected
+ @pytest.mark.parametrize(
+ ("create_mode", "comment", "profile", "instructions",
"expected_params", "expected_payload"),
Review Comment:
`models`, `orchestration`, `tools` and `tool_resources` are never passed
non-`None` in either `test_create_agent` or `test_update_agent`, so four of the
seven branches in the new `_build_agent_payload` are never executed.
From [`AGENTS.md` § Testing
Standards](https://github.com/apache/airflow/blob/main/AGENTS.md#testing-standards):
> Target exactly 100% coverage of what the PR changes — no more, no less.
Every changed or added behaviour must have a test; every test must fail without
the PR's change.
Adding those four to the existing `default` param case (argnames tuple plus
the one case's values and `expected_payload`) covers them without a new test.
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py:
##########
@@ -215,6 +267,114 @@ def run_agent(
response_type="dict",
)
+ def create_agent(
+ self,
+ *,
+ database: str,
+ schema: str,
+ agent_name: str,
+ comment: str | None = None,
+ profile: dict[str, Any] | None = None,
+ models: dict[str, Any] | None = None,
+ instructions: dict[str, Any] | None = None,
+ orchestration: dict[str, Any] | None = None,
+ tools: list[dict[str, Any]] | None = None,
+ tool_resources: dict[str, Any] | None = None,
+ create_mode: CreateMode = CreateMode.ERROR_IF_EXISTS,
Review Comment:
`create_mode.value` below means `create_agent(create_mode="orReplace")`
fails with `AttributeError: 'str' object has no attribute 'value'`.
`CreateMode` subclasses `str` precisely so the raw API strings are usable,
and a plain string is what a config-driven or templated caller naturally
produces. Widening the annotation and normalising at the call site accepts
both, and turns a typo into a clear `ValueError` instead of a malformed
`createMode` query param that Snowflake rejects:
```suggestion
create_mode: CreateMode | str = CreateMode.ERROR_IF_EXISTS,
```
…paired with `params={"createMode": CreateMode(create_mode).value}` in the
`_request` call below.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
