@cloudwarriors-ai/rlm 0.1.7 → 0.1.9

package/src/infrastructure/sandbox/prelude/rlm_prelude.py

@@ -90,6 +90,65 @@ def llm_query(query: str, *context_vars: str) -> str:
     return response.get("result", "")
 
 
+def llm_query_batch(queries: list[tuple[str, str]]) -> list[str]:
+    """
+    Make multiple LLM queries in PARALLEL for significant speedup.
+
+    Per MIT RLM paper: Sequential sub-calls are a runtime bottleneck.
+    Use this instead of multiple llm_query() calls when queries are independent.
+
+    Args:
+        queries: List of (query, context) tuples. Each context is a string.
+
+    Returns:
+        List of responses in the same order as queries.
+
+    Example:
+        # Instead of this (sequential, slow):
+        results = []
+        for chunk in chunks:
+            result = llm_query("Analyze this", chunk)
+            results.append(result)
+
+        # Do this (parallel, fast):
+        queries = [("Analyze this", chunk) for chunk in chunks]
+        results = llm_query_batch(queries)
+    """
+    if not queries:
+        return []
+
+    # Build batch payload
+    batch = []
+    for query, context in queries:
+        # Resolve context if it's a variable name
+        if context in _context_store:
+            context = _context_store[context]
+        batch.append({
+            "query": query,
+            "context": [context] if isinstance(context, str) else list(context)
+        })
+
+    # Send batch command
+    _send_command("llm_query_batch", {"queries": batch})
+
+    # Wait for batch response
+    response = _wait_for_response()
+
+    if "error" in response:
+        raise RuntimeError(f"Batch query failed: {response['error']}")
+
+    results = response.get("results", [])
+
+    # Check for individual errors
+    final_results = []
+    for i, r in enumerate(results):
+        if isinstance(r, dict) and "error" in r:
+            raise RuntimeError(f"Query {i} failed: {r['error']}")
+        final_results.append(r)
+
+    return final_results
+
+
 def set_result(result: str) -> None:
     """
     Set the final result of the RLM execution.
@@ -105,6 +164,30 @@ def set_result(result: str) -> None:
     _send_command("set_result", {"result": _result})
 
 
+def set_result_final(result: str, confidence: float = 1.0) -> None:
+    """
+    Set result AND signal that the answer is complete.
+
+    Use this when you have found a definitive answer and no further
+    processing is needed. This is more efficient than set_result()
+    when you're confident the answer is complete.
+
+    Args:
+        result: The final answer
+        confidence: How confident (0.0-1.0) that this is complete
+
+    Example:
+        # Found the specific file requested
+        set_result_final("The main entry point is src/index.ts", confidence=1.0)
+    """
+    global _result
+    _result = str(result)
+    _send_command("set_result_final", {
+        "result": _result,
+        "confidence": confidence
+    })
+
+
 def set_variable(name: str, value: str) -> None:
     """
     Store a variable for use in subsequent code or queries.
@@ -213,6 +296,500 @@ def count_tokens(text: str) -> int:
     return len(text) // 4
 
 
