polos-sdk 0.1.0__py3-none-any.whl

polos/llm/providers/together.py

@@ -0,0 +1,40 @@
+"""Together provider - routes to OpenAI provider with chat_completions API."""
+
+from .base import register_provider
+from .openai import OpenAIProvider
+
+
+@register_provider("together")
+class TogetherProvider(OpenAIProvider):
+    """Together provider using OpenAI provider with Chat Completions API."""
+
+    def __init__(self, api_key=None):
+        """
+        Initialize Together provider.
+
+        Args:
+            api_key: Together API key. If not provided, uses TOGETHER_API_KEY env var.
+        """
+        import os
+
+        together_api_key = api_key or os.getenv("TOGETHER_API_KEY")
+        if not together_api_key:
+            raise ValueError(
+                "Together API key not provided. Set TOGETHER_API_KEY environment variable "
+                "or pass api_key parameter."
+            )
+
+        try:
+            from openai import AsyncOpenAI  # noqa: F401
+        except ImportError:
+            raise ImportError(
+                "OpenAI SDK not installed. Install it with: pip install 'polos[together]'"
+            ) from None
+
+        # Initialize with Together's base URL and chat_completions API version
+        super().__init__(
+            api_key=together_api_key,
+            base_url="https://api.together.xyz/v1",
+            llm_api="chat_completions",
+        )
+        self.supports_structured_output = False

polos/llm/stream.py

