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

qtype/interpreter/executors/llm_inference_executor.py

@@ -0,0 +1,272 @@
+from typing import AsyncIterator
+
+from llama_cloud import MessageRole as LlamaMessageRole
+from llama_index.core.base.llms.types import ChatResponse, CompletionResponse
+from openinference.semconv.trace import OpenInferenceSpanKindValues
+
+from qtype.base.types import PrimitiveTypeEnum
+from qtype.dsl.domain_types import ChatContent, ChatMessage, MessageRole
+from qtype.interpreter.base.base_step_executor import StepExecutor
+from qtype.interpreter.base.executor_context import ExecutorContext
+from qtype.interpreter.conversions import (
+    from_chat_message,
+    to_chat_message,
+    to_llm,
+    to_memory,
+    variable_to_chat_message,
+)
+from qtype.interpreter.types import FlowMessage
+from qtype.semantic.model import LLMInference
+
+
+class LLMInferenceExecutor(StepExecutor):
+    """Executor for LLMInference steps."""
+
+    # LLM inference spans should be marked as LLM type
+    span_kind = OpenInferenceSpanKindValues.LLM
+
+    def __init__(
+        self, step: LLMInference, context: ExecutorContext, **dependencies
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, LLMInference):
+            raise ValueError(
+                "LLMInferenceExecutor can only execute LLMInference steps."
+            )
+        self.step: LLMInference = step
+
+    def __extract_stream_reasoning_(self, response):
+        raw = response.raw
+        content_block_delta = raw.get("contentBlockDelta")
+        block_index = (
+            content_block_delta.get("contentBlockIndex")
+            if isinstance(content_block_delta, dict)
+            else None
+        )
+
+        reasoning_text = None
+        if block_index == 0:
+            reasoning_text = (
+                content_block_delta.get("delta", {})
+                .get("reasoningContent", {})
+                .get("text")
+            )
+
+        return reasoning_text
+
+    async def process_message(
+        self,
+        message: FlowMessage,
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a single FlowMessage for the LLMInference step.
+
+        Args:
+            message: The FlowMessage to process.
+
+        Yields:
+            FlowMessage with the results of LLM inference.
+        """
+        # Get output variable info
+        output_variable_id = self.step.outputs[0].id
+        output_variable_type = self.step.outputs[0].type
+
+        try:
+            # Determine if this is a chat or completion inference
+            if output_variable_type == ChatMessage:
+                result_message = await self._process_chat(
+                    message, output_variable_id
+                )
+            else:
+                result_message = await self._process_completion(
+                    message, output_variable_id
+                )
+
+            yield result_message
+
+        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
+
+    async def _process_chat(
+        self,
+        message: FlowMessage,
+        output_variable_id: str,
+    ) -> FlowMessage:
+        """Process a chat inference request.
+
+        Args:
+            message: The FlowMessage to process.
+            output_variable_id: The ID of the output variable.
+
+        Returns:
+            FlowMessage with the chat response.
+        """
+        model = to_llm(
+            self.step.model, self.step.system_message, self._secret_manager
+        )
+
+        # Convert input variables to chat messages
+        inputs = []
+        for input_var in self.step.inputs:
+            value = message.variables.get(input_var.id)
+            # Convert any value type to ChatMessage, then to LlamaChatMessage
+            chat_msg = variable_to_chat_message(value, input_var)
+            inputs.append(to_chat_message(chat_msg))
+
+        # Get session ID for memory isolation
+        session_id = message.session.session_id
+
+        # If memory is defined, use it
+        if self.step.memory:
+            memory = to_memory(session_id, self.step.memory)
+
+            # Add the inputs to the memory
+            from llama_index.core.async_utils import asyncio_run
+
+            asyncio_run(memory.aput_messages(inputs))
+            # Use the whole memory state as inputs to the llm
+            inputs = memory.get_all()
+        else:
+            # If memory is not defined, use conversation history from session
+            conversation_history = (
+                message.session.conversation_history
+                if hasattr(message.session, "conversation_history")
+                else []
+            )
+            if conversation_history:
+                inputs = [
+                    to_chat_message(msg) for msg in conversation_history
+                ] + inputs
+
+        # Add system message if needed
+        if (
+            self.step.system_message
+            and inputs
+            and inputs[0].role != LlamaMessageRole.SYSTEM
+        ):
+            system_message = ChatMessage(
+                role=MessageRole.system,
+                blocks=[
+                    ChatContent(
+                        type=PrimitiveTypeEnum.text,
+                        content=self.step.system_message,
+                    )
+                ],
+            )
+            inputs = [to_chat_message(system_message)] + inputs
+
+        chat_result: ChatResponse
+        if self.context.on_stream_event:
+            # Generate a unique stream ID for this inference
+            stream_id = f"llm-{self.step.id}-{id(message)}"
+            async with self.stream_emitter.reasoning_stream(
+                f"llm-{self.step.id}-{id(message)}-reasoning"
+            ) as reasoning:
+                generator = await model.astream_chat(
+                    messages=inputs,
+                    **(
+                        self.step.model.inference_params
+                        if self.step.model.inference_params
+                        else {}
+                    ),
+                )
+                async for complete_response in generator:
+                    reasoning_text = self.__extract_stream_reasoning_(
+                        complete_response
+                    )
+                    if reasoning_text:
+                        await reasoning.delta(reasoning_text)
+
+            async with self.stream_emitter.text_stream(stream_id) as streamer:
+                generator = await model.astream_chat(
+                    messages=inputs,
+                    **(
+                        self.step.model.inference_params
+                        if self.step.model.inference_params
+                        else {}
+                    ),
+                )
+                async for chat_response in generator:
+                    chat_text = chat_response.delta
+                    if chat_text.strip() != "":
+                        await streamer.delta(chat_response.delta)
+            # Get the final result
+            chat_result = chat_response
+        else:
+            chat_result = model.chat(
+                messages=inputs,
+                **(
+                    self.step.model.inference_params
+                    if self.step.model.inference_params
+                    else {}
+                ),
+            )
+
+        # Store result in memory if configured
+        if self.step.memory:
+            memory.put(chat_result.message)
+
+        # Convert result and return
+        result_value = from_chat_message(chat_result.message)
+        return message.copy_with_variables({output_variable_id: result_value})
+
+    async def _process_completion(
+        self,
+        message: FlowMessage,
+        output_variable_id: str,
+    ) -> FlowMessage:
+        """Process a completion inference request.
+
+        Args:
+            message: The FlowMessage to process.
+            output_variable_id: The ID of the output variable.
+
+        Returns:
+            FlowMessage with the completion response.
+        """
+        model = to_llm(
+            self.step.model, self.step.system_message, self._secret_manager
+        )
+
+        # Get input value
+        input_value = message.variables.get(self.step.inputs[0].id)
+        if not isinstance(input_value, str):
+            input_value = str(input_value)
+
+        # Perform inference with streaming if callback provided
+        complete_result: CompletionResponse
+        if self.context.on_stream_event:
+            # Generate a unique stream ID for this inference
+            stream_id = f"llm-{self.step.id}-{id(message)}"
+
+            async with self.stream_emitter.text_stream(stream_id) as streamer:
+                generator = await model.astream_complete(
+                    prompt=input_value,
+                    **(
+                        self.step.model.inference_params
+                        if self.step.model.inference_params
+                        else {}
+                    ),
+                )
+
+                async for complete_response in generator:
+                    text = complete_response.delta
+                    if complete_response.text.strip() != "":
+                        await streamer.delta(text)
+
+            complete_result = complete_response
+        else:
+            complete_result = model.complete(
+                prompt=input_value,
+                **(
+                    self.step.model.inference_params
+                    if self.step.model.inference_params
+                    else {}
+                ),
+            )
+
+        response: dict[str, str] = {output_variable_id: complete_result.text}
+
+        return message.copy_with_variables(response)

qtype/interpreter/executors/prompt_template_executor.py

@@ -0,0 +1,78 @@
+import string
+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 PromptTemplate
+
+
+def get_format_arguments(format_string: str) -> set[str]:
+    formatter = string.Formatter()
+    arguments = []
+    for literal_text, field_name, format_spec, conversion in formatter.parse(
+        format_string
+    ):
+        if field_name:
+            arguments.append(field_name)
+    return set(arguments)
+
+
+class PromptTemplateExecutor(StepExecutor):
+    """Executor for PromptTemplate steps."""
+
+    def __init__(
+        self, step: PromptTemplate, context: ExecutorContext, **dependencies
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, PromptTemplate):
+            raise ValueError(
+                (
+                    "PromptTemplateExecutor can only execute "
+                    "PromptTemplate steps."
+                )
+            )
+        self.step: PromptTemplate = step
+
+    async def process_message(
+        self,
+        message: FlowMessage,
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a single FlowMessage for the PromptTemplate step.
+
+        Args:
+            message: The FlowMessage to process.
+        Yields:
+            FlowMessage with the results of processing.
+        """
+        format_args = get_format_arguments(self.step.template)
+        try:
+            # Read input values from FlowMessage.variables
+            input_map = {}
+            for var in self.step.inputs:
+                if var.id in format_args:
+                    value = message.variables.get(var.id)
+                    if value is not None:
+                        input_map[var.id] = value
+
+            missing = format_args - input_map.keys()
+            if missing:
+                raise ValueError(
+                    (
+                        "The following fields are in the prompt template "
+                        f"but not in the inputs: {missing}"
+                    )
+                )
+            result = self.step.template.format(**input_map)
+            output_var_id = self.step.outputs[0].id
+
+            await self.stream_emitter.status(
+                (f"Processed message with PromptTemplate step {self.step.id}"),
+            )
+            yield message.copy_with_variables({output_var_id: 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/sql_source_executor.py

@@ -0,0 +1,106 @@
+from typing import AsyncIterator
+
+import boto3  # type: ignore[import-untyped]
+import pandas as pd
+import sqlalchemy
+from sqlalchemy import create_engine
+from sqlalchemy.exc import SQLAlchemyError
+
+from qtype.interpreter.auth.generic import auth
+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 SQLSource
+
+
+class SQLSourceExecutor(StepExecutor):
+    """Executor for SQLSource steps."""
+
+    def __init__(
+        self, step: SQLSource, context: ExecutorContext, **dependencies
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, SQLSource):
+            raise ValueError(
+                "SQLSourceExecutor can only execute SQLSource steps."
+            )
+        self.step: SQLSource = step
+
+    async def process_message(
+        self,
+        message: FlowMessage,
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a single FlowMessage for the SQLSource step.
+
+        Args:
+            message: The FlowMessage to process.
+        Yields:
+            FlowMessages with the results of SQL query execution.
+        """
+        # Create a database engine - resolve connection string if it's a SecretReference
+        connection_string = self._resolve_secret(self.step.connection)
+        connect_args = {}
+        if self.step.auth:
+            with auth(self.step.auth) as creds:
+                if isinstance(creds, boto3.Session):
+                    connect_args["session"] = creds
+        engine = create_engine(connection_string, connect_args=connect_args)
+
+        output_columns = {output.id for output in self.step.outputs}
+        step_inputs = {i.id for i in self.step.inputs}
+
+        try:
+            # Make a dictionary of column_name: value from message variables
+            params = {
+                col: message.variables.get(col)
+                for col in step_inputs
+                if col in message.variables
+            }
+
+            await self.stream_emitter.status(
+                f"Executing SQL query with params: {params}",
+            )
+
+            # Execute the query and fetch the results into a DataFrame
+            with engine.connect() as connection:
+                result = connection.execute(
+                    sqlalchemy.text(self.step.query),
+                    parameters=params if params else None,
+                )
+                df = pd.DataFrame(
+                    result.fetchall(), columns=list(result.keys())
+                )
+
+            # Confirm the outputs exist in the dataframe
+            columns = set(df.columns)
+            missing_columns = output_columns - columns
+            if missing_columns:
+                raise ValueError(
+                    (
+                        f"SQL Result was missing expected columns: "
+                        f"{', '.join(missing_columns)}, it has columns: "
+                        f"{', '.join(columns)}"
+                    )
+                )
+
+            # Emit one message per result row
+            for _, row in df.iterrows():
+                # Create a dict with only the output columns
+                row_dict = {
+                    str(k): v
+                    for k, v in row.to_dict().items()
+                    if str(k) in output_columns
+                }
+                # Merge with original message variables
+                yield message.copy_with_variables(new_variables=row_dict)
+
+            await self.stream_emitter.status(
+                f"Emitted {len(df)} rows from SQL query"
+            )
+
+        except SQLAlchemyError as e:
+            # Emit error event to stream so frontend can display it
+            await self.stream_emitter.error(str(e))
+            # Set error on the message and yield it
+            message.set_error(self.step.id, e)
+            yield message

qtype/interpreter/executors/vector_search_executor.py

@@ -0,0 +1,91 @@
+"""Vector search executor for retrieving relevant chunks from vector stores."""
+
+from __future__ import annotations
+
+import logging
+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_node_with_score,
+    to_llama_vector_store_and_retriever,
+)
+from qtype.interpreter.types import FlowMessage
+from qtype.semantic.model import VectorIndex, VectorSearch
+
+logger = logging.getLogger(__name__)
+
+
+class VectorSearchExecutor(StepExecutor):
+    """Executor for VectorSearch steps using LlamaIndex vector stores."""
+
+    def __init__(
+        self, step: VectorSearch, context: ExecutorContext, **dependencies
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, VectorSearch):
+            raise ValueError(
+                "VectorSearchExecutor can only execute VectorSearch steps."
+            )
+        self.step: VectorSearch = step
+
+        if not isinstance(self.step.index, VectorIndex):
+            raise ValueError(
+                f"VectorSearch step {self.step.id} must reference a VectorIndex."
+            )
+        self.index: VectorIndex = self.step.index
+
+        # Get the vector store and retriever
+        self._vector_store, self._retriever = (
+            to_llama_vector_store_and_retriever(
+                self.step.index, self.context.secret_manager
+            )
+        )
+
+    async def process_message(
+        self,
+        message: FlowMessage,
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a single FlowMessage for the VectorSearch step.
+
+        Args:
+            message: The FlowMessage to process.
+
+        Yields:
+            FlowMessage with search results.
+        """
+        try:
+            # Get the query from the input variable
+            # (validated to be exactly one text input)
+            input_var = self.step.inputs[0]
+            query = message.variables.get(input_var.id)
+
+            if not isinstance(query, str):
+                raise ValueError(
+                    f"VectorSearch input must be text, got {type(query)}"
+                )
+
+            # Perform the vector search
+            logger.debug(f"Performing vector search with query: {query}")
+            nodes_with_scores = await self._retriever.aretrieve(query)
+
+            # Convert results to RAGSearchResult objects
+            search_results = [
+                from_node_with_score(node_with_score)
+                for node_with_score in nodes_with_scores
+            ]
+
+            # Set the output variable (validated to be exactly one output
+            # of type list[RAGSearchResult])
+            output_var = self.step.outputs[0]
+            output_vars = {output_var.id: search_results}
+
+            yield message.copy_with_variables(output_vars)
+
+        except Exception as e:
+            logger.error(f"Vector search 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