chore: import upstream snapshot with attribution

runtime/node/agent/thinking/__init__.py

@@ -0,0 +1,8 @@
+from .thinking_manager import ThinkingManagerBase, ThinkingPayload
+from .builtin_thinking import ThinkingManagerFactory
+
+__all__ = [
+    "ThinkingManagerBase",
+    "ThinkingPayload",
+    "ThinkingManagerFactory",
+]

runtime/node/agent/thinking/builtin_thinking.py

@@ -0,0 +1,26 @@
+"""Register built-in thinking modes."""
+
+from entity.configs.node.thinking import ReflectionThinkingConfig, ThinkingConfig
+from runtime.node.agent.thinking.thinking_manager import ThinkingManagerBase
+from runtime.node.agent.thinking.self_reflection import SelfReflectionThinkingManager
+from runtime.node.agent.thinking.registry import (
+    register_thinking_mode,
+    get_thinking_registration,
+)
+
+register_thinking_mode(
+    "reflection",
+    config_cls=ReflectionThinkingConfig,
+    manager_cls=SelfReflectionThinkingManager,
+    summary="LLM reflects on its output and refine its output",
+)
+
+
+class ThinkingManagerFactory:
+    @staticmethod
+    def get_thinking_manager(config: ThinkingConfig) -> ThinkingManagerBase:
+        registration = get_thinking_registration(config.type)
+        typed_config = config.as_config(registration.config_cls)
+        if not typed_config:
+            raise ValueError(f"Invalid thinking config for type '{config.type}'")
+        return registration.manager_cls(typed_config)

runtime/node/agent/thinking/registry.py

@@ -0,0 +1,63 @@
+"""Registry for thinking managers."""
+
+from dataclasses import dataclass
+from importlib import import_module
+from typing import Any, Dict, Type
+
+from schema_registry import register_thinking_schema
+from utils.registry import Registry, RegistryEntry, RegistryError
+from runtime.node.agent.thinking.thinking_manager import ThinkingManagerBase
+
+thinking_registry = Registry("thinking_mode")
+_BUILTINS_LOADED = False
+
+@dataclass(slots=True)
+class ThinkingRegistration:
+    name: str
+    config_cls: Type[Any]
+    manager_cls: Type["ThinkingManagerBase"]
+    summary: str | None = None
+
+
+def _ensure_builtins_loaded() -> None:
+    global _BUILTINS_LOADED
+    if not _BUILTINS_LOADED:
+        import_module("runtime.node.agent.thinking.builtin_thinking")
+        _BUILTINS_LOADED = True
+
+
+def register_thinking_mode(
+    name: str,
+    *,
+    config_cls: Type[Any],
+    manager_cls: Type["ThinkingManagerBase"],
+    summary: str | None = None,
+) -> None:
+    if name in thinking_registry.names():
+        raise RegistryError(f"Thinking mode '{name}' already registered")
+    entry = ThinkingRegistration(name=name, config_cls=config_cls, manager_cls=manager_cls, summary=summary)
+    thinking_registry.register(name, target=entry)
+    register_thinking_schema(name, config_cls=config_cls, summary=summary)
+
+
+def get_thinking_registration(name: str) -> ThinkingRegistration:
+    _ensure_builtins_loaded()
+    entry: RegistryEntry = thinking_registry.get(name)
+    registration = entry.load()
+    if not isinstance(registration, ThinkingRegistration):
+        raise RegistryError(f"Entry '{name}' is not a ThinkingRegistration")
+    return registration
+
+
+def iter_thinking_registrations() -> Dict[str, ThinkingRegistration]:
+    _ensure_builtins_loaded()
+    return {name: entry.load() for name, entry in thinking_registry.items()}
+
+
+__all__ = [
+    "thinking_registry",
+    "ThinkingRegistration",
+    "register_thinking_mode",
+    "get_thinking_registration",
+    "iter_thinking_registrations",
+]

runtime/node/agent/thinking/self_reflection.py