+# Unit index cache for lazy structure extraction
+_unit_index: dict[str, dict] = {}
+
+
+def _extract_codebase_structure(content: str) -> dict:
+    """Extract structure from codebase format (=== FILE: ... ===)."""
+    import re
+    units = []
+    file_pattern = re.compile(r'^=== FILE: (.+?) ===$', re.MULTILINE)
+
+    matches = list(file_pattern.finditer(content))
+    for i, match in enumerate(matches):
+        file_path = match.group(1)
+        start = match.end()
+        end = matches[i + 1].start() if i + 1 < len(matches) else len(content)
+
+        file_content = content[start:end].strip()
+        unit_id = f"file:{file_path}"
+
+        unit_info = {
+            'id': unit_id,
+            'type': 'file',
+            'path': file_path,
+            'start': match.start(),
+            'end': end,
+            'size': len(file_content),
+            'tokens': count_tokens(file_content),
+        }
+        units.append(unit_info)
+        _unit_index[unit_id] = {**unit_info, 'content_start': start, 'content_end': end}
+
+    # Group by directory
+    directories = {}
+    for unit in units:
+        parts = unit['path'].split('/')
+        dir_path = '/'.join(parts[:-1]) if len(parts) > 1 else '.'
+        if dir_path not in directories:
+            directories[dir_path] = []
+        directories[dir_path].append(unit['path'].split('/')[-1])
+
+    return {
+        'type': 'codebase',
+        'total_files': len(units),
+        'total_tokens': sum(u['tokens'] for u in units),
+        'directories': directories,
+        'units': [{'id': u['id'], 'path': u['path'], 'tokens': u['tokens']} for u in units],
+    }
+
+
+def _extract_markdown_structure(content: str) -> dict:
+    """Extract structure from markdown (# headings)."""
+    import re
+    units = []
+    heading_pattern = re.compile(r'^(#{1,6})\s+(.+)$', re.MULTILINE)
+
+    matches = list(heading_pattern.finditer(content))
+    for i, match in enumerate(matches):
+        level = len(match.group(1))
+        title = match.group(2).strip()
+        start = match.start()
+        end = matches[i + 1].start() if i + 1 < len(matches) else len(content)
+
+        section_content = content[start:end].strip()
+        unit_id = f"section:{i}:{title[:50]}"
+
+        unit_info = {
+            'id': unit_id,
+            'type': 'section',
+            'level': level,
+            'title': title,
+            'start': start,
+            'end': end,
+            'size': len(section_content),
+            'tokens': count_tokens(section_content),
+        }
+        units.append(unit_info)
+        _unit_index[unit_id] = {**unit_info, 'content_start': start, 'content_end': end}
+
+    return {
+        'type': 'markdown',
+        'total_sections': len(units),
+        'total_tokens': sum(u['tokens'] for u in units),
+        'units': [{'id': u['id'], 'level': u['level'], 'title': u['title'], 'tokens': u['tokens']} for u in units],
+    }
+
+
+def _extract_generic_structure(content: str) -> dict:
+    """Extract structure by splitting into equal chunks."""
+    chunk_size = 50000  # ~12.5K tokens per chunk
+    units = []
+
+    for i in range(0, len(content), chunk_size):
+        chunk = content[i:i + chunk_size]
+        unit_id = f"chunk:{i // chunk_size}"
+
+        unit_info = {
+            'id': unit_id,
+            'type': 'chunk',
+            'index': i // chunk_size,
+            'start': i,
+            'end': min(i + chunk_size, len(content)),
+            'size': len(chunk),
+            'tokens': count_tokens(chunk),
+        }
+        units.append(unit_info)
+        _unit_index[unit_id] = {**unit_info, 'content_start': i, 'content_end': min(i + chunk_size, len(content))}
+
+    return {
+        'type': 'generic',
+        'total_chunks': len(units),
+        'total_tokens': sum(u['tokens'] for u in units),
+        'units': [{'id': u['id'], 'index': u['index'], 'tokens': u['tokens']} for u in units],
+    }
+
+
+def get_structure(content: str = None) -> dict:
+    """
+    Extract structure lazily from content without loading full content into memory.
+
+    Args:
+        content: Content to analyze (defaults to 'context' variable)
+
+    Returns:
+        Structure dict with type, units summary, and metadata
+
+    Example:
+        structure = get_structure()
+        print(f"Found {structure['total_files']} files")
+        for unit in structure['units']:
+            print(f"  {unit['path']}: {unit['tokens']} tokens")
+    """
+    if content is None:
+        content = _context_store.get('context', '')
+
+    # Clear previous index
+    _unit_index.clear()
+
+    # Detect type and extract structure markers
+    if '=== FILE:' in content:
+        return _extract_codebase_structure(content)
+    elif content.lstrip().startswith('#'):
+        return _extract_markdown_structure(content)
+    else:
+        return _extract_generic_structure(content)
+
+
+def get_unit(unit_id: str, content: str = None) -> str:
+    """
+    Fetch a specific unit by ID.
+
+    Args:
+        unit_id: The unit ID from get_structure()
+        content: Content to extract from (defaults to 'context' variable)
+
+    Returns:
+        The content of the requested unit
+
+    Example:
+        structure = get_structure()
+        file_content = get_unit(structure['units'][0]['id'])
+    """
+    if content is None:
+        content = _context_store.get('context', '')
+
+    # Ensure structure has been extracted
+    if not _unit_index:
+        get_structure(content)
+
+    if unit_id not in _unit_index:
+        raise ValueError(f"Unit not found: {unit_id}. Call get_structure() first.")
+
+    unit = _unit_index[unit_id]
+    return content[unit['content_start']:unit['content_end']].strip()
+
+
+def list_units(pattern: str = None, content: str = None) -> list[dict]:
+    """
+    List units matching an optional pattern.
+
+    Args:
+        pattern: Optional pattern to filter units (substring match on id/path)
+        content: Content to analyze (defaults to 'context' variable)
+
+    Returns:
+        List of unit info dicts with id, type, and tokens
+
+    Example:
+        # List all Python files
+        py_files = list_units('.py')
+        for f in py_files:
+            print(f"{f['id']}: {f['tokens']} tokens")
+    """
+    if content is None:
+        content = _context_store.get('context', '')
+
+    # Ensure structure has been extracted
+    if not _unit_index:
+        get_structure(content)
+
+    units = []
+    for unit_id, info in _unit_index.items():
+        if pattern is None or pattern in unit_id or pattern in info.get('path', ''):
+            units.append({
+                'id': unit_id,
+                'type': info['type'],
+                'tokens': info['tokens'],
+                'path': info.get('path'),
+                'title': info.get('title'),
+            })
+
+    return units
+
+
+def get_directory_tree(content: str = None) -> dict:
+    """
+    Get hierarchical directory structure from codebase content.
+
+    This function parses the codebase content and builds a nested dictionary
+    representing the directory structure.
+
+    Args:
+        content: Content to analyze (defaults to 'context' variable)
+
+    Returns:
+        Nested dict representing directory tree:
+        {
+          'module1': {
+            'files': ['main.py', 'utils.py'],
+            'dirs': {
+              'subdir1': {'files': [...], 'dirs': {...}},
+              ...
+            }
+          },
+          ...
+        }
+
+    Example:
+        tree = get_directory_tree()
+        for module, data in tree.items():
+            print(f"Module: {module}")
+            print(f"  Files: {data['files']}")
+            for subdir in data['dirs']:
+                print(f"  Subdir: {subdir}")
+    """
+    import re
+
+    if content is None:
+        content = _context_store.get('context', '')
+
+    # Parse all file paths
+    file_pattern = re.compile(r'^=== FILE: (.+?) ===$', re.MULTILINE)
+    paths = [m.group(1) for m in file_pattern.finditer(content)]
+
+    # Build tree
+    tree = {}
+
+    def ensure_path(tree_node: dict, parts: list[str], filename: str):
+        """Recursively ensure path exists and add file."""
+        if len(parts) == 0:
+            # At the target directory level
+            if 'files' not in tree_node:
+                tree_node['files'] = []
+            tree_node['files'].append(filename)
+            if 'dirs' not in tree_node:
+                tree_node['dirs'] = {}
+        else:
+            # Need to traverse deeper
+            if 'dirs' not in tree_node:
+                tree_node['dirs'] = {}
+            if 'files' not in tree_node:
+                tree_node['files'] = []
+            dir_name = parts[0]
+            if dir_name not in tree_node['dirs']:
+                tree_node['dirs'][dir_name] = {'files': [], 'dirs': {}}
+            ensure_path(tree_node['dirs'][dir_name], parts[1:], filename)
+
+    for path in paths:
+        parts = path.split('/')
+        if len(parts) == 1:
+            # Root level file
+            if '_root' not in tree:
+                tree['_root'] = {'files': [], 'dirs': {}}
+            tree['_root']['files'].append(parts[0])
+        else:
+            # File in a directory
+            top_level = parts[0]
+            if top_level not in tree:
+                tree[top_level] = {'files': [], 'dirs': {}}
+            if len(parts) == 2:
+                tree[top_level]['files'].append(parts[1])
+            else:
+                ensure_path(tree[top_level], parts[1:-1], parts[-1])
+
+    return tree
+
+
+def get_files_in_directory(directory: str, content: str = None) -> list[dict]:
+    """
+    Get all files in a specific directory with metadata.
+
+    Args:
+        directory: Directory path to list (e.g., "rapture/elastic")
+        content: Content to analyze (defaults to 'context' variable)
+
+    Returns:
+        List of file metadata dicts:
+        [{'id': 'file:path', 'path': 'path', 'name': 'file.py', 'tokens': 123, 'size': 456}, ...]
+
+    Example:
+        files = get_files_in_directory("rapture/api")
+        for f in files:
+            print(f"{f['name']}: {f['tokens']} tokens")
+    """
+    if content is None:
+        content = _context_store.get('context', '')
+
+    # Ensure structure has been extracted
+    if not _unit_index:
+        get_structure(content)
+
+    # Normalize directory path (remove trailing slash)
+    directory = directory.rstrip('/')
+
+    files = []
+    for unit_id, info in _unit_index.items():
+        if info['type'] != 'file':
+            continue
+
+        file_path = info.get('path', '')
+        file_dir = '/'.join(file_path.split('/')[:-1])
+        file_name = file_path.split('/')[-1]
+
+        # Match exact directory (not subdirectories)
+        if file_dir == directory:
+            files.append({
+                'id': unit_id,
+                'path': file_path,
+                'name': file_name,
+                'tokens': info.get('tokens', 0),
+                'size': info.get('size', 0),
+            })
+
+    return sorted(files, key=lambda f: f['name'])
+
+
+def get_directory_content(directory: str, max_tokens: int = 100000, content: str = None) -> tuple[str, list[str]]:
+    """
+    Get combined content of files in a directory up to token limit.
+
+    This is useful for feeding a directory's contents to llm_query.
+
+    Args:
+        directory: Directory path to get content from
+        max_tokens: Maximum total tokens to fetch (default: 100000)
+        content: Content to extract from (defaults to 'context' variable)
+
+    Returns:
+        Tuple of (combined_content, unfetched_file_ids)
+
+    Example:
+        content, remaining = get_directory_content("rapture/elastic", max_tokens=50000)
+        analysis = llm_query("Analyze this directory", content)
+        if remaining:
+            content2, _ = get_directory_content("rapture/elastic", max_tokens=50000)
+            # ... handle remaining files
+    """
+    if content is None:
+        content = _context_store.get('context', '')
+
+    # Get files in directory
+    files = get_files_in_directory(directory, content)
+    file_ids = [f['id'] for f in files]
+
+    # Use get_units_safe to fetch with token budget
+    return get_units_safe(file_ids, max_tokens=max_tokens, content=content)
+
+
+def get_all_directories(content: str = None) -> list[str]:
+    """
+    Get a flat list of all directories in the codebase.
+
+    Args:
+        content: Content to analyze (defaults to 'context' variable)
+
+    Returns:
+        List of directory paths (e.g., ["rapture", "rapture/api", "rapture/elastic", ...])
+
+    Example:
+        dirs = get_all_directories()
+        for d in dirs:
+            print(f"Directory: {d}")
+            files = get_files_in_directory(d)
+            print(f"  Files: {len(files)}")
+    """
+    if content is None:
+        content = _context_store.get('context', '')
+
+    # Ensure structure has been extracted
+    if not _unit_index:
+        get_structure(content)
+
+    # Collect all unique directories
+    directories = set()
+    for unit_id, info in _unit_index.items():
+        if info['type'] != 'file':
+            continue
+
+        file_path = info.get('path', '')
+        parts = file_path.split('/')
+
+        # Add all parent directories
+        for i in range(1, len(parts)):
+            dir_path = '/'.join(parts[:i])
+            directories.add(dir_path)
+
+    return sorted(directories)
+
+
+def get_module_directories(module: str, content: str = None) -> list[str]:
+    """
+    Get all directories within a specific top-level module.
+
+    Args:
+        module: Module name (e.g., "rapture", "rapture-frontend")
+        content: Content to analyze (defaults to 'context' variable)
+
+    Returns:
+        List of directory paths within the module
+
+    Example:
+        dirs = get_module_directories("rapture")
+        # Returns: ["rapture", "rapture/api", "rapture/elastic", ...]
+    """
+    all_dirs = get_all_directories(content)
+    return [d for d in all_dirs if d == module or d.startswith(module + '/')]
+
+
+def get_units_safe(unit_ids: list[str], max_tokens: int = 50000, content: str = None) -> tuple[str, list[str]]:
+    """
+    Fetch units up to a token budget.
+
+    This is the recommended way to fetch multiple units while staying within
+    the llm_query() capacity limit.
+
+    Args:
+        unit_ids: List of unit IDs to fetch
+        max_tokens: Maximum total tokens to fetch (default: 50000 = ~200K chars)
+        content: Content to extract from (defaults to 'context' variable)
+
+    Returns:
+        Tuple of (combined_content, unfetched_unit_ids)
+
+    Example:
+        structure = get_structure()
+        all_ids = [u['id'] for u in structure['units']]
+
+        # Fetch in batches that fit in llm_query()
+        fetched, remaining = get_units_safe(all_ids, max_tokens=100000)
+        result1 = llm_query("Analyze these files", fetched)
+
+        if remaining:
+            fetched2, remaining2 = get_units_safe(remaining, max_tokens=100000)
+            result2 = llm_query("Analyze these files", fetched2)
+    """
+    if content is None:
+        content = _context_store.get('context', '')
+
+    # Ensure structure has been extracted
+    if not _unit_index:
+        get_structure(content)
+
+    fetched_content = []
+    fetched_tokens = 0
+    unfetched = []
+
+    for unit_id in unit_ids:
+        if unit_id not in _unit_index:
+            unfetched.append(unit_id)
+            continue
+
+        unit_info = _unit_index[unit_id]
+        unit_tokens = unit_info['tokens']
+
+        if fetched_tokens + unit_tokens > max_tokens:
+            unfetched.append(unit_id)
+        else:
+            # Include the full unit with header marker so get_structure() works on result
+            unit_content = content[unit_info['start']:unit_info['content_end']].strip()
+            fetched_content.append(unit_content)
+            fetched_tokens += unit_tokens
+
+    return '\n\n'.join(fetched_content), unfetched
+
+
 # Safe modules that can be imported (defined first so _restricted_import can use it)
 _SAFE_MODULES = {
     'json': __import__('json'),
@@ -305,11 +882,24 @@ _SAFE_BUILTINS = {
 
     # RLM functions
     'llm_query': llm_query,
+    'llm_query_batch': llm_query_batch,
     'set_result': set_result,
+    'set_result_final': set_result_final,
     'set_variable': set_variable,
     'chunk_text': chunk_text,
     'filter_lines': filter_lines,
     'count_tokens': count_tokens,
+    # Structure extraction (lazy loading)
+    'get_structure': get_structure,
+    'get_unit': get_unit,
+    'list_units': list_units,
+    'get_units_safe': get_units_safe,
+    # Directory-aware functions (for hierarchical indexing)
+    'get_directory_tree': get_directory_tree,
+    'get_files_in_directory': get_files_in_directory,
+    'get_directory_content': get_directory_content,
+    'get_all_directories': get_all_directories,
+    'get_module_directories': get_module_directories,
 }
 
 
@@ -318,53 +908,59 @@ def main():
     # Read initial context
     try:
         init_line = sys.stdin.readline().strip()
-        if init_line.startswith("__INIT__:"):
+
+        # Handle large payloads via temp file (avoids stdin buffer issues)
+        if init_line.startswith("__INIT_FILE__:"):
+            temp_file = init_line[len("__INIT_FILE__:"):]
+            with open(temp_file, 'r', encoding='utf-8') as f:
+                init_data = json.load(f)
+        elif init_line.startswith("__INIT__:"):
             init_data = json.loads(init_line[len("__INIT__:"):])
+        else:
+            print("__ERROR__:Invalid initialization", flush=True)
+            return 1
 
-            # Load context variables
-            for name, value in init_data.get("context", {}).items():
-                _context_store[name] = value
+        # Load context variables
+        for name, value in init_data.get("context", {}).items():
+            _context_store[name] = value
 
-            # Get the code to execute
-            code = init_data.get("code", "")
+        # Get the code to execute
+        code = init_data.get("code", "")
 
-            if not code:
-                print("__ERROR__:No code provided", flush=True)
-                return 1
+        if not code:
+            print("__ERROR__:No code provided", flush=True)
+            return 1
 
-            # Create restricted execution environment
-            exec_globals = {
-                '__builtins__': _SAFE_BUILTINS,
-                '__name__': '__main__',
-                '__doc__': None,
+        # Create restricted execution environment
+        exec_globals = {
+            '__builtins__': _SAFE_BUILTINS,
+            '__name__': '__main__',
+            '__doc__': None,
+        }
+
+        # Add context variables to globals
+        exec_globals.update(_context_store)
+
+        # Execute the code
+        try:
+            exec(code, exec_globals)
+
+            # Check if result was set
+            if _result is None:
+                print("__WARNING__:No result was set. Call set_result() with your answer.", flush=True)
+
+            print("__DONE__", flush=True)
+            return 0
+
+        except Exception as e:
+            import traceback
+            tb = traceback.format_exc()
+            error_info = {
+                "error": str(e),
+                "type": type(e).__name__,
+                "traceback": tb
             }
-
-            # Add context variables to globals
-            exec_globals.update(_context_store)
-
-            # Execute the code
-            try:
-                exec(code, exec_globals)
-
-                # Check if result was set
-                if _result is None:
-                    print("__WARNING__:No result was set. Call set_result() with your answer.", flush=True)
-
-                print("__DONE__", flush=True)
-                return 0
-
-            except Exception as e:
-                import traceback
-                tb = traceback.format_exc()
-                error_info = {
-                    "error": str(e),
-                    "type": type(e).__name__,
-                    "traceback": tb
-                }
-                print(f"__ERROR__:{json.dumps(error_info)}", flush=True)
-                return 1
-        else:
-            print("__ERROR__:Invalid initialization", flush=True)
+            print(f"__ERROR__:{json.dumps(error_info)}", flush=True)
             return 1
 
     except Exception as e: