軽量 Langchain 的な何か

agent.test.ts

import { defineAgent, ok, err, runAgent, stepAgent, chain, initAgent } from "./agent.ts";
import { Agent, AgentError, AgentErrorCode } from "./types.ts";
import { expect } from "https://deno.land/std@0.214.0/expect/mod.ts";

Deno.test("stepAgent: simple", async () => {
  const simple = defineAgent((_init) => {
    return {
      async invoke(input, _options) {
        return `Hello, ${input}`;
      },
      async parse(invoked) {
        const r = invoked.match(/Hello, (\w+)/)?.[1].trim();
        if (r) {
          return r;
        } else {
          throw new AgentError(AgentErrorCode.ParseError, 'No match');
        }
      },
      async validate(parsed) {
        if (parsed[0].toUpperCase() !== parsed[0]) {
          throw new AgentError(AgentErrorCode.ValidationError, 'Not camelcase');
        }
      }
    }
  })({});

  const s0 = await initAgent(simple, "John", {});
  const s1 = await stepAgent(simple, s0);
  expect(s1).toEqual({
    errorCount: 0,
    step: "Invoked",
    input: "John",
    options: {},
    invoked: { ok: true, value: "Hello, John" }
  });

  const s2 = await stepAgent(simple, s1);
  expect(s2).toEqual({
    errorCount: 0,
    step: "Parsed",
    input: "John",
    options: {},
    invoked: { ok: true, value: "Hello, John" },
    parsed: { ok: true, value: "John" },
  });

  const s3 = await stepAgent(simple, s2);
  expect(s3).toEqual({
    errorCount: 0,
    step: "Done",
    input: "John",
    options: {},
    invoked: { ok: true, value: "Hello, John" },
    parsed: { ok: true, value: "John" },
    validated: { ok: true, value: undefined },
  });
});

Deno.test("chain", async () => {
  const a = defineAgent<{}, string, { invoked: string }, { parsed: string }>((_init) => {
    return {
      async invoke(input, _options) {
        return { invoked: input };
      },
      parse(result) {
        return { parsed: result.invoked };
      }
    }
  })({}) as Agent<{}, string, { invoked: string }, { parsed: string }>;

  const b = defineAgent<{}, { parsed: string }, { invoked2: string }, { parsed2: string }>((_init) => {
    return {
      async invoke(input, _options) {
        return { invoked2: input.parsed };
      },
      parse(result) {
        return { parsed2: result.invoked2 };
      }
    }
  })({}) as Agent<{}, { parsed: string }, { invoked2: string }, { parsed2: string }>;

  const c = defineAgent<{}, { xxxx: string }, { invoked2: string }, { parsed2: string }>((_init) => {
    return {
      async invoke(input, _options) {
        return { invoked2: input.xxxx };
      },
    }
  })({});


  const chained = chain(a, b);

  // @ts-expect-error
  const chained2 = chain(a, c);


  // const result1 = await runAgent(simple, "John");
  // expect(result1).toEqual(ok("John"));

  // // parse error
  // const result2 = await runAgent(simple, "");
  // expect(result2).toEqual(err({
  //   code: AgentErrorCode.ParseError,
  //   message: 'No match',
  // }));

  // // validation error
  // const result3 = await runAgent(simple, "john");
  // expect(result3).toEqual(err({
  //   code: AgentErrorCode.ValidationError,
  //   message: 'Not camelcase',
  // }));
});


Deno.test("runAgent: simple", async () => {
  const simple = defineAgent((_init) => {
    return {
      async invoke(input, _options) {
        return `Hello, ${input}`;
      },
      async parse(result) {
        const r = result.match(/Hello, (\w+)/)?.[1].trim();
        if (r) {
          return r;
        } else {
          throw new AgentError(AgentErrorCode.ParseError, 'No match');
        }
      },
      async validate(parsed) {
        if (parsed[0].toUpperCase() !== parsed[0]) {
          throw new AgentError(AgentErrorCode.ValidationError, 'Not camelcase');
        }
      }
    }
  })({});

  const result1 = await runAgent(simple, "John");
  expect(result1).toEqual(ok("John"));

  // parse error
  const result2 = await runAgent(simple, "");
  expect(result2).toEqual(err({
    code: AgentErrorCode.ParseError,
    message: 'No match',
  }));

  // validation error
  const result3 = await runAgent(simple, "john");
  expect(result3).toEqual(err({
    code: AgentErrorCode.ValidationError,
    message: 'Not camelcase',
  }));
});

Deno.test("runAgent: codegen", async () => {
  const codegen = defineAgent<{ maxRetries: number }>((init) => {
    return {
      async invoke(input, options) {
        return `Test
\`\`\`tsx
const x = 1;
\`\`\`
  `;
      },
      async parse(result) {
        const code = extractCodeBlock(result);
        if (code) {
          return code;
        } else {
          throw new AgentError(AgentErrorCode.ParseError, 'No code block found');
        }
      }
    }

    function extractCodeBlock(str: string): string {
      const result = str.match(/```tsx\n([\s\S]+?)\n```/)?.[1] ?? '';
      return result.trim();
      // return prettier.format(result, { parser: "typescript" });
    }
  });

  const writer = codegen({
    maxRetries: 3,
  });
  const result = await runAgent(writer, "const x = 1");
  expect(result).toEqual(ok("const x = 1;"));
});