@@ -0,0 +1,54 @@
+from entity.configs import ReflectionThinkingConfig
+from entity.messages import Message, MessageRole
+from runtime.node.agent.thinking.thinking_manager import (
+    ThinkingManagerBase,
+    AgentInvoker,
+    ThinkingPayload,
+)
+
+
+class SelfReflectionThinkingManager(ThinkingManagerBase):
+    """
+    A simple implementation of thinking manager, named self-reflection.
+    This part of the code is borrowed from ChatDev (https://github.com/OpenBMB/ChatDev) and adapted.
+    """
+
+    def __init__(self, config: ReflectionThinkingConfig):
+        super().__init__(config)
+        self.before_gen_think_enabled = False
+        self.after_gen_think_enabled = True
+        self.base_prompt = """Here is a conversation between two roles: {conversations} {reflection_prompt}"""
+
+        self.reflection_prompt = config.reflection_prompt or "Reflect on the given information and summarize key points in a few words."
+
+    def _before_gen_think(
+        self,
+        agent_invoker: AgentInvoker,
+        input_payload: ThinkingPayload,
+        agent_role: str,
+        memory: ThinkingPayload | None,
+    ) -> tuple[str, bool]:
+        ...
+
+    def _after_gen_think(
+        self,
+        agent_invoker: AgentInvoker,
+        input_payload: ThinkingPayload,
+        agent_role: str,
+        memory: ThinkingPayload | None,
+        gen_payload: ThinkingPayload,
+    ) -> tuple[str, bool]:
+        conversations = [
+            f"SYSTEM: {agent_role}",
+            f"USER: {input_payload.text}",
+            f"ASSISTANT: {gen_payload.text}",
+        ]
+        if memory and memory.text:
+            conversations = [memory.text] + conversations
+        prompt = self.base_prompt.format(conversations="\n\n".join(conversations),
+                                         reflection_prompt=self.reflection_prompt)
+
+        reflection_message = agent_invoker(
+            [Message(role=MessageRole.USER, content=prompt)]
+        )
+        return reflection_message.text_content(), True

runtime/node/agent/thinking/thinking_manager.py

@@ -0,0 +1,106 @@
+from abc import abstractmethod, ABC
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List
+
+from entity.configs import ThinkingConfig
+from entity.messages import Message, MessageRole, MessageBlock
+
+AgentInvoker = Callable[[List[Message]], Message]
+
+
+@dataclass
+class ThinkingPayload:
+    """Container used to pass multimodal context into thinking managers."""
+
+    text: str
+    blocks: List[MessageBlock] = field(default_factory=list)
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    raw: Any | None = None
+
+    def describe(self) -> str:
+        return self.text
+
+
+
+class ThinkingManagerBase(ABC):
+    def __init__(self, config: ThinkingConfig):
+        self.config = config
+        self.before_gen_think_enabled = False
+        self.after_gen_think_enabled = False
+
+        # you can customize the prompt by override this attribute
+        self.thinking_concat_prompt = "{origin}\n\nThinking Result: {thinking}"
+
+    @abstractmethod
+    def _before_gen_think(
+        self,
+        agent_invoker: AgentInvoker,
+        input_payload: ThinkingPayload,
+        agent_role: str,
+        memory: ThinkingPayload | None,
+    ) -> tuple[str, bool]:
+        """
+        think based on input_data before calling model API for node to generate
+
+        Returns:
+            str: thinking result
+            bool: whether to replace the original input_data with the modified one
+        """
+        ...
+
+    @abstractmethod
+    def _after_gen_think(
+        self,
+        agent_invoker: AgentInvoker,
+        input_payload: ThinkingPayload,
+        agent_role: str,
+        memory: ThinkingPayload | None,
+        gen_payload: ThinkingPayload,
+    ) -> tuple[str, bool]:
+        """
+        think based on input_data and gen_data after calling model API for node to generate
+
+        Returns:
+            str: thinking result
+            bool: whether to replace the original gen_data with the modified one
+        """
+        ...
+
+    def think(
+        self,
+        agent_invoker: AgentInvoker,
+        input_payload: ThinkingPayload,
+        agent_role: str,
+        memory: ThinkingPayload | None,
+        gen_payload: ThinkingPayload | None = None,
+    ) -> str | Message:
+        """
+        think based on input_data and gen_data if provided
+
+        Returns:
+            str: result for next step
+        """
+
+        normalized_input = input_payload.text
+        normalized_gen = gen_payload.text if gen_payload is not None else None
+
+        if gen_payload is None and self.before_gen_think_enabled:
+            think_result, replace_input = self._before_gen_think(
+                agent_invoker, input_payload, agent_role, memory
+            )
+            if replace_input:
+                return think_result
+            else:
+                return self.thinking_concat_prompt.format(origin=normalized_input, thinking=think_result)
+        elif gen_payload is not None and self.after_gen_think_enabled:
+            think_result, replace_gen = self._after_gen_think(
+                agent_invoker, input_payload, agent_role, memory, gen_payload
+            )
+            if replace_gen:
+                return think_result
+            else:
+                return self.thinking_concat_prompt.format(origin=normalized_gen or "", thinking=think_result)
+        else:
+            if gen_payload is not None:
+                return gen_payload.raw if gen_payload.raw is not None else gen_payload.text
+            return input_payload.raw if input_payload.raw is not None else input_payload.text