hdsp-jupyter-extension 2.0.23__py3-none-any.whl → 2.0.25__py3-none-any.whl

agent_server/context_providers/__init__.py

@@ -0,0 +1,22 @@
+"""
+Context Providers Module
+
+Provides context injection for LLM prompts (e.g., @file, @notebook, @selection).
+Also provides action commands (e.g., @reset).
+Inspired by jupyter-ai's context provider architecture.
+"""
+
+from .actions import ResetContextProvider
+from .base import BaseContextProvider, ContextCommand, find_commands
+from .file import FileContextProvider
+from .processor import ContextProcessor, process_context_commands
+
+__all__ = [
+    "ContextCommand",
+    "BaseContextProvider",
+    "FileContextProvider",
+    "ResetContextProvider",
+    "ContextProcessor",
+    "find_commands",
+    "process_context_commands",
+]

agent_server/context_providers/actions.py

@@ -0,0 +1,45 @@
+"""
+Action Context Providers
+
+These providers handle action commands like @reset that perform actions
+rather than providing context to the LLM.
+"""
+
+from .base import BaseContextProvider, ContextCommand, ListOptionsEntry
+
+
+class ResetContextProvider(BaseContextProvider):
+    """
+    Action provider for resetting the session.
+
+    Usage:
+        @reset - Clear conversation history and recreate agent
+    """
+
+    id = "@reset"
+    description = "Clear session and reset agent"
+    requires_arg = False
+    only_start = False
+
+    def __init__(self, base_dir: str = "."):
+        super().__init__(base_dir)
+
+    def get_arg_options(self, arg_prefix: str) -> list[ListOptionsEntry]:
+        """No arguments for @reset."""
+        return []
+
+    def make_context(self, command: ContextCommand) -> str:
+        """@reset doesn't provide context - it's an action command."""
+        return ""
+
+    def replace_command(self, command: ContextCommand) -> str:
+        """Replace @reset with empty string (it's an action, not context)."""
+        return ""
+
+    def is_action_command(self) -> bool:
+        """Indicate this is an action command."""
+        return True
+
+
+# Add more action commands here as needed
+# Example: @clear, @help, @history, etc.

agent_server/context_providers/base.py

