yera 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

infra_mvp/stream/services/session_service.py

@@ -0,0 +1,247 @@
+"""Session service for managing agent sessions."""
+
+import asyncio
+import contextlib
+import datetime
+import json
+import logging
+import uuid
+from collections.abc import AsyncIterator
+from pathlib import Path
+
+from infra_mvp.stream.schemas.session import SessionState
+from infra_mvp.stream.services.event_converter import (
+    output_event_to_json,
+    user_interaction_to_input_event,
+)
+from yera.agents.decorator import AgentFunctionWrapper
+from yera.events.runtime import PyRuntimeExecutor
+from yera.events.stream import OutputEvent
+
+
+class SessionService:
+    """Service for managing agent sessions and generating SSE events."""
+
+    def __init__(self, log_events: bool = False):
+        """Initialise the Python session service.
+
+        Args:
+            log_events: Enable logging of events to file
+        """
+        self._sessions: dict[str, SessionState] = {}
+        self._log_events = log_events
+        self._events_log_file: Path | None = None
+
+        if self._log_events:
+            # Create timestamped events log file
+            timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
+            self._events_log_file = Path(f"yera-event-log-{timestamp}.jsonl")
+            logging.info(
+                f"Events logging enabled: events will be logged to {self._events_log_file}"
+            )
+
+    def generate_session_id(self) -> str:
+        """Generate a unique session ID.
+
+        Returns:
+            UUID string for the session
+        """
+        return str(uuid.uuid4())
+
+    def create_session(self, agent: AgentFunctionWrapper, args: tuple = ()) -> str:
+        """Create a new session for the agent.
+
+        Args:
+            agent: The agent wrapper to execute.
+            args: Positional arguments to pass to the agent's ``invoke`` method.
+
+        Returns:
+            Generated session_id
+        """
+        session_id = self.generate_session_id()
+        executor = PyRuntimeExecutor(agent, args)
+
+        session_state = SessionState(
+            agent=agent,
+            executor=executor,
+            session_id=session_id,
+        )
+
+        self._sessions[session_id] = session_state
+        return session_id
+
+    def get_session(self, session_id: str) -> SessionState:
+        """Retrieve an existing session.
+
+        Args:
+            session_id: Session identifier
+
+        Returns:
+            Python agent session state
+
+        Raises:
+            KeyError: If session not found
+        """
+        if session_id not in self._sessions:
+            raise KeyError(f"Session '{session_id}' not found")
+        return self._sessions[session_id]
+
+    def add_interaction(self, session_id: str, interaction: dict) -> None:
+        """Add an interaction to a session.
+
+        Converts the UserInteraction to an InputEvent and pushes it to the
+        executor's EventStream input queue. Also stores the interaction in
+        session state for debugging.
+
+        Args:
+            session_id: Session identifier
+            interaction: Interaction dictionary (UserInteraction format)
+
+        Raises:
+            KeyError: If session not found
+            RuntimeError: If executor not started
+        """
+        session = self.get_session(session_id)
+        executor = session.executor
+
+        # Store interaction in session state for debugging
+        session.interactions.append(interaction)
+
+        # Push to executor's EventStream if executor is running
+        if not executor.is_running():
+            raise RuntimeError("Cannot add interaction: executor not started")
+
+        # Form submission: expand to one InputEvent per sub-block (same order as agent's await_input calls)
+        if interaction.get("blockType") == "form" and isinstance(
+            interaction.get("value"), list
+        ):
+            for item in interaction["value"]:
+                sub_interaction = {
+                    "blockId": item["sub_block_id"],
+                    "blockType": item["sub_block_type"],
+                    "value": item["value"],
+                }
+                input_event = user_interaction_to_input_event(sub_interaction)
+                executor.stream.push_input(input_event)
+            return
+
+        # Single interaction: one InputEvent
+        input_event = user_interaction_to_input_event(interaction)
+        executor.stream.push_input(input_event)
+
+    async def generate_events(self, session_id: str) -> AsyncIterator[str]:
+        """Generate SSE events from Python agent execution."""
+        session = self.get_session(session_id)
+        executor = session.executor
+        executor.start()
+
+        # Layouts that have had at least one await_user child; stream pauses on layout_end for these
+        layouts_with_await_user: set[str] = set()
+
+        while True:
+            if not executor.is_running():
+                async for sse in self._drain_remaining_events(session_id, executor):
+                    yield sse
+                return
+
+            event = await asyncio.to_thread(executor.stream.pop_output, timeout=None)
+            if event is None:
+                continue
+
+            yield self._format_sse_event(session_id, event)
+
+            if event.block_type == "exit":
+                self.cleanup_session(session_id)
+                return
+
+            if event.message_type == "await_user":
+                if event.parent_block_id is not None:
+                    layouts_with_await_user.add(event.parent_block_id)
+                    # Do not return; keep streaming until layout_end for this layout
+                else:
+                    return
+
+            if (
+                event.block_type == "layout_end"
+                and event.block_id in layouts_with_await_user
+            ):
+                return
+
+    def _format_sse_event(self, session_id: str, event: OutputEvent) -> str:
+        """Convert OutputEvent to SSE format and optionally log."""
+        event_json = output_event_to_json(event)
+        if self._log_events:
+            self._log_event_to_file(session_id, event_json)
+        return f"data: {json.dumps(event_json)}\n\n"
+
+    async def _drain_remaining_events(
+        self, session_id: str, executor: PyRuntimeExecutor
+    ) -> AsyncIterator[str]:
+        """Drain remaining events from a stopped executor."""
+        if executor.stream is None:
+            return
+
+        while executor.stream.has_output():
+            event = await asyncio.to_thread(executor.stream.pop_output, timeout=0.1)
+            if event is None:
+                break
+            yield self._format_sse_event(session_id, event)
+
+    def cleanup_session(self, session_id: str) -> None:
+        """Clean up session-specific resources.
+
+        Stops the executor process if running and removes the session from
+        internal storage.
+
+        Args:
+            session_id: Session identifier to clean up
+
+        Raises:
+            KeyError: If session not found
+        """
+        if session_id not in self._sessions:
+            raise KeyError(f"Session '{session_id}' not found")
+
+        session = self._sessions[session_id]
+
+        # Stop executor if it exists (terminates process)
+        if session.executor is not None:
+            session.executor.stop()
+
+        # Remove session from storage
+        del self._sessions[session_id]
+
+    def cleanup_all(self) -> None:
+        """Clean up all sessions.
+
+        Iterates through all sessions and calls cleanup_session() for each.
+        Called on server shutdown to ensure no orphaned processes remain.
+        """
+        # Create a list of session IDs to avoid modifying dict during iteration
+        session_ids = list(self._sessions.keys())
+
+        for session_id in session_ids:
+            with contextlib.suppress(KeyError):
+                self.cleanup_session(session_id)
+
+    def _log_event_to_file(self, session_id: str, event_json: dict) -> None:
+        """Log an event to the debug log file in JSONL format.
+
+        Args:
+            session_id: Session identifier
+            event_json: Event data as JSON-serializable dict
+        """
+        if self._events_log_file is None:
+            return
+
+        log_entry = {
+            "timestamp": datetime.datetime.now().isoformat(),
+            "session_id": session_id,
+            "event": event_json,
+        }
+
+        try:
+            with open(self._events_log_file, "a", encoding="utf-8") as f:
+                f.write(json.dumps(log_entry) + "\n")
+        except Exception as e:
+            logging.warning(f"Failed to write debug log: {e}")

yera/__init__.py

@@ -1,2 +1,51 @@
-__version__ = "0.0.2"
+from pydantic import Field as Field
 
+from yera.agents import agent as agent
+from yera.dsl.functions import buttons as buttons
+from yera.dsl.functions import chat as chat
+from yera.dsl.functions import date_picker as date_picker
+from yera.dsl.functions import slider as slider
+from yera.dsl.functions import struct as struct
+from yera.dsl.functions import sys_prompt as sys_prompt
+from yera.dsl.functions import text_input as text_input
+from yera.dsl.struct import Struct as Struct
+from yera.dsl.workspace import workspace as workspace
+from yera.events import (
+    action as action,
+)
+from yera.events import (
+    bar_chart as bar_chart,
+)
+from yera.events import (
+    columns as columns,
+)
+from yera.events import (
+    container as container,
+)
+from yera.events import (
+    exit_event as exit,
+)
+from yera.events import (
+    form as form,
+)
+from yera.events import (
+    line_chart as line_chart,
+)
+from yera.events import (
+    markdown as markdown,
+)
+from yera.events import (
+    quit_event as quit,
+)
+from yera.events import (
+    spinner as spinner,
+)
+from yera.events import (
+    table as table,
+)
+from yera.models.llm_atlas_proxy import llm as llm
+from yera.tools.decorator import tool as tool
+from yera.tools.tool_atlas import get_tool_atlas as get_tool_atlas
+from yera.tools.tool_utils import save_tool as save_tool
+
+__all__ = ["Struct", "agent", "exit", "quit"]

yera/agents/__init__.py

@@ -0,0 +1,2 @@
+from .context import get_agent_context as get_agent_context
+from .decorator import agent as agent

yera/agents/context.py

