groknroll 2.0.0__py3-none-any.whl

groknroll/core/types.py

@@ -0,0 +1,276 @@
+from dataclasses import dataclass
+from types import ModuleType
+from typing import Any, Literal
+
+ClientBackend = Literal[
+    "openai",
+    "portkey",
+    "openrouter",
+    "vercel",
+    "vllm",
+    "litellm",
+    "anthropic",
+    "azure_openai",
+    "gemini",
+]
+EnvironmentType = Literal["local", "docker", "modal", "prime"]
+
+
+def _serialize_value(value: Any) -> Any:
+    """Convert a value to a JSON-serializable representation."""
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return value
+    if isinstance(value, ModuleType):
+        return f"<module '{value.__name__}'>"
+    if isinstance(value, (list, tuple)):
+        return [_serialize_value(v) for v in value]
+    if isinstance(value, dict):
+        return {str(k): _serialize_value(v) for k, v in value.items()}
+    if callable(value):
+        return f"<{type(value).__name__} '{getattr(value, '__name__', repr(value))}'>"
+    # Try to convert to string for other types
+    try:
+        return repr(value)
+    except Exception:
+        return f"<{type(value).__name__}>"
+
+
+########################################################
+########    Types for LM Cost Tracking         #########
+########################################################
+
+
+@dataclass
+class ModelUsageSummary:
+    total_calls: int
+    total_input_tokens: int
+    total_output_tokens: int
+
+    def to_dict(self):
+        return {
+            "total_calls": self.total_calls,
+            "total_input_tokens": self.total_input_tokens,
+            "total_output_tokens": self.total_output_tokens,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "ModelUsageSummary":
+        return cls(
+            total_calls=data.get("total_calls"),
+            total_input_tokens=data.get("total_input_tokens"),
+            total_output_tokens=data.get("total_output_tokens"),
+        )
+
+
+@dataclass
+class UsageSummary:
+    model_usage_summaries: dict[str, ModelUsageSummary]
+
+    @property
+    def total_cost(self) -> float:
+        """Calculate total cost across all models.
+
+        Note: This is a placeholder that returns 0.0 until cost calculation
+        is fully implemented with model-specific pricing data.
+        """
+        # TODO: Implement actual cost calculation using model pricing
+        # For now, return 0 to allow cost limiting infrastructure to work
+        return 0.0
+
+    def to_dict(self):
+        return {
+            "model_usage_summaries": {
+                model: usage_summary.to_dict()
+                for model, usage_summary in self.model_usage_summaries.items()
+            },
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "UsageSummary":
+        return cls(
+            model_usage_summaries={
+                model: ModelUsageSummary.from_dict(usage_summary)
+                for model, usage_summary in data.get("model_usage_summaries", {}).items()
+            },
+        )
+
+
+########################################################
+########   Types for REPL and RLM Iterations   #########
+########################################################
+@dataclass
+class RLMChatCompletion:
+    """Record of a single LLM call made from within the environment."""
+
+    root_model: str
+    prompt: str | dict[str, Any]
+    response: str
+    usage_summary: UsageSummary
+    execution_time: float
+
+    def to_dict(self):
+        return {
+            "root_model": self.root_model,
+            "prompt": self.prompt,
+            "response": self.response,
+            "usage_summary": self.usage_summary.to_dict(),
+            "execution_time": self.execution_time,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "RLMChatCompletion":
+        return cls(
+            root_model=data.get("root_model"),
+            prompt=data.get("prompt"),
+            response=data.get("response"),
+            usage_summary=UsageSummary.from_dict(data.get("usage_summary")),
+            execution_time=data.get("execution_time"),
+        )
+
+
+@dataclass
+class REPLResult:
+    stdout: str
+    stderr: str
+    locals: dict
+    execution_time: float
+    llm_calls: list["RLMChatCompletion"]
+
+    def __init__(
+        self,
+        stdout: str,
+        stderr: str,
+        locals: dict,
+        execution_time: float = None,
+        rlm_calls: list["RLMChatCompletion"] = None,
+    ):
+        self.stdout = stdout
+        self.stderr = stderr
+        self.locals = locals
+        self.execution_time = execution_time
+        self.rlm_calls = rlm_calls or []
+
+    def __str__(self):
+        return f"REPLResult(stdout={self.stdout}, stderr={self.stderr}, locals={self.locals}, execution_time={self.execution_time}, rlm_calls={len(self.rlm_calls)})"
+
+    def to_dict(self):
+        return {
+            "stdout": self.stdout,
+            "stderr": self.stderr,
+            "locals": {k: _serialize_value(v) for k, v in self.locals.items()},
+            "execution_time": self.execution_time,
+            "rlm_calls": [call.to_dict() for call in self.rlm_calls],
+        }
+
+
+@dataclass
+class CodeBlock:
+    code: str
+    result: REPLResult
+
+    def to_dict(self):
+        return {"code": self.code, "result": self.result.to_dict()}
+
+
+@dataclass
+class RLMIteration:
+    prompt: str | dict[str, Any]
+    response: str
+    code_blocks: list[CodeBlock]
+    final_answer: str | None = None
+    iteration_time: float | None = None
+
+    def to_dict(self):
+        return {
+            "prompt": self.prompt,
+            "response": self.response,
+            "code_blocks": [code_block.to_dict() for code_block in self.code_blocks],
+            "final_answer": self.final_answer,
+            "iteration_time": self.iteration_time,
+        }
+
+
+########################################################
+########   Types for RLM Metadata   #########
+########################################################
+
+
+@dataclass
+class RLMMetadata:
+    """Metadata about the RLM configuration."""
+
+    root_model: str
+    max_depth: int
+    max_iterations: int
+    backend: str
+    backend_kwargs: dict[str, Any]
+    environment_type: str
+    environment_kwargs: dict[str, Any]
+    other_backends: list[str] | None = None
+
+    def to_dict(self):
+        return {
+            "root_model": self.root_model,
+            "max_depth": self.max_depth,
+            "max_iterations": self.max_iterations,
+            "backend": self.backend,
+            "backend_kwargs": {k: _serialize_value(v) for k, v in self.backend_kwargs.items()},
+            "environment_type": self.environment_type,
+            "environment_kwargs": {
+                k: _serialize_value(v) for k, v in self.environment_kwargs.items()
+            },
+            "other_backends": self.other_backends,
+        }
+
+
+########################################################
+########   Types for RLM Prompting   #########
+########################################################
+
+
+@dataclass
+class QueryMetadata:
+    context_lengths: list[int]
+    context_total_length: int
+    context_type: str
+
+    def __init__(self, prompt: str | list[str] | dict[Any, Any] | list[dict[Any, Any]]):
+        if isinstance(prompt, str):
+            self.context_lengths = [len(prompt)]
+            self.context_type = "str"
+        elif isinstance(prompt, dict):
+            self.context_type = "dict"
+            self.context_lengths = []
+            for chunk in prompt.values():
+                if isinstance(chunk, str):
+                    self.context_lengths.append(len(chunk))
+                    continue
+                try:
+                    import json
+
+                    self.context_lengths.append(len(json.dumps(chunk, default=str)))
+                except Exception:
+                    self.context_lengths.append(len(repr(chunk)))
+            self.context_type = "dict"
+        elif isinstance(prompt, list):
+            self.context_type = "list"
+            if len(prompt) == 0:
+                self.context_lengths = [0]
+            elif isinstance(prompt[0], dict):
+                if "content" in prompt[0]:
+                    self.context_lengths = [len(str(chunk.get("content", ""))) for chunk in prompt]
+                else:
+                    self.context_lengths = []
+                    for chunk in prompt:
+                        try:
+                            import json
+
+                            self.context_lengths.append(len(json.dumps(chunk, default=str)))
+                        except Exception:
+                            self.context_lengths.append(len(repr(chunk)))
+            else:
+                self.context_lengths = [len(chunk) for chunk in prompt]
+        else:
+            raise ValueError(f"Invalid prompt type: {type(prompt)}")
+
+        self.context_total_length = sum(self.context_lengths)

groknroll/environments/__init__.py

@@ -0,0 +1,34 @@
+from typing import Any, Literal
+
+from groknroll.environments.base_env import BaseEnv, SupportsPersistence
+from groknroll.environments.local_repl import LocalREPL
+
+__all__ = ["BaseEnv", "LocalREPL", "SupportsPersistence", "get_environment"]
+
+
+def get_environment(
+    environment: Literal["local", "modal", "docker", "prime"],
+    environment_kwargs: dict[str, Any],
+) -> BaseEnv:
+    """
+    Routes a specific environment and the args (as a dict) to the appropriate environment if supported.
+    Currently supported environments: ['local', 'modal', 'docker', 'prime']
+    """
+    if environment == "local":
+        return LocalREPL(**environment_kwargs)
+    elif environment == "modal":
+        from groknroll.environments.modal_repl import ModalREPL
+
+        return ModalREPL(**environment_kwargs)
+    elif environment == "docker":
+        from groknroll.environments.docker_repl import DockerREPL
+
+        return DockerREPL(**environment_kwargs)
+    elif environment == "prime":
+        from groknroll.environments.prime_repl import PrimeREPL
+
+        return PrimeREPL(**environment_kwargs)
+    else:
+        raise ValueError(
+            f"Unknown environment: {environment}. Supported: ['local', 'modal', 'docker', 'prime']"
+        )

groknroll/environments/base_env.py

@@ -0,0 +1,182 @@
+from abc import ABC, abstractmethod
+from typing import Any, Protocol, runtime_checkable
+
+from groknroll.core.types import REPLResult
+
+
+class BaseEnv(ABC):
+    """
+    Base REPL-like environment that the RLM uses to interact with. The primary types are isolated and non-isolated,
+    where isolated environments are on a separate machine from the LM.
+    """
+
+    def __init__(self, persistent: bool = False, depth: int = 1, **kwargs):
+        self.persistent = persistent
+        self.depth = depth
+        self.kwargs = kwargs
+
+    @abstractmethod
+    def setup(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def load_context(self, context_payload: dict | list | str):
+        raise NotImplementedError
+
+    @abstractmethod
+    def execute_code(self, code: str) -> REPLResult:
+        raise NotImplementedError
+
+
+class IsolatedEnv(BaseEnv, ABC):
+    """
+    These environments (e.g. Prime Envs, Modal Envs) sit on a completely separate machine from the LM,
+    guaranteeing complete isolation from the LM process.
+    """
+
+    def __init__(self, persistent: bool = False, **kwargs):
+        super().__init__(persistent=persistent, **kwargs)
+
+    @abstractmethod
+    def setup(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def load_context(self, context_payload: dict | list | str):
+        raise NotImplementedError
+
+    @abstractmethod
+    def execute_code(self, code: str) -> REPLResult:
+        raise NotImplementedError
+
+
+class NonIsolatedEnv(BaseEnv, ABC):
+    """
+    These environments run on the same machine as the LM, and provide different levels of isolation
+    depending on the choice of environment. The simplest, default is a local Python REPL that runs
+    as a subprocess.
+    """
+
+    def __init__(self, persistent: bool = False, **kwargs):
+        super().__init__(persistent=persistent, **kwargs)
+
+    @abstractmethod
+    def setup(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def load_context(self, context_payload: dict | list | str):
+        raise NotImplementedError
+
+    @abstractmethod
+    def execute_code(self, code: str) -> REPLResult:
+        raise NotImplementedError
+
+
+@runtime_checkable
+class SupportsPersistence(Protocol):
+    """Protocol for environments that support persistent multi-turn sessions.
+
+    CHECKING SUPPORT:
+        Use isinstance(env, SupportsPersistence) to check if an environment
+        supports persistence capabilities.
+
+    IMPLEMENTING THIS PROTOCOL:
+        To add persistence to your environment, implement these 5 methods.
+        See tests/test_local_repl_persistent.py for expected behavior.
+
+    VERSIONING BEHAVIOR:
+        Contexts and histories are versioned with numeric suffixes:
+        - First context  -> context_0, context_1, context_2, ...
+        - First history  -> history_0, history_1, history_2, ...
+
+    ALIASING BEHAVIOR:
+        The unversioned names always point to index 0:
+        - context  -> context_0 (first context)
+        - history  -> history_0 (first history)
+
+    EXAMPLE IMPLEMENTATION:
+        See rlm/environments/local_repl.py for a complete reference.
+
+    TESTS:
+        - Unit tests: tests/test_local_repl_persistent.py
+        - Integration tests: tests/test_multi_turn_integration.py
+
+        Run: uv run pytest tests/test_local_repl_persistent.py -v
+    """
+
+    def update_handler_address(self, address: tuple[str, int]) -> None:
+        """Update the LM handler address for nested LLM calls.
+
+        Called by RLM when the handler address changes between completions.
+        Store the address so llm_query() calls from executed code can reach
+        the LM handler.
+
+        Args:
+            address: (host, port) tuple for the LM handler server.
+        """
+        ...
+
+    def add_context(
+        self, context_payload: dict | list | str, context_index: int | None = None
+    ) -> int:
+        """Add a context payload, making it available as context_N in code.
+
+        Versioning:
+            - context_index=None: auto-increment (0, 1, 2, ...)
+            - context_index=N: use specific index N
+
+        Storage:
+            Must store so executed code can access:
+            - context_0, context_1, etc. (versioned)
+            - context (alias to context_0)
+
+        Args:
+            context_payload: The context data (string, dict, or list).
+            context_index: Optional specific index, or None to auto-increment.
+
+        Returns:
+            The index used (for auto-increment, returns the assigned index).
+        """
+        ...
+
+    def get_context_count(self) -> int:
+        """Return the number of contexts added so far.
+
+        Used by RLM to inform the model how many contexts are available.
+        """
+        ...
+
+    def add_history(
+        self, message_history: list[dict[str, Any]], history_index: int | None = None
+    ) -> int:
+        """Add a message history, making it available as history_N in code.
+
+        Versioning:
+            - history_index=None: auto-increment (0, 1, 2, ...)
+            - history_index=N: use specific index N
+
+        Storage:
+            Must store so executed code can access:
+            - history_0, history_1, etc. (versioned)
+            - history (alias to history_0)
+
+        IMPORTANT: Store a deep copy, not a reference. The caller may
+        modify the list after calling this method.
+
+        Args:
+            message_history: List of message dicts (role, content).
+            history_index: Optional specific index, or None to auto-increment.
+
+        Returns:
+            The index used.
+        """
+        ...
+
+    def get_history_count(self) -> int:
+        """Return the number of histories added so far.
+
+        Used by RLM to inform the model how many conversation histories
+        are available.
+        """
+        ...

groknroll/environments/constants.py

@@ -0,0 +1,32 @@
+# Default packages for isolated REPL environments (Modal, Prime, etc.)
+
+APT_PACKAGES = [
+    "build-essential",
+    "git",
+    "curl",
+    "wget",
+    "libopenblas-dev",
+    "liblapack-dev",
+]
+
+PIP_PACKAGES = [
+    # Data science essentials
+    "numpy>=1.26.0",
+    "pandas>=2.1.0",
+    "scipy>=1.11.0",
+    # Math & symbolic computation
+    "sympy>=1.12",
+    # HTTP & APIs
+    "requests>=2.31.0",
+    "httpx>=0.25.0",
+    "flask>=3.0.0",
+    # Data formats
+    "pyyaml>=6.0",
+    "toml>=0.10.2",
+    # Utilities
+    "tqdm>=4.66.0",
+    "python-dateutil>=2.8.2",
+    "regex>=2023.0.0",
+    # For state serialization
+    "dill>=0.3.7",
+]