Chatbot, LLM companion and data retrieval framework
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Tools for maintaining a chat session (i.e. LLM context with messages)."""

from __future__ import annotations

import asyncio
import enum
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Literal, Protocol

import aioconsole  # type: ignore[import]
import openai
from loguru import logger
from overrides import override

from rrosti.llm_api import openai_api
from rrosti.utils import misc
from rrosti.utils.config import config

ERROR_SLEEP_SECONDS = 2


class Importance(enum.IntEnum):
    """Importance of a message."""

    NOISE = 0
    LOW = 10
    MEDIUM = 20
    HIGH = 30


@dataclass
class Message:
    """A message in a chat session."""

    role: Literal["user", "assistant", "system"]
    text: str

    # To manage the size of the context, we may need to trim messages.
    # 1. Each message has an importance score; when the context is full, we always trim the earliest
    #    message with the lowest importance score.
    # 2. Messages can have an optional ttl (time-to-live) field which gets decremented at each
    #    user input. When the ttl reaches 0, the message is removed.
    importance: Importance

    ttl: int | None = None

    # In an assistant message, the model used to process it. In a user message, the model to use for the next assistant
    # message.
    model: str | None = None

    cost: float | None = None
    time_used: float | None = None

    def __post_init__(self) -> None:
        assert not (self.role == "user" and self.cost is not None), "User messages cannot have a cost"

    def as_user_message(self, model: str | None = None) -> Message:
        """
        Return a user message with the same text as this message, no with cost and with possibly a different model.
        """
        return Message(role="user", text=self.text, importance=self.importance, model=model)

    def to_string(self, agent: str | None = None) -> str:
        meta: list[str] = [self.role]
        if self.model:
            meta.append(self.model)
        if self.cost is not None:
            meta.append(f"{self.cost:.5f} USD")
        if self.time_used is not None:
            meta.append(f"time={self.time_used:.2f}")
        if self.ttl is not None:
            meta.append(f"ttl={self.ttl}")
        meta.append(f"importance={self.importance}")
        meta_str = ", ".join(meta)
        header = f"{agent} ({meta_str})" if agent else meta_str
        lines = []
        lines.append("-" * 5 + " " + header + " " + "-" * (70 - 7 - len(header)))
        lines.append(self.text.strip())
        lines.append("-" * 70)
        return "\n".join(lines)


async def get_user_input(prompt: str = "", end_line: str | None = None) -> str:
    """Get user input from the console, optionally ending on a specific line."""
    lines: list[str] = []
    while True:
        line = await aioconsole.ainput(f"{prompt}> ")
        if end_line is not None and line.strip() == end_line:
            break
        lines.append(line)
        if end_line is None:
            break

    return "\n".join(lines)


class LLM(ABC):
    """A language model."""

    @abstractmethod
    async def chat_completion(self, messages: list[Message], agent_name: str = "", model: str | None = None) -> Message:
        ...


class OpenAI(LLM):
    """OpenAI chat language model."""

    temperature: float
    prompt_tokens: int
    completion_tokens: int
    cost: float
    openai_provider: openai_api.OpenAIApiProvider

    def __init__(self, openai_provider: openai_api.OpenAIApiProvider, temperature: float = 1.0) -> None:
        self.openai_provider = openai_provider
        self.temperature = temperature
        self.prompt_tokens = 0
        self.completion_tokens = 0
        self.cost = 0.0

    @override
    async def chat_completion(
        self,
        messages: list[Message],
        agent_name: str = "",
        model: str | None = None,
        msg_importance: Importance = Importance.MEDIUM,
    ) -> Message:
        del agent_name  # intentionally unused parameter

        if model is None:
            model = config.openai_api.chat_completion_model

        while True:
            try:
                start_time = asyncio.get_event_loop().time()
                resp = await self.openai_provider.acreate_chat_completion(
                    messages=[dict(role=message.role, content=message.text) for message in messages],
                    max_tokens=config.openai_api.completion_max_tokens,
                    model=model,
                )
                end_time = asyncio.get_event_loop().time()
                elapsed_time = end_time - start_time
                break
            except (
                openai.error.RateLimitError,
                openai.error.Timeout,
                openai.error.APIConnectionError,
                openai.error.APIError,
                openai.error.ServiceUnavailableError,
            ) as e:
                logger.error("OpenAI API error: {}. Sleeping for {} seconds...", e, ERROR_SLEEP_SECONDS)
                await asyncio.sleep(ERROR_SLEEP_SECONDS)

        prompt_tokens = resp["usage"]["prompt_tokens"]
        completion_tokens = resp["usage"]["completion_tokens"]

        self.prompt_tokens += prompt_tokens
        self.completion_tokens += completion_tokens

        model_cost = config.openai_api.model_cost[
            model if model is not None else config.openai_api.chat_completion_model
        ]

        # TODO: handle prompts that did not return any content at all (currently it crashes the process and starts a new
        # session with the client without any Warning
        text: str | None = resp["choices"][0]["message"].get("content")
        if not text:
            # If nothing is returned, the query (or answer) got blocked by Azure. This is the identifier so that the
            # state machine can give the user an answer explaining it.
            text = "$$$error$$$"

        return Message(
            role="assistant",
            text=text,
            importance=msg_importance,
            model=model,
            cost=model_cost.calculate(prompt_tokens, completion_tokens),
            time_used=elapsed_time,
        )


class UserInputLLM(LLM):
    """LLM that asks the user for input. For testing."""

    @override
    async def chat_completion(self, messages: list[Message], agent_name: str = "", model: str | None = None) -> Message:
        for message in messages:
            print("-" * 5 + " " + message.role + " " + "-" * (40 - 7 - len(message.role)))
            print(message.text)
        print("-" * 40)
        print("Enter message. '.' on a line alone finishes.")
        text = await get_user_input(agent_name, end_line=".")
        return Message(role="user", importance=Importance.HIGH, text=text, model=model or "user-input")


class MessageCallback(Protocol):
    def __call__(self, message: Message, agent: str | None, quiet: bool) -> None:
        ...


class ChatSession:
    """
    A chat session (i.e. LLM context, consisting of messages).

    Also contains functionality to prune messages to prevent the LLM context from becoming too long.
    """

    messages: list[Message]
    llm: LLM
    _name: str  # A human-readable name. Could be e.g. a name of an agent. May be empty string.

    _message_callback: MessageCallback | None

    def __init__(self, llm: LLM, name: str = "", callback: MessageCallback | None = None) -> None:
        self.messages = []
        self.prompt_tokens = 0
        self.completion_tokens = 0
        self.llm = llm
        self._name = name
        self._message_callback = callback

    def _prune_unimportant_message(self) -> None:
        """Prune the least important message."""
        assert self.messages, "No messages to prune"
        min_importance = min(message.importance for message in self.messages)
        for message in self.messages:
            if message.importance == min_importance:
                logger.info(
                    "Pruning unimportant message (importance={}): {}",
                    message.importance,
                    misc.truncate_string(message.text),
                )
                self.messages.remove(message)
                return

    def decrease_ttls(self) -> None:
        """Decrease ttls and kill messages with expired ttls."""

        new_messages = []

        for message in self.messages:
            if message.ttl is not None:
                message.ttl -= 1
                logger.info("Decreasing ttl of message to {}: {}", message.ttl, misc.truncate_string(message.text))

            if message.ttl is not None and message.ttl <= 0:
                logger.info("Killing message with expired ttl: {}", misc.truncate_string(message.text))
            else:
                new_messages.append(message)

        self.messages = new_messages

    def add_message(
        self,
        message: Message,
        quiet: bool = False,
    ) -> None:
        """Add a message to the chat session."""

        assert message.ttl is None or message.ttl > 0, "llm_ttl must be None or positive"

        self.messages.append(message)
        logger.info("Message({}): {}", message.role, misc.truncate_string(message.text))

        if self._message_callback is not None:
            self._message_callback(message, self._name, quiet)

    async def generate_response(self, model: str | None = None) -> Message:
        while True:
            try:
                resp = await self.llm.chat_completion(self.messages, agent_name=self._name, model=model)
                break
            except openai.error.InvalidRequestError as e:
                if "maximum context length" in str(e):
                    logger.info("Context length exceeded: {}", str(e))
                    self._prune_unimportant_message()
                    continue
                raise

        self.add_message(resp)
        return resp

rrosti/chat/chat_session.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""AST for a dialog state machine."""

from __future__ import annotations

import string
from abc import ABC, abstractmethod, abstractproperty
from dataclasses import dataclass, field
from typing import Any, Literal, Mapping, Sequence
from uuid import uuid4

from overrides import override

from rrosti.utils.misc import JSONValue

MSG_CUT_LEN = 40


class JSONMixin:
    """Mixin for classes that can be converted to JSON."""

    @abstractmethod
    def to_json(self) -> JSONValue:
        """Convert to a dictionary that can be dumped as JSON or YAML."""


class GraphvizMixin(ABC):
    """Mixin for classes that can be converted to Graphviz dot format."""

    @abstractmethod
    def viz_label(self) -> str | None:
        """Return a label for the node in Graphviz dot string format."""


class ActionVisitor(ABC):  # noqa: B024 (intentional abstract class without abstract methods)
    """Visitor for actions."""

    def end_action(self, action: EndAction) -> None:  # noqa: B027 (intentionally empty and not abstract)
        """Visit an end action."""

    def goto_action(self, action: GotoAction) -> None:  # noqa: B027 (intentionally empty and not abstract)
        """Visit a goto action."""

    def message_action(self, action: MessageAction) -> None:  # noqa: B027 (intentionally empty and not abstract)
        """Visit a message action."""

    def send_action(self, action: SendAction) -> None:  # noqa: B027 (intentionally empty and not abstract)
        """Visit a send action."""


class AsyncActionVisitor(ABC):  # noqa: B024 (intentional abstract class without abstract methods)
    """Visitor for actions."""

    async def end_action(self, action: EndAction) -> None:  # noqa: B027 (intentionally empty and not abstract)
        """Visit an end action."""

    async def goto_action(self, action: GotoAction) -> None:  # noqa: B027 (intentionally empty and not abstract)
        """Visit a goto action."""

    async def message_action(self, action: MessageAction) -> None:  # noqa: B027 (intentionally empty and not abstract)
        """Visit a message action."""

    async def send_action(self, action: SendAction) -> None:  # noqa: B027 (intentionally empty and not abstract)
        """Visit a send action."""


@dataclass
class Config:
    model_ast: str | None
    model: str | None = field(init=False)

    # TODO: This should be a non-dataclass and we should have a setter

    def __post_init__(self) -> None:
        self.model = self.model_ast

    # update the config with the given config, returning a new config
    def update(self, config: Config) -> Config:
        return Config(model_ast=config.model_ast or self.model_ast)

    def set_outer_config(self, outer: Config) -> None:
        self.model = self.model_ast or outer.model


class ASTNode(ABC):
    """Base class for all AST nodes."""

    _inherited_config: Config = Config(model_ast=None)

    @abstractproperty
    def _children(self) -> Sequence[ASTNode]:
        """Return a list of children."""

    def _link_configs(self, outer: Config) -> None:
        """Link the config to the outer config."""

        if hasattr(self, "config"):
            object.__setattr__(self, "_inherited_config", self.config)

        self._inherited_config.set_outer_config(outer)
        for child in self._children:
            child._link_configs(self._inherited_config)


class ASTLeaf(ASTNode):
    @property
    def _children(self) -> Sequence[ASTNode]:
        return []


class ActionBase(JSONMixin, GraphvizMixin, ASTNode, ABC):
    """Base class for actions."""

    @abstractproperty
    @override
    def _children(self) -> Sequence[ASTNode]:
        """Return the children of this node."""

    @abstractmethod
    def accept(self, visitor: ActionVisitor) -> None:
        """Accept a visitor."""

    @abstractmethod
    async def aaccept(self, visitor: AsyncActionVisitor) -> None:
        """Accept a visitor."""


class _VarExpansionFormatter(string.Formatter):
    """A string formatter that expands variables and records encountered function calls."""

    variables: dict[str, str]
    used_vars: set[str]
    placeholders: dict[str, str]

    def __init__(self, variables: dict[str, str]) -> None:
        super().__init__()
        self.variables = variables
        self.used_vars = set()
        self.placeholders = {}

    @override
    def get_value(self, key: int | str, args: Sequence[Any], kwargs: Mapping[str, Any]) -> Any:
        assert isinstance(key, str), "Only string keys are supported"
        assert len(args) == 0, "Positional arguments are not supported"
        assert len(kwargs) == 0, "Keyword arguments are not supported"
        if key in self.variables:
            self.used_vars.add(key)
            return self.variables[key]
        if key.endswith(")"):
            assert key.endswith("()"), key
            # add functions to undefineds so this can be redone with their values
            uuid = f"FUNCALL({uuid4()})"
            self.placeholders[uuid] = key
            return uuid
        raise ValueError(f"Undefined variable: {key}")


@dataclass
class _InterpolatedString:
    """
    Represents a string with variables interpolated.

    However, functions may not be resolved yet. Hence, they will be replaced with unique values.

    The `placeholders` field contains the mapping from the unique values to the function calls.
    """

    text: str
    placeholders: dict[str, str]


def _interpolate_vars(vars: dict[str, str], s: str, used_vars: frozenset[str] = frozenset([])) -> _InterpolatedString:
    """
    Interpolate variables in the variable definitions.

    Returns: Intepolated string on success, set of undefined variables on failure.
    """

    formatter = _VarExpansionFormatter(vars)
    formatted = formatter.format(s)
    if formatter.used_vars & used_vars:
        raise ValueError(f"Variable expansion loop detected: {formatter.used_vars & used_vars}")
    if not formatter.used_vars:
        return _InterpolatedString(formatted, formatter.placeholders.copy())

    news = _interpolate_vars(vars, formatted, used_vars | formatter.used_vars)
    news.placeholders.update(formatter.placeholders)
    return news


@dataclass(frozen=True)
class EndAction(ASTLeaf, ActionBase):
    """Represents the end of a conversation."""

    @override
    def to_json(self) -> JSONValue:
        return "end"

    def viz_label(self) -> str:
        return "end"

    @override
    def accept(self, visitor: ActionVisitor) -> None:
        visitor.end_action(self)

    @override
    async def aaccept(self, visitor: AsyncActionVisitor) -> None:
        await visitor.end_action(self)


@dataclass(frozen=True)
class GotoAction(ASTLeaf, ActionBase):
    """Represents an action that goes to another state in the state machine."""

    label: str

    @override
    def to_json(self) -> JSONValue:
        return {"goto": self.label}

    def viz_label(self) -> str:
        return f"GOTO {self.label}"

    @override
    def accept(self, visitor: ActionVisitor) -> None:
        visitor.goto_action(self)

    @override
    async def aaccept(self, visitor: AsyncActionVisitor) -> None:
        await visitor.goto_action(self)


def _interpolations(s: str) -> list[str]:
    """Get all interpolations in a message text."""

    return [fname for _, fname, _, _ in string.Formatter().parse(s) if fname is not None]


@dataclass(frozen=True)
class MessageAction(ASTLeaf, ActionBase):
    """Represents a message action, adding either a user or assistant message to the conversation."""

    # TODO: make the interpolation more robust and runnable only once (and automatically)

    role: Literal["user", "assistant", "system"]
    text: str
    placeholders: dict[str, str] = field(init=False)

    @override
    def to_json(self) -> JSONValue:
        if self.role == "user":
            return {"message": self.text}
        if self.role == "assistant":
            return {"assistant_message": self.text}
        if self.role == "system":
            return {"system_message": self.text}
        assert False, f"Invalid role: {self.role}"

    def viz_label(self) -> str:
        msg = self.text
        msg = msg.replace("\n", r"\n")
        if len(msg) > MSG_CUT_LEN:
            msg = msg[:MSG_CUT_LEN] + "..."
        return f'MSG({self.role}, "{msg}")'

    @override
    def accept(self, visitor: ActionVisitor) -> None:
        visitor.message_action(self)

    @override
    async def aaccept(self, visitor: AsyncActionVisitor) -> None:
        await visitor.message_action(self)

    def __post_init__(self) -> None:
        object.__setattr__(self, "placeholders", {})

    @property
    def interpolations(self) -> list[str]:
        """All interpolations in the message text."""
        return _interpolations(self.text)

    def _interpolate_vars(self, vars: dict[str, str]) -> None:
        """Interpolate variables in the message text."""
        s = _interpolate_vars(vars, self.text)
        object.__setattr__(self, "text", s.text)
        object.__setattr__(self, "placeholders", s.placeholders)


@dataclass(frozen=True)
class SendAction(ASTLeaf, ActionBase):
    """Represents a send action, sending a message (and transfering control) to another agent."""

    to: str
    next_state: str

    @override
    def to_json(self) -> JSONValue:
        return dict(send=dict(to=self.to, next_state=self.next_state))

    def viz_label(self) -> str:
        return f"SEND {self.to}, {self.next_state}"

    @override
    def accept(self, visitor: ActionVisitor) -> None:
        visitor.send_action(self)

    @override
    async def aaccept(self, visitor: AsyncActionVisitor) -> None:
        await visitor.send_action(self)


@dataclass(frozen=True)
class ActionList(JSONMixin, ASTNode):
    """Represents a list of actions to be executed in order."""

    actions: list[ActionBase]

    @property
    def _children(self) -> Sequence[ASTNode]:
        return self.actions

    def __post_init__(self) -> None:
        if len(self.actions) == 0:
            raise ValueError("Action list must contain at least one action")
        for action in self.actions[:-1]:
            if isinstance(action, GotoAction):
                raise ValueError("Only the last action may be a goto.")  # noqa: TRY004 (TypeError is not a good match)
            if isinstance(action, EndAction):
                raise ValueError("Only the last action may be an end.")  # noqa: TRY004 (TypeError is not a good match)
            if isinstance(action, SendAction):
                raise ValueError("Only the last action may be a send.")  # noqa: TRY004 (TypeError is not a good match)

    def target_state_name(self, curr: str | None = None) -> str | None:
        """Returns the target state of the action list, given the current state."""
        if isinstance(self.actions[-1], GotoAction):
            return self.actions[-1].label
        if isinstance(self.actions[-1], EndAction):
            return "END"
        return curr

    @override
    def to_json(self) -> JSONValue:
        return [a.to_json() for a in self.actions]

    def viz_label(self) -> str:
        actions = self.actions
        # If the last action is a goto or end, omit it
        if isinstance(actions[-1], (GotoAction, EndAction)):
            actions = actions[:-1]
        edge_labels = [a.viz_label() for a in actions]
        return "\n".join(label for label in edge_labels if label is not None)

    def accept(self, visitor: ActionVisitor) -> None:
        for action in self.actions:
            action.accept(visitor)

    async def aaccept(self, visitor: AsyncActionVisitor) -> None:
        for action in self.actions:
            await action.aaccept(visitor)

    def _interpolate_vars(self, vars: dict[str, str]) -> None:
        """Interpolate variables in the message text."""
        for action in self.actions:
            if isinstance(action, MessageAction):
                action._interpolate_vars(vars)


class ConditionBase(JSONMixin, ASTNode, ABC):
    """Base class for conditions."""

    action: ActionList

    config: Config

    def target_state(self, curr: str | None = None) -> str | None:
        """Get the target state of the condition. 'curr' is returned if there is no goto."""
        return self.action.target_state_name(curr)

    @abstractmethod
    def condition_label(self) -> str:
        """Get a human-readable label for the condition."""

    @abstractmethod
    def is_triggered_by(self, llm_output: str) -> bool:
        """Returns True if the condition is triggered by the given LLM output."""


@dataclass(frozen=True)
class Condition(ConditionBase):
    """Represents a condition that must be met for the action to be executed."""

    action: ActionList
    contains: str
    config: Config  # the config in the yaml, not necessarily the full config

    @property
    def _children(self) -> Sequence[ASTNode]:
        return [self.action]

    @override
    def to_json(self) -> JSONValue:
        return {"contains": self.contains, "action": self.action.to_json()}

    @override
    def condition_label(self) -> str:
        return f'contains: "{self.contains}"'

    @override
    def is_triggered_by(self, llm_output: str) -> bool:
        return self.contains in llm_output


@dataclass(frozen=True)
class DefaultCondition(ConditionBase):
    """Represents a condition that met if no other conditions are met."""

    action: ActionList
    config: Config  # the config in the yaml, not necessarily the full config

    @property
    def _children(self) -> Sequence[ASTNode]:
        return [self.action]

    @override
    def to_json(self) -> JSONValue:
        return {"default": self.action.to_json()}

    def viz_edge_label(self) -> str:
        return f"DEFAULT / {self.action.viz_label()}"

    @override
    def condition_label(self) -> str:
        return "default"

    @override
    def is_triggered_by(self, llm_output: str) -> bool:
        return True


class State(JSONMixin, ASTNode, ABC):
    """Base class for state machine states."""

    config: Config  # the config in the yaml, not necessarily the full config

    @abstractmethod
    def all_actions(self) -> list[ActionList]:
        """Returns all actions in the state."""

    def all_goto_state_labels(self) -> set[str]:
        """Returns all goto labels in the state."""
        s: set[str] = set()
        for action in self.all_actions():
            if isinstance(action.actions[-1], GotoAction):
                s.add(action.actions[-1].label)
            elif isinstance(action.actions[-1], SendAction):
                s.add(action.actions[-1].next_state)
        return s

    def all_send_agent_labels(self) -> set[str]:
        """Returns all send agent labels in the state."""
        s: set[str] = set()
        for action in self.all_actions():
            if isinstance(action.actions[-1], SendAction):
                s.add(action.actions[-1].to)
        return s


@dataclass(frozen=True)
class InitialState(State):
    """Represents the initial state of the state machine, without conditions."""

    action: ActionList
    config: Config  # the config in the yaml, not necessarily the full config

    @property
    def _children(self) -> Sequence[ASTNode]:
        return [self.action]

    def __post_init__(self) -> None:
        for action in self.action.actions:
            if isinstance(action, SendAction):
                raise ValueError(  # noqa: TRY004 (TypeError is not a good match)
                    "Send action not allowed in initial state."
                )
        if not isinstance(self.action.actions[-1], GotoAction):
            raise ValueError("Initial state must end in a goto action.")  # noqa: TRY004 (TypeError is not a good match)

    @override
    def all_actions(self) -> list[ActionList]:
        return [self.action]

    @override
    def to_json(self) -> JSONValue:
        return {"name": "initial", "action": self.action.to_json()}


@dataclass(frozen=True)
class NonInitialState(State):
    """Represents any other but the initial state in the state machine"""

    name: str
    conditions: list[Condition]
    config: Config  # the config in the yaml, not necessarily the full config

    @property
    def _children(self) -> Sequence[ASTNode]:
        return self.conditions

    def __post_init__(self) -> None:
        if self.name == "initial":
            raise ValueError('Attempt to create NonInitialState with name "initial".')
        if len(self.conditions) == 0:
            raise ValueError("State must contain at least one condition")
        for condition in self.conditions[:-1]:
            if isinstance(condition, DefaultCondition):
                raise ValueError(  # noqa: TRY004 (TypeError is not a good match)
                    "Only the last condition may be a default."
                )
        if not isinstance(self.conditions[-1], DefaultCondition):
            raise ValueError("The last condition must be a default.")  # noqa: TRY004 (TypeError is not a good match)
        if self.name in self.all_send_agent_labels():
            raise ValueError(f"State {self.name} sends to itself.")

    @override
    def all_actions(self) -> list[ActionList]:
        return [c.action for c in self.conditions]

    @override
    def to_json(self) -> JSONValue:
        return {
            "name": self.name,
            "conditions": [c.to_json() for c in self.conditions],
        }

    def triggered_condition(self, llm_output: str) -> Condition:
        for condition in self.conditions:
            if condition.is_triggered_by(llm_output):
                return condition
        # This should never happen since we are supposed to have a default condition
        raise ValueError(f"No condition triggered by {llm_output}")


@dataclass(frozen=True)
class Agent(JSONMixin, ASTNode):
    """Represents an agent."""

    name: str
    initial_state: InitialState
    noninitial_states: list[NonInitialState]
    _states_by_name: dict[str, NonInitialState] = field(init=False, repr=False)
    config: Config  # the config in the yaml, not necessarily the full config

    @property
    def _children(self) -> Sequence[ASTNode]:
        return [self.initial_state, *self.noninitial_states]

    def get_state(self, name: str) -> NonInitialState:
        """Returns the state with the given name."""
        return self._states_by_name[name]

    def __post_init__(self) -> None:
        object.__setattr__(self, "_states_by_name", {s.name: s for s in self.noninitial_states})

        # state names must be unique
        if len(self.noninitial_states) != len(self._states_by_name):
            raise ValueError("State names must be unique.")

        # no agent may send to itself
        if self.name in self.all_send_agent_labels():
            raise ValueError(f"Agent {self.name} sends to itself.")

    @override
    def to_json(self) -> JSONValue:
        return {
            "name": self.name,
            "states": [self.initial_state.to_json()] + [s.to_json() for s in self.noninitial_states],
        }

    def all_send_agent_labels(self) -> set[str]:
        """Returns all send agent labels in the agent."""
        s: set[str] = set()
        for state in self.noninitial_states:
            s.update(state.all_send_agent_labels())
        return s


@dataclass(frozen=True)
class Variables(JSONMixin):
    """Represents the variable definitions."""

    # variables are arbitrary identifiers mapped to strings. In yaml:
    # variables:
    #   foo: bar
    #   baz: quux
    vars: dict[str, str]

    @override
    def to_json(self) -> JSONValue:
        # (mypy does not support variance in mappings)
        return self.vars  # type: ignore[return-value]

    def __post_init__(self) -> None:
        for k, v in self.vars.items():
            s = _interpolate_vars(self.vars, v, frozenset([k]))
            if s.placeholders:
                raise ValueError(f"Undefined variables: {list(s.placeholders.values())}")
            self.vars[k] = s.text


@dataclass(frozen=True)
class StateMachine(ASTNode, JSONMixin):
    """Represents the conversation state machine."""

    variables: Variables
    agents: list[Agent]
    config: Config

    @property
    def _children(self) -> Sequence[ASTNode]:
        return self.agents

    def __post_init__(self) -> None:
        for agent in self.agents:
            for state in [agent.initial_state, *agent.noninitial_states]:
                for action_list in state.all_actions():
                    action_list._interpolate_vars(self.variables.vars)

        agent_names = [a.name for a in self.agents]
        if len(agent_names) != len(set(agent_names)):
            for name in agent_names:
                if agent_names.count(name) > 1:
                    raise ValueError(f"Duplicate agent name: {name}")
            assert False, "Unreachable"

        # unknown agents referenced?
        agent_references: set[str] = set()
        for agent in self.agents:
            agent_references.update(agent.all_send_agent_labels())
        for agent_reference in agent_references:
            if agent_reference not in agent_names:
                raise ValueError(f"Unknown agent: {agent_reference}")

        for child in self._children:
            child._link_configs(self.config)

    def get_agent(self, name: str) -> Agent:
        """Returns the agent with the given name."""
        for agent in self.agents:
            if agent.name == name:
                return agent
        raise ValueError(f"No agent named {name}")

    @override
    def to_json(self) -> JSONValue:
        return {
            "variables": self.variables.to_json(),
            "agents": [a.to_json() for a in self.agents],
        }

rrosti/chat/state_machine/ast.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""State machine execution."""

from __future__ import annotations

import argparse
import asyncio
import sys
import uuid
from enum import Enum

import marshmallow as ma
from loguru import logger
from overrides import override
from ruamel.yaml import YAML

from rrosti.chat import chat_session
from rrosti.chat.chat_session import Importance, Message
from rrosti.chat.state_machine import ast, interpolable, parsing
from rrosti.llm_api import openai_api, openai_api_direct
from rrosti.query import logging as qlog
from rrosti.servers.websocket_query_server import Frontend, UserInputMessage, WebFrontend
from rrosti.utils import misc
from rrosti.utils.config import config
from rrosti.utils.misc import ProgramArgsBase, ProgramArgsMixinProtocol, truncate_json_strings

yaml = YAML(typ=["rt", "string"])

_NO_PYTHON_CODE_BLOCK_FOUND_MSG = """Your message did not contain a Python code block.

It should look like this:

$$$python
print("Hello world!")
$$$

Please try again. Do not apologize."""

_NO_RTFM_CODE_BLOCK_FOUND_MSG = """Your message did not contain an rtfm code block.

It should look like this:

$$$rtfm
block contents here
$$$