agent.ts

import { Agent, AgentBuilder, AgentError, AgentState, AgentStateStep, Err, InvokeError, Ok, ParseError, Result, AgentErrorCode } from "./types.ts";

export function ok<T>(value: T): Ok<T> {
  return { ok: true, value };
}

export function err<T>(error: T): Err<T> {
  return { ok: false, error };
}

export async function initAgent<
  Options,
  Input,
  Invoked,
  Parsed,
>(
  agent: Agent<Options, Input, Invoked, Parsed>,
  input: Input,
  options: Options,
): Promise<AgentState<Options, Input, Invoked, Parsed>> {
  if (agent.init) {
    return await agent.init(input, options);
  }
  // return {} as any;
  const state: AgentState<Options, Input, Invoked, Parsed> = {
    errorCount: 0,
    step: AgentStateStep.Initialized,
    input,
    options,
  };
  return state;
}

export async function stepAgent<
  Options,
  Input,
  Invoked,
  Parsed,
>(
  agent: Agent<Options, Input, Invoked, Parsed>,
  state: AgentState<Options, Input, Invoked, Parsed>,
): Promise<AgentState<Options, Input, Invoked, Parsed>> {
  if (agent.step) {
    return await agent.step(state);
  }

  const newState = {
    ...state,
  };
  switch (state.step) {
    case AgentStateStep.Initialized: {
      // Initialized => Invoke
      const ctx = {
        invoked: state.invoked,
        parsed: state.parsed,
        validated: state.validated,
      };
      try {
        const invoked = await agent.invoke(state.input, { signal: undefined }, ctx);
        newState.step = AgentStateStep.Invoked;
        newState.invoked = ok(invoked);
        return newState;
      } catch (err) {
        newState.invoked = err({
          code: AgentErrorCode.InvokeError,
          message: err.message,
        });
        newState.errorCount++;
        return newState;
      }
    }
    case AgentStateStep.Invoked: {
      // Invoke => Parse
      const invoked = state.invoked!;
      // console.log({ invoked });
      // if (!invoked) {
      //   throw new Error('Invalid state');
      // }
      // if (!agent.parse) {
      //   newState.step = AgentStateStep.Done;
      //   newState.parsed = ok(invoked);
      //   return newState;
      // }
      try {
        if (!agent.parse) throw new Error('No parse');
        if (!invoked.ok) throw new Error('Not ok');
        const parsed = await agent.parse!(invoked!.value);
        newState.step = AgentStateStep.Parsed;
        newState.parsed = ok(parsed);
        return newState;
      } catch (error) {
        if (error instanceof AgentError && error.code === AgentErrorCode.ParseError) {
          newState.parsed = err({
            code: AgentErrorCode.ParseError,
            message: error.message,
          });
          newState.errorCount++;
          newState.step = error.rollbackTo ?? AgentStateStep.Initialized;
          return newState;
        }
        throw error;
      }
    }

    case AgentStateStep.Parsed: {
      // Parse => Validate
      const parsed = state.parsed!;
      if (!parsed!.ok) {
        throw new Error('Invalid state in Parsed');
      }
      if (!agent.validate) {
        newState.step = AgentStateStep.Done;
        newState.validated = ok(undefined);
        return newState;
      }
      try {
        await agent.validate(parsed.value);
        newState.step = AgentStateStep.Done;
        newState.validated = ok(undefined);
        return newState;
      } catch (error) {
        if (error instanceof AgentError) {
          newState.validated = err({
            code: AgentErrorCode.ValidationError,
            message: error.message,
          });
          newState.errorCount++;
          newState.step = error.rollbackTo ?? AgentStateStep.Initialized;
          return newState;
        }
        throw error;
      }
    }
  }
  return state;
}

type RunOptions = {
  maxRetries?: number;
}

export async function runAgent<
  Options,
  Input,
  Invoked,
  Parsed,
>(
  agent: Agent<Options, Input, Invoked, Parsed>,
  input: Input,
  options: RunOptions = {},
): Promise<Result<Parsed, AgentError>> {
  let state = await initAgent(agent, input, {} as Options);
  const maxRetries = options.maxRetries ?? 3;
  while (state.step !== AgentStateStep.Done) {
    state = await stepAgent(agent, state);
    // TODO: fix
    if (state.errorCount > maxRetries) {
      // console.log({ state });
      // throw new Error('Too many retries');
      if (state.validated && state.validated?.ok === false) return state.validated! as any;
      if (state.parsed && state.parsed?.ok === false) return state.parsed! as any;
      if (state.invoked && state.invoked?.ok === false) return state.invoked! as any;
      break;
    }
    if (state.step === AgentStateStep.Done && state.parsed) {
      return state.parsed as Result<Parsed, AgentError>;
    }
    // else {
    //   console.log({ state });
    //   throw new Error('Invalid state in runAgent');
    // }
  }
  throw new Error('Invalid state');
}

