Config

@synalinks_export(
    [
        "synalinks.utils.clear_session",
        "synalinks.backend.clear_session",
        "synalinks.clear_session",
    ]
)
def clear_session(free_memory=True):
    """Resets all state generated by synalinks.

    synalinks manages a global state, which it uses to implement the Functional
    model-building API and to uniquify autogenerated modules names.

    If you are creating many models in a loop, this global state will consume
    an increasing amount of memory over time, and you may want to clear it.
    Calling `clear_session()` releases the global state: this helps avoid
    clutter from old programs and modules, especially when memory is limited.

    Args:
        free_memory (bool): Whether to call Python garbage collection.
            It's usually a good practice to call it to make sure
            memory used by deleted objects is immediately freed.
            However, it may take a few seconds to execute, so
            when using `clear_session()` in a short loop,
            you may want to skip it.
    """
    global GLOBAL_STATE_TRACKER
    global GLOBAL_SETTINGS_TRACKER

    GLOBAL_STATE_TRACKER = threading.local()
    GLOBAL_SETTINGS_TRACKER = threading.local()

    if free_memory:
        # Manually trigger garbage collection.
        gc.collect()

@synalinks_export(
    [
        "synalinks.config.api_base",
        "synalinks.backend.api_base",
        "synalinks.api_base",
    ]
)
def api_base():
    """Returns the Synalinks api base

    Returns:
        (str): The observability api base
    """
    return _SYNALINKS_API_BASE

@synalinks_export(
    [
        "synalinks.config.api_key",
        "synalinks.backend.api_key",
    ]
)
def api_key():
    """Synalinks API key.

    Returns:
        (str): Synalinks API key.

    ```python
    >>> synalinks.config.api_key()
    'my-secret-api-key'
    ```

    """
    return _SYNALINKS_API_KEY

@synalinks_export(["synalinks.config.backend", "synalinks.backend.backend"])
def backend():
    """Publicly accessible method for determining the current backend.

    Returns:
        (str): The name of the backend synalinks is currently using. like
            `"pydantic"`.

    Example:

    ```python
    >>> synalinks.config.backend()
    'pydantic'
    ```
    """
    return _BACKEND

@synalinks_export(
    [
        "synalinks.config.default_embedding_model",
        "synalinks.default_embedding_model",
    ]
)
def default_embedding_model():
    """Return the default `EmbeddingModel` instance, or `None` if unset."""
    global _DEFAULT_EMBEDDING_MODEL
    if (
        _DEFAULT_EMBEDDING_MODEL is None
        and _DEFAULT_EMBEDDING_MODEL_IDENTIFIER is not None
    ):
        from synalinks.src.modules.embedding_models import get as _get_em

        _DEFAULT_EMBEDDING_MODEL = _get_em(_DEFAULT_EMBEDDING_MODEL_IDENTIFIER)
    return _DEFAULT_EMBEDDING_MODEL

@synalinks_export(
    [
        "synalinks.config.default_knowledge_base",
        "synalinks.default_knowledge_base",
    ]
)
def default_knowledge_base():
    """Return the default `KnowledgeBase` instance, or `None` if unset."""
    global _DEFAULT_KNOWLEDGE_BASE
    if _DEFAULT_KNOWLEDGE_BASE is None and _DEFAULT_KNOWLEDGE_BASE_IDENTIFIER is not None:
        from synalinks.src.knowledge_bases import get as _get_kb

        _DEFAULT_KNOWLEDGE_BASE = _get_kb(_DEFAULT_KNOWLEDGE_BASE_IDENTIFIER)
    return _DEFAULT_KNOWLEDGE_BASE

@synalinks_export(
    [
        "synalinks.config.default_language_model",
        "synalinks.default_language_model",
    ]
)
def default_language_model():
    """Return the default `LanguageModel` instance, or `None` if unset.

    The instance is constructed lazily on first call when set from a
    persisted identifier (e.g. via `~/.synalinks/synalinks.json`).
    """
    global _DEFAULT_LANGUAGE_MODEL
    if _DEFAULT_LANGUAGE_MODEL is None and _DEFAULT_LANGUAGE_MODEL_IDENTIFIER is not None:
        from synalinks.src.modules.language_models import get as _get_lm

        _DEFAULT_LANGUAGE_MODEL = _get_lm(_DEFAULT_LANGUAGE_MODEL_IDENTIFIER)
    return _DEFAULT_LANGUAGE_MODEL

@synalinks_export(
    [
        "synalinks.config.enable_logging",
        "synalinks.backend.enable_logging",
        "synalinks.enable_logging",
    ]
)
def enable_logging(log_level="debug", log_to_file=False):
    """
    Configures and enables logging for the application.

    This function sets up the logging configuration for the application,
    allowing logs to be output either to a specified file or to the console.
    The logging level can be set to DEBUG for more verbose logging or INFO
    for standard logging.

    Args:
        log_level (str): The logging level.
        log_to_file (bool): If True save the logs into `synalinks.log`

    The log message format includes the timestamp, log level, and the log message itself.
    """
    global logger

    log_level = log_level.upper()
    if log_level and hasattr(logging, log_level):
        log_level = getattr(logging, log_level)

    logger.setLevel(log_level)

    for handler in logger.handlers[:]:
        logger.removeHandler(handler)

    stream_handler = logging.StreamHandler()
    stream_handler.setLevel(log_level)
    stream_handler.setFormatter(SynalinksLogFormatter())
    logger.addHandler(stream_handler)

    if log_to_file:
        file_handler = logging.FileHandler("synalinks.log", mode="w")
        file_handler.setLevel(log_level)
        formatter = SynalinksLogFormatter("%(asctime)s - %(levelname)s - %(message)s")
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)

@synalinks_export(
    [
        "synalinks.config.enable_observability",
        "synalinks.backend.enable_observability",
        "synalinks.enable_observability",
    ]
)
def enable_observability(tracking_uri=None, experiment_name=None):
    """
    Configures and enables observability for the application using MLflow.

    This function sets up the observability configuration for the application,
    enabling tracing of module calls via MLflow.

    Args:
        tracking_uri (str): Optional. The MLflow tracking server URI.
            If not provided, MLflow will use the default (local ./mlruns
            directory or MLFLOW_TRACKING_URI environment variable).
        experiment_name (str): Optional. The MLflow experiment name.
            Defaults to "synalinks_traces".

    Example:

    ```python
    import synalinks

    # Basic usage with defaults
    synalinks.enable_observability()

    # With custom MLflow tracking server
    synalinks.enable_observability(
        tracking_uri="http://localhost:5000",
        experiment_name="my_experiment"
    )
    ```
    """
    global _MLFLOW_TRACKING_URI
    global _MLFLOW_EXPERIMENT_NAME
    global _ENABLE_OBSERVABILITY

    if tracking_uri:
        _MLFLOW_TRACKING_URI = tracking_uri
    if experiment_name:
        _MLFLOW_EXPERIMENT_NAME = experiment_name
    _ENABLE_OBSERVABILITY = True

@synalinks_export(["synalinks.config.epsilon", "synalinks.backend.epsilon"])
def epsilon():
    """Return the value of the fuzz factor used in numeric expressions.

    Returns:
        (float): The epsilon value.

    Example:

    ```python
    >>> synalinks.config.epsilon()
    1e-07
    ```

    """
    return _EPSILON

@synalinks_export(["synalinks.config.floatx", "synalinks.backend.floatx"])
def floatx():
    """Return the default float type, as a string.

    E.g. `'bfloat16'`, `'float16'`, `'float32'`, `'float64'`.

    Returns:
        (str): The current default float type.

    Example:

    ```python
    >>> synalinks.config.floatx()
    'float32'
    ```

    """
    return _FLOATX

@synalinks_export(
    [
        "synalinks.config.is_observability_enabled",
        "synalinks.backend.is_observability_enabled",
        "synalinks.is_observability_enabled",
    ]
)
def is_observability_enabled():
    """Check if the observability is enabled

    Returns:
        (bool): True if the observability is enabled.
    """
    return _ENABLE_OBSERVABILITY

@synalinks_export(
    [
        "synalinks.config.is_trace_recording_enabled",
        "synalinks.backend.is_trace_recording_enabled",
        "synalinks.is_trace_recording_enabled",
    ]
)
def is_trace_recording_enabled():
    """Check if the trace recording is enabled

    Returns:
        (bool): True if the trace recording is enabled.
    """
    return _ENABLE_TRACE_RECORDING

@synalinks_export(
    [
        "synalinks.config.mlflow_experiment_name",
        "synalinks.backend.mlflow_experiment_name",
    ]
)
def mlflow_experiment_name():
    """Returns the MLflow experiment name for observability.

    Returns:
        (str): The MLflow experiment name.
    """
    return _MLFLOW_EXPERIMENT_NAME

