docker-compose/ai/functions/save_to_disk.py

"""
title: Save Code to Local Disk via SFTP
description: Saves generated code to local filesystem using SFTP
author: Claude
version: 1.0.0
"""

import paramiko
import os
from typing import Optional
from pydantic import BaseModel, Field


class Tools:
    class Valves(BaseModel):
        SFTP_HOST: str = Field(
            default="vps",
            description="SFTP host address"
        )
        SFTP_PORT: int = Field(
            default=22,
            description="SFTP port"
        )
        SFTP_USERNAME: str = Field(
            default="valknar",
            description="SFTP username"
        )
        SFTP_BASE_PATH: str = Field(
            default="/home/valknar/Projects",
            description="Base path for file operations"
        )
        SFTP_USE_KEY: bool = Field(
            default=True,
            description="Use SSH key authentication"
        )
        SFTP_KEY_PATH: str = Field(
            default="/app/.ssh/id_rsa",
            description="Path to SSH private key"
        )

    def __init__(self):
        self.valves = self.Valves()

    def save_file(
        self,
        filepath: str,
        content: str,
        __user__: dict = {},
    ) -> str:
        """
        Save content to a file on the local disk via SFTP.

        :param filepath: Relative path from base directory (e.g. "myproject/main.py")
        :param content: File content to save
        :return: Success message or error
        """

        try:
            # Create SSH client
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

            # Connect using SSH key or password
            if self.valves.SFTP_USE_KEY:
                ssh.connect(
                    hostname=self.valves.SFTP_HOST,
                    port=self.valves.SFTP_PORT,
                    username=self.valves.SFTP_USERNAME,
                    key_filename=self.valves.SFTP_KEY_PATH,
                )
            else:
                # For password auth, would need SFTP_PASSWORD in valves
                raise ValueError("Password auth not configured. Use SSH key authentication.")

            # Open SFTP session
            sftp = ssh.open_sftp()

            # Construct full path
            full_path = os.path.join(self.valves.SFTP_BASE_PATH, filepath)
            directory = os.path.dirname(full_path)

            # Create directories if they don't exist
            try:
                sftp.stat(directory)
            except FileNotFoundError:
                # Create parent directories recursively
                dirs = []
                temp_dir = directory
                while temp_dir and temp_dir != '/':
                    try:
                        sftp.stat(temp_dir)
                        break
                    except FileNotFoundError:
                        dirs.append(temp_dir)
                        temp_dir = os.path.dirname(temp_dir)

                for dir_path in reversed(dirs):
                    sftp.mkdir(dir_path)

            # Write file
            with sftp.open(full_path, 'w') as f:
                f.write(content)

            sftp.close()
            ssh.close()

            return f"✅ File saved successfully to: {full_path}"

        except Exception as e:
            return f"❌ Error saving file: {str(e)}"

    def read_file(
        self,
        filepath: str,
        __user__: dict = {},
    ) -> str:
        """
        Read content from a file on the local disk via SFTP.

        :param filepath: Relative path from base directory (e.g. "myproject/main.py")
        :return: File content or error message
        """

        try:
            # Create SSH client
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

            # Connect using SSH key
            ssh.connect(
                hostname=self.valves.SFTP_HOST,
                port=self.valves.SFTP_PORT,
                username=self.valves.SFTP_USERNAME,
                key_filename=self.valves.SFTP_KEY_PATH,
            )

            # Open SFTP session
            sftp = ssh.open_sftp()

            # Construct full path
            full_path = os.path.join(self.valves.SFTP_BASE_PATH, filepath)

            # Read file
            with sftp.open(full_path, 'r') as f:
                content = f.read()

            sftp.close()
            ssh.close()

            return content

        except Exception as e:
            return f"❌ Error reading file: {str(e)}"

    def list_files(
        self,
        directory: str = ".",
        __user__: dict = {},
    ) -> str:
        """
        List files in a directory on the local disk via SFTP.

        :param directory: Relative directory path from base (e.g. "myproject/")
        :return: List of files or error message
        """

        try:
            # Create SSH client
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

            # Connect using SSH key
            ssh.connect(
                hostname=self.valves.SFTP_HOST,
                port=self.valves.SFTP_PORT,
                username=self.valves.SFTP_USERNAME,
                key_filename=self.valves.SFTP_KEY_PATH,
            )

            # Open SFTP session
            sftp = ssh.open_sftp()

            # Construct full path
            full_path = os.path.join(self.valves.SFTP_BASE_PATH, directory)

            # List files
            files = sftp.listdir(full_path)

            sftp.close()
            ssh.close()

            return f"📁 Files in {full_path}:\n" + "\n".join(f"  - {f}" for f in files)

        except Exception as e:
            return f"❌ Error listing files: {str(e)}"