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

qtype/interpreter/conversions.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import importlib
+import uuid
 from typing import Any
 
 from llama_index.core.base.embeddings.base import BaseEmbedding
@@ -17,10 +18,15 @@ from llama_index.core.base.llms.types import (
 from llama_index.core.memory import Memory as LlamaMemory
 from llama_index.core.schema import Document as LlamaDocument
 from llama_index.core.vector_stores.types import BasePydanticVectorStore
-from opensearchpy import AWSV4SignerAuth, OpenSearch
+from opensearchpy import AsyncOpenSearch, AWSV4SignerAuth
 
 from qtype.base.types import PrimitiveTypeEnum
-from qtype.dsl.domain_types import ChatContent, ChatMessage, RAGDocument
+from qtype.dsl.domain_types import (
+    ChatContent,
+    ChatMessage,
+    RAGDocument,
+    RAGSearchResult,
+)
 from qtype.dsl.model import Memory
 from qtype.interpreter.auth.aws import aws
 from qtype.interpreter.auth.generic import auth
@@ -305,7 +311,8 @@ def to_embedding_model(model: Model) -> BaseEmbedding:
         )
 
         bedrock_embedding: BaseEmbedding = BedrockEmbedding(
-            model_name=model.model_id if model.model_id else model.id
+            model_name=model.model_id if model.model_id else model.id,
+            max_retries=100,
         )
         return bedrock_embedding
     elif model.provider == "openai":
@@ -326,7 +333,7 @@ def to_embedding_model(model: Model) -> BaseEmbedding:
 @cached_resource
 def to_opensearch_client(
     index: DocumentIndex, secret_manager: SecretManagerBase
-) -> OpenSearch:
+) -> AsyncOpenSearch:
     """
     Convert a DocumentIndex to an OpenSearch/Elasticsearch client.
 
@@ -375,7 +382,7 @@ def to_opensearch_client(
                 f"Unsupported authentication type for DocumentIndex: {type(index.auth)}"
             )
 
-    return OpenSearch(**client_kwargs)
+    return AsyncOpenSearch(**client_kwargs)
 
 
 def to_content_block(content: ChatContent) -> ContentBlock:
@@ -506,26 +513,30 @@ def to_text_splitter(splitter: DocumentSplitter) -> Any:
     Raises:
         InterpreterError: If the splitter class cannot be found or instantiated.
     """
-    from llama_index.core.node_parser import SentenceSplitter
 
-    # Map common splitter names to their classes
-    splitter_classes = {
-        "SentenceSplitter": SentenceSplitter,
-    }
+    module_path = "llama_index.core.node_parser"
+    class_name = splitter.splitter_name
+    try:
+        reader_module = importlib.import_module(module_path)
+        splitter_class = getattr(reader_module, class_name)
+    except (ImportError, AttributeError) as e:
+        raise ImportError(
+            f"Failed to import reader class '{class_name}' from '{module_path}': {e}"
+        ) from e
+    from llama_index.core.schema import BaseNode
 
-    # Get the splitter class
-    splitter_class = splitter_classes.get(splitter.splitter_name)
+    # TODO: let the user specify a custom ID namespace
+    namespace = uuid.UUID("12345678-1234-5678-1234-567812345678")
 
-    if splitter_class is None:
-        raise InterpreterError(
-            f"Unsupported text splitter: {splitter.splitter_name}. "
-            f"Supported splitters: {', '.join(splitter_classes.keys())}"
-        )
+    def id_func(i: int, doc: BaseNode) -> str:
+        u = uuid.uuid5(namespace, f"{doc.node_id}_{i}")
+        return str(u)
 
     # Prepare arguments for the splitter
     splitter_args = {
         "chunk_size": splitter.chunk_size,
         "chunk_overlap": splitter.chunk_overlap,
+        "id_func": id_func,
         **splitter.args,
     }
 
