qtype 0.0.16__py3-none-any.whl → 0.1.1__py3-none-any.whl

qtype/dsl/model.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
 import inspect
+import sys
 from abc import ABC
 from enum import Enum
-from typing import Any, Literal, Type, Union
+from functools import partial
+from typing import Annotated, Any, Literal, Type, Union
 
 from pydantic import (
     BaseModel,
@@ -14,20 +16,23 @@ from pydantic import (
 )
 
 import qtype.dsl.domain_types as domain_types
-from qtype.dsl.base_types import (
+from qtype.base.types import (
+    BatchableStepMixin,
+    BatchConfig,
+    CachedStepMixin,
+    ConcurrentStepMixin,
     PrimitiveTypeEnum,
+    Reference,
     StepCardinality,
     StrictBaseModel,
 )
-from qtype.dsl.domain_types import ChatContent, ChatMessage, Embedding
-
-
-class StructuralTypeEnum(str, Enum):
-    """Represents a structured type that can be used in the DSL."""
-
-    object = "object"
-    array = "array"
-
+from qtype.dsl.domain_types import (
+    ChatContent,
+    ChatMessage,
+    Embedding,
+    RAGChunk,
+    RAGDocument,
+)
 
 DOMAIN_CLASSES = {
     name: obj
@@ -36,63 +41,164 @@ DOMAIN_CLASSES = {
 }
 
 
+def _resolve_list_type(
+    element_type_str: str, custom_type_registry: dict[str, Type[BaseModel]]
+) -> ListType:
+    """
+    Resolve a list element type and return a ListType.
+
+    Args:
+        element_type_str: The element type string (e.g., "text", "ChatMessage")
+        custom_type_registry: Registry of custom types
+
+    Returns:
+        ListType with resolved element type
+
+    Raises:
+        ValueError: If element type is invalid for lists
+    """
+    # Recursively resolve the element type
+    element_type = _resolve_variable_type(
+        element_type_str, custom_type_registry
+    )
+
+    # Allow both primitive types and custom types (but no nested lists)
+    if isinstance(element_type, PrimitiveTypeEnum):
+        return ListType(element_type=element_type)
+    elif isinstance(element_type, str):
+        # This is a custom type reference - store as string for later resolution
+        return ListType(element_type=element_type)
+    elif element_type in DOMAIN_CLASSES.values():
+        # Domain class - store its name as string reference
+        for name, cls in DOMAIN_CLASSES.items():
+            if cls == element_type:
+                return ListType(element_type=name)
+        return ListType(element_type=str(element_type))
+    else:
+        raise ValueError(
+            (
+                "List element type must be a primitive or custom type "
+                f"reference, got: {element_type}"
+            )
+        )
+
+
+def _resolve_primitive_type(type_str: str) -> PrimitiveTypeEnum | None:
+    """
+    Try to resolve a string as a primitive type.
+
+    Args:
+        type_str: The type string to resolve
+
+    Returns:
+        PrimitiveTypeEnum if it matches, None otherwise
+    """
+    try:
+        return PrimitiveTypeEnum(type_str)
+    except ValueError:
+        return None
+
+
+def _resolve_domain_type(type_str: str) -> Type[BaseModel] | None:
+    """
+    Try to resolve a string as a built-in domain entity class.
+
+    Args:
+        type_str: The type string to resolve
+
+    Returns:
+        Domain class if found, None otherwise
+    """
+    return DOMAIN_CLASSES.get(type_str)
+
+
+def _resolve_custom_type(
+    type_str: str, custom_type_registry: dict[str, Type[BaseModel]]
+) -> Type[BaseModel] | None:
+    """
+    Try to resolve a string as a custom type from the registry.
+
+    Args:
+        type_str: The type string to resolve
+        custom_type_registry: Registry of custom types
+
+    Returns:
+        Custom type class if found, None otherwise
+    """
+    return custom_type_registry.get(type_str)
+
+
 def _resolve_variable_type(
     parsed_type: Any, custom_type_registry: dict[str, Type[BaseModel]]
 ) -> Any:
-    """Resolve a type string to its corresponding PrimitiveTypeEnum or return as is."""
+    """
+    Resolve a type to its corresponding representation.
+
+    Handles primitive types, list types, domain types, and custom types.
+
+    Args:
+        parsed_type: The type to resolve (can be string or already resolved)
+        custom_type_registry: Registry of dynamically created custom types
+
+    Returns:
+        Resolved type (PrimitiveTypeEnum, ListType, domain class, or string)
+    """
     # If the type is already resolved or is a structured definition, pass it through.
     if not isinstance(parsed_type, str):
         return parsed_type
 
-    # --- Case 1: The type is a string ---
     # Check if it's a list type (e.g., "list[text]")
     if parsed_type.startswith("list[") and parsed_type.endswith("]"):
-        # Extract the element type from "list[element_type]"
         element_type_str = parsed_type[5:-1]  # Remove "list[" and "]"
+        return _resolve_list_type(element_type_str, custom_type_registry)
 
-        # Recursively resolve the element type
-        element_type = _resolve_variable_type(
-            element_type_str, custom_type_registry
-        )
+    # Try to resolve as primitive type
+    primitive = _resolve_primitive_type(parsed_type)
+    if primitive is not None:
+        return primitive
 
-        # Allow both primitive types and custom types (but no nested lists)
-        if isinstance(element_type, PrimitiveTypeEnum):
-            return ListType(element_type=element_type)
-        elif isinstance(element_type, str):
-            # This is a custom type reference - store as string for later resolution
-            return ListType(element_type=element_type)
-        elif element_type in DOMAIN_CLASSES.values():
-            # Domain class - store its name as string reference
-            for name, cls in DOMAIN_CLASSES.items():
-                if cls == element_type:
-                    return ListType(element_type=name)
-            return ListType(element_type=str(element_type))
-        else:
-            raise ValueError(
-                f"List element type must be a primitive type or custom type reference, got: {element_type}"
-            )
+    # Try to resolve as built-in domain entity class
+    domain = _resolve_domain_type(parsed_type)
+    if domain is not None:
+        return domain
 
-    # Try to resolve it as a primitive type first.
-    try:
-        return PrimitiveTypeEnum(parsed_type)
-    except ValueError:
-        pass  # Not a primitive, continue to the next check.
+    # Try to resolve as custom type
+    custom = _resolve_custom_type(parsed_type, custom_type_registry)
+    if custom is not None:
+        return custom
 
-    # Try to resolve it as a built-in Domain Entity class.
-    # (Assuming domain_types and inspect are defined elsewhere)
-    if parsed_type in DOMAIN_CLASSES:
-        return DOMAIN_CLASSES[parsed_type]
+    # If it's not any known type, return it as a string.
+    # This assumes it might be a forward reference to a custom type.
+    return parsed_type
 
-    # Check the registry of dynamically created custom types
-    if parsed_type in custom_type_registry:
-        return custom_type_registry[parsed_type]
 
-    # If it's not a primitive or a known domain entity, return it as a string.
-    # This assumes it might be a reference ID to another custom type.
-    return parsed_type
+def _resolve_type_field_validator(data: Any, info: ValidationInfo) -> Any:
+    """
+    Shared validator for resolving 'type' fields in models.
 
+    This validator resolves string-based type references using the custom
+    type registry from the validation context.
 
-class Variable(BaseModel):
+    Args:
+        data: The data dict being validated
+        info: Pydantic validation info containing context
+
+    Returns:
+        Updated data dict with resolved type field
+    """
+    if (
+        isinstance(data, dict)
+        and "type" in data
+        and isinstance(data["type"], str)
+    ):
+        # Get the registry of custom types from the validation context.
+        custom_types = (info.context or {}).get("custom_types", {})
+        resolved = _resolve_variable_type(data["type"], custom_types)
+        data["type"] = resolved
+    return data
+
+
+class Variable(StrictBaseModel):
     """Schema for a variable that can serve as input, output, or parameter within the DSL."""
 
     id: str = Field(
@@ -109,21 +215,24 @@ class Variable(BaseModel):
     @model_validator(mode="before")
     @classmethod
     def resolve_type(cls, data: Any, info: ValidationInfo) -> Any:
-        """
-        This validator runs during the main validation pass. It uses the
-        context to resolve string-based type references.
-        """
-        if (
-            isinstance(data, dict)
-            and "type" in data
-            and isinstance(data["type"], str)
-        ):
-            # Get the registry of custom types from the validation context.
-            custom_types = (info.context or {}).get("custom_types", {})
-            resolved = _resolve_variable_type(data["type"], custom_types)
-            # {'id': 'user_message', 'type': 'ChatMessage'}
-            data["type"] = resolved
-        return data
+        """Resolve string-based type references using the shared validator."""
+        return _resolve_type_field_validator(data, info)
+
+
+class SecretReference(StrictBaseModel):
+    """
+    A reference to a secret in the application's configured SecretManager.
+    This value is resolved at runtime by the interpreter.
+    """
+
+    secret_name: str = Field(
+        ...,
+        description="The name, ID, or ARN of the secret to fetch (e.g., 'my-project/db-password').",
+    )
+    key: str | None = Field(
+        default=None,
+        description="Optional key if the secret is a JSON blob or map (e.g., a specific key in a K8s secret).",
+    )
 
 
 class CustomType(StrictBaseModel):
@@ -145,20 +254,8 @@ class ToolParameter(BaseModel):
     @model_validator(mode="before")
     @classmethod
     def resolve_type(cls, data: Any, info: ValidationInfo) -> Any:
-        """
-        This validator runs during the main validation pass. It uses the
-        context to resolve string-based type references.
-        """
-        if (
-            isinstance(data, dict)
-            and "type" in data
-            and isinstance(data["type"], str)
-        ):
-            # Get the registry of custom types from the validation context.
-            custom_types = (info.context or {}).get("custom_types", {})
-            resolved = _resolve_variable_type(data["type"], custom_types)
-            data["type"] = resolved
-        return data
+        """Resolve string-based type references using the shared validator."""
+        return _resolve_type_field_validator(data, info)
 
 
 class ListType(BaseModel):
@@ -183,6 +280,8 @@ VariableType = (
     | Type[ChatMessage]
     | Type[ChatContent]
     | Type[BaseModel]
+    | Type[RAGDocument]
+    | Type[RAGChunk]
     | ListType
 )
 
@@ -190,28 +289,31 @@ VariableType = (
 class Model(StrictBaseModel):
     """Describes a generative model configuration, including provider and model ID."""
 
+    type: Literal["Model"] = "Model"
     id: str = Field(..., description="Unique ID for the model.")
-    auth: AuthProviderType | str | None = Field(
+    auth: Reference[AuthProviderType] | str | None = Field(
         default=None,
         description="AuthorizationProvider used for model access.",
     )
-    inference_params: dict[str, Any] | None = Field(
-        default=None,
+    inference_params: dict[str, Any] = Field(
+        default_factory=dict,
         description="Optional inference parameters like temperature or max_tokens.",
     )
     model_id: str | None = Field(
         default=None,
         description="The specific model name or ID for the provider. If None, id is used",
     )
-    # TODO(maybe): Make this an enum?
-    provider: str = Field(
-        ..., description="Name of the provider, e.g., openai or anthropic."
+    provider: Literal["openai", "anthropic", "aws-bedrock", "gcp-vertex"] = (
+        Field(
+            ..., description="Name of the provider, e.g., openai or anthropic."
+        )
     )
 
 
 class EmbeddingModel(Model):
     """Describes an embedding model configuration, extending the base Model class."""
 
+    type: Literal["EmbeddingModel"] = "EmbeddingModel"
     dimensions: int = Field(
         ...,
         description="Dimensionality of the embedding vectors produced by this model.",
@@ -243,20 +345,22 @@ class Memory(StrictBaseModel):
 #
 
 
-class Step(StrictBaseModel, ABC):
+class Step(CachedStepMixin, StrictBaseModel, ABC):
     """Base class for components that take inputs and produce outputs."""
 
     id: str = Field(..., description="Unique ID of this component.")
+    type: str = Field(..., description="Type of the step 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.",
+    inputs: list[Reference[Variable] | str] = Field(
+        default_factory=list,
+        description="References to the variables required by this step.",
     )
-    outputs: list[Variable | str] | None = Field(
-        default=None, description="Variable where output is stored."
+    outputs: list[Reference[Variable] | str] = Field(
+        default_factory=list,
+        description="References to the variables where output is stored.",
     )
 
 
@@ -264,50 +368,12 @@ class PromptTemplate(Step):
     """Defines a prompt template with a string format and variable bindings.
     This is used to generate prompts dynamically based on input variables."""
 
+    type: Literal["PromptTemplate"] = "PromptTemplate"  # type: ignore
     template: str = Field(
         ...,
         description="String template for the prompt with variable placeholders.",
     )
 
-    @model_validator(mode="after")
-    def set_default_outputs(self) -> "PromptTemplate":
-        """Set default output variable if none provided."""
-        if self.outputs is None:
-            self.outputs = [
-                Variable(id=f"{self.id}.prompt", type=PrimitiveTypeEnum.text)
-            ]
-        if len(self.outputs) != 1:
-            raise ValueError(
-                "PromptTemplate steps must have exactly one output variable -- the result of applying the template."
-            )
-        return self
-
-
-class Condition(Step):
-    """Conditional logic gate within a flow. Supports branching logic for execution based on variable values."""
-
-    # TODO: Add support for more complex conditions
-    else_: StepType | str | None = Field(
-        default=None,
-        alias="else",
-        description="Optional step to run if condition fails.",
-    )
-    equals: Variable | str | None = Field(
-        default=None, description="Match condition for equality check."
-    )
-    then: StepType | str = Field(
-        ..., description="Step to run if condition matches."
-    )
-
-    @model_validator(mode="after")
-    def set_default_outputs(self) -> "Condition":
-        """Set default output variable if none provided."""
-        if not self.inputs or len(self.inputs) != 1:
-            raise ValueError(
-                "Condition steps must have exactly one input variable."
-            )
-        return self
-
 
 class Tool(StrictBaseModel, ABC):
     """
@@ -319,12 +385,12 @@ class Tool(StrictBaseModel, ABC):
     description: str = Field(
         ..., description="Description of what the tool does."
     )
-    inputs: dict[str, ToolParameter] | None = Field(
-        default=None,
+    inputs: dict[str, ToolParameter] = Field(
+        default_factory=dict,
         description="Input parameters required by this tool.",
     )
-    outputs: dict[str, ToolParameter] | None = Field(
-        default=None,
+    outputs: dict[str, ToolParameter] = Field(
+        default_factory=dict,
         description="Output parameters produced by this tool.",
     )
 
@@ -332,6 +398,7 @@ class Tool(StrictBaseModel, ABC):
 class PythonFunctionTool(Tool):
     """Tool that calls a Python function."""
 
+    type: Literal["PythonFunctionTool"] = "PythonFunctionTool"
     function_name: str = Field(
         ..., description="Name of the Python function to call."
     )
@@ -344,34 +411,36 @@ class PythonFunctionTool(Tool):
 class APITool(Tool):
     """Tool that invokes an API endpoint."""
 
+    type: Literal["APITool"] = "APITool"
     endpoint: str = Field(..., description="API endpoint URL to call.")
     method: str = Field(
         default="GET",
         description="HTTP method to use (GET, POST, PUT, DELETE, etc.).",
     )
-    auth: AuthProviderType | str | None = Field(
+    auth: Reference[AuthProviderType] | str | None = Field(
         default=None,
         description="Optional AuthorizationProvider for API authentication.",
     )
-    headers: dict[str, str] | None = Field(
-        default=None,
+    headers: dict[str, str] = Field(
+        default_factory=dict,
         description="Optional HTTP headers to include in the request.",
     )
-    parameters: dict[str, ToolParameter] | None = Field(
-        default=None,
+    parameters: dict[str, ToolParameter] = Field(
+        default_factory=dict,
         description="Output parameters produced by this tool.",
     )
 
 
-class LLMInference(Step):
+class LLMInference(Step, ConcurrentStepMixin):
     """Defines a step that performs inference using a language model.
     It can take input variables and produce output variables based on the model's response."""
 
-    memory: Memory | str | None = Field(
+    type: Literal["LLMInference"] = "LLMInference"
+    memory: Reference[Memory] | str | None = Field(
         default=None,
-        description="Memory object to retain context across interactions.",
+        description="A reference to a Memory object to retain context across interactions.",
     )
-    model: ModelType | str = Field(
+    model: Reference[Model] | str = Field(
         ..., description="The model to use for inference."
     )
     system_message: str | None = Field(
@@ -379,43 +448,70 @@ class LLMInference(Step):
         description="Optional system message to set the context for the model.",
     )
 
-    @model_validator(mode="after")
-    def set_default_outputs(self) -> "LLMInference":
-        """Set default output variable if none provided."""
-        if self.outputs is None:
-            self.outputs = [
-                Variable(id=f"{self.id}.response", type=PrimitiveTypeEnum.text)
-            ]
-        return self
+
+class InvokeEmbedding(Step, ConcurrentStepMixin):
+    """Defines a step that generates embeddings using an embedding model.
+    It takes input variables and produces output variables containing the embeddings."""
+
+    type: Literal["InvokeEmbedding"] = "InvokeEmbedding"
+    model: Reference[EmbeddingModel] | str = Field(
+        ..., description="The embedding model to use."
+    )
 
 
 class Agent(LLMInference):
     """Defines an agent that can perform tasks and make decisions based on user input and context."""
 
-    tools: list[ToolType | str] = Field(
-        ..., description="List of tools available to the agent."
+    type: Literal["Agent"] = "Agent"
+
+    tools: list[Reference[ToolType] | str] = Field(
+        default_factory=list,
+        description="List of tools available to the agent.",
     )
 
 
-class Flow(Step):
+class Flow(StrictBaseModel):
     """Defines a flow of steps that can be executed in sequence or parallel.
     If input or output variables are not specified, they are inferred from
-    the first and last step, respectively.
-    """
+    the first and last step, respectively."""
 
+    id: str = Field(..., description="Unique ID of the flow.")
+    type: Literal["Flow"] = "Flow"
     description: str | None = Field(
         default=None, description="Optional description of the flow."
     )
+    steps: list[StepType | Reference[StepType]] = Field(
+        default_factory=list,
+        description="List of steps or references to steps",
+    )
 
-    cardinality: StepCardinality = Field(
-        default=StepCardinality.auto,
-        description="The cardinality of the flow, inferred from its steps when set to 'auto'.",
+    interface: FlowInterface | None = Field(default=None)
+    variables: list[Variable] = Field(
+        default_factory=list,
+        description="List of variables available at the application scope.",
+    )
+    inputs: list[Reference[Variable] | str] = Field(
+        default_factory=list,
+        description="Input variables required by this step.",
     )
+    outputs: list[Reference[Variable] | str] = Field(
+        default_factory=list, description="Resulting variables"
+    )
+
+
+class FlowInterface(StrictBaseModel):
+    """
+    Defines the public-facing contract for a Flow, guiding the UI
+    and session management.
+    """
 
-    mode: Literal["Complete", "Chat"] = "Complete"
+    # 1. Tells the UI how to render this flow
+    type: Literal["Complete", "Conversational"] = "Complete"
 
-    steps: list[StepType | str] = Field(
-        default_factory=list, description="List of steps or step IDs."
+    # 2. Declares which inputs are "sticky" and persisted in the session
+    session_inputs: list[Reference[Variable] | str] = Field(
+        default_factory=list,
+        description="A list of input variable IDs that are set once and then persisted across a session.",
     )
 
 
@@ -430,53 +526,117 @@ class Decoder(Step):
     """Defines a step that decodes string data into structured outputs.
 
     If parsing fails, the step will raise an error and halt execution.
-    Use conditional logic in your flow to handle potential parsing errors.
-    """
+    Use conditional logic in your flow to handle potential parsing errors."""
+
+    type: Literal["Decoder"] = "Decoder"
 
     format: DecoderFormat = Field(
         DecoderFormat.json,
         description="Format in which the decoder processes data. Defaults to JSON.",
     )
 
-    @model_validator(mode="after")
-    def set_default_outputs(self) -> "Decoder":
-        """Set default output variable if none provided."""
-
-        if (
-            self.inputs is None
-            or len(self.inputs) != 1
-            or (
-                isinstance(self.inputs[0], Variable)
-                and self.inputs[0].type != PrimitiveTypeEnum.text
-            )
-        ):
-            raise ValueError(
-                f"Decoder steps must have exactly one input variable of type 'text'. Found: {self.inputs}"
-            )
-        if self.outputs is None:
-            raise ValueError(
-                "Decoder steps must have at least one output variable defined."
-            )
-        return self
 
+class Echo(Step):
+    """Defines a step that echoes its inputs as outputs.
+
+    Useful for debugging flows by inspecting variable values at a specific
+    point in the execution pipeline. The step simply passes through all input
+    variables as outputs without modification.
+    """
 
-class Invoke(Step):
+    type: Literal["Echo"] = "Echo"
+
+
+class FieldExtractor(Step):
+    """Extracts specific fields from input data using JSONPath expressions.
+
+    This step uses JSONPath syntax to extract data from structured inputs
+    (Pydantic models, dicts, lists). The input is first converted to a dict
+    using model_dump() if it's a Pydantic model, then the JSONPath expression
+    is evaluated.
+
+    If the JSONPath matches multiple values, the step yields multiple output
+    messages (1-to-many cardinality). If it matches a single value, it yields
+    one output message. If it matches nothing, it raises an error.
+
+    The extracted data is used to construct the output variable by passing it
+    as keyword arguments to the output type's constructor.
+
+    Example JSONPath expressions:
+    - `$.field_name` - Extract a single field
+    - `$.items[*]` - Extract all items from a list
+    - `$.items[?(@.price > 10)]` - Filter items by condition
+    """
+
+    type: Literal["FieldExtractor"] = "FieldExtractor"
+    json_path: str = Field(
+        ...,
+        description="JSONPath expression to extract data from the input. Uses jsonpath-ng syntax.",
+    )
+
+
+class InvokeTool(Step, ConcurrentStepMixin):
     """Invokes a tool with input and output bindings."""
 
-    tool: ToolType | str = Field(
+    type: Literal["InvokeTool"] = "InvokeTool"
+
+    tool: Reference[ToolType] | str = Field(
         ...,
         description="Tool to invoke.",
     )
     input_bindings: dict[str, str] = Field(
         ...,
-        description="Mapping from step input IDs to tool input parameter names.",
+        description="Mapping from variable references to tool input parameter names.",
     )
     output_bindings: dict[str, str] = Field(
         ...,
-        description="Mapping from tool output parameter names to step output IDs.",
+        description="Mapping from variable references to tool output parameter names.",
     )
 
 
+class InvokeFlow(Step):
+    """Invokes a flow with input and output bindings."""
+
+    type: Literal["InvokeFlow"] = "InvokeFlow"
+
+    flow: Reference[Flow] | str = Field(
+        ...,
+        description="Flow to invoke.",
+    )
+    input_bindings: dict[Reference[Variable], str] = Field(
+        ...,
+        description="Mapping from variable references to flow input variable IDs.",
+    )
+    output_bindings: dict[Reference[Variable], str] = Field(
+        ...,
+        description="Mapping from variable references to flow output variable IDs.",
+    )
+
+
+#
+# ---------------- Secret Manager Component ----------------
+#
+
+
+class SecretManager(StrictBaseModel, ABC):
+    """Base class for secret manager configurations."""
+
+    id: str = Field(
+        ..., description="Unique ID for this secret manager configuration."
+    )
+    type: str = Field(..., description="The type of secret manager.")
+    auth: Reference[AuthProviderType] | str = Field(
+        ...,
+        description="AuthorizationProvider used to access this secret manager.",
+    )
+
+
+class AWSSecretManager(SecretManager):
+    """Configuration for AWS Secrets Manager."""
+
+    type: Literal["aws_secret_manager"] = "aws_secret_manager"
+
+
 #
 # ---------------- Observability and Authentication Components ----------------
 #
@@ -495,7 +655,9 @@ class APIKeyAuthProvider(AuthorizationProvider):
     """API key-based authentication provider."""
 
     type: Literal["api_key"] = "api_key"
-    api_key: str = Field(..., description="API key for authentication.")
+    api_key: str | SecretReference = Field(
+        ..., description="API key for authentication."
+    )
     host: str | None = Field(
         default=None, description="Base URL or domain of the provider."
     )
@@ -505,7 +667,9 @@ class BearerTokenAuthProvider(AuthorizationProvider):
     """Bearer token authentication provider."""
 
     type: Literal["bearer_token"] = "bearer_token"
-    token: str = Field(..., description="Bearer token for authentication.")
+    token: str | SecretReference = Field(
+        ..., description="Bearer token for authentication."
+    )
 
 
 class OAuth2AuthProvider(AuthorizationProvider):
@@ -513,26 +677,60 @@ class OAuth2AuthProvider(AuthorizationProvider):
 
     type: Literal["oauth2"] = "oauth2"
     client_id: str = Field(..., description="OAuth2 client ID.")
-    client_secret: str = Field(..., description="OAuth2 client secret.")
+    client_secret: str | SecretReference = Field(
+        ..., description="OAuth2 client secret."
+    )
     token_url: str = Field(..., description="Token endpoint URL.")
-    scopes: list[str] | None = Field(
-        default=None, description="OAuth2 scopes required."
+    scopes: list[str] = Field(
+        default_factory=list, description="OAuth2 scopes required."
     )
 
 
+class VertexAuthProvider(AuthorizationProvider):
+    """Google Vertex authentication provider supporting gcloud profile or service account."""
+
+    type: Literal["vertex"] = "vertex"
+    profile_name: str | None = Field(
+        default=None,
+        description="Local gcloud profile name (if using existing CLI credentials).",
+    )
+    project_id: str | None = Field(
+        default=None,
+        description="Explicit GCP project ID override (if different from profile).",
+    )
+    service_account_file: str | None = Field(
+        default=None,
+        description="Path to a service account JSON key file.",
+    )
+    region: str | None = Field(
+        default=None,
+        description="Vertex region (e.g., us-central1).",
+    )
+
+    @model_validator(mode="after")
+    def validate_vertex_auth(self) -> VertexAuthProvider:
+        """Ensure at least one credential source is provided."""
+        if not (self.profile_name or self.service_account_file):
+            raise ValueError(
+                "VertexAuthProvider requires either a profile_name or a "
+                "service_account_file."
+            )
+        return self
+
+
 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(
+    access_key_id: str | SecretReference | None = Field(
         default=None, description="AWS access key ID."
     )
-    secret_access_key: str | None = Field(
+    secret_access_key: str | SecretReference | None = Field(
         default=None, description="AWS secret access key."
     )
-    session_token: str | None = Field(
+    session_token: str | SecretReference | None = Field(
         default=None,
         description="AWS session token for temporary credentials.",
     )
@@ -557,7 +755,7 @@ class AWSAuthProvider(AuthorizationProvider):
     region: str | None = Field(default=None, description="AWS region.")
 
     @model_validator(mode="after")
-    def validate_aws_auth(self) -> "AWSAuthProvider":
+    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
@@ -585,13 +783,18 @@ class TelemetrySink(StrictBaseModel):
     id: str = Field(
         ..., description="Unique ID of the telemetry sink configuration."
     )
-    auth: AuthProviderType | str | None = Field(
+    provider: Literal["Phoenix", "Langfuse"] = "Phoenix"
+    auth: Reference[AuthProviderType] | str | None = Field(
         default=None,
         description="AuthorizationProvider used to authenticate telemetry data transmission.",
     )
-    endpoint: str = Field(
+    endpoint: str | SecretReference = Field(
         ..., description="URL endpoint where telemetry data will be sent."
     )
+    args: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Additional configuration arguments specific to the telemetry sink type.",
+    )
 
 
 #
@@ -616,48 +819,53 @@ class Application(StrictBaseModel):
     )
 
     # Core components
-    memories: list[Memory] | None = Field(
-        default=None,
+    memories: list[Memory] = Field(
+        default_factory=list,
         description="List of memory definitions used in this application.",
     )
-    models: list[ModelType] | None = Field(
-        default=None, description="List of models used in this application."
+    models: list[ModelType] = Field(
+        default_factory=list,
+        description="List of models used in this application.",
     )
-    types: list[CustomType] | None = Field(
-        default=None,
+    types: list[CustomType] = Field(
+        default_factory=list,
         description="List of custom types defined in this application.",
     )
-    variables: list[Variable] | None = Field(
-        default=None, description="List of variables used in this application."
-    )
 
     # Orchestration
-    flows: list[Flow] | None = Field(
-        default=None, description="List of flows defined in this application."
+    flows: list[Flow] = Field(
+        default_factory=list,
+        description="List of flows defined in this application.",
     )
 
     # External integrations
-    auths: list[AuthProviderType] | None = Field(
-        default=None,
+    auths: list[AuthProviderType] = Field(
+        default_factory=list,
         description="List of authorization providers used for API access.",
     )
-    tools: list[ToolType] | None = Field(
-        default=None,
+    tools: list[ToolType] = Field(
+        default_factory=list,
         description="List of tools available in this application.",
     )
-    indexes: list[IndexType] | None = Field(
-        default=None,
+    indexes: list[IndexType] = Field(
+        default_factory=list,
         description="List of indexes available for search operations.",
     )
 
+    # Secret management
+    secret_manager: SecretManagerType | None = Field(
+        default=None,
+        description="Optional secret manager configuration for the application.",
+    )
+
     # Observability
     telemetry: TelemetrySink | None = Field(
         default=None, description="Optional telemetry sink for observability."
     )
 
     # Extensibility
-    references: list[Document] | None = Field(
-        default=None,
+    references: list[Document] = Field(
+        default_factory=list,
         description="List of other q-type documents you may use. This allows modular composition and reuse of components across applications.",
     )
 
@@ -667,6 +875,14 @@ class Application(StrictBaseModel):
 #
 
 
+class ConstantPath(StrictBaseModel):
+    uri: str = Field(..., description="A constant Fsspec URI.")
+
+
+# Let's the user use a constant path or reference a variable
+PathType = ConstantPath | Reference[Variable] | str
+
+
 class Source(Step):
     """Base class for data sources"""
 
@@ -680,135 +896,161 @@ class Source(Step):
 class SQLSource(Source):
     """SQL database source that executes queries and emits rows."""
 
+    type: Literal["SQLSource"] = "SQLSource"
     query: str = Field(
         ..., description="SQL query to execute. Inputs are injected as params."
     )
-    connection: str = Field(
+    connection: str | SecretReference = Field(
         ...,
         description="Database connection string or reference to auth provider. Typically in SQLAlchemy format.",
     )
-    auth: AuthProviderType | str | None = Field(
+    auth: Reference[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 FileSource(Source):
     """File source that reads data from a file using fsspec-compatible URIs."""
 
-    path: str | None = Field(
-        default=None,
-        description="fsspec-compatible URI to read from. If None, expects 'path' input variable.",
+    type: Literal["FileSource"] = "FileSource"
+    path: PathType = Field(
+        default=...,
+        description="Reference to a variable with an fsspec-compatible URI to read from, or the uri itself.",
     )
 
-    @model_validator(mode="after")
-    def validate_file_source(self) -> "FileSource":
-        """Validate that either path is specified or 'path' input variable exists."""
-        if self.path is None:
-            # Check if 'path' input variable exists
-            if self.inputs is None:
-                raise ValueError(
-                    "FileSource must either specify 'path' field or have a 'path' input variable."
-                )
-
-            path_input_exists = any(
-                (isinstance(inp, Variable) and inp.id == "path")
-                or (isinstance(inp, str) and inp == "path")
-                for inp in self.inputs
-            )
 
-            if not path_input_exists:
-                raise ValueError(
-                    "FileSource must either specify 'path' field or have a 'path' input variable."
-                )
+class Writer(Step, BatchableStepMixin):
+    """Base class for things that write data in batches."""
 
-        return self
+    id: str = Field(..., description="Unique ID of the data writer.")
 
 
-class Sink(Step):
-    """Base class for data sinks"""
+class FileWriter(Writer, BatchableStepMixin):
+    """File writer that writes data to a file using fsspec-compatible URIs."""
 
-    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.",
+    type: Literal["FileWriter"] = "FileWriter"
+    path: PathType = Field(
+        default=...,
+        description="Reference to a variable with an fsspec-compatible URI to read from, or the uri itself.",
+    )
+    batch_config: BatchConfig = Field(
+        default_factory=partial(BatchConfig, batch_size=sys.maxsize),
+        description="Configuration for processing the input stream in batches. If omitted, the step processes items one by one.",
+    )
+
+
+class Aggregate(Step):
+    """
+    A terminal step that consumes an entire input stream and produces a single
+    summary message with success/error counts.
+    """
+
+    type: Literal["Aggregate"] = "Aggregate"
+    cardinality: Literal[StepCardinality.one] = StepCardinality.one
+
+    # Outputs are now optional. The user can provide 0, 1, 2, or 3 names.
+    # The order will be: success_count, error_count, total_count
+    outputs: list[Reference[Variable] | str] = Field(
+        default_factory=list,
+        description="References to the variables for the output. There should be one and only one output with type AggregateStats",
     )
 
 
-class FileSink(Sink):
-    """File sink that writes data to a file using fsspec-compatible URIs."""
+#
+# ---------------- Retrieval Augmented Generation Components ----------------
+#
 
-    path: str | None = Field(
+
+class DocumentSource(Source):
+    """A source of documents that will be used in retrieval augmented generation.
+    It uses LlamaIndex readers to load one or more raw Documents
+    from a specified path or system (e.g., Google Drive, web page).
+    See https://github.com/run-llama/llama_index/tree/main/llama-index-integrations/readers
+    """
+
+    type: Literal["DocumentSource"] = "DocumentSource"
+    reader_module: str = Field(
+        ...,
+        description="Module path of the LlamaIndex Reader).",
+    )
+    args: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Reader-specific arguments to pass to the Reader constructor.",
+    )
+    loader_args: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Loader-specific arguments to pass to the load_data method.",
+    )
+    auth: Reference[AuthProviderType] | str | None = Field(
         default=None,
-        description="fsspec-compatible URI to write to. If None, expects 'path' input variable.",
+        description="AuthorizationProvider for accessing the source.",
     )
 
-    @model_validator(mode="after")
-    def validate_file_sink(self) -> "FileSink":
-        """Validate that either path is specified or 'path' input variable exists."""
-        # Ensure user does not set any output variables
-        if self.outputs is not None and len(self.outputs) > 0:
-            raise ValueError(
-                "FileSink outputs are automatically generated. Do not specify outputs."
-            )
 
-        # Automatically set the output variable
-        self.outputs = [Variable(id=f"{self.id}-file-uri", type="text")]
+class DocToTextConverter(Step, ConcurrentStepMixin):
+    """Defines a step to convert raw documents (e.g., PDF, DOCX) loaded by a DocumentSource into plain text
+    using an external tool like Docling or LlamaParse for pre-processing before chunking.
+    The input and output are both RAGDocument, but the output after processing with have content of type markdown.
+    """
 
-        if self.path is None:
-            # Check if 'path' input variable exists
-            if self.inputs is None:
-                raise ValueError(
-                    "FileSink must either specify 'path' field or have a 'path' input variable."
-                )
+    type: Literal["DocToTextConverter"] = "DocToTextConverter"
 
-            path_input_exists = any(
-                (isinstance(inp, Variable) and inp.id == "path")
-                or (isinstance(inp, str) and inp == "path")
-                for inp in self.inputs
-            )
 
-            if not path_input_exists:
-                raise ValueError(
-                    "FileSink must either specify 'path' field or have a 'path' input variable."
-                )
+class DocumentSplitter(Step, ConcurrentStepMixin):
+    """Configuration for chunking/splitting documents into embeddable nodes/chunks."""
 
-        return self
+    type: Literal["DocumentSplitter"] = "DocumentSplitter"
+    cardinality: Literal[StepCardinality.many] = Field(
+        default=StepCardinality.many,
+        description="Consumes one document and emits 0...N nodes/chunks.",
+    )
+
+    splitter_name: str = Field(
+        default="SentenceSplitter",
+        description="Name of the LlamaIndex TextSplitter class.",
+    )
+    chunk_size: int = Field(default=1024, description="Size of each chunk.")
+    chunk_overlap: int = Field(
+        default=20, description="Overlap between consecutive chunks."
+    )
+    args: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Additional arguments specific to the chosen splitter class.",
+    )
 
 
-#
-# ---------------- Retrieval Augmented Generation Components ----------------
-#
+class DocumentEmbedder(Step, ConcurrentStepMixin):
+    """Embeds document chunks using a specified embedding model."""
+
+    type: Literal["DocumentEmbedder"] = "DocumentEmbedder"
+    cardinality: Literal[StepCardinality.many] = Field(
+        default=StepCardinality.many,
+        description="Consumes one chunk and emits one embedded chunk.",
+    )
+    model: Reference[EmbeddingModel] | str = Field(
+        ..., description="Embedding model to use for vectorization."
+    )
 
 
 class Index(StrictBaseModel, ABC):
     """Base class for searchable indexes that can be queried by search steps."""
 
     id: str = Field(..., description="Unique ID of the index.")
-    args: dict[str, Any] | None = Field(
-        default=None,
+    args: dict[str, Any] = Field(
+        default_factory=dict,
         description="Index-specific configuration and connection parameters.",
     )
-    auth: AuthProviderType | str | None = Field(
+    auth: Reference[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(
+class IndexUpsert(Writer):
+    type: Literal["IndexUpsert"] = "IndexUpsert"
+    index: Reference[IndexType] | str = Field(
         ..., description="Index to upsert into (object or ID reference)."
     )
 
@@ -816,7 +1058,12 @@ class IndexUpsert(Sink):
 class VectorIndex(Index):
     """Vector database index for similarity search using embeddings."""
 
-    embedding_model: EmbeddingModel | str = Field(
+    type: Literal["VectorIndex"] = "VectorIndex"
+    module: str = Field(
+        ...,
+        description="Python module path for the vector store implementation (e.g., 'llama_index.vector_stores.qdrant.QdrantVectorStore').",
+    )
+    embedding_model: Reference[EmbeddingModel] | str = Field(
         ...,
         description="Embedding model used to vectorize queries and documents.",
     )
@@ -825,67 +1072,53 @@ class VectorIndex(Index):
 class DocumentIndex(Index):
     """Document search index for text-based search (e.g., Elasticsearch, OpenSearch)."""
 
-    # TODO: add anything that is needed for document search indexes
-    pass
+    type: Literal["DocumentIndex"] = "DocumentIndex"
+    endpoint: str = Field(
+        ...,
+        description="URL endpoint for the search cluster (e.g., https://my-cluster.es.amazonaws.com).",
+    )
 
 
 class Search(Step, ABC):
     """Base class for search operations against indexes."""
 
-    filters: dict[str, Any] | None = Field(
-        default=None, description="Optional filters to apply during search."
+    filters: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Optional filters to apply during search.",
     )
-    index: IndexType | str = Field(
+    index: Reference[IndexType] | str = Field(
         ..., description="Index to search against (object or ID reference)."
     )
 
 
-class VectorSearch(Search):
+class VectorSearch(Search, BatchableStepMixin):
     """Performs vector similarity search against a vector index."""
 
+    type: Literal["VectorSearch"] = "VectorSearch"
     default_top_k: int | None = Field(
         default=50,
         description="Number of top results to retrieve if not provided in the inputs.",
     )
 
-    @model_validator(mode="after")
-    def set_default_inputs_outputs(self) -> "VectorSearch":
-        """Set default input and output variables if none provided."""
-        if self.inputs is None:
-            self.inputs = [
-                Variable(id="top_k", type=PrimitiveTypeEnum.int),
-                Variable(id="query", type=PrimitiveTypeEnum.text),
-            ]
-
-        if self.outputs is None:
-            self.outputs = [Variable(id=f"{self.id}.results", type=Embedding)]
-        return self
-
 
-class DocumentSearch(Search):
+class DocumentSearch(Search, ConcurrentStepMixin):
     """Performs document search against a document index."""
 
-    @model_validator(mode="after")
-    def set_default_inputs_outputs(self) -> "DocumentSearch":
-        """Set default input and output variables if none provided."""
-        if self.inputs is None:
-            self.inputs = [Variable(id="query", type=PrimitiveTypeEnum.text)]
-
-        if self.outputs is None:
-            self.outputs = [
-                Variable(id=f"{self.id}.results", type=PrimitiveTypeEnum.text)
-            ]
-        return self
+    type: Literal["DocumentSearch"] = "DocumentSearch"
 
 
 # Create a union type for all tool types
-ToolType = Union[
-    APITool,
-    PythonFunctionTool,
+ToolType = Annotated[
+    Union[
+        APITool,
+        PythonFunctionTool,
+    ],
+    Field(discriminator="type"),
 ]
 
 # Create a union type for all source types
 SourceType = Union[
+    DocumentSource,
     FileSource,
     SQLSource,
 ]
@@ -896,36 +1129,61 @@ AuthProviderType = Union[
     BearerTokenAuthProvider,
     AWSAuthProvider,
     OAuth2AuthProvider,
+    VertexAuthProvider,
+]
+
+# Create a union type for all secret manager types
+SecretManagerType = Annotated[
+    Union[
+        AWSSecretManager
+        # Add future managers like KubernetesSecretManager here
+    ],
+    Field(discriminator="type"),
 ]
 
 # Create a union type for all step types
-StepType = Union[
-    Agent,
-    Condition,
-    Decoder,
-    DocumentSearch,
-    FileSink,
-    FileSource,
-    Flow,
-    IndexUpsert,
-    Invoke,
-    LLMInference,
-    PromptTemplate,
-    SQLSource,
-    Sink,
-    VectorSearch,
+StepType = Annotated[
+    Union[
+        Agent,
+        Aggregate,
+        Decoder,
+        DocToTextConverter,
+        DocumentEmbedder,
+        DocumentSearch,
+        DocumentSplitter,
+        DocumentSource,
+        Echo,
+        FieldExtractor,
+        FileSource,
+        FileWriter,
+        IndexUpsert,
+        InvokeEmbedding,
+        InvokeFlow,
+        InvokeTool,
+        LLMInference,
+        PromptTemplate,
+        SQLSource,
+        VectorSearch,
+    ],
+    Field(discriminator="type"),
 ]
 
 # Create a union type for all index types
-IndexType = Union[
-    DocumentIndex,
-    VectorIndex,
+IndexType = Annotated[
+    Union[
+        DocumentIndex,
+        VectorIndex,
+    ],
+    Field(discriminator="type"),
 ]
 
 # Create a union type for all model types
-ModelType = Union[
-    EmbeddingModel,
-    Model,
+ModelType = Annotated[
+    Union[
+        EmbeddingModel,
+        Model,
+    ],
+    Field(discriminator="type"),
 ]
 
 #
@@ -940,12 +1198,6 @@ class AuthorizationProviderList(RootModel[list[AuthProviderType]]):
     root: list[AuthProviderType]
 
 
-class IndexList(RootModel[list[IndexType]]):
-    """Schema for a standalone list of indexes."""
-
-    root: list[IndexType]
-
-
 class ModelList(RootModel[list[ModelType]]):
     """Schema for a standalone list of models."""
 
@@ -971,11 +1223,8 @@ class VariableList(RootModel[list[Variable]]):
 
 
 DocumentType = Union[
-    Agent,
     Application,
     AuthorizationProviderList,
-    Flow,
-    IndexList,
     ModelList,
     ToolList,
     TypeList,