Please try again. Do not apologize."""


_IMPORTANCE_USER_INPUT_INTERPOLABLE = Importance.MEDIUM
_TTL_USER_INPUT_INTERPOLABLE: int | None = None

_IMPORTANCE_PYTHON_INTERPOLABLE = Importance.LOW
_TTL_PYTHON_INTERPOLABLE: int | None = None

_IMPORTANCE_RTFM_INTERPOLABLE = Importance.NOISE
_TTL_RTFM_INTERPOLABLE: int | None = 2


class ProgramArgsMixin(ProgramArgsMixinProtocol):
    user_simulate_llm: bool

    @classmethod
    def _add_args(cls, parser: argparse.ArgumentParser) -> None:
        super()._add_args(parser)
        parser.add_argument(
            "--user-simulate-llm",
            action="store_true",
            help="Instead of using an actual LLM, ask for user input in the console.",
        )


class _Python(interpolable.Interpolable):
    """Implementation of {python()}."""

    _show_output = "Python Output:\n\n"
    _python_vars: dict[str, object]

    def __init__(self, python_vars: dict[str, object]) -> None:
        self._python_vars = python_vars

    @override
    def is_for_me(self, code: str) -> bool:
        return code == "python()"

    # TODO: make async
    @override
    async def execute(self, last_msg: str | None, frontend: Frontend) -> interpolable.InterpolableOutput:
        if last_msg is None:
            logger.error("python() executed without a message in context")
            assert False
        try:
            python_out = interpolable.execute_python_in_msg(last_msg, self._python_vars)
        except interpolable.NoCodeBlockFoundError:
            return interpolable.InterpolableOutput(
                output=_NO_PYTHON_CODE_BLOCK_FOUND_MSG,
                info_prefix="ERROR:\n\n",
                importance=_IMPORTANCE_PYTHON_INTERPOLABLE,
                ttl=_TTL_PYTHON_INTERPOLABLE,
            )
        else:
            frontend.handle_python_output(python_out)
            return interpolable.InterpolableOutput(
                output=python_out._output,
                info_prefix=self._show_output,
                importance=_IMPORTANCE_PYTHON_INTERPOLABLE,
                ttl=_TTL_PYTHON_INTERPOLABLE,
            )


class _Rtfm(interpolable.Interpolable):
    """Implementation of {rtfm()}."""

    _show_output = "RTFM output:\n\n"
    _openai_provider: openai_api.OpenAIApiProvider

    def __init__(self, openai_provider: openai_api.OpenAIApiProvider) -> None:
        self._openai_provider = openai_provider

    @override
    def is_for_me(self, code: str) -> bool:
        return code == "rtfm()"

    @override
    async def execute(self, last_msg: str | None, frontend: Frontend) -> interpolable.InterpolableOutput:
        if last_msg is None:
            logger.error("rtfm() executed without a message in context")
            assert False
        try:
            snippets = await interpolable.execute_rtfm_in_msg(self._openai_provider, last_msg)
        except interpolable.NoCodeBlockFoundError:
            return interpolable.InterpolableOutput(
                output=_NO_RTFM_CODE_BLOCK_FOUND_MSG,
                info_prefix="ERROR:\n\n",
                importance=_IMPORTANCE_RTFM_INTERPOLABLE,
                ttl=_TTL_RTFM_INTERPOLABLE,
            )
        else:
            start_index = frontend.handle_rtfm_output(snippets)
            output_str = "\n-----\n".join(
                f"Extract #{i + start_index}:\n\n{snip.text}" for i, snip in enumerate(snippets)
            ).strip()
            return interpolable.InterpolableOutput(
                output=output_str,
                info_prefix=self._show_output,
                importance=_IMPORTANCE_RTFM_INTERPOLABLE,
                ttl=_TTL_RTFM_INTERPOLABLE,
            )


class _ActionHandler(ast.AsyncActionVisitor):
    runner: _AgentRunner
    config: ast.Config | None

    def __init__(self, runner: _AgentRunner) -> None:
        self.runner = runner
        self.config = None

    @override
    async def message_action(self, action: ast.MessageAction) -> None:
        assert self.config is not None
        logger.info(
            "[Agent {}]: Message action ({}): {}",
            self.runner.name,
            action.role,
            truncate_json_strings(action.text),
        )

        text = action.text

        importances: set[Importance] = set()
        ttls: set[int | None] = set()

        # handle function invocations
        for placeholder, code in action.placeholders.items():
            logger.info("[Agent {}]: Handling function invocation: {}", self.runner.name, code)
            interpolable_output: interpolable.InterpolableOutput
            for interp in self.runner._interpolables:
                if interp.is_for_me(code):
                    interpolable_output = await interp.execute(
                        self.runner._session.messages[-1].text,
                        frontend=self.runner._sm_runner.frontend,
                    )
                    if interpolable_output.info_prefix:
                        await self.runner._sm_runner.frontend.send_message(
                            interpolable_output.info_prefix + interpolable_output.output
                        )
                    break
            else:
                raise ValueError(f"Don't know what to do with interpolation: {code}")

            # We'll probably give the state machine writer the best tools if we strip whitespace
            # from the output, whatever it is.
            replacement = interpolable_output.output.strip()
            text = text.replace(placeholder, replacement)
            importances.add(interpolable_output.importance)
            ttls.add(interpolable_output.ttl)

        if config.state_machine.debug_detect_unresolved_funcalls:
            assert "FUNCALL(" not in text, "Unresolved function call(s)"

        ttl = None
        if not importances:
            # This means there were no function invocations.
            # As a heuristic, we assume that user and system messages added in the initial state
            # are of high importance, and messages added in other states are of medium importance.
            if isinstance(self.runner._curr_state, ast.InitialState):
                importance = Importance.HIGH
            else:
                importance = Importance.MEDIUM
        else:
            # Otherwise, if there's more than one of either, warn and take the min.
            if len(importances) > 1:
                logger.warning(
                    "Multiple importances for message action: {} (taking min)",
                    importances,
                )
            importance = min(importances)
            if len(ttls) > 1:
                logger.warning(
                    "Multiple ttls for message action: {} (taking None)",
                    ttls,
                )
            else:
                ttl = next(iter(ttls))

        self.runner._session.add_message(
            chat_session.Message(role=action.role, text=text, ttl=ttl, importance=importance)
        )

    @override
    async def goto_action(self, action: ast.GotoAction) -> None:
        assert self.config is not None
        logger.info("[Agent {}]: Goto action: {}", self.runner.name, action.label)
        self.runner._curr_state = self.runner._agent.get_state(action.label)

    @override
    async def end_action(self, action: ast.EndAction) -> None:
        assert self.config is not None
        logger.info("[Agent {}]: End action", self.runner.name)
        self.runner._runner_state = RunnerState.TERMINATED

    @override
    async def send_action(self, action: ast.SendAction) -> None:
        assert self.config is not None
        # FIXME: make sure send is the last action.
        logger.info("[Agent {}]: Send action: to: {}; next: {}", self.runner.name, action.to, action.next_state)
        self.runner._runner_state = RunnerState.WAITING
        await self.runner._sm_runner._get_agent_runner(action.to).add_message(
            self.runner._session.messages[-1].as_user_message(model=self.config.model), quiet=True
        )


class RunnerState(Enum):
    """
    The state of the agent runner.

    The runner starts in the NOT_STARTED state. The first agent's runner is started by
    StateMachineRunner, putting it in the RUNNING state.

    Apart from this, the only state changes are:
    - Sending a message causes a RUNNING -> WAITING transition.
    - Receiving a message causes a NOT_STARTED/WAITING -> RUNNING transition.
    - Excuting an `end` action causes a RUNNING -> TERMINATED transition.
    """

    NOT_STARTED = "not_started"
    RUNNING = "running"
    WAITING = "waiting"
    TERMINATED = "terminated"


class _AgentRunner:
    """
    The agents are sort-of independent actors that run like coroutines.

    Do not use directly. Use StateMachineRunner.
    """

    _session: chat_session.ChatSession
    _agent: ast.Agent
    _curr_state: ast.State
    _sm_runner: StateMachineRunner
    _inbox: asyncio.Queue[chat_session.Message]
    __runner_state: RunnerState
    _handler: _ActionHandler
    _interpolables: list[interpolable.Interpolable]

    @property
    def _runner_state(self) -> RunnerState:
        return self.__runner_state

    @_runner_state.setter
    def _runner_state(self, value: RunnerState) -> None:
        logger.info("[Agent {}]: State transition: {} -> {}", self.name, self._runner_state, value)
        self.__runner_state = value

    @property
    def _sm(self) -> ast.StateMachine:
        return self._sm_runner._sm

    @property
    def name(self) -> str:
        return self._agent.name

    def __init__(
        self,
        sm_runner: StateMachineRunner,
        agent: ast.Agent,
        llm: chat_session.LLM,
        python_vars: dict[str, object] | None = None,
    ) -> None:
        self._sm_runner = sm_runner
        self._inbox = asyncio.Queue()
        self._agent = agent
        self._session = chat_session.ChatSession(llm, name=self._agent.name, callback=self._sm_runner._message_callback)
        self._curr_state = self._agent.initial_state
        self._handler = _ActionHandler(self)
        self.__runner_state = RunnerState.NOT_STARTED

        async def _get_user_input() -> str:
            # Expire messages with expired ttls.
            self._session.decrease_ttls()
            return (await self._sm_runner.frontend.get_user_input()).content

        self._interpolables = [
            interpolable.SimpleInterpolable(
                code="user_input()",
                importance=_IMPORTANCE_USER_INPUT_INTERPOLABLE,
                ttl=_TTL_USER_INPUT_INTERPOLABLE,
                coro=_get_user_input,
            ),
            _Python(python_vars or {}),
            _Rtfm(self._sm_runner._openai_provider),
        ]

    async def _execute_action(self, action: ast.ActionBase) -> None:
        """Execute an action."""

        assert self._runner_state == RunnerState.RUNNING, self._runner_state
        await action.aaccept(self._handler)

    async def _execute_actions(self, actions: ast.ActionList, config: ast.Config) -> None:
        self._handler.config = config
        await actions.aaccept(self._handler)

    async def start(self) -> None:
        """Start the agent's state machine."""

        assert self._runner_state == RunnerState.NOT_STARTED, self._runner_state
        assert isinstance(self._curr_state, ast.InitialState)
        self._runner_state = RunnerState.RUNNING
        # For the first step, we need to get out of the initial state.
        # For it, we just execute actions. No LLM is invoked here.
        # The initial state is used to set up the environment for the first LLM invocation.
        logger.info("[Agent {}]: Starting.", self._agent.name)
        await self._execute_actions(self._curr_state.action, self._curr_state.config)
        assert isinstance(self._curr_state, ast.NonInitialState), "Did not get out of initial state"
        # self._process_inbox()  # add any messages from inbox

    def _add_message_to_session(self, message: chat_session.Message) -> None:
        """Add a message to the agent's session."""
        assert self._runner_state not in (RunnerState.NOT_STARTED, RunnerState.TERMINATED), self._runner_state
        self._session.add_message(message)
        logger.info("[Agent {}]: Received message: {}", self._agent.name, message)

    def _process_inbox(self) -> None:
        """Process the agent's inbox, adding to the session."""

        if self._inbox.empty():
            return
        assert self._inbox.qsize() == 1, "Sus: Several messages in inbox."
        self._add_message_to_session(self._inbox.get_nowait())

    async def add_message(self, message: chat_session.Message, quiet: bool = False) -> None:
        """Add a message to the agent's inbox."""
        logger.info("[Agent {}]: Adding message to inbox: {}", self._agent.name, message)
        if self._runner_state == RunnerState.NOT_STARTED:
            await self.start()
        elif self._runner_state == RunnerState.WAITING:
            self._runner_state = RunnerState.RUNNING
        self._session.add_message(message, quiet=quiet)

    async def step(self) -> None:
        """Execute one step of the agent's state machine."""

        assert self._runner_state != RunnerState.TERMINATED, "Cannot step after termination"
        if self._runner_state in (RunnerState.NOT_STARTED, RunnerState.WAITING):
            self._add_message_to_session(await self._inbox.get())
        assert self._runner_state == RunnerState.RUNNING, self._runner_state
        assert isinstance(self._curr_state, ast.NonInitialState)
        logger.info("[Agent {}]: Non-initial state: {}", self._agent.name, self._curr_state.name)

        self._process_inbox()  # add any messages from inbox

        msg = await self._session.generate_response(
            model=(self._session.messages[-1].model if self._session.messages else None)
            or self._curr_state.config.model
            or self._agent.config.model
            or self._sm.config.model
        )
        await self._sm_runner.frontend.send_message(msg)
        cond = self._curr_state.triggered_condition(msg.text)
        await self._execute_actions(cond.action, cond.config)

    async def run(self) -> None:
        """Run the agent's state machine until it terminates."""
        assert self._runner_state != RunnerState.TERMINATED, "Cannot run after termination"
        while self._runner_state != RunnerState.TERMINATED:  # type: ignore[comparison-overlap] # (mypy 1.4.1 bug)
            await self.step()

    def __repr__(self) -> str:
        return f"<_AgentRunner {self._agent.name}>"


class StateMachineRunner:
    """Execute a network of agents (essentially a hierarchical state machine)."""

    _sm: ast.StateMachine
    _llm: chat_session.LLM
    _python_vars: dict[str, object]
    _agent_runners: list[_AgentRunner]
    frontend: Frontend
    _message_callback: chat_session.MessageCallback | None
    _openai_provider: openai_api.OpenAIApiProvider

    @property
    def _is_terminated(self) -> bool:
        return any(runner._runner_state == RunnerState.TERMINATED for runner in self._agent_runners)

    @property
    def _running_agent(self) -> _AgentRunner:
        assert not self._is_terminated, "Cannot get running agents after termination"
        rs = [runner for runner in self._agent_runners if runner._runner_state == RunnerState.RUNNING]
        assert len(rs) == 1, f"Exactly one agent should be running, got {len(rs)}: {rs}"
        return rs[0]

    def __init__(
        self,
        *,
        sm: ast.StateMachine,
        llm: chat_session.LLM,
        frontend: Frontend,
        openai_provider: openai_api.OpenAIApiProvider,
        python_vars: dict[str, object] | None = None,
        message_callback: chat_session.MessageCallback | None = None,
    ) -> None:
        """
        Args:
            sm: The state machine to execute.
            llm: The LLM to use.
            frontend: an Object that defines and handles communication with the Frontend
            openai_provider: The OpenAI provider to use.
            python_vars: The Python variables to use.
            message_callback: A callback to call when a message is sent or received.
        """
        self._openai_provider = openai_provider
        self._sm = sm
        self._llm = llm
        self._python_vars = python_vars or {}
        self._message_callback = message_callback
        self._agent_runners = [_AgentRunner(self, agent, self._llm, self._python_vars) for agent in sm.agents]
        self.frontend = frontend

    def _get_agent_runner(self, name: str) -> _AgentRunner:
        """Get the agent runner with the given name."""
        for runner in self._agent_runners:
            if runner.name == name:
                return runner
        raise ValueError(f"Agent {name} not found")

    async def step(self) -> None:
        """Execute one step of the state machine."""

        # The first agent is the one that starts the state machine.
        if self._agent_runners[0]._runner_state == RunnerState.NOT_STARTED:
            await self._agent_runners[0].start()

        assert not self._is_terminated, "Cannot step after termination"
        await self._running_agent.step()

    async def run(self) -> None:
        """Run the state machine until it terminates."""
        assert not self._is_terminated, "Cannot run after termination"
        while not self._is_terminated:
            await self.step()


class MessageObserver:
    total_cost = 0.0

    def __call__(self, message: chat_session.Message, agent: str | None, quiet: bool) -> None:
        if not quiet:
            print(message.to_string(agent=agent))
        if message.cost:
            self.total_cost += message.cost
            logger.info("Total cost so far: {:.5f}", self.total_cost)


@misc.async_once_in_thread
def _load_state_machine() -> tuple[str, ast.StateMachine]:
    yaml_data = config.state_machine.yaml_path.read_text()
    sm = parsing.loads_from_yaml(yaml_data)
    yaml.indent(mapping=4, sequence=4, offset=0)
    logger.info(
        "Loaded state machine:\n{}",
        yaml.dump_to_string(truncate_json_strings(sm.to_json())),  # type: ignore[attr-defined]
    )
    return yaml_data, sm


async def ensure_loaded() -> None:
    await _load_state_machine()


async def load_and_run(
    *,
    openai_provider: openai_api.OpenAIApiProvider,
    llm: chat_session.LLM,
    frontend: Frontend,
) -> None:
    yaml_text, sm = await _load_state_machine()

    # Get the first user input as the query, just for logging purposes.
    query = await frontend.peek_user_input()
    logger.info("User conversation starter: {}", query)

    with qlog.QueryEvent.section(text=query.content, uuid=query.uuid, prompt=yaml_text):
        runner = StateMachineRunner(
            sm=sm,
            llm=llm,
            python_vars=dict(),
            message_callback=MessageObserver(),
            frontend=frontend,
            openai_provider=openai_provider,
        )

        await runner.run()


class ProgramArgs(ProgramArgsBase, ProgramArgsMixin):
    @classmethod
    def _add_args(cls, parser: argparse.ArgumentParser) -> None:
        super()._add_args(parser)


class FakeWebFrontend(WebFrontend):
    def __init__(self) -> None:
        super().__init__(None, False)

    @override
    async def send_message(self, msg: Message | str) -> None:
        # TODO: implement console version, if needed
        pass

    @override
    async def _get_user_input_impl(self) -> UserInputMessage:
        # TODO: implement console version, if needed
        return UserInputMessage(content="", uuid=str(uuid.uuid4()))


async def main() -> None:
    misc.setup_logging()
    openai_provider = openai_api_direct.DirectOpenAIApiProvider()
    args = ProgramArgs.parse_args()
    llm: chat_session.LLM
    if args.user_simulate_llm:
        llm = chat_session.UserInputLLM()
    else:
        llm = chat_session.OpenAI(openai_provider)
    try:
        await load_and_run(llm=llm, frontend=FakeWebFrontend(), openai_provider=openai_provider)
    except ma.exceptions.ValidationError as e:
        print(e.messages, sys.stderr)
        raise


if __name__ == "__main__":
    asyncio.run(main())

rrosti/chat/state_machine/execution.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
Contains the logic for interpolables, i.e. functions that can be called from within a message, like
{user_input()} or {python()}.

The basic idea is that a user can enter a message like

    ```
    {user_input()}
    ```

The message is then parsed and the code block is extracted. In this case, the code block is

    {user_input()}

The code block is then checked against all interpolables to see which one is responsible for it. In this case, the
`UserInputInterpolable` would be responsible. The interpolable is then executed and the output is returned. In this
case, the user would be asked for input.
"""

import asyncio
import io
import re
from abc import ABC, abstractmethod
from contextlib import redirect_stdout
from typing import Awaitable, Callable

import attrs
import pandas as pd
from overrides import override

from rrosti.chat.chat_session import Importance
from rrosti.llm_api import openai_api
from rrosti.servers.websocket_query_server import Frontend, PythonItem
from rrosti.snippets import document_sync
from rrosti.snippets.abstract_snippet_database import AbstractSnippetDatabase
from rrosti.snippets.sklearn_snippet_database import SklearnSnippetDatabase
from rrosti.snippets.snippet import Snippet
from rrosti.utils import misc
from rrosti.utils.config import config


class NoCodeBlockFoundError(ValueError):
    pass


@attrs.frozen
class InterpolableOutput:
    output: str

    # For debug purposes, shown in frontend. If None, then the actual output is not shown at all.
    # If it should be shown but without prefix, simply use ''
    info_prefix: str | None

    importance: Importance  # Influences pruning order

    # None = do not prune; positive integer = prune after this many user inputs
    ttl: int | None = attrs.field()

    @ttl.validator
    def _check_ttl(self, _attribute: str, value: int | None) -> None:
        if value is not None and value <= 0:
            raise ValueError("ttl must be None or a positive integer")


# TODO: Is this name descriptive? Can we think of something better?
class Interpolable(ABC):
    """An interpolable function invocation, like {user_input()}, {python()} or {rtfm()}."""

    @abstractmethod
    def is_for_me(self, code: str) -> bool:
        ...

    @abstractmethod
    async def execute(self, last_msg: str | None, frontend: Frontend) -> InterpolableOutput:
        ...


# A simple interpolable that wraps a coroutine
class SimpleInterpolable(Interpolable):
    _code: str
    _coro: Callable[[], Awaitable[str]]
    _importance: Importance
    _ttl: int | None

    def __init__(self, code: str, importance: Importance, ttl: int | None, coro: Callable[[], Awaitable[str]]) -> None:
        self._code = code
        self._coro = coro
        self._importance = importance
        self._ttl = ttl

    @override
    def is_for_me(self, code: str) -> bool:
        return code == self._code

    @override
    async def execute(self, last_msg: str | None, frontend: Frontend) -> InterpolableOutput:
        return InterpolableOutput(
            output=await self._coro(), info_prefix=None, importance=self._importance, ttl=self._ttl
        )


def extract_code_block(language: str, msg_text: str) -> str:
    # First check that we don't have multiple matches for ```{language}
    if msg_text.count(f"```{language}") > 1:
        raise ValueError("Multiple code blocks found in message.")

    code_block_pattern = re.compile(r"\$\$\$" + language + r"\r?\n(.*?)\$\$\$", re.DOTALL)
    match = code_block_pattern.search(msg_text)
    if not match:
        raise NoCodeBlockFoundError("No full and terminated code block found in message")
    return match.group(1)


def execute_python_in_msg(msg_text: str, vars: dict[str, object]) -> PythonItem:
    code = extract_code_block("python", msg_text)

    try:
        f = io.StringIO()
        with redirect_stdout(f):
            exec(code, dict(pd=pd, **vars))
        out = f.getvalue().strip()
    except Exception as e:
        out = "Exception: " + str(e)
    return PythonItem(_code=code, _output=out)


@misc.async_once_blocking
async def _get_database() -> AbstractSnippetDatabase:
    snippets = await document_sync.sync_and_get_snippets()
    return await asyncio.to_thread(SklearnSnippetDatabase, snippets)


async def ensure_loaded() -> None:
    await _get_database()


async def execute_rtfm_in_msg(openai_provider: openai_api.OpenAIApiProvider, message: str) -> list[Snippet]:
    text = extract_code_block("rtfm", message)

    # We need the encoding for the text
    snippet = Snippet.from_query(text)
    await snippet.async_ensure_embedding(openai_provider)

    # Now execute the nearest neighbors search
    db = await _get_database()
    return await db.find_nearest_merged(
        openai_provider, snippet, config.state_machine.rtfm_max_tokens, config.state_machine.rtfm_merge_candidates
    )

rrosti/chat/state_machine/interpolable.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Read a description of a dialog state machine from a .yaml file and constructs an AST."""

from __future__ import annotations

from pathlib import Path
from typing import Any, Mapping, TextIO, cast

import jsonschema
import marshmallow as ma
from loguru import logger
from overrides import override
from ruamel.yaml import YAML

from rrosti.chat.state_machine import ast

yaml = YAML(typ=["rt", "string"])

# We don't really use this for anything but validating, but it helps us give better error messages.
_SCHEMA = {
    "$schema": "http://json-schema.org/draft-07/schema#",
    "type": "object",
    "additionalProperties": False,
    "properties": {
        "variables": {
            "type": "object",
            "additionalProperties": {"type": "string"},  # Variables are arbitrary string values
        },
        "agents": {
            "type": "array",
            "items": {"$ref": "#/definitions/agent"},
            "minItems": 1,
        },
        "config": {
            "model": {"type": "string"},
        },
        "required": ["agents"],
    },
    "required": ["agents"],
    "definitions": {
        "initial_state": {
            "type": "object",
            "additionalProperties": False,
            "properties": {
                "name": {"const": "initial"},
                "action": {
                    "$ref": "#/definitions/action_list",
                },
                "model": {"type": "string"},
            },
            "required": ["name", "action"],
        },
        "noninitial_state": {
            "type": "object",
            "additionalProperties": False,
            "properties": {
                "name": {"type": "string", "not": {"const": "initial"}},
                "conditions": {
                    "type": "array",
                    "items": {"$ref": "#/definitions/condition"},
                },
                "model": {"type": "string"},
            },
            "required": ["name", "conditions"],
        },
        "states_list": {
            "type": "array",
            "minItems": 1,  # Must have at least one state
            "prefixItems": [
                {"$ref": "#/definitions/initial_state"},
            ],
            "items": {
                "$ref": "#/definitions/noninitial_state",
            },
        },
        "agent": {
            "type": "object",
            "additionalProperties": False,
            "properties": {
                "name": {"type": "string"},
                "states": {
                    "$ref": "#/definitions/states_list",
                },
                "model": {"type": "string"},
            },
            "required": ["name", "states"],
        },
        "atomic_actions": {
            "end": {"const": "end"},
            "goto": {
                "type": "object",
                "additionalProperties": False,
                "properties": {
                    "goto": {"type": "string"},
                },
                "required": ["goto"],
            },
            "message": {
                "type": "object",
                "additionalProperties": False,
                "properties": {
                    "message": {"type": "string"},
                },
                "required": ["message"],
            },
            "assistant_message": {
                "type": "object",
                "additionalProperties": False,
                "properties": {
                    "assistant_message": {"type": "string"},
                },
                "required": ["assistant_message"],
            },
            "system_message": {
                "type": "object",
                "additionalProperties": False,
                "properties": {
                    "system_message": {"type": "string"},
                },
                "required": ["system_message"],
            },
            "send": {
                "type": "object",
                "additionalProperties": False,
                "properties": {
                    "send": {
                        "type": "object",
                        "additionalProperties": False,
                        "properties": {
                            "to": {"type": "string"},
                            "next_state": {"type": "string"},
                        },
                    },
                },
                "required": ["send"],
            },
        },
        "action": {
            "oneOf": [
                {"$ref": "#/definitions/atomic_actions/end"},
                {"$ref": "#/definitions/atomic_actions/goto"},
                {"$ref": "#/definitions/atomic_actions/message"},
                {"$ref": "#/definitions/atomic_actions/assistant_message"},
                {"$ref": "#/definitions/atomic_actions/system_message"},
                {"$ref": "#/definitions/atomic_actions/send"},
            ],
        },
        "action_list": {
            "type": "array",
            "items": {"$ref": "#/definitions/action"},
            "minItems": 1,
        },
        "if_cond_expr": {
            "type": "object",
            "additionalProperties": False,
            "properties": {
                "contains": {"type": "string"},
            },
            "required": ["contains"],
        },
        "then_object": {
            "type": "object",
            "additionalProperties": False,
            "properties": {
                "model": {"type": "string"},
                "action": {"$ref": "#/definitions/action_list"},
            },
            "required": ["action"],
        },
        "if_then": {
            "type": "object",
            "additionalProperties": False,
            "properties": {
                "if": {"$ref": "#/definitions/if_cond_expr"},
                "then": {"$ref": "#/definitions/then_object"},
            },
            "required": ["if", "then"],
        },
        "default": {
            "type": "object",
            "additionalProperties": False,
            "properties": {
                "default": {"$ref": "#/definitions/then_object"},
            },
            "required": ["default"],
        },
        "condition": {
            "oneOf": [
                {"$ref": "#/definitions/default"},
                {"$ref": "#/definitions/if_then"},
            ],
        },
    },
}


def _parse_action_one_param(dic: Mapping[str, Any]) -> ast.ActionBase:
    assert len(dic) == 1, f"Invalid action: {dic}"
    key, value = next(iter(dic.items()))
    assert isinstance(value, str), f"Invalid action: {dic}"
    assert isinstance(key, str), f"Invalid action: {dic}"
    value = value.strip()
    match key:
        case "goto":
            return ast.GotoAction(value)
        case "message":
            return ast.MessageAction("user", value)
        case "assistant_message":
            return ast.MessageAction("assistant", value)
        case "system_message":
            return ast.MessageAction("system", value)
        case _:
            raise ma.ValidationError(f"Invalid action: {dic}")


def _parse_action(value: Any) -> ast.ActionBase:
    match value:
        case "end":
            return ast.EndAction()
        case str():
            raise ma.ValidationError(f"Invalid action: {value}")
        case {"send": dict() as send}:
            return ast.SendAction(send["to"], send["next_state"])
        case dict():
            return _parse_action_one_param(value)
        case _:
            raise ma.ValidationError(f"Invalid action: {value}")


class ActionListField(ma.fields.Field):
    """
    Parse a list of actions of the format
    ["end", {"goto": "label"}, {"message": "text"}, {"assistant_message": "text"}, ...]
    """

    @override
    def _deserialize(
        self, value: Any, attr: str | None, data: Mapping[str, Any] | None, **kwargs: Any
    ) -> ast.ActionList:
        logger.info("ActionListField: value={}, attr={}, data={}, kwargs={}", value, attr, data, kwargs)
        if not isinstance(value, list):
            raise ma.ValidationError(f"Invalid action: {value}")
        try:
            return ast.ActionList([_parse_action(v) for v in value])
        except ValueError as e:
            raise ma.ValidationError(*e.args) from e


class ConfigSchema(ma.Schema):
    model = ma.fields.String(required=False)

    class Meta:
        unknown = ma.EXCLUDE

    @ma.post_load
    def make_config(self, data: dict[str, Any], **kwargs: Any) -> ast.Config:
        logger.info("make_config: data={}, kwargs={}", data, kwargs)
        return ast.Config(model_ast=data.get("model", None))


class IfSchema(ma.Schema):
    contains = ma.fields.String(required=True)

    class Meta:
        unknown = ma.EXCLUDE


class ThenSchema(ma.Schema):
    action = ActionListField(required=True)
    model = ma.fields.String(required=False)

    class Meta:
        unknown = ma.EXCLUDE


class IfThenSchema(ma.Schema):
    if_ = ma.fields.Nested(IfSchema, required=True, data_key="if")
    then = ma.fields.Nested(ThenSchema, required=True)

    class Meta:
        unknown = ma.EXCLUDE

    @ma.post_load
    def make_if_then(self, data: dict[str, Any], **kwargs: Any) -> ast.Condition:
        logger.info("make_if_then: data={}, kwargs={}", data, kwargs)
        return ast.Condition(
            contains=data["if"]["contains"],
            action=data["then"]["action"],
            config=ast.Config(model_ast=data["then"].get("model")),
        )


class DefaultSchema(ma.Schema):
    # default = ma.fields.Nested(ThenSchema, required=True)
    action = ActionListField(required=True)
    model = ma.fields.String(required=False)

    class Meta:
        unknown = ma.EXCLUDE

    @ma.post_load
    def make_default(self, data: dict[str, Any], **kwargs: Any) -> ast.DefaultCondition:
        logger.info("make_default: data={}, kwargs={}", data, kwargs)
        return ast.DefaultCondition(action=data["action"], config=ast.Config(model_ast=data.get("model")))


class ConditionSchema(ma.Schema):
    # We have either:
    # if:
    #     contains: "foo"
    # then:
    #     action: ["end", {"goto": "label"}, ...]
    #     model: gpt-4   # optional
    #
    # OR:
    #
    # default:
    #     action: ["end", {"goto": "label"}, ...]
    #     model: gpt-4   # optional
    if_ = ma.fields.Nested(IfSchema, required=False, data_key="if")
    then = ma.fields.Nested(ThenSchema, required=False)
    default = ma.fields.Nested(DefaultSchema, required=False)

    class Meta:
        unknown = ma.EXCLUDE

    @ma.post_load
    def make_condition(self, data: dict[str, Any], **kwargs: Any) -> ast.Condition | ast.DefaultCondition:
        logger.info("make_condition: data={}, kwargs={}", data, kwargs)
        if "if_" in data:
            assert "default" not in data, "Cannot have both if and default"
            assert "then" in data, "if without then"
            return ast.Condition(
                contains=data["if_"]["contains"],
                action=data["then"]["action"],
                config=ast.Config(model_ast=data["then"].get("model")),
            )
        assert data.get("default"), "Must have either if or default"
        assert "then" not in data, "default cannot have then"
        return data["default"]  # type: ignore[no-any-return]


