snowflake-cli-labs 2.3.0rc1__py3-none-any.whl → 2.4.0__py3-none-any.whl

snowflake/cli/plugins/cortex/commands.py

@@ -0,0 +1,312 @@
+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 SnowTyper
+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 = SnowTyper(
+    name="cortex",
+    help="Provides access to Snowflake Cortex.",
+)
+
+SEARCH_COMMAND_ENABLED = sys.version_info < (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

@@ -0,0 +1,3 @@
+from snowflake.cli.plugins.cortex.types import Model
+
+DEFAULT_MODEL: Model = Model("snowflake-arctic")

snowflake/cli/plugins/cortex/manager.py

@@ -0,0 +1,175 @@
+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

@@ -0,0 +1,16 @@
+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,
+    )

snowflake/cli/plugins/cortex/types.py

@@ -0,0 +1,8 @@
+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/commands.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import List, Optional
 
@@ -16,6 +18,10 @@ from snowflake.cli.api.constants import ObjectType
 from snowflake.cli.api.output.types import CollectionResult, CommandResult, QueryResult
 from snowflake.cli.api.utils.path_utils import is_stage_path
 from snowflake.cli.plugins.git.manager import GitManager
+from snowflake.cli.plugins.object.command_aliases import (
+    add_object_command_aliases,
+    scope_option,
+)
 from snowflake.cli.plugins.object.manager import ObjectManager
 from snowflake.cli.plugins.stage.commands import get
 from snowflake.cli.plugins.stage.manager import OnErrorType
@@ -50,6 +56,15 @@ RepoPathArgument = typer.Argument(
     ),
     callback=_repo_path_argument_callback,
 )
+add_object_command_aliases(
+    app=app,
+    object_type=ObjectType.GIT_REPOSITORY,
+    name_argument=RepoNameArgument,
+    like_option=like_option(
+        help_example='`list --like "my%"` lists all git repositories with name that begin with “my”',
+    ),
+    scope_option=scope_option(help_example="`list --in database my_db`"),
+)
 
 
 def _assure_repository_does_not_exist(om: ObjectManager, repository_name: str) -> None: