qtype 0.0.9__py3-none-any.whl → 0.0.11__py3-none-any.whl

qtype/interpreter/auth/aws.py

@@ -0,0 +1,234 @@
+"""
+AWS authentication context manager for QType interpreter.
+
+This module provides a context manager for creating boto3 sessions using
+AWSAuthProvider configuration from the semantic model.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Any, Generator
+
+import boto3  # type: ignore[import-untyped]
+from botocore.exceptions import (  # type: ignore[import-untyped]
+    ClientError,
+    NoCredentialsError,
+)
+
+from qtype.interpreter.auth.cache import cache_auth, get_cached_auth
+from qtype.semantic.model import AWSAuthProvider
+
+
+class AWSAuthenticationError(Exception):
+    """Raised when AWS authentication fails."""
+
+    pass
+
+
+def _is_session_valid(session: boto3.Session) -> bool:
+    """
+    Check if a boto3 session is still valid by testing credential access.
+
+    Args:
+        session: The boto3 session to validate
+
+    Returns:
+        bool: True if the session is valid, False otherwise
+    """
+    try:
+        credentials = session.get_credentials()
+        if credentials is None:
+            return False
+
+        # For temporary credentials, check if they're still valid
+        if hasattr(credentials, "token") and credentials.token:
+            # Create a test STS client to verify the credentials
+            sts_client = session.client("sts")
+            sts_client.get_caller_identity()
+
+        return True
+    except (ClientError, NoCredentialsError):
+        return False
+    except Exception:
+        # Any other exception means the session is likely invalid
+        return False
+
+
+@contextmanager
+def aws(aws_provider: AWSAuthProvider) -> Generator[boto3.Session, None, None]:
+    """
+    Create a boto3 Session using AWS authentication provider configuration.
+
+    This context manager creates a boto3 Session based on the authentication
+    method specified in the AWSAuthProvider. Sessions are cached using an LRU
+    cache to avoid recreating them unnecessarily. The cache size can be configured
+    via the AUTH_CACHE_MAX_SIZE environment variable (default: 128).
+
+    It supports:
+    - Direct credentials (access key + secret key + optional session token)
+    - AWS profiles from shared credentials/config files
+    - Role assumption (with optional external ID and MFA)
+    - Environment-based authentication (when no explicit credentials provided)
+
+    Caching behavior:
+    - Sessions are cached based on the AWSAuthProvider configuration
+    - Cached sessions are validated before reuse to check for expiration
+    - Invalid or expired sessions are evicted and recreated
+
+    Args:
+        aws_provider: AWSAuthProvider instance containing authentication configuration
+
+    Yields:
+        boto3.Session: Configured boto3 session ready for creating AWS service clients
+
+    Raises:
+        AWSAuthenticationError: When authentication fails or configuration is invalid
+
+    Example:
+        ```python
+        from qtype.semantic.model import AWSAuthProvider
+        from qtype.interpreter.auth.aws import aws
+
+        aws_auth = AWSAuthProvider(
+            id="my-aws-auth",
+            type="aws",
+            access_key_id="AKIA...",
+            secret_access_key="...",
+            region="us-east-1"
+        )
+
+        with aws(aws_auth) as session:
+            athena_client = session.client("athena")
+            s3_client = session.client("s3")
+        ```
+    """
+    try:
+        # Check cache first - use provider object directly as cache key
+        cached_session = get_cached_auth(aws_provider)
+
+        if cached_session is not None and _is_session_valid(cached_session):
+            # Cache hit with valid session
+            yield cached_session
+            return
+
+        # Cache miss or invalid session - create new session
+        session = _create_session(aws_provider)
+
+        # Validate the session by attempting to get credentials
+        credentials = session.get_credentials()
+        if credentials is None:
+            raise AWSAuthenticationError(
+                f"Failed to obtain AWS credentials for provider '{aws_provider.id}'"
+            )
+
+        # Cache the valid session using provider object as key
+        cache_auth(aws_provider, session)
+
+        yield session
+
+    except (ClientError, NoCredentialsError) as e:
+        raise AWSAuthenticationError(
+            f"AWS authentication failed for provider '{aws_provider.id}': {e}"
+        ) from e
+    except Exception as e:
+        raise AWSAuthenticationError(
+            f"Unexpected error during AWS authentication for provider '{aws_provider.id}': {e}"
+        ) from e
+
+
+def _create_session(aws_provider: AWSAuthProvider) -> boto3.Session:
+    """
+    Create a boto3 Session based on the AWS provider configuration.
+
+    Args:
+        aws_provider: AWSAuthProvider with authentication details
+
+    Returns:
+        boto3.Session: Configured session
+
+    Raises:
+        AWSAuthenticationError: If configuration is invalid
+    """
+    session_kwargs: dict[str, Any] = {}
+
+    # Add region if specified
+    if aws_provider.region:
+        session_kwargs["region_name"] = aws_provider.region
+
+    # Handle different authentication methods
+    if aws_provider.profile_name:
+        # Use AWS profile from shared credentials/config files
+        session_kwargs["profile_name"] = aws_provider.profile_name
+
+    elif aws_provider.access_key_id and aws_provider.secret_access_key:
+        # Use direct credentials
+        session_kwargs["aws_access_key_id"] = aws_provider.access_key_id
+        session_kwargs["aws_secret_access_key"] = (
+            aws_provider.secret_access_key
+        )
+
+        if aws_provider.session_token:
+            session_kwargs["aws_session_token"] = aws_provider.session_token
+
+    # Create the base session
+    session = boto3.Session(**session_kwargs)
+
+    # Handle role assumption if specified
+    if aws_provider.role_arn:
+        session = _assume_role_session(session, aws_provider)
+
+    return session
+
+
+def _assume_role_session(
+    base_session: boto3.Session, aws_provider: AWSAuthProvider
+) -> boto3.Session:
+    """
+    Create a new session by assuming an IAM role.
+
+    Args:
+        base_session: The base session to use for assuming the role
+        aws_provider: AWSAuthProvider with role configuration
+
+    Returns:
+        boto3.Session: New session with assumed role credentials
+
+    Raises:
+        AWSAuthenticationError: If role assumption fails
+    """
+    if not aws_provider.role_arn:
+        raise AWSAuthenticationError(
+            "role_arn is required for role assumption"
+        )
+
+    try:
+        sts_client = base_session.client("sts")
+
+        # Prepare AssumeRole parameters
+        assume_role_params: dict[str, Any] = {
+            "RoleArn": aws_provider.role_arn,
+            "RoleSessionName": aws_provider.role_session_name
+            or f"qtype-session-{aws_provider.id}",
+        }
+
+        if aws_provider.external_id:
+            assume_role_params["ExternalId"] = aws_provider.external_id
+
+        # Assume the role
+        response = sts_client.assume_role(**assume_role_params)
+        credentials = response["Credentials"]
+
+        # Create new session with temporary credentials
+        return boto3.Session(
+            aws_access_key_id=credentials["AccessKeyId"],
+            aws_secret_access_key=credentials["SecretAccessKey"],
+            aws_session_token=credentials["SessionToken"],
+            region_name=aws_provider.region or base_session.region_name,
+        )
+
+    except ClientError as e:
+        error_code = e.response.get("Error", {}).get("Code", "Unknown")
+        raise AWSAuthenticationError(
+            f"Failed to assume role '{aws_provider.role_arn}': {error_code} - {e}"
+        ) from e

qtype/interpreter/auth/cache.py

@@ -0,0 +1,67 @@
+"""
+Authorization cache for QType interpreter.
+
+This module provides a shared LRU cache for authorization sessions and tokens
+across different authentication providers (AWS, OAuth2, API keys, etc.).
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+from cachetools import LRUCache
+
+# Global LRU cache for authorization sessions with configurable size
+_AUTH_CACHE_MAX_SIZE = int(os.environ.get("AUTH_CACHE_MAX_SIZE", 128))
+_AUTHORIZATION_CACHE: LRUCache[Any, Any] = LRUCache(
+    maxsize=_AUTH_CACHE_MAX_SIZE
+)
+
+
+def get_cached_auth(auth_provider: Any) -> Any | None:
+    """
+    Get a cached authorization session for the given provider.
+
+    Args:
+        auth_provider: Authorization provider instance (must be hashable)
+
+    Returns:
+        Cached session/token or None if not found
+    """
+    return _AUTHORIZATION_CACHE.get(auth_provider)
+
+
+def cache_auth(auth_provider: Any, session: Any) -> None:
+    """
+    Cache an authorization session for the given provider.
+
+    Args:
+        auth_provider: Authorization provider instance (must be hashable)
+        session: Session or token to cache
+    """
+    _AUTHORIZATION_CACHE[auth_provider] = session
+
+
+def clear_auth_cache() -> None:
+    """
+    Clear all cached authorization sessions.
+
+    This can be useful for testing or when credential configurations change.
+    """
+    _AUTHORIZATION_CACHE.clear()
+
+
+def get_cache_info() -> dict[str, Any]:
+    """
+    Get information about the current state of the authorization cache.
+
+    Returns:
+        Dictionary with cache statistics and configuration
+    """
+    return {
+        "max_size": _AUTH_CACHE_MAX_SIZE,
+        "current_size": len(_AUTHORIZATION_CACHE),
+        "hits": getattr(_AUTHORIZATION_CACHE, "hits", 0),
+        "misses": getattr(_AUTHORIZATION_CACHE, "misses", 0),
+    }

