Introduce Plugin Architecture to ax/llm for Enhanced Agent Flow Tracing

[TYRONEMICHAEL]

Introduce Plugin Architecture to ax/llm with Strongly Typed Contexts

Overview

This proposes adding a simple, strongly-typed plugin architecture to ax/llm, allowing developers to extend functionality at specific points in the code without modifying the core library.

Motivation

Currently, features like debugging and telemetry are tightly coupled with the core library. A plugin system with strongly typed contexts would:

1. Provide flexibility for custom behavior at specific execution points.
2. Ensure type safety and improve developer experience when creating plugins.
3. Decouple auxiliary features (e.g., logging, telemetry) from the core library.
4. Enable easier extension and customization of the library's behavior.

'm currently working on a system to compose agents using react-flow, and I need a way to hook into specific events in the ax/llm execution flow in order to "react" when composing agents.

Proposal

Introduce a lightweight plugin system that allows hooking into key points of the ax/llm execution flow, with strongly typed contexts for each execution point.

Plugin Interface with Typed Contexts

type PluginContext =
  | { point: 'beforeChatExecution'; data: BeforeChatData }
  | { point: 'afterChatExecution'; data: AfterChatData }
  // Add other execution points as needed

interface Plugin {
  execute(context: PluginContext): Promise<void>;
}

interface BeforeChatData {
  req: AxChatRequest;
  options?: AxAIPromptConfig & AxAIServiceActionOptions;
  model: string;
  modelConfig: AxModelConfig;
}

interface AfterChatData extends BeforeChatData {
  response: AxChatResponse | ReadableStream<AxChatResponse>;
}

Plugin Manager

const pluginManager = {
  plugins: [] as Plugin[],
  async register(plugin: Plugin) {}
  async executePlugins(context: PluginContext) {}
};

Example: Refactoring the chat method

Here's a simplified example of how the current chat method could be refactored to use the strongly typed plugin system:

async function chat(req: AxChatRequest, options?: AxAIPromptConfig & AxAIServiceActionOptions) {
  const model = req.model ?? this.models.model;
  const mc = this.getModelConfig();

  await pluginManager.executePlugins({
    point: 'beforeChatExecution',
    data: { req, options, model, modelConfig: mc }
  });

  const res = await this._chat(model, req, options);

  await pluginManager.executePlugins({
    point: 'afterChatExecution',
    data: { req, options, model, modelConfig: mc, response: res }
  });

  return res;
}

This approach would replace the current tracer implementation, allowing for more flexible and type-safe logging and telemetry solutions.

Benefits

1. Type Safety: Strongly typed contexts ensure correct usage and improve developer experience.
2. Flexibility: Easy injection of custom behavior at specific execution points.
3. Modularity: Separation of auxiliary features from core library code.
4. Extensibility: Addition of new functionality without core code modification.

Next Steps

If there's interest in this concept, we could:

1. Finalize the list of execution points and their corresponding typed contexts.
2. Implement the basic plugin system with type safety.
3. Create initial plugins for logging and telemetry using the typed contexts.
4. Update documentation for plugin usage and creation, emphasizing the type safety benefits.

I welcome your thoughts on this proposal. This architecture could significantly enhance ax/llm's flexibility, extensibility, with type safety.

[dosco]

sorry was a bit busy. yes i'm open to a plugin system, i'm thinking of it as hooks that can be used for whatever. not sure as yet if telemetry should be a hook (maybe it should) technically we use an open standard so it's kinda a plugin.