class StateSchema(ma.Schema):
    name = ma.fields.String(required=True)
    model = ma.fields.String(required=False)

    # Conditions are mandatory for non-initial states and forbidden for initial states.
    conditions = ma.fields.List(ma.fields.Nested(ConditionSchema), required=False)
    # The opposite applies to action.
    action = ActionListField(required=False)

    class Meta:
        unknown = ma.EXCLUDE

    @ma.post_load
    def make_state(self, data: dict[str, Any], **kwargs: Any) -> ast.InitialState | ast.NonInitialState:
        logger.info("make_state: data={}, kwargs={}", data, kwargs)
        if data["name"] == "initial":
            assert "conditions" not in data, "Initial state cannot have conditions"
            assert "action" in data, "Initial state must have action"
            try:
                return ast.InitialState(action=data["action"], config=ast.Config(model_ast=data.get("model")))
            except ValueError as e:
                raise ma.ValidationError(*e.args) from e
        assert "conditions" in data, "Non-initial state must have conditions"
        assert "action" not in data, "Non-initial state cannot have action"
        try:
            return ast.NonInitialState(
                name=data["name"], conditions=data["conditions"], config=ast.Config(model_ast=data.get("model"))
            )
        except ValueError as e:
            raise ma.ValidationError(*e.args) from e


def _validate_states(states: list[ast.State]) -> None:
    """Validate that all state names are unique."""
    noninitials: list[ast.NonInitialState]

    match states:
        case []:
            raise ma.ValidationError("There must be at least one state")
        case [ast.InitialState(), *noninitials_]:
            if not all(isinstance(s, ast.NonInitialState) for s in noninitials_):
                raise ma.ValidationError("Only the first state can be initial")
            noninitials = cast(list[ast.NonInitialState], noninitials_)
        case [ast.NonInitialState(), *_]:
            raise ma.ValidationError("The first state must be initial")
        case _:
            assert False, "Unreachable"

    names = [s.name for s in noninitials]
    if any(s == "initial" for s in names):
        raise ma.ValidationError("State name 'initial' is reserved")
    if len(set(names)) != len(names):
        dups = {name for name in names if names.count(name) > 1}
        raise ma.ValidationError(f"Duplicate state name: {dups if len(dups) != 1 else dups.pop()}")

    # all gotoed states must exist
    for state in states:
        for label in state.all_goto_state_labels():
            if label not in names:
                raise ma.ValidationError(f"State '{label}' does not exist")

    # all states must be reachable
    reachable = {"initial"}
    reachable.update(states[0].all_goto_state_labels())
    while True:
        new_reachable = reachable.copy()
        for state in noninitials:
            if state.name in reachable:
                new_reachable.update(state.all_goto_state_labels())
        if new_reachable == reachable:
            break
        reachable = new_reachable
    unreachable = set(names) - reachable
    if unreachable:
        raise ma.ValidationError(f"States {unreachable} are unreachable")


class AgentSchema(ma.Schema):
    name = ma.fields.String(required=True)
    states = ma.fields.List(ma.fields.Nested(StateSchema), required=True, validate=_validate_states)
    model = ma.fields.String(required=False)

    class Meta:
        unknown = ma.EXCLUDE

    @ma.post_load
    def make_agent(self, data: dict[str, Any], **kwargs: Any) -> ast.Agent:
        logger.info("make_agent: data={}, kwargs={}", data, kwargs)
        states = data["states"]
        assert len(states) >= 1, "There must be at least one state"
        try:
            return ast.Agent(
                name=data["name"],
                initial_state=states[0],
                noninitial_states=states[1:],
                config=ast.Config(model_ast=data.get("model")),
            )
        except ValueError as e:
            raise ma.ValidationError(*e.args) from e


class StateMachineSchema(ma.Schema):
    variables = ma.fields.Dict(keys=ma.fields.String(), values=ma.fields.String(), required=False)
    agents = ma.fields.List(ma.fields.Nested(AgentSchema), required=True)
    config = ma.fields.Nested(ConfigSchema, required=False)

    class Meta:
        unknown = ma.EXCLUDE

    # Use jsonschema to validate the dict
    @ma.pre_load
    def validate(self, data: dict[str, Any], **kwargs: Any) -> dict[str, Any]:
        v = jsonschema.Draft202012Validator(_SCHEMA)
        errors = sorted(v.iter_errors(data), key=lambda e: e.path)
        if errors:
            for e in errors:
                logger.error("Validation error:\n{}", e)
            raise ma.ValidationError("\n-------\n".join(str(e) for e in errors))
        return data

    @ma.post_load
    def make_state_machine(self, data: dict[str, Any], **kwargs: Any) -> ast.StateMachine:
        variables = data.get("variables", {})
        variables = {k: v.strip() for k, v in variables.items()}
        config = data.get("config", ast.Config(model_ast=None))

        try:
            return ast.StateMachine(variables=ast.Variables(variables), agents=data["agents"], config=config)
        except ValueError as e:
            raise ma.ValidationError(*e.args) from e


state_machine_schema = StateMachineSchema()


def loads_from_yaml(s: str) -> ast.StateMachine:
    """Load a state machine from a YAML string."""
    d = yaml.load(s)
    sm: ast.StateMachine = state_machine_schema.load(d)
    return sm


def load_from_yaml(f: TextIO) -> ast.StateMachine:
    """Load a state machine from a YAML file."""
    d = yaml.load(f)
    sm: ast.StateMachine = state_machine_schema.load(d)
    return sm


def load_from_yaml_file(path: Path) -> ast.StateMachine:
    """Load a state machine from a YAML file."""
    with open(path) as f:
        return load_from_yaml(f)

rrosti/chat/state_machine/parsing.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from flask import Flask
from jinja2 import Environment, PackageLoader, select_autoescape
from werkzeug.middleware.proxy_fix import ProxyFix

import rrosti.utils.config
from rrosti.utils.config import config

app = Flask(__name__)
app.wsgi_app = ProxyFix(app.wsgi_app, x_for=1, x_proto=1, x_host=1, x_prefix=1)  # type: ignore[method-assign]


@app.route("/")
def serve_client() -> str:
    rrosti.utils.config.load()
    env = Environment(loader=PackageLoader("rrosti.frontend.client"), autoescape=select_autoescape())

    template = env.get_template("index.html")

    return template.render(title=config.frontend.title, websocketUrl=config.frontend.websocket_url)

rrosti/frontend/client.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


"""
Interfaces to query the OpenAI API.

The abstract interface OpenAIApiProvider has implementations that either directly call the OpenAI API
(OpenAIApiProvider) or proxy through a centralized server (OpenAIApiProxyProvider).
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Literal, TypedDict, TypeVar

import numpy as np

from rrosti.utils.config import FromConfig, ModelCost, config, uses_config
from rrosti.utils.misc import FloatArray

T = TypeVar("T")


# Here's an example of an API response:
# {
#   "data": [
#     {
#       "embedding": [
#         -0.00239578727632761,
#         # ...
#         -0.02301216684281826
#       ],
#       "index": 0,
#       "object": "embedding"
#     }
#   ],
#   "model": "text-embedding-ada-002-v2",
#   "object": "list",
#   "usage": {
#     "prompt_tokens": 6,
#     "total_tokens": 6
#   }
# }


def format_cost(cost: float) -> str:
    # if cost < 0.01 and cost != 0:
    #     return f"{cost * 100:.5f}".rstrip("0").rstrip(".") + " ¢ (USD)"

    return f"{cost:.8f}".rstrip("0").rstrip(".") + " USD"


def _get_model_cost_per_1k_tokens(model: str) -> ModelCost:
    try:
        return config.openai_api.model_cost[model]
    except KeyError:
        for k, v in config.openai_api.model_cost.items():
            if model.startswith(k):
                return v
    raise ValueError(f"Unknown model: {model}")


class _UsageDict(TypedDict):
    prompt_tokens: int
    completion_tokens: int  # only present for completions
    total_tokens: int


class _ChatCompletionMessageDict(TypedDict):
    role: Literal["system", "user", "assistant"]
    content: str


class _ChatCompletionResponseItemDict(TypedDict):
    index: int
    message: _ChatCompletionMessageDict
    finish_reason: Literal["stop", "length"]


class _ChatCompletionResponseDict(TypedDict):
    id: str
    object: Literal["chat.completion"]
    created: int
    model: str
    choices: list[_ChatCompletionResponseItemDict]
    usage: _UsageDict


class _EmbeddingResponseItem(TypedDict):
    object: Literal["embedding"]
    embedding: list[float]
    index: int


class _EmbeddingResponseDict(TypedDict):
    object: Literal["list"]
    data: list[_EmbeddingResponseItem]
    model: str
    usage: _UsageDict


@dataclass(frozen=True)
class EmbeddingResponse:
    snippets: list[str]  # (N,)
    embeddings: FloatArray  # (N, emb_size)
    model: str
    prompt_tokens: int

    @property
    def cost(self) -> float:
        return _get_model_cost_per_1k_tokens(self.model).calculate(self.prompt_tokens, 0)

    @property
    def cost_str(self) -> str:
        return format_cost(self.cost)

    @staticmethod
    def from_dict(snippets: list[str], raw: _EmbeddingResponseDict) -> EmbeddingResponse:
        model = raw["model"]
        prompt_tokens = raw["usage"]["prompt_tokens"]
        # Verify that the data has the right number of elements
        assert len(raw["data"]) == len(snippets), raw
        # The API docs warn us that the order may be different, so sort by the index field
        raw["data"].sort(key=lambda d: d["index"])
        assert [d["index"] for d in raw["data"]] == list(range(len(snippets))), raw
        embeddings = np.array([raw["data"][i]["embedding"] for i in range(len(snippets))], dtype=np.float32)
        return EmbeddingResponse(
            snippets=snippets,
            embeddings=embeddings,
            model=model,
            prompt_tokens=prompt_tokens,
        )

    def __repr__(self) -> str:
        return (
            f"EmbeddingResponse(snippets={self.snippets!r}, model={self.model!r}, "
            f"prompt_tokens={self.prompt_tokens!r}), cost={format_cost(self.cost)}"
        )


class OpenAIApiProvider(ABC):
    """An abstract class that provides the OpenAI API."""

    @abstractmethod
    def _create_chat_completion_raw(
        self,
        *,
        model: str,
        messages: list[_ChatCompletionMessageDict],
        max_tokens: int,
        temperature: float,
    ) -> _ChatCompletionResponseDict:
        """
        Create a chat completion synchronously.

        Wraps openai.ChatCompletion.create, handling Azure etc. transparently.

        This is the raw method that subclasses should override.
        """

    @uses_config
    def create_chat_completion(
        self,
        *,
        model: str,
        messages: list[_ChatCompletionMessageDict],
        max_tokens: int = FromConfig("openai_api.completion_max_tokens"),
        temperature: float = FromConfig("openai_api.completion_temperature"),
    ) -> _ChatCompletionResponseDict:
        """
        Create a chat completion synchronously.

        Wraps openai.ChatCompletion.create, handling Azure etc. transparently.
        """

        assert messages, "Must have at least one message"

        return self._create_chat_completion_raw(
            model=model,
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
        )

    @abstractmethod
    async def _acreate_chat_completion_raw(
        self,
        *,
        model: str,
        messages: list[_ChatCompletionMessageDict],
        max_tokens: int,
        temperature: float,
    ) -> _ChatCompletionResponseDict:
        """
        Create a chat completion asynchronously.

        Wraps openai.ChatCompletion.create, handling Azure etc. transparently.

        This is the raw method that subclasses should override.
        """

    @uses_config
    async def acreate_chat_completion(
        self,
        *,
        model: str,
        messages: list[_ChatCompletionMessageDict],
        max_tokens: int = FromConfig("openai_api.completion_max_tokens"),
        temperature: float = FromConfig("openai_api.completion_temperature"),
    ) -> _ChatCompletionResponseDict:
        """
        Create a chat completion asynchronously.

        Wraps openai.ChatCompletion.create, handling Azure etc. transparently.
        """

        assert messages, "Must have at least one message"

        if max_tokens is None:
            max_tokens = config.openai_api.completion_max_tokens

        if temperature is None:
            temperature = config.openai_api.completion_temperature

        return await self._acreate_chat_completion_raw(
            model=model,
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
        )

    @abstractmethod
    async def _acreate_embedding_raw(
        self,
        input: list[str],
        *,
        model: str,
    ) -> _EmbeddingResponseDict:
        """
        Create an embedding asynchronously.

        Wraps openai.Embedding.create, handling Azure etc. transparently.

        Returns an _EmbeddingResponseDict, which is a TypedDict of what is returned by the API.
        """

    @abstractmethod
    def _create_embedding_raw(
        self,
        input: list[str],
        *,
        model: str,
    ) -> _EmbeddingResponseDict:
        """
        Create an embedding synchronously.

        Wraps openai.Embedding.create, handling Azure etc. transparently.

        Returns an _EmbeddingResponseDict, which is a TypedDict of what is returned by the API.
        """

    async def acreate_embedding(
        self,
        input: list[str],
        *,
        model: str | None = None,
    ) -> EmbeddingResponse:
        """
        Create an embedding asynchronously.

        Wraps openai.Embedding.create, handling Azure etc. transparently.
        """

        assert input, "Must have at least one prompt"

        if model is None:
            model = config.openai_api.embedding_model

        return EmbeddingResponse.from_dict(input, await self._acreate_embedding_raw(input, model=model))

    def create_embedding(
        self,
        input: list[str],
        *,
        model: str | None = None,
    ) -> EmbeddingResponse:
        """
        Create an embedding synchronously.

        Wraps openai.Embedding.create, handling Azure etc. transparently.
        """

        assert input, "Must have at least one prompt"

        if model is None:
            model = config.openai_api.embedding_model

        return EmbeddingResponse.from_dict(input, self._create_embedding_raw(input, model=model))

rrosti/llm_api/openai_api.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""An implementation of OpenAIApiProvider that directly calls the OpenAI API."""

from __future__ import annotations

import asyncio
import concurrent
import threading
import time as _unwrapped_time
from abc import ABC, abstractmethod
from typing import Any, ClassVar, Generic, Literal, TypeVar, cast

import attrs
import janus
import openai
from loguru import logger
from overrides import override
from typing_extensions import Self

from rrosti.llm_api import openai_api as api
from rrosti.llm_api.rate_limiter import RateLimiter
from rrosti.query import logging as qlog
from rrosti.utils.config import config


class _SystemAdapter:
    """
    This static class wraps some system operations that we need to mock in tests without influencing the rest of the
    system.
    """

    class time:
        @staticmethod
        def time() -> float:
            return _unwrapped_time.time()

    class asyncio:
        @staticmethod
        async def sleep(seconds: float) -> None:
            await asyncio.sleep(seconds)


@attrs.frozen
class ChatCompletionRequest:
    model: str
    messages: list[api._ChatCompletionMessageDict]
    max_tokens: int | None
    temperature: float


@attrs.frozen
class EmbeddingRequest:
    model: str
    input: list[str]


class DirectOpenAIApiProvider(api.OpenAIApiProvider):
    """An OpenAIApiProvider that directly calls the OpenAI API."""

    _chat_worker: _ChatCompletionWorker
    _embedding_worker: _EmbeddingWorker

    def __init__(self) -> None:
        self._chat_worker = _ChatCompletionWorker.get_instance()
        self._embedding_worker = _EmbeddingWorker.get_instance()

        endpoint = config.openai_api.endpoint
        openai.api_key = endpoint.api_key
        if endpoint.api_type is not None:
            openai.api_type = endpoint.api_type
        if endpoint.api_version is not None:
            openai.api_version = endpoint.api_version
        if endpoint.api_base is not None:
            openai.api_base = endpoint.api_base

    @override
    def _create_chat_completion_raw(
        self,
        *,
        model: str,
        messages: list[api._ChatCompletionMessageDict],
        max_tokens: int,
        temperature: float,
    ) -> api._ChatCompletionResponseDict:
        return self._chat_worker.work(
            ChatCompletionRequest(model=model, messages=messages, max_tokens=max_tokens, temperature=temperature)
        )

    @override
    async def _acreate_chat_completion_raw(
        self,
        *,
        model: str,
        messages: list[api._ChatCompletionMessageDict],
        max_tokens: int,
        temperature: float,
    ) -> api._ChatCompletionResponseDict:
        return await self._chat_worker.awork(
            ChatCompletionRequest(model=model, messages=messages, max_tokens=max_tokens, temperature=temperature)
        )

    @override
    def _create_embedding_raw(
        self,
        input: list[str],
        *,
        model: str,
    ) -> api._EmbeddingResponseDict:
        return self._embedding_worker.work(EmbeddingRequest(model=model, input=input))

    @override
    async def _acreate_embedding_raw(
        self,
        input: list[str],
        *,
        model: str,
    ) -> api._EmbeddingResponseDict:
        return await self._embedding_worker.awork(EmbeddingRequest(model=model, input=input))


def _map_model(model: str) -> dict[str, str]:
    engine_map = config.openai_api.endpoint.engine_map
    if engine_map is not None:
        return dict(engine=engine_map[model])
    return dict(model=model)


RequestT = TypeVar("RequestT")
ResponseT = TypeVar("ResponseT")


class WorkItem(Generic[RequestT, ResponseT]):
    """
    A work item that can be used synchronously or asynchronously.

    For example, a request to the embedding API or the chat completion API.
    """

    request: RequestT
    fut: _SyncFuture[ResponseT] | _AsyncFuture[ResponseT]
    time_created: float
    time_dequeued: float | None
    time_started: float | None
    time_completed: float | None

    _private_object = object()

    def __init__(self, private_object: object) -> None:
        """Private constructor. Do not call."""
        assert private_object is self._private_object, "Do not call this constructor"

    def _initialize(
        self,
        request: RequestT,
        fut: concurrent.futures.Future[ResponseT] | asyncio.Future[ResponseT],
    ) -> None:
        self.request = request
        self.fut = fut
        self.time_created = _SystemAdapter.time.time()
        self.time_dequeued = None
        self.time_completed = None

    @classmethod
    def make_sync(cls, request: RequestT) -> WorkItem[RequestT, ResponseT]:
        """Make a synchronous work item."""
        item = cls(cls._private_object)
        item._initialize(request, _SyncFuture())
        return item

    @classmethod
    def make_async(cls, request: RequestT) -> WorkItem[RequestT, ResponseT]:
        """Make an asynchronous work item."""
        item = cls(cls._private_object)
        item._initialize(request, _AsyncFuture())
        return item

    def set_result(self, result: ResponseT) -> None:
        """Set the result of the future."""
        self.time_completed = _SystemAdapter.time.time()
        if isinstance(self.fut, _SyncFuture):
            self.fut.set_result(result)
        else:
            assert isinstance(self.fut, _AsyncFuture)
            self.fut.get_loop().call_soon_threadsafe(self.fut.set_result, result)

    def set_exception(self, exc: Exception) -> None:
        """Set the exception of the future."""
        self.time_completed = _SystemAdapter.time.time()
        if isinstance(self.fut, _SyncFuture):
            self.fut.set_exception(exc)
        else:
            assert isinstance(self.fut, _AsyncFuture)
            self.fut.get_loop().call_soon_threadsafe(self.fut.set_exception, exc)


_RequestWorkerT = TypeVar("_RequestWorkerT", bound="RequestWorker[Any, Any]")


class RequestWorker(ABC, Generic[RequestT, ResponseT]):
    """A worker with a queue of requests that can be used synchronously or asynchronously."""

    _instances: ClassVar[dict[type, Any]] = {}

    _queue: janus.Queue[WorkItem[RequestT, ResponseT] | Literal["stop"]]
    _threads: list[threading.Thread]
    _rate_limiter: RateLimiter

    @classmethod
    def get_instance(cls) -> Self:
        """Get an instance of the given class, creating it if necessary."""
        if cls is RequestWorker:
            raise TypeError("Cannot instantiate abstract class RequestWorker")
        if cls not in RequestWorker._instances:
            # This requires cls to have a no-argument constructor
            RequestWorker._instances[cls] = cls()  # type: ignore[call-arg]
        return cast(Self, RequestWorker._instances[cls])

    def _finalize_on_delete(self) -> None:
        pass

    @classmethod
    def delete_instance(cls: type[_RequestWorkerT]) -> None:
        """Delete the instance of the given class, if it exists."""
        if cls in RequestWorker._instances:
            RequestWorker._instances[cls]._finalize_on_delete()
            del RequestWorker._instances[cls]

    class RateLimitException(Exception):
        pass

    def stop(self) -> None:
        """Stop the worker. Useful for testing and orderly shutdown."""
        logger.info("Stopping {}...", type(self).__name__)
        for _ in self._threads:
            self._queue.sync_q.put("stop")
        for thread in self._threads:
            thread.join()
        self._rate_limiter.stop()
        logger.info("Stopped {}.", type(self).__name__)

    def __init__(self, num_threads: int) -> None:
        if type(self) is RequestWorker:
            raise TypeError("Cannot instantiate abstract class RequestWorker")

        logger.info("Initializing {} with {} threads", type(self).__name__, num_threads)
        self._queue = janus.Queue()
        self._threads = [
            threading.Thread(target=self._thread_main, daemon=True, name="{i}-worker-{type(self)}")
            for i in range(num_threads)
        ]
        self._rate_limiter = RateLimiter()
        for thread in self._threads:
            thread.start()

    def _check_request(self, req: RequestT) -> None:
        pass

    def work(self, request: RequestT) -> ResponseT:
        """Work on a request synchronously."""
        self._check_request(request)
        item = WorkItem[RequestT, ResponseT].make_sync(request)
        self._queue.sync_q.put(item)
        return item.fut.result()

    async def awork(self, request: RequestT) -> ResponseT:
        """Work on a request asynchronously."""
        self._check_request(request)
        item = WorkItem[RequestT, ResponseT].make_async(request)
        await self._queue.async_q.put(item)
        await cast(Any, item.fut)
        return item.fut.result()

    @abstractmethod
    def _process_one(self, request: RequestT) -> ResponseT:
        ...

    def _thread_main(self) -> None:
        # logger.info("Thread {} starting", threading.current_thread().name)
        while True:
            item = self._queue.sync_q.get()
            if item == "stop":
                return
            item.time_dequeued = _SystemAdapter.time.time()
            while True:  # loop until we get a response
                with self._rate_limiter() as scope:
                    item.time_started = _SystemAdapter.time.time()
                    try:
                        response = self._process_one(item.request)
                    except RequestWorker.RateLimitException:
                        logger.info("{}: Rate limit reached.", str(type(self).__name__))
                        scope.signal_error()
                        continue
                    except Exception as e:
                        item.time_completed = _SystemAdapter.time.time()
                        item.set_exception(e)
                    else:
                        item.time_completed = _SystemAdapter.time.time()
                        item.set_result(response)
                    break


class _ChatCompletionWorker(RequestWorker[ChatCompletionRequest, api._ChatCompletionResponseDict]):
    """A worker with a queue of chat completion requests that can be used synchronously or asynchronously."""

    def __init__(self) -> None:
        super().__init__(num_threads=config.openai_api.num_chat_completion_worker_threads)

    @override
    def _check_request(self, req: ChatCompletionRequest) -> None:
        if not req.messages:
            raise ValueError("Must have at least one message")

    @override
    def _process_one(self, request: ChatCompletionRequest) -> api._ChatCompletionResponseDict:
        logger.info("Getting chat completion for {} messages", len(request.messages))
        try:
            start_time = _SystemAdapter.time.time()
            resp = _create_chat_completion_internal(
                model=request.model,
                messages=request.messages,
                max_tokens=request.max_tokens,
                temperature=request.temperature,
            )
        except openai.error.RateLimitError as e:
            raise RequestWorker.RateLimitException from e
        except Exception as e:
            logger.error("Error during chat completion: {}", e)
            raise

        end_time = _SystemAdapter.time.time()
        elapsed_time = end_time - start_time
        logger.info("Got chat completion in {:.3f} seconds.", elapsed_time)
        return resp


class _EmbeddingWorker(RequestWorker[EmbeddingRequest, api._EmbeddingResponseDict]):
    """A worker with a queue of embedding requests that can be used synchronously or asynchronously."""

    def __init__(self) -> None:
        super().__init__(num_threads=config.openai_api.num_embedding_worker_threads)

    @override
    def _check_request(self, req: EmbeddingRequest) -> None:
        if not req.input:
            raise ValueError("Must have at least one prompt")

    @override
    def _process_one(self, request: EmbeddingRequest) -> api._EmbeddingResponseDict:
        logger.info("Getting embedding for {} prompts", len(request.input))
        while True:
            try:
                start_time = _SystemAdapter.time.time()
                resp = _create_embedding_internal(
                    model=request.model,
                    input=request.input,
                )
            except openai.error.RateLimitError as e:
                raise RequestWorker.RateLimitException from e
            except Exception as e:
                logger.error("Error during embedding: {}", e)
                raise

            end_time = _SystemAdapter.time.time()
            elapsed_time = end_time - start_time
            logger.info("Got {} embeddings in {:.3f} seconds.", len(request.input), elapsed_time)
            return resp


def _create_chat_completion_internal(
    *, model: str, messages: list[api._ChatCompletionMessageDict], max_tokens: int | None, temperature: float
) -> api._ChatCompletionResponseDict:
    assert messages, "Must have at least one message"

    with qlog.ChatCompletionCallEvent.section(model, messages, openai.api_base):
        return cast(
            api._ChatCompletionResponseDict,
            openai.ChatCompletion.create(  # type: ignore[no-untyped-call]
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
                **_map_model(model),
            ),
        )


_SyncFuture = concurrent.futures.Future
_AsyncFuture = asyncio.Future


def _create_embedding_internal(
    *,
    model: str,
    input: list[str],
) -> api._EmbeddingResponseDict:
    """
    Create an embedding synchronously.

    Wraps openai.Embedding.create, handling Azure etc. transparently.
    """
    assert input, "Must have at least one prompt"

    return cast(
        api._EmbeddingResponseDict,
        openai.Embedding.create(input=input, **_map_model(model)),  # type: ignore[no-untyped-call]
    )

rrosti/llm_api/openai_api_direct.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from __future__ import annotations

import queue
import threading
import time as _unwrapped_time
from typing import Literal

import attrs
from loguru import logger
from typing_extensions import Self

from rrosti.utils.config import config


class _SystemAdapter:
    """
    This static class wraps some system operations that we need to mock in tests without influencing the rest of the
    system.
    """

    class time:
        @staticmethod
        def time() -> float:
            return _unwrapped_time.time()


class RateLimiter:
    """
    A tool to rate limit anything.

    Enforces a minimum time between operations.
    """

    @attrs.define
    class Params:
        initial_interval: float
        min_interval: float
        max_interval: float
        error_multiplier: float
        success_multiplier: float
        decrease_after: float

        @classmethod
        def from_config(cls) -> Self:
            return cls(
                initial_interval=config.openai_api.rate_limit_initial_interval,
                min_interval=config.openai_api.rate_limit_min_interval,
                max_interval=config.openai_api.rate_limit_max_interval,
                error_multiplier=config.openai_api.rate_limit_error_multiplier,
                success_multiplier=config.openai_api.rate_limit_success_multiplier,
                decrease_after=config.openai_api.rate_limit_decrease_after_seconds,
            )

        def __attrs_post_init__(self) -> None:
            assert self.min_interval <= self.initial_interval <= self.max_interval, (
                self.min_interval,
                self.initial_interval,
                self.max_interval,
            )
            assert self.min_interval > 0, self.min_interval
            assert self.error_multiplier >= 1, self.error_multiplier
            assert 0 < self.success_multiplier <= 1, self.success_multiplier
            assert self.min_interval <= self.max_interval, (self.min_interval, self.max_interval)

    class _Scope:
        _rate_limiter: RateLimiter
        _ok: bool

        def __init__(self, rate_limiter: RateLimiter) -> None:
            self._rate_limiter = rate_limiter
            self._ok = True

        def signal_error(self) -> None:
            assert self._ok, "Cannot signal error twice"
            self._ok = False
            self._rate_limiter._signal_rate_limit_hit()

        def __enter__(self) -> RateLimiter._Scope:
            return self

        def __exit__(self, _exc_type: object, _exc_value: object, _traceback: object) -> None:
            if self._ok:
                self._rate_limiter._signal_success()

    _params: Params

    _ticket_queue: queue.Queue[None | Literal["stop"]]
    _started_queue: queue.Queue[None | Literal["stop"]]
    _result_queue: queue.Queue[bool | Literal["stop"]]

    # The thread that issues permissions to proceed
    _ticket_thread: threading.Thread

    def __init__(
        self,
        params: Params | None = None,
    ) -> None:
        if params is None:
            params = RateLimiter.Params.from_config()
        self._params = params
        self._result_queue = queue.Queue()
        self._started_queue = queue.Queue()
        self._ticket_queue = queue.Queue()

        self._ticket_thread = threading.Thread(target=self._ticket_thread_main, daemon=True)
        self._ticket_thread.start()

    def _ticket_thread_main(self) -> None:
        last_request = 0.0
        curr_interval = self._params.initial_interval

        while True:
            # Allow a call to proceed
            self._ticket_queue.put(None)
            # Wait for someone to take the ticket, or for decrease_after seconds to pass
            while True:
                try:
                    r = self._started_queue.get(timeout=self._params.decrease_after)
                    if r == "stop":
                        return
                except queue.Empty:
                    # timed out; gravitate towards initial interval
                    curr_interval = (3 * curr_interval + self._params.initial_interval) / 4
                    continue
                break

            # Now, we wait for self._curr_interval seconds, except if we get a result, in which
            # case we adjust the timeout.
            last_request = _SystemAdapter.time.time()
            wakeup_time = last_request + curr_interval
            sleep_seconds = wakeup_time - last_request
            while True:
                try:
                    result = self._result_queue.get(timeout=sleep_seconds if sleep_seconds > 0 else 0)
                    if result == "stop":
                        return
                except queue.Empty:
                    # timed out; allow next call to proceed
                    break
                if result:
                    curr_interval *= self._params.success_multiplier
                    curr_interval = max(curr_interval, self._params.min_interval)
                    logger.info("Rate limiter: success, new interval: {:g}", curr_interval)
                else:
                    curr_interval *= self._params.error_multiplier
                    curr_interval = min(curr_interval, self._params.max_interval)
                    logger.info("Rate limiter: error, new interval: {:g}", curr_interval)

    def _signal_rate_limit_hit(self) -> None:
        self._result_queue.put(False)

    def _signal_success(self) -> None:
        self._result_queue.put(True)

    def stop(self) -> None:
        """Stop the rate limiter."""
        self._started_queue.put("stop")
        self._result_queue.put("stop")

        # Empty the ticket queue, then put enough stops
        while True:
            try:
                self._ticket_queue.get_nowait()
            except queue.Empty:
                break
        for _ in range(10):
            # Each of these also requeue the stop, so one in theory is enough, but more can be
            # faster.
            self._ticket_queue.put("stop")

        self._ticket_thread.join()

    def __call__(self) -> _Scope:
        """Get a context manager that can be used to rate limit an operation."""
        # Wait for a ticket
        r = self._ticket_queue.get()
        if r == "stop":
            # Requeue; we don't know how many clients are waiting
            self._ticket_queue.put("stop")
            raise KeyboardInterrupt("Rate limiter stopped")
        self._started_queue.put(None)
        return self._Scope(self)

