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

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,232 @@
+"""Index upsert executor for inserting documents/chunks into indexes."""
+
+from __future__ import annotations
+
+import logging
+import uuid
+from typing import AsyncIterator
+
+from llama_index.core.schema import TextNode
+from opensearchpy import AsyncOpenSearch
+from pydantic import BaseModel
+
+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.index_type = "vector"
+        elif isinstance(self.step.index, DocumentIndex):
+            # Document index for text-based search
+            self._opensearch_client: AsyncOpenSearch = 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
+            self._document_index: DocumentIndex = self.step.index
+        else:
+            raise ValueError(
+                f"Unsupported index type: {type(self.step.index)}"
+            )
+
+    async def finalize(self) -> AsyncIterator[FlowMessage]:
+        """Clean up resources after all messages are processed."""
+        if hasattr(self, "_opensearch_client") and self._opensearch_client:
+            try:
+                await self._opensearch_client.close()
+            except Exception:
+                pass
+        # Make this an async generator
+        return
+        yield  # type: ignore[unreachable]
+
+    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.debug(
+            f"Executing IndexUpsert step: {self.step.id} with batch size: {len(batch)}"
+        )
+        if len(batch) == 0:
+            return
+
+        try:
+            if self.index_type == "vector":
+                result_iter = self._upsert_to_vector_store(batch)
+            else:
+                result_iter = self._upsert_to_document_index(batch)
+            async for message in result_iter:
+                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, batch: list[FlowMessage]
+    ) -> AsyncIterator[FlowMessage]:
+        """Upsert items to vector store.
+
+        Args:
+            items: List of RAGChunk or RAGDocument objects
+        """
+        # safe since semantic validation checks input length
+        input_var = self.step.inputs[0]
+
+        # Collect all RAGChunks or RAGDocuments from the batch inputs
+        items = []
+        for message in batch:
+            input_data = message.variables.get(input_var.id)
+            if not isinstance(input_data, (RAGChunk, RAGDocument)):
+                raise ValueError(
+                    f"IndexUpsert only supports RAGChunk or RAGDocument "
+                    f"inputs. Got: {type(input_data)}"
+                )
+            items.append(input_data)
+
+        # 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)
+        num_inserted = len(items)
+
+        # Emit status update
+        await self.stream_emitter.status(
+            f"Upserted {num_inserted} items to index {self.step.index.name}"
+        )
+        for message in batch:
+            yield message
+
+    async def _upsert_to_document_index(
+        self, batch: list[FlowMessage]
+    ) -> AsyncIterator[FlowMessage]:
+        """Upsert items to document index using bulk API.
+
+        Args:
+            batch: List of FlowMessages containing documents to upsert
+        """
+
+        bulk_body = []
+        message_by_id: dict[str, FlowMessage] = {}
+
+        for message in batch:
+            # Collect all input variables into a single document dict
+            doc_dict = {}
+            for input_var in self.step.inputs:
+                value = message.variables.get(input_var.id)
+
+                # Convert to dict if it's a Pydantic model
+                if isinstance(value, BaseModel):
+                    value = value.model_dump()
+
+                # Merge into document dict
+                if isinstance(value, dict):
+                    doc_dict.update(value)
+                else:
+                    # Primitive types - use variable name as field name
+                    doc_dict[input_var.id] = value
+
+            # Determine the document id field
+            id_field = None
+            if self._document_index.id_field is not None:
+                id_field = self._document_index.id_field
+                if id_field not in doc_dict:
+                    raise ValueError(
+                        f"Specified id_field '{id_field}' not found in inputs"
+                    )
+            else:
+                # Auto-detect with fallback
+                for field in ["_id", "id", "doc_id", "document_id"]:
+                    if field in doc_dict:
+                        id_field = field
+                        break
+            if id_field is not None:
+                doc_id = str(doc_dict[id_field])
+            else:
+                # Generate a UUID if no id field found
+                doc_id = str(uuid.uuid4())
+
+            # Add bulk action and document
+            bulk_body.append(
+                {"index": {"_index": self.index_name, "_id": doc_id}}
+            )
+            bulk_body.append(doc_dict)
+            message_by_id[doc_id] = message
+
+        # Execute bulk request asynchronously
+        response = await self._opensearch_client.bulk(body=bulk_body)
+
+        num_inserted = 0
+        for item in response["items"]:
+            doc_id = item["index"]["_id"]
+            message = message_by_id[doc_id]
+            if "error" in item.get("index", {}):
+                message.set_error(
+                    self.step.id,
+                    Exception(item["index"]["error"]),
+                )
+            else:
+                num_inserted += 1
+            yield message
+        await self.stream_emitter.status(
+            f"Upserted {num_inserted} items to index {self.step.index.name}, {len(batch) - num_inserted} errors occurred."
+        )