@synalinks_export(
    [
        "synalinks.config.mlflow_tracking_uri",
        "synalinks.backend.mlflow_tracking_uri",
    ]
)
def mlflow_tracking_uri():
    """Returns the MLflow tracking URI for observability.

    Returns:
        (str): The MLflow tracking URI, or None if not set.
    """
    return _MLFLOW_TRACKING_URI

@synalinks_export(
    [
        "synalinks.config.record_traces",
        "synalinks.backend.record_traces",
        "synalinks.record_traces",
    ]
)
def record_traces(base_dir=None):
    """Enables trace recording of `LanguageModel` calls.

    This function enables the `Recorder` hook for every module, recording
    each LM call (chat messages, completion, token usage, cost ...) as one
    JSON line under `base_dir`, organized as one folder per program and one
    subfolder per originating module. Useful to collect training data.

    Call it at the beginning of your scripts, before creating your modules.

    Args:
        base_dir (str): Optional. The root folder where the records are
            written. If not provided, uses `synalinks_home()`
            (`$SYNALINKS_HOME` or `~/.synalinks`).

    Example:

    ```python
    import synalinks

    # Basic usage: records under ~/.synalinks
    synalinks.record_traces()

    # With a custom folder
    synalinks.record_traces(base_dir="./traces")
    ```
    """
    global _ENABLE_TRACE_RECORDING
    global _TRACE_RECORDING_DIR

    if base_dir:
        _TRACE_RECORDING_DIR = base_dir
    _ENABLE_TRACE_RECORDING = True

@synalinks_export(
    [
        "synalinks.config.set_api_base",
        "synalinks.backend.set_api_base",
        "synalinks.set_api_base",
    ]
)
def set_api_base(api_base):
    """Set the observability api base

    Args:
        api_base (str): The observability api base
    """
    global _SYNALINKS_API_BASE
    _SYNALINKS_API_BASE = api_base

@synalinks_export(
    [
        "synalinks.config.set_api_key",
        "synalinks.backend.set_api_key",
    ]
)
def set_api_key(key):
    """Set Synalinks API key.

    Args:
        key (str): The API key value.

    The API key is retrieved from the env variables at start.

    ```python
    >>> os.environ["SYNALINKS_API_KEY"] = 'my-secret-api-key'
    ```

    Or you can setup it using the config

    ```python
    >>> synalinks.config.set_api_key('my-secret-api-key')
    >>> synalinks.config.api_key()
    'my-secret-api-key'
    ```

    Args:
        key (str): Synalinks API key.
    """
    global _SYNALINKS_API_KEY
    _SYNALINKS_API_KEY = key

@synalinks_export(
    [
        "synalinks.config.set_default_embedding_model",
        "synalinks.set_default_embedding_model",
    ]
)
def set_default_embedding_model(identifier: "str | dict | object | None"):
    """Set the default `EmbeddingModel`.

    Args:
        identifier (str | dict | EmbeddingModel | None): A model string
            (e.g. `"openai/text-embedding-3-small"`), a config dict, an
            existing `EmbeddingModel` instance, or `None` to clear.
            Strings persist into the on-disk config; instances do not.
    """
    global _DEFAULT_EMBEDDING_MODEL, _DEFAULT_EMBEDDING_MODEL_IDENTIFIER
    if identifier is None:
        _DEFAULT_EMBEDDING_MODEL = None
        _DEFAULT_EMBEDDING_MODEL_IDENTIFIER = None
        _persist_config()
        return
    from synalinks.src.modules.embedding_models import get as _get_em

    _DEFAULT_EMBEDDING_MODEL = _get_em(identifier)
    _DEFAULT_EMBEDDING_MODEL_IDENTIFIER = (
        identifier if isinstance(identifier, str) else None
    )
    _persist_config()