@@ -0,0 +1,231 @@
+"""
+Base Context Provider Module
+
+Defines the base classes and utilities for context providers.
+Based on jupyter-ai's context provider architecture.
+"""
+
+import re
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class ContextCommand(BaseModel):
+    """
+    Represents a context command in user input (e.g., @file:path/to/file.py).
+
+    Properties:
+        cmd: Full command string (e.g., "@file:path/to/file.py")
+        id: Command identifier (e.g., "@file")
+        arg: Optional argument after the colon (e.g., "path/to/file.py")
+    """
+
+    cmd: str
+
+    @property
+    def id(self) -> str:
+        """Extract command ID (e.g., '@file' from '@file:path')"""
+        return self.cmd.partition(":")[0]
+
+    @property
+    def arg(self) -> Optional[str]:
+        """
+        Extract argument after colon.
+        Handles quoted paths and escaped spaces.
+        """
+        if ":" not in self.cmd:
+            return None
+
+        arg = self.cmd.partition(":")[2]
+        if not arg:
+            return None
+
+        # Remove surrounding quotes if present
+        arg = arg.strip("'\"")
+        # Handle escaped spaces
+        arg = arg.replace("\\ ", " ")
+
+        return arg
+
+
+class ListOptionsEntry(BaseModel):
+    """Entry for autocomplete options list."""
+
+    id: str  # e.g., "@file"
+    label: str  # e.g., "@file:" or "@file:path/to/file.py"
+    description: str  # e.g., "Include file contents" or "Python file"
+    only_start: bool = False  # Whether command only works at start of prompt
+    is_complete: bool = True  # Whether this option completes the command
+
+
+class ListOptionsResponse(BaseModel):
+    """Response for autocomplete options."""
+
+    options: list[ListOptionsEntry] = []
+
+
+class ContextProviderException(Exception):
+    """Exception raised by context providers."""
+
+    pass
+
+
+class BaseContextProvider(ABC):
+    """
+    Base class for context providers.
+
+    Context providers handle @commands in user input, providing context
+    to the LLM from external sources (files, notebooks, etc.).
+    """
+
+    # Unique identifier for this provider (e.g., "@file")
+    id: str = ""
+
+    # Human-readable description
+    description: str = ""
+
+    # Whether this provider requires an argument after the colon
+    requires_arg: bool = True
+
+    # Whether this command only works at the start of the prompt
+    only_start: bool = False
+
+    def __init__(self, base_dir: str = "."):
+        """
+        Initialize the context provider.
+
+        Args:
+            base_dir: Base directory for resolving relative paths
+        """
+        self.base_dir = base_dir
+
+    @property
+    def command_id(self) -> str:
+        """Get the command ID (e.g., '@file')."""
+        return self.id
+
+    @property
+    def pattern(self) -> str:
+        """
+        Regex pattern to match this context command.
+
+        Handles:
+        - @file:path/to/file
+        - @file:'path with spaces'
+        - @file:"path with spaces"
+        - @file:path\\ with\\ escaped\\ spaces
+        """
+        if self.requires_arg:
+            # Pattern for commands with arguments
+            # Matches: @file:path, @file:'quoted path', @file:"quoted path", @file:escaped\\ path
+            return (
+                rf"(?<![^\s.]){re.escape(self.command_id)}:"
+                rf"(?:'[^']+'|\"[^\"]+\"|[^\s\\]+(?:\\ [^\s\\]*)*)"
+            )
+        else:
+            # Pattern for commands without arguments
+            return rf"(?<![^\s.]){re.escape(self.command_id)}(?![^\s.])"
+
+    @abstractmethod
+    def get_arg_options(self, arg_prefix: str) -> list[ListOptionsEntry]:
+        """
+        Get autocomplete options for the command argument.
+
+        Args:
+            arg_prefix: Current argument prefix (e.g., "src/" for "@file:src/")
+
+        Returns:
+            List of autocomplete options
+        """
+        pass
+
+    @abstractmethod
+    def make_context(self, command: ContextCommand) -> str:
+        """
+        Generate context content for this command.
+
+        Args:
+            command: The parsed context command
+
+        Returns:
+            Context string to inject into the prompt
+
+        Raises:
+            ContextProviderException: If context cannot be generated
+        """
+        pass
+
+    def replace_command(self, command: ContextCommand) -> str:
+        """
+        Get the replacement text for this command in the cleaned prompt.
+
+        By default, replaces @file:path with 'path'.
+
+        Args:
+            command: The parsed context command
+
+        Returns:
+            Replacement text
+        """
+        if command.arg:
+            return f"'{command.arg}'"
+        return ""
+
+    def get_provider_option(self) -> ListOptionsEntry:
+        """Get the autocomplete option for this provider itself."""
+        return ListOptionsEntry(
+            id=self.command_id,
+            label=f"{self.command_id}:" if self.requires_arg else self.command_id,
+            description=self.description,
+            only_start=self.only_start,
+            is_complete=not self.requires_arg,
+        )
+
+
+def _is_inside_code_block(match: re.Match, text: str) -> bool:
+    """
+    Check if a regex match is inside a code block (backticks).
+
+    Args:
+        match: Regex match object
+        text: Full text being searched
+
+    Returns:
+        True if match is inside backticks
+    """
+    start = match.start()
+    before = text[:start]
+
+    # Count backticks before the match
+    # If odd number, we're inside a code block
+    single_backticks = before.count("`") - before.count("```") * 3
+    triple_backticks = before.count("```")
+
+    # Inside code block if odd number of backticks/triple backticks
+    return single_backticks % 2 == 1 or triple_backticks % 2 == 1
+
+
+def find_commands(provider: BaseContextProvider, text: str) -> list[ContextCommand]:
+    """
+    Find all context commands for a provider in the given text.
+
+    Args:
+        provider: The context provider to match
+        text: Text to search
+
+    Returns:
+        List of found context commands
+    """
+    matches = list(re.finditer(provider.pattern, text))
+    results = []
+
+    for match in matches:
+        # Skip if inside code block
+        if _is_inside_code_block(match, text):
+            continue
+
+        results.append(ContextCommand(cmd=match.group()))
+
+    return results

agent_server/context_providers/file.py

