Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

use any-cloud-storage to abstract artifact storage #104

Open
wants to merge 12 commits into
base: main
Choose a base branch
from
13,390 changes: 11,163 additions & 2,227 deletions package-lock.json

Large diffs are not rendered by default.

10 changes: 10 additions & 0 deletions package.json
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,22 @@
}
},
"devDependencies": {
"@types/express": "^4.17.21",
"@types/firebase": "^3.2.1",
"@types/redis": "^4.0.11",
"@typescript-eslint/eslint-plugin": "^5.62.0",
"eslint": "^8.45.0",
"eslint-config-prettier": "^8.8.0",
"eslint-config-standard-with-typescript": "^36.1.0",
"husky": "^7.0.0",
"prettier": "^3.0.0",
"typescript": "^5.1.6"
},
"optionalDependencies": {
"any-cloud-storage": "^1.0.3"
},
"dependencies": {
"express": "^4.18.2",
"express-openapi-validator": "^5.1.5"
}
}
21 changes: 20 additions & 1 deletion packages/sdk/js/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,26 @@ async function taskHandler(taskInput: TaskInput | null): Promise<StepHandler> {
return stepHandler
}

Agent.handleTask(taskHandler).start()
const config = {
// port: 8000,
// workspace: './workspace'
}
Agent.handleTask(taskHandler, config).start()
```

Note: By default, artifacts will be saved/read from disk, but you can configure other options using the [any-cloud-storage](https://github.com/nalbion/any-cloud-storage) library:

```typescript
import { ArtifactStorageFactory } from 'agent-protocol/artifacts'
artifactStorage = ArtifactStorageFactory.create({
type: 's3',
bucket: 'my-bucket',
region: 'us-west-2',
// other AWS S3 configuration options...
})

