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

qtype/interpreter/endpoints.py

@@ -0,0 +1,355 @@
+"""Unified API endpoint for flow execution with streaming and REST support."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, AsyncIterator
+
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import StreamingResponse
+
+from qtype.dsl.domain_types import ChatMessage, MessageRole
+from qtype.interpreter.base.executor_context import ExecutorContext
+from qtype.interpreter.flow import run_flow
+from qtype.interpreter.stream.chat import format_stream_events_as_sse
+from qtype.interpreter.stream.chat.ui_request_to_domain_type import (
+    completion_request_to_input_model,
+    ui_request_to_domain_type,
+)
+from qtype.interpreter.stream.chat.vercel import ChatRequest, CompletionRequest
+from qtype.interpreter.stream.utils import callback_to_async_iterator
+from qtype.interpreter.types import FlowMessage, StreamEvent
+from qtype.interpreter.typing import (
+    create_input_shape,
+    create_output_container_type,
+    create_output_shape,
+    flow_results_to_output_container,
+    request_to_flow_message,
+)
+from qtype.semantic.model import Flow
+
+logger = logging.getLogger(__name__)
+
+
+async def _execute_flow_with_streaming(
+    flow: Flow,
+    initial_message: FlowMessage,
+    context: ExecutorContext,
+) -> AsyncIterator[StreamEvent]:
+    """
+    Execute flow and yield StreamEvents as they occur.
+
+    This function converts run_flow's callback-based streaming to an
+    async iterator. Executors emit StreamEvents (including ErrorEvents)
+    through the on_stream_event callback.
+
+    Args:
+        flow: The flow to execute
+        initial_message: Initial FlowMessage with inputs
+
+    Yields:
+        StreamEvent instances from executors
+    """
+
+    async def execute_with_callback(callback):  # type: ignore[no-untyped-def]
+        """Execute flow with streaming callback."""
+        # Update context with streaming callback
+        stream_context = ExecutorContext(
+            secret_manager=context.secret_manager,
+            on_stream_event=callback,
+            on_progress=context.on_progress,
+            tracer=context.tracer,
+        )
+        await run_flow(
+            flow,
+            initial_message,
+            context=stream_context,
+        )
+
+    # Convert callback-based streaming to async iterator
+    async for event in callback_to_async_iterator(execute_with_callback):
+        yield event
+
+
+async def _stream_sse_response(
+    flow: Flow,
+    initial_message: FlowMessage,
+    context: ExecutorContext,
+    output_metadata: dict[str, Any] | None = None,
+) -> AsyncIterator[str]:
+    """
+    Execute flow and stream Server-Sent Events using Vercel AI SDK protocol.
+
+    This function orchestrates flow execution and formats the resulting
+    StreamEvents as SSE for the Vercel AI SDK frontend.
+
+    Args:
+        flow: The flow to execute
+        initial_message: Initial FlowMessage with inputs
+        output_metadata: Optional metadata to include in FinishChunk
+
+    Yields:
+        SSE formatted strings (data: {json}\\n\\n)
+    """
+    # Execute flow and get event stream
+    event_stream = _execute_flow_with_streaming(flow, initial_message, context)
+
+    # Format events as SSE with metadata
+    async for sse_line in format_stream_events_as_sse(
+        event_stream, output_metadata=output_metadata
+    ):
+        yield sse_line
+
+
+def _create_chat_streaming_endpoint(
+    app: FastAPI, flow: Flow, context: ExecutorContext
+) -> None:
+    """
+    Create streaming endpoint for Conversational flows.
+
+    Args:
+        app: FastAPI application instance
+        flow: Flow with Conversational interface
+        secret_manager: Optional secret manager for resolving secrets
+    """
+    flow_id = flow.id
+
+    # Type narrowing for mypy
+    if flow.interface is None:
+        raise ValueError(f"Flow {flow_id} has no interface defined")
+    interface = flow.interface
+
+    async def stream_chat(request: ChatRequest) -> StreamingResponse:
+        """Stream conversational flow with Vercel AI SDK protocol."""
+        try:
+            # Convert Vercel ChatRequest to ChatMessages
+            messages = ui_request_to_domain_type(request)
+            if not messages:
+                raise ValueError("No input messages received")
+
+            current_input = messages.pop()
+            if current_input.role != MessageRole.user:
+                raise ValueError(
+                    f"Expected user message, got {current_input.role}"
+                )
+
+            # Find ChatMessage input variable
+            chat_input_var = next(
+                (var for var in flow.inputs if var.type == ChatMessage),
+                None,
+            )
+            if not chat_input_var:
+                raise ValueError("No ChatMessage input found in flow inputs")
+
+            input_data = {chat_input_var.id: current_input}
+
+            # Add session_inputs from interface
+            for session_var in interface.session_inputs:
+                if session_var.value is not None:
+                    input_data[session_var.id] = session_var.value
+
+            # Create a dynamic request model with the input data
+            RequestModel = create_input_shape(flow)
+            request_obj = RequestModel(**input_data)
+
+            initial_message = request_to_flow_message(
+                request=request_obj,
+                session_id=request.id,
+                conversation_history=messages,
+            )
+
+            return StreamingResponse(
+                _stream_sse_response(
+                    flow,
+                    initial_message,
+                    output_metadata=None,
+                    context=context,
+                ),
+                media_type="text/plain; charset=utf-8",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                    "X-Accel-Buffering": "no",
+                    "x-vercel-ai-ui-message-stream": "v1",
+                },
+            )
+
+        except ValueError as ve:
+            logger.error(f"Validation error: {ve}")
+            raise HTTPException(status_code=400, detail=str(ve)) from ve
+        except Exception as e:
+            logger.error(
+                f"Flow streaming failed for {flow_id}: {e}", exc_info=True
+            )
+            raise HTTPException(
+                status_code=500,
+                detail=f"Flow streaming failed: {str(e)}",
+            ) from e
+
+    # Register streaming endpoint
+    app.post(
+        f"/flows/{flow_id}/stream",
+        tags=["flows"],
+        summary=f"Stream {flow_id} flow (Chat)",
+        description=(
+            flow.description or f"Stream the {flow_id} conversational flow"
+        ),
+    )(stream_chat)
+
+
+def _create_completion_streaming_endpoint(
+    app: FastAPI, flow: Flow, context: ExecutorContext
+) -> None:
+    """
+    Create streaming endpoint for Complete flows.
+
+    Args:
+        app: FastAPI application instance
+        flow: Flow with Complete interface
+        secret_manager: Optional secret manager for resolving secrets
+    """
+    flow_id = flow.id
+
+    async def stream_completion(
+        request: CompletionRequest,
+    ) -> StreamingResponse:
+        """Stream completion flow with Vercel AI SDK protocol."""
+        try:
+            # Complete flows: convert CompletionRequest to input model
+            InputModel = create_input_shape(flow)
+            request_obj = completion_request_to_input_model(
+                request, InputModel
+            )
+            initial_message = request_to_flow_message(request=request_obj)
+
+            return StreamingResponse(
+                _stream_sse_response(
+                    flow,
+                    initial_message,
+                    output_metadata=None,
+                    context=context,
+                ),
+                media_type="text/plain; charset=utf-8",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                    "X-Accel-Buffering": "no",
+                    "x-vercel-ai-ui-message-stream": "v1",
+                },
+            )
+
+        except ValueError as ve:
+            logger.error(f"Validation error: {ve}")
+            raise HTTPException(status_code=400, detail=str(ve)) from ve
+        except Exception as e:
+            logger.error(
+                f"Flow streaming failed for {flow_id}: {e}", exc_info=True
+            )
+            raise HTTPException(
+                status_code=500,
+                detail=f"Flow streaming failed: {str(e)}",
+            ) from e
+
+    # Register streaming endpoint
+    app.post(
+        f"/flows/{flow_id}/stream",
+        tags=["flows"],
+        summary=f"Stream {flow_id} flow (Complete)",
+        description=(
+            flow.description or f"Stream the {flow_id} completion flow"
+        ),
+    )(stream_completion)
+
+
+def create_streaming_endpoint(
+    app: FastAPI, flow: Flow, context: ExecutorContext
+) -> None:
+    """
+    Create streaming endpoint for flow execution.
+
+    Args:
+        app: FastAPI application instance
+        flow: Flow to create endpoint for
+        secret_manager: Optional secret manager for resolving secrets
+    """
+    if flow.interface is None:
+        raise ValueError(f"Flow {flow.id} has no interface defined")
+
+    # Dispatch based on interface type
+    interface_type = flow.interface.type
+    if interface_type == "Conversational":
+        _create_chat_streaming_endpoint(app, flow, context)
+    elif interface_type == "Complete":
+        _create_completion_streaming_endpoint(app, flow, context)
+    else:
+        raise ValueError(
+            f"Unknown interface type for flow {flow.id}: {interface_type}"
+        )
+
+
+def create_rest_endpoint(
+    app: FastAPI, flow: Flow, context: ExecutorContext
+) -> None:
+    """
+    Create only the REST endpoint for flow execution.
+
+    Args:
+        app: FastAPI application instance
+        flow: Flow to create endpoint for
+        secret_manager: Optional secret manager for resolving secrets
+    """
+    RequestModel = create_input_shape(flow)
+    ResultShape = create_output_shape(flow)
+    ResponseModel = create_output_container_type(flow)
+
+    async def execute_flow_rest(
+        body: RequestModel,  # type: ignore[valid-type]
+        request: Request,
+    ) -> ResponseModel:  # type: ignore[valid-type]
+        """Execute the flow and return JSON response."""
+        try:
+            # Only pass session_id if it's provided in headers
+            kwargs = {}
+            if "session_id" in request.headers:
+                kwargs["session_id"] = request.headers["session_id"]
+
+            initial_message = request_to_flow_message(request=body, **kwargs)
+
+            # Execute flow
+            results = await run_flow(flow, initial_message, context=context)
+
+            if not results:
+                raise HTTPException(
+                    status_code=500, detail="No results returned"
+                )
+
+            return flow_results_to_output_container(
+                results,
+                output_shape=ResultShape,
+                output_container=ResponseModel,
+            )
+        except ValueError as ve:
+            logger.error(f"Validation error: {ve}")
+            raise HTTPException(status_code=400, detail=str(ve)) from ve
+        except Exception as e:
+            logger.error(
+                f"Flow execution failed for {flow.id}: {e}", exc_info=True
+            )
+            raise HTTPException(
+                status_code=500, detail=f"Flow execution failed: {str(e)}"
+            ) from e
+
+    # Set annotations for REST endpoint
+    execute_flow_rest.__annotations__ = {
+        "body": RequestModel,
+        "request": Request,
+        "return": ResponseModel,
+    }
+
+    # Register REST endpoint
+    app.post(
+        f"/flows/{flow.id}",
+        tags=["flows"],
+        description=flow.description or f"Execute the {flow.id} flow",
+        response_model=ResponseModel,
+    )(execute_flow_rest)

qtype/interpreter/executors/agent_executor.py

@@ -0,0 +1,242 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, AsyncIterator, cast
+
+from llama_index.core.agent import ReActAgent
+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 to_chat_message, to_llm, to_memory
+from qtype.interpreter.executors.invoke_tool_executor import ToolExecutionMixin
+from qtype.interpreter.tools import FunctionToolHelper
+from qtype.interpreter.types import FlowMessage
+from qtype.semantic.model import Agent, APITool, PythonFunctionTool
+
+logger = logging.getLogger(__name__)
+
+
+class AgentExecutor(StepExecutor, ToolExecutionMixin, FunctionToolHelper):
+    """Executor for Agent steps using LlamaIndex ReActAgent."""
+
+    # Agent execution should be marked as AGENT type (similar to LLM)
+    span_kind = OpenInferenceSpanKindValues.AGENT
+
+    def __init__(self, step: Agent, context: ExecutorContext, **dependencies):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, Agent):
+            raise ValueError("AgentExecutor can only execute Agent steps.")
+        self.step: Agent = step
+        self._agent = self._create_agent()
+
+    def _create_agent(self) -> ReActAgent:
+        """Create the ReActAgent instance.
+
+        Returns:
+            Configured ReActAgent instance.
+        """
+        # Convert QType tools to LlamaIndex FunctionTools
+        llama_tools = [
+            self._create_function_tool(
+                cast(APITool | PythonFunctionTool, tool)
+            )
+            for tool in self.step.tools
+        ]
+
+        # Get the LLM for the agent
+        llm = to_llm(
+            self.step.model, self.step.system_message, self._secret_manager
+        )
+
+        # Create ReActAgent
+        return ReActAgent(
+            name=self.step.id,
+            description=f"Agent for step {self.step.id}",
+            system_prompt=self.step.system_message,
+            tools=llama_tools,  # type: ignore[arg-type]
+            llm=llm,
+        )
+
+    async def process_message(
+        self,
+        message: FlowMessage,
+    ) -> AsyncIterator[FlowMessage]:
+        """Process a single FlowMessage for the Agent step.
+
+        Args:
+            message: The FlowMessage to process.
+
+        Yields:
+            FlowMessage with agent execution results.
+        """
+        # 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:
+            logger.error(f"Agent execution 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
+
+    async def _process_chat(
+        self,
+        message: FlowMessage,
+        output_variable_id: str,
+    ) -> FlowMessage:
+        """Process a chat-based agent request.
+
+        Args:
+            message: The FlowMessage to process.
+            output_variable_id: The ID of the output variable.
+
+        Returns:
+            FlowMessage with the agent chat response.
+        """
+        # Convert input variables to chat messages
+        inputs = []
+        for input_var in self.step.inputs:
+            value = message.variables.get(input_var.id)
+            if value and isinstance(value, ChatMessage):
+                inputs.append(to_chat_message(value))
+
+        # Get session ID for memory isolation
+        session_id = message.session.session_id
+
+        # Handle memory if configured
+        if self.step.memory:
+            memory = to_memory(session_id, self.step.memory)
+
+            # Add the inputs to the memory
+            await memory.aput_messages(inputs)
+            # Use the whole memory state as inputs
+            inputs = memory.get_all()
+        else:
+            # Use conversation history from session if no memory
+            conversation_history = getattr(
+                message.session, "conversation_history", []
+            )
+            if conversation_history:
+                inputs = [
+                    to_chat_message(msg) for msg in conversation_history
+                ] + inputs
+
+        # Prepare the user query (last message in inputs)
+        if not inputs:
+            raise ValueError("No input messages provided to agent")
+
+        # Get the last user message as the query
+        user_msg = inputs[-1] if inputs else None
+        chat_hist = inputs[:-1] if len(inputs) > 1 else []
+
+        # Execute agent (ReActAgent.run returns a WorkflowHandler)
+        handler = self._agent.run(
+            user_msg=user_msg,
+            chat_history=chat_hist,
+        )
+
+        # Stream or await the result
+        agent_response = await self._execute_handler(handler, message)
+
+        # Store result in memory if configured
+        if self.step.memory:
+            memory = to_memory(session_id, self.step.memory)
+            # Convert agent response to chat message
+            assistant_message = ChatMessage(
+                role=MessageRole.assistant,
+                blocks=[
+                    ChatContent(
+                        type=PrimitiveTypeEnum.text, content=agent_response
+                    )
+                ],
+            )
+            memory.put(to_chat_message(assistant_message))
+
+        # Convert result to ChatMessage
+        result_value = ChatMessage(
+            role=MessageRole.assistant,
+            blocks=[
+                ChatContent(
+                    type=PrimitiveTypeEnum.text, content=agent_response
+                )
+            ],
+        )
+
+        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-based agent request.
+
+        Args:
+            message: The FlowMessage to process.
+            output_variable_id: The ID of the output variable.
+
+        Returns:
+            FlowMessage with the agent completion response.
+        """
+        # Get input value (expecting text)
+        input_value = message.variables.get(self.step.inputs[0].id)
+        if not isinstance(input_value, str):
+            input_value = str(input_value)
+
+        # Execute agent with the input as a simple message
+        handler = self._agent.run(user_msg=input_value)
+
+        # Stream or await the result
+        agent_response = await self._execute_handler(handler, message)
+
+        # Return result as text
+        return message.copy_with_variables(
+            {output_variable_id: agent_response}
+        )
+
+    async def _execute_handler(
+        self, handler: Any, message: FlowMessage
+    ) -> str:
+        """Execute the agent handler and return the response.
+
+        Args:
+            handler: The WorkflowHandler from ReActAgent.run
+            message: The FlowMessage for stream ID generation
+
+        Returns:
+            The agent's response as a string
+        """
+        if self.context.on_stream_event:
+            # Generate a unique stream ID for this inference
+            stream_id = f"agent-{self.step.id}-{id(message)}"
+
+            async with self.stream_emitter.text_stream(stream_id) as streamer:
+                # Stream the agent response
+                async for event in handler.stream_events():
+                    if hasattr(event, "delta") and event.delta:
+                        await streamer.delta(event.delta)
+
+                # Get the final result
+                result = await handler
+        else:
+            # Non-streaming execution
+            result = await handler
+
+        return str(result)