@@ -569,7 +580,7 @@ def to_llama_vector_store_and_retriever(
     return vector_store, retriever
 
 
-def from_node_with_score(node_with_score) -> Any:
+def from_node_with_score(node_with_score) -> RAGSearchResult:
     """Convert a LlamaIndex NodeWithScore to a RAGSearchResult.
 
     Args:
@@ -597,4 +608,8 @@ def from_node_with_score(node_with_score) -> Any:
     )
 
     # Wrap in RAGSearchResult with score
-    return RAGSearchResult(chunk=chunk, score=node_with_score.score or 0.0)
+    return RAGSearchResult(
+        content=chunk,
+        doc_id=chunk.document_id,
+        score=node_with_score.score or 0.0,
+    )

qtype/interpreter/converters.py

@@ -2,17 +2,21 @@
 
 from __future__ import annotations
 
+from collections.abc import AsyncIterator
+from typing import Any, cast
+
 import pandas as pd
+from pydantic import BaseModel
 
 from qtype.interpreter.types import FlowMessage, Session
 from qtype.semantic.model import Flow
 
 
-def dataframe_to_flow_messages(
+async def dataframe_to_flow_messages(
     df: pd.DataFrame, session: Session
-) -> list[FlowMessage]:
+) -> AsyncIterator[FlowMessage]:
     """
-    Convert a DataFrame to a list of FlowMessages.
+    Convert a DataFrame to an async generator of FlowMessages.
 
     Each row in the DataFrame becomes a FlowMessage with the same session.
 
@@ -20,14 +24,15 @@ def dataframe_to_flow_messages(
         df: DataFrame where each row represents one set of inputs
         session: Session object to use for all messages
 
-    Returns:
-        List of FlowMessages, one per DataFrame row
+    Yields:
+        FlowMessages, one per DataFrame row
     """
-    messages = []
-    for _, row in df.iterrows():
-        variables = row.to_dict()
-        messages.append(FlowMessage(session=session, variables=variables))
-    return messages
+    # Use to_dict with orient='records' - much faster than iterrows
+    # This returns a list of dicts directly without Series overhead
+    records = cast(list[dict[str, Any]], df.to_dict(orient="records"))
+
+    for record in records:
+        yield FlowMessage(session=session, variables=record)
 
 
 def flow_messages_to_dataframe(
@@ -45,8 +50,6 @@ def flow_messages_to_dataframe(
     Returns:
         DataFrame with one row per message, columns for each output variable
     """
-    from typing import Any
-
     results = []
     for idx, message in enumerate(messages):
         row_data: dict[str, Any] = {"row": idx}
@@ -54,7 +57,10 @@ def flow_messages_to_dataframe(
         # Extract output variables
         for var in flow.outputs:
             if var.id in message.variables:
-                row_data[var.id] = message.variables[var.id]
+                value = message.variables[var.id]
+                if isinstance(value, BaseModel):
+                    value = value.model_dump()
+                row_data[var.id] = value
             else:
                 row_data[var.id] = None
 

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/document_embedder_executor.py

@@ -1,5 +1,14 @@
 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
@@ -8,6 +17,13 @@ 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."""
 
@@ -24,7 +40,25 @@ class DocumentEmbedderExecutor(StepExecutor):
             )
         self.step: DocumentEmbedder = step
         # Initialize the embedding model once for the executor
-        self.embedding_model = to_embedding_model(self.step.model)
+        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,
@@ -52,9 +86,7 @@ class DocumentEmbedderExecutor(StepExecutor):
                 )
 
             # Generate embedding for the chunk content
-            vector = self.embedding_model.get_text_embedding(
-                text=str(chunk.content)
-            )
+            vector = await self._embed(str(chunk.content))
 
             # Create the output chunk with the vector
             embedded_chunk = RAGChunk(

qtype/interpreter/executors/document_search_executor.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 from typing import AsyncIterator
 
-from qtype.dsl.domain_types import RAGChunk, RAGSearchResult
+from qtype.dsl.domain_types import SearchResult
 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
@@ -29,6 +31,17 @@ class DocumentSearchExecutor(StepExecutor):
         )
         self.index_name = self.step.index.name
 
+    async def finalize(self) -> AsyncIterator[FlowMessage]:
+        """Clean up resources after all messages are processed."""
+        if hasattr(self, "client") and self.client:
+            try:
+                await self.client.close()
+            except Exception:
+                pass
+        # Make this an async generator
+        return
+        yield  # type: ignore[unreachable]
+
     async def process_message(
         self,
         message: FlowMessage,
@@ -39,7 +52,7 @@ class DocumentSearchExecutor(StepExecutor):
             message: The FlowMessage to process.
 
         Yields:
-            FlowMessage with search results as RAGSearchResult instances.
+            A list of dictionaries with _source, _search_score, and _search_id fields.
         """
         input_id = self.step.inputs[0].id
         output_id = self.step.outputs[0].id
@@ -58,62 +71,40 @@ class DocumentSearchExecutor(StepExecutor):
             # Build the search query
             search_body = {
                 "query": {
-                    "multi_match": {
-                        "query": query_text,
-                        "fields": ["content^2", "title", "*"],
-                        "type": "best_fields",
-                    }
+                    "multi_match": {"query": query_text} | self.step.query_args
                 },
-                "size": 10,  # Default top 10 results
+                "size": self.step.default_top_k,
             }
 
             # 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()
-                            ],
-                        }
+                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(
+            # Execute the search asynchronously using AsyncOpenSearch
+            response = await self.client.search(
                 index=self.index_name, body=search_body
             )
 
-            # Process each hit and yield as RAGSearchResult
+            # Process each hit and yield as SearchResult
+            # TODO: add support for decomposing a RAGSearchResult for hybrid search
+            search_results = []
             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,
+                search_results.append(
+                    SearchResult(
+                        content=hit["_source"],
+                        doc_id=hit["_id"],
+                        score=hit["_score"],
+                    )
                 )
-
-                # 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})
+            yield message.copy_with_variables({output_id: search_results})
 
         except Exception as e:
             # Emit error event to stream so frontend can display it

qtype/interpreter/executors/document_splitter_executor.py

@@ -72,7 +72,7 @@ class DocumentSplitterExecutor(StepExecutor):
             llama_doc = LlamaDocument(
                 text=content_text,
                 metadata=document.metadata or {},
-                id_=document.file_id,
+                doc_id=document.file_id,
             )
 
             # Split the document using the LlamaIndex splitter

qtype/interpreter/executors/field_extractor_executor.py

@@ -132,12 +132,17 @@ class FieldExtractorExecutor(StepExecutor):
             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"
+                if self.step.fail_on_missing:
+                    raise ValueError(
+                        (
+                            f"JSONPath expression '{self.step.json_path}' "
+                            f"did not match any data in input"
+                        )
                     )
-                )
+                else:
+                    # Yield message with None output
+                    yield message.copy_with_variables({output_id: None})
+                    return
 
             await self.stream_emitter.status(
                 f"JSONPath matched {len(matches)} value(s)"