# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.

from __future__ import annotations

from typing import Dict, List, Union, Iterable, Optional
from typing_extensions import Literal, Required, TypeAlias, TypedDict

from .._types import SequenceNotStr
from .jinja_node_template_param import JinjaNodeTemplateParam

__all__ = [
    "NodeItemInputParam",
    "Config",
    "ConfigNodeConfig",
    "ConfigBatchedNodeConfig",
    "ConfigChunkEvaluationNodeConfig",
    "ConfigRerankerNodeConfig",
    "ConfigRetrieverNodeConfig",
    "ConfigCitationNodeConfig",
    "ConfigCitationNodeConfigCitationContext",
    "ConfigSearchCitationNodeConfig",
    "ConfigCodeExecutionConfig",
    "ConfigConditionNodeConfigInput",
    "ConfigDataTransformNodeConfig",
    "ConfigCreateMessagesNodeConfig",
    "ConfigCreateMessagesNodeConfigMessageConfig",
    "ConfigCreateMessagesNodeConfigMessageConfigAlternatingRoleMessages",
    "ConfigCreateMessagesNodeConfigMessageConfigSingleRoleMessages",
    "ConfigInsertMessagesConfig",
    "ConfigRemoveMessageConfig",
    "ConfigGetMessageConfig",
    "ConfigTokenizerChatTemplateConfig",
    "ConfigGenerationNodeConfig",
    "ConfigGenerationNodeConfigRetryConfig",
    "ConfigChatGenerationNodeConfig",
    "ConfigChatGenerationNodeConfigRetryConfig",
    "ConfigGenerationWithCitationsNodeConfig",
    "ConfigGenerationWithCitationsNodeConfigRetryConfig",
    "ConfigChatGenerationWithCitationsNodeConfig",
    "ConfigChatGenerationWithCitationsNodeConfigRetryConfig",
    "ConfigResponseParserNodeConfig",
    "ConfigJinjaNodeConfig",
    "ConfigProcessingNodeConfig",
    "ConfigProcessingNodeConfigFunctionSpecs",
    "ConfigRegexMatchNodeConfig",
    "ConfigSqlExecutorNodeConfig",
    "ConfigStaticNodeConfig",
    "ConfigLlmEngineNodeConfig",
    "ConfigLlmEngineNodeConfigBatchSysKwargs",
]


class ConfigNodeConfig(TypedDict, total=False):
    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigBatchedNodeConfig(TypedDict, total=False):
    batch_size: Required[int]
    """The size of the batches to split the input into."""

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    split_approx_evenly: bool
    """Whether to split the input into batches of approximately equal size."""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigChunkEvaluationNodeConfig(TypedDict, total=False):
    top_k_thresholds: Required[Iterable[int]]
    """The top-k values you'd like to evaluate retrieval against."""

    fuzzy_match_threshold: float
    """The threshold for fuzzy string matching."""

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    require_all: bool
    """Whether all sources must be present in the top k chunks."""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigRerankerNodeConfig(TypedDict, total=False):
    num_to_return: Required[int]
    """The number of chunks to return after reranking."""

    scorers: Required[Iterable[object]]
    """The scorers to use to rerank the chunks.

    It is a dict with one of the following keys: dict_keys(['bm25', 'recency',
    'cross-encoder', 'sgp-custom-model', 'llm-reranker'])
    """

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    score_threshold: Optional[float]
    """The threshold of the score to filter the chunks."""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigRetrieverNodeConfig(TypedDict, total=False):
    num_to_return: Required[int]
    """The number of chunks to return after retrieval."""

    exact_knn_search: Optional[bool]
    """Whether to use exact KNN search.

    If not provided, will be set to True if num_to_return < 1000.
    """

    knowledge_base_id: Optional[str]
    """The ID of the knowledge base to query."""

    knowledge_base_name: Optional[str]
    """The name of the knowledge base to query."""

    metadata: Optional[Dict[str, Optional[str]]]
    """Metadata filters to apply to the chunks retrieved."""

    min_results_per_knowledge_base: int
    """The minimum number of results to return per knowledge base."""

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_retriever_workers: int
    """The number of concurrent retriever workers to use."""

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigCitationNodeConfigCitationContext(TypedDict, total=False):
    """The citation context (e.g. configuration of how make the citations) to use."""

    generate_with_llm: bool
    """
    When True, we prompt the LLM to include citations in it's response and parse
    them during post-processing rather than identifying cited chunks entirely in
    post-processing like we do for other methods.
    """

    metric: Optional[str]
    """The ROUGE metric to use (rouge2 by default)"""

    min_similarity: Optional[float]
    """The similarity threshold at which a candidate match is returned as a citation"""

    score: Optional[Literal["precision", "recall", "fmeasure"]]
    """The value to use with ROUGE citation matching (recall is the default)"""


