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

qtype/interpreter/executors/bedrock_reranker_executor.py

@@ -0,0 +1,195 @@
+"""BedrockReranker executor for reordering search results by relevance."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import AsyncIterator
+
+from pydantic import BaseModel
+
+from qtype.base.types import PrimitiveTypeEnum
+from qtype.dsl.domain_types import RAGChunk, SearchResult
+from qtype.interpreter.auth.aws import aws
+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 BedrockReranker, ListType
+
+logger = logging.getLogger(__name__)
+
+
+class BedrockRerankerExecutor(StepExecutor):
+    """Executor for BedrockReranker steps that reorder search results by relevance."""
+
+    def __init__(
+        self, step: BedrockReranker, context: ExecutorContext, **dependencies
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, BedrockReranker):
+            raise ValueError(
+                "BedrockRerankerExecutor can only execute BedrockReranker steps."
+            )
+        self.step: BedrockReranker = step
+
+    async def process_message(
+        self,
+        message: FlowMessage,
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a single FlowMessage for the BedrockReranker step.
+
+        Args:
+            message: The FlowMessage to process.
+
+        Yields:
+            FlowMessage with reranked results.
+        """
+        try:
+            # Get the inputs
+            query = self._query(message)
+            docs = self._docs(message)
+
+            if len(docs) == 0:
+                # No documents to rerank, yield original message
+                yield message.copy_with_variables(
+                    {self.step.outputs[0].id: docs}
+                )
+                return
+
+            # Get session for region info
+            if self.step.auth is not None:
+                with aws(self.step.auth, self.context.secret_manager) as s:
+                    region_name = s.region_name
+            else:
+                import boto3
+
+                region_name = boto3.Session().region_name
+
+            # Convert the types
+            queries = [
+                {
+                    "type": "TEXT",
+                    "textQuery": {"text": query},
+                }
+            ]
+            documents = []
+
+            for doc in docs:
+                if isinstance(doc.content, RAGChunk):
+                    documents.append(
+                        {
+                            "type": "INLINE",
+                            "inlineDocumentSource": {
+                                "type": "TEXT",
+                                "textDocument": {"text": str(doc.content)},
+                            },
+                        }
+                    )
+                elif isinstance(doc.content, dict):
+                    documents.append(
+                        {
+                            "type": "INLINE",
+                            "inlineDocumentSource": {
+                                "type": "JSON",
+                                "jsonDocument": doc.content,
+                            },
+                        }
+                    )
+                elif isinstance(doc.content, BaseModel):
+                    documents.append(
+                        {
+                            "type": "INLINE",
+                            "inlineDocumentSource": {
+                                "type": "JSON",
+                                "jsonDocument": doc.content.model_dump(),
+                            },
+                        }
+                    )
+                else:
+                    raise ValueError(
+                        f"Unsupported document content type for BedrockReranker: {type(doc.content)}"
+                    )
+
+            reranking_configuration = {
+                "type": "BEDROCK_RERANKING_MODEL",
+                "bedrockRerankingConfiguration": {
+                    "numberOfResults": self.step.num_results or len(docs),
+                    "modelConfiguration": {
+                        "modelArn": f"arn:aws:bedrock:{region_name}::foundation-model/{self.step.model_id}"
+                    },
+                },
+            }
+
+            def _call_bedrock_rerank():
+                """Create client and call rerank in executor thread."""
+                if self.step.auth is not None:
+                    with aws(self.step.auth, self.context.secret_manager) as s:
+                        client = s.client("bedrock-agent-runtime")
+                        return client.rerank(
+                            queries=queries,
+                            sources=documents,
+                            rerankingConfiguration=reranking_configuration,
+                        )
+                else:
+                    import boto3
+
+                    session = boto3.Session()
+                    client = session.client("bedrock-agent-runtime")
+                    return client.rerank(
+                        queries=queries,
+                        sources=documents,
+                        rerankingConfiguration=reranking_configuration,
+                    )
+
+            loop = asyncio.get_running_loop()
+            response = await loop.run_in_executor(
+                self.context.thread_pool, _call_bedrock_rerank
+            )
+
+            results = []
+            for d in response["results"]:
+                doc = docs[d["index"]]
+                new_score = d["relevanceScore"]
+                results.append(doc.copy(update={"score": new_score}))
+
+            # Update the message with reranked results
+            yield message.copy_with_variables(
+                {self.step.outputs[0].id: results}
+            )
+        except Exception as e:
+            logger.error(f"Reranking failed: {e}", exc_info=True)
+            # 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 _query(self, message: FlowMessage) -> str:
+        """Extract the query string from the FlowMessage.
+
+        Args:
+            message: The FlowMessage containing the query variable.
+        Returns:
+            The query string.
+        """
+        for i in self.step.inputs:
+            if i.type == PrimitiveTypeEnum.text:
+                return message.variables[i.id]
+        raise ValueError(
+            f"No text input found for BedrockReranker step {self.step.id}"
+        )
+
+    def _docs(self, message: FlowMessage) -> list[SearchResult]:
+        """Extract the list of SearchResult documents from the FlowMessage.
+
+        Args:
+            message: The FlowMessage containing the document variable.
+        Returns:
+            The list of SearchResult documents.
+        """
+        for i in self.step.inputs:
+            if i.type == ListType(element_type="SearchResult"):
+                docs = message.variables[i.id]
+                return docs
+        raise ValueError(
+            f"No list of SearchResults input found for BedrockReranker step {self.step.id}"
+        )

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,123 @@
+import asyncio
+import logging
+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, context.secret_manager
+        )
+
+    # 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.
+        """
+
+        # TODO: switch back to async once aws auth supports it.
+        # https://github.com/bazaarvoice/qtype/issues/108
+        def _call():
+            return self.embedding_model.get_text_embedding(text=text)
+
+        loop = asyncio.get_running_loop()
+        response = await loop.run_in_executor(self.context.thread_pool, _call)
+
+        return response
+        # 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))
+            logging.error(
+                f"Error processing DocumentEmbedder step {self.step.id}",
+                exc_info=e,
+            )
+            message.set_error(self.step.id, e)
+            yield message