@synalinks_export(
    [
        "synalinks.config.set_default_knowledge_base",
        "synalinks.set_default_knowledge_base",
    ]
)
def set_default_knowledge_base(identifier: "str | dict | object | None"):
    """Set the default `KnowledgeBase`.

    Args:
        identifier (str | dict | KnowledgeBase | None): A URI string
            (e.g. `"duckdb://./my_database.db"`), a config dict, an
            existing `KnowledgeBase` instance, or `None` to clear.
            Strings persist into the on-disk config; instances do not.
    """
    global _DEFAULT_KNOWLEDGE_BASE, _DEFAULT_KNOWLEDGE_BASE_IDENTIFIER
    if identifier is None:
        _DEFAULT_KNOWLEDGE_BASE = None
        _DEFAULT_KNOWLEDGE_BASE_IDENTIFIER = None
        _persist_config()
        return
    from synalinks.src.knowledge_bases import get as _get_kb

    _DEFAULT_KNOWLEDGE_BASE = _get_kb(identifier)
    _DEFAULT_KNOWLEDGE_BASE_IDENTIFIER = (
        identifier if isinstance(identifier, str) else None
    )
    _persist_config()

@synalinks_export(
    [
        "synalinks.config.set_default_language_model",
        "synalinks.set_default_language_model",
    ]
)
def set_default_language_model(identifier: "str | dict | object | None"):
    """Set the default `LanguageModel`.

    Args:
        identifier (str | dict | LanguageModel | None): A model string
            (e.g. `"openai/gpt-4o-mini"`), a config dict, an existing
            `LanguageModel` instance, or `None` to clear. Strings persist
            into the on-disk config; instances do not.
    """
    global _DEFAULT_LANGUAGE_MODEL, _DEFAULT_LANGUAGE_MODEL_IDENTIFIER
    if identifier is None:
        _DEFAULT_LANGUAGE_MODEL = None
        _DEFAULT_LANGUAGE_MODEL_IDENTIFIER = None
        _persist_config()
        return
    from synalinks.src.modules.language_models import get as _get_lm

    _DEFAULT_LANGUAGE_MODEL = _get_lm(identifier)
    _DEFAULT_LANGUAGE_MODEL_IDENTIFIER = (
        identifier if isinstance(identifier, str) else None
    )
    _persist_config()

@synalinks_export(["synalinks.config.set_epsilon", "synalinks.backend.set_epsilon"])
def set_epsilon(value):
    """Set the value of the fuzz factor used in numeric expressions.

    Args:
        value (float): The new value of epsilon.

    Examples:

    ```python
    >>> synalinks.config.epsilon()
    1e-07
    ```

    ```python
    >>> synalinks.config.set_epsilon(1e-5)
    >>> synalinks.config.epsilon()
    1e-05
    ```

    ```python
    >>> # Set it back to the default value.
    >>> synalinks.config.set_epsilon(1e-7)
    ```

    """
    global _EPSILON
    _EPSILON = value

@synalinks_export(["synalinks.config.set_floatx", "synalinks.backend.set_floatx"])
def set_floatx(value):
    """Set the default float dtype.

    Note: It is not recommended to set this to `"float16"`,
    as this will likely cause numeric stability issues.
    Instead, use `float64` or `float32`.

    Args:
        value (str): The float type between `'bfloat16'`, `'float16'`, `'float32'`,
            or `'float64'`.

    Examples:

    ```python
    >>> synalinks.config.floatx()
    'float32'
    ```

    ```python
    >>> synalinks.config.set_floatx('float64')
    >>> synalinks.config.floatx()
    'float64'
    ```

    ```python
    >>> # Set it back to float32
    >>> synalinks.config.set_floatx('float32')
    ```

    Raises:
        ValueError: In case of invalid value.
    """
    global _FLOATX
    accepted_dtypes = {"bfloat16", "float16", "float32", "float64"}
    if value not in accepted_dtypes:
        raise ValueError(
            f"Unknown `floatx` value: {value}. Expected one of {accepted_dtypes}"
        )
    _FLOATX = str(value)

@synalinks_export(
    [
        "synalinks.config.set_seed",
        "synalinks.backend.set_seed",
        "synalinks.set_seed",
    ]
)
def set_seed(value):
    """Set the value of the random seed for reproductability.

    Args:
        value (float): The new value of epsilon.
    """
    global _RANDOM_SEED
    _RANDOM_SEED = value
    np.random.seed(_RANDOM_SEED)
    random.seed(_RANDOM_SEED)

@synalinks_export(
    [
        "synalinks.config.trace_recording_dir",
        "synalinks.backend.trace_recording_dir",
    ]
)
def trace_recording_dir():
    """Returns the root folder where the trace records are written.

    Returns:
        (str): The folder set via `record_traces()`, or
            `synalinks_home()` if none was set.
    """
    if _TRACE_RECORDING_DIR:
        return _TRACE_RECORDING_DIR
    return synalinks_home()