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

qtype/interpreter/executors/decoder_executor.py

@@ -0,0 +1,163 @@
+import json
+import xml.etree.ElementTree as ET
+from typing import Any, AsyncIterator
+
+from qtype.dsl.model import DecoderFormat
+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 Decoder
+
+
+class DecoderExecutor(StepExecutor):
+    """Executor for Decoder steps."""
+
+    def __init__(
+        self, step: Decoder, context: ExecutorContext, **dependencies
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, Decoder):
+            raise ValueError("DecoderExecutor can only execute Decoder steps.")
+        self.step: Decoder = step
+
+    def _parse_json(self, input_str: str) -> dict[str, Any]:
+        """Parse a JSON string into a Python object.
+
+        Args:
+            input_str: The JSON string to parse.
+
+        Returns:
+            A dictionary parsed from the JSON.
+
+        Raises:
+            ValueError: If the JSON is invalid or not an object.
+        """
+        try:
+            cleaned_response = input_str.strip()
+            # Remove markdown code fences if present
+            if cleaned_response.startswith("```json"):
+                cleaned_response = cleaned_response[7:]
+            if cleaned_response.endswith("```"):
+                cleaned_response = cleaned_response[:-3]
+            cleaned_response = cleaned_response.strip()
+
+            # Parse the JSON
+            parsed = json.loads(cleaned_response)
+            if not isinstance(parsed, dict):
+                raise ValueError(f"Parsed JSON is not an object: {parsed}")
+            return parsed
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON input: {e}") from e
+
+    def _parse_xml(self, input_str: str) -> dict[str, Any]:
+        """Parse an XML string into a Python object.
+
+        Args:
+            input_str: The XML string to parse.
+
+        Returns:
+            A dictionary with tag names as keys and text content as values.
+
+        Raises:
+            ValueError: If the XML is invalid.
+        """
+        try:
+            cleaned_response = input_str.strip()
+            # Remove markdown code fences if present
+            if cleaned_response.startswith("```xml"):
+                cleaned_response = cleaned_response[6:]
+            if cleaned_response.endswith("```"):
+                cleaned_response = cleaned_response[:-3]
+            cleaned_response = cleaned_response.strip()
+
+            # Escape ampersands
+            cleaned_response = cleaned_response.replace("&", "&")
+            tree = ET.fromstring(cleaned_response)
+            result = {c.tag: c.text for c in tree}
+
+            return result
+        except Exception as e:
+            raise ValueError(f"Invalid XML input: {e}") from e
+
+    def _parse(self, input_str: str) -> dict[str, Any]:
+        """Parse input string based on the decoder format.
+
+        Args:
+            input_str: The string to parse.
+
+        Returns:
+            A dictionary parsed from the input.
+
+        Raises:
+            ValueError: If the format is unsupported or parsing fails.
+        """
+        if self.step.format == DecoderFormat.json:
+            return self._parse_json(input_str)
+        elif self.step.format == DecoderFormat.xml:
+            return self._parse_xml(input_str)
+        else:
+            raise ValueError(
+                (
+                    f"Unsupported decoder format: {self.step.format}. "
+                    f"Supported formats are: {DecoderFormat.json}, "
+                    f"{DecoderFormat.xml}."
+                )
+            )
+
+    async def process_message(
+        self,
+        message: FlowMessage,
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a single FlowMessage for the Decoder step.
+
+        Args:
+            message: The FlowMessage to process.
+
+        Yields:
+            A FlowMessage with decoded outputs or an error.
+        """
+        input_id = self.step.inputs[0].id
+
+        try:
+            # Get the input string to decode
+            input_value = message.variables.get(input_id)
+            if not isinstance(input_value, str):
+                raise ValueError(
+                    (
+                        f"Input to decoder step {self.step.id} must be "
+                        f"a string, found {type(input_value).__name__}."
+                    )
+                )
+
+            await self.stream_emitter.status(
+                f"Decoding {self.step.format.value} input"
+            )
+
+            # Parse the input
+            result_dict = self._parse(input_value)
+
+            # Extract output variables from the parsed result
+            output_vars = {}
+            for output in self.step.outputs:
+                if output.id in result_dict:
+                    output_vars[output.id] = result_dict[output.id]
+                else:
+                    raise ValueError(
+                        (
+                            f"Output variable {output.id} not found in "
+                            f"decoded result: {result_dict}"
+                        )
+                    )
+
+            await self.stream_emitter.status(
+                f"Decoded {len(output_vars)} output variables"
+            )
+
+            # Yield the result
+            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/doc_to_text_executor.py

@@ -0,0 +1,112 @@
+from io import BytesIO
+from typing import AsyncIterator
+
+from docling.document_converter import DocumentConverter
+from docling_core.types.io import DocumentStream
+
+from qtype.base.types import PrimitiveTypeEnum
+from qtype.dsl.domain_types import RAGDocument
+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 DocToTextConverter
+
+
+class DocToTextConverterExecutor(StepExecutor):
+    """Executor for DocToTextConverter steps."""
+
+    def __init__(
+        self,
+        step: DocToTextConverter,
+        context: ExecutorContext,
+        **dependencies,
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, DocToTextConverter):
+            raise ValueError(
+                (
+                    "DocToTextConverterExecutor can only execute "
+                    "DocToTextConverter steps."
+                )
+            )
+        self.step: DocToTextConverter = step
+        # Initialize the Docling converter once for the executor
+        self.docling_converter = DocumentConverter()
+
+    async def process_message(
+        self,
+        message: FlowMessage,
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a single FlowMessage for the DocToTextConverter step.
+
+        Args:
+            message: The FlowMessage to process.
+        Yields:
+            FlowMessage with converted document.
+        """
+        input_id = self.step.inputs[0].id
+        output_id = self.step.outputs[0].id
+
+        try:
+            # Get the input document
+            if input_id not in message.variables:
+                raise ValueError(f"Input variable '{input_id}' is missing")
+            doc = message.variables.get(input_id)
+            if not isinstance(doc, RAGDocument):
+                raise ValueError(
+                    f"Input variable '{input_id}' must be a RAGDocument"
+                )
+
+            await self.stream_emitter.status(
+                f"Converting document: {doc.file_name}",
+            )
+
+            # Convert the document
+            converted_doc = self._convert_doc(doc)
+
+            await self.stream_emitter.status(
+                f"Converted {doc.file_name} to markdown text",
+            )
+
+            # Yield the result
+            yield message.copy_with_variables({output_id: converted_doc})
+
+        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
+
+    def _convert_doc(self, doc: RAGDocument) -> RAGDocument:
+        """Convert a RAGDocument to text/markdown format.
+
+        Args:
+            doc: The document to convert.
+
+        Returns:
+            A RAGDocument with markdown text content.
+        """
+        # If already text, no conversion needed
+        if doc.type == PrimitiveTypeEnum.text:
+            return doc
+
+        # Convert based on content type
+        if isinstance(doc.content, bytes):
+            # Use DocumentStream for bytes content
+            stream = DocumentStream(
+                name=doc.file_name, stream=BytesIO(doc.content)
+            )
+            document = self.docling_converter.convert(stream).document
+        else:
+            # Convert string content directly
+            document = self.docling_converter.convert(doc.content).document
+
+        # Export to markdown
+        markdown = document.export_to_markdown()
+
+        # Return new RAGDocument with markdown content
+        return RAGDocument(
+            **doc.model_dump(exclude={"content", "type"}),
+            content=markdown,
+            type=PrimitiveTypeEnum.text,
+        )

qtype/interpreter/executors/document_embedder_executor.py

@@ -0,0 +1,107 @@
+from typing import AsyncIterator
+
+from botocore.exceptions import ClientError
+from llama_index.core.base.embeddings.base import BaseEmbedding
+from tenacity import (
+    retry,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_exponential,
+)
+
+from qtype.dsl.domain_types import RAGChunk
+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 DocumentEmbedder
+
+
+def is_throttling_error(e):
+    return (
+        isinstance(e, ClientError)
+        and e.response["Error"]["Code"] == "ThrottlingException"
+    )
+
+
+class DocumentEmbedderExecutor(StepExecutor):
+    """Executor for DocumentEmbedder steps."""
+
+    def __init__(
+        self, step: DocumentEmbedder, context: ExecutorContext, **dependencies
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, DocumentEmbedder):
+            raise ValueError(
+                (
+                    "DocumentEmbedderExecutor can only execute "
+                    "DocumentEmbedder steps."
+                )
+            )
+        self.step: DocumentEmbedder = step
+        # Initialize the embedding model once for the executor
+        self.embedding_model: BaseEmbedding = to_embedding_model(
+            self.step.model
+        )
+
+    # TODO: properly abstract this into a mixin
+    @retry(
+        retry=retry_if_exception(is_throttling_error),
+        wait=wait_exponential(multiplier=0.5, min=1, max=30),
+        stop=stop_after_attempt(10),
+    )
+    async def _embed(self, text: str) -> list[float]:
+        """Generate embedding for the given text using the embedding model.
+
+        Args:
+            text: The text to embed.
+        Returns:
+            The embedding vector as a list of floats.
+        """
+        return await self.embedding_model.aget_text_embedding(text=text)
+
+    async def process_message(
+        self,
+        message: FlowMessage,
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a single FlowMessage for the DocumentEmbedder step.
+
+        Args:
+            message: The FlowMessage to process.
+        Yields:
+            FlowMessage with embedded chunk.
+        """
+        input_id = self.step.inputs[0].id
+        output_id = self.step.outputs[0].id
+
+        try:
+            # Get the input chunk
+            chunk = message.variables.get(input_id)
+            if not isinstance(chunk, RAGChunk):
+                raise ValueError(
+                    (
+                        f"Input variable '{input_id}' must be a RAGChunk, "
+                        f"got {type(chunk)}"
+                    )
+                )
+
+            # Generate embedding for the chunk content
+            vector = await self._embed(str(chunk.content))
+
+            # Create the output chunk with the vector
+            embedded_chunk = RAGChunk(
+                vector=vector,
+                content=chunk.content,
+                chunk_id=chunk.chunk_id,
+                document_id=chunk.document_id,
+                metadata=chunk.metadata,
+            )
+
+            # Yield the result
+            yield message.copy_with_variables({output_id: embedded_chunk})
+
+        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/document_search_executor.py

@@ -0,0 +1,122 @@
+from typing import AsyncIterator
+
+from qtype.dsl.domain_types import RAGChunk, RAGSearchResult
+from qtype.interpreter.base.base_step_executor import StepExecutor
+from qtype.interpreter.base.executor_context import ExecutorContext
+from qtype.interpreter.conversions import to_opensearch_client
+from qtype.interpreter.types import FlowMessage
+from qtype.semantic.model import DocumentSearch
+
+
+class DocumentSearchExecutor(StepExecutor):
+    """Executor for DocumentSearch steps using OpenSearch/Elasticsearch."""
+
+    def __init__(
+        self, step: DocumentSearch, context: ExecutorContext, **dependencies
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, DocumentSearch):
+            raise ValueError(
+                (
+                    "DocumentSearchExecutor can only execute "
+                    "DocumentSearch steps."
+                )
+            )
+        self.step: DocumentSearch = step
+        # Initialize the OpenSearch client once for the executor
+        self.client = to_opensearch_client(
+            self.step.index, self._secret_manager
+        )
+        self.index_name = self.step.index.name
+
+    async def process_message(
+        self,
+        message: FlowMessage,
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a single FlowMessage for the DocumentSearch step.
+
+        Args:
+            message: The FlowMessage to process.
+
+        Yields:
+            FlowMessage with search results as RAGSearchResult instances.
+        """
+        input_id = self.step.inputs[0].id
+        output_id = self.step.outputs[0].id
+
+        try:
+            # Get the search query text
+            query_text = message.variables.get(input_id)
+            if not isinstance(query_text, str):
+                raise ValueError(
+                    (
+                        f"Input variable '{input_id}' must be a string "
+                        f"(text query), got {type(query_text)}"
+                    )
+                )
+
+            # Build the search query
+            search_body = {
+                "query": {
+                    "multi_match": {
+                        "query": query_text,
+                        "fields": ["content^2", "title", "*"],
+                        "type": "best_fields",
+                    }
+                },
+                "size": 10,  # Default top 10 results
+            }
+
+            # Apply any filters if specified
+            if self.step.filters:
+                if "query" in search_body:
+                    search_body["query"] = {
+                        "bool": {
+                            "must": [search_body["query"]],
+                            "filter": [
+                                {"term": {k: v}}
+                                for k, v in self.step.filters.items()
+                            ],
+                        }
+                    }
+
+            # Execute the search
+            response = self.client.search(
+                index=self.index_name, body=search_body
+            )
+
+            # Process each hit and yield as RAGSearchResult
+            for hit in response["hits"]["hits"]:
+                source = hit["_source"]
+                doc_id = hit["_id"]
+                score = hit["_score"]
+
+                # Extract content (adjust field name based on your schema)
+                content = source.get("content", "")
+
+                # Build metadata from the source, excluding content field
+                metadata = {
+                    k: v for k, v in source.items() if k not in ["content"]
+                }
+
+                # Create a RAGChunk from the search result
+                # Use the document ID as both chunk_id and document_id
+                chunk = RAGChunk(
+                    content=content,
+                    chunk_id=doc_id,
+                    document_id=source.get("document_id", doc_id),
+                    vector=None,  # Document search doesn't return embeddings
+                    metadata=metadata,
+                )
+
+                # Wrap in RAGSearchResult with the score
+                search_result = RAGSearchResult(chunk=chunk, score=score)
+
+                # Yield result for each document
+                yield message.copy_with_variables({output_id: search_result})
+
+        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/document_source_executor.py

@@ -0,0 +1,118 @@
+import importlib
+from typing import AsyncIterator
+
+from qtype.interpreter.base.base_step_executor import StepExecutor
+from qtype.interpreter.base.executor_context import ExecutorContext
+from qtype.interpreter.conversions import from_llama_document
+from qtype.interpreter.types import FlowMessage
+from qtype.semantic.model import DocumentSource
+
+
+class DocumentSourceExecutor(StepExecutor):
+    """Executor for DocumentSource steps."""
+
+    def __init__(
+        self, step: DocumentSource, context: ExecutorContext, **dependencies
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, DocumentSource):
+            raise ValueError(
+                (
+                    "DocumentSourceExecutor can only execute "
+                    "DocumentSource steps."
+                )
+            )
+        self.step: DocumentSource = step
+        # Initialize the reader class once for the executor
+        self.reader_class = self._load_reader_class()
+
+    def _load_reader_class(self) -> type:
+        """Load the LlamaIndex reader class dynamically.
+
+        Returns:
+            The reader class.
+
+        Raises:
+            ImportError: If the reader class cannot be imported.
+        """
+        # Parse the reader module path
+        # Format: 'file.SimpleDirectoryReader' -> llama_index.readers.file + SimpleDirectoryReader
+        # Special case: 'file.SimpleDirectoryReader' is actually in llama_index.core
+        parts = self.step.reader_module.split(".")
+        module_path = ".".join(parts[:-1])
+        class_name = parts[-1]
+
+        # Dynamically import the reader module and get the class
+        try:
+            module = importlib.import_module(module_path)
+            reader_class = getattr(module, class_name)
+            return reader_class
+        except (ImportError, AttributeError) as e:
+            raise ImportError(
+                (
+                    f"Failed to import reader class '{class_name}' "
+                    f"from '{module_path}': {e}"
+                )
+            ) from e
+
+    async def process_message(
+        self,
+        message: FlowMessage,
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a single FlowMessage for the DocumentSource step.
+
+        Args:
+            message: The FlowMessage to process.
+        Yields:
+            FlowMessages with loaded documents.
+        """
+        output_id = self.step.outputs[0].id
+
+        try:
+            # Resolve any SecretReferences in step args
+            context = f"step '{self.step.id}'"
+            resolved_args = self._secret_manager.resolve_secrets_in_dict(
+                self.step.args, context
+            )
+
+            # Combine resolved step args with message variables as runtime args
+            runtime_args = {
+                key: message.variables.get(key)
+                for key in message.variables.keys()
+            }
+            combined_args = {**resolved_args, **runtime_args}
+
+            # Instantiate the reader with combined arguments
+            loader = self.reader_class(**combined_args)
+
+            # Load documents using the loader
+            if not hasattr(loader, "load_data"):
+                raise AttributeError(
+                    (
+                        f"Reader class '{self.reader_class.__name__}' "
+                        "does not have a 'load_data' method"
+                    )
+                )
+            load_args = self.step.loader_args or {}
+
+            llama_documents = loader.load_data(**load_args)
+
+            # Convert LlamaIndex Documents to RAGDocuments
+            rag_documents = [
+                from_llama_document(doc) for doc in llama_documents
+            ]
+
+            # Emit feedback about total documents loaded
+            await self.stream_emitter.status(
+                f"Loaded {len(rag_documents)} documents"
+            )
+
+            # Yield one message per document (fan-out)
+            for doc in rag_documents:
+                yield message.copy_with_variables({output_id: doc})
+
+        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