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

qtype/interpreter/executors/echo_executor.py

@@ -0,0 +1,63 @@
+from typing import AsyncIterator
+
+from qtype.interpreter.base.base_step_executor import StepExecutor
+from qtype.interpreter.base.executor_context import ExecutorContext
+from qtype.interpreter.types import FlowMessage
+from qtype.semantic.model import Echo
+
+
+class EchoExecutor(StepExecutor):
+    """Executor for Echo steps.
+
+    Passes through input variables as outputs without modification.
+    Useful for debugging flows by inspecting variable values at specific
+    points in the execution pipeline.
+    """
+
+    def __init__(
+        self,
+        step: Echo,
+        context: ExecutorContext,
+        **dependencies: object,
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, Echo):
+            raise ValueError("EchoExecutor can only execute Echo steps.")
+        self.step: Echo = step
+
+    async def process_message(
+        self,
+        message: FlowMessage,
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a single FlowMessage for the Echo step.
+
+        Reads all input variables from the message and copies them to
+        the output variables with the same IDs.
+
+        Args:
+            message: The FlowMessage to process.
+
+        Yields:
+            FlowMessage with the echoed variables.
+        """
+        try:
+            # Build a dict of output variable values by reading from inputs
+            output_vars = {}
+            for input_var in self.step.inputs:
+                value = message.variables.get(input_var.id)
+                # Find the corresponding output variable ID (should match)
+                for output_var in self.step.outputs:
+                    if output_var.id == input_var.id:
+                        output_vars[output_var.id] = value
+                        break
+
+            await self.stream_emitter.status(
+                f"Echoed {len(output_vars)} variable(s) in step {self.step.id}",
+            )
+            yield message.copy_with_variables(output_vars)
+
+        except Exception as e:
+            # Emit error event to stream so frontend can display it
+            await self.stream_emitter.error(str(e))
+            message.set_error(self.step.id, e)
+            yield message

qtype/interpreter/executors/field_extractor_executor.py

@@ -0,0 +1,160 @@
+from typing import Any, AsyncIterator
+
+from jsonpath_ng.ext import parse  # type: ignore[import-untyped]
+from pydantic import BaseModel
+
+from qtype.base.types import PrimitiveTypeEnum
+from qtype.dsl.model import ListType
+from qtype.interpreter.base.base_step_executor import StepExecutor
+from qtype.interpreter.base.executor_context import ExecutorContext
+from qtype.interpreter.types import FlowMessage
+from qtype.semantic.model import FieldExtractor
+
+
+class FieldExtractorExecutor(StepExecutor):
+    """Executor for FieldExtractor steps.
+
+    Extracts fields from input data using JSONPath expressions and
+    constructs output instances. Supports 1-to-many cardinality when
+    the JSONPath matches multiple values.
+    """
+
+    def __init__(
+        self,
+        step: FieldExtractor,
+        context: ExecutorContext,
+        **dependencies: object,
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, FieldExtractor):
+            raise ValueError(
+                "FieldExtractorExecutor can only execute FieldExtractor steps."
+            )
+        self.step: FieldExtractor = step
+
+        # Parse the JSONPath expression once at initialization
+        try:
+            self.jsonpath_expr = parse(self.step.json_path)
+        except Exception as e:
+            raise ValueError(
+                f"Invalid JSONPath expression '{self.step.json_path}': {e}"
+            ) from e
+
+    def _to_dict(self, value: Any) -> Any:
+        """Convert value to dict representation for JSONPath processing.
+
+        Args:
+            value: The value to convert (could be BaseModel, dict, list, etc.)
+
+        Returns:
+            Dict representation suitable for JSONPath processing
+        """
+        if isinstance(value, BaseModel):
+            return value.model_dump()
+        return value
+
+    def _construct_output(self, extracted_data: Any) -> Any:
+        """Construct the output value from extracted data.
+
+        Args:
+            extracted_data: The data extracted by JSONPath
+
+        Returns:
+            Constructed output value based on the output variable type
+        """
+        output_var = self.step.outputs[0]
+        output_type = output_var.type
+
+        # Handle primitive types - just return the extracted data
+        if isinstance(output_type, PrimitiveTypeEnum):
+            return extracted_data
+
+        # Handle list types
+        if isinstance(output_type, ListType):
+            # The extracted_data should already be a list
+            if not isinstance(extracted_data, list):
+                extracted_data = [extracted_data]
+            return extracted_data
+
+        # Handle BaseModel types (domain types and custom types)
+        if isinstance(output_type, type) and issubclass(
+            output_type, BaseModel
+        ):
+            # If extracted_data is a dict, use it as kwargs
+            if isinstance(extracted_data, dict):
+                return output_type(**extracted_data)
+            # If it's already the right type, return it
+            elif isinstance(extracted_data, output_type):
+                return extracted_data
+            else:
+                raise ValueError(
+                    (
+                        f"Cannot construct {output_type.__name__} from "
+                        f"{type(extracted_data).__name__}"
+                    )
+                )
+
+        # Fallback - return as-is
+        return extracted_data
+
+    async def process_message(
+        self,
+        message: FlowMessage,
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a single FlowMessage for the FieldExtractor step.
+
+        Args:
+            message: The FlowMessage to process.
+
+        Yields:
+            FlowMessage(s) with extracted and constructed output values.
+            Multiple messages may be yielded if JSONPath matches multiple values.
+        """
+        input_id = self.step.inputs[0].id
+        output_id = self.step.outputs[0].id
+
+        try:
+            # Get the input value
+            input_value = message.variables.get(input_id)
+            if input_value is None:
+                raise ValueError(
+                    f"Input variable '{input_id}' is not set or is None"
+                )
+
+            await self.stream_emitter.status(
+                f"Extracting fields using JSONPath: {self.step.json_path}"
+            )
+
+            # Convert input to dict for JSONPath processing
+            input_dict = self._to_dict(input_value)
+
+            # Apply JSONPath expression
+            matches = self.jsonpath_expr.find(input_dict)
+
+            if not matches:
+                raise ValueError(
+                    (
+                        f"JSONPath expression '{self.step.json_path}' "
+                        f"did not match any data in input"
+                    )
+                )
+
+            await self.stream_emitter.status(
+                f"JSONPath matched {len(matches)} value(s)"
+            )
+
+            # Yield one message per match (1-to-many)
+            for match in matches:
+                extracted_data = match.value
+
+                # Construct the output value
+                output_value = self._construct_output(extracted_data)
+
+                # Yield message with the constructed output
+                yield message.copy_with_variables({output_id: output_value})
+
+        except Exception as e:
+            # Emit error event to stream so frontend can display it
+            await self.stream_emitter.error(str(e))
+            message.set_error(self.step.id, e)
+            yield message

qtype/interpreter/executors/file_source_executor.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import AsyncIterator
+
+import fsspec
+import pandas as pd
+
+from qtype.interpreter.base.base_step_executor import StepExecutor
+from qtype.interpreter.base.executor_context import ExecutorContext
+from qtype.interpreter.types import FlowMessage
+from qtype.semantic.model import ConstantPath, FileSource
+
+
+class FileSourceExecutor(StepExecutor):
+    """Executor for FileSource steps."""
+
+    def __init__(
+        self, step: FileSource, context: ExecutorContext, **dependencies
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, FileSource):
+            raise ValueError(
+                "FileSourceExecutor can only execute FileSource steps."
+            )
+        self.step = step
+
+    async def process_message(
+        self,
+        message: FlowMessage,
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a single FlowMessage for the FileSource step.
+
+        Args:
+            message: The FlowMessage to process.
+
+        Yields:
+            FlowMessages with the results of processing.
+        """
+        output_columns = {output.id for output in self.step.outputs}
+
+        # get the path
+        if isinstance(self.step.path, ConstantPath):
+            file_path = self.step.path
+        else:
+            file_path = message.variables.get(self.step.path.id)
+            if not file_path:
+                raise ValueError(
+                    (
+                        f"FileSource step {self.step.id} requires a path "
+                        "variable."
+                    )
+                )
+            await self.stream_emitter.status(
+                f"Reading file from path: {file_path}"
+            )
+
+        # Determine file format from extension
+        file_path_str = (
+            file_path.uri if isinstance(file_path, ConstantPath) else file_path
+        )
+        extension = Path(file_path_str).suffix.lower()
+
+        # Use fsspec to open the file and read with pandas
+        with fsspec.open(file_path_str, "rb") as file_handle:
+            if extension == ".csv":
+                df = pd.read_csv(file_handle)  # type: ignore[arg-type]
+            elif extension == ".parquet":
+                df = pd.read_parquet(file_handle)  # type: ignore[arg-type]
+            elif extension == ".json":
+                df = pd.read_json(file_handle)  # type: ignore[arg-type]
+            elif extension == ".jsonl":
+                df = pd.read_json(
+                    file_handle,
+                    lines=True,  # type: ignore[arg-type]
+                )
+            else:
+                # Default to parquet if no extension or unknown
+                df = pd.read_parquet(file_handle)  # type: ignore[arg-type]
+
+        # confirm the outputs exist in the dataframe
+        columns = set(df.columns)
+        missing_columns = output_columns - columns
+        if missing_columns:
+            raise ValueError(
+                (
+                    f"File {file_path_str} missing expected columns: "
+                    f"{', '.join(missing_columns)}. Available columns: "
+                    f"{', '.join(columns)}"
+                )
+            )
+
+        for row in df.to_dict(orient="records"):
+            # Filter to only the expected output columns if they exist
+            row = {
+                str(k): v for k, v in row.items() if str(k) in output_columns
+            }
+            yield message.copy_with_variables(new_variables=row)
+        await self.stream_emitter.status(
+            f"Emitted {len(df)} rows from: {file_path_str}"
+        )

qtype/interpreter/executors/file_writer_executor.py

@@ -0,0 +1,110 @@
+from itertools import groupby
+from typing import AsyncIterator
+
+import fsspec
+import pandas as pd
+
+from qtype.interpreter.base.batch_step_executor import BatchedStepExecutor
+from qtype.interpreter.base.executor_context import ExecutorContext
+from qtype.interpreter.types import FlowMessage
+from qtype.semantic.model import ConstantPath, FileWriter, Variable
+
+
+class FileWriterExecutor(BatchedStepExecutor):
+    """Executor for FileWriter steps."""
+
+    def __init__(
+        self,
+        step: FileWriter,
+        context: ExecutorContext,
+        **dependencies,
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, FileWriter):
+            raise ValueError(
+                "FileWriterExecutor can only execute FileWriter steps."
+            )
+        self.step = step
+
+    def to_pandas(self, messages: list[FlowMessage]) -> pd.DataFrame:
+        """Convert a list of FlowMessages to a pandas DataFrame."""
+        records = [msg.variables for msg in messages]
+        return pd.DataFrame.from_records(records)
+
+    async def process_batch(
+        self,
+        batch: list[FlowMessage],
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a batch of FlowMessages for the FileWriter step.
+
+        Args:
+            batch: A list of FlowMessages to process.
+
+        Yields:
+            FlowMessages with the results of processing.
+        """
+        output_name = None
+        if len(self.step.outputs):
+            output_name = self.step.outputs[0].id
+
+        if isinstance(self.step.path, ConstantPath):
+            file_path = self.step.path.uri
+            df = self.to_pandas(batch)
+            # A fixed path is provided -- just write all of the data
+            await self.stream_emitter.status(
+                f"Writing {len(df)} records to {file_path}"
+            )
+            with fsspec.open(file_path, "wb") as file_handle:
+                df.to_parquet(file_handle, index=False)  # type: ignore[arg-type]
+            await self.stream_emitter.status(
+                f"Wrote {len(df)} records to {file_path}"
+            )
+            for msg in batch:
+                yield (
+                    msg
+                    if not output_name
+                    else msg.copy_with_variables({output_name: file_path})
+                )
+        else:
+            # Group messages by file path (path is a Variable in this branch)
+            if not isinstance(self.step.path, Variable):
+                raise ValueError(
+                    "Expected path to be a Variable in dynamic path case."
+                )
+
+            path_var_id = self.step.path.id
+
+            # Sort messages by file path for groupby
+            sorted_batch = sorted(
+                batch, key=lambda msg: msg.variables[path_var_id]
+            )
+
+            # Group messages by file path
+            grouped_messages = groupby(
+                sorted_batch, key=lambda msg: msg.variables[path_var_id]
+            )
+
+            distinct_paths = list(
+                set(msg.variables[path_var_id] for msg in batch)
+            )
+            await self.stream_emitter.status(
+                f"There are {len(distinct_paths)} different files to write."
+            )
+            for file_path, msg_group in grouped_messages:
+                msg_list = list(msg_group)
+                df_group = self.to_pandas(msg_list)
+                await self.stream_emitter.status(
+                    f"Writing {len(df_group)} records to {file_path}"
+                )
+                with fsspec.open(file_path, "wb") as file_handle:
+                    df_group.to_parquet(file_handle, index=False)  # type: ignore[arg-type]
+                await self.stream_emitter.status(
+                    f"Wrote {len(df_group)} records to {file_path}"
+                )
+                # Re-yield the original messages for this group
+                for msg in batch:
+                    yield (
+                        msg
+                        if not output_name
+                        else msg.copy_with_variables({output_name: file_path})
+                    )

qtype/interpreter/executors/index_upsert_executor.py

@@ -0,0 +1,228 @@
+"""Index upsert executor for inserting documents/chunks into indexes."""
+
+from __future__ import annotations
+
+import logging
+from typing import AsyncIterator
+
+from llama_index.core.schema import TextNode
+
+from qtype.dsl.domain_types import RAGChunk, RAGDocument
+from qtype.interpreter.base.batch_step_executor import BatchedStepExecutor
+from qtype.interpreter.base.executor_context import ExecutorContext
+from qtype.interpreter.conversions import (
+    to_llama_vector_store_and_retriever,
+    to_opensearch_client,
+)
+from qtype.interpreter.types import FlowMessage
+from qtype.semantic.model import DocumentIndex, IndexUpsert, VectorIndex
+
+logger = logging.getLogger(__name__)
+
+
+class IndexUpsertExecutor(BatchedStepExecutor):
+    """Executor for IndexUpsert steps supporting both vector and document indexes."""
+
+    def __init__(
+        self, step: IndexUpsert, context: ExecutorContext, **dependencies
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, IndexUpsert):
+            raise ValueError(
+                "IndexUpsertExecutor can only execute IndexUpsert steps."
+            )
+        self.step: IndexUpsert = step
+
+        # Determine index type and initialize appropriate client
+        if isinstance(self.step.index, VectorIndex):
+            # Vector index for RAGChunk embeddings
+            self._vector_store, _ = to_llama_vector_store_and_retriever(
+                self.step.index, self.context.secret_manager
+            )
+            self._opensearch_client = None
+            self.index_type = "vector"
+        elif isinstance(self.step.index, DocumentIndex):
+            # Document index for text-based search
+            self._opensearch_client = to_opensearch_client(
+                self.step.index, self.context.secret_manager
+            )
+            self._vector_store = None
+            self.index_type = "document"
+            self.index_name = self.step.index.name
+        else:
+            raise ValueError(
+                f"Unsupported index type: {type(self.step.index)}"
+            )
+
+    async def process_batch(
+        self, batch: list[FlowMessage]
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a batch of FlowMessages for the IndexUpsert step.
+
+        Args:
+            batch: A list of FlowMessages to process.
+
+        Yields:
+            FlowMessages: Success messages after upserting to the index
+        """
+        logger.info(
+            f"Executing IndexUpsert step: {self.step.id} with batch size: {len(batch)}"
+        )
+
+        try:
+            # Get the input variable (exactly one as validated by checker)
+            if not self.step.inputs:
+                raise ValueError("IndexUpsert step requires exactly one input")
+
+            input_var = self.step.inputs[0]
+
+            # Collect all RAGChunks or RAGDocuments from the batch
+            items_to_upsert = []
+            for message in batch:
+                input_data = message.variables.get(input_var.id)
+
+                if input_data is None:
+                    logger.warning(
+                        f"No data found for input: {input_var.id} in message"
+                    )
+                    continue
+
+                if not isinstance(input_data, (RAGChunk, RAGDocument)):
+                    raise ValueError(
+                        f"IndexUpsert only supports RAGChunk or RAGDocument "
+                        f"inputs. Got: {type(input_data)}"
+                    )
+
+                items_to_upsert.append(input_data)
+
+            # Upsert to appropriate index type
+            if items_to_upsert:
+                if self.index_type == "vector":
+                    await self._upsert_to_vector_store(items_to_upsert)
+                else:  # document index
+                    await self._upsert_to_document_index(items_to_upsert)
+
+                logger.info(
+                    f"Successfully upserted {len(items_to_upsert)} items "
+                    f"to {self.index_type} index in batch"
+                )
+
+                # Emit status update
+                index_type_display = (
+                    "vector index"
+                    if self.index_type == "vector"
+                    else "document index"
+                )
+                await self.stream_emitter.status(
+                    f"Upserted {len(items_to_upsert)} items to "
+                    f"{index_type_display}"
+                )
+
+            # Yield all input messages back (IndexUpsert typically doesn't have outputs)
+            for message in batch:
+                yield message
+
+        except Exception as e:
+            logger.error(f"Error in IndexUpsert step {self.step.id}: {e}")
+            # Emit error event to stream so frontend can display it
+            await self.stream_emitter.error(str(e))
+
+            # Mark all messages with the error and yield them
+            for message in batch:
+                message.set_error(self.step.id, e)
+                yield message
+
+    async def _upsert_to_vector_store(
+        self, items: list[RAGChunk | RAGDocument]
+    ) -> None:
+        """Upsert items to vector store.
+
+        Args:
+            items: List of RAGChunk or RAGDocument objects
+        """
+        # Convert to LlamaIndex TextNode objects
+        nodes = []
+        for item in items:
+            if isinstance(item, RAGChunk):
+                node = TextNode(
+                    id_=item.chunk_id,
+                    text=str(item.content),
+                    metadata=item.metadata,
+                    embedding=item.vector,
+                )
+            else:  # RAGDocument
+                # For documents, use file_id and convert content to string
+                node = TextNode(
+                    id_=item.file_id,
+                    text=str(item.content),
+                    metadata=item.metadata,
+                    embedding=None,  # Documents don't have embeddings
+                )
+            nodes.append(node)
+
+        # Batch upsert all nodes to the vector store
+        await self._vector_store.async_add(nodes)
+
+    async def _upsert_to_document_index(
+        self, items: list[RAGChunk | RAGDocument]
+    ) -> None:
+        """Upsert items to document index using bulk API.
+
+        Args:
+            items: List of RAGChunk or RAGDocument objects
+        """
+        # Build bulk request body
+        bulk_body = []
+        for item in items:
+            if isinstance(item, RAGChunk):
+                # Add index action
+                bulk_body.append(
+                    {
+                        "index": {
+                            "_index": self.index_name,
+                            "_id": item.chunk_id,
+                        }
+                    }
+                )
+                # Add document content
+                doc = {
+                    "text": str(item.content),
+                    "metadata": item.metadata,
+                }
+                # Include embedding if available
+                if item.vector:
+                    doc["embedding"] = item.vector
+                bulk_body.append(doc)
+            else:  # RAGDocument
+                # Add index action
+                bulk_body.append(
+                    {
+                        "index": {
+                            "_index": self.index_name,
+                            "_id": item.file_id,
+                        }
+                    }
+                )
+                # Add document content
+                doc = {
+                    "text": str(item.content),
+                    "metadata": item.metadata,
+                    "file_name": item.file_name,
+                }
+                if item.uri:
+                    doc["uri"] = item.uri
+                bulk_body.append(doc)
+
+        # Execute bulk request
+        response = self._opensearch_client.bulk(body=bulk_body)
+
+        # Check for errors
+        if response.get("errors"):
+            error_items = [
+                item
+                for item in response["items"]
+                if "error" in item.get("index", {})
+            ]
+            logger.warning(
+                f"Bulk upsert had {len(error_items)} errors: {error_items}"
+            )