SameerMesiah97 commented on code in PR #70101:
URL: https://github.com/apache/airflow/pull/70101#discussion_r4027186742
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py:
##########
@@ -164,6 +209,111 @@ def run_agent(
endpoint=endpoint,
payload=payload,
timeout=timeout,
+ response_type="dict",
+ )
+
+ def describe_agent(
+ self,
+ *,
+ database: str,
+ schema: str,
+ agent_name: str,
+ timeout: int | None = 600,
+ ) -> JsonDict:
+ """
+ Describe a Snowflake Cortex Agent.
+
+ :param database: Database containing the Cortex Agent.
+ :param schema: Schema containing the Cortex Agent.
+ :param agent_name: Name of the Cortex Agent.
+ :param timeout: Maximum time in seconds to wait for the Cortex Agent
+ request to complete. Optional. Defaults to ``600``.
+ :return: JSON description of the Cortex Agent.
+ """
+ endpoint =
f"/api/v2/databases/{database}/schemas/{schema}/agents/{agent_name}"
+
+ return self._request(
+ method="GET",
+ endpoint=endpoint,
+ timeout=timeout,
+ response_type="dict",
+ )
+
+ def list_agents(
+ self,
+ *,
+ database: str,
+ schema: str,
+ like: str | None = None,
+ from_name: str | None = None,
+ show_limit: int | None = None,
+ timeout: int | None = 600,
+ ) -> JsonList:
+ """
+ List Snowflake Cortex Agents.
+
+ :param database: Database containing the Cortex Agents.
+ :param schema: Schema containing the Cortex Agents.
+ :param like: Case-insensitive name filter. Optional.
+ Defaults to ``None``.
+ :param from_name: Pagination starting point. Optional.
+ Defaults to ``None``.
+ :param show_limit: Maximum number of agents to return. Optional.
+ Defaults to ``None``.
+ :param timeout: Maximum time in seconds to wait for the Cortex Agent
+ request to complete. Optional. Defaults to ``600``.
+ :return: List of Cortex Agents.
+ """
+ endpoint = f"/api/v2/databases/{database}/schemas/{schema}/agents"
+
+ params: dict[str, Any] = {}
+
+ if like is not None:
+ params["like"] = like
+
+ if from_name is not None:
+ params["fromName"] = from_name
+
+ if show_limit is not None:
+ params["showLimit"] = show_limit
+
+ return self._request(
+ method="GET",
+ endpoint=endpoint,
+ params=params or None,
+ timeout=timeout,
+ response_type="list",
+ )
+
+ def delete_agent(
+ self,
+ *,
+ database: str,
+ schema: str,
+ agent_name: str,
+ if_exists: bool = False,
+ timeout: int | None = 600,
+ ) -> JsonDict:
+ """
+ Delete a Snowflake Cortex Agent.
+
+ :param database: Database containing the Cortex Agent.
+ :param schema: Schema containing the Cortex Agent.
+ :param agent_name: Name of the Cortex Agent.
+ :param if_exists: If ``True``, do not fail when the agent does not
exist.
+ Optional. Defaults to ``False``.
+ :param timeout: Maximum time in seconds to wait for the Cortex Agent
request
+ to complete. Optional. Defaults to ``600``.
+ :return: JSON response confirming deletion.
+ """
+ endpoint =
f"/api/v2/databases/{database}/schemas/{schema}/agents/{agent_name}"
Review Comment:
I have wrapped the various parts of each endpoint with
`urllib.parse.quote(x, safe="")`
--
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]
