qtype 0.1.1__py3-none-any.whl → 0.1.2__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/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/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)"

qtype/interpreter/executors/index_upsert_executor.py

@@ -3,9 +3,12 @@
 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
@@ -39,21 +42,32 @@ class IndexUpsertExecutor(BatchedStepExecutor):
             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._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]:
@@ -68,58 +82,15 @@ class IndexUpsertExecutor(BatchedStepExecutor):
         logger.debug(
             f"Executing IndexUpsert step: {self.step.id} with batch size: {len(batch)}"
         )
+        if len(batch) == 0:
+            return
 
         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.debug(
-                    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:
+            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:
@@ -133,13 +104,27 @@ class IndexUpsertExecutor(BatchedStepExecutor):
                 yield message
 
     async def _upsert_to_vector_store(
-        self, items: list[RAGChunk | RAGDocument]
-    ) -> None:
+        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:
@@ -162,67 +147,86 @@ class IndexUpsertExecutor(BatchedStepExecutor):
 
         # 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, items: list[RAGChunk | RAGDocument]
-    ) -> None:
+        self, batch: list[FlowMessage]
+    ) -> AsyncIterator[FlowMessage]:
         """Upsert items to document index using bulk API.
 
         Args:
-            items: List of RAGChunk or RAGDocument objects
+            batch: List of FlowMessages containing documents to upsert
         """
-        # 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}"
+        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."
+        )