qtype/interpreter/executors/aggregate_executor.py

@@ -0,0 +1,93 @@
+from typing import AsyncIterator
+
+from qtype.dsl.domain_types import AggregateStats
+from qtype.interpreter.base.batch_step_executor import BatchedStepExecutor
+from qtype.interpreter.base.executor_context import ExecutorContext
+from qtype.interpreter.types import FlowMessage, Session
+from qtype.semantic.model import Aggregate
+
+
+class AggregateExecutor(BatchedStepExecutor):
+    """
+    Executor for the Aggregate step.
+
+    This is a terminal, many-to-one operation that reduces an entire stream
+    to a single summary message containing counts of successful and failed
+    messages. It processes all messages without modification during the
+    processing phase, then emits a single aggregate summary during finalization.
+    """
+
+    def __init__(
+        self, step: Aggregate, context: ExecutorContext, **dependencies
+    ):
+        super().__init__(step, context, **dependencies)
+        if not isinstance(step, Aggregate):
+            raise ValueError(
+                "AggregateExecutor can only execute Aggregate steps."
+            )
+        self.step: Aggregate = step
+        self.last_session: Session | None = None
+
+    async def process_batch(
+        self,
+        batch: list[FlowMessage],
+    ) -> AsyncIterator[FlowMessage]:
+        """
+        Process messages by passing them through unchanged.
+
+        The aggregate step doesn't modify messages - it just passes them
+        through. All counting is handled by the base class's ProgressTracker,
+        which is updated as messages flow through execute(). The actual
+        aggregation happens in finalize(), which runs after all progress
+        tracking is complete.
+
+        Note: Failed messages are filtered out by the base class before
+        reaching this method, so all messages in the batch are successful.
+
+        Args:
+            batch: List of messages to process (all successful)
+
+        Yields:
+            Each message unchanged (pass-through)
+        """
+        for msg in batch:
+            # Track the last session seen for use in finalize
+            self.last_session = msg.session
+
+            # Pass message through unchanged
+            yield msg
+
+    async def finalize(
+        self,
+    ) -> AsyncIterator[FlowMessage]:
+        """
+        Emit a single summary message with aggregate statistics.
+
+        This runs after all messages have been processed and counted by the
+        base class's ProgressTracker. The summary includes counts of
+        successful, failed, and total messages.
+
+        Yields:
+            A single FlowMessage containing the aggregate statistics
+        """
+        if not self.last_session:
+            # If no messages were processed, create a summary with all zeros
+            # using a default session
+            session = Session(session_id="aggregate-no-input")
+        else:
+            session = self.last_session
+
+        # Create the aggregate stats output using counts from ProgressTracker.
+        # Since finalize() now runs AFTER all progress tracking is complete,
+        # self.progress has accurate counts of all messages processed.
+        variable_name = self.step.outputs[0].id
+        yield FlowMessage(
+            session=session,
+            variables={
+                variable_name: AggregateStats(
+                    num_successful=self.progress.items_succeeded,
+                    num_failed=self.progress.items_in_error,
+                    num_total=self.progress.items_processed,
+                )
+            },
+        )