planar 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

planar/ai/pydantic_ai_agent.py

@@ -0,0 +1,329 @@
+import inspect
+from dataclasses import dataclass
+from typing import Any, Type, cast
+
+from pydantic import BaseModel
+from pydantic_ai import models
+
+from planar.ai.agent import AgentBase
+from planar.ai.agent_utils import (
+    AgentEventType,
+    ToolCallResult,
+    create_tool_definition,
+    extract_files_from_model,
+    get_agent_config,
+    render_template,
+)
+from planar.ai.models import (
+    AgentRunResult,
+    AssistantMessage,
+    ModelMessage,
+    SystemMessage,
+    ToolDefinition,
+    ToolMessage,
+    ToolResponse,
+    UserMessage,
+)
+from planar.ai.providers import ModelSpec
+from planar.ai.pydantic_ai import ModelRunResponse, model_run
+from planar.logging import get_logger
+from planar.utils import utc_now
+from planar.workflows.models import StepType
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class Agent[
+    TInput: BaseModel | str,
+    TOutput: BaseModel | str,
+](AgentBase[TInput, TOutput]):
+    model: models.KnownModelName | models.Model = "openai:gpt-4o"
+
+    async def run_step(
+        self,
+        input_value: TInput,
+    ) -> AgentRunResult[TOutput]:
+        """Execute the agent with the provided inputs.
+
+        Args:
+            input_value: The primary input value to the agent, can be a string or Pydantic model
+            **kwargs: Alternative way to pass inputs as keyword arguments
+
+        Returns:
+            AgentRunResult containing the agent's response
+        """
+        event_emitter = self.event_emitter
+        logger.debug(
+            "agent run_step called", agent_name=self.name, input_type=type(input_value)
+        )
+        result = None
+
+        config = await get_agent_config(self.name, self.to_config())
+        logger.debug("agent using config", agent_name=self.name, config=config)
+
+        input_map: dict[str, str | dict[str, Any]] = {}
+
+        files = extract_files_from_model(input_value)
+        logger.debug(
+            "extracted files from input for agent",
+            num_files=len(files),
+            agent_name=self.name,
+        )
+        match input_value:
+            case BaseModel():
+                if self.input_type and not isinstance(input_value, self.input_type):
+                    logger.warning(
+                        "input value type mismatch for agent",
+                        agent_name=self.name,
+                        expected_type=self.input_type,
+                        got_type=type(input_value),
+                    )
+                    raise ValueError(
+                        f"Input value must be of type {self.input_type}, but got {type(input_value)}"
+                    )
+                input_map["input"] = cast(BaseModel, input_value).model_dump()
+            case str():
+                input_map["input"] = input_value
+            case _:
+                logger.warning(
+                    "unexpected input value type for agent",
+                    agent_name=self.name,
+                    type=type(input_value),
+                )
+                raise ValueError(f"Unexpected input value type: {type(input_value)}")
+
+        # Add built-in variables
+        # TODO: Make deterministic or step
+        built_in_vars = {
+            "datetime_now": utc_now().isoformat(),
+            "date_today": utc_now().date().isoformat(),
+        }
+        input_map.update(built_in_vars)
+
+        # Format the prompts with the provided arguments using Jinja templates
+        try:
+            formatted_system_prompt = (
+                render_template(config.system_prompt, input_map)
+                if config.system_prompt
+                else ""
+            )
+            formatted_user_prompt = (
+                render_template(config.user_prompt, input_map)
+                if config.user_prompt
+                else ""
+            )
+        except ValueError as e:
+            logger.exception("error formatting prompts for agent", agent_name=self.name)
+            raise ValueError(f"Missing required parameter for prompt formatting: {e}")
+
+        # Get the LLM provider and model
+        if isinstance(self.model, str):
+            model = models.infer_model(self.model)
+        else:
+            model = self.model
+
+        # Apply model parameters if specified
+        model_settings = None
+        if config.model_parameters:
+            model_settings = config.model_parameters
+
+        # Prepare structured messages
+        messages: list[ModelMessage] = []
+        if formatted_system_prompt:
+            messages.append(SystemMessage(content=formatted_system_prompt))
+
+        if formatted_user_prompt:
+            messages.append(UserMessage(content=formatted_user_prompt, files=files))
+
+        # Prepare tools if provided
+        tool_definitions = None
+        if self.tools:
+            tool_definitions = [create_tool_definition(tool) for tool in self.tools]
+
+        # Determine output type for the agent call
+        # Pass the Pydantic model type if output_type is a subclass of BaseModel,
+        # otherwise pass None (indicating string output is expected).
+        output_type: Type[BaseModel] | None = None
+        # Use issubclass safely by checking if output_type is a type first
+        if inspect.isclass(self.output_type) and issubclass(
+            self.output_type, BaseModel
+        ):
+            output_type = cast(Type[BaseModel], self.output_type)
+
+        # Execute the LLM call
+        max_turns = config.max_turns
+
+        # We use this inner function to pass "model" and "event_emitter",
+        # which are not serializable as step parameters.
+        async def agent_run_step(
+            model_spec: ModelSpec,
+            messages: list[ModelMessage],
+            turns_left: int,
+            tools: list[ToolDefinition] | None = None,
+            output_type: Type[BaseModel] | None = None,
+        ):
+            logger.debug(
+                "agent running",
+                agent_name=self.name,
+                model=model_spec,
+                model_settings=model_settings,
+                output_type=output_type,
+            )
+            if output_type is None:
+                return await model_run(
+                    model=model,
+                    max_extra_turns=turns_left,
+                    model_settings=model_settings,
+                    messages=messages,
+                    tools=tools or [],
+                    event_handler=cast(Any, event_emitter),
+                )
+            else:
+                return await model_run(
+                    model=model,
+                    max_extra_turns=turns_left,
+                    model_settings=model_settings,
+                    messages=messages,
+                    output_type=output_type,
+                    tools=tools or [],
+                    event_handler=cast(Any, event_emitter),
+                )
+
+        model_spec = ModelSpec(
+            model_id=str(model),
+            parameters=config.model_parameters,
+        )
+        result = None
+        logger.debug(
+            "agent performing multi-turn completion with tools",
+            agent_name=self.name,
+            max_turns=max_turns,
+        )
+        turns_left = max_turns
+        while turns_left > 0:
+            turns_left -= 1
+            logger.debug("agent turn", agent_name=self.name, turns_left=turns_left)
+
+            # Get model response
+            run_response = await self.as_step_if_durable(
+                agent_run_step,
+                step_type=StepType.AGENT,
+                return_type=ModelRunResponse[output_type or str],
+            )(
+                model_spec=model_spec,
+                messages=messages,
+                turns_left=turns_left,
+                output_type=output_type,
+                tools=tool_definitions or [],
+            )
+            response = run_response.response
+            turns_left -= run_response.extra_turns_used
+
+            # Emit response event if event_emitter is provided
+            if event_emitter:
+                event_emitter.emit(AgentEventType.RESPONSE, response.content)
+
+            # If no tool calls or last turn, return content
+            if not response.tool_calls or turns_left == 0:
+                logger.debug(
+                    "agent completion: no tool calls or last turn",
+                    agent_name=self.name,
+                    has_content=response.content is not None,
+                )
+                result = response.content
+                break
+
+            # Process tool calls
+            logger.debug(
+                "agent received tool calls",
+                agent_name=self.name,
+                num_tool_calls=len(response.tool_calls),
+            )
+            assistant_message = AssistantMessage(
+                content=None,
+                tool_calls=response.tool_calls,
+            )
+            messages.append(assistant_message)
+
+            # Execute each tool and add tool responses to messages
+            for tool_call_idx, tool_call in enumerate(response.tool_calls):
+                logger.debug(
+                    "agent processing tool call",
+                    agent_name=self.name,
+                    tool_call_index=tool_call_idx + 1,
+                    tool_call_id=tool_call.id,
+                    tool_call_name=tool_call.name,
+                )
+                # Find the matching tool function
+                tool_fn = next(
+                    (t for t in self.tools if t.__name__ == tool_call.name),
+                    None,
+                )
+
+                if not tool_fn:
+                    tool_result = f"Error: Tool '{tool_call.name}' not found."
+                    logger.warning(
+                        "tool not found for agent",
+                        tool_name=tool_call.name,
+                        agent_name=self.name,
+                    )
+                else:
+                    # Execute the tool with the provided arguments
+                    tool_result = await self.as_step_if_durable(
+                        tool_fn,
+                        step_type=StepType.TOOL_CALL,
+                    )(**tool_call.arguments)
+                    logger.info(
+                        "tool executed by agent",
+                        tool_name=tool_call.name,
+                        agent_name=self.name,
+                        result_type=type(tool_result),
+                    )
+
+                # Create a tool response
+                tool_response = ToolResponse(
+                    tool_call_id=tool_call.id or "call_1", content=str(tool_result)
+                )
+
+                # Emit tool response event if event_emitter is provided
+                if event_emitter:
+                    event_emitter.emit(
+                        AgentEventType.TOOL_RESPONSE,
+                        ToolCallResult(
+                            tool_call_id=tool_call.id or "call_1",
+                            tool_call_name=tool_call.name,
+                            content=tool_result,
+                        ),
+                    )
+
+                tool_message = ToolMessage(
+                    content=tool_response.content,
+                    tool_call_id=tool_response.tool_call_id or "call_1",
+                )
+                messages.append(tool_message)
+
+            # Continue to next turn
+
+        if result is None:
+            logger.warning(
+                "agent completed tool interactions but result is none",
+                agent_name=self.name,
+                expected_type=self.output_type,
+            )
+            raise ValueError(
+                f"Expected result of type {self.output_type} but got none after tool interactions."
+            )
+
+        if event_emitter:
+            event_emitter.emit(AgentEventType.COMPLETED, result)
+
+        logger.info(
+            "agent completed",
+            agent_name=self.name,
+            final_result_type=type(result),
+        )
+        return AgentRunResult[TOutput](output=cast(TOutput, result))
+
+    def get_model_str(self) -> str:
+        return str(self.model)