rrosti/llm_api/rate_limiter.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
Log queries and other pertinent information to a file.

The general loguru logger is used for general application logging, in the interests of software
development and debugging, often in a very ad hoc fashion.

The query logger, instead, is used for more planned and carefully specified logging of queries,
LLM executions, the times they took, user feedback (thumbs up/down), and so on.
"""


from rrosti.query.logging.query_logging import (
    ChatCompletionCallEvent,
    ConnectionEvent,
    LoggingNotInitializedError,
    QueryEvent,
    ReceivedMessageFromFrontendEvent,
    SendMessageToFrontendEvent,
    ServerStartedEvent,
    UserDislikedMessageEvent,
    UserInputReceivedEvent,
    UserLikedMessageEvent,
    WaitForUserInputEvent,
    init_disabled_logging,
)

__all__ = [
    "ChatCompletionCallEvent",
    "ConnectionEvent",
    "LoggingNotInitializedError",
    "QueryEvent",
    "ReceivedMessageFromFrontendEvent",
    "SendMessageToFrontendEvent",
    "ServerStartedEvent",
    "UserDislikedMessageEvent",
    "UserInputReceivedEvent",
    "UserLikedMessageEvent",
    "WaitForUserInputEvent",
    "init_disabled_logging",
]

rrosti/query/logging/__init__.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
Log queries and other pertinent information to a file.

The general loguru logger is used for general application logging, in the interests of software
development and debugging, often in a very ad hoc fashion.

The query logger, instead, is used for more planned and carefully specified logging of queries,
LLM executions, the times they took, user feedback (thumbs up/down), and so on.
"""

from __future__ import annotations

import contextlib
import copy
import dataclasses
import os
import platform
import signal
import sys
import threading
import uuid
from abc import ABC, abstractmethod
from contextlib import suppress
from contextvars import ContextVar
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from pathlib import Path
from types import FrameType
from typing import Any, BinaryIO, Callable, Iterator

import loguru
import orjson
import zstandard as zstd
from overrides import override
from typing_extensions import Self

# This is used just as a hint to make sure users don't instantiate some classes.
_PRIVATE_OBJECT = object()

# This is a silly chicken and egg problem. We globally start logging before we have read the config, so we cannot
# really use the config here.
_ROTATION_SIZE_MEGABYTES = 250


class _SystemAdapter:
    """Wrap some filesystem operations that we need to mock in tests."""

    @staticmethod
    def mkdir(path: Path) -> None:
        """Make a directory."""
        path.mkdir(parents=True, exist_ok=True)

    @staticmethod
    def open_xb(path: Path) -> BinaryIO:
        """Open a file."""
        if path == Path(os.devnull):
            return open(path, "wb")  # noqa: SIM115 - cannot use context handler here
        return open(path, "xb")  # noqa: SIM115 - cannot use context handler here

    @staticmethod
    def close(f: BinaryIO) -> None:
        f.close()

    @staticmethod
    def rename(original_filename: Path, new_filename: Path) -> None:
        """Rename a File"""
        original_filename.rename(new_filename)

    @staticmethod
    def exists(path: Path) -> bool:
        """Check if a given file exists"""
        return os.path.exists(path)

    @staticmethod
    def getsize(path: Path) -> int:
        """Get the Size of a given File in bytes"""
        return os.path.getsize(path)


def _get_timestamp() -> str:
    """Get a timestamp for the current time in the format YYYYMMDD_HHMMSS."""
    return datetime.now().strftime("%Y%m%d_%H%M%S")


class LoggingNotInitializedError(Exception):
    """Raised when logging is used before it's initialized."""


class _LoggingEnvironment:
    """
    Base class representing a logging environment. This one logs to /dev/null and can thus also
    be used directly in tests.

    There is (at most) one per execution. It contains things like the run UUID and the log file.

    For server runs, this is initialized by logging a ServerStartedEvent. For test and script runs,
    something else should usually be done in order to not pollute the server logs.
    """

    run_uuid = str(uuid.uuid4())
    _query_logger: loguru.Logger

    _log_file: RotatingLog
    _run_log_fname_stem: Path  # This Filename should not include the ending ".log.json.zst"

    LOG_PATH: str = os.devnull  # override in subclasses

    def _make_query_filename(cls, query_uuid: str) -> Path:
        """Make a query log file name."""
        if cls._run_log_fname_stem == Path(os.devnull):
            return Path(os.devnull)
        return cls._run_log_fname_stem.parent / "queries" / f"{_get_timestamp()}__{query_uuid}"

    @property
    def log_file(self) -> RotatingLog:
        """Get the log file."""
        return self._log_file

    @classmethod
    def _add_regular_logger_handlers(cls) -> None:
        """Add regular logger handlers."""

        Path(cls.LOG_PATH).mkdir(parents=True, exist_ok=True)

        rotation_size = f"{_ROTATION_SIZE_MEGABYTES} MB"

        # Log the regular logs as JSON...
        loguru.logger.add(
            Path(cls.LOG_PATH) / "debug.log.json",
            level="DEBUG",
            serialize=True,
            enqueue=True,
            rotation=rotation_size,
            compression="gz",
        )

        # ... and, for convenience, also as text.
        loguru.logger.add(
            Path(cls.LOG_PATH) / "debug.log.txt",
            level="DEBUG",
            enqueue=True,
            rotation=rotation_size,
            compression="gz",
        )

    def __init__(self) -> None:
        """Initialize the logging environment."""

        loguru.logger.info("Starting logging to '{}'", self.LOG_PATH)

        if self.LOG_PATH != os.devnull:
            loguru.logger.remove()
            loguru.logger.add(sys.stderr, level=os.environ.get("LOGURU_LEVEL", "INFO"))
            self._add_regular_logger_handlers()

            self._run_log_fname_stem = (
                Path(self.LOG_PATH) / f"{_get_timestamp()}__{_LoggingEnvironment.run_uuid}" / "run"
            )

            _SystemAdapter.mkdir(self._run_log_fname_stem.parent)
        else:
            self._run_log_fname_stem = Path(os.devnull)

        self._query_logger = loguru.logger.bind(logger="query_logger")
        self._query_logger.add(_handler, filter=lambda r: r["message"] == _RawQueryLogEntry._MAGIC)

        self._log_file = RotatingLog(self._run_log_fname_stem)


class _ServerLoggingEnvironment(_LoggingEnvironment):
    """A logging environment for server runs."""

    LOG_PATH = "logs"


_logging_environment: _LoggingEnvironment | None = None
_query_logger_context: ContextVar[_Context | None] = ContextVar("query_logger_context", default=None)


def _cleanup() -> None:
    """Reset the global state of the module. Used in test fixtures."""

    global _logging_environment
    global _query_logger_context
    _logging_environment = None
    _query_logger_context = ContextVar("query_logger_context", default=None)
    loguru.logger.remove()
    loguru.logger.add(sys.stderr, level=os.environ.get("LOGURU_LEVEL", "INFO"))


def _env() -> _LoggingEnvironment:
    """Get the logging environment."""
    if _logging_environment is None:
        raise LoggingNotInitializedError(
            "Logging environment not initialized by ServerStartedEvent (for servers) or similar"
        )
    return _logging_environment


@dataclass
class _Context:
    """
    A context for logging. This is used to keep track of the current logfile and sections we're in, in a stack-like
    fashion.
    """

    logfile: RotatingLog
    vars: dict[str, Any]
    outer: _Context | None = None
    depth: int = field(init=False)
    _old_context: _Context | None = field(default=None, repr=False, init=False)

    def __post_init__(self) -> None:
        self.depth = 0 if self.outer is None else self.outer.depth + 1

    def get(self, key: str, default: Any = None) -> Any:
        if not isinstance(key, str):
            raise TypeError
        """Get a context variable."""
        if key in self.vars:
            return self.vars[key]
        if self.outer is not None:
            return self.outer.get(key, default)
        return default

    def as_dict(self) -> dict[str, Any]:
        """Convert the context to a dictionary."""
        return dict(depth=self.depth, vars=self.vars, outer=self.outer.as_dict() if self.outer else None)

    def nest(self, logfile: RotatingLog | None = None, **kwargs: Any) -> _Context:
        """Nest a new context."""
        if logfile is None:
            logfile = self.logfile
        return _Context(logfile=logfile, vars=kwargs, outer=self)

    def __enter__(self) -> Self:
        """Enter the context."""
        self._old_context = _query_logger_context.get()
        _query_logger_context.set(self)
        return self

    def __exit__(self, _exc_type: object, _exc_value: object, _traceback: object) -> None:
        """Exit the context."""
        _query_logger_context.set(self._old_context)


def _get_query_logger_context() -> _Context:
    """
    Get the current query logger context.

    For creating new contexts, prefer _context().
    """
    context = _query_logger_context.get()
    if context is None:
        context = _Context(logfile=_env().log_file, vars={})
        _query_logger_context.set(context)
    return context


@contextlib.contextmanager
def _context(**kwargs: Any) -> Iterator[_Context]:
    """
    Add contextual information to the query log.

    When creating a new context, this context manager takes care of the cleanup of the old context.
    """
    outer_context = _get_query_logger_context()
    with outer_context.nest(**kwargs) as inner_context:
        yield inner_context


@dataclass(frozen=True)
class _Metadata:
    """Metadata about a query log entry."""

    level: str
    time: datetime
    elapsed: timedelta
    run_uuid: str = _LoggingEnvironment.run_uuid

    @staticmethod
    def from_loguru_record(record: loguru.Record) -> _Metadata:
        assert isinstance(record["time"], datetime)
        assert isinstance(record["elapsed"], timedelta)
        return _Metadata(
            elapsed=record["elapsed"],
            level=record["level"].name,
            time=record["time"],
        )


@dataclass(frozen=True)
class _RawQueryLogEntry:
    """A single entry in the query log, encompassing metadata and log data."""

    rec: dict[str, Any]
    meta: _Metadata
    _MAGIC = "$$query_log_entry"

    @staticmethod
    def from_loguru_record(record: loguru.Record) -> _RawQueryLogEntry:
        if record["message"] != _RawQueryLogEntry._MAGIC:
            raise ValueError(f"Invalid log entry: {record}")
        extra = copy.copy(record["extra"])
        if "logger" in extra and extra["logger"] == "query_logger":
            del extra["logger"]
        return _RawQueryLogEntry(
            meta=_Metadata.from_loguru_record(record),
            rec=extra,
        )

    def as_dict(self) -> dict[str, Any]:
        """Convert the query log entry to a dictionary."""
        return dict(
            rec=self.rec,
            meta=dataclasses.asdict(self.meta),
        )


def _to_json(obj: Any) -> str | float:
    if isinstance(obj, timedelta):
        return obj.total_seconds()
    if isinstance(obj, datetime):
        return obj.isoformat()
    if isinstance(obj, Path):
        return str(obj)
    raise TypeError(f"Type is not JSON serializable: {type(obj)}")


class RotatingLog:
    _file: BinaryIO
    _compressor: zstd.ZstdCompressionWriter
    filename: str
    filename_without_ending: str
    _file_ending: str
    _file_lock: threading.Lock

    _original_sigterm_handler: Callable[[int, FrameType | None], Any] | int | None
    _original_sigint_handler: Callable[[int, FrameType | None], Any] | int | None

    @property
    def max_size(self) -> int:
        return _ROTATION_SIZE_MEGABYTES * 1_000_000

    def __init__(self, filename: str | Path) -> None:
        if Path(filename) == Path(os.devnull):
            self._file_ending = ""
        else:
            self._file_ending = ".log.json.zst"
        self.filename_without_ending = str(filename)
        self.filename = self.filename_without_ending + self._file_ending
        self._open()
        self._file_lock = threading.Lock()

        self._original_sigterm_handler = signal.getsignal(signal.SIGTERM)
        self._original_sigint_handler = signal.getsignal(signal.SIGINT)

        # Register signal handler for SIGTERM and SIGINT (Ctrl+C)
        signal.signal(signal.SIGTERM, self._handle_exit)
        signal.signal(signal.SIGINT, self._handle_exit)

    def __del__(self) -> None:
        # Close File (and flush compressor) if not already closed
        with suppress(ValueError):
            self._close()

    def _handle_exit(self, signum: int, frame: FrameType | None) -> None:
        # Close File (and flush compressor) if not already closed
        with suppress(ValueError):
            self._close()
        # Call the original signal handler to exit gracefully
        original_handler = self._original_sigterm_handler if signum == signal.SIGTERM else self._original_sigint_handler
        if original_handler is None:
            return
        if isinstance(original_handler, int):
            assert original_handler == signal.SIG_DFL or original_handler == signal.SIG_IGN, original_handler
            return
        original_handler(signum, frame)

    def _get_rotation_filename(self) -> Path:
        formatted_datetime = datetime.now().strftime("%Y-%m-%d_%H-%M-%S_%f")
        return Path(self.filename_without_ending + "." + formatted_datetime + self._file_ending)

    def __rotation_needed(self) -> bool:
        return (
            _SystemAdapter.exists(Path(self.filename)) and _SystemAdapter.getsize(Path(self.filename)) >= self.max_size
        )

    def __rotate_if_needed(self) -> None:
        if self.__rotation_needed():
            self._close()
            _SystemAdapter.rename(Path(self.filename), self._get_rotation_filename())
            self._open()

    def write(self, msg: "loguru.Message") -> None:  # noqa: UP037  (quotes necessary because loguru doesn't export)
        with self._file_lock:
            try:
                self.__rotate_if_needed()
                self._compressor.write(
                    orjson.dumps(_RawQueryLogEntry.from_loguru_record(msg.record).as_dict(), default=_to_json)
                )
                self._compressor.write(b"\n")
                self._compressor.flush(
                    zstd.FLUSH_FRAME
                )  # TODO: This is costly, so for the future we should try to move all the logging to a seperate thread
            except ValueError:
                loguru.logger.error(f"Trying to write to closed file {self.filename}.")

    def _open(self) -> None:
        self._file = _SystemAdapter.open_xb(Path(self.filename))
        self._compressor = zstd.ZstdCompressor().stream_writer(self._file)

    def _close(self) -> None:
        self._compressor.flush(zstd.FLUSH_FRAME)
        _SystemAdapter.close(self._file)


def _handler(msg: "loguru.Message") -> None:  # noqa: UP037  (quotes necessary because loguru doesn't export)
    """
    Handle the processing and writing of a log message to a log file.

    Converts the log message to JSON and writes it to the appropriate logfile.
    """
    _get_query_logger_context().logfile.write(msg)


def _get_platform_info() -> dict[str, Any]:
    """Get information about the current platform."""
    return dict(
        hostname=platform.node(),
        system=platform.system(),
        platform=platform.platform(),
        python_buildno=platform.python_build()[0],
        python_builddate=platform.python_build()[1],
        processor=platform.processor(),
        python_version=platform.python_version(),
        python_path=sys.executable,
    )


@contextlib.contextmanager
def _log_section(event: _LogEvent) -> Iterator[None]:
    """Add contextual information to the query log."""
    event_dict = event.to_dict()
    event_name = event_dict.pop("event")
    with _context(event=event_name, args=event_dict) as inner_context:
        _env()._query_logger.bind(
            event=event_name, args=event_dict, section=True, context=inner_context.as_dict()
        ).info(_RawQueryLogEntry._MAGIC)

        exception: str | None = None
        try:
            yield
        except Exception as e:
            try:
                exception = f"{type(e).__name__}: {e}"
            except Exception:
                exception = f"Exception of type {type(e)} that cannot be converted to string"
            raise
        finally:
            _env()._query_logger.bind(
                event="end_section", name=event_name, section=True, context=inner_context.as_dict(), exception=exception
            ).info(_RawQueryLogEntry._MAGIC)


class _LogEvent(ABC):
    """
    An abstract base class representing a log event.
    Classes derived from this class represent specific kinds of log

    There are two ways to use the events:

    1. query_logging.ServerStartedEvent.log(args=sys.argv)
    2. with query_logging.ServerStartedEvent.section(args=sys.argv): ...

    Most subclasses only support one or the other.
    """

    @abstractmethod
    def to_dict(self) -> dict[str, Any]:
        """Convert the event to a dictionary."""


def _log(event: _LogEvent) -> None:
    """Log a query log entry."""
    d = event.to_dict()
    event_name = d.pop("event")

    _env()._query_logger.bind(
        event=event_name, args=d, section=False, context=_get_query_logger_context().as_dict()
    ).info(_RawQueryLogEntry._MAGIC)


@contextlib.contextmanager
def _switch_logfile(logfile: Path) -> Iterator[None]:
    """Switch to a new logfile."""
    with _LogfileSwitchEvent.section(logfile):
        old_context = _get_query_logger_context()
        old_logfile = old_context.logfile
        assert old_logfile is not None
        _SystemAdapter.mkdir(logfile.parent)
        with _context(logfile=RotatingLog(logfile), logfile_name=logfile):
            _LogfileSwitchedFromEvent.log(parent=str(old_logfile.filename))
            yield


## These are the only public API of this module.


def init_disabled_logging() -> None:
    """Initialize logging for a test run."""
    global _logging_environment
    _logging_environment = _LoggingEnvironment()


class ServerStartedEvent(_LogEvent):
    """A server started event."""

    args: list[str]

    def __init__(self, private_object: object, args: list[str]) -> None:
        """Private constructor. Use the log() method instead."""
        assert private_object is _PRIVATE_OBJECT, "Cannot instantiate ServerStartedEvent"
        self.args = args

    # We need to repeat to_dict and log in every subclass, sadly, to get mypy and type hinting
    # in IDEs to work well.
    @override
    def to_dict(self) -> dict[str, Any]:
        """Convert the event to a dictionary."""
        return dict(event="server_started", args=self.args, platform=_get_platform_info())

    @staticmethod
    def log(args: list[str]) -> None:
        """Log the event. Also initializes the logging environment."""

        global _logging_environment
        if _logging_environment is None:
            _logging_environment = _ServerLoggingEnvironment()
        else:
            raise RuntimeError("ServerStartedEvent already logged. There should be only one.")

        _log(ServerStartedEvent(_PRIVATE_OBJECT, args=args))


class ConnectionEvent(_LogEvent):
    """A connection event."""

    id: str
    local_addr: str
    remote_addr: str

    def __init__(self, private_object: object, id: str, local_addr: str, remote_addr: str) -> None:
        """Private constructor. Use the section() method instead."""
        assert private_object is _PRIVATE_OBJECT, "Cannot instantiate ConnectionEvent"
        self.id = id
        self.local_addr = local_addr
        self.remote_addr = remote_addr

    @override
    def to_dict(self) -> dict[str, Any]:
        """Convert the event to a dictionary."""
        return dict(event="connection", id=self.id, local_addr=self.local_addr, remote_addr=self.remote_addr)

    @staticmethod
    @contextlib.contextmanager
    def section(id: str, local_addr: str, remote_addr: str) -> Iterator[None]:
        with _log_section(ConnectionEvent(_PRIVATE_OBJECT, id=id, local_addr=local_addr, remote_addr=remote_addr)):
            yield


class QueryEvent(_LogEvent):
    """A query event."""

    text: str
    uuid: str
    prompt: str

    def __init__(self, private_object: object, text: str, uuid: str, prompt: str) -> None:
        """Private constructor. Use the section() method instead."""
        assert private_object is _PRIVATE_OBJECT, "Cannot instantiate QueryEvent"
        self.text = text
        self.uuid = uuid
        self.prompt = prompt

    @override
    def to_dict(self) -> dict[str, Any]:
        """Convert the event to a dictionary."""
        return dict(event="query", text=self.text, uuid=self.uuid, prompt=self.prompt)

    @staticmethod
    @contextlib.contextmanager
    def section(text: str, uuid: str, prompt: str) -> Iterator[None]:
        query_fname = _env()._make_query_filename(uuid)
        with _log_section(QueryEvent(_PRIVATE_OBJECT, text=text, uuid=uuid, prompt=prompt)), _switch_logfile(
            query_fname
        ):
            yield


class ChatCompletionCallEvent(_LogEvent):
    """An LLM call event."""

    model: str
    text: str
    endpoint: str  # OPENAI_API_BASE or similar

    def __init__(self, private_object: object, model: str, text: str, endpoint: str) -> None:
        """Private constructor. Use the section() method instead."""
        assert private_object is _PRIVATE_OBJECT, "Cannot instantiate ChatCompletionCallEvent"
        self.model = model
        self.text = text
        self.endpoint = endpoint

    @override
    def to_dict(self) -> dict[str, Any]:
        """Convert the event to a dictionary."""
        return dict(event="chat_completion_call", model=self.model, text=self.text, endpoint=self.endpoint)

    @staticmethod
    @contextlib.contextmanager
    def section(model: str, text: list[Any], endpoint: str) -> Iterator[None]:
        with _log_section(ChatCompletionCallEvent(_PRIVATE_OBJECT, model=model, text=str(text), endpoint=endpoint)):
            yield


class _LogfileSwitchEvent(_LogEvent):
    """A logfile switch event."""

    logfile: Path

    def __init__(self, private_object: object, logfile: Path) -> None:
        """Private constructor. Use the section() method instead."""
        assert private_object is _PRIVATE_OBJECT, "Cannot instantiate LogfileSwitchEvent"
        self.logfile = logfile

    @override
    def to_dict(self) -> dict[str, Any]:
        """Convert the event to a dictionary."""
        return dict(event="switch_to_logfile", logfile=str(self.logfile))

    @staticmethod
    @contextlib.contextmanager
    def section(logfile: Path) -> Iterator[None]:
        with _log_section(_LogfileSwitchEvent(_PRIVATE_OBJECT, logfile=logfile)):
            yield


class _LogfileSwitchedFromEvent(_LogEvent):
    """A logfile switch event."""

    parent: str

    def __init__(self, private_object: object, parent: str) -> None:
        """Private constructor. Use the log() method instead."""
        assert private_object is _PRIVATE_OBJECT, "Cannot instantiate LogfileSwitchedFromEvent"
        self.parent = parent

    @override
    def to_dict(self) -> dict[str, Any]:
        """Convert the event to a dictionary."""
        return dict(event="switched_from_logfile", parent=self.parent)

    @staticmethod
    def log(parent: str) -> None:
        """Log the event."""
        _log(_LogfileSwitchedFromEvent(_PRIVATE_OBJECT, parent=parent))


class WaitForUserInputEvent(_LogEvent):
    """A wait for user input event."""

    def __init__(self, private_object: object) -> None:
        """Private constructor. Use the section() method instead."""
        assert private_object is _PRIVATE_OBJECT, "Cannot instantiate WaitForUserInputEvent"

    @override
    def to_dict(self) -> dict[str, Any]:
        """Convert the event to a dictionary."""
        return dict(event="wait_for_user_input")

    @staticmethod
    @contextlib.contextmanager
    def section() -> Iterator[None]:
        with _log_section(WaitForUserInputEvent(_PRIVATE_OBJECT)):
            yield


class UserInputReceivedEvent(_LogEvent):
    """A user input event."""

    text: str
    uuid: str

    def __init__(self, private_object: object, text: str, uuid: str) -> None:
        """Private constructor. Use the log() method instead."""
        assert private_object is _PRIVATE_OBJECT, "Cannot instantiate UserInputReceivedEvent"
        self.text = text
        self.uuid = uuid

    @override
    def to_dict(self) -> dict[str, Any]:
        """Convert the event to a dictionary."""
        return dict(event="user_input_received", text=self.text, uuid=self.uuid)

    @staticmethod
    def log(text: str, uuid: str) -> None:
        """Log the event."""
        _log(UserInputReceivedEvent(_PRIVATE_OBJECT, text=text, uuid=uuid))


class SendMessageToFrontendEvent(_LogEvent):
    """A message sent event."""

    uuid: str
    msg: str
    time: float | None
    cost: float | None

    def __init__(self, private_object: object, uuid: str, msg: str, time: float | None, cost: float | None) -> None:
        """Private constructor. Use the log() method instead."""
        assert private_object is _PRIVATE_OBJECT, "Cannot instantiate SendMessageToFrontendEvent"
        self.uuid = uuid
        self.msg = msg
        self.time = time
        self.cost = cost

    @override
    def to_dict(self) -> dict[str, Any]:
        """Convert the event to a dictionary."""
        return dict(event="send_message_frontend", uuid=self.uuid, msg=self.msg, time_used=self.time, cost=self.cost)

    @staticmethod
    def log(uuid: str, msg: str, time_used: float | None, cost: float | None) -> None:
        """Log the event."""
        _log(SendMessageToFrontendEvent(_PRIVATE_OBJECT, uuid=uuid, msg=msg, time=time_used, cost=cost))


class ReceivedMessageFromFrontendEvent(_LogEvent):
    """A message received event."""

    uuid: str
    msg: str

    def __init__(self, private_object: object, uuid: str, msg: str) -> None:
        """Private constructor. Use the log() method instead."""
        assert private_object is _PRIVATE_OBJECT, "Cannot instantiate ReceivedMessageFromFrontendEvent"
        self.uuid = uuid
        self.msg = msg

    @override
    def to_dict(self) -> dict[str, Any]:
        """Convert the event to a dictionary."""
        return dict(event="received_message_frontend", uuid=self.uuid, msg=self.msg)

    @staticmethod
    def log(uuid: str, msg: str) -> None:
        """Log the event."""
        _log(ReceivedMessageFromFrontendEvent(_PRIVATE_OBJECT, uuid=uuid, msg=msg))


class UserLikedMessageEvent(_LogEvent):
    """A user liked message event."""

    uuid: str
    text: str | None
    ref_uuid: str

    def __init__(self, private_object: object, uuid: str, text: str | None, ref_uuid: str) -> None:
        """Private constructor. Use the log() method instead."""
        assert private_object is _PRIVATE_OBJECT, "Cannot instantiate UserLikedMessageEvent"
        self.uuid = uuid
        self.text = text
        self.ref_uuid = ref_uuid

    @override
    def to_dict(self) -> dict[str, Any]:
        """Convert the event to a dictionary."""
        return dict(event="user_liked_message", uuid=self.uuid, text=self.text, ref_uuid=self.ref_uuid)

    @staticmethod
    def log(uuid: str, text: str | None, ref_uuid: str) -> None:
        """Log the event."""
        _log(UserLikedMessageEvent(_PRIVATE_OBJECT, uuid=uuid, text=text, ref_uuid=ref_uuid))


class UserDislikedMessageEvent(_LogEvent):
    """A user disliked message event."""

    uuid: str
    text: str | None
    ref_uuid: str

    def __init__(self, private_object: object, uuid: str, text: str | None, ref_uuid: str) -> None:
        """Private constructor. Use the log() method instead."""
        assert private_object is _PRIVATE_OBJECT, "Cannot instantiate UserDislikedMessageEvent"
        self.uuid = uuid
        self.text = text
        self.ref_uuid = ref_uuid

    @override
    def to_dict(self) -> dict[str, Any]:
        """Convert the event to a dictionary."""
        return dict(event="user_disliked_message", uuid=self.uuid, text=self.text, ref_uuid=self.ref_uuid)

    @staticmethod
    def log(uuid: str, text: str | None, ref_uuid: str) -> None:
        """Log the event."""
        _log(UserDislikedMessageEvent(_PRIVATE_OBJECT, uuid=uuid, text=text, ref_uuid=ref_uuid))

rrosti/query/logging/query_logging.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
Allow users to ask questions about data. Answer those questions by delegating them to an LLM,
which in turn can execute specific functions like rtfm to answer them.
"""

from __future__ import annotations

import asyncio
import sys

import marshmallow as ma
from overrides import override

from rrosti.chat import chat_session
from rrosti.chat.state_machine import execution, interpolable
from rrosti.chat.state_machine.execution import load_and_run
from rrosti.llm_api import openai_api
from rrosti.servers.websocket_query_server import Frontend, QueryEngineBase


class DataQueryEngine(QueryEngineBase):
    _llm: chat_session.LLM
    _openai_provider: openai_api.OpenAIApiProvider

    def __init__(self, llm: chat_session.LLM, openai_provider: openai_api.OpenAIApiProvider) -> None:
        self._llm = llm
        self._openai_provider = openai_provider

    @override
    async def ensure_loaded(self) -> None:
        await asyncio.gather(interpolable.ensure_loaded(), execution.ensure_loaded())

    @override
    async def arun(self, frontend: Frontend) -> None:
        try:
            await load_and_run(llm=self._llm, frontend=frontend, openai_provider=self._openai_provider)
        except ma.exceptions.ValidationError as e:
            print(e.messages, sys.stderr)
            raise

rrosti/servers/data_retrieval_engine.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""A websocket query server that uses the data retrieval engine to answer questions."""

import argparse
import asyncio
import logging
import sys
import tracemalloc

from loguru import logger