export function defineAgent<
  Options = {},
  Input = string,
  Invoked = Input,
  Parsed = Invoked,
>(agentFn: (options: Options) => Agent<Options, Input, Invoked, Parsed>): AgentBuilder<Options, Input, Invoked, Parsed> {
  return (options: Options) => agentFn(options);
}

export function chain_<
  // FirstInput,
  // FirstInvoked,
  // FirstParsed,
  // LastInvoked,
  // LastParsed,
  A extends Agent<any, any, any, any>,
  B extends Agent<any, any, any, any>,
  FirstInput = A extends Agent<any, infer FirstInput, any, any> ? FirstInput : never,
// LastParsed = B extends Agent<any, any, any, infer LastParsed> ? LastParsed : never,
>(
  a: A,
  b: B,
  options: RunOptions = {},
) {
  return async (input: FirstInput) => {
    const result = await runAgent(a, input);
    if (!result.ok) return result;
    return runAgent(b, result.value, options);
  }
}


// export function chain<
//   FirstInput,
//   FirstParsed,
//   SecondParsed,
// >(
//   a: Agent<any, FirstInput, any, FirstParsed>,
//   b: Agent<any, FirstParsed, any, SecondParsed>,
//   options: RunOptions = {},
// ): (input: FirstInput) => Promise<Result<SecondParsed, AgentError>> {
//   return async (input: FirstInput) => {
//     const result = await runAgent(a, input);
//     if (!result.ok) return result;
//     return await runAgent(b, result.value, options);
//   }
// }

export function chain<
  A extends Agent<any, any, any, any>,
  B extends Agent<any, A extends Agent<any, any, any, infer FirstParsed> ? FirstParsed : never, any, any>,
>(
  a: A,
  b: B,
  options: RunOptions = {},
): (input: A extends Agent<any, infer FirstInput, any, any> ? FirstInput : never) => Promise<
  Result<B extends Agent<any, any, any, infer SecondParsed> ? SecondParsed : never, AgentError>
> {
  return async (input) => {
    const result = await runAgent(a, input);
    if (!result.ok) return result;
    return await runAgent(b, result.value, options);
  }
}

types.ts

export type Promiseable<T> = T | Promise<T>;
export type Result<T, E = string> = Ok<T> | Err<E>;
export type Ok<T = {}> = { ok: true; value: T };
export type Err<E> = { ok: false; error: E };


export const enum AgentErrorCode {
  InvokeError = 'InvokeError',
  ParseError = 'ParseError',
  ValidationError = 'ValidationError',
}


export class AgentError extends Error {
  constructor(public code: AgentErrorCode, message: string, public rollbackTo?: AgentStateStep) {
    super(message);
  }
}


export const enum RuntimeErrorCode {
  NotExists = 'NotExists',
}

export class RuntimeError extends Error {
  constructor(public code: RuntimeErrorCode, message: string, public rollbackTo?: AgentStateStep) {
    super(message);
  }
}

export const enum AgentStateStep {
  Initialized = 'Initialized',
  Invoked = 'Invoked',
  Validated = 'Validated',
  Parsed = 'Parsed',
  Done = 'Done',
}

export type ValidationError = {
  code: AgentErrorCode.ValidationError;
  message: string;
  rollbackTo?: AgentStateStep;
};

export type ParseError = {
  code: AgentErrorCode.ParseError;
  message: string;
};

export type InvokeError = {
  code: AgentErrorCode.InvokeError;
  message: string;
}

// export type AgentError = ValidationError | ParseError | InvokeError;
export type AgentInvokeOptions<Options, E> = {
  signal?: AbortSignal;
  override?: Partial<Options>;
};

export type AgentState<
  Options,
  Input,
  Invoked,
  Parsed,
> = {
  errorCount: number;
  input: Input;
  step: AgentStateStep;
  invoked?: Result<Invoked, InvokeError>;
  parsed?: Result<Parsed, ParseError>;
  validated?: Result<void, ValidationError>
  options: Options;
};

export type Agent<
  Options,
  Input,
  Invoked,
  Parsed,
> = {
  description?: string;
  init?: (input: Input, options: Options) => Promiseable<
    AgentState<Options, Input, Invoked, Parsed>
  >;
  step?: (state: AgentState<Options, Input, Invoked, Parsed>) => Promiseable<AgentState<Options, Input, Invoked, Parsed>>;
  invoke(input: Input, options: AgentInvokeOptions<Options, any>, context?: {
    invoked?: Result<Invoked, InvokeError>;
    parsed?: Result<Parsed, ParseError>;
    validated?: Result<void, ValidationError>
  }): Promiseable<Invoked>;
  parse?(invoked: Invoked): Promiseable<Parsed>;
  validate?(parsed: Parsed): Promiseable<void>;
};

export type AgentBuilder<
  Options,
  Input,
  Invoked,
  Parsed,
> = (init: Options) => Agent<Options, Input, Invoked, Parsed>;