@@ -0,0 +1,41 @@
+from contextvars import ContextVar
+from typing import Optional
+
+from yera.agents.dataclasses import AgentMetadata
+
+_current_agent_context: ContextVar[Optional["AgentContext"]] = ContextVar(
+    "agent_context", default=None
+)
+
+
+class AgentContext:
+    def __init__(self, metadata: AgentMetadata):
+        self.metadata = metadata
+        self.instance_number = 0
+        self._token = None
+        self._graph = None
+        self._is_top_of_agent_stack = None
+
+    def __enter__(self):
+        self._is_top_of_agent_stack = _current_agent_context.get() is None
+        self._token = _current_agent_context.set(self)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        _current_agent_context.reset(self._token)
+
+    def top_level(self) -> bool:
+        if self._is_top_of_agent_stack is None:
+            raise RuntimeError("This context has not been initialised")
+        return self._is_top_of_agent_stack
+
+
+def no_active_agent_context():
+    return _current_agent_context.get() is None
+
+
+def get_agent_context():
+    agent = _current_agent_context.get()
+    if agent is None:
+        raise RuntimeError("No agent context")
+    return agent

yera/agents/dataclasses.py

@@ -0,0 +1,69 @@
+from dataclasses import dataclass
+from inspect import Parameter
+from typing import Any
+
+
+@dataclass
+class AgentInput:
+    name: str
+    type: str
+    kind: str
+    has_default: bool
+    default_value: Any
+
+    @staticmethod
+    def from_parameter(p: Parameter) -> "AgentInput":
+        return AgentInput(
+            name=p.name,
+            type=str(p.annotation),
+            kind=str(p.kind),
+            has_default=p.default is not Parameter.empty,
+            default_value=None if p.default is Parameter.empty else p.default,
+        )
+
+
+@dataclass
+class AgentInstance:
+    agent_id: str
+    instance_id: int
+
+
+@dataclass
+class AgentMetadata:
+    name: str
+    module: str
+    identifier: str
+    params: list[AgentInput]
+    return_type: str
+    docs: str | None = None
+    description: str | None = None
+
+    def make_instance_id(self, ix: int) -> AgentInstance:
+        return AgentInstance(
+            agent_id=f"{self.module}.{self.identifier}", instance_id=ix
+        )
+
+    def pretty_print(self) -> str:
+        lines = [
+            f"Agent: {self.name}",
+            f"Module: {self.module}",
+            f"Identifier: {self.identifier}",
+        ]
+
+        if self.description:
+            lines.append(f"Description: {self.description}")
+
+        if self.docs:
+            lines.append(f"Docs: {self.docs}")
+
+        lines.append(f"\nParameters ({len(self.params)}):")
+        if self.params:
+            for param in self.params:
+                default = f" = {param.default_value}" if param.has_default else ""
+                lines.append(f"  - {param.name}: {param.type}{default}")
+        else:
+            lines.append("  (none)")
+
+        lines.append(f"\nReturn Type: {self.return_type}")
+
+        return "\n".join(lines)

yera/agents/decorator.py