import rrosti.utils.config
from rrosti.chat.chat_session import OpenAI, UserInputLLM
from rrosti.chat.state_machine import execution
from rrosti.llm_api import openai_api_direct
from rrosti.query import logging as qlog
from rrosti.servers import data_retrieval_engine, websocket_query_server
from rrosti.utils import misc
from rrosti.utils.config import config
from rrosti.utils.misc import ProgramArgsBase


class ProgramArgs(websocket_query_server.ProgramArgsMixin, execution.ProgramArgsMixin, ProgramArgsBase):
    temperature: float

    @classmethod
    def _add_args(cls, parser: argparse.ArgumentParser) -> None:
        super()._add_args(parser)
        parser.add_argument(
            "--temperature",
            type=float,
            default=config.openai_api.completion_temperature,
            help="The temperature of the OpenAI Model",
        )


async def main() -> None:
    rrosti.utils.config.load()
    openai_provider = openai_api_direct.DirectOpenAIApiProvider()

    tracemalloc.start()
    misc.setup_logging()

    args = ProgramArgs.parse_args()

    logging.getLogger("matplotlib").setLevel(logging.INFO)
    logging.getLogger("asyncio").setLevel(logging.INFO if not args.debug_asyncio else logging.DEBUG)
    logging.getLogger("urllib3").setLevel(logging.INFO)
    logging.getLogger("websockets").setLevel(logging.INFO)
    # logging.getLogger("openai").setLevel(logging.INFO)

    qlog.ServerStartedEvent.log(args=sys.argv)

    engine: websocket_query_server.QueryEngineBase
    if args.user_simulate_llm:
        logger.info("Starting query server, using a user simulated LLM")
        engine = data_retrieval_engine.DataQueryEngine(llm=UserInputLLM(), openai_provider=openai_provider)
    else:
        logger.info("Starting query server")
        engine = data_retrieval_engine.DataQueryEngine(
            llm=OpenAI(openai_api_direct.DirectOpenAIApiProvider(), temperature=args.temperature),
            openai_provider=openai_provider,
        )

    asyncio.get_event_loop().set_debug(args.debug_asyncio)
    await engine.aserve_forever(args)


if __name__ == "__main__":
    asyncio.run(main())

rrosti/servers/serve_data_retrieval_ws.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Implementation of a websocket query server."""

from __future__ import annotations

import argparse
import asyncio
import json
import re
import uuid
from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum, auto
from pathlib import Path
from typing import Any

import websockets
from loguru import logger
from overrides import override

from rrosti.chat.chat_session import Message
from rrosti.query import logging as qlog
from rrosti.snippets.snippet import Snippet
from rrosti.utils.config import config
from rrosti.utils.misc import ProgramArgsMixinProtocol

_FAKE_DELAY = 1.5

_ALLOWED_CHARACTERS_RE = re.compile(
    r"[^a-zA-ZäÄëËïÏöÖüÜâÂêÊîÎôÔûÛéÉèÈàÀùÙñÑßçÇ0-9" + re.escape(r",.?!\-_: ();<>=+*°") + "]"
)


class ProgramArgsMixin(ProgramArgsMixinProtocol):
    port: int
    public: bool
    debug_send_intermediates: bool
    debug_asyncio: bool

    @classmethod
    def _add_args(cls, parser: argparse.ArgumentParser) -> None:
        super()._add_args(parser)
        parser.add_argument("--port", type=int, default=config.backend.listen_port, help="Port to listen on")
        parser.add_argument("--public", action="store_true", help="Listen on all interfaces")
        parser.add_argument("--debug-asyncio", action="store_true", help="Debug coroutine execution")
        parser.add_argument(
            "--debug-send-intermediates",
            action="store_true",
            help="Send intermediate messages and more detailed output to the frontend",
        )


@dataclass(frozen=True)
class UserInputMessage:
    content: str
    uuid: str


class MessageType(Enum):
    question = auto()
    like = auto()
    dislike = auto()
    new_question = auto()
    unknown = auto()


class UnknownMessageType(Exception):
    def __init__(self, message: str) -> None:
        self.message = message
        super().__init__(message)


class Frontend(ABC):
    _peek_buffer: UserInputMessage | None = None

    @abstractmethod
    async def _get_user_input_impl(self) -> UserInputMessage:
        ...

    async def get_user_input(self) -> UserInputMessage:
        msg = await self.peek_user_input()
        self._peek_buffer = None
        return msg

    async def peek_user_input(self) -> UserInputMessage:
        if self._peek_buffer is None:
            self._peek_buffer = await self._get_user_input_impl()
        return self._peek_buffer

    @abstractmethod
    async def send_message(self, msg: Message | str) -> None:
        ...

    @abstractmethod
    async def read_message(self, message: str) -> MessageType:
        ...

    @abstractmethod
    def handle_python_output(self, python_output: PythonItem) -> None:
        ...

    @abstractmethod
    def handle_rtfm_output(self, snippets: list[Snippet]) -> int:
        ...


class WebFrontend(Frontend):
    __user_message_queue: asyncio.Queue[UserInputMessage]
    answer: Answer
    __debug_send_intermediates: bool
    __websocket: websockets.WebSocketServerProtocol  # type: ignore[name-defined]

    def __init__(
        self,
        websocket: websockets.WebSocketServerProtocol,  # type: ignore[name-defined]
        debug_send_intermediates: bool = False,
    ) -> None:
        self.__websocket = websocket
        self.__debug_send_intermediates = debug_send_intermediates
        self.__user_message_queue = asyncio.Queue()
        self.answer = Answer()

    @staticmethod
    def _sanitize_str(s: str) -> str:
        return _ALLOWED_CHARACTERS_RE.sub("", s)

    @override
    async def _get_user_input_impl(self) -> UserInputMessage:
        with qlog.WaitForUserInputEvent.section():
            if not self.answer.is_empty():
                await self.__send_msg(
                    dict(
                        id=str(uuid.uuid4()),
                        type="answer",
                        content=self.answer.to_message(self.__debug_send_intermediates),
                    ),
                    time=None,
                    cost=None,
                )
            await self.__send_msg(dict(id=str(uuid.uuid4()), type="user_turn"), time=None, cost=None)
            reply = await self.__user_message_queue.get()
            qlog.UserInputReceivedEvent.log(text=reply.content, uuid=reply.uuid)
        return reply

    @override
    async def send_message(self, msg: Message | str) -> None:
        self.answer.append_answer(msg)
        if isinstance(msg, Message):
            await self.__send_msg(
                dict(id=str(uuid.uuid4()), type="intermediate", content=msg.text.replace("\n", "<br>")),
                time=msg.time_used,
                cost=msg.cost,
            )
        else:
            await self.__send_msg(
                dict(id=str(uuid.uuid4()), type="intermediate", content=msg.replace("\n", "<br>")),
                time=None,
                cost=None,
            )

    @override
    async def read_message(self, message: str) -> MessageType:
        message_data = json.loads(message)
        logger.info("Received query: {}", message_data)
        message_content = self._sanitize_str(message_data["content"])
        if message_data["type"] == "question":
            await self.__user_message_queue.put(UserInputMessage(content=message_content, uuid=message_data["id"]))
            return MessageType.question
        if message_data["type"] == "like":
            qlog.UserLikedMessageEvent.log(
                uuid=message_data["id"], text=message_content, ref_uuid=message_data["refId"]
            )
            return MessageType.like
        if message_data["type"] == "dislike":
            qlog.UserDislikedMessageEvent.log(
                uuid=message_data["id"], text=message_content, ref_uuid=message_data["refId"]
            )
            return MessageType.dislike
        if message_data["type"] == "new_question":
            return MessageType.new_question
        return MessageType.unknown

    @override
    def handle_python_output(self, python_output: PythonItem) -> None:
        self.answer.append_python(python_output._code, python_output._output)

    @override
    def handle_rtfm_output(self, snippets: list[Snippet]) -> int:
        start_index = self.answer._n_excerpts
        for snippet in snippets:
            if not snippet.page_start:
                page_range = ""
            elif snippet.page_start == snippet.page_end:
                page_range = f", p. {snippet.page_start}"
            else:
                page_range = f", p. {snippet.page_start}-{snippet.page_end}"
            link_path = str(Path("source") / Path(snippet.source_filename).name)
            if snippet.page_start:
                link_path += "#page=" + str(snippet.page_start)
            self.answer.append_rtfm(
                content=snippet.text,
                title=f"{Path(snippet.source_filename).stem}{page_range}",
                link=link_path,
            )
        return start_index

    async def __send_msg(self, msg: dict[str, Any], time: float | None, cost: float | None) -> None:
        """
        Make sure that a message is logged (even if not sent), and then send it to the frontend (over the open
        websocket), if either it is not just an intermediate message or if the server is run with
        --debug-send-intermediates.
        """
        msg_str = json.dumps(msg)
        # FIXME: We very much should record these at the level of the messages that were shown to the user,
        # not at the level of the messages that were sent to the frontend.
        qlog.SendMessageToFrontendEvent.log(uuid=msg["id"], msg=msg_str, time_used=time, cost=cost)
        if msg["type"] != "intermediate" or self.__debug_send_intermediates:
            await self.__websocket.send(msg_str)


@dataclass
class InterpolableItem(ABC):
    """
    A base class for Items that should be sent together with an Answer and were previously generated from an
    Interpolable.

    For example either a PythonItem including python code and output of a piece of python code that was run from an
    interpolable.
    """

    @abstractmethod
    def to_dict(self) -> dict[str, Any]:
        ...


@dataclass
class PythonItem(InterpolableItem):
    _code: str
    _output: str

    @override
    def to_dict(self) -> dict[str, str]:
        return {"code": self._code, "output": self._output}


@dataclass
class RTFMExcerpt(InterpolableItem):
    _index: int
    _content: str
    _title: str | None
    _link: str | None

    @override
    def to_dict(self, debug: bool = False) -> dict[str, str | int | None]:
        return_dict: dict[str, str | int | None] = {
            "title": self._title,
            "link": self._link,
        }
        if debug:
            return_dict["index"] = str(self._index)
            return_dict["content"] = self._content
        return return_dict


class Answer:
    """
    A class that collects all intermediate answers, rtfm excerpts and python code that was run to then be sent all
    together in a final answer to the frontend.
    """

    _intermediate_answers: list[Message | str]
    _rtfm_excerpts: list[RTFMExcerpt]
    _python: list[PythonItem]
    _n_excerpts: int

    def __init__(self) -> None:
        self._intermediate_answers = []
        self._rtfm_excerpts = []
        self._python = []
        self._n_excerpts = 0

    def is_empty(self) -> bool:
        return len(self._intermediate_answers) < 1

    def append_answer(self, answer: Message | str) -> None:
        self._intermediate_answers.append(answer)

    def append_rtfm(self, content: str, title: str | None = None, link: str | None = None) -> None:
        self._rtfm_excerpts.append(RTFMExcerpt(_index=self._n_excerpts, _content=content, _title=title, _link=link))
        self._n_excerpts += 1

    def append_python(self, code: str, output: str) -> None:
        self._python.append(PythonItem(_code=code, _output=output))

    def __intermediates_to_list_of_dicts(
        self,
    ) -> tuple[list[dict[str, str | float | None]], dict[str, str | float | None]]:
        messages = []
        final_answer = {}
        total_time: float = 0
        total_cost: float = 0
        for msg in self._intermediate_answers:
            if isinstance(msg, Message):
                messages.append(dict(content=msg.text.replace("\n", "<br>"), time_used=msg.time_used, cost=msg.cost))
                total_time += msg.time_used if isinstance(msg.time_used, float) else 0
                total_cost += msg.cost if isinstance(msg.cost, float) else 0
            else:
                messages.append(dict(content=msg, time_used=None, cost=None))
        final_answer["content"] = messages[-1]["content"]
        final_answer["time_used"] = total_time
        final_answer["cost"] = total_cost
        return messages, final_answer

    def to_message(self, debug: bool = False) -> dict[str, Any]:
        intermediates_dict, final_answer = self.__intermediates_to_list_of_dicts()
        data: dict[str, Any] = {"final_answer": final_answer}
        if debug:
            data["intermediate_answers"] = intermediates_dict
            data["python"] = [python.to_dict() for python in self._python]
        data["rtfm_excerpts"] = [excerpt.to_dict(debug) for excerpt in self._rtfm_excerpts]
        return data


class QueryEngineBase(ABC):
    """
    A base class for query engines.

    These need to be stateless; there is a single instance of an engine that is used to serve all
    user sessions.
    """

    async def ensure_loaded(self) -> None:  # noqa: B027 (intentional empty non-abstract method)
        """Called at initialization time before the first query is handled."""

    @abstractmethod
    async def arun(self, frontend: Frontend) -> None:
        """
        Starts a process that can send messages and request information/questions from the user. At the beginning
        get_user_input has to be called to initiate the process with the frontend.
        """

    @logger.catch
    async def aserve_forever(self, args: ProgramArgsMixin) -> None:
        asyncio.get_event_loop().slow_callback_duration = 0.05  # 50 ms

        # Used to create and close a query logging context.
        @logger.catch
        async def ahandle_query(
            websocket: websockets.WebSocketServerProtocol, path: str  # type: ignore[name-defined]
        ) -> None:
            nonlocal args

            async def awebsocket_message_handler(frontend: Frontend) -> None:
                try:
                    async for message in websocket:
                        message_type = await frontend.read_message(message)
                        if message_type is MessageType.new_question:
                            query_handler_task.cancel()
                            # by closing the websocket, the whole engine in the back is relaunched for a new question
                            # (not keeping the old questions)
                            await websocket.close()
                        elif message_type is MessageType.unknown:
                            raise UnknownMessageType(f'Error: Unknown message type "{message["type"]}".')
                except (
                    websockets.exceptions.ConnectionClosedError,  # type: ignore[attr-defined]
                    websockets.exceptions.ConnectionClosedOK,  # type: ignore[attr-defined]
                ):
                    logger.info("WebSocket connection closed")
                finally:
                    query_handler_task.cancel()
                    # by closing the websocket, the whole engine in the back is relaunched for a new question
                    # (not keeping the old questions)
                    await websocket.close()

            async def aquery_handler(frontend: Frontend) -> None:
                await self.arun(frontend)

            frontend = WebFrontend(websocket=websocket, debug_send_intermediates=args.debug_send_intermediates)

            try:
                remote_address: tuple[Any, ...] = (
                    websocket.request_headers.get("X-Real-IP"),
                    int(websocket.request_headers.get("X-Remote-Port")),
                )
            except TypeError:
                remote_address = websocket.remote_address
            if remote_address[0] is None:
                remote_address = websocket.remote_address
            with qlog.ConnectionEvent.section(
                id=websocket.id, local_addr=websocket.local_address, remote_addr=str(remote_address)
            ):
                # create tasks for the websocket message handler and the query handler
                query_handler_task = asyncio.create_task(aquery_handler(frontend))
                user_message_handler_task = asyncio.create_task(awebsocket_message_handler(frontend))

                # wait for both tasks to finish (they won't, under normal conditions)
                await asyncio.gather(user_message_handler_task, query_handler_task)

        logger.info("Loading all engine resources...")
        await self.ensure_loaded()
        logger.info("Starting websocket server...")

        try:
            async with websockets.serve(  # type: ignore[attr-defined]
                ahandle_query, "0.0.0.0" if args.public else "127.0.0.1", port=args.port
            ):
                await asyncio.Future()
        except websockets.exceptions.ConnectionClosedOK as e:  # type: ignore[attr-defined]
            logger.info("Connection closed gracefully:", e)

    def serve_forever(self, args: ProgramArgsMixin) -> None:
        asyncio.run(self.aserve_forever(args), debug=args.debug_asyncio)

rrosti/servers/websocket_query_server.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import attrs


@attrs.frozen
class Page:
    """A single page of a document. Contains the page text and any text OCRd from images."""

    page_number: int = attrs.field(validator=attrs.validators.instance_of(int))
    text: str = attrs.field(validator=attrs.validators.instance_of(str))

    # Image_text may be the empty string for images without text.
    image_text: list[str] = attrs.field(
        validator=attrs.validators.deep_iterable(member_validator=attrs.validators.instance_of(str))
    )

rrosti/snippets/_page.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
Base class for vector databases of snippets.

A snippet database is a collection of snippets that can be searched by vector similarity.
This is used to find the most relevant snippets for a query.
"""

from __future__ import annotations

import asyncio
from abc import ABC, abstractmethod
from typing import Iterable

from loguru import logger

from rrosti.llm_api import openai_api
from rrosti.snippets.snippet import Snippet
from rrosti.utils.token_count import token_count


def _count_tokens(question: str, snippets: list[str]) -> int:
    """Returns the token count of the query that would be sent to the LLM."""
    # FIXME: This estimates the token count in some way by procuding something that looks a bit like a prompt.
    # It used to be sillier, using an old jinja template that was not used anywhere else.
    # Now it tries to emulate that somewhat for compatibility.
    sysmsg = "\n\n".join(f"## Extract #{i+1}:\n\n{snippet}" for i, snippet in enumerate(snippets))
    return token_count(sysmsg) + token_count(question) + 80


class AbstractSnippetDatabase(ABC):
    """
    Base class for vector databases of snippets.

    A snippet database is a collection of snippets that can be searched by vector similarity.
    This is used to find the most relevant snippets for a query.
    """

    @abstractmethod
    def has_id(self, id: str) -> bool:
        """Check if the id (text hash) is in the database."""

    @abstractmethod
    def __contains__(self, snippet: Snippet) -> bool:
        """Check if the snippet is in the database."""

    @abstractmethod
    def get_by_id(self, id: str) -> Snippet | None:
        """Get a snippet by its id (text hash)."""

    @abstractmethod
    def add_snippet(self, snippet: Snippet) -> Snippet:
        """
        Add a snippet to the database.

        Returns the snippet that was added, which may be different from the input snippet
        if the snippet was already in the database.
        """

    @abstractmethod
    def add_snippets(self, snippets: Iterable[Snippet]) -> None:
        """Add multiple snippets to the database. Upserts."""

    @abstractmethod
    async def find_nearest_raw(
        self, openai_provider: openai_api.OpenAIApiProvider, query: Snippet, n_results: int
    ) -> list[Snippet]:
        """
        Find the nearest raw (unmerged) snippets to the query.

        Returns a list of snippets sorted by distance.
        """

    async def find_nearest_merged(
        self, openai_provider: openai_api.OpenAIApiProvider, query: Snippet, max_tokens: int, n_merge_candidates: int
    ) -> list[Snippet]:
        """
        query: The query to find snippets for. Used to determine the total length of snippets.
        max_tokens: The maximum number of tokens to return.
        n_merge_candidates: How many raw snippets to fetch and consider for merging.

        Returns (snippets: (1..M,)), where M <= n_merge_candidates.

        Distances are not returned because they are not well defined for merged snippets.
        """

        await query.async_ensure_embedding(openai_provider)

        candidates = await self.find_nearest_raw(openai_provider, query, n_merge_candidates)
        while True:
            merged = Snippet.merge_list(candidates)
            logger.debug("{} candidates merged into {} snippets", len(candidates), len(merged))
            # FIXME: This looks fishy: I think only qa_token_count is used from that module, and it relies on a prompt
            # template that we don't use otherwise.
            token_count = await asyncio.to_thread(_count_tokens, query.text, [m.text for m in merged])
            if token_count > max_tokens:
                # too long
                logger.debug("{} snippets is too long: {} > {}", len(candidates), token_count, max_tokens)
                candidates = candidates[:-1]
            else:
                logger.info("Going to query with {} snippets, {} after merging", len(candidates), len(merged))
                logger.debug("Tokens: {} <= {}", token_count, max_tokens)
                return merged

rrosti/snippets/abstract_snippet_database.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
Parse documents from `config.document_sync.source_docs_path` into `config.document_sync.parsed_docs_path`.

Generate snippets and maintain an embedding cache.
"""

import asyncio
import hashlib
from pathlib import Path
from typing import Awaitable, Iterable

import orjson
from aiopath import AsyncPath  # type: ignore[import]
from loguru import logger
from tqdm import tqdm

from rrosti.llm_api import openai_api, openai_api_direct
from rrosti.snippets.parsed_document import ParsedDocument
from rrosti.snippets.snippet import EmbeddingCache, Snippet
from rrosti.utils.config import config


async def _handle_source_document(doc_path: Path) -> bool:
    """
    Parse a source document and output it to `config.document_sync.parsed_docs_path`.

    Returns True if something was updated.
    """

    if doc_path.suffix.lower() not in [".pdf", ".txt"]:
        logger.warning(f"Unknown file type, ignoring document: {doc_path}")
        return False

    target_apath = AsyncPath(config.document_sync.parsed_docs_path / (doc_path.name + ".json"))

    doc_apath = AsyncPath(doc_path)
    doc_data = await doc_apath.read_bytes()

    # If the target exists, we check (using the sha256) if it's the same document.
    if await target_apath.exists():
        json_text = await target_apath.read_bytes()
        # TODO: handle corrupted JSON
        parsed_document = ParsedDocument.from_dict(orjson.loads(json_text))
        new_sha = hashlib.sha256(doc_data).hexdigest()
        if parsed_document.sha256 == new_sha:
            logger.trace(f"Document already parsed: {doc_path}")
            return False
        logger.info("Document changed: {}; old={}, new={}", doc_path, parsed_document.sha256, new_sha)
    else:
        logger.info("New document: {}", doc_path)

    # Look at the suffix. We handle pdf and txt.
    if doc_path.suffix.lower() == ".pdf":
        raise NotImplementedError("PDF parsing is not yet implemented")
        # parsed_document = await pdf_parse.parse_pdf(doc_path)
    elif doc_path.suffix.lower() == ".txt":  # noqa: RET506 (unnecessary elif after raise)
        parsed_document = ParsedDocument.from_textfile_bytes(doc_data, name=doc_path.name, path=str(doc_path))
    else:
        assert False

    # Write the parsed document to `config.document_sync.parsed_docs_path`. Add .json to the end of the filename.
    output_apath = AsyncPath(config.document_sync.parsed_docs_path / (doc_path.name + ".json"))
    await output_apath.parent.mkdir(parents=True, exist_ok=True)
    await output_apath.write_bytes(orjson.dumps(parsed_document.to_dict()))

    return True


def _snippetize_document(path: Path) -> list[Snippet]:
    """Take in a parsed document and output a list of snippets."""
    # logger.info("Reading bytes from {}", path)
    data = path.read_bytes()
    # logger.info("Parsing JSON")
    json_dict = orjson.loads(data)
    # logger.info("Creating ParsedDocument")
    doc = ParsedDocument.from_dict(json_dict)
    # logger.info("Getting snippets")
    return doc.get_snippets(images=True)


async def _snippetize_documents(paths: Iterable[Path]) -> list[Snippet]:
    """Take in a list of parsed documents and output a list of snippets."""
    aws: list[Awaitable[list[Snippet]]] = [asyncio.to_thread(_snippetize_document, path) for path in paths]

    snippets: list[Snippet] = []
    for aw in aws:
        snippets.extend(await aw)

    return snippets


async def sync_and_get_snippets() -> list[Snippet]:
    config.document_sync.data_gen_path.mkdir(parents=True, exist_ok=True)
    openai_provider = openai_api_direct.DirectOpenAIApiProvider()  # TODO

    logger.info("Looking for source documents...")
    source_documents = [p for p in config.document_sync.source_docs_path.glob("*") if p.is_file()]
    logger.info(f"Found {len(source_documents)} source documents")

    if not source_documents:
        logger.error("No source documents found")
        raise RuntimeError("No source documents found")

    # Parse all source documents.

    # Ok, we could do this in parallel, but there's something fishy in the PDF parsing library;
    # some PDFs seem to randomly cause it to segfault. Doing it sequentially mitigates this,
    # at least in the sense that we will compute one PDF at a time and save the result.

    updated = [await _handle_source_document(path) for path in source_documents]

    num_updated = sum(updated)
    if num_updated == 0:
        logger.info("No documents updated")
        # TODO: Detect removed documents so we don't need to always regenerate the snippets?
    else:
        logger.info("Updated {} documents.", num_updated)

    # Snippetize all parsed documents.
    logger.info("Generating snippets...")
    snippets = await _snippetize_documents(config.document_sync.parsed_docs_path.glob("*.json"))

    # Load cached embeddings so we won't recompute for snippets we've already seen.
    cache = await asyncio.to_thread(EmbeddingCache.load)
    await asyncio.to_thread(cache.sync_with_snippets, snippets)

    no_embedding = [s for s in snippets if not s.has_embedding]
    logger.info("Found {} (out of {}) snippets without embeddings.", len(no_embedding), len(snippets))

    async def ensure_embeddings_and_sync(
        openai_provider: openai_api.OpenAIApiProvider, snippets: list[Snippet]
    ) -> None:
        await Snippet.async_ensure_embeddings(openai_provider, snippets)
        cache.sync_with_snippets(snippets)

    if no_embedding:
        logger.info("Querying embeddings...")
        aws: list[Awaitable[None]] = []
        for i in tqdm(range(0, len(no_embedding), config.openai_api.endpoint.max_embedding_requests_per_query)):
            chunk = no_embedding[i : i + config.openai_api.endpoint.max_embedding_requests_per_query]
            aws.append(ensure_embeddings_and_sync(openai_provider, chunk))

        await asyncio.gather(*aws)
        await asyncio.to_thread(cache.save)

    return snippets


if __name__ == "__main__":
    asyncio.run(sync_and_get_snippets())

rrosti/snippets/document_sync.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from __future__ import annotations

import hashlib
from pathlib import Path
from typing import Any

import attrs

from rrosti.snippets import snippet_windows
from rrosti.snippets._page import Page
from rrosti.snippets.snippet import Snippet
from rrosti.utils.config import config


@attrs.frozen
class ParsedDocument:
    """
    A parsed document contains our computed data about a document (such as a PDF file).

    Identifies the document by name and sha256 checksum. Contains the extracted snippets.
    """

    name: str | None = attrs.field(validator=attrs.validators.optional(attrs.validators.instance_of(str)))
    path: str = attrs.field(validator=attrs.validators.instance_of(str))
    sha256: str = attrs.field(validator=attrs.validators.instance_of(str))
    pages: list[Page] = attrs.field(
        validator=attrs.validators.deep_iterable(member_validator=attrs.validators.instance_of(Page))
    )

    @sha256.validator
    def _validate_sha256(self, _attribute: str, value: str) -> None:
        if len(value) != 64 or not all(c in "0123456789abcdef" for c in value):
            raise ValueError(f"Invalid sha256 checksum: {value}")

    def to_dict(self) -> dict[str, str | list[dict[str, str | list[str] | int]] | None]:
        """Create a dict representation of the document, suitable for JSON serialization."""

        return dict(
            name=self.name,
            path=self.path,
            sha256=self.sha256,
            pages=[
                dict(
                    page_number=page.page_number,
                    text=page.text,
                    image_text=page.image_text,
                )
                for page in self.pages
            ],
        )

    @classmethod
    def from_textfile_bytes(cls, data: bytes, name: str, path: str | Path, encoding: str = "utf-8") -> ParsedDocument:
        return ParsedDocument(
            name=name,
            path=str(path),
            pages=[Page(page_number=0, text=data.decode(encoding).strip(), image_text=[])],
            sha256=hashlib.sha256(data).hexdigest(),
        )

    @classmethod
    def from_textfile(cls, path: Path, encoding: str = "utf-8") -> ParsedDocument:
        return cls.from_textfile_bytes(path.read_bytes(), name=path.name, path=path, encoding=encoding)

    @classmethod
    def from_dict(cls, data: dict[str, str | list[Any] | None]) -> ParsedDocument:
        """Create a ParsedDocument from a dict representation, as returned by to_dict()."""

        assert isinstance(data["name"], str) or data["name"] is None
        assert isinstance(data["path"], str)
        assert isinstance(data["sha256"], str)
        assert isinstance(data["pages"], list)

        return cls(
            name=data["name"],
            path=data["path"],
            sha256=data["sha256"],
            pages=[
                Page(
                    page_number=page["page_number"],
                    text=page["text"],
                    image_text=page["image_text"],
                )
                for page in data["pages"]
            ],
        )

    def get_snippets(self, images: bool = True) -> list[Snippet]:
        r"""
        Fetch the full texts of the document. That is:

        - the body of the document, pages joined with "\n\n"
        - each image's text, separately
        """

        snippets = []

        for snip in snippet_windows.windows_with_source(
            pages=[page for page in self.pages if page.text],
            join_str="\n\n",
            window_size=config.document_sync.snippet_window_size,
            step_size=config.document_sync.snippet_step_size,
        ):
            if len(snip.text) < config.document_sync.min_snippet_size:
                continue
            snippets.append(
                Snippet(
                    text=snip.text,
                    source_filename=f"{self.path}",
                    start_offset=snip.start,
                    page_start=snip.page_start,
                    page_end=snip.page_end,
                )
            )

        if not images:
            return snippets

        for page in self.pages:
            for image_text in page.image_text:
                if image_text:
                    for snip in snippet_windows.windows(
                        image_text,
                        window_size=config.document_sync.snippet_window_size,
                        step_size=config.document_sync.snippet_step_size,
                    ):
                        if len(snip.text) < config.document_sync.min_snippet_size:
                            continue
                        snippets.append(
                            Snippet(
                                text=snip.text,
                                source_filename=f"{self.path}",
                                start_offset=snip.start,
                                page_start=snip.page_start,
                                page_end=snip.page_end,
                            )
                        )
        return snippets

rrosti/snippets/parsed_document.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""An in-memory snippet vector database using sklearn's NearestNeighbors."""

from __future__ import annotations

import asyncio
from typing import Iterable

from loguru import logger
from overrides import override
from sklearn.neighbors import NearestNeighbors  # type: ignore[import]

from rrosti.llm_api import openai_api
from rrosti.snippets.abstract_snippet_database import AbstractSnippetDatabase
from rrosti.snippets.snippet import Snippet
from rrosti.utils.misc import FloatArray


