groknroll 2.0.0__py3-none-any.whl

groknroll/core/lm_handler.py

@@ -0,0 +1,206 @@
+"""
+LMHandler - Routes LLM requests from the RLM process and environment subprocesses.
+
+Uses a multi-threaded socket server. Protocol: 4-byte length prefix + JSON payload.
+"""
+
+import asyncio
+import time
+from socketserver import StreamRequestHandler, ThreadingTCPServer
+from threading import Thread
+
+from groknroll.clients.base_lm import BaseLM
+from groknroll.core.comms_utils import LMRequest, LMResponse, socket_recv, socket_send
+from groknroll.core.types import RLMChatCompletion, UsageSummary
+
+
+class LMRequestHandler(StreamRequestHandler):
+    """Socket handler for LLM completion requests."""
+
+    def handle(self):
+        try:
+            request_data = socket_recv(self.connection)
+            if not isinstance(request_data, dict):
+                response = LMResponse.error_response("Request must be a JSON object")
+                socket_send(self.connection, response.to_dict())
+                return
+
+            request = LMRequest.from_dict(request_data)
+            handler: LMHandler = self.server.lm_handler  # type: ignore
+
+            if request.is_batched:
+                # Batched request: process multiple prompts concurrently
+                response = self._handle_batched(request, handler)
+            elif request.prompt:
+                # Single request: process one prompt
+                response = self._handle_single(request, handler)
+            else:
+                response = LMResponse.error_response("Missing 'prompt' or 'prompts' in request.")
+
+            socket_send(self.connection, response.to_dict())
+
+        except Exception as e:
+            response = LMResponse.error_response(str(e))
+            socket_send(self.connection, response.to_dict())
+
+    def _handle_single(self, request: LMRequest, handler: "LMHandler") -> LMResponse:
+        """Handle a single prompt request."""
+        client = handler.get_client(request.model, request.depth)
+
+        start_time = time.perf_counter()
+        content = client.completion(request.prompt)
+        end_time = time.perf_counter()
+
+        usage_summary = client.get_last_usage()
+        return LMResponse.success_response(
+            chat_completion=RLMChatCompletion(
+                root_model=request.model or client.model_name,
+                prompt=request.prompt,
+                response=content,
+                usage_summary=usage_summary,
+                execution_time=end_time - start_time,
+            )
+        )
+
+    def _handle_batched(self, request: LMRequest, handler: "LMHandler") -> LMResponse:
+        """Handle a batched prompts request using async for concurrency."""
+        client = handler.get_client(request.model, request.depth)
+
+        start_time = time.perf_counter()
+
+        async def run_all():
+            tasks = [client.acompletion(prompt) for prompt in request.prompts]
+            return await asyncio.gather(*tasks)
+
+        results = asyncio.run(run_all())
+        end_time = time.perf_counter()
+
+        total_time = end_time - start_time
+        usage_summary = client.get_last_usage()
+
+        chat_completions = [
+            RLMChatCompletion(
+                root_model=request.model or client.model_name,
+                prompt=prompt,
+                response=content,
+                usage_summary=usage_summary,
+                execution_time=total_time / len(request.prompts),  # approximate per-prompt time
+            )
+            for prompt, content in zip(request.prompts, results, strict=True)
+        ]
+
+        return LMResponse.batched_success_response(chat_completions=chat_completions)
+
+
+class ThreadingLMServer(ThreadingTCPServer):
+    """Multi-threaded TCP server for LM requests."""
+
+    daemon_threads = True
+    allow_reuse_address = True
+
+
+class LMHandler:
+    """
+    Handles all LM calls from the RLM main process and environment subprocesses.
+
+    Uses a multi-threaded socket server for concurrent requests.
+    Protocol: 4-byte big-endian length prefix + JSON payload.
+    """
+
+    def __init__(
+        self,
+        client: BaseLM,
+        host: str = "127.0.0.1",
+        port: int = 0,  # auto-assign available port
+        other_backend_client: BaseLM | None = None,
+    ):
+        self.default_client = client
+        self.other_backend_client = other_backend_client
+        self.clients: dict[str, BaseLM] = {}
+        self.host = host
+        self._server: ThreadingLMServer | None = None
+        self._thread: Thread | None = None
+        self._port = port
+
+        self.register_client(client.model_name, client)
+
+    def register_client(self, model_name: str, client: BaseLM) -> None:
+        """Register a client for a specific model name."""
+        self.clients[model_name] = client
+
+    def get_client(self, model: str | None = None, depth: int = 0) -> BaseLM:
+        """Get client by model name or depth, or return default.
+
+        Routing logic:
+        - depth=0: use default_client (main backend)
+        - depth=1: use other_backend_client if it exists, otherwise default_client
+        - If model is specified and exists in clients, use that (overrides depth routing)
+        """
+        if model and model in self.clients:
+            return self.clients[model]
+
+        # Route based on depth
+        if depth == 1 and self.other_backend_client is not None:
+            return self.other_backend_client
+
+        return self.default_client
+
+    @property
+    def port(self) -> int:
+        """Get the actual port (useful when auto-assigned)."""
+        if self._server:
+            return self._server.server_address[1]
+        return self._port
+
+    @property
+    def address(self) -> tuple[str, int]:
+        """Get (host, port) tuple for connecting."""
+        return (self.host, self.port)
+
+    def start(self) -> tuple[str, int]:
+        """Start the socket server in a background thread. Returns (host, port)."""
+        if self._server is not None:
+            return self.address
+
+        self._server = ThreadingLMServer((self.host, self._port), LMRequestHandler)
+        self._server.lm_handler = self  # type: ignore
+
+        self._thread = Thread(target=self._server.serve_forever, daemon=True)
+        self._thread.start()
+
+        return self.address
+
+    def stop(self):
+        """Stop the socket server."""
+        if self._server:
+            self._server.shutdown()
+            self._server = None
+            self._thread = None
+
+    def completion(self, prompt: str, model: str | None = None) -> str:
+        """Direct completion call (for main process use)."""
+        return self.get_client(model).completion(prompt)
+
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop()
+        return False
+
+    def get_usage_summary(self) -> UsageSummary:
+        """Get the usage summary for all clients, merged into a single dict."""
+        merged = {}
+        # Include default client
+        default_summary = self.default_client.get_usage_summary()
+        merged.update(default_summary.model_usage_summaries)
+        # Include other backend client if it exists
+        if self.other_backend_client is not None:
+            other_summary = self.other_backend_client.get_usage_summary()
+            merged.update(other_summary.model_usage_summaries)
+        # Include all registered clients
+        for client in self.clients.values():
+            client_summary = client.get_usage_summary()
+            merged.update(client_summary.model_usage_summaries)
+        return UsageSummary(model_usage_summaries=merged)

