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

yera/events/runtime.py

@@ -0,0 +1,187 @@
+import multiprocessing as mp
+import shutil
+import string
+from time import sleep
+
+from tabulate import tabulate
+from termcolor import colored
+
+from yera.events.models import TableData
+from yera.events.stream import EventStream, OutputEvent, push_text_input
+
+
+def pretty_print_table(data: TableData):
+    print(tabulate(data.rows, headers=data.columns, tablefmt="fancy_grid"))
+
+
+class PyRuntimeExecutor:
+    """Execute an agent in a separate process using an EventStream."""
+
+    def __init__(self, agent, args):
+        """Initialise the runtime executor.
+
+        Args:
+            agent: The agent to execute.
+            args: Positional arguments to pass to the agent's ``invoke`` method.
+        """
+        self.process: mp.Process | None = None
+        self.agent = agent
+        self.args = args
+        self.stream: EventStream | None = None
+
+    @staticmethod
+    def _agent_process(agent, stream: EventStream, *args):
+        """Target function for the subprocess.
+
+        Sets the current EventStream, marks the agent as top-level, and invokes it.
+        """
+        stream.set_current()
+        agent.top_level = True
+        agent.invoke(*args)
+
+    def start(self) -> None:
+        """Start the agent execution in a new process with an internally created EventStream."""
+        if self.process is not None and self.process.is_alive():
+            # Already running; no-op
+            return
+
+        ctx = mp.get_context("spawn")
+        self.stream = EventStream.build()
+        self.process = ctx.Process(
+            target=self._agent_process,
+            args=(self.agent, self.stream, *self.args),
+        )
+        self.process.start()
+
+    def stop(self) -> None:
+        """Stop the agent process if it is running and clean up resources."""
+        # Capture process locally to avoid race conditions
+        process = self.process
+        if process is None:
+            return
+
+        process.join(timeout=2)
+        if process.is_alive():
+            process.terminate()
+            process.join()
+
+        self.process = None
+        self.stream = None
+
+    def is_running(self) -> bool:
+        """Check if the executor process is running.
+
+        Returns:
+            True if process is alive and stream is available
+        """
+        return (
+            self.process is not None
+            and self.process.is_alive()
+            and self.stream is not None
+        )
+
+    def __enter__(self):
+        """Start execution with a fresh EventStream and return self."""
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Ensure the subprocess is terminated when leaving the context."""
+        self.stop()
+        return False
+
+
+class TokenWrapper:
+    def __init__(self, max_width=100):
+        terminal_width = shutil.get_terminal_size().columns
+        self.width = min(max_width, terminal_width)  # Cap at max_width
+        self.current_line_length = 0
+
+    def print_token(self, token):
+        # Handle newlines in token
+        if "\n" in token:
+            print(token, end="", flush=True)
+            self.current_line_length = len(token.split("\n")[-1])
+            return
+
+        token_length = len(token)
+
+        # Check if token is only punctuation
+        is_punctuation = token.strip() and all(
+            c in string.punctuation for c in token.strip()
+        )
+
+        # Would this token overflow the line?
+        longer = self.current_line_length + token_length > self.width
+        if longer and not is_punctuation and self.current_line_length > 0:
+            # Break to new line
+            print("\n", end="", flush=True)
+            token = token.lstrip(" ")  # Remove leading spaces
+            token_length = len(token)
+            self.current_line_length = 0
+
+        print(token, end="", flush=True)
+        self.current_line_length += token_length
+
+    def reset(self, new_lines: int = 1):
+        """Reset line counter (call when starting new paragraph/section)"""
+        self.current_line_length = 0
+        print("\n" * (new_lines - 1))
+
+
+def stream_handler(metadata) -> OutputEvent:
+    name = colored(metadata.name, "cyan")
+    print(f"Starting {name}:")
+    print(metadata.pretty_print())
+    stream = EventStream.build()
+
+    exit_event = None
+
+    current_agent_id = None
+    current_block_id = None
+
+    wrapper = TokenWrapper(max_width=100)  # Cap at 100 columns
+
+    while exit_event is None:
+        for event in stream.iter_output():
+            if event.block_type == "exit":
+                exit_event = event
+                break
+
+            is_new_agent = event.agent_instance.agent_id != current_agent_id
+            if is_new_agent:
+                wrapper.reset(new_lines=2)
+                label = colored(event.agent_instance.agent_id.split(".")[-1], "yellow")
+                print(f"[{label}]", end="")
+
+            is_new_block = event.block_id != current_block_id
+            if is_new_block:
+                wrapper.reset()
+                label = colored(event.block_type.upper(), "green")
+                print(f"[{label}]")
+
+            # current_block_type = event.block_type
+            current_agent_id = event.agent_instance.agent_id
+            current_block_id = event.block_id
+
+            if event.block_type == "input_request":
+                print()
+                res = input(event.data.message)
+                push_text_input(res)
+
+            elif event.block_type in ["input_prompt", "system_prompt", "markdown"]:
+                wrapper.print_token(event.data.content)
+            elif event.block_type == "table":
+                pretty_print_table(event.data)
+            else:
+                print(event.data.model_dump_json(indent=2))
+
+        sleep(0.1)
+
+    exit_data = exit_event.data
+    colour = "cyan" if exit_data.exit_code == 0 else "red"
+    exit_msg = colored(f"Yera exit code {exit_data.exit_code}".upper(), colour)
+    print(f"\n\n[{exit_msg}]")
+    print(exit_data.model_dump_json(indent=2))
+    list(stream.iter_output())
+    return exit_event

yera/events/stream.py

@@ -0,0 +1,91 @@
+import multiprocessing as mp
+from collections.abc import Iterator
+from contextvars import ContextVar
+from typing import Optional
+
+from .models.in_event import InputEvent
+from .models.out_event import OutputEvent
+
+_event_stream_context: ContextVar[Optional["EventStream"]] = ContextVar(
+    "event_stream", default=None
+)
+
+
+def push_output(event: OutputEvent):
+    stream = EventStream.get_current()
+    stream.push_output(event)
+
+
+def await_input(timeout=None) -> InputEvent:
+    stream = EventStream.get_current()
+    return stream.pop_input(timeout)
+
+
+def push_input(in_event: InputEvent):
+    stream = EventStream.get_current()
+    stream.push_input(in_event)
+
+
+def push_text_input(text: str):
+    in_event = InputEvent(identifier="ABC", type="text", data=text)
+    push_input(in_event)
+
+
+class EventStream:
+    def __init__(self, out_event_queue, in_event_queue):
+        self.out_event_queue = out_event_queue
+        self.in_event_queue = in_event_queue
+        super().__init__()
+
+    def push_output(self, event: OutputEvent) -> None:
+        if not isinstance(event, OutputEvent):
+            raise TypeError(f"Expected OutputEvent got {type(event).__name__}")
+        self.out_event_queue.put(event)
+
+    def push_input(self, in_event: InputEvent) -> None:
+        if not isinstance(in_event, InputEvent):
+            raise TypeError(f"Expected InputEvent got {type(in_event).__name__}")
+        self.in_event_queue.put(in_event)
+
+    def pop_output(self, timeout: float | None = None) -> OutputEvent:
+        return self.out_event_queue.get(timeout=timeout)
+
+    def pop_input(self, timeout: float | None = None) -> InputEvent:
+        return self.in_event_queue.get(timeout=timeout)
+
+    def has_output(self) -> bool:
+        return not self.out_event_queue.empty()
+
+    def has_input(self) -> bool:
+        return not self.in_event_queue.empty()
+
+    def iter_output(self) -> Iterator[OutputEvent]:
+        while self.has_output():
+            yield self.pop_output()
+
+    def set_current(self):
+        if _event_stream_context.get() is not None:
+            raise RuntimeError("Event stream is already set in context.")
+        _event_stream_context.set(self)
+
+    @classmethod
+    def get_current(cls):
+        reg = _event_stream_context.get()
+        if reg is None:
+            raise RuntimeError(
+                "No event stream has been set in current context. "
+                "Call EventStream's set_current method first."
+            )
+        return reg
+
+    @staticmethod
+    def build() -> "EventStream":
+        reg = _event_stream_context.get()
+        if reg is not None:
+            return reg
+
+        ctx = mp.get_context("spawn")
+        manager = ctx.Manager()
+        stream = EventStream(manager.Queue(), manager.Queue())
+        stream.set_current()
+        return stream

yera/models/data_classes.py

@@ -0,0 +1,20 @@
+from typing import Literal
+
+from pydantic import BaseModel
+
+
+class Message(BaseModel):
+    role: Literal["user", "system", "assistant"]
+    content: dict | str
+
+    @staticmethod
+    def user(content: str) -> "Message":
+        return Message(role="user", content=content)
+
+    @staticmethod
+    def system(content: str) -> "Message":
+        return Message(role="system", content=content)
+
+    @staticmethod
+    def assistant(content: str) -> "Message":
+        return Message(role="assistant", content=content)

yera/models/llm_atlas_proxy.py

@@ -0,0 +1,44 @@
+from yera.config2.paths import has_global_yera_toml, has_local_yera_toml
+from yera.models.model_atlas import ModelAtlas
+
+
+class _ModelsProxy:
+    _atlas: ModelAtlas | None = None
+
+    def __init__(self):
+        self._atlas = None
+
+    def _ensure_atlas(self) -> ModelAtlas:
+        from ..config2.read import read_config
+
+        if has_global_yera_toml() or has_local_yera_toml():
+            yera_config = read_config()
+            self._atlas = ModelAtlas.from_yera_config(yera_config)
+        else:
+            self._atlas = ModelAtlas()
+        return self._atlas
+
+    def __getattr__(self, name):
+        return getattr(self._ensure_atlas(), name)
+
+    def __getitem__(self, key):
+        return self._ensure_atlas()[key]
+
+    def __dir__(self):
+        """Forward to atlas for tab completion"""
+        return dir(self._ensure_atlas())
+
+    def _ipython_key_completions_(self):
+        """Forward to atlas for IPython/Jupyter completion"""
+        return self._ensure_atlas()._ipython_key_completions_()
+
+    def list_models(self):
+        """Forward the list_models method"""
+        return self._ensure_atlas().list_models()
+
+    def _repr_html_(self):
+        """Forward HTML representation for Jupyter notebooks"""
+        return self._ensure_atlas()._repr_html_()
+
+
+llm: ModelAtlas = _ModelsProxy()

yera/models/llm_context.py

@@ -0,0 +1,99 @@
+from contextvars import ContextVar
+from typing import Optional
+
+from yera.events import input_echo, markdown, system_prompt
+from yera.models.llm_interfaces.base import BaseLLMInterface, TBaseModel
+from yera.models.llm_workspace import LLMWorkspace
+
+_current_llm_context: ContextVar[Optional["LLMContext"]] = ContextVar(
+    "llm_context", default=None
+)
+
+
+class LLMContext:
+    def __init__(
+        self,
+        *,
+        interface: BaseLLMInterface,
+        workspace: LLMWorkspace,
+    ):
+        self.workspace = workspace
+        self.interface = interface
+        self._token = None
+
+    def _push_event(self, role, content):
+        if role == "user":
+            input_echo(content)
+        elif role == "system":
+            system_prompt(content)
+        elif role == "agent":
+            markdown(content)
+        else:
+            raise ValueError(f"Unknown role: {role}")
+
+    def add_sys_line(self, content):
+        self._push_event("system", content)
+        self.workspace.add_sys_message(content)
+
+    def add_user_line(self, content):
+        self._push_event("user", content)
+        self.workspace.add_user_message(content)
+
+    def add_assistant_line(self, content):
+        self.workspace.add_assistant_message(content)
+
+    def set(self, key, value):
+        self.workspace.variables[key] = value
+
+    def get(self, key, default=None):
+        return self.workspace.variables.get(key, default)
+
+    def __getitem__(self, key):
+        return self.workspace.variables[key]
+
+    def prompt_chat(self, prompt: str, **kwargs) -> str:
+        self.add_user_line(prompt)
+        stream = self.interface.chat(self.workspace.messages, **kwargs)
+
+        tokens = []
+        # Use markdown as context manager to create a single stream
+        # All tokens will go to the same block with sequential chunk_ids
+        with markdown() as md:
+            for token in stream:
+                md.append(token)
+                tokens.append(token)
+
+        response = "".join(tokens)
+        self.add_assistant_line(response)
+        return response
+
+    def prompt_struct(self, prompt: str, cls: type[TBaseModel], **kwargs) -> TBaseModel:
+        self.add_user_line(prompt)
+        stream = self.interface.make_struct(self.workspace.messages, cls, **kwargs)
+
+        tokens = []
+        # Use markdown as context manager to create a single stream
+        # All tokens will go to the same block with sequential chunk_ids
+        with markdown() as md:
+            for token in stream:
+                md.append(token)
+                tokens.append(token)
+
+        response_json = "".join(tokens)
+        self.add_assistant_line(response_json)
+        return cls.model_validate_json(response_json)
+
+    def __enter__(self):
+        self._token = _current_llm_context.set(self)
+        self.interface.start()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.interface.stop()
+        _current_llm_context.reset(self._token)
+
+
+def llm_context() -> LLMContext:
+    llm = _current_llm_context.get()
+    if llm is None:
+        raise RuntimeError("No current LLM")
+    return llm

yera/models/llm_interfaces/anthropic.py

@@ -0,0 +1,153 @@
+import re
+from collections.abc import Iterator
+
+from anthropic import Anthropic, transform_schema
+
+from yera.config2.dataclasses import LLMConfig
+from yera.config2.keyring import DevKeyring
+from yera.models.data_classes import Message
+from yera.models.llm_interfaces.base import BaseLLMInterface, TBaseModel
+
+
+class AnthropicLLM(BaseLLMInterface):
+    def __init__(self, config: LLMConfig, overrides, creds_map):
+        kws = {**config.inference.model_dump(), **overrides}
+        self.client = None
+        self.model_id = config.model_id
+        self.struct_gen = self._supports_structured_outputs(self.model_id)
+        self.creds_map = creds_map
+        super().__init__(**kws)
+
+    @staticmethod
+    def _supports_structured_outputs(model_id: str) -> bool:
+        pattern = r"claude-(?:sonnet|opus|haiku)-(\d+)-(\d+)"
+        match = re.search(pattern, model_id)
+
+        if not match:
+            return False
+
+        major = int(match.group(1))
+        minor = int(match.group(2))
+
+        # Structured outputs supported on 4.5+
+        return (major, minor) >= (4, 5)
+
+    def start(self):
+        api_key = DevKeyring.get("providers." + self.creds_map["api_key"])
+        self.client = Anthropic(api_key=api_key)
+
+    def stop(self):
+        self.client = None
+
+    def chat(
+        self, messages: list[Message], *, max_tokens: int = 4096, **anthropic_kw
+    ) -> Iterator[str]:
+        # todo: claude does not accept system lines, but one a single system param
+        #   this means sys lines in the middle of the conversation have to be converted
+        #   into user lines. Content: f"<instruction>\n{msg['content']}\n</instruction>"
+
+        system_lines = [msg.content for msg in messages if msg.role == "system"]
+        messages = [
+            {"role": msg.role, "content": msg.content}
+            for msg in messages
+            if msg.role != "system"
+        ]
+
+        system = "\n\n".join(system_lines) if system_lines else None
+
+        kws = {
+            "model": self.model_id,
+            "max_tokens": max_tokens,
+            "messages": messages,
+            **anthropic_kw,
+            "thinking": {"type": "enabled", "budget_tokens": 1024},
+        }
+        if system:
+            kws["system"] = system
+
+        with self.client.messages.stream(
+            **kws,
+        ) as stream:
+            for event in stream:
+                if event.type == "content_block_start":
+                    if event.content_block.type == "thinking":
+                        yield "<thinking>"
+                    elif event.content_block.type == "text":
+                        yield "\n<thinking>\n\n"
+
+                elif event.type == "content_block_delta":
+                    if event.delta.type == "thinking_delta":
+                        yield event.delta.thinking
+                    elif event.delta.type == "text_delta":
+                        yield event.delta.text
+
+    def make_struct(
+        self,
+        messages: list[Message],
+        cls: type[TBaseModel],
+        *,
+        max_tokens: int = 4096,
+        **anthropic_kw,
+    ) -> Iterator[str]:
+        system_lines = [msg.content for msg in messages if msg.role == "system"]
+        conversation_messages = [
+            {"role": msg.role, "content": msg.content}
+            for msg in messages
+            if msg.role != "system"
+        ]
+
+        system = "\n\n".join(system_lines) if system_lines else None
+
+        if self.struct_gen:
+            kws = {
+                "model": self.model_id,
+                "max_tokens": max_tokens,
+                "messages": conversation_messages,
+                "betas": ["structured-outputs-2025-11-13"],
+                "output_format": {
+                    "type": "json_schema",
+                    "schema": transform_schema(cls),
+                },
+                **anthropic_kw,
+            }
+            if system:
+                kws["system"] = system
+
+            with self.client.beta.messages.stream(**kws) as stream:
+                yield from stream.text_stream
+        else:
+            tool_schema = {
+                "name": "return_structured_output",
+                "description": f"Returns structured data conforming to the {cls.__name__} schema",
+                "input_schema": cls.model_json_schema(),
+            }
+
+            kws = {
+                "model": self.model_id,
+                "max_tokens": max_tokens,
+                "messages": conversation_messages,
+                "tools": [tool_schema],
+                "tool_choice": {"type": "tool", "name": "return_structured_output"},
+                **anthropic_kw,
+            }
+            if system:
+                kws["system"] = system
+
+            tool_use_started = False
+
+            with self.client.messages.stream(**kws) as stream:
+                for event in stream:
+                    if event.type == "content_block_start":
+                        if (
+                            hasattr(event.content_block, "type")
+                            and event.content_block.type == "tool_use"
+                        ):
+                            tool_use_started = True  # We got the tool call!
+
+                    elif event.type == "content_block_delta":
+                        delta = event.delta
+                        if hasattr(delta, "type") and delta.type == "input_json_delta":
+                            yield delta.partial_json
+
+                if not tool_use_started:
+                    raise ValueError("Model did not call the structured output tool")

yera/models/llm_interfaces/aws_bedrock.py

@@ -0,0 +1,14 @@
+from collections.abc import Iterator
+
+from yera.models.data_classes import Message
+from yera.models.llm_interfaces.base import BaseLLMInterface, TBaseModel
+
+
+class AwsBedrockLLM(BaseLLMInterface):
+    def chat(self, messages: list[Message], **kwargs) -> Iterator[str]:
+        pass
+
+    def make_struct(
+        self, prompt: str, cls: type[TBaseModel], **kwargs
+    ) -> Iterator[str]:
+        pass