class SklearnSnippetDatabase(AbstractSnippetDatabase):
    """An in-memory snippet vector database using sklearn's NearestNeighbors."""

    _nbrs: NearestNeighbors
    _snippets: list[Snippet]
    _embeddings: FloatArray
    _id_to_index_map: dict[str, int]

    def __init__(self, snippets: list[Snippet]) -> None:
        self._snippets = Snippet._drop_duplicates(snippets)
        self._id_to_index_map = {snippet.hash: i for i, snippet in enumerate(self._snippets)}
        self._embeddings = Snippet.consolidate_embeddings(self._snippets)
        self._nbrs = NearestNeighbors(n_neighbors=1, algorithm="brute", metric="cosine", n_jobs=-1).fit(
            self._embeddings
        )
        logger.info(f"Initialized SklearnSnippetDatabase with {len(self._snippets)} snippets")

    @override
    def has_id(self, id: str) -> bool:
        return id in self._id_to_index_map

    @override
    def __contains__(self, snippet: Snippet) -> bool:
        return snippet.hash in self._id_to_index_map

    @override
    def get_by_id(self, id: str) -> Snippet | None:
        if id not in self._id_to_index_map:
            return None
        return self._snippets[self._id_to_index_map[id]]

    @override
    def add_snippet(self, snippet: Snippet) -> Snippet:
        raise NotImplementedError("TODO: implement this")

    @override
    def add_snippets(self, snippets: Iterable[Snippet]) -> None:
        raise NotImplementedError("TODO: implement this")

    @override
    async def find_nearest_raw(
        self, openai_provider: openai_api.OpenAIApiProvider, query: Snippet, n_results: int
    ) -> list[Snippet]:
        emb = (await query.async_get_embedding(openai_provider)).reshape(1, -1)
        logger.info("Finding {} neighbors...", n_results)
        _, indices = await asyncio.to_thread(self._nbrs.kneighbors, emb, n_results)
        logger.info("Found neighbors")
        return [self._snippets[i] for i in indices[0]]

rrosti/snippets/sklearn_snippet_database.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Snippet type that supports merging and optionally has an embedding."""

from __future__ import annotations

import hashlib
from functools import cached_property
from pathlib import Path
from typing import IO, Iterable, Sequence, cast

import attrs
import numpy as np
import orjson
from attrs import field
from attrs.setters import frozen, validate
from loguru import logger

from rrosti.llm_api import openai_api
from rrosti.utils.config import config
from rrosti.utils.misc import FloatArray


def _hash_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def _hash_str(data: str) -> str:
    return _hash_bytes(data.encode("utf-8"))


def _embedding_cache_path() -> Path:
    return config.document_sync.data_gen_path / "embeddings.map"


# TODO: Some of this logic could be simplified, by e.g. having a transparent-ish cache
# that queries the API when it doesn't have an embedding.
class EmbeddingCache:
    _private_tag = object()
    _map: dict[str, FloatArray]
    _sorted: bool

    def _ensure_sorted(self) -> None:
        if not self._sorted:
            self._map = dict(sorted(self._map.items()))
            self._sorted = True

    def __init__(self, private_tag: object) -> None:
        """Internal constructor."""
        assert private_tag is self._private_tag, "EmbeddingCache should not be instantiated directly."
        self._map = {}
        self._sorted = True

    @classmethod
    def load(cls) -> EmbeddingCache:
        cache = cls(cls._private_tag)

        path = _embedding_cache_path()

        if not path.exists():
            logger.info("No embedding cache found, creating new one.")
            return cache

        logger.info("Loading embedding cache from {}...", path)

        hash_list: list[str] = []
        all_bytes: list[bytes] = []

        with open(path) as f:
            for line in f:
                hash, bytes_hex = line.strip().split("\t")
                hash_list.append(hash)
                all_bytes.append(bytes.fromhex(bytes_hex))

        embeddings = np.frombuffer(b"".join(all_bytes), dtype=np.float32).reshape(len(hash_list), -1)

        cache._map = dict(zip(hash_list, embeddings))
        cache._sorted = False

        assert all(len(k) == 64 for k in cache._map)
        assert len({len(v) for v in cache._map.values()}) == 1

        logger.info("Loaded a cache of {} embeddings.", len(cache._map))

        return cache

    def save(self) -> None:
        self._ensure_sorted()

        path = _embedding_cache_path()

        # Atomically replace
        fname_new = Path(str(path) + ".new")
        with open(fname_new, "w") as f:
            for hash, embedding in self._map.items():
                f.write(f"{hash}\t{embedding.tobytes().hex()}\n")

        fname_new.rename(path)

    def _assert_consistency(self, snippet: Snippet) -> None:
        if snippet._embedding is None:
            return
        cached = self._map.get(snippet.hash)
        if cached is None:
            return

        assert snippet._embedding.shape == cached.shape, (snippet._embedding.shape, cached.shape)
        if not np.all(snippet._embedding == cached):
            # Tolerate a small cosine distance between the two embeddings.
            cos_dist = 1 - np.dot(snippet._embedding, cached) / np.linalg.norm(snippet._embedding) / np.linalg.norm(
                cached
            )

            if cos_dist > 0.01:
                logger.error("Embedding cache is inconsistent with snippet.")
                logger.error("Snippet:\n{}", snippet.text)
                logger.error("Snippet embedding:\n{}", snippet._embedding)
                logger.error("Cached embedding:\n{}", cached)
                logger.error("Cosine distance: {}", cos_dist)
                assert False

            # Close enough. Just use the cached embedding.
            snippet._embedding = cached

    def _copy_to_snippet(self, snippet: Snippet) -> None:
        if snippet._embedding is None and snippet.hash in self._map:
            snippet._embedding = self._map[snippet.hash].copy()
            return

    def _copy_from_snippet(self, snippet: Snippet) -> None:
        if snippet._embedding is not None and snippet.hash not in self._map:
            self._map[snippet.hash] = snippet._embedding.copy()

    def sync_with_snippet(self, snippet: Snippet) -> None:
        """
        If the snippet has an embedding, add it to the cache.

        If the cache has an embedding for the snippet, add it to the snippet.
        """

        self._assert_consistency(snippet)
        self._copy_from_snippet(snippet)
        self._copy_to_snippet(snippet)

    def sync_with_snippets(self, snippets: Iterable[Snippet]) -> None:
        for snippet in snippets:
            self.sync_with_snippet(snippet)

    @classmethod
    def from_snippets(cls, snippets: Iterable[Snippet]) -> EmbeddingCache:
        cache = cls(cls._private_tag)
        cache.sync_with_snippets(snippets)
        return cache


# TODO: Split into Snippet and SnippetWithEmbedding
@attrs.define(slots=False)  # slots=False because we use cached_property
class Snippet:
    """
    A snippet type that supports merging. Can optionally have an embedding.

    Embeddings are dense vector representations in a lower-dimensional space
    that capture semantic meaning, allowing for operations like similarity checks.
    """

    source_filename: str = field(
        on_setattr=frozen, kw_only=True, validator=[attrs.validators.instance_of(str)]
    )  # typically filename or similar identifier
    start_offset: int = field(on_setattr=frozen, kw_only=True, validator=[attrs.validators.instance_of(int)])
    text: str = field(on_setattr=frozen, validator=[attrs.validators.instance_of(str)])
    _embedding: FloatArray | None = field(on_setattr=validate, default=None, repr=False, kw_only=True)
    page_start: int | None
    page_end: int | None

    @classmethod
    def _drop_duplicates(cls, snippets: Iterable[Snippet]) -> list[Snippet]:
        hashes_seen: dict[str, int] = {}
        out: list[Snippet] = []
        dup_count = 0

        for snippet in snippets:
            if snippet.hash not in hashes_seen:
                hashes_seen[snippet.hash] = len(out)
                out.append(snippet)
            else:
                # If this one has an embedding and the previous one not, add the embedding to the previous one
                if snippet._embedding is not None and out[hashes_seen[snippet.hash]]._embedding is None:
                    out[hashes_seen[snippet.hash]]._embedding = snippet._embedding
                dup_count += 1

        if dup_count:
            logger.warning("Dropped {} duplicate snippets.", dup_count)

        return out

    @_embedding.validator
    def _validate_embedding(self, _attribute: attrs.Attribute[Snippet], value: FloatArray | None) -> None:
        if value is None:
            return
        assert isinstance(value, np.ndarray)
        assert value.dtype == np.float32
        assert value.ndim == 1
        assert value.shape[0] > 0

    def sync_ensure_embedding(self, openai_provider: openai_api.OpenAIApiProvider) -> None:
        if self._embedding is None:
            self._embedding = openai_provider.create_embedding(input=[self.text]).embeddings[0]

    @classmethod
    def ensure_embeddings(cls, openai_provider: openai_api.OpenAIApiProvider, snippets: Sequence[Snippet]) -> None:
        """
        Ensure that all snippets have embeddings.

        This method will query OpenAI's API for embeddings for all snippets that do not have them.
        """
        snippets_without_embeddings = [s for s in snippets if s._embedding is None]
        if not snippets_without_embeddings:
            return

        logger.info(f"Querying embeddings for {len(snippets_without_embeddings)} snippets")
        embeddings = openai_provider.create_embedding([s.text for s in snippets_without_embeddings]).embeddings
        for snippet, embedding in zip(snippets_without_embeddings, embeddings):
            snippet._embedding = embedding

    async def async_ensure_embedding(self, openai_provider: openai_api.OpenAIApiProvider) -> None:
        if self._embedding is None:
            self._embedding = (await openai_provider.acreate_embedding([self.text])).embeddings[0]

    @classmethod
    async def async_ensure_embeddings(
        cls, openai_provider: openai_api.OpenAIApiProvider, snippets: Sequence[Snippet]
    ) -> None:
        """
        Ensure that all snippets have embeddings.

        This method will query OpenAI's API for embeddings for all snippets that do not have them.
        """
        snippets_without_embeddings = [s for s in snippets if s._embedding is None]
        if not snippets_without_embeddings:
            return

        logger.info(f"Querying embeddings for {len(snippets_without_embeddings)} snippets")
        embeddings = (await openai_provider.acreate_embedding([s.text for s in snippets_without_embeddings])).embeddings
        for snippet, embedding in zip(snippets_without_embeddings, embeddings):
            snippet._embedding = embedding

    @property
    def has_embedding(self) -> bool:
        return self._embedding is not None

    def sync_get_embedding(self, openai_provider: openai_api.OpenAIApiProvider) -> FloatArray:
        """Get the embedding for this snippet, requesting it synchronously if necessary."""
        self.sync_ensure_embedding(openai_provider)
        assert self._embedding is not None
        return self._embedding

    async def async_get_embedding(self, openai_provider: openai_api.OpenAIApiProvider) -> FloatArray:
        """Get the embedding for this snippet, requesting it asynchronously if necessary."""
        await self.async_ensure_embedding(openai_provider)
        assert self._embedding is not None
        return self._embedding

    @cached_property
    def hash(self) -> str:
        """The hash of the snippet text."""
        return _hash_str(self.text)

    @property
    def end_offset(self) -> int:
        return self.start_offset + len(self.text)

    @property
    def length(self) -> int:
        return len(self.text)

    def try_merge(self, other: Snippet) -> Snippet | None:
        """Merge with a potentially overlapping later snippet, unless there is a gap."""
        if self.source_filename != other.source_filename:
            return None

        assert self.start_offset <= other.start_offset, "Snippets must be sorted"
        if other.start_offset > self.end_offset:
            return None
        added_text = other.text[self.end_offset - other.start_offset :]
        return Snippet(
            text=self.text + added_text,
            start_offset=self.start_offset,
            source_filename=self.source_filename,
            embedding=None,
            page_start=self.page_start,
            page_end=other.page_end,
        )

    @staticmethod
    def from_query(query: str) -> Snippet:
        return Snippet(
            text=query, start_offset=0, source_filename="$$query", embedding=None, page_start=None, page_end=None
        )

    @staticmethod
    def merge_list(snippets: Iterable[Snippet]) -> list[Snippet]:
        """
        Given a list of snippets, merge snippets that are next to each other or that overlap.

        Changes snippet order.
        """

        snippets = sorted(snippets, key=lambda s: (s.source_filename, s.start_offset))

        merged_snippets: list[Snippet] = []
        for snippet in snippets:
            if not merged_snippets:
                merged_snippets.append(snippet)
                continue

            last_snippet = merged_snippets[-1]
            merged_snippet = last_snippet.try_merge(snippet)
            if merged_snippet is not None:
                merged_snippets[-1] = merged_snippet
            else:
                merged_snippets.append(snippet)

        return merged_snippets

    @staticmethod
    def from_dict(data: dict[str, int | str | None]) -> Snippet:
        """
        Load a snippet from a dictionary.

        These snippets are required to have embeddings.
        """
        assert isinstance(data["start"], int)
        assert isinstance(data["doc"], str)
        assert isinstance(data["content"], str)
        assert isinstance(data["page_start"], int | None)
        assert isinstance(data["page_end"], int | None)
        if data["embedding"] is None:
            emb: FloatArray | None = None
        else:
            assert isinstance(data["embedding"], str)
            emb = np.frombuffer(bytes.fromhex(data["embedding"]), dtype=np.float32)
        return Snippet(
            source_filename=data["doc"],
            start_offset=data["start"],
            text=data["content"],
            embedding=emb,
            page_start=data["page_start"],  # type: ignore[arg-type]
            page_end=data["page_end"],  # type: ignore[arg-type]
        )

    def to_dict(self) -> dict[str, int | str | None]:
        """Convert a snippet to a dictionary."""
        return {
            "start": self.start_offset,
            "doc": self.source_filename,
            "content": self.text,
            "embedding": self._embedding.tobytes().hex() if self._embedding is not None else None,
            "page_start": self.page_start,
            "page_end": self.page_end,
        }

    @staticmethod
    def consolidate_embeddings(snippets: Sequence[Snippet]) -> FloatArray:
        """
        Consolidate the embeddings of a list of snippets into a single numpy array, replacing
        the embeddings in the snippets with slices of the array.

        The snippets must have embeddings.

        The larger array is useful for computing distances between snippets, while the slices
        are useful for taking less memory than fully copying the embeddings.
        """

        assert all(s._embedding is not None for s in snippets), "All snippets must have embeddings"
        assert all(
            cast(FloatArray, s._embedding).shape == cast(FloatArray, snippets[0]._embedding).shape for s in snippets
        ), "All embeddings must have the same shape"

        embeddings = np.stack([cast(FloatArray, s._embedding) for s in snippets])
        for s, e in zip(snippets, embeddings):
            s._embedding = e
        return embeddings

    @staticmethod
    def load_from_jsonl(file_or_filename: str | Path | IO[bytes]) -> list[Snippet]:
        """Load snippets from a JSONL file."""

        logger.info("Loading snippets from {}", file_or_filename)

        def _do(file: IO[bytes]) -> list[Snippet]:
            return [Snippet.from_dict(orjson.loads(line)) for line in file]

        if isinstance(file_or_filename, (str, Path)):
            with open(file_or_filename, "rb") as f:
                snippets = _do(f)
        else:
            snippets = _do(file_or_filename)

        logger.info("Loaded {} snippets from {}", len(snippets), file_or_filename)
        return snippets

rrosti/snippets/snippet.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


"""Tools to split a text into overlapping windows of a given size."""

import re
from dataclasses import dataclass
from typing import Iterable

from rrosti.snippets._page import Page


@dataclass
class WindowedSnippetInfo:
    text: str
    start: int
    end: int

    # These are 0 for windows()
    page_start: int
    page_end: int


# TODO: This is apparently a bit slow.
def windows(text: str, window_size: int, step_size: int) -> Iterable[WindowedSnippetInfo]:
    """Split a text snippet into windows of a given size, with a given stride."""
    BOUNDARY_RE = re.compile(r"\b", re.DOTALL)
    start = 0
    prev_end = -1
    while start < len(text):
        # Start and end at word boundary. Find the largest window that fits.
        m = BOUNDARY_RE.search(text, start)
        if m is None:
            break
        start = m.start()
        # eat whitespace
        while start < len(text) and text[start].isspace():
            start += 1
        if start >= len(text):
            break
        assert start == 0 or not (text[start - 1].isspace() and text[start].isspace()), text[start - 1 : start + 1]

        end = None
        for match in BOUNDARY_RE.finditer(text, start):
            if match.start() - start > window_size:
                break
            end = match.start()
        if end is None:
            break

        assert end - start <= window_size, (start, end, window_size)
        assert start == 0 or not (text[start - 1].isspace() and text[start].isspace()), text[start - 1 : start + 1]

        if end > prev_end:
            yield WindowedSnippetInfo(text=text[start:end], start=start, end=end, page_start=0, page_end=0)
        prev_end = end

        start += step_size


def windows_with_source(
    pages: list[Page],
    join_str: str,
    window_size: int,
    step_size: int,
) -> Iterable[WindowedSnippetInfo]:
    """
    Split a text snippet into windows of a given size, with a given stride.

    Like windows(), but takes pages separately and fills the page_start and page_end fields.
    """

    # Construct the text body and the indices of the page boundaries
    page_boundaries = [0]
    texts = []
    page_numbers = [1]
    for page in pages:
        page_boundaries.append(page_boundaries[-1] + len(page.text) + len(join_str))
        texts.append(page.text)
        page_numbers.append(page.page_number)
    text = join_str.join(texts)

    for window in windows(text, window_size, step_size):
        # Find the page boundaries that this window overlaps
        page_start = None
        page_end = None
        for i, boundary in enumerate(page_boundaries):
            if boundary > window.start:
                page_start = page_numbers[i]
                break
        for i, boundary in enumerate(page_boundaries):
            if boundary > window.end:
                page_end = page_numbers[i]
                break

        assert page_start is not None
        assert page_end is not None

        yield WindowedSnippetInfo(
            text=window.text,
            start=window.start,
            end=window.end,
            page_start=page_start,
            page_end=page_end,
        )

rrosti/snippets/snippet_windows.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Consume a text file from stdin and output a JSON representation of a ParsedDocument."""

import argparse
import sys
from pathlib import Path

import orjson
from loguru import logger

from rrosti.snippets import parsed_document
from rrosti.snippets._page import Page
from rrosti.utils.misc import file_sha256


class ProgramArgs:
    input: list[Path]


def parse_args() -> ProgramArgs:
    parser = argparse.ArgumentParser()
    parser.add_argument("input", nargs="+", help="Input file(s)", type=Path)
    return parser.parse_args(namespace=ProgramArgs())


def main() -> None:
    args = parse_args()
    for fname in args.input:
        logger.info(f"Parsing {fname}")
        with open(fname) as f:
            text = f.read().strip()
        doc = parsed_document.ParsedDocument(
            name=fname.name,
            path=str(fname),
            pages=[Page(page_number=0, text=text, image_text=[])],
            sha256=file_sha256(fname),
        )
        sys.stdout.buffer.write(orjson.dumps(doc.to_dict()))


if __name__ == "__main__":
    main()

rrosti/tools/text_to_document.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Configuration facility."""

import contextlib
import functools
import inspect
from pathlib import Path
from typing import IO, Any, Callable, Iterator, ParamSpec, TypeVar, cast

import attrs
import omegaconf
from attrs import validators as AV
from loguru import logger
from omegaconf import OmegaConf
from typing_extensions import Self

from rrosti.utils import misc

CONFIG_PATH = Path("config.yaml")
CONFIG_DEFAULTS_PATH = Path("assets/config.defaults.yaml")


@attrs.define
class FrontendConfig:
    websocket_url: str
    title: str


@attrs.define
class ModelCost:
    """Model usage costs per 1k tokens."""

    prompt_tokens: float
    completion_tokens: float = omegaconf.SI("${prompt_tokens}")

    def calculate(self, n_prompt_tokens: int, n_completion_tokens: int) -> float:
        return (n_prompt_tokens / 1000) * self.prompt_tokens + (n_completion_tokens / 1000) * self.completion_tokens


@attrs.define
class Endpoint:
    """Configuration items for an OpenAI API endpoint."""

    api_key: str = attrs.field(repr=False)
    max_embedding_requests_per_query: int
    engine_map: dict[str, str] | None = None
    api_base: str | None = None
    api_type: str | None = None
    api_version: str | None = None


@attrs.define
class OpenAI:
    embedding_model: str
    chat_completion_model: str
    max_tokens_per_model: dict[str, int]
    completion_max_tokens: int = attrs.field(validator=AV.gt(0))

    # We refuse to query if we can request fewer tokens than this
    completion_min_tokens: int = attrs.field(validator=AV.gt(0))

    completion_temperature: float = attrs.field(validator=AV.ge(0.0))
    chat_models: list[str]
    model_cost: dict[str, ModelCost]
    use_endpoint: str
    endpoints: dict[str, Endpoint]

    num_chat_completion_worker_threads: int = attrs.field(validator=AV.gt(0))
    num_embedding_worker_threads: int = attrs.field(validator=AV.gt(0))

    rate_limit_initial_interval: float = attrs.field(validator=AV.gt(0.0))
    rate_limit_min_interval: float = attrs.field(validator=AV.gt(0.0))
    rate_limit_max_interval: float = attrs.field(validator=AV.gt(0.0))
    rate_limit_error_multiplier: float = attrs.field(validator=AV.gt(1.0))
    rate_limit_success_multiplier: float = attrs.field(validator=AV.gt(0.0))
    rate_limit_decrease_after_seconds: float = attrs.field(validator=AV.gt(0.0))

    @property
    def endpoint(self) -> Endpoint:
        return self.endpoints[self.use_endpoint]

    def __attrs_post_init__(self) -> None:
        assert self.completion_min_tokens <= self.completion_max_tokens
        assert self.chat_completion_model in self.chat_models


@attrs.define
class DocumentSync:
    data_gen_path: Path
    source_docs_path: Path
    parsed_docs_path: Path
    snippet_window_size: int
    snippet_step_size: int
    min_snippet_size: int


@attrs.define
class StateMachine:
    yaml_path: Path = attrs.field()
    debug_detect_unresolved_funcalls: bool
    rtfm_max_tokens: int = attrs.field(validator=AV.gt(0))
    rtfm_merge_candidates: int = attrs.field(validator=AV.gt(0))

    @yaml_path.validator
    def _check_yaml_path(self, _attribute: str, value: Path) -> None:
        if not value.exists():
            raise ValueError(f"YAML file {value} does not exist")


@attrs.define
class Backend:
    listen_port: int


@attrs.define
class _Config:
    backend: Backend
    document_sync: DocumentSync
    frontend: FrontendConfig
    openai_api: OpenAI
    state_machine: StateMachine


@contextlib.contextmanager
def _with_resolver(resolver: str, func: Callable[..., str]) -> Iterator[None]:
    OmegaConf.register_new_resolver(resolver, func)
    try:
        yield
    finally:
        OmegaConf.clear_resolver(resolver)


class ConfigurationNotLoadedError(Exception):
    pass


class _ConfigProxy:
    global config

    def __getattribute__(self, name: str) -> Any:
        if name.startswith("__"):
            return super().__getattribute__(name)

        # This can happen if this has been "imported from" before the config was loaded
        if isinstance(config, _Config):
            return getattr(config, name)

        logger.error("Configuration not loaded yet")
        raise ConfigurationNotLoadedError("Configuration not loaded yet")


config: _Config = cast(_Config, _ConfigProxy())


class _ConfigContext:
    old_value: _Config
    new_value: _Config

    def __init__(self, new_value: _Config) -> None:
        global config

        logger.info("Config context: old={}({}), new={}({})", type(config), id(config), type(new_value), id(new_value))
        self.old_value = config
        self.new_value = new_value
        config = new_value

    def __enter__(self) -> Self:
        return self

    def __exit__(self, exc_type: object, exc_value: object, traceback: object) -> None:
        global config
        if config is not self.new_value:
            logger.warning("Config context: config changed during context; not restoring")
            return
        logger.info("Config context: restoring config to {}({})", type(self.old_value), id(self.old_value))
        config = self.old_value


def load(
    *,
    reload: bool = False,
    defaults_file: Path | IO[Any] = CONFIG_DEFAULTS_PATH,
    config_file: Path | IO[Any] | None = CONFIG_PATH,
    allow_module_level: bool = False,
) -> _ConfigContext:
    """
    Load the configuration.

    Do not call from module level (at import time).
    """

    global config

    if isinstance(config, _Config) and not reload:
        logger.warning("Configuration already loaded")
        return _ConfigContext(config)

    if not allow_module_level and misc.is_import_in_progress():
        raise ConfigurationNotLoadedError("Configuration loaded during import")

    def _read_file(fname: str, default: str) -> str:
        path = Path(fname).expanduser()
        if path.exists():
            return path.read_text().strip()
        return default

    with _with_resolver("file", _read_file):
        cfg = OmegaConf.structured(_Config)
        cfg = OmegaConf.merge(cfg, OmegaConf.load(defaults_file))

        do_load_config = config_file is not None
        if isinstance(config_file, Path) and not config_file.exists():
            logger.warning("Configuration file {} does not exist", config_file)
            do_load_config = False

        if do_load_config:
            logger.info("Loading configuration from {}", config_file)
            assert config_file is not None
            cfg = OmegaConf.merge(cfg, OmegaConf.load(config_file))
        else:
            logger.info("Not loading config: {}", config_file)

        cfg = OmegaConf.to_object(cfg)

    assert isinstance(cfg, _Config)

    # Some further validation
    assert (
        cfg.openai_api.completion_max_tokens
        <= cfg.openai_api.max_tokens_per_model[cfg.openai_api.chat_completion_model]
    )

    return _ConfigContext(cfg)


def load_test_config(config_file: Path | IO[Any] = Path("/nonexistent")) -> _ConfigContext:
    ctx = load(config_file=config_file, allow_module_level=True, reload=True)
    assert isinstance(config, _Config)
    return ctx


@attrs.frozen
class _FromConfig:
    key: str


FromConfig: Callable[[str], Any] = _FromConfig


@functools.cache
def _get_config_item(path: str) -> Any:
    cfg = config
    for part in path.split("."):
        cfg = getattr(cfg, part)
    return cfg


_T = TypeVar("_T")


def _check_config_key(path: str) -> None:
    cls: type[Any] = _Config
    for part in path.split("."):
        with contextlib.suppress(attrs.exceptions.NotAnAttrsClassError):
            cls = getattr(attrs.fields(cls), part).type


_Params = ParamSpec("_Params")


def uses_config(func: Callable[_Params, _T]) -> Callable[_Params, _T]:
    """Replace uses of FromConfig by values from config."""

    config_args: dict[str, Any] | None = None

    # Validate the config keys
    for arg_value in inspect.signature(func).parameters.values():
        if isinstance(arg_value.default, _FromConfig):
            _check_config_key(arg_value.default.key)

    def init() -> None:
        nonlocal config_args

        if misc.is_import_in_progress():
            raise ConfigurationNotLoadedError("Function with @uses_config decorator called during import")

        config_args = {}
        for arg_name, arg_value in inspect.signature(func).parameters.items():
            if isinstance(arg_value.default, _FromConfig):
                config_args[arg_name] = _get_config_item(arg_value.default.key)

    @functools.wraps(func)
    def wrapper(*args: _Params.args, **kwargs: _Params.kwargs) -> _T:
        if config_args is None:
            init()
        assert config_args is not None
        kwargs = {**config_args, **kwargs}  # type: ignore[assignment]
        return func(*args, **kwargs)

    return wrapper

rrosti/utils/config.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Miscellaneous utility functions and classes."""

from __future__ import annotations

import argparse
import asyncio
import functools
import hashlib
import inspect
import logging
import sys
from pathlib import Path
from typing import Any, Callable, Coroutine, Mapping, NoReturn, Protocol, TypeVar, no_type_check

import numpy as np
import numpy.typing as npt
from loguru import logger
from typing_extensions import Self

FloatArray = npt.NDArray[np.float32]

JSONPrimitive = str | int | float | bool | None
JSONValue = JSONPrimitive | Mapping[JSONPrimitive, "JSONValue"] | list["JSONValue"]


class _ProgramArgsBaseProtocol(Protocol):
    @classmethod
    def parse_args(cls) -> Self:
        """Parse program arguments and return an instance of this class."""
        ...


class ProgramArgsMixinProtocol(Protocol):
    @classmethod
    def _add_args(cls, parser: argparse.ArgumentParser) -> None:
        ...


class _ProgramArgsTypeProtocol(_ProgramArgsBaseProtocol, ProgramArgsMixinProtocol, Protocol):
    pass


_ProgramArgsType = TypeVar("_ProgramArgsType", bound=_ProgramArgsTypeProtocol)


class ProgramArgsBase:
    """
    A base class for parsing program arguments.

    The idea is that your main program can collect arguments from different modules like this:

    class ProgramArgs(ProgramArgsBase, module1.ProgramArgsMixin, module2.ProgramArgsMixin):
        main_program_arg: int

        @classmethod
        def _add_args(cls, parser: argparse.ArgumentParser) -> None:
            super()._add_args(parser)
            parser.add_argument("--main-program-arg", type=int, default=42)

    args = ProgramArgs.parse_args()

    The mixins look like this:

    class ProgramArgsMixin(ProgramArgsMixinProtocol):
        mixin_arg: int

        @classmethod
        def _add_args(cls, parser: argparse.ArgumentParser) -> None:
            parser.add_argument("--mixin-arg", type=int, default=42)
    """

    @classmethod
    def parse_args(cls: type[_ProgramArgsType]) -> _ProgramArgsType:  # noqa: PYI019 (cannot return Self here)
        parser = argparse.ArgumentParser()
        cls._add_args(parser)
        return parser.parse_args(namespace=cls())


class InterceptHandler(logging.Handler):
    @no_type_check
    def emit(self, record: Any) -> None:
        if logger is None:
            # Somehow we get called with this at the end of the program.
            return
        # Get corresponding Loguru level if it exists.
        try:
            level = logger.level(record.levelname).name
        except ValueError:
            level = record.levelno

        # Find the frame that originated the log message.
        frame, depth = sys._getframe(6), 6
        while frame and frame.f_code.co_filename == logging.__file__:
            frame = frame.f_back
            depth += 1

        logger.opt(depth=depth, exception=record.exc_info).log(level, record.getMessage())


def setup_logging(level: str = "DEBUG") -> None:
    logging.basicConfig(handlers=[InterceptHandler()], level=level)