groknroll/core/rlm.py

@@ -0,0 +1,446 @@
+import time
+from contextlib import contextmanager
+from typing import Any
+
+from groknroll.clients import BaseLM, get_client
+from groknroll.core.exceptions import (
+    CompletionTimeoutError,
+    CostLimitExceededError,
+    FinalAnswerNotFoundError,
+    IterationLimitExceededError,
+)
+from groknroll.core.lm_handler import LMHandler
+from groknroll.core.types import (
+    ClientBackend,
+    CodeBlock,
+    EnvironmentType,
+    REPLResult,
+    RLMChatCompletion,
+    RLMIteration,
+    RLMMetadata,
+)
+from groknroll.environments import BaseEnv, SupportsPersistence, get_environment
+from groknroll.logger import RLMLogger, VerbosePrinter
+from groknroll.utils.parsing import (
+    find_code_blocks,
+    find_final_answer,
+    format_iteration,
+)
+from groknroll.utils.prompts import (
+    RLM_SYSTEM_PROMPT,
+    QueryMetadata,
+    build_rlm_system_prompt,
+    build_user_prompt,
+)
+from groknroll.utils.rlm_utils import filter_sensitive_keys
+
+
+class RLM:
+    """
+    Recursive Language Model class that the user instantiates and runs on their tasks.
+
+    Each completion() call spawns its own environment and LM handler, which are
+    cleaned up when the call completes.
+    """
+
+    def __init__(
+        self,
+        backend: ClientBackend = "openai",
+        backend_kwargs: dict[str, Any] | None = None,
+        environment: EnvironmentType = "local",
+        environment_kwargs: dict[str, Any] | None = None,
+        depth: int = 0,
+        max_depth: int = 1,
+        max_iterations: int = 30,
+        custom_system_prompt: str | None = None,
+        other_backends: list[ClientBackend] | None = None,
+        other_backend_kwargs: list[dict[str, Any]] | None = None,
+        logger: RLMLogger | None = None,
+        verbose: bool = False,
+        persistent: bool = False,
+        max_cost: float | None = None,
+        timeout_seconds: float | None = None,
+        iteration_timeout_seconds: float | None = None,
+    ):
+        """
+        Args:
+            backend: The backend to use for the RLM.
+            backend_kwargs: The kwargs to pass to the backend.
+            environment: The environment to use for the RLM.
+            environment_kwargs: The kwargs to pass to the environment.
+            depth: The current depth of the RLM (0-indexed).
+            max_depth: The maximum depth of the RLM. Currently, only depth 1 is supported.
+            max_iterations: The maximum number of iterations of the RLM.
+            custom_system_prompt: The custom system prompt to use for the RLM.
+            other_backends: A list of other client backends that the environments can use to make sub-calls.
+            other_backend_kwargs: The kwargs to pass to the other client backends (ordered to match other_backends).
+            logger: The logger to use for the RLM.
+            verbose: Whether to print verbose output in rich to console.
+            persistent: If True, reuse the environment across completion() calls for multi-turn conversations.
+            max_cost: Maximum allowed cost (in USD) for a single completion. If exceeded, raises CostLimitExceededError.
+            timeout_seconds: Maximum total time (in seconds) for a single completion. If exceeded, raises CompletionTimeoutError.
+            iteration_timeout_seconds: Maximum time (in seconds) for a single iteration. If exceeded, raises CompletionTimeoutError.
+        """
+        # Store config for spawning per-completion
+        self.backend = backend
+        self.backend_kwargs = backend_kwargs
+        self.environment_type = environment
+        self.environment_kwargs = (
+            environment_kwargs.copy() if environment_kwargs is not None else {}
+        )
+        # Validate other_backends: currently only support one additional backend
+        if other_backends is not None:
+            if len(other_backends) != 1:
+                raise ValueError(
+                    "We currently only support one additional backend for the recursive sub-calls! "
+                    "This model will be the model used for recursive sub-calls, but this will change in the future"
+                )
+
+        self.other_backends = other_backends
+        self.other_backend_kwargs = other_backend_kwargs
+
+        self.depth = depth
+        self.max_depth = max_depth
+        self.max_iterations = max_iterations
+        self.system_prompt = custom_system_prompt if custom_system_prompt else RLM_SYSTEM_PROMPT
+        self.logger = logger
+        self.verbose = VerbosePrinter(enabled=verbose)
+
+        # Cost and timeout limits
+        self.max_cost = max_cost
+        self.timeout_seconds = timeout_seconds
+        self.iteration_timeout_seconds = iteration_timeout_seconds
+
+        # Persistence support
+        self.persistent = persistent
+        self._persistent_env: SupportsPersistence | None = None
+
+        # Validate persistence support at initialization
+        if self.persistent:
+            self._validate_persistent_environment_support()
+
+        # Log metadata if logger is provided
+        if self.logger or verbose:
+            metadata = RLMMetadata(
+                root_model=backend_kwargs.get("model_name", "unknown")
+                if backend_kwargs
+                else "unknown",
+                max_depth=max_depth,
+                max_iterations=max_iterations,
+                backend=backend,
+                backend_kwargs=filter_sensitive_keys(backend_kwargs) if backend_kwargs else {},
+                environment_type=environment,
+                environment_kwargs=filter_sensitive_keys(environment_kwargs)
+                if environment_kwargs
+                else {},
+                other_backends=other_backends,
+            )
+            if self.logger:
+                self.logger.log_metadata(metadata)
+            self.verbose.print_metadata(metadata)
+
+    @contextmanager
+    def _spawn_completion_context(self, prompt: str | dict[str, Any]):
+        """
+        Spawn an LM handler and environment for a single completion call.
+
+        When persistent=True, the environment is reused across calls.
+        When persistent=False (default), creates fresh environment each call.
+        """
+        # Create client and wrap in handler
+        client: BaseLM = get_client(self.backend, self.backend_kwargs)
+
+        # Create other_backend_client if provided (for depth=1 routing)
+        other_backend_client: BaseLM | None = None
+        if self.other_backends and self.other_backend_kwargs:
+            other_backend_client = get_client(self.other_backends[0], self.other_backend_kwargs[0])
+
+        lm_handler = LMHandler(client, other_backend_client=other_backend_client)
+
+        # Register other clients to be available as sub-call options (by model name)
+        if self.other_backends and self.other_backend_kwargs:
+            for backend, kwargs in zip(self.other_backends, self.other_backend_kwargs, strict=True):
+                other_client: BaseLM = get_client(backend, kwargs)
+                lm_handler.register_client(other_client.model_name, other_client)
+
+        lm_handler.start()
+
+        # Environment: reuse if persistent, otherwise create fresh
+        if self.persistent and self._persistent_env is not None:
+            environment = self._persistent_env
+            # Defensive check: ensure environment supports persistence methods
+            if not self._env_supports_persistence(environment):
+                raise RuntimeError(
+                    f"Persistent environment of type '{type(environment).__name__}' does not "
+                    f"implement required methods (update_handler_address, add_context, get_context_count). "
+                    f"This should have been caught at initialization."
+                )
+            environment.update_handler_address((lm_handler.host, lm_handler.port))
+            environment.add_context(prompt)
+        else:
+            env_kwargs = self.environment_kwargs.copy()
+            env_kwargs["lm_handler_address"] = (lm_handler.host, lm_handler.port)
+            env_kwargs["context_payload"] = prompt
+            env_kwargs["depth"] = self.depth + 1  # Environment depth is RLM depth + 1
+            environment: BaseEnv = get_environment(self.environment_type, env_kwargs)
+
+            if self.persistent:
+                self._persistent_env = environment
+
+        try:
+            yield lm_handler, environment
+        finally:
+            lm_handler.stop()
+            if not self.persistent and hasattr(environment, "cleanup"):
+                environment.cleanup()
+
+    def _setup_prompt(self, prompt: str | dict[str, Any]) -> list[dict[str, Any]]:
+        """
+        Setup the system prompt for the RLM. Also include metadata about the prompt and build
+        up the initial message history.
+        """
+        metadata = QueryMetadata(prompt)
+        message_history = build_rlm_system_prompt(
+            system_prompt=self.system_prompt, query_metadata=metadata
+        )
+
+        return message_history
+
+    def completion(
+        self, prompt: str | dict[str, Any], root_prompt: str | None = None
+    ) -> RLMChatCompletion:
+        """
+        Recursive Language Model completion call. This is the main entry point for querying an RLM, and
+        can replace a regular LM completion call.
+
+        Spawns its own environment and LM handler for the duration of this call.
+
+        Args:
+            prompt: A single string or dictionary of messages to pass as context to the model.
+            root_prompt: We allow the RLM's root LM to see a (small) prompt that the user specifies. A common example of this
+            is if the user is asking the RLM to answer a question, we can pass the question as the root prompt.
+        Returns:
+            A final answer as a string.
+        """
+        time_start = time.perf_counter()
+
+        # If we're at max depth, the RLM is an LM, so we fallback to the regular LM.
+        if self.depth >= self.max_depth:
+            return self._fallback_answer(prompt)
+
+        with self._spawn_completion_context(prompt) as (lm_handler, environment):
+            message_history = self._setup_prompt(prompt)
+
+            for i in range(self.max_iterations):
+                # Check timeout
+                if self.timeout_seconds is not None:
+                    elapsed = time.perf_counter() - time_start
+                    if elapsed > self.timeout_seconds:
+                        raise CompletionTimeoutError(
+                            f"Completion exceeded timeout of {self.timeout_seconds}s",
+                            timeout_seconds=self.timeout_seconds,
+                            elapsed_seconds=elapsed,
+                        )
+
+                # Check cost limit
+                if self.max_cost is not None:
+                    current_cost = lm_handler.get_usage_summary().total_cost
+                    if current_cost > self.max_cost:
+                        raise CostLimitExceededError(
+                            f"Completion cost ${current_cost:.4f} exceeded limit ${self.max_cost:.4f}",
+                            current_cost=current_cost,
+                            cost_limit=self.max_cost,
+                        )
+
+                # Current prompt = message history + additional prompt suffix
+                context_count = (
+                    environment.get_context_count()
+                    if isinstance(environment, SupportsPersistence)
+                    else 1
+                )
+                history_count = (
+                    environment.get_history_count()
+                    if isinstance(environment, SupportsPersistence)
+                    else 0
+                )
+                current_prompt = message_history + [
+                    build_user_prompt(root_prompt, i, context_count, history_count)
+                ]
+
+                iteration: RLMIteration = self._completion_turn(
+                    prompt=current_prompt,
+                    lm_handler=lm_handler,
+                    environment=environment,
+                )
+
+                # Check if RLM is done and has a final answer.
+                final_answer = find_final_answer(iteration.response, environment=environment)
+                iteration.final_answer = final_answer
+
+                # If logger is used, log the iteration.
+                if self.logger:
+                    self.logger.log(iteration)
+
+                # Verbose output for this iteration
+                self.verbose.print_iteration(iteration, i + 1)
+
+                if final_answer is not None:
+                    time_end = time.perf_counter()
+                    usage = lm_handler.get_usage_summary()
+                    self.verbose.print_final_answer(final_answer)
+                    self.verbose.print_summary(i + 1, time_end - time_start, usage.to_dict())
+
+                    # Store message history in persistent environment
+                    if self.persistent and isinstance(environment, SupportsPersistence):
+                        environment.add_history(message_history)
+
+                    return RLMChatCompletion(
+                        root_model=self.backend_kwargs.get("model_name", "unknown")
+                        if self.backend_kwargs
+                        else "unknown",
+                        prompt=prompt,
+                        response=final_answer,
+                        usage_summary=usage,
+                        execution_time=time_end - time_start,
+                    )
+
+                # Format the iteration for the next prompt.
+                new_messages = format_iteration(iteration)
+
+                # Update message history with the new messages.
+                message_history.extend(new_messages)
+
+            # Default behavior: we run out of iterations, provide one final answer
+            time_end = time.perf_counter()
+            final_answer = self._default_answer(message_history, lm_handler)
+            usage = lm_handler.get_usage_summary()
+            self.verbose.print_final_answer(final_answer)
+            self.verbose.print_summary(self.max_iterations, time_end - time_start, usage.to_dict())
+
+            # Store message history in persistent environment
+            if self.persistent and isinstance(environment, SupportsPersistence):
+                environment.add_history(message_history)
+
+            return RLMChatCompletion(
+                root_model=self.backend_kwargs.get("model_name", "unknown")
+                if self.backend_kwargs
+                else "unknown",
+                prompt=prompt,
+                response=final_answer,
+                usage_summary=usage,
+                execution_time=time_end - time_start,
+            )
+
+    def _completion_turn(
+        self,
+        prompt: str | dict[str, Any],
+        lm_handler: LMHandler,
+        environment: BaseEnv,
+    ) -> RLMIteration:
+        """
+        Perform a single iteration of the RLM, including prompting the model
+        and code execution + tool execution.
+        """
+        iter_start = time.perf_counter()
+        response = lm_handler.completion(prompt)
+        code_block_strs = find_code_blocks(response)
+        code_blocks = []
+
+        for code_block_str in code_block_strs:
+            # Check iteration timeout
+            if self.iteration_timeout_seconds is not None:
+                elapsed = time.perf_counter() - iter_start
+                if elapsed > self.iteration_timeout_seconds:
+                    raise CompletionTimeoutError(
+                        f"Iteration exceeded timeout of {self.iteration_timeout_seconds}s",
+                        timeout_seconds=self.iteration_timeout_seconds,
+                        elapsed_seconds=elapsed,
+                    )
+
+            code_result: REPLResult = environment.execute_code(code_block_str)
+            code_blocks.append(CodeBlock(code=code_block_str, result=code_result))
+
+        iteration_time = time.perf_counter() - iter_start
+        return RLMIteration(
+            prompt=prompt,
+            response=response,
+            code_blocks=code_blocks,
+            iteration_time=iteration_time,
+        )
+
+    def _default_answer(self, message_history: list[dict[str, Any]], lm_handler: LMHandler) -> str:
+        """
+        Default behavior if the RLM runs out of iterations and does not find a final answer.
+        It will take the message history, and try to generate a final answer from it.
+        """
+        current_prompt = message_history + [
+            {
+                "role": "assistant",
+                "content": "Please provide a final answer to the user's question based on the information provided.",
+            }
+        ]
+        response = lm_handler.completion(current_prompt)
+
+        if self.logger:
+            self.logger.log(
+                RLMIteration(
+                    prompt=current_prompt,
+                    response=response,
+                    final_answer=response,
+                    code_blocks=[],
+                )
+            )
+
+        return response
+
+    def _fallback_answer(self, message: str | dict[str, Any]) -> str:
+        """
+        Fallback behavior if the RLM is actually at max depth, and should be treated as an LM.
+        """
+        client: BaseLM = get_client(self.backend, self.backend_kwargs)
+        response = client.completion(message)
+        return response
+
+    def _validate_persistent_environment_support(self) -> None:
+        """
+        Validate that the configured environment type supports persistent mode.
+
+        Persistent mode requires environments to implement:
+        - update_handler_address(address): Update LM handler address between calls
+        - add_context(payload, index): Add new context for multi-turn conversations
+        - get_context_count(): Return the number of loaded contexts
+
+        Currently only 'local' (LocalREPL) supports these methods.
+
+        Raises:
+            ValueError: If the environment type does not support persistent mode.
+        """
+        # Known environments that support persistence
+        persistent_supported_environments = {"local"}
+
+        if self.environment_type not in persistent_supported_environments:
+            raise ValueError(
+                f"persistent=True is not supported for environment type '{self.environment_type}'. "
+                f"Persistent mode requires environments that implement update_handler_address(), "
+                f"add_context(), and get_context_count(). "
+                f"Supported environments: {sorted(persistent_supported_environments)}"
+            )
+
+    @staticmethod
+    def _env_supports_persistence(env: BaseEnv) -> bool:
+        """Check if an environment instance supports persistent mode methods."""
+        return isinstance(env, SupportsPersistence)
+
+    def close(self) -> None:
+        """Clean up persistent environment. Call when done with multi-turn conversations."""
+        if self._persistent_env is not None:
+            if hasattr(self._persistent_env, "cleanup"):
+                self._persistent_env.cleanup()
+            self._persistent_env = None
+
+    def __enter__(self) -> "RLM":
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
+        self.close()
+        return False