@@ -0,0 +1,207 @@
+import functools
+import inspect
+import traceback
+from collections.abc import Callable
+from typing import TYPE_CHECKING, get_type_hints
+
+import cloudpickle
+
+from yera.agents.context import (
+    AgentContext,
+    no_active_agent_context,
+)
+from yera.agents.dataclasses import AgentInput, AgentMetadata
+from yera.agents.typing.coerce import (
+    coerce_value,
+)
+from yera.agents.typing.utils import get_type_name
+from yera.agents.typing.validate import assert_value_is_return_type, validate_signature
+from yera.dsl.functions import sys_prompt
+from yera.events import exit_event
+from yera.events.runtime import PyRuntimeExecutor, stream_handler
+from yera.events.stream import EventStream
+from yera.models.llm_atlas_proxy import llm
+
+if TYPE_CHECKING:
+    from yera import LLMContext
+
+
+class AgentFunctionWrapper:
+    def __init__(
+        self,
+        agent_function: Callable,
+        name: str | None,
+        description: str | None,
+        model: "LLMContext | None",
+        sys_prompt_str: str | None,
+    ):
+        fn_name = agent_function.__name__
+        module = agent_function.__module__
+        identifier = f"{module}.{fn_name}"
+
+        self.agent_function = agent_function
+        self.model = model or llm.get_default()
+        self.sys_prompt = sys_prompt_str
+
+        sig = inspect.signature(agent_function)
+        type_hints = get_type_hints(agent_function)
+        validate_signature(sig, type_hints, identifier)
+        self.signature = sig
+        self.type_hints = type_hints
+
+        params = _get_agent_inputs(sig, type_hints)
+
+        self.return_cls = _get_return_type(sig, type_hints)
+        self.metadata = AgentMetadata(
+            name=name or fn_name,
+            description=description,
+            module=module,
+            docs=agent_function.__doc__,
+            identifier=identifier,
+            params=params,
+            return_type=get_type_name(self.return_cls),
+        )
+
+        functools.update_wrapper(self, agent_function)
+        self.__signature__ = sig
+
+    def _validate_and_coerce_args(self, args: tuple, kwargs: dict) -> dict:
+        try:
+            bound = self.signature.bind(*args, **kwargs)
+        except TypeError as e:
+            raise TypeError(
+                f"Invalid arguments for {self.name} ({self.metadata.identifier}): {e}"
+            ) from None
+
+        bound.apply_defaults()
+
+        coerced_args = {}
+
+        for name, arg in bound.arguments.items():
+            if arg is None:
+                continue
+
+            try:
+                coerced_args[name] = coerce_value(arg, self.type_hints[name])
+            except (TypeError, ValueError) as e:
+                raise TypeError(
+                    f"Argument '{name}' for agent '{self.metadata.identifier}': "
+                    f"expected {get_type_name(self.type_hints[name])}, "
+                    f"got {type(arg).__name__}. Error: {e}"
+                ) from None
+
+        return coerced_args
+
+    def invoke(self, *args, **kwargs):
+        with AgentContext(self.metadata) as ctx:
+            model_entered = True
+
+            try:
+                c_args = self._validate_and_coerce_args(args, kwargs)
+
+                if self.model and ctx.top_level():
+                    self.model.__enter__()
+                    model_entered = True
+
+                    if self.sys_prompt:
+                        sys_prompt(self.sys_prompt)
+
+                ret_value = self.agent_function(**c_args)
+
+                assert_value_is_return_type(ret_value, self.return_cls, self.metadata)
+
+                if ctx.top_level():
+                    exit_event(
+                        exit_code=0,
+                        reason="Yera program completed successfully.",
+                        return_value=ret_value,
+                    )
+                return ret_value
+            except Exception as e:
+                if ctx.top_level():
+                    exit_event(
+                        exit_code=1,
+                        reason=f"Error during agent execution: {e!s}",
+                        return_value=traceback.format_exc(),
+                    )
+                else:
+                    raise
+            finally:
+                if model_entered and self.model and ctx.top_level():
+                    self.model.__exit__(None, None, None)
+
+    @property
+    def name(self) -> str:
+        return self.metadata.name
+
+    def __call__(self, *args, **kwargs):
+        if len(kwargs) != 0:
+            raise NotImplementedError("Agent kwargs are currently not supported.")
+
+        EventStream.build()
+
+        if no_active_agent_context():  # top of the agent stack, about to start
+            with PyRuntimeExecutor(self, args):
+                exit_evt = stream_handler(self.metadata)
+
+            if exit_evt.data.exit_code != 0:
+                raise RuntimeError(
+                    f"Yera program exited with code: {exit_evt.data.exit_code},\n"
+                    f"AGENT: {exit_evt.agent_instance.agent_id}\n"
+                    f"REASON: {exit_evt.data.reason}\n"
+                    f"TRACEBACK: {exit_evt.data.return_value}\n"
+                )
+
+            return coerce_value(exit_evt.data.return_value, self.return_cls)
+        return self.invoke(*args)
+
+    def __reduce__(self):
+        args = (
+            self.metadata.name,
+            self.metadata.description,
+            self.model,
+            self.sys_prompt,
+        )
+        return (
+            _unpickle_agent_wrapper,
+            (cloudpickle.dumps(self.agent_function), args),
+        )
+
+
+def _get_return_type(sig, type_hints):
+    ret_type = type_hints.get("return", None)
+    if ret_type is None and sig.return_annotation is not inspect.Signature.empty:
+        ret_type = sig.return_annotation
+    return ret_type
+
+
+def _get_agent_inputs(sig, type_hints) -> list[AgentInput]:
+    params = []
+    for param_name, param in sig.parameters.items():
+        if param_name in type_hints:
+            param = param.replace(annotation=type_hints[param_name])
+        params.append(AgentInput.from_parameter(param))
+    return params
+
+
+def _unpickle_agent_wrapper(pickled_func, args):
+    agent_function = cloudpickle.loads(pickled_func)
+    return AgentFunctionWrapper(agent_function, *args)
+
+
+def agent(
+    fn=None,
+    *,
+    name=None,
+    description=None,
+    model=None,
+    sys_prompt=None,
+) -> AgentFunctionWrapper:
+    args = (name, description, model, sys_prompt)
+
+    def decorator(func) -> AgentFunctionWrapper:
+        return AgentFunctionWrapper(func, *args)
+
+    if fn is not None and callable(fn):
+        return decorator(fn)
+    return decorator