def indent(text: str, prefix: str | int) -> str:
    if isinstance(prefix, int):
        prefix = " " * prefix
    return "\n".join(prefix + line for line in text.splitlines())


def truncate_string(s: str, length: int = 68) -> str:
    """Truncate a string to a maximum length."""
    if len(s) <= length:
        return s
    return s[:length] + "..."


def truncate_json_strings(obj: JSONValue, length: int = 68) -> JSONValue:
    """Truncate all strings in a JSON object to a maximum length."""
    if isinstance(obj, str):
        return truncate_string(obj, length)
    if isinstance(obj, list):
        return [truncate_json_strings(item, length) for item in obj]
    if isinstance(obj, dict):
        return {key: truncate_json_strings(value, length) for key, value in obj.items()}
    return obj


_T = TypeVar("_T")


def file_sha256(path: Path) -> str:
    sha256 = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            sha256.update(chunk)
    return sha256.hexdigest()


def async_once_blocking(
    func: Callable[[], Coroutine[NoReturn, NoReturn, _T]]
) -> Callable[[], Coroutine[NoReturn, NoReturn, _T]]:
    """
    Decorator for a coroutine that should only be run once, and then cached.

    The coroutine is run in the event loop, so it should be fast.
    """

    task: asyncio.Task[_T] | None = None

    @functools.wraps(func)
    async def wrapper() -> _T:
        nonlocal task
        if task is None:
            task = asyncio.create_task(func())
        return await task

    return wrapper


def async_once_in_thread(func: Callable[[], _T]) -> Callable[[], Coroutine[NoReturn, NoReturn, _T]]:
    """
    Decorator for a coroutine that should only be run once, and then cached.

    The coroutine is run in a thread, so it can be slow and/or I/O bound.
    """

    task: asyncio.Task[_T] | None = None

    @functools.wraps(func)
    async def wrapper() -> _T:
        nonlocal task
        if task is None:
            task = asyncio.create_task(asyncio.to_thread(func))
        return await task

    return wrapper


def is_import_in_progress() -> bool:
    return any("importlib._bootstrap" in frame_info.filename for frame_info in inspect.stack())

rrosti/utils/misc.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import tiktoken

PRELOAD = True


def token_count(text: str) -> int:
    enc = tiktoken.get_encoding("cl100k_base")
    return len(enc.encode(text))


if PRELOAD:
    tiktoken.get_encoding("cl100k_base")

rrosti/utils/token_count.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from typing import Iterable

import aioresponses
import numpy as np
import pytest
import requests_mock as rmock
from pytest_mock import MockerFixture

from rrosti.chat import chat_session
from rrosti.chat.state_machine import execution, interpolable, parsing
from rrosti.llm_api import openai_api, openai_api_direct
from rrosti.servers.websocket_query_server import Frontend
from rrosti.snippets.abstract_snippet_database import AbstractSnippetDatabase
from rrosti.snippets.snippet import Snippet


def make_state_machine_runner(mocker: MockerFixture, code: str) -> execution.StateMachineRunner:
    openai_provider = openai_api_direct.DirectOpenAIApiProvider()
    llm: chat_session.LLM = mocker.Mock(spec=chat_session.LLM, openai_provider=openai_provider)
    llm.chat_completion.return_value = chat_session.Message(  # type: ignore[attr-defined]
        role="assistant",
        importance=chat_session.Importance.LOW,
        ttl=None,
        text="Hello, world from chat completion!",
    )

    frontend: Frontend = mocker.Mock(spec=Frontend)
    sm = parsing.loads_from_yaml(code)
    return execution.StateMachineRunner(sm=sm, llm=llm, frontend=frontend, openai_provider=openai_provider)


@pytest.fixture(autouse=True)
def _no_network(requests_mock: rmock.Mocker) -> Iterable[None]:
    with aioresponses.aioresponses():
        yield


async def test_execution(mocker: MockerFixture) -> None:
    runner = make_state_machine_runner(
        mocker,
        r"""
    agents:
    -   name: agent1
        states:
        -   name: initial
            action:
            -   goto: some_state
        -   name: some_state
            conditions:
            -   default:
                    action:
                    -   end
        """,
    )

    await runner.run()


SIMPLE_PYTHON_YAML = r"""
agents:
-   name: agent1
    states:
    -   name: initial
        action:
        -   goto: some_state
    -   name: some_state
        conditions:
        -   default:
                action:
                -   message: "The Python output is:\n\n{python()}\n\nEnd of Python output."
                -   end
    """


async def test_python(mocker: MockerFixture) -> None:
    runner = make_state_machine_runner(mocker, SIMPLE_PYTHON_YAML)

    # The LLM gives us a message with a Python block.
    runner._llm.chat_completion.return_value = chat_session.Message(  # type: ignore[attr-defined]
        role="assistant",
        importance=chat_session.Importance.LOW,
        ttl=1,
        text=f"""
Let's execute some Python.

$$$python
print(f"Hello, world! 1+1={1+1}")
$$$

How does that look?
        """.strip(),
    )

    await runner.run()

    # The runner should have appended a message with the Python output.
    msg = runner._agent_runners[0]._session.messages[-1]
    assert msg.role == "user"
    assert msg.text == "The Python output is:\n\nHello, world! 1+1=2\n\nEnd of Python output."


def assert_no_python_block_found_error(runner: execution.StateMachineRunner) -> None:
    assert len(runner._agent_runners[0]._session.messages) == 2
    msg = runner._agent_runners[0]._session.messages[-1].text
    assert msg.startswith("The Python output is:\n\n")
    assert msg.endswith("\n\nEnd of Python output.")
    msg = msg.removeprefix("The Python output is:\n\n").removesuffix("\n\nEnd of Python output.")
    assert msg.startswith("Your message did not contain a Python code block.")


async def test_no_python_block(mocker: MockerFixture) -> None:
    runner = make_state_machine_runner(mocker, SIMPLE_PYTHON_YAML)

    # The LLM gives us a message with a Python block.
    runner._llm.chat_completion.return_value = chat_session.Message(  # type: ignore[attr-defined]
        role="assistant",
        ttl=None,
        importance=chat_session.Importance.LOW,
        text="""
Let's execute some Python.

How does that look?
        """.strip(),
    )

    await runner.run()
    assert_no_python_block_found_error(runner)


async def test_unterminated_python_block(mocker: MockerFixture) -> None:
    runner = make_state_machine_runner(mocker, SIMPLE_PYTHON_YAML)

    # The LLM gives us a message with a Python block.
    runner._llm.chat_completion.return_value = chat_session.Message(  # type: ignore[attr-defined]
        role="assistant",
        importance=chat_session.Importance.LOW,
        ttl=1,
        text="""
Let's execute some Python.

$$$python
print("foo")

How does that look?
        """.strip(),
    )

    await runner.run()
    assert_no_python_block_found_error(runner)


SIMPLE_RTFM_YAML = r"""
agents:
-   name: agent1
    states:
    -   name: initial
        action:
        -   system_message: "Here are your instructions."
        -   goto: some_state
    -   name: some_state
        conditions:
        -   default:
                action:
                -   message: "The RTFM output is:\n\n{rtfm()}\n\nEnd of RTFM output."
                -   end
    """


@pytest.fixture
def snippet_db(mocker: MockerFixture) -> AbstractSnippetDatabase:
    mocker.patch("rrosti.chat.state_machine.interpolable._get_database")
    db = mocker.Mock(spec=AbstractSnippetDatabase)
    db.find_nearest_merged.return_value = [
        Snippet("hello, world", source_filename="source1", start_offset=0, page_start=1, page_end=1),
        Snippet("goodbye, cruel world", source_filename="source2", start_offset=10, page_start=10, page_end=10),
    ]

    interpolable._get_database.return_value = db  # type: ignore[attr-defined]
    return db  # type: ignore[no-any-return]


def assert_rtfm_output_is(runner: execution.StateMachineRunner, expected: str) -> None:
    assert len(runner._agent_runners[0]._session.messages) == 3
    msg = runner._agent_runners[0]._session.messages[-1].text
    assert msg.startswith("The RTFM output is:\n\n")
    assert msg.endswith("\n\nEnd of RTFM output.")
    msg = msg.removeprefix("The RTFM output is:\n\n").removesuffix("\n\nEnd of RTFM output.")
    assert msg == expected


async def mock_query_embedding_async(
    openai_provider: openai_api.OpenAIApiProvider, snippets: list[str]
) -> openai_api.EmbeddingResponse:
    return openai_api.EmbeddingResponse(
        snippets=snippets[:],
        embeddings=np.ones((len(snippets), 123), dtype=np.float32),
        model="fake_model",
        prompt_tokens=42,
    )


async def test_rtfm_stub(mocker: MockerFixture, snippet_db: AbstractSnippetDatabase) -> None:
    runner = make_state_machine_runner(mocker, SIMPLE_RTFM_YAML)

    # The LLM gives us a message with an rtfm block
    runner._llm.chat_completion.return_value = chat_session.Message(  # type: ignore[attr-defined]
        role="assistant",
        ttl=1,
        importance=chat_session.Importance.LOW,
        text="""
Let's execute some RTFM.

$$$rtfm
Something here.
$$$

How does that look?
        """.strip(),
    )

    EXPECTED_OUTPUT = """
Extract #123:

hello, world
-----
Extract #124:

goodbye, cruel world
""".strip()

    mocker.patch.object(openai_api_direct.DirectOpenAIApiProvider, "acreate_embedding", new=mock_query_embedding_async)
    runner.frontend.handle_rtfm_output.return_value = 123  # type: ignore[attr-defined]

    await runner.run()

    assert_rtfm_output_is(runner, EXPECTED_OUTPUT)

tests/chat/state_machine/test_execution.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from pprint import pformat

import pytest
from loguru import logger
from marshmallow import ValidationError

import rrosti.utils.config
from rrosti.chat.state_machine import ast, parsing
from rrosti.utils.config import config

# This loads the default config
rrosti.utils.config.load_test_config()


def test_minimal_example() -> None:
    dsm = parsing.loads_from_yaml(
        """
    agents:
    -   name: agent1
        states:
        -   name: initial
            action:
            -   goto: some_state
        -   name: some_state
            conditions:
            -   default:
                    action:
                    -   end
        """
    )

    assert isinstance(dsm, ast.StateMachine)
    assert isinstance(dsm.variables, ast.Variables)
    assert dsm.variables.vars == {}
    assert len(dsm.agents) == 1
    assert dsm.agents[0].name == "agent1"
    agent = dsm.agents[0]
    assert isinstance(agent.initial_state, ast.InitialState)
    assert agent.initial_state.action == ast.ActionList([ast.GotoAction(label="some_state")])
    assert len(agent.noninitial_states) == 1

    state = agent.noninitial_states[0]
    assert isinstance(state, ast.NonInitialState)
    assert state.name == "some_state"
    assert len(state.conditions) == 1

    cond = state.conditions[0]
    assert isinstance(cond, ast.DefaultCondition)
    assert cond.action == ast.ActionList([ast.EndAction()])


def test_goto_does_not_exist() -> None:
    with pytest.raises(ValidationError) as excinfo:
        parsing.loads_from_yaml(
            """
        agents:
        -   name: agent1
            states:
            -   name: initial
                action:
                -   goto: nonexistent_state
            -   name: some_state
                conditions:
                -   default:
                        action:
                        -   end
        """
        )

    assert "State 'nonexistent_state' does not exist" in str(excinfo.value)


def test_goto() -> None:
    dsm = parsing.loads_from_yaml(
        """
    agents:
    -   name: agent1
        states:
        -   name: initial
            action:
            -   goto: some_state
        -   name: some_state
            conditions:
            -   default:
                    action:
                    -   end
    """
    )

    assert isinstance(dsm, ast.StateMachine)
    assert isinstance(dsm.variables, ast.Variables)
    assert dsm.variables.vars == {}
    assert len(dsm.agents) == 1
    assert dsm.agents[0].name == "agent1"
    agent = dsm.agents[0]
    assert isinstance(agent.initial_state, ast.InitialState)
    assert agent.initial_state.action == ast.ActionList([ast.GotoAction(label="some_state")])
    assert len(agent.noninitial_states) == 1

    state = agent.noninitial_states[0]
    assert isinstance(state, ast.NonInitialState)
    assert state.name == "some_state"
    assert len(state.conditions) == 1

    cond = state.conditions[0]
    assert isinstance(cond, ast.DefaultCondition)
    assert cond.action == ast.ActionList([ast.EndAction()])


def test_undefined_variable() -> None:
    with pytest.raises(ValidationError) as excinfo:
        parsing.loads_from_yaml(
            """
        agents:
        -   name: agent1
            states:
            -   name: initial
                action:
                -   message: "{undefined_variable}"
                -   goto: some_state
            -   name: some_state
                conditions:
                -   default:
                        action:
                        -   end
        """
        )

    assert "Undefined variable" in str(excinfo.value)


def test_expand_variable_in_variable() -> None:
    dsm = parsing.loads_from_yaml(
        """
    variables:
        variable1: "default value"
        variable2: "{variable1}"
    agents:
    -   name: agent1
        states:
        -   name: initial
            action:
            -   message: "{variable2}"
            -   goto: some_state
        -   name: some_state
            conditions:
            -   default:
                    action:
                    -   end
    """
    )

    assert isinstance(dsm, ast.StateMachine)
    assert isinstance(dsm.variables, ast.Variables)
    assert dsm.variables.vars == {"variable1": "default value", "variable2": "default value"}
    agent = dsm.agents[0]
    assert isinstance(agent.initial_state, ast.InitialState)
    assert agent.initial_state.action == ast.ActionList(
        [ast.MessageAction(role="user", text="default value"), ast.GotoAction(label="some_state")]
    )


def test_self_referencing_variable() -> None:
    with pytest.raises(ValidationError) as excinfo:
        parsing.loads_from_yaml(
            """
        variables:
            variable1: "{variable1}"
        agents:
        -   name: agent1
            states:
            -   name: initial
                action:
                -   system_message: "Hello system!"
                -   message: "{variable1}"
                -   goto: some_state
            -   name: some_state
                conditions:
                -   default:
                        action:
                        -   end
            """
        )

    assert "Variable expansion loop detected" in str(excinfo.value)


def test_circular_variable() -> None:
    with pytest.raises(ValidationError) as excinfo:
        parsing.loads_from_yaml(
            """
        variables:
            variable1: "{variable2}"
            variable2: "{variable1}"
        agents:
        -   name: agent1
            states:
            -   name: initial
                action:
                -   message: "{variable1}"
                -   goto: some_state
            -   name: some_state
                conditions:
                -   default:
                        action:
                        -   end
        """
        )

    assert "Variable expansion loop detected" in str(excinfo.value)


def test_variable() -> None:
    dsm = parsing.loads_from_yaml(
        """
    variables:
        variable1: "default value"
    agents:
    -   name: agent1
        states:
        -   name: initial
            action:
            -   message: "{variable1}"
            -   goto: some_state
        -   name: some_state
            conditions:
            -   default:
                    action:
                    -   end
    """
    )

    assert isinstance(dsm, ast.StateMachine)
    assert isinstance(dsm.variables, ast.Variables)
    assert dsm.variables.vars == {"variable1": "default value"}
    agent = dsm.agents[0]
    assert isinstance(agent.initial_state, ast.InitialState)
    assert agent.initial_state.action == ast.ActionList(
        [ast.MessageAction(role="user", text="default value"), ast.GotoAction(label="some_state")]
    )
    assert len(agent.noninitial_states) == 1

    state = agent.noninitial_states[0]
    assert isinstance(state, ast.NonInitialState)
    assert state.name == "some_state"
    assert len(state.conditions) == 1

    cond = state.conditions[0]
    assert isinstance(cond, ast.DefaultCondition)
    assert cond.action == ast.ActionList([ast.EndAction()])


def test_state_not_reachable() -> None:
    with pytest.raises(ValidationError) as excinfo:
        parsing.loads_from_yaml(
            """
        agents:
        -   name: agent1
            states:
            -   name: initial
                action:
                -   goto: state1
            -   name: state1
                conditions:
                -   default:
                        action:
                        -   end
            -   name: unreachable_state
                conditions:
                -   default:
                        action:
                        -   end
        """
        )

    assert " are unreachable" in str(excinfo.value)


# Also test with two unreachable states that are connected to each other
def test_state_not_reachable2() -> None:
    with pytest.raises(ValidationError) as excinfo:
        parsing.loads_from_yaml(
            """
        agents:
        -   name: agent1
            states:
            -   name: initial
                action:
                -   goto: state1
            -   name: state1
                conditions:
                -   default:
                        action:
                        -   goto: state1
            -   name: unreachable_state1
                conditions:
                -   default:
                        action:
                        -   goto: unreachable_state2
            -   name: unreachable_state2
                conditions:
                -   default:
                        action:
                        -   goto: unreachable_state1
        """
        )

    assert " are unreachable" in str(excinfo.value)


def test_duplicate_state_name() -> None:
    with pytest.raises(ValidationError) as excinfo:
        parsing.loads_from_yaml(
            """
        agents:
        -   name: agent1
            states:
            -   name: initial
                action:
                -   goto: state1
            -   name: state1
                conditions:
                -   default:
                        action:
                        -   goto: state2
            -   name: state2
                conditions:
                -   default:
                        action:
                        -   goto: state3
            -   name: state3
                conditions:
                -   default:
                        action:
                        -   goto: state1
            -   name: state1
                conditions:
                -   default:
                        action:
                        -   goto: state2
        """
        )

    assert "Duplicate state name: state1" in str(excinfo.value)


def test_duplicate_agent_name() -> None:
    with pytest.raises(ValidationError) as excinfo:
        parsing.loads_from_yaml(
            """
        agents:
        -   name: agent1
            states:
            -   name: initial
                action:
                -   goto: state1
            -   name: state1
                conditions:
                -   default:
                        action:
                        -   goto: state2
            -   name: state2
                conditions:
                -   default:
                        action:
                        -   goto: state3
            -   name: state3
                conditions:
                -   default:
                        action:
                        -   goto: state1
        -   name: agent1
            states:
            -   name: initial
                action:
                -   goto: state1
            -   name: state1
                conditions:
                -   default:
                        action:
                        -   goto: state2
            -   name: state2
                conditions:
                -   default:
                        action:
                        -   goto: state3
            -   name: state3
                conditions:
                -   default:
                        action:
                        -   goto: state1
        """
        )

    assert "Duplicate agent name: agent1" in str(excinfo.value)


def test_first_state_not_initial() -> None:
    with pytest.raises(ValidationError) as excinfo:
        parsing.loads_from_yaml(
            """
        agents:
        -   name: agent1
            states:
            -   name: state1
                conditions:
                -   default:
                        action:
                        -   goto: state2
            -   name: state2
                conditions:
                -   default:
                        action:
                        -   goto: state3
            -   name: state3
                conditions:
                -   default:
                        action:
                        -   goto: state1
        """
        )

    # assert "The first state must be initial" in str(excinfo.value)
    assert "'initial' was expected" in str(excinfo.value)


def test_no_states() -> None:
    with pytest.raises(ValidationError) as excinfo:
        parsing.loads_from_yaml(
            """
        agents:
        -   states: []
        """
        )

    # assert "There must be at least one state" in str(excinfo.value)
    assert "[] is too short" in str(excinfo.value)


def test_send_action() -> None:
    dsm = parsing.loads_from_yaml(
        """
    agents:
    -   name: agent1
        states:
        -   name: initial
            action:
            -   goto: some_state
        -   name: some_state
            conditions:
            -   default:
                    action:
                    -   send:
                            to: agent2
                            next_state: some_state
    -   name: agent2
        states:
        -   name: initial
            action:
            -   goto: some_state
        -   name: some_state
            conditions:
            -   default:
                    action:
                    -   send:
                            to: agent1
                            next_state: some_state
    """
    )

    assert isinstance(dsm, ast.StateMachine)
    agent = dsm.agents[0]
    assert isinstance(agent.initial_state, ast.InitialState)
    assert agent.initial_state.action == ast.ActionList([ast.GotoAction(label="some_state")])
    assert agent._states_by_name["some_state"].conditions[0].action == ast.ActionList(
        [
            ast.SendAction(
                to="agent2",
                next_state="some_state",
            )
        ]
    )


def test_state_with_model() -> None:
    dsm = parsing.loads_from_yaml(
        """
    agents:
    -   name: agent1
        states:
        -   name: initial
            action:
            -   goto: some_state
        -   name: some_state
            model: some_model
            conditions:
            -   default:
                    action:
                    -   end
    """
    )

    assert isinstance(dsm, ast.StateMachine)
    agent = dsm.agents[0]
    assert isinstance(agent.initial_state, ast.InitialState)
    assert agent.initial_state.action == ast.ActionList([ast.GotoAction(label="some_state")])
    assert agent._states_by_name["some_state"].config.model_ast == "some_model"


def test_if_then() -> None:
    dsm = parsing.loads_from_yaml(
        """
    agents:
    -   name: agent1
        states:
        -   name: initial
            action:
            -   message: some_message
            -   goto: some_state
        -   name: some_state
            conditions:
            -   if:
                    contains: some_substring
                then:
                    model: some_model
                    action:
                    -   goto: some_state
            -   default:
                    action:
                    -   end
    """
    )

    assert isinstance(dsm, ast.StateMachine)
    agent = dsm.agents[0]
    assert isinstance(agent.initial_state, ast.InitialState)
    assert agent.initial_state.action == ast.ActionList(
        [
            ast.MessageAction(text="some_message", role="user"),
            ast.GotoAction("some_state"),
        ]
    )
    assert len(agent._states_by_name["some_state"].conditions) == 2
    assert agent._states_by_name["some_state"].conditions[0].action == ast.ActionList(
        [
            ast.GotoAction("some_state"),
        ]
    )
    assert agent._states_by_name["some_state"].conditions[0].config.model_ast == "some_model"
    assert agent._states_by_name["some_state"].conditions[0].config.model == "some_model"
    assert agent._states_by_name["some_state"].conditions[1].action == ast.ActionList(
        [
            ast.EndAction(),
        ]
    )
    assert agent._states_by_name["some_state"].conditions[1].config.model_ast is None
    assert agent._states_by_name["some_state"].conditions[1].config.model is None


def test_send_to_self_is_error() -> None:
    with pytest.raises(ValidationError) as excinfo:
        parsing.loads_from_yaml(
            """
        agents:
        -   name: agent1
            states:
            -   name: initial
                action:
                -   goto: some_state
            -   name: some_state
                conditions:
                -   default:
                        action:
                        -   send:
                                to: agent1
                                next_state: some_state
        """
        )

    assert "Agent agent1 sends to itself." in str(excinfo.value)


def test_send_to_unknown_agent_is_error() -> None:
    with pytest.raises(ValidationError) as excinfo:
        parsing.loads_from_yaml(
            """
        agents:
        -   name: agent1
            states:
            -   name: initial
                action:
                -   goto: some_state
            -   name: some_state
                conditions:
                -   default:
                        action:
                        -   send:
                                to: agent2
                                next_state: some_state
        """
        )

    assert "Unknown agent: agent2" in str(excinfo.value)


def test_send_to_unknown_state_is_error() -> None:
    with pytest.raises(ValidationError) as excinfo:
        parsing.loads_from_yaml(
            """
        agents:
        -   name: agent1
            states:
            -   name: initial
                action:
                -   goto: some_state
            -   name: some_state
                conditions:
                -   default:
                        action:
                        -   send:
                                to: agent1
                                next_state: some_state2
        """
        )

    assert "State 'some_state2' does not exist" in str(excinfo.value)


def test_send_must_be_last_action() -> None:
    with pytest.raises(ValidationError) as excinfo:
        parsing.loads_from_yaml(
            """
        agents:
        -   name: agent1
            states:
            -   name: initial
                action:
                -   message: some_message
                -   goto: some_state
            -   name: some_state
                conditions:
                -   default:
                        action:
                        -   send:
                                to: agent2
                                next_state: some_state
                        -   end
        -   name: agent2
            states:
            -   name: initial
                action:
                -   goto: some_state
            -   name: some_state
                conditions:
                -   default:
                        action:
                        -   end
        """
        )

    assert "Only the last action may be a send." in str(excinfo.value)


def test_send_not_allowed_in_initial() -> None:
    with pytest.raises(ValidationError) as excinfo:
        parsing.loads_from_yaml(
            """
        agents:
        -   name: agent1
            states:
            -   name: initial
                action:
                -   send:
                        to: agent2
                        next_state: some_state
            -   name: some_state
                conditions:
                -   default:
                        action:
                        -   end
        -   name: agent2
            states:
            -   name: initial
                action:
                -   goto: some_state
            -   name: some_state
                conditions:
                -   default:
                        action:
                        -   end
        """
        )

    assert "Send action not allowed in initial state" in str(excinfo.value)


def test_model_in_global_config() -> None:
    dsm = parsing.loads_from_yaml(
        """
    config:
        model: some_model
    agents:
    -   name: agent1
        states:
        -   name: initial
            action:
            -   goto: some_state
        -   name: some_state
            conditions:
            -   default:
                    action:
                    -   end
    """
    )

    assert isinstance(dsm, ast.StateMachine)
    assert dsm.config.model_ast == "some_model"
    agent = dsm.agents[0]
    assert isinstance(agent.initial_state, ast.InitialState)
    assert agent.initial_state.action == ast.ActionList([ast.GotoAction(label="some_state")])
    assert agent._states_by_name["some_state"].config.model_ast is None
    assert agent._states_by_name["some_state"].config.model == "some_model"


def test_override_model_in_agent() -> None:
    dsm = parsing.loads_from_yaml(
        """
    config:
        model: some_model
    agents:
    -   name: agent1
        model: some_model2
        states:
        -   name: initial
            action:
            -   goto: some_state
        -   name: some_state
            conditions:
            -   default:
                    action:
                    -   end
    """
    )

    logger.info("dsm:\n{}", pformat(dsm))

    assert isinstance(dsm, ast.StateMachine)
    assert dsm.config.model_ast == "some_model"
    agent = dsm.agents[0]
    assert isinstance(agent.initial_state, ast.InitialState)
    assert agent.initial_state.action == ast.ActionList([ast.GotoAction(label="some_state")])
    assert agent.config.model_ast == "some_model2"
    assert agent.config.model == "some_model2"
    assert agent._states_by_name["some_state"].config.model == "some_model2"


def test_override_modle_in_condition() -> None:
    dsm = parsing.loads_from_yaml(
        """
    config:
        model: some_model
    agents:
    -   name: agent1
        model: some_model2
        states:
        -   name: initial
            action:
            -   goto: some_state
        -   name: some_state
            conditions:
            -   default:
                    model: some_model3
                    action:
                    -   end
    """
    )

    logger.info("dsm:\n{}", pformat(dsm))

    assert isinstance(dsm, ast.StateMachine)
    assert dsm.config.model_ast == "some_model"
    agent = dsm.agents[0]
    assert isinstance(agent.initial_state, ast.InitialState)
    assert agent.initial_state.action == ast.ActionList([ast.GotoAction(label="some_state")])
    assert agent.config.model_ast == "some_model2"
    assert agent.config.model == "some_model2"
    assert agent._states_by_name["some_state"].config


def test_load_assets_prompt() -> None:
    """Test that the default prompt.yaml can be loaded."""

    yaml_data = config.state_machine.yaml_path.read_text()
    dsm = parsing.loads_from_yaml(yaml_data)
    assert isinstance(dsm, ast.StateMachine)


# -*- mode: python; indent-tabs-mode: nil; tab-width: 4 -*-

tests/chat/state_machine/test_parsing.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import threading
from typing import AsyncIterable, Iterable

import aioresponses
import openai
import pytest
from pytest_mock import MockerFixture

import rrosti.utils.config
from rrosti.llm_api import openai_api, openai_api_direct
from rrosti.utils.config import config

rrosti.utils.config.load_test_config()


@pytest.fixture(autouse=True)
def _no_network() -> Iterable[None]:
    with aioresponses.aioresponses():
        yield


FAKE_CHAT_COMPLETION: openai_api._ChatCompletionResponseDict = {
    "id": "chatcmpl-abc123",
    "object": "chat.completion",
    "created": 1677858242,
    "model": "gpt-3.5-turbo-0613",
    "usage": {"prompt_tokens": 13, "completion_tokens": 7, "total_tokens": 20},
    "choices": [
        {"message": {"role": "assistant", "content": "\n\nThis is a test!"}, "finish_reason": "stop", "index": 0}
    ],
}


FAKE_EMBEDDING_RESPONSE: openai_api._EmbeddingResponseDict = {
    "object": "list",
    "data": [{"object": "embedding", "embedding": [0.0] * 512, "index": 0}],
    "model": "text-embedding-ada-002",
    "usage": {"prompt_tokens": 8, "total_tokens": 8, "completion_tokens": 0},
}


@pytest.fixture
async def mock_chat_worker(mocker: MockerFixture) -> AsyncIterable[openai_api_direct._ChatCompletionWorker]:
    openai_api_direct._ChatCompletionWorker.delete_instance()

    # Patch the number of threads just to make backtraces nicer
    mocker.patch.object(config.openai_api, "num_chat_completion_worker_threads", 3)

    # No need to sleep so long in tests
    mocker.patch.object(config.openai_api, "rate_limit_initial_interval", 0.02)

    mocker.patch.object(openai_api_direct, "_create_chat_completion_internal", return_value=FAKE_CHAT_COMPLETION)
    yield openai_api_direct._ChatCompletionWorker.get_instance()
    openai_api_direct._ChatCompletionWorker.delete_instance()


@pytest.fixture
async def mock_embedding_worker(mocker: MockerFixture) -> AsyncIterable[openai_api_direct._EmbeddingWorker]:
    openai_api_direct._EmbeddingWorker.delete_instance()

    # Patch the number of threads just to make backtraces nicer
    mocker.patch.object(config.openai_api, "num_embedding_worker_threads", 3)

    # No need to sleep so long in tests
    mocker.patch.object(config.openai_api, "rate_limit_initial_interval", 0.02)
    mocker.patch.object(openai_api_direct, "_create_embedding_internal", return_value=FAKE_EMBEDDING_RESPONSE)
    yield openai_api_direct._EmbeddingWorker.get_instance()
    openai_api_direct._EmbeddingWorker.delete_instance()