qtype/interpreter/auth/generic.py

@@ -0,0 +1,103 @@
+"""
+Generic authorization context manager for QType interpreter.
+
+This module provides a unified context manager that can handle any AuthorizationProvider
+type and return the appropriate session or provider instance.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Generator
+
+import boto3  # type: ignore[import-untyped]
+
+from qtype.interpreter.auth.aws import aws
+from qtype.semantic.model import (
+    APIKeyAuthProvider,
+    AuthorizationProvider,
+    AWSAuthProvider,
+    OAuth2AuthProvider,
+)
+
+
+class UnsupportedAuthProviderError(Exception):
+    """Raised when an unsupported authorization provider type is used."""
+
+    pass
+
+
+@contextmanager
+def auth(
+    auth_provider: AuthorizationProvider,
+) -> Generator[boto3.Session | APIKeyAuthProvider, None, None]:
+    """
+    Create an appropriate session or provider instance based on the auth provider type.
+
+    This context manager dispatches to the appropriate authentication handler based
+    on the type of AuthorizationProvider:
+    - AWSAuthProvider: Returns a configured boto3.Session
+    - APIKeyAuthProvider: Returns the provider instance (contains the API key)
+    - OAuth2AuthProvider: Raises NotImplementedError (not yet supported)
+
+    Args:
+        auth_provider: AuthorizationProvider instance of any supported type
+
+    Yields:
+        boto3.Session | APIKeyAuthProvider: The appropriate session or provider instance
+
+    Raises:
+        UnsupportedAuthProviderError: When an unsupported provider type is used
+        NotImplementedError: When OAuth2AuthProvider is used (not yet implemented)
+
+    Example:
+        ```python
+        from qtype.semantic.model import AWSAuthProvider, APIKeyAuthProvider
+        from qtype.interpreter.auth.generic import auth
+
+        # AWS provider - returns boto3.Session
+        aws_auth = AWSAuthProvider(
+            id="my-aws-auth",
+            type="aws",
+            access_key_id="AKIA...",
+            secret_access_key="...",
+            region="us-east-1"
+        )
+
+        with auth(aws_auth) as session:
+            s3_client = session.client("s3")
+
+        # API Key provider - returns the provider itself
+        api_auth = APIKeyAuthProvider(
+            id="my-api-auth",
+            type="api_key",
+            api_key="sk-...",
+            host="api.openai.com"
+        )
+
+        with auth(api_auth) as provider:
+            headers = {"Authorization": f"Bearer {provider.api_key}"}
+        ```
+    """
+    if isinstance(auth_provider, AWSAuthProvider):
+        # Use AWS-specific context manager
+        with aws(auth_provider) as session:
+            yield session
+
+    elif isinstance(auth_provider, APIKeyAuthProvider):
+        # For API key providers, just return the provider itself
+        # The caller can access provider.api_key and provider.host
+        yield auth_provider
+
+    elif isinstance(auth_provider, OAuth2AuthProvider):
+        # OAuth2 not yet implemented
+        raise NotImplementedError(
+            f"OAuth2 authentication is not yet implemented for provider '{auth_provider.id}'"
+        )
+
+    else:
+        # Unknown provider type
+        raise UnsupportedAuthProviderError(
+            f"Unsupported authorization provider type: {type(auth_provider).__name__} "
+            f"for provider '{auth_provider.id}'"
+        )

qtype/interpreter/batch/flow.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Tuple
+
+import pandas as pd
+
+from qtype.interpreter.batch.step import batch_execute_step
+from qtype.interpreter.batch.types import BatchConfig
+from qtype.interpreter.batch.utils import reconcile_results_and_errors
+from qtype.semantic.model import Flow, Sink
+
+logger = logging.getLogger(__name__)
+
+
+def batch_execute_flow(
+    flow: Flow,
+    inputs: pd.DataFrame,
+    batch_config: BatchConfig,
+    **kwargs: dict[Any, Any],
+) -> Tuple[pd.DataFrame, pd.DataFrame]:
+    """Executes a flow in a batch context.
+
+    Args:
+        flow: The flow to execute.
+        batch_config: The batch configuration to use.
+        **kwargs: Additional keyword arguments to pass to the flow.
+
+    Returns:
+        A list of output variables produced by the flow.
+    """
+
+    previous_outputs = inputs
+
+    all_errors = []
+
+    # Iterate over each step in the flow
+    for step in flow.steps:
+        results: list[pd.DataFrame] = []
+        errors: list[pd.DataFrame] = []
+
+        if isinstance(step, Sink):
+            # Send the entire batch to the sink
+            batch_results, batch_errors = batch_execute_step(
+                step, previous_outputs, batch_config
+            )
+            results.append(batch_results)
+            if len(batch_errors) > 1:
+                errors.append(batch_errors)
+        else:
+            # batch the current data into dataframes of max size batch_size
+            batch_size = batch_config.batch_size
+            for start in range(0, len(previous_outputs), batch_size):
+                end = start + batch_size
+                batch = previous_outputs.iloc[start:end]
+                # Execute the step with the current batch
+                batch_results, batch_errors = batch_execute_step(
+                    step, batch, batch_config
+                )
+
+                results.append(batch_results)
+                if len(batch_errors) > 1:
+                    errors.append(batch_errors)
+
+        previous_outputs, errors_df = reconcile_results_and_errors(
+            results, errors
+        )
+
+        if len(errors_df):
+            all_errors.append(errors_df)
+            if batch_config.write_errors_to:
+                output_file = (
+                    f"{batch_config.write_errors_to}/{step.id}.errors.parquet"
+                )
+                try:
+                    errors_df.to_parquet(
+                        output_file, engine="pyarrow", compression="snappy"
+                    )
+                    logging.info(
+                        f"Saved errors for step {step.id} to {output_file}"
+                    )
+                except Exception as e:
+                    logging.warning(
+                        f"Could not save errors step {step.id} to {output_file}",
+                        exc_info=e,
+                        stack_info=True,
+                    )
+
+    # Return the last steps results and errors
+    rv_errors = (
+        pd.concat(all_errors, ignore_index=True)
+        if len(all_errors)
+        else pd.DataFrame({})
+    )
+    return previous_outputs, rv_errors

qtype/interpreter/batch/sql_source.py

@@ -0,0 +1,95 @@
+from typing import Any, Tuple
+
+import boto3  # type: ignore[import-untyped]
+import pandas as pd
+import sqlalchemy
+from sqlalchemy import create_engine
+from sqlalchemy.exc import SQLAlchemyError
+
+from qtype.base.exceptions import InterpreterError
+from qtype.interpreter.auth.generic import auth
+from qtype.interpreter.batch.types import BatchConfig, ErrorMode
+from qtype.interpreter.batch.utils import (
+    reconcile_results_and_errors,
+    validate_inputs,
+)
+from qtype.semantic.model import SQLSource
+
+
+def to_output_columns(
+    df: pd.DataFrame, output_columns: set[str]
+) -> pd.DataFrame:
+    """Filters the DataFrame to only include specified output columns.
+
+    Args:
+        df: The input DataFrame.
+        output_columns: A set of column names to retain in the DataFrame.
+
+    Returns:
+        A DataFrame containing only the specified output columns.
+    """
+    if len(df) == 0:
+        return df
+    missing = output_columns - set(df.columns)
+    if missing:
+        raise InterpreterError(
+            f"SQL Result was missing expected columns: {','.join(missing)}, it has columns: {','.join(df.columns)}"
+        )
+
+    return df[[col for col in df.columns if col in output_columns]]
+
+
+def execute_sql_source(
+    step: SQLSource,
+    inputs: pd.DataFrame,
+    batch_config: BatchConfig,
+    **kwargs: dict[Any, Any],
+) -> Tuple[pd.DataFrame, pd.DataFrame]:
+    """Executes a SQLSource step to retrieve data from a SQL database.
+
+    Args:
+        step: The SQLSource step to execute.
+
+    Returns:
+        A tuple containing two DataFrames:
+            - The first DataFrame contains the successfully retrieved data.
+            - The second DataFrame contains rows that encountered errors with an 'error' column.
+    """
+    # Create a database engine
+    validate_inputs(inputs, step)
+
+    connect_args = {}
+    if step.auth:
+        with auth(step.auth) as creds:
+            if isinstance(creds, boto3.Session):
+                connect_args["session"] = creds
+    engine = create_engine(step.connection, connect_args=connect_args)
+
+    output_columns = {output.id for output in step.outputs}
+
+    results = []
+    errors = []
+    step_inputs = {i.id for i in step.inputs}
+    for _, row in inputs.iterrows():
+        try:
+            # Make a dictionary of column_name: value from row
+            params = {col: row[col] for col in row.index if col in step_inputs}
+            # Execute the query and fetch the results into a DataFrame
+            with engine.connect() as connection:
+                result = connection.execute(
+                    sqlalchemy.text(step.query),
+                    parameters=params if len(params) else None,
+                )
+                df = pd.DataFrame(
+                    result.fetchall(), columns=list(result.keys())
+                )
+            df = to_output_columns(df, output_columns)
+            results.append(df)
+        except SQLAlchemyError as e:
+            if batch_config.error_mode == ErrorMode.FAIL:
+                raise e
+            # If there's an error, return an empty DataFrame and the error message
+            error_df = pd.DataFrame([{"error": str(e)}])
+            errors.append(error_df)
+
+    return reconcile_results_and_errors(results, errors)

qtype/interpreter/batch/step.py

@@ -0,0 +1,63 @@
+from functools import partial
+from typing import Any, Tuple
+
+import pandas as pd
+
+from qtype.interpreter.batch.sql_source import execute_sql_source
+from qtype.interpreter.batch.types import BatchConfig
+from qtype.interpreter.batch.utils import (
+    batch_iterator,
+    single_step_adapter,
+    validate_inputs,
+)
+from qtype.interpreter.exceptions import InterpreterError
+from qtype.semantic.model import (
+    Condition,
+    Decoder,
+    Flow,
+    PromptTemplate,
+    Search,
+    SQLSource,
+    Step,
+    Tool,
+)
+
+SINGLE_WRAP_STEPS = {Decoder, Condition, PromptTemplate, Search, Tool}
+
+
+def batch_execute_step(
+    step: Step,
+    inputs: pd.DataFrame,
+    batch_config: BatchConfig,
+    **kwargs: dict[str, Any],
+) -> Tuple[pd.DataFrame, pd.DataFrame]:
+    """
+    Executes a given step in a batch processing pipeline.
+
+    Args:
+        step (Step): The step to be executed.
+        inputs (pd.DataFrame): The input data for the step.
+        batch_config (BatchConfig): Configuration for batch processing.
+        **kwargs: Additional keyword arguments.
+
+    Returns:
+        Tuple[pd.DataFrame, pd.DataFrame]: A tuple containing the output results and any rows that returned errors.
+    """
+
+    validate_inputs(inputs, step)
+
+    if isinstance(step, Flow):
+        from qtype.interpreter.batch.flow import batch_execute_flow
+
+        return batch_execute_flow(step, inputs, batch_config, **kwargs)
+    elif isinstance(step, SQLSource):
+        return execute_sql_source(step, inputs, batch_config, **kwargs)
+    elif step in SINGLE_WRAP_STEPS:
+        return batch_iterator(
+            f=partial(single_step_adapter, step=step),
+            batch=inputs,
+            batch_config=batch_config,
+        )
+    # TODO: implement batching for multi-row steps. For example, llm inference can be sped up in batch...
+    else:
+        raise InterpreterError(f"Unsupported step type: {type(step).__name__}")

qtype/interpreter/batch/types.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from enum import Enum
+
+from pydantic import BaseModel, Field
+
+
+class ErrorMode(str, Enum):
+    """Error handling mode for batch processing."""
+
+    FAIL = "fail"
+    DROP = "drop"
+
+
+class BatchConfig(BaseModel):
+    """Configuration for batch execution.
+
+    Attributes:
+        num_workers: Number of async workers for batch operations.
+        batch_size: Maximum number of rows to send to a step at a time.
+        error_mode: Error handling mode for batch processing.
+    """
+
+    num_workers: int = Field(
+        default=4,
+        description="Number of async workers for batch operations",
+        gt=0,
+    )
+    batch_size: int = Field(
+        default=512,
+        description="Max number of rows to send to a step at a time",
+        gt=0,
+    )
+    error_mode: ErrorMode = Field(
+        default=ErrorMode.FAIL,
+        description="Error handling mode for batch processing",
+    )
+    write_errors_to: str | None = Field(
+        default=None,
+        description="If error mode is DROP, the errors for any step are saved to this directory",
+    )