snowflake-cli-labs 2.8.0rc1__py3-none-any.whl → 2.8.2__py3-none-any.whl

snowflake/cli/plugins/cortex/commands.py

@@ -1,327 +0,0 @@
-# Copyright (c) 2024 Snowflake Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import sys
-from pathlib import Path
-from typing import List, Optional
-
-import click
-import typer
-from click import UsageError
-from snowflake.cli.api.cli_global_context import cli_context
-from snowflake.cli.api.commands.flags import readable_file_option
-from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
-from snowflake.cli.api.constants import PYTHON_3_12
-from snowflake.cli.api.output.types import (
-    CollectionResult,
-    CommandResult,
-    MessageResult,
-)
-from snowflake.cli.api.secure_path import SecurePath
-from snowflake.cli.plugins.cortex.constants import DEFAULT_MODEL
-from snowflake.cli.plugins.cortex.manager import CortexManager
-from snowflake.cli.plugins.cortex.types import (
-    Language,
-    Model,
-    Question,
-    SourceDocument,
-    Text,
-)
-
-app = SnowTyperFactory(
-    name="cortex",
-    help="Provides access to Snowflake Cortex.",
-)
-
-SEARCH_COMMAND_ENABLED = sys.version_info < PYTHON_3_12
-
-
-@app.command(
-    requires_connection=True,
-    hidden=not SEARCH_COMMAND_ENABLED,
-)
-def search(
-    query: str = typer.Argument(help="The search query string"),
-    service: str = typer.Option(
-        help="Cortex search service to be used. Example: --service my_cortex_service",
-    ),
-    columns: Optional[List[str]] = typer.Option(
-        help='Columns that will be returned with the results. If none is provided, only search column will be included in results. Example --columns "foo" --columns "bar"',
-        default=None,
-    ),
-    limit: int = typer.Option(help="Maximum number of results retrieved", default=1),
-    **options,
-):
-    """
-    Performs query search using Cortex Search Services.
-    """
-
-    if not SEARCH_COMMAND_ENABLED:
-        raise click.ClickException(
-            "Cortex Search uses Snowflake Python API that currently does not support your Python version"
-        )
-
-    from snowflake.core import Root
-
-    if not columns:
-        columns = []
-
-    conn = cli_context.connection
-
-    search_service = (
-        Root(conn)
-        .databases[conn.database]
-        .schemas[conn.schema]
-        .cortex_search_services[service]
-    )
-
-    response = search_service.search(
-        query=query, columns=columns, limit=limit, filter={}
-    )
-
-    return CollectionResult(response.results)
-
-
-@app.command(
-    name="complete",
-    requires_connection=True,
-)
-def complete(
-    text: Optional[str] = typer.Argument(
-        None,
-        help="Prompt to be used to generate a completion. Cannot be combined with --file option.",
-        show_default=False,
-    ),
-    model: Optional[str] = typer.Option(
-        DEFAULT_MODEL,
-        "--model",
-        help="String specifying the model to be used.",
-    ),
-    file: Optional[Path] = readable_file_option(
-        param_name="--file",
-        help_str="JSON file containing conversation history to be used to generate a completion. Cannot be combined with TEXT argument.",
-    ),
-    **options,
-) -> CommandResult:
-    """
-    Given a prompt, the command generates a response using your choice of language model.
-    In the simplest use case, the prompt is a single string.
-    You may also provide a JSON file with conversation history including multiple prompts and responses for interactive chat-style usage.
-    """
-
-    manager = CortexManager()
-
-    if text and file:
-        raise UsageError("--file option cannot be used together with TEXT argument.")
-    if text:
-        result_text = manager.complete_for_prompt(
-            text=Text(text),
-            model=Model(model),
-        )
-    elif file:
-        result_text = manager.complete_for_conversation(
-            conversation_json_file=SecurePath(file),
-            model=Model(model),
-        )
-    else:
-        raise UsageError("Either --file option or TEXT argument has to be provided.")
-
-    return MessageResult(result_text.strip())
-
-
-@app.command(
-    name="extract-answer",
-    requires_connection=True,
-)
-def extract_answer(
-    question: str = typer.Argument(
-        None,
-        help="String containing the question to be answered.",
-        show_default=False,
-    ),
-    source_document_text: Optional[str] = typer.Argument(
-        None,
-        help="String containing the plain-text or JSON document that contains the answer to the question. Cannot be combined with --file option.",
-        show_default=False,
-    ),
-    file: Optional[Path] = readable_file_option(
-        param_name="--file",
-        help_str="File containing the plain-text or JSON document that contains the answer to the question. Cannot be combined with SOURCE_DOCUMENT_TEXT argument.",
-    ),
-    **options,
-) -> CommandResult:
-    """
-    Extracts an answer to a given question from a text document.
-    The document may be a plain-English document or a string representation of a semi-structured (JSON) data object.
-    """
-
-    manager = CortexManager()
-
-    if source_document_text and file:
-        raise UsageError(
-            "--file option cannot be used together with SOURCE_DOCUMENT_TEXT argument."
-        )
-    if source_document_text:
-        result_text = manager.extract_answer_from_source_document(
-            source_document=SourceDocument(source_document_text),
-            question=Question(question),
-        )
-    elif file:
-        result_text = manager.extract_answer_from_source_document_file(
-            source_document_input_file=SecurePath(file),
-            question=Question(question),
-        )
-    else:
-        raise UsageError(
-            "Either --file option or SOURCE_DOCUMENT_TEXT argument has to be provided."
-        )
-
-    return MessageResult(result_text.strip())
-
-
-@app.command(
-    name="sentiment",
-    requires_connection=True,
-)
-def sentiment(
-    text: Optional[str] = typer.Argument(
-        None,
-        help="String containing the text for which a sentiment score should be calculated. Cannot be combined with --file option.",
-        show_default=False,
-    ),
-    file: Optional[Path] = readable_file_option(
-        param_name="--file",
-        help_str="File containing the text for which a sentiment score should be calculated. Cannot be combined with TEXT argument.",
-    ),
-    **options,
-) -> CommandResult:
-    """
-    Returns sentiment as a score between -1 to 1
-    (with -1 being the most negative and 1 the most positive,
-    with values around 0 neutral) for the given English-language input text.
-    """
-
-    manager = CortexManager()
-
-    if text and file:
-        raise UsageError("--file option cannot be used together with TEXT argument.")
-    if text:
-        result_text = manager.calculate_sentiment_for_text(
-            text=Text(text),
-        )
-    elif file:
-        result_text = manager.calculate_sentiment_for_text_file(
-            text_file=SecurePath(file),
-        )
-    else:
-        raise UsageError("Either --file option or TEXT argument has to be provided.")
-
-    return MessageResult(result_text.strip())
-
-
-@app.command(
-    name="summarize",
-    requires_connection=True,
-)
-def summarize(
-    text: Optional[str] = typer.Argument(
-        None,
-        help="String containing the English text from which a summary should be generated. Cannot be combined with --file option.",
-        show_default=False,
-    ),
-    file: Optional[Path] = readable_file_option(
-        param_name="--file",
-        help_str="File containing the English text from which a summary should be generated. Cannot be combined with TEXT argument.",
-    ),
-    **options,
-) -> CommandResult:
-    """
-    Summarizes the given English-language input text.
-    """
-
-    manager = CortexManager()
-
-    if text and file:
-        raise UsageError("--file option cannot be used together with TEXT argument.")
-    if text:
-        result_text = manager.summarize_text(
-            text=Text(text),
-        )
-    elif file:
-        result_text = manager.summarize_text_file(
-            text_file=SecurePath(file),
-        )
-    else:
-        raise UsageError("Either --file option or TEXT argument has to be provided.")
-
-    return MessageResult(result_text.strip())
-
-
-@app.command(
-    name="translate",
-    requires_connection=True,
-)
-def translate(
-    text: Optional[str] = typer.Argument(
-        None,
-        help="String containing the text to be translated. Cannot be combined with --file option.",
-        show_default=False,
-    ),
-    from_language: Optional[str] = typer.Option(
-        None,
-        "--from",
-        help="String specifying the language code for the language the text is currently in. See Snowflake Cortex documentation for a list of supported language codes.",
-        show_default=False,
-    ),
-    to_language: str = typer.Option(
-        ...,
-        "--to",
-        help="String specifying the language code into which the text should be translated. See Snowflake Cortex documentation for a list of supported language codes.",
-        show_default=False,
-    ),
-    file: Optional[Path] = readable_file_option(
-        param_name="--file",
-        help_str="File containing the text to be translated. Cannot be combined with TEXT argument.",
-    ),
-    **options,
-) -> CommandResult:
-    """
-    Translates text from the indicated or detected source language to a target language.
-    """
-
-    manager = CortexManager()
-
-    source_language = None if from_language is None else Language(from_language)
-    target_language = Language(to_language)
-
-    if text and file:
-        raise UsageError("--file option cannot be used together with TEXT argument.")
-    if text:
-        result_text = manager.translate_text(
-            text=Text(text),
-            source_language=source_language,
-            target_language=target_language,
-        )
-    elif file:
-        result_text = manager.translate_text_file(
-            text_file=SecurePath(file),
-            source_language=source_language,
-            target_language=target_language,
-        )
-    else:
-        raise UsageError("Either --file option or TEXT argument has to be provided.")
-
-    return MessageResult(result_text.strip())