planar/ai/test_agent.py

@@ -435,7 +435,7 @@ async def test_agent_with_input_validation(
 async def test_agent_with_tools(
     mock_providers,
     client: PlanarTestClient,
-    app,
+    app: PlanarApp,
 ):
     """Test agent with tools for multi-turn conversations."""
     openai_mock, anthropic_mock = mock_providers
@@ -475,7 +475,7 @@ async def test_agent_with_tools(
     ):
         # Start and execute the workflow
         wf = await tools_workflow.start("complex problem")
-        result = await execute(wf)
+        result = await app.orchestrator.wait_for_completion(wf.id)
 
         # Verify the result
         assert isinstance(result, dict)

planar/app.py

@@ -1,6 +1,8 @@
 import asyncio
+import signal
 from asyncio import CancelledError
 from contextlib import asynccontextmanager
+from types import FrameType
 from typing import Any, Callable, Coroutine, Type
 
 from fastapi import APIRouter, FastAPI, HTTPException, Request
@@ -11,7 +13,7 @@ from sqlalchemy.ext.asyncio import AsyncEngine
 from typing_extensions import TypeVar
 
 from planar.ai import Agent
-from planar.config import PlanarConfig, load_environment_aware_config
+from planar.config import Environment, PlanarConfig, load_environment_aware_config
 from planar.db import DatabaseManager
 from planar.files.storage.base import Storage
 from planar.files.storage.config import create_from_config
@@ -30,8 +32,8 @@ from planar.routers.entity_router import create_entities_router
 from planar.routers.object_config_router import create_object_config_router
 from planar.routers.rule import create_rule_router
 from planar.rules.decorator import RULE_REGISTRY
+from planar.security.auth_middleware import AuthMiddleware
 from planar.security.authorization import PolicyService, policy_service_context
-from planar.security.jwt_middleware import JWTMiddleware
 from planar.session import config_var, session_context
 from planar.sse.proxy import SSEProxy
 from planar.workflows import (
@@ -92,7 +94,7 @@ class PlanarApp:
         setup_orchestrator_middleware(self)
         setup_workflow_notification_middleware(self)
         setup_tracer_middleware(self)
-        setup_jwt_middleware(self)
+        setup_auth_middleware(self)
         setup_http_exception_handler(self)
         setup_authorization_policy_service(self)
 
@@ -202,6 +204,27 @@ class PlanarApp:
 
     @asynccontextmanager
     async def _lifespan(self, app: FastAPI):
+        # We manually capture SIGINT/SIGTERM to trigger our own graceful shutdown.
+        # This is necessary because long-lived connections, such as from the SSE
+        # proxy, can cause uvicorn's default graceful shutdown to hang, preventing
+        # the lifespan shutdown logic (after the yield) from ever being reached.
+        # Our handler starts the shutdown of these components and then chains to the
+        # original uvicorn handler to allow it to proceed with its own shutdown.
+        original_handlers = {
+            signal.SIGINT: signal.getsignal(signal.SIGINT),
+            signal.SIGTERM: signal.getsignal(signal.SIGTERM),
+        }
+
+        def terminate_now(signum: int, frame: FrameType | None = None):
+            asyncio.create_task(self.graceful_shutdown())
+            handler = original_handlers.get(signal.Signals(signum))
+            if callable(handler):
+                handler(signum, frame)
+
+        signal.signal(signal.SIGINT, terminate_now)
+        signal.signal(signal.SIGTERM, terminate_now)
+
+        # Begin the normal lifespan logic
         self.db_manager.connect()
         await self.db_manager.migrate(
             self.config.use_alembic if self.config.use_alembic is not None else True
@@ -240,6 +263,10 @@ class PlanarApp:
             config_var.reset(config_tok)
 
         await self.db_manager.disconnect()
+
+        if self.storage:
+            await self.storage.close()
+
         logger.info("stopping sse")
         await self.stop_sse()
         logger.info("lifespan completed")
@@ -435,15 +462,15 @@ def setup_http_exception_handler(app: PlanarApp):
 
 def setup_cors_middleware(app: PlanarApp):
     opts = {
-        "allow_headers": app.config.cors.allow_headers,
-        "allow_methods": app.config.cors.allow_methods,
-        "allow_credentials": app.config.cors.allow_credentials,
+        "allow_headers": app.config.security.cors.allow_headers,
+        "allow_methods": app.config.security.cors.allow_methods,
+        "allow_credentials": app.config.security.cors.allow_credentials,
     }
 
-    if isinstance(app.config.cors.allow_origins, str):
-        opts["allow_origin_regex"] = app.config.cors.allow_origins
+    if isinstance(app.config.security.cors.allow_origins, str):
+        opts["allow_origin_regex"] = app.config.security.cors.allow_origins
     else:
-        opts["allow_origins"] = app.config.cors.allow_origins
+        opts["allow_origins"] = app.config.security.cors.allow_origins
 
     app.fastapi.add_middleware(
         CORSMiddleware,
@@ -451,32 +478,49 @@ def setup_cors_middleware(app: PlanarApp):
     )
 
 
-def setup_jwt_middleware(app: PlanarApp):
-    if app.config.jwt and app.config.jwt.enabled and app.config.jwt.client_id:
-        client_id = app.config.jwt.client_id
-        org_id = app.config.jwt.org_id
-        additional_exclusion_paths = app.config.jwt.additional_exclusion_paths
+def setup_auth_middleware(app: PlanarApp):
+    if (
+        app.config.security
+        and app.config.security.jwt
+        and app.config.security.jwt.client_id
+        and app.config.security.jwt.org_id
+    ):
+        client_id = app.config.security.jwt.client_id
+        org_id = app.config.security.jwt.org_id
+        additional_exclusion_paths = app.config.security.jwt.additional_exclusion_paths
         app.fastapi.add_middleware(
-            JWTMiddleware,  # type: ignore
+            AuthMiddleware,  # type: ignore
             client_id,
             org_id,
             additional_exclusion_paths,
+            service_token=app.config.security.service_token.token
+            if app.config.security.service_token
+            and app.config.security.service_token.token
+            else None,
         )
         logger.info(
-            "jwt middleware enabled",
+            "Auth middleware enabled",
             client_id=client_id,
             org_id=org_id,
             additional_exclusion_paths=additional_exclusion_paths,
         )
+    elif app.config.environment == Environment.PROD:
+        raise ValueError(
+            "Auth middleware is required in production. Please set the JWT config and optionally service token config."
+        )
     else:
-        logger.warning("JWT middleware disabled")
+        logger.warning("Auth middleware disabled")
 
 
 def setup_authorization_policy_service(app: PlanarApp):
-    if app.config.authz and app.config.authz.enabled:
+    if (
+        app.config.security
+        and app.config.security.authz
+        and app.config.security.authz.enabled
+    ):
         app.policy_service = PolicyService(
-            policy_file_path=app.config.authz.policy_file
-            if app.config.authz.policy_file
+            policy_file_path=app.config.security.authz.policy_file
+            if app.config.security.authz.policy_file
             else None
         )
         logger.info(

planar/cli.py

@@ -13,28 +13,6 @@ from planar.config import Environment
 app = typer.Typer(help="Planar CLI tool")
 
 
-class PlanarServer(uvicorn.Server):
-    """Intercept SIGINT/SIGTERM to trigger early shutdown on the app."""
-
-    def __init__(self, config: uvicorn.Config, app_import_string: str):
-        super().__init__(config)
-        self.app_import_string = app_import_string
-
-    def handle_exit(self, sig, frame):
-        # Import the PlanarApp instance and fire its early-shutdown hook
-        import asyncio
-        import importlib
-
-        module_name, var_name = self.app_import_string.split(":")
-        app_module = importlib.import_module(module_name)
-        planar_app = getattr(app_module, var_name, None)
-        if planar_app and hasattr(planar_app, "graceful_shutdown"):
-            asyncio.create_task(planar_app.graceful_shutdown())
-
-        # Continue with Uvicorn's normal shutdown procedure
-        super().handle_exit(sig, frame)
-
-
 def find_default_app_path() -> Path:
     """Checks for default app file paths (app.py, then main.py)."""
     for filename in ["app.py", "main.py"]:
@@ -94,9 +72,25 @@ def dev_command(
         "--script",
         help="Run as a script with 'uv run' instead of starting a server",
     ),
+    ssl_keyfile: str | None = typer.Option(
+        None, "--ssl-keyfile", help="Path to SSL key file"
+    ),
+    ssl_certfile: str | None = typer.Option(
+        None, "--ssl-certfile", help="Path to SSL cert file"
+    ),
 ):
     """Run Planar in development mode"""
-    run_command(Environment.DEV, port, host, config, path, app_name, script)
+    run_command(
+        Environment.DEV,
+        port,
+        host,
+        config,
+        path,
+        app_name,
+        script,
+        ssl_keyfile,
+        ssl_certfile,
+    )
 
 
 @app.command("prod")
@@ -119,9 +113,25 @@ def prod_command(
         "--script",
         help="Run as a script with 'uv run' instead of starting a server",
     ),
+    ssl_keyfile: str | None = typer.Option(
+        None, "--ssl-keyfile", help="Path to SSL key file"
+    ),
+    ssl_certfile: str | None = typer.Option(
+        None, "--ssl-certfile", help="Path to SSL cert file"
+    ),
 ):
     """Run Planar in production mode"""
-    run_command(Environment.PROD, port, host, config, path, app_name, script)
+    run_command(
+        Environment.PROD,
+        port,
+        host,
+        config,
+        path,
+        app_name,
+        script,
+        ssl_keyfile,
+        ssl_certfile,
+    )
 
 
 def run_command(
@@ -132,6 +142,8 @@ def run_command(
     path: Path | None,
     app_name: str,
     script: bool = False,
+    ssl_keyfile: str | None = None,
+    ssl_certfile: str | None = None,
 ):
     """Common logic for both dev and prod commands"""
     os.environ["PLANAR_ENV"] = env.value
@@ -188,15 +200,15 @@ def run_command(
         typer.echo(f"Starting Planar in {env.value} mode")
 
         try:
-            config = uvicorn.Config(
+            uvicorn.run(
                 app_import_string,
                 host=host or ("127.0.0.1" if env == Environment.DEV else "0.0.0.0"),
                 port=port or 8000,
                 reload=True if env == Environment.DEV else False,
                 timeout_graceful_shutdown=4,
+                ssl_keyfile=ssl_keyfile,
+                ssl_certfile=ssl_certfile,
             )
-
-            PlanarServer(config, app_import_string).run()
         except Exception as e:
             # Provide more context on import errors
             if isinstance(e, (ImportError, AttributeError)):