@@ -0,0 +1,183 @@
+"""Built-in LLM streaming function."""
+
+from typing import Any
+
+from ..core.context import WorkflowContext
+from ..core.workflow import _execution_context
+from ..types.types import AgentConfig
+from ..utils.agent import convert_input_to_messages
+from ..utils.client_context import get_client_or_raise
+from .providers import get_provider
+
+
+async def _llm_stream(ctx: WorkflowContext, payload: dict[str, Any]) -> dict[str, Any]:
+    """
+    Durable function for LLM streaming.
+
+    Must be executed within a workflow execution context. Uses step_outputs for durability.
+
+    Args:
+        ctx: WorkflowContext for the current execution
+        payload: Dictionary containing:
+            - agent_run_id: str
+            - agent_config: Dict with provider, model, tools, system_prompt, etc.
+            - input: Union[str, List[Dict]] - Input data
+            - agent_step: int - Step in agent conversation (1 = first, 2 = after tools, etc.)
+
+    Returns:
+        Dictionary with streaming result containing content, tool_calls, usage, etc.
+    """
+    # Check we're in a workflow execution context
+    exec_context = _execution_context.get()
+    if not exec_context or not exec_context.get("execution_id"):
+        raise ValueError("_llm_stream must be executed within an agent")
+
+    # Extract payload
+    agent_run_id = payload["agent_run_id"]
+    agent_config = AgentConfig.model_validate(payload["agent_config"])
+    input_data = payload["input"]
+    agent_step = payload.get("agent_step", 1)
+    tool_results = payload.get("tool_results")  # Optional tool results in OpenAI format
+
+    # Get LLM provider
+    provider_kwargs = {}
+    if agent_config.provider_base_url:
+        provider_kwargs["base_url"] = agent_config.provider_base_url
+    if agent_config.provider_llm_api:
+        provider_kwargs["llm_api"] = agent_config.provider_llm_api
+    provider = get_provider(agent_config.provider, **provider_kwargs)
+
+    # Convert input to messages format (without system_prompt - provider will handle it)
+    messages = convert_input_to_messages(input_data, system_prompt=None)
+
+    topic = f"workflow:{agent_run_id}"
+
+    # Stream from LLM API and publish events
+    # Call helper function to handle streaming using step.run() for durable execution
+    streaming_result = await ctx.step.run(
+        f"llm_stream:{agent_step}",
+        _stream_from_provider,
+        ctx=ctx,
+        provider=provider,
+        messages=messages,
+        agent_config=agent_config,
+        tool_results=tool_results,
+        topic=topic,
+        agent_step=agent_step,
+        agent_run_id=agent_run_id,
+    )
+
+    return streaming_result
+
+
+async def _stream_from_provider(
+    ctx: WorkflowContext,
+    provider: Any,
+    messages: list[dict[str, Any]],
+    agent_config: AgentConfig,
+    tool_results: list[dict[str, Any]] | None,
+    topic: str,
+    agent_step: int,
+    agent_run_id: str,
+) -> dict[str, Any]:
+    """
+    Helper function to stream from provider and publish events.
+
+    Returns:
+        Dictionary with chunk_index, response_content, response_tool_calls, usage, raw_output
+    """
+    from ..features.events import publish as publish_event
+
+    chunk_index = 0
+    response_content = None
+    response_tool_calls = None
+    usage = None
+    raw_output = None
+    polos_client = get_client_or_raise()
+
+    # Publish start event
+    # This is needed for invalidating events in the case of failures during
+    # agent streaming
+    # If the consumer seems stream_start event for the same agent_step,
+    # discard previous events for that agent_step
+    await publish_event(
+        client=polos_client,
+        topic=topic,
+        event_type="stream_start",
+        data={"step": agent_step},
+    )
+
+    # Stream from provider
+    # Pass agent_config and tool_results to provider
+    # Include provider_kwargs if provided
+    provider_kwargs = agent_config.provider_kwargs or {}
+    async for event in provider.stream(
+        messages=messages,
+        model=agent_config.model,
+        tools=agent_config.tools,
+        temperature=agent_config.temperature,
+        max_tokens=agent_config.max_output_tokens,
+        top_p=agent_config.top_p,
+        agent_config=agent_config.model_dump(mode="json"),
+        tool_results=tool_results,
+        output_schema=agent_config.output_schema,
+        output_schema_name=agent_config.output_schema_name,
+        **provider_kwargs,
+    ):
+        # Handle error events
+        if event.get("type") == "error":
+            error_msg = event.get("data", {}).get("error", "Unknown error")
+            raise RuntimeError(f"LLM streaming error: {error_msg}")
+
+        # Event is already in normalized format from provider
+        normalized_chunk = event
+        event_type = None
+
+        # Accumulate response data for llm_calls update
+        if normalized_chunk["type"] == "text_delta":
+            event_type = "text_delta"
+            if response_content is None:
+                response_content = ""
+            content = normalized_chunk["data"].get("content", "")
+            if content:
+                response_content += content
+        elif normalized_chunk["type"] == "tool_call":
+            event_type = "tool_call"
+            if response_tool_calls is None:
+                response_tool_calls = []
+            tool_call = normalized_chunk["data"].get("tool_call")
+            if tool_call:
+                response_tool_calls.append(tool_call)
+        elif normalized_chunk["type"] == "done":
+            usage = normalized_chunk["data"].get("usage")
+            raw_output = normalized_chunk["data"].get("raw_output")
+
+        # Publish chunk as event (skip "done" events)
+        if normalized_chunk["type"] != "done" and event_type:
+            await publish_event(
+                client=polos_client,
+                topic=topic,
+                event_type=event_type,
+                data={
+                    "step": agent_step,
+                    "chunk_index": chunk_index,
+                    "content": normalized_chunk["data"].get("content"),
+                    "tool_call": normalized_chunk["data"].get("tool_call"),
+                    "usage": normalized_chunk["data"].get("usage"),
+                    "_metadata": {
+                        "execution_id": agent_run_id,
+                        "workflow_id": ctx.workflow_id,
+                    },
+                },
+            )
+            chunk_index += 1
+
+    return {
+        "agent_run_id": agent_run_id,
+        "chunk_count": chunk_index,
+        "status": "completed",
+        "content": response_content,
+        "tool_calls": response_tool_calls if response_tool_calls else None,
+        "usage": usage if usage else None,
+        "raw_output": raw_output if raw_output else None,
+    }

polos/middleware/guardrail.py

