SameerMesiah97 commented on code in PR #70930:
URL: https://github.com/apache/airflow/pull/70930#discussion_r4200451978
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowflake_cortex_agent.py:
##########
@@ -139,3 +139,275 @@ def execute(self, context: Context) -> dict[str, Any]:
tool_resources=self.tool_resources,
timeout=self.timeout,
)
+
+
+class SnowflakeCortexAgentCreateOperator(BaseOperator):
+ """
+ Create a Snowflake Cortex Agent.
+
+ :param database: Database in which to create the Cortex Agent.
+ :param schema: Schema in which to create the Cortex Agent.
+ :param agent_name: Name of the Cortex Agent.
+ :param comment: Optional comment.
+ :param profile: Agent profile configuration. Optional.
+ :param models: Model configuration. Optional.
+ :param instructions: Agent instructions. Optional.
+ :param orchestration: Orchestration configuration. Optional.
+ :param tools: Agent tools. Optional.
+ :param tool_resources: Tool resource configuration. Optional.
+ :param create_mode: Resource creation mode. Defaults to
+ ``CreateMode.ERROR_IF_EXISTS``.
+ :param timeout: Maximum time in seconds to wait for the request to
+ complete. Defaults to ``600``.
+ :param snowflake_conn_id: Snowflake connection ID. Defaults to
+ ``snowflake_default``.
+ """
+
+ template_fields: Sequence[str] = (
+ "database",
+ "schema",
+ "agent_name",
+ "comment",
+ "profile",
+ "models",
+ "instructions",
+ "orchestration",
+ "tools",
+ "tool_resources",
+ "snowflake_conn_id",
+ )
+
+ ui_color = "#29B5E8"
+
+ def __init__(
+ 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 | str = CreateMode.ERROR_IF_EXISTS,
+ timeout: int | None = 600,
+ snowflake_conn_id: str = "snowflake_default",
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+
+ self.database = database
+ self.schema = schema
+ self.agent_name = agent_name
+ self.comment = comment
+ self.profile = profile
+ self.models = models
+ self.instructions = instructions
+ self.orchestration = orchestration
+ self.tools = tools
+ self.tool_resources = tool_resources
+ self.create_mode = create_mode
Review Comment:
The argument for create_mode is now being normalized. An additional test has
been added for this as well.
--
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]