snowflake/cli/plugins/cortex/constants.py

@@ -1,17 +0,0 @@
-# Copyright (c) 2024 Snowflake Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from snowflake.cli.plugins.cortex.types import Model
-
-DEFAULT_MODEL: Model = Model("snowflake-arctic")

snowflake/cli/plugins/cortex/manager.py

@@ -1,189 +0,0 @@
-# Copyright (c) 2024 Snowflake Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import json
-import logging
-from typing import Callable, Optional
-
-from click import ClickException
-from snowflake.cli.api.constants import DEFAULT_SIZE_LIMIT_MB
-from snowflake.cli.api.exceptions import SnowflakeSQLExecutionError
-from snowflake.cli.api.secure_path import SecurePath
-from snowflake.cli.api.sql_execution import SqlExecutionMixin
-from snowflake.cli.plugins.cortex.types import (
-    Language,
-    Model,
-    Question,
-    SourceDocument,
-    Text,
-)
-from snowflake.connector import ProgrammingError
-from snowflake.connector.cursor import DictCursor
-
-log = logging.getLogger(__name__)
-
-
-class CortexManager(SqlExecutionMixin):
-    def complete_for_prompt(
-        self,
-        text: Text,
-        model: Model,
-    ) -> str:
-        query = f"""\
-            SELECT SNOWFLAKE.CORTEX.COMPLETE(
-                '{model}',
-                '{self._escape_input(text)}'
-            ) AS CORTEX_RESULT;"""
-        return self._query_cortex_result_str(query)
-
-    def complete_for_conversation(
-        self,
-        conversation_json_file: SecurePath,
-        model: Model,
-    ) -> str:
-        json_content = conversation_json_file.read_text(
-            file_size_limit_mb=DEFAULT_SIZE_LIMIT_MB
-        )
-        query = f"""\
-            SELECT SNOWFLAKE.CORTEX.COMPLETE(
-                '{model}',
-                PARSE_JSON('{self._escape_input(json_content)}'),
-                {{}}
-            ) AS CORTEX_RESULT;"""
-        raw_result = self._query_cortex_result_str(query)
-        json_result = json.loads(raw_result)
-        return self._extract_text_result_from_json_result(
-            lambda: json_result["choices"][0]["messages"]
-        )
-
-    def extract_answer_from_source_document(
-        self,
-        source_document: SourceDocument,
-        question: Question,
-    ) -> str:
-        query = f"""\
-            SELECT SNOWFLAKE.CORTEX.EXTRACT_ANSWER(
-                '{self._escape_input(source_document)}',
-                '{self._escape_input(question)}'
-            ) AS CORTEX_RESULT;"""
-        raw_result = self._query_cortex_result_str(query)
-        json_result = json.loads(raw_result)
-        return self._extract_text_result_from_json_result(
-            lambda: json_result[0]["answer"]
-        )
-
-    def extract_answer_from_source_document_file(
-        self,
-        source_document_input_file: SecurePath,
-        question: Question,
-    ) -> str:
-        source_document_content = source_document_input_file.read_text(
-            file_size_limit_mb=DEFAULT_SIZE_LIMIT_MB
-        )
-        return self.extract_answer_from_source_document(
-            source_document=SourceDocument(source_document_content), question=question
-        )
-
-    def calculate_sentiment_for_text(
-        self,
-        text: Text,
-    ) -> str:
-        query = f"""\
-            SELECT SNOWFLAKE.CORTEX.SENTIMENT(
-                '{self._escape_input(text)}'
-            ) AS CORTEX_RESULT;"""
-        return self._query_cortex_result_str(query)
-
-    def calculate_sentiment_for_text_file(
-        self,
-        text_file: SecurePath,
-    ) -> str:
-        file_content = text_file.read_text(file_size_limit_mb=DEFAULT_SIZE_LIMIT_MB)
-        return self.calculate_sentiment_for_text(
-            text=Text(file_content),
-        )
-
-    def summarize_text(
-        self,
-        text: Text,
-    ) -> str:
-        query = f"""\
-            SELECT SNOWFLAKE.CORTEX.SUMMARIZE(
-                '{self._escape_input(text)}'
-            ) AS CORTEX_RESULT;"""
-        return self._query_cortex_result_str(query)
-
-    def summarize_text_file(
-        self,
-        text_file: SecurePath,
-    ) -> str:
-        file_content = text_file.read_text(file_size_limit_mb=DEFAULT_SIZE_LIMIT_MB)
-        return self.summarize_text(
-            text=Text(file_content),
-        )
-
-    def translate_text(
-        self,
-        text: Text,
-        source_language: Optional[Language],
-        target_language: Language,
-    ) -> str:
-        query = f"""\
-            SELECT SNOWFLAKE.CORTEX.TRANSLATE(
-                '{self._escape_input(text)}',
-                '{source_language or ""}',
-                '{target_language}'
-            ) AS CORTEX_RESULT;"""
-        return self._query_cortex_result_str(query)
-
-    def translate_text_file(
-        self,
-        text_file: SecurePath,
-        source_language: Optional[Language],
-        target_language: Language,
-    ) -> str:
-        file_content = text_file.read_text(file_size_limit_mb=DEFAULT_SIZE_LIMIT_MB)
-        return self.translate_text(
-            text=Text(file_content),
-            source_language=source_language,
-            target_language=target_language,
-        )
-
-    @staticmethod
-    def _escape_input(plain_input: str):
-        # escape backslashes to not escape too much, this replace has to be the first one
-        # escape single quotes because they are wrapping the whole string in SQL
-        return plain_input.replace("\\", "\\\\").replace("'", "\\'")
-
-    @staticmethod
-    def _extract_text_result_from_json_result(
-        extract_function: Callable[[], str]
-    ) -> str:
-        try:
-            return extract_function()
-        except (KeyError, IndexError) as ex:
-            log.debug("Cannot find Cortex result message in a response", exc_info=ex)
-            raise ClickException("Unexpected format of response from Snowflake")
-
-    def _query_cortex_result_str(self, query: str) -> str:
-        try:
-            cursor = self._execute_query(query, cursor_class=DictCursor)
-            if cursor.rowcount is None:
-                raise SnowflakeSQLExecutionError(query)
-            return str(cursor.fetchone()["CORTEX_RESULT"])
-        except ProgrammingError as ex:
-            log.debug("ProgrammingError occurred during SQL execution", exc_info=ex)
-            raise ClickException(str(ex))

snowflake/cli/plugins/cortex/plugin_spec.py

@@ -1,30 +0,0 @@
-# Copyright (c) 2024 Snowflake Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from snowflake.cli.api.plugins.command import (
-    SNOWCLI_ROOT_COMMAND_PATH,
-    CommandSpec,
-    CommandType,
-    plugin_hook_impl,
-)
-from snowflake.cli.plugins.cortex import commands
-
-
-@plugin_hook_impl
-def command_spec():
-    return CommandSpec(
-        parent_command_path=SNOWCLI_ROOT_COMMAND_PATH,
-        command_type=CommandType.COMMAND_GROUP,
-        typer_instance=commands.app.create_instance(),
-    )

snowflake/cli/plugins/cortex/types.py

@@ -1,22 +0,0 @@
-# Copyright (c) 2024 Snowflake Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import NewType
-
-Text = NewType("Text", str)
-Model = NewType("Model", str)
-Language = NewType("Language", str)
-InputString = NewType("InputString", str)
-Question = NewType("Question", str)
-SourceDocument = NewType("SourceDocument", str)

snowflake/cli/plugins/git/__init__.py

@@ -1,13 +0,0 @@
-# Copyright (c) 2024 Snowflake Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.