@@ -0,0 +1,148 @@
+"""Guardrail classes for validating/modifying LLM responses before tool execution.
+
+Guardrails are executed after LLM calls but before tool execution.
+They can validate, filter, or modify the LLM content and tool_calls.
+"""
+
+import inspect
+from collections.abc import Callable
+from typing import Any
+
+from pydantic import BaseModel
+
+from ..types.types import AgentConfig, Step, ToolCall
+from .hook import HookAction, HookResult
+
+
+class GuardrailContext(BaseModel):
+    """Context specific to guardrails - what they receive.
+
+    Guardrails receive the LLM response (content and tool_calls) along with
+    execution context to make validation/modification decisions.
+    """
+
+    # LLM response data
+    content: Any | None = None  # LLM response content
+    tool_calls: list[ToolCall] | None = None  # LLM tool calls
+
+    # Execution context (for guardrail to make decisions)
+    agent_workflow_id: str | None = ""
+    agent_run_id: str | None = ""
+    session_id: str | None = None
+    user_id: str | None = None
+    llm_config: AgentConfig = AgentConfig(name="", provider="", model="")
+    steps: list[Step] = []  # Previous conversation steps
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dict for serialization."""
+        return self.model_dump(mode="json")
+
+    @classmethod
+    def from_dict(cls, data: Any) -> "GuardrailContext":
+        """Create GuardrailContext from dictionary."""
+        if isinstance(data, GuardrailContext):
+            return data
+        if isinstance(data, dict):
+            return cls.model_validate(data)
+        raise TypeError(f"Cannot create GuardrailContext from {type(data)}")
+
+
+class GuardrailResult(HookResult):
+    """Result from guardrail execution.
+
+    Inherits from HookResult but adds guardrail-specific modifications
+    for LLM content and tool_calls.
+    """
+
+    # Modify LLM response
+    modified_content: Any | None = None  # Modified content
+    modified_tool_calls: list[ToolCall] | None = None  # Modified tool calls
+    modified_llm_config: AgentConfig | None = None
+
+    # If modified_tool_calls is empty list [], no tools will be executed
+    # If modified_tool_calls is None, original tool_calls are used
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dict, including guardrail-specific fields."""
+        return self.model_dump(mode="json")
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "GuardrailResult":
+        """Create GuardrailResult from dictionary."""
+        return cls.model_validate(data)
+
+    @classmethod
+    def continue_with(cls, **modifications) -> "GuardrailResult":
+        """Continue with optional modifications.
+
+        Args:
+            **modifications: Can include modified_content, modified_tool_calls,
+                           modified_agent_config, etc.
+        """
+        return cls(action=HookAction.CONTINUE, **modifications)
+
+    @classmethod
+    def fail(cls, message: str) -> "GuardrailResult":
+        """Fail processing with an error message.
+
+        Args:
+            message: Error message to return
+        """
+        return cls(action=HookAction.FAIL, error_message=message)
+
+
+def _validate_guardrail_signature(func: Callable) -> None:
+    """Validate that guardrail function has correct signature.
+
+    Expected: (ctx: WorkflowContext, guardrail_context: GuardrailContext) -> GuardrailResult
+
+    Raises:
+        TypeError: If signature is invalid
+    """
+    sig = inspect.signature(func)
+    params = list(sig.parameters.values())
+
+    # Must have exactly 2 parameters: ctx and guardrail_context
+    if len(params) != 2:
+        raise TypeError(
+            f"Guardrail function '{func.__name__}' must have exactly 2 parameters: "
+            f"(ctx: WorkflowContext, guardrail_context: GuardrailContext). "
+            f"Got {len(params)} parameters."
+        )
+
+
+def guardrail(func: Callable | None = None):
+    """
+    Decorator to mark a function as a guardrail.
+
+    Guardrail functions must have the signature:
+        (ctx: WorkflowContext, guardrail_context: GuardrailContext) -> GuardrailResult
+
+    Usage:
+        @guardrail
+        def my_guardrail(
+            ctx: WorkflowContext, guardrail_context: GuardrailContext
+        ) -> GuardrailResult:
+            return GuardrailResult.continue_with()
+
+    Args:
+        func: The function to decorate (when used as @guardrail)
+
+    Returns:
+        The function itself (validated)
+
+    Raises:
+        TypeError: If function signature is invalid
+    """
+
+    def decorator(f: Callable) -> Callable:
+        # Validate function signature
+        _validate_guardrail_signature(f)
+        return f
+
+    # Handle @guardrail (without parentheses) - the function is passed as the first argument
+    if func is not None:
+        return decorator(func)
+
+    # Handle @guardrail() - return decorator
+    return decorator

polos/middleware/guardrail_executor.py