qtype/interpreter/executors/invoke_embedding_executor.py

@@ -0,0 +1,104 @@
+import asyncio
+from typing import AsyncIterator
+
+from openinference.semconv.trace import OpenInferenceSpanKindValues
+
+from qtype.base.types import PrimitiveTypeEnum
+from qtype.dsl.domain_types import Embedding
+from qtype.interpreter.base.base_step_executor import StepExecutor
+from qtype.interpreter.base.executor_context import ExecutorContext
+from qtype.interpreter.conversions import to_embedding_model
+from qtype.interpreter.types import FlowMessage
+from qtype.semantic.model import InvokeEmbedding
+
+
+class InvokeEmbeddingExecutor(StepExecutor):
+    """Executor for InvokeEmbedding steps."""
+
+    # Embedding operations should be marked as EMBEDDING type
+    span_kind = OpenInferenceSpanKindValues.EMBEDDING
+
+    def __init__(
+        self, step: InvokeEmbedding, context: ExecutorContext, **dependencies
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, InvokeEmbedding):
+            raise ValueError(
+                (
+                    "InvokeEmbeddingExecutor can only execute "
+                    "InvokeEmbedding steps."
+                )
+            )
+        self.step: InvokeEmbedding = step
+        # Initialize the embedding model once for the executor
+        self.embedding_model = to_embedding_model(
+            self.step.model, context.secret_manager
+        )
+
+    async def process_message(
+        self,
+        message: FlowMessage,
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a single FlowMessage for the InvokeEmbedding step.
+
+        Args:
+            message: The FlowMessage to process.
+        Yields:
+            FlowMessage with embedding.
+        """
+        input_id = self.step.inputs[0].id
+        input_type = self.step.inputs[0].type
+        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 missing")
+
+            def _call(input_value=input_value):
+                # Generate embedding based on input type
+                if input_type == PrimitiveTypeEnum.text:
+                    if not isinstance(input_value, str):
+                        input_value = str(input_value)
+                    vector = self.embedding_model.get_text_embedding(
+                        text=input_value
+                    )
+                    content = input_value
+                elif input_type == PrimitiveTypeEnum.image:
+                    # For image embeddings
+                    vector = self.embedding_model.get_image_embedding(
+                        image_path=input_value
+                    )
+                    content = input_value
+                else:
+                    raise ValueError(
+                        (
+                            f"Unsupported input type for embedding: "
+                            f"{input_type}. Must be 'text' or 'image'."
+                        )
+                    )
+
+                # Create the Embedding object
+                embedding = Embedding(
+                    vector=vector,
+                    content=content,
+                )
+                return embedding
+
+            # TODO: switch back to async once aws auth supports it.
+            # https://github.com/bazaarvoice/qtype/issues/108
+            loop = asyncio.get_running_loop()
+            embedding = await loop.run_in_executor(
+                self.context.thread_pool, _call
+            )
+
+            # Yield the result
+            yield message.copy_with_variables({output_id: embedding})
+
+        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/invoke_flow_executor.py

@@ -0,0 +1,51 @@
+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 InvokeFlow
+
+
+class InvokeFlowExecutor(StepExecutor):
+    """Executor for InvokeFlow steps."""
+
+    def __init__(
+        self, step: InvokeFlow, context: ExecutorContext, **dependencies
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, InvokeFlow):
+            raise ValueError(
+                ("InvokeFlowExecutor can only execute InvokeFlow steps.")
+            )
+        self.step: InvokeFlow = step
+
+    async def process_message(
+        self, message: FlowMessage
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a single FlowMessage for the InvokeFlow step.
+
+        Args:
+            message: The FlowMessage to process.
+        Yields:
+            FlowMessage with results from the invoked flow.
+        """
+        from qtype.interpreter.flow import run_flow
+
+        initial = message.copy_with_variables(
+            {
+                id: message.variables.get(var.id)
+                for var, id in self.step.input_bindings.items()
+            }
+        )
+        # Pass through context (already available as self.context)
+        result = await run_flow(
+            self.step.flow, [initial], context=self.context
+        )
+
+        for msg in result:
+            yield msg.copy_with_variables(
+                {
+                    var.id: msg.variables.get(id)
+                    for var, id in self.step.output_bindings.items()
+                }
+            )