class ConfigCitationNodeConfig(TypedDict, total=False):
    """A data model describing parameters for back-citation using ROUGE similarity.

    metric is the ROUGE metric to use (e.g. rouge1, rouge2, rougeLsum)
    score is one of "precision", "recall", "fmeasure"

    NOTE (john): copied directly from generation.py in order to subclass from NodeConfig.
    """

    citation_type: Required[Literal["rouge", "model_defined"]]
    """The type of citation to use."""

    citation_context: ConfigCitationNodeConfigCitationContext
    """The citation context (e.g. configuration of how make the citations) to use."""

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    s3_path_override: Optional[str]
    """
    If provided, the S3 path to use for the attachment URL in the citation to deal
    with auth challenges.
    """

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigSearchCitationNodeConfig(TypedDict, total=False):
    end_search_regex: Required[str]
    """The regex to search for the end of the search result in the response."""

    search_regex: Required[str]
    """The regex to search for the search result in the response."""

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigCodeExecutionConfig(TypedDict, total=False):
    files: Required[Dict[str, str]]
    """A mapping of sandboxed file paths to their source paths.

    Keys should be relative paths at which the files will be saved within the
    sandbox, and values should be source file identifiers compatible with
    smart_open.
    """

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    output_artifacts_dir: Optional[str]
    """If set, output files will be saved to this directory"""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigConditionNodeConfigInput(TypedDict, total=False):
    condition: Required["CompoundConditionInputParam"]
    """
    A compound condition that will be evaluated to determine whether the node should
    run. The condition has the following structure: logical_operator: Literal["ALL",
    "ANY", "NOT"] = "ALL" conditions: Optional[List[Union[UnaryCondition,
    "CompoundCondition"]]] = None input_names: List[str] = None # populated by model
    validator

        Each condition in the conditions list can be either a UnaryCondition or a CompoundCondition. A Unary Condition looks like;
            Representation of a boolean function with a single input
            e.g. the condition specified by
                input_name: x
                operator: 'contains'
                ref_value: 'c'
            would evaluate to True if x == 'cat'
            Operators are defined in the constant function store dict_keys(['greater_than', 'equals', 'less_than', 'is_in', 'contains', 'has_duplicates', 'is_null', 'bool'])
    """

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigDataTransformNodeConfig(TypedDict, total=False):
    action: Required[
        Literal[
            "load_json_chunk_vals",
            "to_key_value_dict",
            "get_json_val",
            "to_list",
            "strip",
            "df_sample_str",
            "till_stop_token",
            "identity",
            "select_keys_from_dict",
            "int_adder",
            "json_loads",
            "clean_sql",
            "len",
            "index",
            "slice",
            "get_pydantic_attribute",
            "get",
            "read_text_file",
            "split_to_chunks_nltk",
        ]
    ]
    """The action to perform on the input data.

    Must be one of the keys in dict_keys(['load_json_chunk_vals',
    'to_key_value_dict', 'get_json_val', 'to_list', 'strip', 'df_sample_str',
    'till_stop_token', 'identity', 'select_keys_from_dict', 'int_adder',
    'json_loads', 'clean_sql', 'len', 'index', 'slice', 'get_pydantic_attribute',
    'get', 'read_text_file', 'split_to_chunks_nltk']).
    """

    additional_inputs: object
    """Additional inputs for the action."""

    apply_to_dictlist_leaves: bool
    """Whether the action should be applied to the leaves of a nested dict/list."""

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigCreateMessagesNodeConfigMessageConfigAlternatingRoleMessages(TypedDict, total=False):
    """A class representing multiple roles that alternate between each other"""

    role_value_pairs: Required[Iterable[Dict[str, str]]]
    """
    A list of dictionaries with a role and a value (where value is a string or a
    string representing an input)
    """


