justinpakzad commented on code in PR #70101:
URL: https://github.com/apache/airflow/pull/70101#discussion_r3792441132
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py:
##########
@@ -159,11 +161,124 @@ def run_agent(
endpoint =
f"/api/v2/databases/{database}/schemas/{schema}/agents/{agent_name}:run"
- return self._request(
- method="POST",
- endpoint=endpoint,
- payload=payload,
- timeout=timeout,
+ return cast(
+ "dict[str, Any]",
+ self._request(
+ method="POST",
+ endpoint=endpoint,
+ payload=payload,
+ timeout=timeout,
+ ),
+ )
+
+ def describe_agent(
+ self,
+ *,
+ database: str,
+ schema: str,
+ agent_name: str,
+ timeout: int | None = 600,
+ ) -> dict[str, Any]:
+ """
+ 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. Defaults to ``600``.
+ :return: JSON description of the Cortex Agent.
+ """
+ endpoint =
f"/api/v2/databases/{database}/schemas/{schema}/agents/{agent_name}"
+
+ return cast(
+ "dict[str, Any]",
+ self._request(
+ method="GET",
+ endpoint=endpoint,
+ timeout=timeout,
+ ),
+ )
+
+ 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,
+ ) -> list[dict[str, Any]]:
+ """
+ List Snowflake Cortex Agents.
+
+ :param database: Database containing the Cortex Agents.
+ :param schema: Schema containing the Cortex Agents.
+ :param like: Optional case-insensitive name filter. Optional.
+ Defaults to ``None``.
+ :param from_name: Optional pagination starting point. Optional.
+ Defaults to ``None``.
Review Comment:
Nit: "Optional" is repeated here. I think we can drop the first occurrence.
--
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]