const config = { artifactStorage }
Agent.handleTask(taskHandler, config).start()
```

## Docs
Expand Down
112 changes: 60 additions & 52 deletions packages/sdk/js/src/agent.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,4 @@
import { v4 as uuid } from 'uuid'
import * as fs from 'fs'
import * as path from 'path'

import {
type TaskInput,
Expand All @@ -14,13 +12,13 @@ import {
type TaskRequestBody,
StepStatus,
} from './models'
import {
createApi,
type ApiConfig,
type RouteRegisterFn,
type RouteContext,
} from './api'
import { type Router } from 'express'
import { createApi, type ApiConfig, type RouteRegisterFn } from './api'
import { type Router, type Express } from 'express'
import { FileStorage, type ArtifactStorage } from './artifacts'

export interface RouteContext {
agent: Agent
}

/**
* A function that handles a step in a task.
Expand Down Expand Up @@ -304,39 +302,16 @@ const registerGetArtifacts: RouteRegisterFn = (router: Router) => {
})
}

/**
* Get path of an artifact associated to a task
* @param taskId Task associated with artifact
* @param artifact Artifact associated with the path returned
* @returns Absolute path of the artifact
*/
export const getArtifactPath = (
taskId: string,
workspace: string,
artifact: Artifact
): string => {
const rootDir = path.isAbsolute(workspace)
? workspace
: path.join(process.cwd(), workspace)

return path.join(
rootDir,
taskId,
artifact.relative_path ?? '',
artifact.file_name
)
}

/**
* Creates an artifact for a task
* @param task Task associated with new artifact
* @param file File that will be added as artifact
* @param relativePath Relative path where the artifact might be stored. Can be undefined
*/
export const createArtifact = async (
workspace: string,
task: Task,
file: any,
agent: Agent,
file: Express.Multer.File,
relativePath?: string
): Promise<Artifact> => {
const artifactId = uuid()
Expand All @@ -353,16 +328,16 @@ export const createArtifact = async (
: []
task.artifacts.push(artifact)

const artifactFolderPath = getArtifactPath(task.task_id, workspace, artifact)

// Save file to server's file system
fs.mkdirSync(path.join(artifactFolderPath, '..'), { recursive: true })
fs.writeFileSync(artifactFolderPath, file.buffer)
// Save the file
const [storage, workspace] = agent.getArtifactStorageAndWorkspace(
task.task_id
)
await storage.writeArtifact(task.task_id, workspace, artifact, file)
return artifact
}
const registerCreateArtifact: RouteRegisterFn = (
router: Router,
context: RouteContext
agent: Agent
) => {
router.post('/agent/tasks/:task_id/artifacts', (req, res) => {
void (async () => {
Expand All @@ -381,13 +356,18 @@ const registerCreateArtifact: RouteRegisterFn = (

const files = req.files as Express.Multer.File[]
const file = files.find(({ fieldname }) => fieldname === 'file')
const artifact = await createArtifact(
context.workspace,
task[0],
file,
relativePath
)
res.status(200).json(artifact)

if (file == null) {
res.status(400).json({ message: 'No file found in the request' })
} else {
const artifact = await createArtifact(
task[0],
agent,
file,
relativePath
)
res.status(200).json(artifact)
}
} catch (err: Error | any) {
console.error(err)
res.status(500).json({ error: err.message })
Expand Down Expand Up @@ -417,17 +397,19 @@ export const getTaskArtifact = async (
}
const registerGetTaskArtifact: RouteRegisterFn = (
router: Router,
context: RouteContext
agent: Agent
) => {
router.get('/agent/tasks/:task_id/artifacts/:artifact_id', (req, res) => {
void (async () => {
const taskId = req.params.task_id
const artifactId = req.params.artifact_id
try {
const artifact = await getTaskArtifact(taskId, artifactId)
const artifactPath = getArtifactPath(
const [storage, workspace] =
agent.getArtifactStorageAndWorkspace(taskId)
const artifactPath = storage.getArtifactPath(
taskId,
context.workspace,
workspace,
artifact
)
res.status(200).sendFile(artifactPath)
Expand All @@ -442,18 +424,26 @@ const registerGetTaskArtifact: RouteRegisterFn = (
export interface AgentConfig {
port: number
workspace: string
artifactStorage: ArtifactStorage
}

export const defaultAgentConfig: AgentConfig = {
port: 8000,
workspace: './workspace',
artifactStorage: new FileStorage(),
}

export class Agent {
private workspace: string
private artifactStorage: ArtifactStorage

constructor(
public taskHandler: TaskHandler,
public config: AgentConfig
) {}
) {
this.artifactStorage = config.artifactStorage
this.workspace = config.workspace
}

static handleTask(
_taskHandler: TaskHandler,
Expand All @@ -463,6 +453,8 @@ export class Agent {
return new Agent(_taskHandler, {
workspace: config.workspace ?? defaultAgentConfig.workspace,
port: config.port ?? defaultAgentConfig.port,
artifactStorage:
config.artifactStorage ?? defaultAgentConfig.artifactStorage,
})
}

Expand All @@ -484,10 +476,26 @@ export class Agent {
console.log(`Agent listening at http://localhost:${this.config.port}`)
},
context: {
workspace: this.config.workspace,
agent: this,
},
}

createApi(config)
}

/**
* @param taskId (potentially) POST /agent/tasks { additional_input } could configure the artifactStorage and/or workspace for a Task
*/
getArtifactStorageAndWorkspace(taskId: string): [ArtifactStorage, string] {
return [this.artifactStorage, this.workspace]
}

/** It's easier for Serverless apps to configure artifactStorage after the app has been created */
setArtifactStorage(storage: ArtifactStorage): void {
this.artifactStorage = storage
}

setWorkspace(workspace: string): void {
this.workspace = workspace
}
}
9 changes: 3 additions & 6 deletions packages/sdk/js/src/api.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,11 @@ import express, { Router } from 'express' // <-- Import Router
import type * as core from 'express-serve-static-core'

import spec from '../../../../schemas/openapi.yml'
import { type Agent, type RouteContext } from './agent'

export type ApiApp = core.Express

export interface RouteContext {
workspace: string
}

