qtype 0.0.10__py3-none-any.whl → 0.0.12__py3-none-any.whl

qtype/dsl/model.py

@@ -14,7 +14,11 @@ from pydantic import (
 )
 
 import qtype.dsl.domain_types as domain_types
-from qtype.dsl.base_types import PrimitiveTypeEnum, StrictBaseModel
+from qtype.dsl.base_types import (
+    PrimitiveTypeEnum,
+    StepCardinality,
+    StrictBaseModel,
+)
 from qtype.dsl.domain_types import ChatContent, ChatMessage, Embedding
 
 
@@ -116,7 +120,7 @@ class Model(StrictBaseModel):
     """Describes a generative model configuration, including provider and model ID."""
 
     id: str = Field(..., description="Unique ID for the model.")
-    auth: AuthorizationProvider | str | None = Field(
+    auth: AuthProviderType | str | None = Field(
         default=None,
         description="AuthorizationProvider used for model access.",
     )
@@ -172,6 +176,10 @@ class Step(StrictBaseModel, ABC):
     """Base class for components that take inputs and produce outputs."""
 
     id: str = Field(..., description="Unique ID of this component.")
+    cardinality: StepCardinality = Field(
+        default=StepCardinality.one,
+        description="Does this step emit 1 (one) or 0...N (many) instances of the outputs?",
+    )
     inputs: list[Variable | str] | None = Field(
         default=None,
         description="Input variables required by this step.",
@@ -261,7 +269,7 @@ class APITool(Tool):
         default="GET",
         description="HTTP method to use (GET, POST, PUT, DELETE, etc.).",
     )
-    auth: AuthorizationProvider | str | None = Field(
+    auth: AuthProviderType | str | None = Field(
         default=None,
         description="Optional AuthorizationProvider for API authentication.",
     )
@@ -315,6 +323,11 @@ class Flow(Step):
         default=None, description="Optional description of the flow."
     )
 
+    cardinality: StepCardinality = Field(
+        default=StepCardinality.auto,
+        description="The cardinality of the flow, inferred from its steps when set to 'auto'.",
+    )
+
     mode: Literal["Complete", "Chat"] = "Complete"
 
     steps: list[StepType | str] = Field(
@@ -368,34 +381,95 @@ class Decoder(Step):
 #
 
 
-class AuthorizationProvider(StrictBaseModel):
-    """Defines how tools or providers authenticate with APIs, such as OAuth2 or API keys."""
+class AuthorizationProvider(StrictBaseModel, ABC):
+    """Base class for authentication providers."""
 
     id: str = Field(
         ..., description="Unique ID of the authorization configuration."
     )
-    api_key: str | None = Field(
-        default=None, description="API key if using token-based auth."
-    )
-    client_id: str | None = Field(
-        default=None, description="OAuth2 client ID."
-    )
-    client_secret: str | None = Field(
-        default=None, description="OAuth2 client secret."
-    )
+    type: str = Field(..., description="Authorization method type.")
+
+
+class APIKeyAuthProvider(AuthorizationProvider):
+    """API key-based authentication provider."""
+
+    type: Literal["api_key"] = "api_key"
+    api_key: str = Field(..., description="API key for authentication.")
     host: str | None = Field(
         default=None, description="Base URL or domain of the provider."
     )
+
+
+class OAuth2AuthProvider(AuthorizationProvider):
+    """OAuth2 authentication provider."""
+
+    type: Literal["oauth2"] = "oauth2"
+    client_id: str = Field(..., description="OAuth2 client ID.")
+    client_secret: str = Field(..., description="OAuth2 client secret.")
+    token_url: str = Field(..., description="Token endpoint URL.")
     scopes: list[str] | None = Field(
         default=None, description="OAuth2 scopes required."
     )
-    token_url: str | None = Field(
-        default=None, description="Token endpoint URL."
+
+
+class AWSAuthProvider(AuthorizationProvider):
+    """AWS authentication provider supporting multiple credential methods."""
+
+    type: Literal["aws"] = "aws"
+
+    # Method 1: Access key/secret/session
+    access_key_id: str | None = Field(
+        default=None, description="AWS access key ID."
     )
-    type: str = Field(
-        ..., description="Authorization method, e.g., 'oauth2' or 'api_key'."
+    secret_access_key: str | None = Field(
+        default=None, description="AWS secret access key."
+    )
+    session_token: str | None = Field(
+        default=None,
+        description="AWS session token for temporary credentials.",
     )
 
+    # Method 2: Profile
+    profile_name: str | None = Field(
+        default=None, description="AWS profile name from credentials file."
+    )
+
+    # Method 3: Role assumption
+    role_arn: str | None = Field(
+        default=None, description="ARN of the role to assume."
+    )
+    role_session_name: str | None = Field(
+        default=None, description="Session name for role assumption."
+    )
+    external_id: str | None = Field(
+        default=None, description="External ID for role assumption."
+    )
+
+    # Common AWS settings
+    region: str | None = Field(default=None, description="AWS region.")
+
+    @model_validator(mode="after")
+    def validate_aws_auth(self) -> "AWSAuthProvider":
+        """Validate AWS authentication configuration."""
+        # At least one auth method must be specified
+        has_keys = self.access_key_id and self.secret_access_key
+        has_profile = self.profile_name
+        has_role = self.role_arn
+
+        if not (has_keys or has_profile or has_role):
+            raise ValueError(
+                "AWSAuthProvider must specify at least one authentication method: "
+                "access keys, profile name, or role ARN."
+            )
+
+        # If assuming a role, need either keys or profile for base credentials
+        if has_role and not (has_keys or has_profile):
+            raise ValueError(
+                "Role assumption requires base credentials (access keys or profile)."
+            )
+
+        return self
+
 
 class TelemetrySink(StrictBaseModel):
     """Defines an observability endpoint for collecting telemetry data from the QType runtime."""
@@ -403,7 +477,7 @@ class TelemetrySink(StrictBaseModel):
     id: str = Field(
         ..., description="Unique ID of the telemetry sink configuration."
     )
-    auth: AuthorizationProvider | str | None = Field(
+    auth: AuthProviderType | str | None = Field(
         default=None,
         description="AuthorizationProvider used to authenticate telemetry data transmission.",
     )
@@ -455,7 +529,7 @@ class Application(StrictBaseModel):
     )
 
     # External integrations
-    auths: list[AuthorizationProvider] | None = Field(
+    auths: list[AuthProviderType] | None = Field(
         default=None,
         description="List of authorization providers used for API access.",
     )
@@ -480,6 +554,58 @@ class Application(StrictBaseModel):
     )
 
 
+#
+# ---------------- Data Pipeline Components ----------------
+#
+
+
+class Source(Step):
+    """Base class for data sources"""
+
+    id: str = Field(..., description="Unique ID of the data source.")
+    cardinality: Literal[StepCardinality.many] = Field(
+        default=StepCardinality.many,
+        description="Sources always emit 0...N instances of the outputs.",
+    )
+
+
+class SQLSource(Source):
+    """SQL database source that executes queries and emits rows."""
+
+    query: str = Field(
+        ..., description="SQL query to execute. Inputs are injected as params."
+    )
+    connection: str = Field(
+        ...,
+        description="Database connection string or reference to auth provider. Typically in SQLAlchemy format.",
+    )
+    auth: AuthProviderType | str | None = Field(
+        default=None,
+        description="Optional AuthorizationProvider for database authentication.",
+    )
+
+    @model_validator(mode="after")
+    def validate_sql_source(self) -> "SQLSource":
+        """Validate SQL source configuration."""
+        if self.outputs is None:
+            raise ValueError(
+                "SQLSource must define output variables that match the result columns."
+            )
+        return self
+
+
+class Sink(Step):
+    """Base class for data sinks"""
+
+    id: str = Field(..., description="Unique ID of the data sink.")
+    # Remove cardinality field - it's always one for sinks
+    # ...existing code...
+    cardinality: Literal[StepCardinality.one] = Field(
+        default=StepCardinality.one,
+        description="Flows always emit exactly one instance of the outputs.",
+    )
+
+
 #
 # ---------------- Retrieval Augmented Generation Components ----------------
 #
@@ -493,13 +619,19 @@ class Index(StrictBaseModel, ABC):
         default=None,
         description="Index-specific configuration and connection parameters.",
     )
-    auth: AuthorizationProvider | str | None = Field(
+    auth: AuthProviderType | str | None = Field(
         default=None,
         description="AuthorizationProvider for accessing the index.",
     )
     name: str = Field(..., description="Name of the index/collection/table.")
 
 
+class IndexUpsert(Sink):
+    index: IndexType | str = Field(
+        ..., description="Index to upsert into (object or ID reference)."
+    )
+
+
 class VectorIndex(Index):
     """Vector database index for similarity search using embeddings."""
 
@@ -531,6 +663,7 @@ class VectorSearch(Search):
     """Performs vector similarity search against a vector index."""
 
     default_top_k: int | None = Field(
+        default=50,
         description="Number of top results to retrieve if not provided in the inputs.",
     )
 
@@ -539,7 +672,7 @@ class VectorSearch(Search):
         """Set default input and output variables if none provided."""
         if self.inputs is None:
             self.inputs = [
-                Variable(id="top_k", type=PrimitiveTypeEnum.number),
+                Variable(id="top_k", type=PrimitiveTypeEnum.int),
                 Variable(id="query", type=PrimitiveTypeEnum.text),
             ]
 
@@ -570,6 +703,16 @@ ToolType = Union[
     PythonFunctionTool,
 ]
 
+# Create a union type for all source types
+SourceType = Union[SQLSource,]
+
+# Create a union type for all authorization provider types
+AuthProviderType = Union[
+    APIKeyAuthProvider,
+    AWSAuthProvider,
+    OAuth2AuthProvider,
+]
+
 # Create a union type for all step types
 StepType = Union[
     Agent,
@@ -578,9 +721,12 @@ StepType = Union[
     Decoder,
     DocumentSearch,
     Flow,
+    IndexUpsert,
     LLMInference,
     PromptTemplate,
     PythonFunctionTool,
+    SQLSource,
+    Sink,
     VectorSearch,
 ]
 
@@ -602,10 +748,10 @@ ModelType = Union[
 #
 
 
-class AuthorizationProviderList(RootModel[list[AuthorizationProvider]]):
+class AuthorizationProviderList(RootModel[list[AuthProviderType]]):
     """Schema for a standalone list of authorization providers."""
 
-    root: list[AuthorizationProvider]
+    root: list[AuthProviderType]
 
 
 class IndexList(RootModel[list[IndexType]]):
@@ -638,35 +784,24 @@ class VariableList(RootModel[list[Variable]]):
     root: list[Variable]
 
 
-class Document(
-    RootModel[
-        Union[
-            Agent,
-            Application,
-            AuthorizationProviderList,
-            Flow,
-            IndexList,
-            ModelList,
-            ToolList,
-            TypeList,
-            VariableList,
-        ]
-    ]
-):
+DocumentType = Union[
+    Agent,
+    Application,
+    AuthorizationProviderList,
+    Flow,
+    IndexList,
+    ModelList,
+    ToolList,
+    TypeList,
+    VariableList,
+]
+
+
+class Document(RootModel[DocumentType]):
     """Schema for any valid QType document structure.
 
     This allows validation of standalone lists of components, individual components,
     or full QType application specs. Supports modular composition and reuse.
     """
 
-    root: Union[
-        Agent,
-        Application,
-        AuthorizationProviderList,
-        Flow,
-        IndexList,
-        ModelList,
-        ToolList,
-        TypeList,
-        VariableList,
-    ]
+    root: DocumentType

qtype/dsl/validator.py

@@ -80,7 +80,7 @@ def _update_map_with_unique_check(
             # This is a special case where we do not want to add the string itself.
             continue
         # Note: There is no current abstraction for the `id` field, so we assume it exists.
-        obj_id = obj.id
+        obj_id = obj.id  # type: ignore[attr-defined]
         # If the object already exists in the map, we check if it is the same object.
         # If it is not the same object, we raise an error.
         # This ensures that we do not have duplicate components with the same ID.
@@ -449,9 +449,14 @@ def validate(
                     f"Chat flow {flow.id} must have at least one input variable of type ChatMessage."
                 )
             if (
-                len(flow.outputs) != 1
-                or flow.outputs[0].type != qtype.dsl.domain_types.ChatMessage
-            ):  # type: ignore
+                not flow.outputs
+                or len(flow.outputs) != 1
+                or (
+                    isinstance(flow.outputs[0], dsl.Variable)
+                    and flow.outputs[0].type
+                    != qtype.dsl.domain_types.ChatMessage
+                )
+            ):
                 raise QTypeValidationError(
                     f"Chat flow {flow.id} must have exactly one output variable of type ChatMessage."
                 )

qtype/interpreter/api.py

@@ -1,11 +1,16 @@
 from __future__ import annotations
 
+import logging
 from pathlib import Path
 
-from fastapi import FastAPI, HTTPException
+import pandas as pd
+from fastapi import FastAPI, HTTPException, Query
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 
+from qtype.dsl.base_types import StepCardinality
+from qtype.interpreter.batch.flow import batch_execute_flow
+from qtype.interpreter.batch.types import BatchConfig, ErrorMode
 from qtype.interpreter.flow import execute_flow
 from qtype.interpreter.typing import (
     create_input_type_model,
@@ -82,52 +87,103 @@ class APIExecutor:
         """Create a dynamic POST endpoint for a specific flow."""
         flow_id = flow.id
 
+        # determine if this is a batch inference
+        is_batch = flow.cardinality == StepCardinality.many
+
         # Create dynamic request and response models for this flow
-        RequestModel = create_input_type_model(flow)
-        ResponseModel = create_output_type_model(flow)
+        RequestModel = create_input_type_model(flow, is_batch)
+        ResponseModel = create_output_type_model(flow, is_batch)
 
         # Create the endpoint function with proper model binding
-        def execute_flow_endpoint(request: RequestModel) -> ResponseModel:  # type: ignore
-            try:
-                # Make a copy of the flow to avoid modifying the original
-                # TODO: Use session to ensure memory is not used across requests.
-                flow_copy = flow.model_copy(deep=True)
-                # Set input values on the flow variables
-                if flow_copy.inputs:
-                    for var in flow_copy.inputs:
-                        # Get the value from the request using the variable ID
-                        request_dict = request.model_dump()  # type: ignore
-                        if var.id in request_dict:
-                            var.value = getattr(request, var.id)
-                        elif not var.is_set():
-                            raise HTTPException(
-                                status_code=400,
-                                detail=f"Required input '{var.id}' not provided",
-                            )
-
-                # Execute the flow
-                result_vars = execute_flow(flow_copy)
-
-                # Extract output values
-                outputs = {var.id: var.value for var in result_vars}
-
-                response_data = {
-                    "flow_id": flow_id,
-                    "outputs": outputs,
-                    "status": "success",
-                }
-
-                # Return the response using the dynamic model
-                return ResponseModel(**response_data)  # type: ignore
-
-            except Exception as e:
-                raise HTTPException(
-                    status_code=500, detail=f"Flow execution failed: {str(e)}"
-                )
+        if is_batch:
+
+            def execute_flow_endpoint(  # type: ignore
+                request: RequestModel,  # type: ignore
+                error_mode: ErrorMode = Query(
+                    default=ErrorMode.FAIL,
+                    description="Error handling mode for batch processing",
+                ),
+            ) -> ResponseModel:  # type: ignore
+                try:
+                    # Make a copy of the flow to avoid modifying the original
+                    # TODO: Use session to ensure memory is not used across requests.
+                    flow_copy = flow.model_copy(deep=True)
+                    # convert the inputs into a dataframe with a single row
+                    inputs = pd.DataFrame(
+                        [i.model_dump() for i in request.inputs]  # type: ignore
+                    )
+
+                    # Execute the flow
+                    results, errors = batch_execute_flow(
+                        flow_copy,
+                        inputs,
+                        batch_config=BatchConfig(error_mode=error_mode),
+                    )
+
+                    response_data = {
+                        "flow_id": flow_id,
+                        "outputs": results.to_dict(orient="records"),
+                        "errors": errors.to_dict(orient="records"),
+                        "num_results": len(results),
+                        "num_errors": len(errors),
+                        "num_inputs": len(inputs),
+                        "status": "success" if len(errors) == 0 else "partial",
+                    }
+
+                    # Return the response using the dynamic model
+                    return ResponseModel(**response_data)  # type: ignore
+
+                except Exception as e:
+                    logging.error("Batch Flow Execution Failed", exc_info=e)
+                    raise HTTPException(
+                        status_code=500,
+                        detail=f"Batch flow execution failed: {str(e)}",
+                    )
+        else:
+
+            def execute_flow_endpoint(request: RequestModel) -> ResponseModel:  # type: ignore
+                try:
+                    # Make a copy of the flow to avoid modifying the original
+                    # TODO: Use session to ensure memory is not used across requests.
+                    flow_copy = flow.model_copy(deep=True)
+                    # Set input values on the flow variables
+                    if flow_copy.inputs:
+                        for var in flow_copy.inputs:
+                            # Get the value from the request using the variable ID
+                            request_dict = request.model_dump()  # type: ignore
+                            if var.id in request_dict:
+                                var.value = getattr(request, var.id)
+                            elif not var.is_set():
+                                raise HTTPException(
+                                    status_code=400,
+                                    detail=f"Required input '{var.id}' not provided",
+                                )
+                    return flow_copy
+                    # Execute the flow
+                    result_vars = execute_flow(flow_copy)
+
+                    # Extract output values
+                    outputs = {var.id: var.value for var in result_vars}
+
+                    response_data = {
+                        "flow_id": flow_id,
+                        "outputs": outputs,
+                        "status": "success",
+                    }
+
+                    # Return the response using the dynamic model
+                    return ResponseModel(**response_data)  # type: ignore
+
+                except Exception as e:
+                    raise HTTPException(
+                        status_code=500,
+                        detail=f"Flow execution failed: {str(e)}",
+                    )
 
         # Set the function annotations properly for FastAPI
         execute_flow_endpoint.__annotations__ = {
             "request": RequestModel,
+            "error_mode": ErrorMode,
             "return": ResponseModel,
         }
 

qtype/interpreter/auth/__init__.py

@@ -0,0 +1,3 @@
+"""
+Authorization context managers for QType interpreter.
+"""