@pytest.fixture
async def mock_provider(
    mock_chat_worker: openai_api_direct._ChatCompletionWorker, mock_embedding_worker: openai_api_direct._EmbeddingWorker
) -> openai_api_direct.DirectOpenAIApiProvider:
    return openai_api_direct.DirectOpenAIApiProvider()


async def test_chat_completion_worker_async(mock_provider: openai_api_direct.DirectOpenAIApiProvider) -> None:
    result = await mock_provider.acreate_chat_completion(
        model="test_model",
        messages=[{"role": "system", "content": "Hello, world!"}],
    )

    assert result is FAKE_CHAT_COMPLETION


async def test_chat_completion_worker_sync(mock_provider: openai_api_direct.DirectOpenAIApiProvider) -> None:
    # janus.Queue needs an event loop, hence this is async

    result = mock_provider.create_chat_completion(
        model="test_model",
        messages=[{"role": "system", "content": "Hello, world!"}],
    )

    assert result is FAKE_CHAT_COMPLETION


async def test_embedding_worker_async(
    mock_provider: openai_api_direct.DirectOpenAIApiProvider,
) -> None:
    result = await mock_provider.acreate_embedding(input=["Hello, world!"])

    assert result.embeddings.shape == (1, 512)


async def test_embedding_worker_sync(mock_provider: openai_api_direct.DirectOpenAIApiProvider) -> None:
    result = mock_provider.create_embedding(input=["Hello, world!"])

    assert result.embeddings.shape == (1, 512)


class MyTestException(Exception):
    pass


async def test_chat_completion_worker_exception_async(mock_provider: openai_api_direct.DirectOpenAIApiProvider) -> None:
    openai_api_direct._create_chat_completion_internal.side_effect = MyTestException()  # type: ignore[attr-defined]
    with pytest.raises(MyTestException):
        await mock_provider.acreate_chat_completion(
            model="test_model",
            messages=[{"role": "system", "content": "Hello, world!"}],
        )


async def test_chat_completion_worker_exception_sync(
    mock_provider: openai_api_direct.DirectOpenAIApiProvider,
) -> None:
    openai_api_direct._create_chat_completion_internal.side_effect = MyTestException()  # type: ignore[attr-defined]
    with pytest.raises(MyTestException):
        mock_provider.create_chat_completion(
            model="test_model",
            messages=[{"role": "system", "content": "Hello, world!"}],
        )


class MockTime:
    _now: float

    def __init__(self) -> None:
        self._now = 1000.0

    def time(self) -> float:
        return self._now

    def sleep(self, duration: float) -> None:
        self._now += duration

    async def sleep_async(self, duration: float) -> None:
        self._now += duration

    def wait_cond_with_timeout(self, cond: threading.Condition, timeout: float) -> bool:
        timed_out = cond.wait(timeout=timeout)
        if timed_out:
            self._now += timeout
        return timed_out


@pytest.fixture
def _mock_time(mocker: MockerFixture) -> None:
    _mock_time = MockTime()
    mocker.patch.object(openai_api_direct._SystemAdapter.time, "time", wraps=_mock_time.time)


@pytest.mark.timeout(3)
@pytest.mark.usefixtures("_mock_time")
async def test_ratelimit_error_async(mock_provider: openai_api_direct.DirectOpenAIApiProvider) -> None:
    openai_api_direct._create_chat_completion_internal.side_effect = [  # type: ignore[attr-defined]
        openai.error.RateLimitError(),  # type: ignore[no-untyped-call]
        FAKE_CHAT_COMPLETION,
    ]
    await mock_provider.acreate_chat_completion(
        model="test_model",
        messages=[{"role": "system", "content": "Hello, world!"}],
    )

    # TODO: a proper test
    # openai_api._SystemAdapter.time.sleep.assert_called_once_with(  # type: ignore[attr-defined]
    #     openai_api.params.RATE_LIMIT_INITIAL_INTERVAL * openai_api.params.RATE_LIMIT_ERROR_MULTIPLIER
    # )


@pytest.mark.timeout(3)
@pytest.mark.usefixtures("_mock_time")
async def test_ratelimit_error_sync(mock_provider: openai_api_direct.DirectOpenAIApiProvider) -> None:
    openai_api_direct._create_chat_completion_internal.side_effect = [  # type: ignore[attr-defined]
        openai.error.RateLimitError(),  # type: ignore[no-untyped-call]
        FAKE_CHAT_COMPLETION,
    ]
    mock_provider.create_chat_completion(
        model="test_model",
        messages=[{"role": "system", "content": "Hello, world!"}],
    )

    # TODO: a proper test
    # openai_api._SystemAdapter.time.sleep.assert_called_once_with(  # type: ignore[attr-defined]
    #     openai_api.params.RATE_LIMIT_INITIAL_INTERVAL * openai_api.params.RATE_LIMIT_ERROR_MULTIPLIER
    # )


@pytest.mark.timeout(3)
@pytest.mark.usefixtures("_mock_time")
async def test_ratelimit_success_async(mock_provider: openai_api_direct.DirectOpenAIApiProvider) -> None:
    openai_api_direct._create_chat_completion_internal.side_effect = [  # type: ignore[attr-defined]
        FAKE_CHAT_COMPLETION,
        FAKE_CHAT_COMPLETION,
    ]
    await mock_provider.acreate_chat_completion(
        model="test_model",
        messages=[{"role": "system", "content": "Hello, world!"}],
    )
    await mock_provider.acreate_chat_completion(
        model="test_model",
        messages=[{"role": "system", "content": "Hello, world!"}],
    )

    # TODO: a proper test
    # openai_api._SystemAdapter.time.sleep.assert_called_once_with(  # type: ignore[attr-defined]
    #     openai_api.params.RATE_LIMIT_INITIAL_INTERVAL * openai_api.params.RATE_LIMIT_SUCCESS_MULTIPLIER
    # )


@pytest.mark.timeout(3)
@pytest.mark.usefixtures("_mock_time")
async def test_ratelimit_success_sync(mock_provider: openai_api_direct.DirectOpenAIApiProvider) -> None:
    openai_api_direct._create_chat_completion_internal.side_effect = [  # type: ignore[attr-defined]
        FAKE_CHAT_COMPLETION,
        FAKE_CHAT_COMPLETION,
    ]
    mock_provider.create_chat_completion(
        model="test_model",
        messages=[{"role": "system", "content": "Hello, world!"}],
    )
    mock_provider.create_chat_completion(
        model="test_model",
        messages=[{"role": "system", "content": "Hello, world!"}],
    )

    # TODO: a proper test
    # openai_api._SystemAdapter.time.sleep.assert_called_once_with(  # type: ignore[attr-defined]
    #     openai_api.params.RATE_LIMIT_INITIAL_INTERVAL * openai_api.params.RATE_LIMIT_SUCCESS_MULTIPLIER
    # )

tests/llm_api/test_openai_api.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import io
import os
import pprint
from dataclasses import dataclass
from pathlib import Path
from typing import Any, BinaryIO, Iterable
from unittest.mock import ANY, _Call, call

import orjson
import overrides
import pytest
import zstandard as zstd
from loguru import logger
from pytest_mock import MockerFixture

import rrosti.query.logging.query_logging as qlog
from tests import testing_utils

FAKE_RUN_UUID = "123"
FAKE_TIME = "20230101_123456"

LOG_DIR = Path("logs")
RUN_LOG_DIR = LOG_DIR / f"{FAKE_TIME}__{FAKE_RUN_UUID}"


FileDict = dict[Path, io.BytesIO]


@pytest.fixture
def query_logging_cleanup(mocker: MockerFixture) -> Iterable[FileDict]:
    """Reset the query_logging module state and mock out the filesystem operations."""

    global qlog

    # Reload the module as a guard in addition to calling _cleanup() (which removes the loguru
    # handlers and is necessary in any case.)
    (qlog,) = testing_utils.reload_modules("rrosti.query.logging.query_logging")

    files: FileDict = {}
    rotation_number: int = 0

    def mock_open_xb(path: Path) -> io.BytesIO:
        assert path not in files, path
        f = io.BytesIO()
        files[path] = f
        return f

    def mock_close(f: BinaryIO) -> None:
        pass

    def mock_exists(path: Path) -> bool:
        return path in files

    def mock_getsize(path: Path) -> int:
        assert path in files, path
        return len(files[path].getbuffer())

    def mock_rename(original_filename: Path, new_filename: Path) -> None:
        assert new_filename not in files, new_filename
        assert original_filename in files, original_filename
        files[new_filename] = files[original_filename]
        del files[original_filename]

    class MockRotatingLog(qlog.RotatingLog):
        rotation_number: int

        def __init__(self, filename: str | Path) -> None:
            super().__init__(filename)
            self.rotation_number = 0

        @overrides.override
        def _get_rotation_filename(self) -> Path:
            nonlocal rotation_number
            rotation_number += 1
            return Path(self.filename_without_ending + "." + str(rotation_number) + self._file_ending)

    mocker.patch.object(qlog, "_SystemAdapter", autospec=True)
    mocker.patch.object(qlog._LoggingEnvironment, "_add_regular_logger_handlers")
    mocker.patch.object(qlog._LoggingEnvironment, "run_uuid", FAKE_RUN_UUID)
    mocker.patch.object(qlog, "_get_timestamp", return_value=FAKE_TIME)

    qlog.RotatingLog = MockRotatingLog  # type: ignore[misc]

    qlog._SystemAdapter.open_xb.side_effect = mock_open_xb  # type: ignore[attr-defined]
    qlog._SystemAdapter.exists.side_effect = mock_exists  # type: ignore[attr-defined]
    qlog._SystemAdapter.getsize.side_effect = mock_getsize  # type: ignore[attr-defined]
    qlog._SystemAdapter.rename.side_effect = mock_rename  # type: ignore[attr-defined]
    qlog._SystemAdapter.close.side_effect = mock_close  # type: ignore[attr-defined]

    yield files
    qlog._cleanup()
    del qlog


@pytest.fixture
def query_logging_cleanup_always_rotate(query_logging_cleanup: MockerFixture, mocker: MockerFixture) -> MockerFixture:
    mocker.patch.object(qlog.RotatingLog, "_RotatingLog__rotation_needed", return_value=True)
    return query_logging_cleanup


def decompress_log_contents(contents: bytes) -> list[bytes]:
    cctx = zstd.ZstdDecompressor()
    decompressed_data = b""
    with cctx.stream_reader(contents) as reader:
        while True:
            chunk = reader.read(1024)
            if not chunk:
                break
            decompressed_data += chunk
    return decompressed_data.splitlines()


def test_no_logging_without_server_start_event(query_logging_cleanup: FileDict) -> None:
    with pytest.raises(qlog.LoggingNotInitializedError):
        qlog.UserInputReceivedEvent.log(text="Hello", uuid=FAKE_RUN_UUID)
    assert not qlog._SystemAdapter.open_xb.called  # type: ignore[attr-defined]
    assert not qlog._SystemAdapter.mkdir.called  # type: ignore[attr-defined]
    assert not qlog._LoggingEnvironment._add_regular_logger_handlers.called  # type: ignore[attr-defined]


def test_logging_with_server_start_event(query_logging_cleanup: FileDict) -> None:
    qlog.ServerStartedEvent.log(args=["--fake"])
    qlog.UserInputReceivedEvent.log(text="Hello", uuid=FAKE_RUN_UUID)

    assert qlog._SystemAdapter.open_xb.called  # type: ignore[attr-defined]
    assert qlog._SystemAdapter.mkdir.called  # type: ignore[attr-defined]
    assert qlog._LoggingEnvironment._add_regular_logger_handlers.called  # type: ignore[attr-defined]

    logger.info(qlog._SystemAdapter.open_xb.call_args_list)  # type: ignore[attr-defined]
    qlog._SystemAdapter.mkdir.assert_called_once_with(RUN_LOG_DIR)  # type: ignore[attr-defined]
    qlog._SystemAdapter.open_xb.assert_called_once_with(RUN_LOG_DIR / "run.log.json.zst")  # type: ignore[attr-defined]

    logf = query_logging_cleanup[RUN_LOG_DIR / "run.log.json.zst"]

    log_contents = logf.getvalue()

    # We could do various checks here
    lines = decompress_log_contents(log_contents)
    assert len(lines) == 2
    assert orjson.loads(lines[0])["rec"]["event"] == "server_started"
    assert orjson.loads(lines[1])["rec"]["event"] == "user_input_received"


def test_no_fs_ops_with_disabled_logging(query_logging_cleanup: FileDict) -> None:
    qlog.init_disabled_logging()
    qlog.UserInputReceivedEvent.log(text="Hello", uuid=FAKE_RUN_UUID)
    # Any opens should be of os.devnull
    assert all(
        args == (Path(os.devnull),)
        for args, _ in qlog._SystemAdapter.open_xb.call_args_list  # type: ignore[attr-defined]
    )
    assert not qlog._SystemAdapter.mkdir.called  # type: ignore[attr-defined]


def test_log_chat_completion_call(query_logging_cleanup: FileDict) -> None:
    qlog.ServerStartedEvent.log(args=["--fake"])
    with qlog.ChatCompletionCallEvent.section(model="fake_model", text=["hello, world"], endpoint="fake_endpoint"):
        pass

    qlog._SystemAdapter.open_xb.assert_called_once_with(RUN_LOG_DIR / "run.log.json.zst")  # type: ignore[attr-defined]
    logf = query_logging_cleanup[RUN_LOG_DIR / "run.log.json.zst"]
    log_contents = logf.getvalue()
    recs = [orjson.loads(line) for line in decompress_log_contents(log_contents)]
    logger.info(pprint.pformat(recs))
    assert len(recs) == 3
    assert recs[0]["rec"]["event"] == "server_started"
    assert recs[1]["rec"] == dict(
        event="chat_completion_call",
        args=dict(endpoint="fake_endpoint", model="fake_model", text="['hello, world']"),
        section=True,
        context=ANY,
    )
    assert recs[2]["rec"] == dict(
        event="end_section", name="chat_completion_call", section=True, context=ANY, exception=None
    )


def test_log_chat_completion_call_with_exception(query_logging_cleanup: FileDict) -> None:
    qlog.ServerStartedEvent.log(args=["--fake"])

    class DummyException(Exception):
        pass

    with pytest.raises(DummyException), qlog.ChatCompletionCallEvent.section(
        model="fake_model", text=["hello, world"], endpoint="fake_endpoint"
    ):
        raise DummyException("oops")

    qlog._SystemAdapter.open_xb.assert_called_once_with(RUN_LOG_DIR / "run.log.json.zst")  # type: ignore[attr-defined]
    logf = query_logging_cleanup[RUN_LOG_DIR / "run.log.json.zst"]
    log_contents = logf.getvalue()
    recs = [orjson.loads(line) for line in decompress_log_contents(log_contents)]
    logger.info(pprint.pformat(recs))
    assert len(recs) == 3, "Did not get all of the expected log lines"
    assert recs[0]["rec"]["event"] == "server_started"
    assert recs[1]["rec"] == dict(
        event="chat_completion_call",
        args=dict(endpoint="fake_endpoint", model="fake_model", text="['hello, world']"),
        section=True,
        context=ANY,
    )
    assert recs[2]["rec"] == dict(
        event="end_section",
        name="chat_completion_call",
        section=True,
        context=ANY,
        exception="DummyException: oops",
    )


class NotPicklableException(Exception):
    def __init__(self) -> None:
        super().__init__("Do not pickle me!")

    def __getstate__(self) -> None:
        raise RuntimeError("Do not pickle me!")


def test_log_chat_completion_call_with_unpicklable_exception(query_logging_cleanup: FileDict) -> None:
    qlog.ServerStartedEvent.log(args=["--fake"])
    with pytest.raises(NotPicklableException), qlog.ChatCompletionCallEvent.section(
        model="fake_model", text=["hello, world"], endpoint="fake_endpoint"
    ):
        raise NotPicklableException

    qlog._SystemAdapter.open_xb.assert_called_once_with(RUN_LOG_DIR / "run.log.json.zst")  # type: ignore[attr-defined]
    logf = query_logging_cleanup[RUN_LOG_DIR / "run.log.json.zst"]
    log_contents = logf.getvalue()
    recs = [orjson.loads(line) for line in decompress_log_contents(log_contents)]
    logger.info(pprint.pformat(recs))
    assert len(recs) == 3, "Did not get all of the expected log lines"
    assert recs[0]["rec"]["event"] == "server_started"
    assert recs[1]["rec"] == dict(
        event="chat_completion_call",
        args=dict(endpoint="fake_endpoint", model="fake_model", text="['hello, world']"),
        section=True,
        context=ANY,
    )
    assert recs[2]["rec"] == dict(
        event="end_section",
        name="chat_completion_call",
        section=True,
        context=ANY,
        exception="NotPicklableException: Do not pickle me!",
    )


# TODO: More of the test logic could be in this class and not duplicated in the tests that use this
@dataclass
class ExpectedLog:
    event_name: str
    message: str
    args: dict[str, Any]
    calls: list[_Call]
    run_log_rows: list[dict[str, Any]]

    def compare_rows(self, rows: list[dict[str, Any]]) -> None:
        recs = [row["rec"] for row in rows]
        for rec, expected_rec in zip(recs, self.run_log_rows):
            assert rec == expected_rec
        assert recs == self.run_log_rows


@pytest.mark.parametrize(
    "expected",
    [
        ExpectedLog(
            event_name="_LogfileSwitchedFromEvent",
            message="switched_from_logfile",
            args={"parent": "parent_log"},
            calls=[
                call(RUN_LOG_DIR / "run.log.json.zst"),
            ],
            run_log_rows=[
                dict(event="server_started", args=ANY, context=ANY, section=False),
                dict(
                    event="switched_from_logfile",
                    args=dict(parent="parent_log"),
                    section=False,
                    context=ANY,
                ),
            ],
        ),
        ExpectedLog(
            event_name="UserInputReceivedEvent",
            message="user_input_received",
            args={"uuid": "007", "text": "user_input"},
            calls=[
                call(RUN_LOG_DIR / "run.log.json.zst"),
            ],
            run_log_rows=[
                dict(event="server_started", args=ANY, context=ANY, section=False),
                dict(
                    event="user_input_received",
                    args=dict(uuid="007", text="user_input"),
                    section=False,
                    context=ANY,
                ),
            ],
        ),
        ExpectedLog(
            event_name="SendMessageToFrontendEvent",
            message="send_message_frontend",
            args={"uuid": "007", "msg": "ai_message", "time_used": 0.5, "cost": 0.15},
            calls=[
                call(RUN_LOG_DIR / "run.log.json.zst"),
            ],
            run_log_rows=[
                dict(event="server_started", args=ANY, context=ANY, section=False),
                dict(
                    event="send_message_frontend",
                    args=dict(uuid="007", msg="ai_message", time_used=0.5, cost=0.15),
                    section=False,
                    context=ANY,
                ),
            ],
        ),
        ExpectedLog(
            event_name="ReceivedMessageFromFrontendEvent",
            message="received_message_frontend",
            args={"uuid": "007", "msg": "user_message"},
            calls=[
                call(RUN_LOG_DIR / "run.log.json.zst"),
            ],
            run_log_rows=[
                dict(event="server_started", args=ANY, context=ANY, section=False),
                dict(
                    event="received_message_frontend",
                    args=dict(uuid="007", msg="user_message"),
                    section=False,
                    context=ANY,
                ),
            ],
        ),
        ExpectedLog(
            event_name="UserLikedMessageEvent",
            message="user_liked_message",
            args={"uuid": "007", "text": None, "ref_uuid": "008"},
            calls=[
                call(RUN_LOG_DIR / "run.log.json.zst"),
            ],
            run_log_rows=[
                dict(event="server_started", args=ANY, context=ANY, section=False),
                dict(
                    event="user_liked_message",
                    args=dict(uuid="007", text=None, ref_uuid="008"),
                    section=False,
                    context=ANY,
                ),
            ],
        ),
        ExpectedLog(
            event_name="UserDislikedMessageEvent",
            message="user_disliked_message",
            args={"uuid": "007", "text": None, "ref_uuid": "008"},
            calls=[
                call(RUN_LOG_DIR / "run.log.json.zst"),
            ],
            run_log_rows=[
                dict(event="server_started", args=ANY, context=ANY, section=False),
                dict(
                    event="user_disliked_message",
                    args=dict(uuid="007", text=None, ref_uuid="008"),
                    section=False,
                    context=ANY,
                ),
            ],
        ),
    ],
)
def test_log_events(query_logging_cleanup: FileDict, expected: ExpectedLog) -> None:
    qlog.ServerStartedEvent.log(args=["--fake"])
    getattr(qlog, expected.event_name).log(**expected.args)

    assert qlog._SystemAdapter.open_xb.call_args_list == expected.calls  # type: ignore[attr-defined]
    logf = query_logging_cleanup[RUN_LOG_DIR / "run.log.json.zst"]
    log_contents = logf.getvalue()
    recs = [orjson.loads(line) for line in decompress_log_contents(log_contents)]
    logger.info(pprint.pformat(recs))
    expected.compare_rows(recs)


@pytest.mark.parametrize(
    "expected",
    [
        ExpectedLog(
            event_name="ConnectionEvent",
            message="connection",
            args={"id": "007", "local_addr": "127.0.0.1", "remote_addr": "127.0.0.1"},
            calls=[
                call(RUN_LOG_DIR / "run.log.json.zst"),
            ],
            run_log_rows=[
                dict(event="server_started", args=ANY, context=ANY, section=False),
                dict(
                    event="connection",
                    args=dict(
                        id="007",
                        local_addr="127.0.0.1",
                        remote_addr="127.0.0.1",
                    ),
                    section=True,
                    context=ANY,
                ),
                dict(event="end_section", name="connection", section=True, context=ANY, exception=None),
            ],
        ),
        ExpectedLog(
            event_name="QueryEvent",
            message="query",
            args={"text": "text", "uuid": "007", "prompt": "FAKE_PROMPT"},
            calls=[
                call(RUN_LOG_DIR / "run.log.json.zst"),
                call(RUN_LOG_DIR / "queries" / f"{FAKE_TIME}__007.log.json.zst"),
            ],
            run_log_rows=[
                dict(event="server_started", args=ANY, context=ANY, section=False),
                dict(
                    event="query",
                    args=dict(text="text", uuid="007", prompt="FAKE_PROMPT"),
                    section=True,
                    context=ANY,
                ),
                dict(
                    event="switch_to_logfile",
                    args=dict(logfile=str(Path("logs/20230101_123456__123/queries/20230101_123456__007"))),
                    section=True,
                    context=ANY,
                ),
                dict(event="end_section", name="switch_to_logfile", section=True, context=ANY, exception=None),
                dict(event="end_section", name="query", section=True, context=ANY, exception=None),
            ],
        ),
        ExpectedLog(
            event_name="ChatCompletionCallEvent",
            message="chat_completion_call",
            args={"model": "gpt4", "text": "text", "endpoint": "OPENAI_API_BASE"},
            calls=[
                call(RUN_LOG_DIR / "run.log.json.zst"),
            ],
            run_log_rows=[
                dict(event="server_started", args=ANY, context=ANY, section=False),
                dict(
                    event="chat_completion_call",
                    args=dict(model="gpt4", text="text", endpoint="OPENAI_API_BASE"),
                    section=True,
                    context=ANY,
                ),
                dict(event="end_section", name="chat_completion_call", section=True, context=ANY, exception=None),
            ],
        ),
        ExpectedLog(
            event_name="_LogfileSwitchEvent",
            message="switch_to_logfile",
            args={"logfile": "./query/query.log.json.zst"},
            calls=[
                call(RUN_LOG_DIR / "run.log.json.zst"),
            ],
            run_log_rows=[
                dict(event="server_started", args=ANY, context=ANY, section=False),
                dict(
                    event="switch_to_logfile",
                    args=dict(logfile="./query/query.log.json.zst"),
                    section=True,
                    context=ANY,
                ),
                dict(event="end_section", name="switch_to_logfile", section=True, context=ANY, exception=None),
            ],
        ),
        ExpectedLog(
            event_name="WaitForUserInputEvent",
            message="wait_for_user_input",
            args={},
            calls=[
                call(RUN_LOG_DIR / "run.log.json.zst"),
            ],
            run_log_rows=[
                dict(event="server_started", args=ANY, context=ANY, section=False),
                dict(event="wait_for_user_input", args={}, section=True, context=ANY),
                dict(event="end_section", name="wait_for_user_input", section=True, context=ANY, exception=None),
            ],
        ),
    ],
)
def test_log_section_events(
    query_logging_cleanup: FileDict,
    expected: ExpectedLog,
) -> None:
    # TODO: Also check the query log file
    qlog.ServerStartedEvent.log(args=["--fake"])
    with getattr(qlog, expected.event_name).section(**expected.args):
        pass

    assert qlog._SystemAdapter.open_xb.call_args_list == expected.calls  # type: ignore[attr-defined]
    logf = query_logging_cleanup[RUN_LOG_DIR / "run.log.json.zst"]
    log_contents = logf.getvalue()
    rows = [orjson.loads(line) for line in decompress_log_contents(log_contents)]
    logger.info(pprint.pformat(rows))
    expected.compare_rows(rows)


@pytest.mark.usefixtures("query_logging_cleanup_always_rotate")
def test_log_rotation(query_logging_cleanup_always_rotate: FileDict) -> None:
    file_ending = ".log.json.zst"
    expected_result_file_dict: dict[Path, Any] = {
        RUN_LOG_DIR / ("run.1" + file_ending): [],
        RUN_LOG_DIR / ("run.2" + file_ending): {"event": "server_started"},
        RUN_LOG_DIR / ("run.3" + file_ending): {"event": "user_input_received", "args": {"text": "Hello the first"}},
        RUN_LOG_DIR / ("run" + file_ending): {"event": "user_input_received", "args": {"text": "Hello the second"}},
    }

    qlog.ServerStartedEvent.log(args=["--fake"])
    qlog.UserInputReceivedEvent.log(text="Hello the first", uuid=FAKE_RUN_UUID)
    qlog.UserInputReceivedEvent.log(text="Hello the second", uuid=FAKE_RUN_UUID)

    for key, value in expected_result_file_dict.items():
        assert key in query_logging_cleanup_always_rotate
        logf = query_logging_cleanup_always_rotate[key]
        log_contents = logf.getvalue()
        recs = [orjson.loads(line) for line in decompress_log_contents(log_contents)]
        if value:
            assert value["event"] == recs[0]["rec"]["event"]
            if recs[0]["rec"]["event"] == "user_input_received":
                assert value["args"]["text"] == recs[0]["rec"]["args"]["text"]
        else:
            assert recs == []

tests/query/logging/test_query_logging.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import asyncio
import json
from typing import Callable

import pytest
import websockets
from pytest_mock import MockerFixture

import rrosti.query.logging as qlog
from rrosti.servers import websocket_query_server as wqs
from rrosti.utils import config

config.load_test_config()

WebSocketMocker = Callable[[], websockets.WebSocketServerProtocol]  # type: ignore[name-defined]


@pytest.fixture
def websocket_mocker(mocker: MockerFixture) -> WebSocketMocker:
    def _run() -> websockets.WebSocketServerProtocol:  # type: ignore[name-defined]
        return mocker.Mock(spec=websockets.WebSocketServerProtocol)  # type: ignore[attr-defined]

    return _run


async def test_web_frontend_sessions_simple(websocket_mocker: WebSocketMocker, mocker: MockerFixture) -> None:
    qlog.init_disabled_logging()

    ws1 = websocket_mocker()
    ws2 = websocket_mocker()

    front1 = wqs.WebFrontend(ws1)
    front2 = wqs.WebFrontend(ws2)

    mocker.patch.object(front1, "_get_user_input_impl", return_value=wqs.UserInputMessage("Hello front1", "1"))

    _, _, front2_msg, front1_msg = await asyncio.gather(
        front1.read_message(json.dumps(dict(type="question", content="Hello front1", id="1"))),
        front2.read_message(json.dumps(dict(type="question", content="Hello front2", id="2"))),
        front2.get_user_input(),
        front1.get_user_input(),
    )

    assert front1_msg == wqs.UserInputMessage("Hello front1", "1")
    assert front2_msg == wqs.UserInputMessage("Hello front2", "2")

tests/servers/test_websocket_query_server.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import sys
from typing import Any


def reload_modules(*module_names: str) -> list[Any]:
    """(Try to) reload a list of modules."""

    for module_name in module_names:
        del sys.modules[module_name]

    return [__import__(module_name, fromlist=[""]) for module_name in module_names]

tests/testing_utils.py
rocket-science-ch/rocketrosti
# Copyright (c) 2023 Rocket Science AG, Switzerland

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import io
from typing import Iterable

import pytest
from pytest_mock import MockerFixture

from rrosti.utils import config

TEST_CONFIG = r"""
backend:
    listen_port: 314
"""


@pytest.fixture
def _mock_config(mocker: MockerFixture) -> Iterable[None]:
    with config.load_test_config(config_file=io.StringIO(TEST_CONFIG)):
        yield


@pytest.mark.usefixtures("_mock_config")
def test_default_values_from_config() -> None:
    @config.uses_config
    def example_function(arg1: int = config.FromConfig("backend.listen_port")) -> None:
        assert arg1 == 314

    example_function()


@pytest.mark.usefixtures("_mock_config")
def test_invalid_config_variable() -> None:
    def example_function(arg1: int = config.FromConfig("backend.loguru_rotation_size_megabytes2")) -> None:
        pass

    with pytest.raises(AttributeError):
        config.uses_config(example_function)


@pytest.mark.usefixtures("_mock_config")
def test_config_not_loaded_at_wrapper(mocker: MockerFixture) -> None:
    config.CONFIG_PATH = config.CONFIG_PATH.parent / "config.invalid.yaml"

    @config.uses_config
    def example_function(arg1: int = config.FromConfig("backend.listen_port")) -> None:
        assert arg1 == 314

    example_function()
    config.CONFIG_PATH = config.CONFIG_PATH.parent / "config.defaults.yaml"

tests/utils/test_config.py
rocket-science-ch/rocketrosti