class ConfigCreateMessagesNodeConfigMessageConfigSingleRoleMessages(TypedDict, total=False):
    """A class represent a single role. It is a dictionary with a role"""

    content: Required[str]
    """This can either be a string representing the content (e.g.

    system message) or a string representing an input (either a str or a list of
    strings)
    """

    role: Required[str]
    """The role of the message (e.g. user, assistant, system)"""


ConfigCreateMessagesNodeConfigMessageConfig: TypeAlias = Union[
    ConfigCreateMessagesNodeConfigMessageConfigAlternatingRoleMessages,
    ConfigCreateMessagesNodeConfigMessageConfigSingleRoleMessages,
]


class ConfigCreateMessagesNodeConfig(TypedDict, total=False):
    """
    A config to create a list of messages that can be ingested in various chat templates
    """

    message_configs: Required[Iterable[ConfigCreateMessagesNodeConfigMessageConfig]]
    """
    We take as input a list of configs that represent either a single user or two
    users we want to alternate between
    """

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigInsertMessagesConfig(TypedDict, total=False):
    """
    Config to determine how you to insert a list of messages at a certain index in a list of messages

    This is most commonly useful when adding a system message to convo thus far or adding the user query to the convo
    """

    index: Required[int]
    """The index to insert the new messages at.

    If index = -1, then the messages will be appended to the end of the list
    """

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigRemoveMessageConfig(TypedDict, total=False):
    """Config to determine how you to remove a message from a list of messages

    This is most helpful when you are trying to edit a message (e.g. replace a user question with a new string that includes retrieved context)
    """

    index: Required[int]
    """The index of the message to remove.

    If index = -1, then the last message will be removed
    """

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigGetMessageConfig(TypedDict, total=False):
    index: Required[int]
    """The index of the message to get"""

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigTokenizerChatTemplateConfig(TypedDict, total=False):
    llm_model: Required[
        Literal[
            "llama-2-7b-chat",
            "llama-2-13b-chat",
            "llama-2-70b-chat",
            "falcon-7b-insruct",
            "falcon-40b-instruct",
            "mistral-7b-instruct",
            "mixtral-8x7b-instruct",
            "mixtral-8x22b-instruct",
            "codellama-7b-instruct",
            "codellama-13b-instruct",
            "codellama-34b-instruct",
            "codellama-70b-instruct",
            "zephyr-7b-alpha",
            "zephyr-7b-beta",
            "gemma-2b-instruct",
            "gemma-7b-instruct",
            "llama-3-8b-instruct",
            "llama-3-70b-instruct",
            "llama-3-1-8b-instruct",
            "llama-3-1-70b-instruct",
            "llama-3-1-405b-instruct",
            "dolphin-2.2-70b",
            "dolphin-2.9-llama3-70b",
            "dolphin-2.9-llama3-8b",
        ]
    ]
    """The name of the model to use for tokenization."""

    add_generation_prompt: bool
    """
    Whether to end the prompt with the token(s) that indicate the start of an
    assistant message. This is useful when you want to generate a response from the
    model. Note that this argument will be passed to the chat template, and so it
    must be supported in the template for this argument to have any effect.
    """

    kwargs: object
    """Additional kwargs to pass to the template renderer.

    Will be accessible by the chat template.
    """

    max_length: Optional[int]
    """Maximum length (in tokens) to use for padding or truncation.

    Has no effect if tokenize is `False`.
    """

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    padding: bool
    """Whether to pad sequences to the maximum length.

    Has no effect if tokenize is `False`.
    """

    truncation: bool
    """Whether to truncate sequences at the maximum length.

    Has no effect if tokenize is `False`.
    """

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigGenerationNodeConfigRetryConfig(TypedDict, total=False):
    """The retry configuration for the node.

    Can be used to specify the number of retries, delay, backoff, and exceptions to retry on.
    """

    backoff: int
    """The factor by which to increase the delay between retries."""

    delay: int
    """The number of seconds to wait before retrying the operation."""

    exceptions: List[Literal["SGPClientError", "APITimeoutError", "InternalServerError", "RateLimitError", "Exception"]]
    """The exceptions to retry on."""

    tries: int
    """The number of times to retry the operation before giving up."""


