SameerMesiah97 commented on code in PR #70101:
URL: https://github.com/apache/airflow/pull/70101#discussion_r4027237364
##########
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(
Review Comment:
The pagination mechanism has now been explicitly documented in the docstring
for `list_agents`.
--
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]
