SameerMesiah97 commented on code in PR #70703:
URL: https://github.com/apache/airflow/pull/70703#discussion_r4097489109
##########
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:
I tested the API and PUT supports partial updates so for e.g. updating the
only comment field is not going to overwrite the other fields.
--
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]