class ConfigGenerationNodeConfig(TypedDict, total=False):
    """A data model describing LLM generation parameters."""

    llm_model: str
    """The name of the LLM model to use for generation.

    If not provided, llm_model_instance and llm_model_deployment must be set.
    """

    llm_model_deployment: Optional[str]
    """The name of the custom LLM model deployment to use for generation.

    If not provided, llm_model must be set.
    """

    llm_model_instance: Optional[str]
    """The name of the custom LLM model instance to use for generation.

    If not provided, llm_model must be set.
    """

    max_tokens: int
    """The maximum number of tokens to generate in the completion.

    The maximum value is 2048.
    """

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    retry_config: ConfigGenerationNodeConfigRetryConfig
    """The retry configuration for the node.

    Can be used to specify the number of retries, delay, backoff, and exceptions to
    retry on.
    """

    stop_sequences: Optional[SequenceNotStr[str]]
    """A list of strings that will stop the completion when encountered."""

    strip_whitespace: bool
    """Whether to strip leading and trailing whitespace from the response."""

    temperature: float
    """The sampling temperature.

    Lower values will result in more deterministic completions.
    """

    tool_name: Optional[str]
    """The name of the tool to use for generation.

    Will automatically get invoked. If not provided, no tool will be used.
    """

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigChatGenerationNodeConfigRetryConfig(TypedDict, total=False):
    """The retry configuration for the node.

    Can be used to specify the number of retries, delay, backoff, and exceptions to retry on.
    """

    backoff: int
    """The factor by which to increase the delay between retries."""

    delay: int
    """The number of seconds to wait before retrying the operation."""

    exceptions: List[Literal["SGPClientError", "APITimeoutError", "InternalServerError", "RateLimitError", "Exception"]]
    """The exceptions to retry on."""

    tries: int
    """The number of times to retry the operation before giving up."""


class ConfigChatGenerationNodeConfig(TypedDict, total=False):
    """A data model describing LLM generation parameters for chat completions."""

    memory_strategy: Required[object]
    """
    The memory strategy to prevent exceeding the LLM's context limit from being
    exceeded. Currently SGP only supports the Last K memory strategy (keep the last
    k messages), but will be adding new strategies soon. (e.g. {'type': 'last_k',
    'params: {'k': 1000} })
    """

    instructions: Optional[str]
    """Instructions for the system prompt.

    If not provided, there will be no system message. It can also be provided at
    runtime.
    """

    llm_model: str
    """The name of the LLM model to use for generation.

    If not provided, llm_model_instance and llm_model_deployment must be set.
    """

    llm_model_deployment: Optional[str]
    """The name of the custom LLM model deployment to use for generation.

    If not provided, llm_model must be set.
    """

    llm_model_instance: Optional[str]
    """The name of the custom LLM model instance to use for generation.

    If not provided, llm_model must be set.
    """

    max_tokens: int
    """The maximum number of tokens to generate in the completion.

    The maximum value is 2048.
    """

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    retry_config: ConfigChatGenerationNodeConfigRetryConfig
    """The retry configuration for the node.

    Can be used to specify the number of retries, delay, backoff, and exceptions to
    retry on.
    """

    stop_sequences: Optional[SequenceNotStr[str]]
    """A list of strings that will stop the completion when encountered."""

    strip_whitespace: bool
    """Whether to strip leading and trailing whitespace from the response."""

    temperature: float
    """The sampling temperature.

    Lower values will result in more deterministic completions.
    """

    tool_name: Optional[str]
    """The name of the tool to use for generation.

    Will automatically get invoked. If not provided, no tool will be used.
    """

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigGenerationWithCitationsNodeConfigRetryConfig(TypedDict, total=False):
    """The retry configuration for the node.

    Can be used to specify the number of retries, delay, backoff, and exceptions to retry on.
    """

    backoff: int
    """The factor by which to increase the delay between retries."""

    delay: int
    """The number of seconds to wait before retrying the operation."""

    exceptions: List[Literal["SGPClientError", "APITimeoutError", "InternalServerError", "RateLimitError", "Exception"]]
    """The exceptions to retry on."""

    tries: int
    """The number of times to retry the operation before giving up."""