export type RouteRegisterFn = (app: Router, context: RouteContext) => void
export type RouteRegisterFn = (app: Router, agent: Agent) => void

export interface ApiConfig {
context: RouteContext
Expand Down Expand Up @@ -44,7 +41,7 @@ export const createApi = (config: ApiConfig): void => {
const router = Router()

config.routes.forEach((route) => {
route(router, config.context)
route(router, config.context.agent)
})

app.use('/ap/v1', router)
Expand Down
19 changes: 19 additions & 0 deletions packages/sdk/js/src/artifacts/AnyCloudArtifactStorage.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
import { type FileStorage } from 'any-cloud-storage'
import ArtifactStorage from './ArtifactStorage'

export default class AnyCloudArtifactStorage extends ArtifactStorage {
constructor(private readonly storage: FileStorage) {
super()
}

protected override async saveFile(
artifactPath: string,
data: Buffer
): Promise<void> {
await this.storage.saveFile(artifactPath, data)
}

protected getAbsolutePath(filePath: string): string | Promise<string> {
return this.storage.getAbsolutePath(filePath)
}
}
57 changes: 57 additions & 0 deletions packages/sdk/js/src/artifacts/ArtifactStorage.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
import path from 'path'
import { type Artifact } from '../models'

/**
* @see ArtifactStorageFactory
*/
export default abstract class ArtifactStorage {
/**
* Save an artifact associated to a task
* @param taskId Task associated with artifact
* @param artifact Artifact associated with the path returned
* @returns Absolute path of the artifact
*/
async writeArtifact(
taskId: string,
workspace: string,
artifact: Artifact,
file: Express.Multer.File
): Promise<void> {
const artifactFolderPath = await this.getArtifactPath(
taskId,
workspace,
artifact
)
await this.saveFile(artifactFolderPath, file.buffer)
}

/**
* Get path of an artifact associated to a task
* @param taskId Task associated with artifact
* @param workspace The path to the workspace, defaults to './workspace'
* @param artifact Artifact associated with the path returned
* @returns Absolute path of the artifact
*/
async getArtifactPath(
taskId: string,
workspace: string,
artifact: Artifact
): Promise<string> {
const rootDir = await this.getAbsolutePath(workspace)

return path.join(
rootDir,
taskId,
artifact.relative_path ?? '',
artifact.file_name
)
}

protected getAbsolutePath(filePath: string): string | Promise<string> {
return path.isAbsolute(filePath)
? filePath
: path.join(process.cwd(), filePath)
}

protected async saveFile(artifactPath: string, data: Buffer): Promise<void> {}
}
13 changes: 13 additions & 0 deletions packages/sdk/js/src/artifacts/ArtifactStorageFactory.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
import createStorage, { type StorageType } from 'any-cloud-storage'
import type ArtifactStorage from './ArtifactStorage'
import { AnyCloudArtifactStorage } from './AnyCloudArtifactStorage'

export default class ArtifactStorageFactory {
static async create(
config: { type: StorageType } & Record<string, any>
): Promise<ArtifactStorage> {
const storage = await createStorage(config)

return new AnyCloudArtifactStorage(storage)
}
}
13 changes: 13 additions & 0 deletions packages/sdk/js/src/artifacts/FileStorage.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
import * as fs from 'fs'
import * as path from 'path'
import ArtifactStorage from './ArtifactStorage'

export default class FileStorage extends ArtifactStorage {
protected override async saveFile(
artifactPath: string,
data: Buffer
): Promise<void> {
fs.mkdirSync(path.join(artifactPath, '..'), { recursive: true })
fs.writeFileSync(artifactPath, data)
}
}
4 changes: 4 additions & 0 deletions packages/sdk/js/src/artifacts/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
export type { default as ArtifactStorage } from './ArtifactStorage'
export { default as ArtifactStorageFactory } from './ArtifactStorageFactory'
export { default as AnyCloudArtifactStorage } from './AnyCloudArtifactStorage'
export { default as FileStorage } from './FileStorage'