@@ -0,0 +1,253 @@
+"""Guardrail execution infrastructure.
+
+Guardrails are executed sequentially within a workflow execution context
+and support durable execution.
+"""
+
+import json
+from collections.abc import Callable
+from typing import Any
+
+from ..core.context import WorkflowContext
+from ..core.workflow import _execution_context
+from ..types.types import AgentConfig
+from .guardrail import GuardrailContext, GuardrailResult
+from .hook import HookAction
+
+
+def _get_guardrail_identifier(guardrail: Callable | str, index: int) -> str:
+    """Get a unique identifier for a guardrail.
+
+    Uses function name if callable, otherwise uses a string identifier.
+    """
+    if isinstance(guardrail, str):
+        # For string guardrails, use a truncated version
+        return f"guardrail_string_{guardrail[:50]}"
+    if hasattr(guardrail, "__name__") and guardrail.__name__ != "<lambda>":
+        return guardrail.__name__
+    return f"guardrail_{index}"
+
+
+async def execute_guardrails(
+    guardrail_name: str,
+    guardrails: list[Callable | str],
+    guardrail_context: GuardrailContext,
+    ctx: WorkflowContext,
+    agent_config: AgentConfig | None = None,
+) -> GuardrailResult:
+    """
+    Execute a list of guardrails sequentially and return the combined result.
+
+    Guardrails are executed within a workflow execution context. Each guardrail execution:
+    1. Checks for cached result (for durable execution)
+    2. If cached, returns cached result
+    3. If not cached, executes guardrail and stores result
+
+    Guardrails can be:
+    - Callable: Functions decorated with @guardrail
+    - str: String prompts evaluated using LLM with structured output
+
+    Each guardrail can:
+    - Return CONTINUE to proceed to the next guardrail
+    - Return FAIL to stop execution with an error
+
+    Modifications from guardrails are accumulated and applied in order.
+
+    Args:
+        guardrails: List of guardrail callables or strings
+        guardrail_context: Context to pass to guardrails
+        ctx: WorkflowContext for the current execution
+        agent_config: Optional agent config (model, provider, etc.) for string guardrails
+
+    Returns:
+        GuardrailResult with action and any modifications
+
+    Raises:
+        ValueError: If not executed within a workflow execution context
+    """
+    if not guardrails:
+        return GuardrailResult.continue_with()
+
+    # Check we're in a workflow execution context
+    exec_context = _execution_context.get()
+    if not exec_context or not exec_context.get("execution_id"):
+        raise ValueError("Guardrails must be executed within a workflow execution context")
+
+    # Accumulated modifications
+    if guardrail_context.content is not None:
+        if isinstance(guardrail_context.content, str):
+            modified_content = guardrail_context.content
+        else:
+            modified_content = guardrail_context.content.copy()
+    else:
+        modified_content = None
+
+    modified_tool_calls = (
+        guardrail_context.tool_calls.copy() if guardrail_context.tool_calls else []
+    )
+    modified_llm_config = guardrail_context.llm_config.copy()
+
+    # Execute guardrails sequentially
+    for index, guardrail in enumerate(guardrails):
+        # Get identifier for durable execution
+        guardrail_id = _get_guardrail_identifier(guardrail, index)
+
+        # Update guardrail context with accumulated modifications
+        guardrail_context.content = modified_content
+        guardrail_context.tool_calls = modified_tool_calls
+        guardrail_context.llm_config = modified_llm_config
+
+        # Execute guardrail (callable or string)
+        if isinstance(guardrail, str):
+            # String guardrail: evaluate using LLM
+            if not agent_config:
+                raise ValueError("agent_config is required for string guardrails")
+
+            guardrail_result = await _execute_string_guardrail(
+                guardrail_name=f"{guardrail_name}:{index}",
+                guardrail_string=guardrail,
+                guardrail_context=guardrail_context,
+                ctx=ctx,
+                agent_config=agent_config,
+                index=index,
+            )
+        elif callable(guardrail):
+            guardrail_result = await ctx.step.run(
+                f"{guardrail_name}.{guardrail_id}.{index}", guardrail, ctx, guardrail_context
+            )
+        else:
+            raise TypeError(
+                f"Guardrail at index {index} is neither callable nor string: {type(guardrail)}"
+            )
+
+        # Ensure result is GuardrailResult
+        if not isinstance(guardrail_result, GuardrailResult):
+            guardrail_result = GuardrailResult.fail(
+                f"Guardrail '{guardrail_id}' returned invalid result type: "
+                f"{type(guardrail_result)}. Expected GuardrailResult."
+            )
+
+        # Apply modifications
+        if guardrail_result.modified_content is not None:
+            modified_content = guardrail_result.modified_content
+
+        if guardrail_result.modified_tool_calls is not None:
+            modified_tool_calls = guardrail_result.modified_tool_calls
+
+        if guardrail_result.modified_llm_config is not None:
+            modified_llm_config.update(guardrail_result.modified_llm_config)
+
+        # Check action
+        if guardrail_result.action == HookAction.FAIL:
+            # Fail execution - return error
+            return guardrail_result
+
+        # CONTINUE - proceed to next guardrail
+
+    # All guardrails completed with CONTINUE - return accumulated modifications
+    return GuardrailResult.continue_with(
+        modified_content=modified_content,
+        modified_tool_calls=modified_tool_calls,
+        modified_llm_config=modified_llm_config,
+    )
+
+
+async def _execute_string_guardrail(
+    guardrail_name: str,
+    guardrail_string: str,
+    guardrail_context: GuardrailContext,
+    ctx: WorkflowContext,
+    agent_config: dict[str, Any],
+    index: int,
+) -> GuardrailResult:
+    """Execute a string guardrail using LLM with structured output.
+
+    Args:
+        guardrail_string: The guardrail prompt/instruction
+        guardrail_context: Context containing LLM response to validate
+        ctx: WorkflowContext for the current execution
+        agent_config: Agent configuration (model, provider, etc.)
+
+    Returns:
+        GuardrailResult indicating pass/fail
+    """
+    import json
+
+    from ..llm import _llm_generate
+
+    # Create JSON schema for structured output: {passed: bool, reason: Optional[str]}
+    guardrail_output_schema = {
+        "type": "object",
+        "properties": {"passed": {"type": "boolean"}, "reason": {"type": "string", "default": ""}},
+        "required": ["passed", "reason"],
+        "additionalProperties": False,
+    }
+
+    # Build the evaluation prompt
+    # Include the guardrail instruction and the LLM response to validate
+    llm_content = guardrail_context.content
+    if isinstance(llm_content, str):
+        content_to_validate = llm_content
+    else:
+        content_to_validate = json.dumps(llm_content) if llm_content else ""
+
+    if not content_to_validate:
+        return GuardrailResult.continue_with()
+
+    evaluation_prompt = f"""{guardrail_string}
+
+Please evaluate the following LLM response against the criteria above.
+Return a JSON object with "passed" (boolean) and optionally "reason"
+(string if the guardrail fails).
+
+LLM Response to evaluate:
+{content_to_validate}"""
+
+    # Prepare agent config for guardrail evaluation
+    # Use the agent's model, provider, etc., but override output_schema
+    config_dict = agent_config.model_dump(mode="json")
+    config_dict["output_schema"] = guardrail_output_schema
+    config_dict["output_schema_name"] = "GuardrailEvaluationResult"
+    guardrail_agent_config = AgentConfig.model_validate(config_dict)
+
+    # Call _llm_generate to evaluate the guardrail
+    llm_result = await _llm_generate(
+        ctx,
+        {
+            "agent_run_id": guardrail_context.agent_run_id,
+            "agent_config": guardrail_agent_config.model_dump(mode="json"),
+            "input": evaluation_prompt,
+            "agent_step": 0,  # Guardrail evaluation doesn't count as agent step
+            "guardrails": None,  # Don't recurse guardrails on guardrail evaluation
+            "guardrail_max_retries": 0,
+        },
+    )
+
+    return await ctx.step.run(
+        f"{guardrail_name}.{guardrail_string[:50]}.{index}", _parse_llm_guardrail_result, llm_result
+    )
+
+
+async def _parse_llm_guardrail_result(llm_result: dict[str, Any]) -> GuardrailResult:
+    """Parse the LLM result and return a GuardrailResult."""
+    # Parse the structured output
+    result_content = llm_result.get("content", "")
+    if not result_content:
+        return GuardrailResult.fail("Guardrail evaluation returned empty response")
+
+    try:
+        evaluation_result = (
+            json.loads(result_content) if isinstance(result_content, str) else result_content
+        )
+
+        passed = evaluation_result.get("passed", False)
+        reason = evaluation_result.get("reason")
+
+        if passed:
+            return GuardrailResult.continue_with()
+        else:
+            error_message = reason or "Guardrail validation failed"
+            return GuardrailResult.fail(error_message)
+    except (json.JSONDecodeError, TypeError, AttributeError) as e:
+        # If parsing fails, treat as failure
+        return GuardrailResult.fail(f"Failed to parse guardrail evaluation result: {str(e)}")