indent 0.1.13__py3-none-any.whl → 0.1.28__py3-none-any.whl

exponent/core/remote_execution/cli_rpc_types.py

@@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Any
 
 import msgspec
 import yaml
+from pydantic_ai.format_as_xml import format_as_xml
 
 if TYPE_CHECKING:
     from exponent_server.core.tools.edit_tool import (
@@ -28,6 +29,7 @@ class ToolInput(msgspec.Struct, tag_field="tool_name", omit_defaults=True):
         return msgspec.to_builtins(self)  # type: ignore[no-any-return]
 
     def to_text(self) -> str:
+        """This text is purely used for human debugging, it's not fed to the llm"""
         d = msgspec.to_builtins(self)
         del d["tool_name"]
         return yaml.dump(d)
@@ -41,14 +43,16 @@ class ToolResult(msgspec.Struct, tag_field="tool_name", omit_defaults=True):
         This provides a default textual representation of the tool result. Override it as needed for your tool."""
         d = msgspec.to_builtins(self)
         del d["tool_name"]
-        return yaml.dump(d)
+        return format_as_xml(d, include_root_tag=False, item_tag="item")
 
 
 class ErrorToolResult(ToolResult, tag="error"):
     error_message: str
+    is_assistant_error: bool = False
 
 
 READ_TOOL_NAME = "read"
+READ_TOOL_ARTIFACT_NAME = "read_tool_artifact"
 
 
 class ReadToolInput(ToolInput, tag=READ_TOOL_NAME):
@@ -62,14 +66,28 @@ class FileMetadata(msgspec.Struct):
     file_mode: str
 
 
+class ReadToolArtifactResult(ToolResult, tag=READ_TOOL_ARTIFACT_NAME):
+    s3_uri: str
+    file_path: str
+    media_type: str
+
+    def to_text(self) -> str:
+        return f"[Image artifact uploaded to {self.s3_uri}]"
+
+
 class ReadToolResult(ToolResult, tag=READ_TOOL_NAME):
-    content: str
-    num_lines: int
-    start_line: int
-    total_lines: int
+    content: str | None = None
+    num_lines: int | None = None
+    start_line: int | None = None
+    total_lines: int | None = None
     metadata: FileMetadata | None = None
+    artifact: ReadToolArtifactResult | None = None
 
     def to_text(self) -> str:
+        if self.artifact:
+            return self.artifact.to_text()
+        assert self.content is not None
+        assert self.start_line is not None
         lines = self.content.splitlines()
         lines = [
             f"{str(i).rjust(6)}→{line}"
@@ -187,6 +205,42 @@ class BashToolResult(ToolResult, tag=BASH_TOOL_NAME):
     stopped_by_user: bool
 
 
+DOWNLOAD_ARTIFACT_TOOL_NAME = "download_artifact"
+
+
+class DownloadArtifactToolInput(ToolInput, tag=DOWNLOAD_ARTIFACT_TOOL_NAME):
+    presigned_url: str
+    file_path: str
+    artifact_id: str
+    overwrite: bool = True
+
+
+class DownloadArtifactToolResult(ToolResult, tag=DOWNLOAD_ARTIFACT_TOOL_NAME):
+    file_path: str
+    artifact_id: str
+    file_size_bytes: int
+    content_preview: str | None = None
+    num_lines: int | None = None
+    total_lines: int | None = None
+    truncated: bool = False
+
+
+UPLOAD_ARTIFACT_TOOL_NAME = "upload_artifact"
+
+
+class UploadArtifactToolInput(ToolInput, tag=UPLOAD_ARTIFACT_TOOL_NAME):
+    presigned_url: str
+    file_path: str
+    artifact_id: str
+    content_type: str
+
+
+class UploadArtifactToolResult(ToolResult, tag=UPLOAD_ARTIFACT_TOOL_NAME):
+    artifact_id: str
+    file_size_bytes: int
+    content_type: str
+
+
 class HttpRequest(msgspec.Struct, tag="http_fetch_cli"):
     url: str
     method: str = "GET"
@@ -210,17 +264,22 @@ ToolInputType = (
     | GrepToolInput
     | EditToolInput
     | BashToolInput
+    | DownloadArtifactToolInput
+    | UploadArtifactToolInput
 )
 PartialToolResultType = PartialBashToolResult
 
 ToolResultType = (
     ReadToolResult
+    | ReadToolArtifactResult
     | WriteToolResult
     | ListToolResult
     | GlobToolResult
     | GrepToolResult
     | EditToolResult
     | BashToolResult
+    | DownloadArtifactToolResult
+    | UploadArtifactToolResult
     | ErrorToolResult
 )
 
@@ -259,6 +318,10 @@ class TerminateResponse(msgspec.Struct, tag="terminate"):
     pass
 
 
+class TimeoutResponse(msgspec.Struct, tag="timeout"):
+    pass
+
+
 class BatchToolExecutionResponse(msgspec.Struct, tag="batch_tool_execution"):
     tool_results: list[ToolResultType]
 
@@ -271,6 +334,90 @@ class KeepAliveCliChatResponse(msgspec.Struct, tag="keep_alive_cli_chat"):
     pass
 
 
+class GenerateUploadUrlRequest(msgspec.Struct, tag="generate_upload_url"):
+    s3_key: str
+    content_type: str
+
+
+class GenerateUploadUrlResponse(msgspec.Struct, tag="generate_upload_url"):
+    upload_url: str
+    s3_uri: str
+
+
+# Terminal session management
+class StartTerminalRequest(msgspec.Struct, tag="start_terminal"):
+    """Start a new terminal session with PTY"""
+
+    session_id: str
+    command: list[str] | None = None  # None = default shell, or specific command
+    cols: int = 80
+    rows: int = 24
+    env: dict[str, str] | None = None  # Additional environment variables
+
+
+class StartTerminalResponse(msgspec.Struct, tag="start_terminal"):
+    """Response after starting terminal"""
+
+    session_id: str
+    success: bool
+    error_message: str | None = None
+
+
+# Terminal input (user typing)
+class TerminalInputRequest(msgspec.Struct, tag="terminal_input"):
+    """Send user input to terminal"""
+
+    session_id: str
+    data: str  # Raw input data from user
+
+
+class TerminalInputResponse(msgspec.Struct, tag="terminal_input"):
+    """Acknowledge input received"""
+
+    session_id: str
+    success: bool
+    error_message: str | None = None
+
+
+# Terminal resize
+class TerminalResizeRequest(msgspec.Struct, tag="terminal_resize"):
+    """Resize terminal dimensions"""
+
+    session_id: str
+    cols: int
+    rows: int
+
+
+class TerminalResizeResponse(msgspec.Struct, tag="terminal_resize"):
+    """Acknowledge resize"""
+
+    session_id: str
+    success: bool
+    error_message: str | None = None
+
+
+# Terminal stop
+class StopTerminalRequest(msgspec.Struct, tag="stop_terminal"):
+    """Stop a terminal session"""
+
+    session_id: str
+
+
+class StopTerminalResponse(msgspec.Struct, tag="stop_terminal"):
+    """Acknowledge terminal stopped"""
+
+    session_id: str
+    success: bool
+    error_message: str | None = None
+
+
+class StreamingCodeExecutionRequest(msgspec.Struct, tag="streaming_code_execution"):
+    correlation_id: str
+    language: str  # "python" or "shell"
+    content: str
+    timeout: int
+
+
 class CliRpcRequest(msgspec.Struct):
     request_id: str
     request: (
@@ -281,6 +428,12 @@ class CliRpcRequest(msgspec.Struct):
         | BatchToolExecutionRequest
         | SwitchCLIChatRequest
         | KeepAliveCliChatRequest
+        | GenerateUploadUrlRequest
+        | StartTerminalRequest
+        | TerminalInputRequest
+        | TerminalResizeRequest
+        | StopTerminalRequest
+        | StreamingCodeExecutionRequest
     )
 
 
@@ -292,6 +445,40 @@ class ErrorResponse(msgspec.Struct, tag="error"):
     error_message: str
 
 
+class StreamingCodeExecutionResponseChunk(
+    msgspec.Struct, tag="streaming_code_execution_chunk"
+):
+    correlation_id: str
+    content: str
+    truncated: bool = False
+
+    def add(
+        self, new_chunk: "StreamingCodeExecutionResponseChunk"
+    ) -> "StreamingCodeExecutionResponseChunk":
+        """Aggregates content of this and a new chunk."""
+        assert self.correlation_id == new_chunk.correlation_id
+        return StreamingCodeExecutionResponseChunk(
+            correlation_id=self.correlation_id,
+            content=self.content + new_chunk.content,
+            truncated=self.truncated or new_chunk.truncated,
+        )
+
+
+class StreamingCodeExecutionResponse(msgspec.Struct, tag="streaming_code_execution"):
+    correlation_id: str
+    content: str
+    truncated: bool = False
+    # Only present for shell code execution
+    cancelled_for_timeout: bool = False
+    exit_code: int | None = None
+    halted: bool = False
+
+
+class StreamingErrorResponse(msgspec.Struct, tag="streaming_error"):
+    correlation_id: str
+    error_message: str
+
+
 class CliRpcResponse(msgspec.Struct):
     request_id: str
     response: (
@@ -299,8 +486,17 @@ class CliRpcResponse(msgspec.Struct):
         | GetAllFilesResponse
         | ErrorResponse
         | TerminateResponse
+        | TimeoutResponse
         | BatchToolExecutionResponse
         | HttpResponse
         | SwitchCLIChatResponse
         | KeepAliveCliChatResponse
+        | GenerateUploadUrlResponse
+        | StartTerminalResponse
+        | TerminalInputResponse
+        | TerminalResizeResponse
+        | StopTerminalResponse
+        | StreamingCodeExecutionResponseChunk
+        | StreamingCodeExecutionResponse
+        | StreamingErrorResponse
     )