class ConfigGenerationWithCitationsNodeConfig(TypedDict, total=False):
    regex_pattern: Required[str]
    """The regex pattern to match citations in the text.

    The first group must be the citation index!
    """

    regex_replace: Required[str]
    """The format string to replace the citation with.

    The first argument is the citation index and the second is the fragment UUID.
    """

    llm_model: str
    """The name of the LLM model to use for generation.

    If not provided, llm_model_instance and llm_model_deployment must be set.
    """

    llm_model_deployment: Optional[str]
    """The name of the custom LLM model deployment to use for generation.

    If not provided, llm_model must be set.
    """

    llm_model_instance: Optional[str]
    """The name of the custom LLM model instance to use for generation.

    If not provided, llm_model must be set.
    """

    max_tokens: int
    """The maximum number of tokens to generate in the completion.

    The maximum value is 2048.
    """

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    retry_config: ConfigGenerationWithCitationsNodeConfigRetryConfig
    """The retry configuration for the node.

    Can be used to specify the number of retries, delay, backoff, and exceptions to
    retry on.
    """

    stop_sequences: Optional[SequenceNotStr[str]]
    """A list of strings that will stop the completion when encountered."""

    strip_whitespace: bool
    """Whether to strip leading and trailing whitespace from the response."""

    temperature: float
    """The sampling temperature.

    Lower values will result in more deterministic completions.
    """

    tool_name: Optional[str]
    """The name of the tool to use for generation.

    Will automatically get invoked. If not provided, no tool will be used.
    """

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigChatGenerationWithCitationsNodeConfigRetryConfig(TypedDict, total=False):
    """The retry configuration for the node.

    Can be used to specify the number of retries, delay, backoff, and exceptions to retry on.
    """

    backoff: int
    """The factor by which to increase the delay between retries."""

    delay: int
    """The number of seconds to wait before retrying the operation."""

    exceptions: List[Literal["SGPClientError", "APITimeoutError", "InternalServerError", "RateLimitError", "Exception"]]
    """The exceptions to retry on."""

    tries: int
    """The number of times to retry the operation before giving up."""


class ConfigChatGenerationWithCitationsNodeConfig(TypedDict, total=False):
    memory_strategy: Required[object]
    """
    The memory strategy to prevent exceeding the LLM's context limit from being
    exceeded. Currently SGP only supports the Last K memory strategy (keep the last
    k messages), but will be adding new strategies soon. (e.g. {'type': 'last_k',
    'params: {'k': 1000} })
    """

    regex_pattern: Required[str]
    """The regex pattern to match citations in the text.

    The first group must be the citation index!
    """

    regex_replace: Required[str]
    """The format string to replace the citation with.

    The first argument is the citation index and the second is the fragment UUID.
    """

    instructions: Optional[str]
    """Instructions for the system prompt.

    If not provided, there will be no system message. It can also be provided at
    runtime.
    """

    llm_model: str
    """The name of the LLM model to use for generation.

    If not provided, llm_model_instance and llm_model_deployment must be set.
    """

    llm_model_deployment: Optional[str]
    """The name of the custom LLM model deployment to use for generation.

    If not provided, llm_model must be set.
    """

    llm_model_instance: Optional[str]
    """The name of the custom LLM model instance to use for generation.

    If not provided, llm_model must be set.
    """

    max_tokens: int
    """The maximum number of tokens to generate in the completion.

    The maximum value is 2048.
    """

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    retry_config: ConfigChatGenerationWithCitationsNodeConfigRetryConfig
    """The retry configuration for the node.

    Can be used to specify the number of retries, delay, backoff, and exceptions to
    retry on.
    """

    stop_sequences: Optional[SequenceNotStr[str]]
    """A list of strings that will stop the completion when encountered."""

    strip_whitespace: bool
    """Whether to strip leading and trailing whitespace from the response."""

    temperature: float
    """The sampling temperature.

    Lower values will result in more deterministic completions.
    """

    tool_name: Optional[str]
    """The name of the tool to use for generation.

    Will automatically get invoked. If not provided, no tool will be used.
    """

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigResponseParserNodeConfig(TypedDict, total=False):
    action: Required[
        Literal[
            "greater_than",
            "equals",
            "less_than",
            "is_in",
            "contains",
            "has_duplicates",
            "is_null",
            "bool",
            "to_integer",
            "to_key_value_dict",
        ]
    ]
    """The action to perform on the response.

    Must be one of the keys: dict_keys(['greater_than', 'equals', 'less_than',
    'is_in', 'contains', 'has_duplicates', 'is_null', 'bool', 'to_integer',
    'to_key_value_dict'])
    """

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    reference_value: object
    """A reference value to use in the action.

    The meaning of this value depends on the action.
    """

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigJinjaNodeConfig(TypedDict, total=False):
    context_chunks_key: Optional[str]
    """
    The key in the templated data that contains the context chunks to be
    token-limited.
    """

    data_transformations: Dict[str, JinjaNodeTemplateParam]
    """A dictionary of data transformations to apply to the input data."""

    llm_model: Optional[
        Literal[
            "llama-2-7b-chat",
            "llama-2-13b-chat",
            "llama-2-70b-chat",
            "falcon-7b-insruct",
            "falcon-40b-instruct",
            "mistral-7b-instruct",
            "mixtral-8x7b-instruct",
            "mixtral-8x22b-instruct",
            "codellama-7b-instruct",
            "codellama-13b-instruct",
            "codellama-34b-instruct",
            "codellama-70b-instruct",
            "zephyr-7b-alpha",
            "zephyr-7b-beta",
            "gemma-2b-instruct",
            "gemma-7b-instruct",
            "llama-3-8b-instruct",
            "llama-3-70b-instruct",
            "llama-3-1-8b-instruct",
            "llama-3-1-70b-instruct",
            "llama-3-1-405b-instruct",
            "dolphin-2.2-70b",
            "dolphin-2.9-llama3-70b",
            "dolphin-2.9-llama3-8b",
        ]
    ]
    """The name of the model to use for tokenization."""

    log_output: bool
    """Whether to log the output of the node to the logger."""

    log_prefix: str
    """A prefix to add to the log output.

    Useful for distinguishing between multiple nodes.
    """

    max_tokens: Optional[int]
    """The maximum number of tokens to use for the output template.

    If set, the output template will be token-limited.
    """

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    output_template: JinjaNodeTemplateParam
    """The output template to use."""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""

    verbose: bool
    """Whether to log additional information about the node."""


