feat: docker sandbox function

Author: dosco

README.md

@@ -415,6 +415,14 @@ const res = await program.forward({
 console.log(res);
 ```
 
+## Built-in Functions
+
+| Function           | Name               | Description                                  |
+| ------------------ | ------------------ | -------------------------------------------- |
+| JS Interpreter     | AxJSInterpreter    | Execute JS code in a sandboxed env           |
+| Docker Sandbox     | AxDockerSession    | Execute commands within a docker environment |
+| Embeddings Adapter | AxEmbeddingAdapter | Fetch and pass embedding to your function    |
+
 ## Checkout all the examples
 
 Use the `tsx` command to run the examples it makes node run typescript code. It also support using a `.env` file to pass the AI API Keys as opposed to putting them in the commandline.
@@ -443,14 +451,7 @@ OPENAI_APIKEY=openai_key npm run tsx ./src/examples/marketing.ts
 | smart-hone.ts       | Agent looks for dog in smart home                       |
 | multi-modal.ts      | Use an image input along with other text inputs         |
 | balancer.ts         | Balance between various llm's based on cost, etc        |
-
-## Built-in Functions
-
-| Function           | Description                                            |
-| ------------------ | ------------------------------------------------------ |
-| JS Interpreter     | Used by the LLM to execute JS code in a sandboxed env. |
-| Docker Sandbox     | LLM can Execute commands within a docker environment   |
-| Embeddings Adapter | Wrapper to fetch and pass embedding to your function   |
+| docker.ts           | Use the docker sandbox to find files by description     |
 
 ## Our Goal
 

src/ax/dsp/generate.ts

@@ -77,7 +77,7 @@ export class AxGenerate<
   private pt: AxPromptTemplate;
   private asserts: AxAssertion[];
   private streamingAsserts: AxStreamingAssertion[];
-  private options?: AxGenerateOptions;
+  private options?: Omit<AxGenerateOptions, 'functions'>;
 
   private functions?: AxFunction[];
   private funcProc?: AxFunctionProcessor;
@@ -90,7 +90,7 @@ export class AxGenerate<
   ) {
     super(signature);
 
-    this.functions = this.options?.functions?.map((f) => {
+    this.functions = options?.functions?.map((f) => {
       if ('toFunction' in f) {
         return f.toFunction();
       }

src/ax/funcs/code.ts

@@ -70,7 +70,8 @@ export class AxJSInterpreter {
         required: ['code']
       },
 
-      func: this.codeInterpreterJavascript
+      func: ({ code }: Readonly<{ code: string }>) =>
+        this.codeInterpreterJavascript(code)
     };
   }
 }

src/ax/funcs/docker.ts

@@ -1,5 +1,58 @@
 import type { AxFunction } from '../ai/types.js';
 
+export interface AxDockerContainer {
+  Id: string;
+  Names: string[];
+  Image: string;
+  ImageID: string;
+  Command: string;
+  Created: number;
+  State: {
+    Status: string;
+    Running: boolean;
+    Paused: boolean;
+    Restarting: boolean;
+    OOMKilled: boolean;
+    Dead: boolean;
+    Pid: number;
+    ExitCode: number;
+    Error: string;
+    StartedAt: Date;
+    FinishedAt: Date;
+  };
+  Status: string;
+  Ports: Array<{
+    IP: string;
+    PrivatePort: number;
+    PublicPort: number;
+    Type: string;
+  }>;
+  Labels: { [key: string]: string };
+  SizeRw: number;
+  SizeRootFs: number;
+  HostConfig: {
+    NetworkMode: string;
+  };
+  NetworkSettings: {
+    Networks: {
+      [key: string]: {
+        IPAddress: string;
+        IPPrefixLen: number;
+        Gateway: string;
+        MacAddress: string;
+      };
+    };
+  };
+  Mounts: Array<{
+    Type: string;
+    Source: string;
+    Destination: string;
+    Mode: string;
+    RW: boolean;
+    Propagation: string;
+  }>;
+}
+
 export class AxDockerSession {
   private readonly apiUrl: string;
   private containerId: string | null = null;
@@ -26,31 +79,40 @@ export class AxDockerSession {
 
   async createContainer({
     imageName,
-    volumes = [], // Example format: [{ hostPath: '/host/path', containerPath: '/container/path' }]
-    doNotPullImage
+    volumes = [],
+    doNotPullImage,
+    tag
   }: Readonly<{
     imageName: string;
     volumes?: Array<{ hostPath: string; containerPath: string }>;
     doNotPullImage?: boolean;
+    tag?: string;
   }>) {
     const binds = volumes.map((v) => `${v.hostPath}:${v.containerPath}`);
 
     if (!doNotPullImage) {
       await this.pullImage(imageName);
     }
 
+    const containerConfig = {
+      Image: imageName,
+      Tty: true,
+      OpenStdin: false,
+      AttachStdin: false,
+      AttachStdout: false,
+      AttachStderr: false,
+      HostConfig: { Binds: binds },
+      Labels: {} as Record<string, string>
+    };
+
+    if (tag) {
+      containerConfig.Labels['com.example.tag'] = tag;
+    }
+
     const response = await this.fetchDockerAPI(`/containers/create`, {
       method: 'POST',
       headers: { 'Content-Type': 'application/json' },
-      body: JSON.stringify({
-        Image: imageName,
-        Tty: true,
-        OpenStdin: false,
-        AttachStdin: false,
-        AttachStdout: false,
-        AttachStderr: false,
-        HostConfig: { Binds: binds }
-      })
+      body: JSON.stringify(containerConfig)
     });
 
     if (!response.ok) {
@@ -63,6 +125,47 @@ export class AxDockerSession {
     return data;
   }
 
+  async findOrCreateContainer({
+    imageName,
+    volumes = [],
+    doNotPullImage,
+    tag
+  }: Readonly<{
+    imageName: string;
+    volumes?: Array<{ hostPath: string; containerPath: string }>;
+    doNotPullImage?: boolean;
+    tag: string;
+  }>): Promise<{ Id: string; isNew: boolean }> {
+    // First, try to find existing containers with the given tag
+    const existingContainers = await this.listContainers(true);
+    const matchingContainers = existingContainers.filter(
+      (container) =>
+        container.Labels && container.Labels['com.example.tag'] === tag
+    );
+
+    if (matchingContainers && matchingContainers.length > 0) {
+      // Randomly select a container from the matching ones
+      const randomIndex = Math.floor(Math.random() * matchingContainers.length);
+      const selectedContainer = matchingContainers[randomIndex];
+
+      if (selectedContainer) {
+        // Connect to the selected container
+        await this.connectToContainer(selectedContainer.Id);
+        return { Id: selectedContainer.Id, isNew: false };
+      }
+    }
+
+    // If no container with the tag exists, create a new one
+    const newContainer = await this.createContainer({
+      imageName,
+      volumes,
+      doNotPullImage,
+      tag
+    });
+
+    return { Id: newContainer.Id, isNew: true };
+  }
+
   async startContainer(): Promise<void> {
     if (!this.containerId) {
       throw new Error('No container created or connected');
@@ -92,11 +195,99 @@ export class AxDockerSession {
     this.containerId = containerId;
   }
 
+  async stopContainers({
+    tag,
+    remove,
+    timeout = 10
+  }: Readonly<{ tag?: string; remove?: boolean; timeout?: number }>): Promise<
+    Array<{ Id: string; Action: 'stopped' | 'removed' }>
+  > {
+    const results: Array<{ Id: string; Action: 'stopped' | 'removed' }> = [];
+
+    // List all containers
+    const containers = await this.listContainers(true);
+
+    // Filter containers by tag if provided
+    const targetContainers = tag
+      ? containers.filter(
+          (container) => container.Labels['com.example.tag'] === tag
+        )
+      : containers;
+
+    for (const container of targetContainers) {
+      // Stop the container if it's running
+      if (container.State.Status === 'running') {
+        const stopResponse = await this.fetchDockerAPI(
+          `/containers/${container.Id}/stop?t=${timeout}`,
+          { method: 'POST' }
+        );
+
+        if (!stopResponse.ok) {
+          console.warn(
+            `Failed to stop container ${container.Id}: ${stopResponse.statusText}`
+          );
+          continue;
+        }
+
+        results.push({ Id: container.Id, Action: 'stopped' });
+      }
+
+      // Remove the container if the remove flag is set
+      if (remove) {
+        const removeResponse = await this.fetchDockerAPI(
+          `/containers/${container.Id}`,
+          { method: 'DELETE' }
+        );
+
+        if (!removeResponse.ok) {
+          console.warn(
+            `Failed to remove container ${container.Id}: ${removeResponse.statusText}`
+          );
+          continue;
+        }
+
+        results.push({ Id: container.Id, Action: 'removed' });
+      }
+    }
+
+    return results;
+  }
+
+  async listContainers(all: boolean = false): Promise<AxDockerContainer[]> {
+    const response = await this.fetchDockerAPI(`/containers/json?all=${all}`, {
+      method: 'GET'
+    });
+    return response.json() as Promise<AxDockerContainer[]>;
+  }
+
+  async getContainerLogs(): Promise<string> {
+    if (!this.containerId) {
+      throw new Error('No container created or connected');
+    }
+    const response = await this.fetchDockerAPI(
+      `/containers/${this.containerId}/logs?stdout=true&stderr=true`,
+      { method: 'GET' }
+    );
+    return response.text();
+  }
+
   async executeCommand(command: string) {
+    console.log('Executing command:', command);
+
     if (!this.containerId) {
       throw new Error('No container created or connected');
     }
 
+    // Check container state
+    const containerInfo = await this.getContainerInfo(this.containerId);
+
+    if (containerInfo.State.Status !== 'running') {
+      await this.startContainer();
+
+      // Wait for the container to be in the "running" state
+      await this.waitForContainerToBeRunning(this.containerId);
+    }
+
     // Create exec instance
     const createResponse = await this.fetchDockerAPI(
       `/containers/${this.containerId}/exec`,
@@ -142,39 +333,33 @@ export class AxDockerSession {
     return await startResponse.text();
   }
 
-  async stopContainer() {
-    if (!this.containerId) {
-      throw new Error('No container created or connected');
-    }
+  // Add these new methods to the class:
 
+  private async getContainerInfo(
+    containerId: string
+  ): Promise<AxDockerContainer> {
     const response = await this.fetchDockerAPI(
-      `/containers/${this.containerId}/stop`,
-      {
-        method: 'POST'
-      }
+      `/containers/${containerId}/json`
     );
-
     if (!response.ok) {
-      throw new Error(`Failed to stop container: ${response.statusText}`);
+      throw new Error(`Failed to get container info: ${response.statusText}`);
     }
+    return response.json() as Promise<AxDockerContainer>;
   }
 
-  async listContainers(all: boolean = false) {
-    const response = await this.fetchDockerAPI(`/containers/json?all=${all}`, {
-      method: 'GET'
-    });
-    return response.json();
-  }
-
-  async getContainerLogs(): Promise<string> {
-    if (!this.containerId) {
-      throw new Error('No container created or connected');
+  private async waitForContainerToBeRunning(
+    containerId: string,
+    timeout: number = 30000
+  ): Promise<void> {
+    const startTime = Date.now();
+    while (Date.now() - startTime < timeout) {
+      const containerInfo = await this.getContainerInfo(containerId);
+      if (containerInfo.State.Status === 'running') {
+        return;
+      }
+      await new Promise((resolve) => setTimeout(resolve, 1000)); // Wait for 1 second before checking again
     }
-    const response = await this.fetchDockerAPI(
-      `/containers/${this.containerId}/logs?stdout=true&stderr=true`,
-      { method: 'GET' }
-    );
-    return response.text();
+    throw new Error('Timeout waiting for container to start');
   }
 
   private async fetchDockerAPI(
@@ -202,7 +387,8 @@ export class AxDockerSession {
         required: ['command']
       },
 
-      func: this.executeCommand
+      func: async ({ command }: Readonly<{ command: string }>) =>
+        await this.executeCommand(command)
     };
   }
 }

src/ax/funcs/embed.ts

@@ -70,7 +70,8 @@ export class AxEmbeddingAdapter {
         },
         required: ['text']
       },
-      func: this.embedAdapter
+      func: ({ text }: Readonly<{ text: string }>, options) =>
+        this.embedAdapter(text, options)
     };
   }
 }