@@ -0,0 +1,316 @@
+"""
+File Context Provider
+
+Provides @file: context injection for including file contents in prompts.
+Based on jupyter-ai's FileContextProvider.
+"""
+
+import glob
+import os
+from typing import Optional
+
+import nbformat
+
+from .base import (
+    BaseContextProvider,
+    ContextCommand,
+    ContextProviderException,
+    ListOptionsEntry,
+)
+
+# Supported file extensions
+SUPPORTED_EXTS = {
+    ".py",
+    ".md",
+    ".txt",
+    ".json",
+    ".yaml",
+    ".yml",
+    ".toml",
+    ".ini",
+    ".cfg",
+    ".sh",
+    ".bash",
+    ".ipynb",
+    ".js",
+    ".ts",
+    ".jsx",
+    ".tsx",
+    ".html",
+    ".css",
+    ".scss",
+    ".sql",
+    ".r",
+    ".R",
+    ".Rmd",
+    ".jl",
+    ".csv",
+    ".xml",
+    ".env",
+    ".gitignore",
+    ".dockerignore",
+    "Dockerfile",
+    "Makefile",
+    ".tex",
+}
+
+# Template for file context
+FILE_CONTEXT_TEMPLATE = """The following is the content of the file `{filepath}`:
+
+```
+{content}
+```"""
+
+
+class FileContextProvider(BaseContextProvider):
+    """
+    Context provider for including file contents in prompts.
+
+    Usage:
+        @file:path/to/file.py
+        @file:'path with spaces/file.py'
+        @file:"another/path.py"
+    """
+
+    id = "@file"
+    description = "Include file contents in the prompt"
+    requires_arg = True
+    only_start = False
+
+    def __init__(self, base_dir: str = "."):
+        super().__init__(base_dir)
+
+    def get_arg_options(self, arg_prefix: str) -> list[ListOptionsEntry]:
+        """
+        Get autocomplete options for file paths.
+
+        Args:
+            arg_prefix: Current path prefix (e.g., "src/" or "data/train")
+
+        Returns:
+            List of matching file/directory options
+        """
+        options = []
+
+        # Resolve the path prefix
+        if os.path.isabs(arg_prefix):
+            path_prefix = arg_prefix
+        else:
+            path_prefix = os.path.join(self.base_dir, arg_prefix)
+
+        # Normalize path (handle .. and .)
+        path_prefix = os.path.normpath(path_prefix)
+
+        # If the original prefix ends with / and it's a directory, list contents
+        if arg_prefix.endswith("/") and os.path.isdir(path_prefix):
+            glob_pattern = os.path.join(path_prefix, "*")
+        else:
+            # Otherwise, use the prefix to match
+            glob_pattern = path_prefix + "*"
+
+        # Use glob to find matching paths
+        try:
+            path_matches = glob.glob(glob_pattern)
+        except Exception:
+            path_matches = []
+
+        # Sort: directories first, then files
+        path_matches = sorted(
+            path_matches,
+            key=lambda p: (not os.path.isdir(p), os.path.basename(p).lower()),
+        )
+
+        # Limit results
+        for path in path_matches[:20]:
+            is_dir = os.path.isdir(path)
+
+            # For files, check if extension is supported
+            if not is_dir:
+                ext = os.path.splitext(path)[1]
+                basename = os.path.basename(path)
+                # Check extension or special filenames
+                if ext not in SUPPORTED_EXTS and basename not in SUPPORTED_EXTS:
+                    continue
+
+            # Make relative path for display
+            try:
+                if os.path.isabs(arg_prefix):
+                    display_path = path
+                else:
+                    display_path = os.path.relpath(path, self.base_dir)
+            except ValueError:
+                display_path = path
+
+            # Add trailing slash for directories
+            label_path = display_path + "/" if is_dir else display_path
+
+            # Handle paths with spaces
+            if " " in label_path:
+                label_path = f"'{label_path}'"
+
+            options.append(
+                ListOptionsEntry(
+                    id=self.command_id,
+                    label=f"{self.command_id}:{label_path}",
+                    description="Directory"
+                    if is_dir
+                    else self._get_file_description(path),
+                    only_start=self.only_start,
+                    is_complete=not is_dir,  # Complete only for files
+                )
+            )
+
+        return options
+
+    def _get_file_description(self, filepath: str) -> str:
+        """Get a description for a file based on its extension."""
+        ext = os.path.splitext(filepath)[1].lower()
+        basename = os.path.basename(filepath)
+
+        descriptions = {
+            ".py": "Python file",
+            ".ipynb": "Jupyter notebook",
+            ".md": "Markdown file",
+            ".txt": "Text file",
+            ".json": "JSON file",
+            ".yaml": "YAML file",
+            ".yml": "YAML file",
+            ".toml": "TOML file",
+            ".js": "JavaScript file",
+            ".ts": "TypeScript file",
+            ".tsx": "TypeScript React file",
+            ".jsx": "JavaScript React file",
+            ".html": "HTML file",
+            ".css": "CSS file",
+            ".sql": "SQL file",
+            ".sh": "Shell script",
+            ".csv": "CSV file",
+            ".r": "R file",
+            ".R": "R file",
+            ".jl": "Julia file",
+        }
+
+        # Check special filenames
+        if basename == "Dockerfile":
+            return "Dockerfile"
+        if basename == "Makefile":
+            return "Makefile"
+
+        return descriptions.get(ext, "File")
+
+    def make_context(self, command: ContextCommand) -> str:
+        """
+        Read file contents and format for LLM.
+
+        Args:
+            command: The @file command
+
+        Returns:
+            Formatted file context
+
+        Raises:
+            ContextProviderException: If file cannot be read
+        """
+        filepath = command.arg
+        if not filepath:
+            raise ContextProviderException("No file path specified")
+
+        # Resolve path
+        if not os.path.isabs(filepath):
+            filepath = os.path.join(self.base_dir, filepath)
+
+        filepath = os.path.normpath(filepath)
+
+        # Validation
+        if not os.path.exists(filepath):
+            raise ContextProviderException(f"File not found: {filepath}")
+
+        if os.path.isdir(filepath):
+            raise ContextProviderException(f"Cannot read directory: {filepath}")
+
+        # Check extension
+        ext = os.path.splitext(filepath)[1]
+        basename = os.path.basename(filepath)
+        if ext not in SUPPORTED_EXTS and basename not in SUPPORTED_EXTS:
+            raise ContextProviderException(
+                f"Unsupported file type: {ext}. "
+                f"Supported types: {', '.join(sorted(SUPPORTED_EXTS))}"
+            )
+
+        # Read file
+        try:
+            with open(filepath, "r", encoding="utf-8") as f:
+                content = f.read()
+        except UnicodeDecodeError:
+            raise ContextProviderException(
+                f"Cannot read file (encoding error): {filepath}"
+            )
+        except PermissionError:
+            raise ContextProviderException(f"Permission denied: {filepath}")
+        except Exception as e:
+            raise ContextProviderException(f"Error reading file: {e}")
+
+        # Process content (special handling for notebooks)
+        content = self._process_file(content, filepath)
+
+        # Truncate if too long (100K chars max)
+        max_chars = 100_000
+        if len(content) > max_chars:
+            content = (
+                content[:max_chars]
+                + f"\n\n... [truncated, {len(content) - max_chars} more characters]"
+            )
+
+        # Format with template
+        return FILE_CONTEXT_TEMPLATE.format(
+            filepath=filepath,
+            content=content,
+        )
+
+    def _process_file(self, content: str, filepath: str) -> str:
+        """
+        Process file content, with special handling for certain file types.
+
+        Args:
+            content: Raw file content
+            filepath: Path to the file
+
+        Returns:
+            Processed content
+        """
+        if filepath.endswith(".ipynb"):
+            # Extract cell contents from Jupyter notebook
+            try:
+                nb = nbformat.reads(content, as_version=4)
+                cells = []
+                for i, cell in enumerate(nb.cells):
+                    cell_type = cell.cell_type
+                    source = cell.source
+                    cells.append(f"# Cell {i + 1} ({cell_type}):\n{source}")
+                return "\n\n".join(cells)
+            except Exception:
+                # If parsing fails, return raw content
+                return content
+
+        return content
+
+    def replace_command(self, command: ContextCommand) -> str:
+        """
+        Get replacement text for the command in the cleaned prompt.
+
+        Replaces @file:path with 'path'.
+        """
+        filepath = command.arg or ""
+        return f"'{filepath}'"
+
+
+# Singleton instance with default base directory
+_file_provider: Optional[FileContextProvider] = None
+
+
+def get_file_provider(base_dir: str = ".") -> FileContextProvider:
+    """Get or create the file context provider singleton."""
+    global _file_provider
+    if _file_provider is None or _file_provider.base_dir != base_dir:
+        _file_provider = FileContextProvider(base_dir=base_dir)
+    return _file_provider