class ConfigProcessingNodeConfigFunctionSpecs(TypedDict, total=False):
    """
    Define the specification of function execution
    Key: A a unique name for the execution of a function and the kwargs.
    Values: A dictionary mapping:
        - path (str): path to the function
        - kwargs (Dict[str, Any]): Keyword arguments for this execution of the function.

    Function Key & Path Representation:

        1. Basic Usage
            {
                "first_usage": {
                    "path": "package.module.function",
                    ... # your kwargs
                },
                "second_usage": {
                    "path": "package.module.function",
                    ... # your kwargs
                }
            }

        2. Passing results between functions
            {
                "first_function": {
                    ... # your kwargs
                },
                "second_function" " {
                    input_from_first_function: "first_function"
                }
            }
    """

    kwargs: Required[object]
    """Keyword arguments for the function."""

    path: Required[str]
    """Path to the function."""


class ConfigProcessingNodeConfig(TypedDict, total=False):
    function_specs: Required[Dict[str, ConfigProcessingNodeConfigFunctionSpecs]]
    """Specifications for executing a function."""

    return_key: Required[str]
    """The function to return results from."""

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigRegexMatchNodeConfig(TypedDict, total=False):
    pattern: Required[str]
    """The regex pattern to match against."""

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigSqlExecutorNodeConfig(TypedDict, total=False):
    connector_kwargs: Required[Dict[str, str]]
    """A dictionary of keyword arguments to pass to the DB connector"""

    connector_type: Literal["snowflake"]
    """The type of connector to use"""

    log_queries: bool
    """Whether to log the queries executed by the node"""

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    return_type: Literal["df", "dicts", "markdown", "json", "str"]
    """The type of return value to return"""

    schema_remapping_file: Optional[str]
    """A file containing a schema remapping for tables"""

    secrets: SequenceNotStr[str]
    """A list of keys in connector_kwargs that should be fetched from secrets"""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


class ConfigStaticNodeConfig(TypedDict, total=False):
    from_file: Union[Iterable[object], str, object, None]
    """A local/s3 path or a nested dict/list where each leaf is a local/s3 path.

    Use when you need to compile multiple data sources into one.
    """

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""

    value: object
    """The value to return from the node.

    Exactly one of `value` and `from_file` must be provided.
    """


class ConfigLlmEngineNodeConfigBatchSysKwargs(TypedDict, total=False):
    """System-level kwargs for batch completions."""

    checkpoint_path: Optional[str]
    """The checkpoint path to use for the model."""

    labels: Optional[Dict[str, str]]
    """The labels to use for the model (e.g. team/product)"""

    num_shards: Optional[int]
    """The number of shards to split the model into"""

    seed: Optional[int]
    """The seed to use for the model's generation."""


class ConfigLlmEngineNodeConfig(TypedDict, total=False):
    """Data model for LLMEngine-based generation calls"""

    frequency_penalty: Required[Optional[float]]
    """The frequency penalty to use for generation."""

    llm_model: Required[str]
    """The LLM model to use for generation. This should be a Scale-hosted LLM model."""

    presence_penalty: Required[Optional[float]]
    """The presence penalty to use for generation."""

    top_k: Required[Optional[int]]
    """The top-k value to use for generation."""

    top_p: Required[Optional[float]]
    """The top-p value to use for generation."""

    batch_run_mode: Literal["sync", "async"]
    """Whether to run the batch completions synchronously or asynchronously."""

    batch_sys_kwargs: ConfigLlmEngineNodeConfigBatchSysKwargs
    """System-level kwargs for batch completions."""

    guided_choice: Optional[SequenceNotStr[str]]
    """A list of strings to guide the completion."""

    guided_json: Optional[object]
    """A JSON object to guide the completion."""

    guided_regex: Optional[str]
    """A regex string to guide the completion."""

    include_stop_str_in_output: Optional[bool]
    """Whether to include the stop string in the output."""

    max_tokens: Optional[int]
    """The maximum number of tokens to generate in the completion."""

    node_metadata: Optional[SequenceNotStr[str]]
    """Metadata to attach to the output.

    Can be any of {'\\__output_token_usage', '\\__start_timestamp',
    '\\__input_token_usage', '\\__time_taken'}.
    """

    num_workers: Optional[int]
    """The number of workers to use for parallel processing."""

    stop_sequences: Optional[SequenceNotStr[str]]
    """A list of strings that, if generated, will stop the completion."""

    temperature: Optional[float]
    """The temperature to use for generation. Higher values lead to more randomness."""

    timeout: int
    """The maximum time in seconds to wait for the completion.

    Used for max_runtime_sec in batch completions call.
    """

    type_hints: Optional[object]
    """Type hints for the input parameters of the node in JSON schema."""


Config: TypeAlias = Union[
    ConfigNodeConfig,
    ConfigBatchedNodeConfig,
    ConfigChunkEvaluationNodeConfig,
    ConfigRerankerNodeConfig,
    ConfigRetrieverNodeConfig,
    ConfigCitationNodeConfig,
    ConfigSearchCitationNodeConfig,
    ConfigCodeExecutionConfig,
    ConfigConditionNodeConfigInput,
    ConfigDataTransformNodeConfig,
    ConfigCreateMessagesNodeConfig,
    ConfigInsertMessagesConfig,
    ConfigRemoveMessageConfig,
    ConfigGetMessageConfig,
    ConfigTokenizerChatTemplateConfig,
    ConfigGenerationNodeConfig,
    ConfigChatGenerationNodeConfig,
    ConfigGenerationWithCitationsNodeConfig,
    ConfigChatGenerationWithCitationsNodeConfig,
    ConfigResponseParserNodeConfig,
    ConfigJinjaNodeConfig,
    ConfigProcessingNodeConfig,
    ConfigRegexMatchNodeConfig,
    ConfigSqlExecutorNodeConfig,
    ConfigStaticNodeConfig,
    ConfigLlmEngineNodeConfig,
]


class NodeItemInputParam(TypedDict, total=False):
    """
    A base model purely for validating that each component of a
    workflow config has the requisite parts (name, type, config, inputs).
    Checks that the provided type is allowed as per NodeTypes
    and that the provided config is of the right NodeConfig subtype.

    Attribute Details:
        inputs: inputs can be a mapping from:
                    1) str to str
                    2) str to dict(str, str)
                    3) str to dict(str, dict(str, str))
    """

    config: Required[Config]
    """A data model describing parameters for back-citation using ROUGE similarity.

    metric is the ROUGE metric to use (e.g. rouge1, rouge2, rougeLsum) score is one
    of "precision", "recall", "fmeasure"

    NOTE (john): copied directly from generation.py in order to subclass from
    NodeConfig.
    """

    name: Required[str]

    type: Required[str]

    inputs: Dict[str, Union[str, Dict[str, Union[str, object]]]]


from .compound_condition_input_param import CompoundConditionInputParam
