Overview

The LLM package in Agenite provides utilities for working with language models, including message formatting, content types, and provider interfaces.

Message types

Base message

interface BaseMessage {
  role: 'user' | 'assistant' | 'system';
  content: ContentBlock[];
}

// Helper to create a simple text message
export const userTextMessage = (message: string): BaseMessage => {
  return {
    role: 'user',
    content: [
      {
        type: 'text',
        text: message,
      },
    ],
  };
};

Content blocks

Messages can contain different types of content:

type ContentBlock =
  | TextBlock
  | ImageBlock
  | ToolUseBlock
  | ToolResultBlock
  | DocumentBlock
  | ThinkingBlock
  | RedactedThinkingBlock;

interface TextBlock {
  type: 'text';
  text: string;
  [key: string]: unknown;
}

interface ImageBlock {
  type: 'image';
  source: ImageSource;
  [key: string]: unknown;
}

interface ToolUseBlock {
  type: 'toolUse';
  id: string;
  name: string;
  input: unknown;
  [key: string]: unknown;
}

interface ToolResultBlock {
  type: 'toolResult';
  toolUseId: string;
  toolName: string;
  content?: string | Array<ToolResponseBlock>;
  isError: boolean;
  [key: string]: unknown;
}

Provider interface

The core LLM Provider interface:

interface LLMProvider {
  name: string;
  version?: string;

  /**
   * Simple text generation with full response
   */
  generate(
    input: string | BaseMessage[],
    options?: Partial<GenerateOptions>
  ): Promise<GenerateResponse>;

  /**
   * Simple streaming with partial returns
   */
  stream(
    input: string | BaseMessage[],
    options?: Partial<GenerateOptions>
  ): AsyncGenerator<PartialReturn, GenerateResponse, unknown>;

  /**
   * Low-level generation API with full control
   */
  iterate(
    input: string | BaseMessage[],
    options: IterateGenerateOptions
  ): AsyncGenerator<PartialReturn, GenerateResponse, unknown>;
}

Generate options

interface GenerateOptions {
  tools?: ToolDefinition[];
  temperature?: number;
  maxTokens?: number;
  stopSequences?: string[];
  systemPrompt?: string;
}

interface IterateGenerateOptions extends GenerateOptions {
  stream: boolean;
}

Response types

interface GenerateResponse {
  content: Array<ContentBlock>;
  tokens: TokenUsage[];
  duration: number;
  stopReason?: StopReason;
}

interface TokenUsage {
  inputTokens: number;
  outputTokens: number;
  modelId: string;
}

type StopReason = 'toolUse' | 'maxTokens' | 'stopSequence' | 'endTurn';

Base provider implementation

The BaseLLMProvider class provides a base implementation of the LLM provider interface:

abstract class BaseLLMProvider implements LLMProvider {
  abstract name: string;
  abstract version?: string;

  abstract generate(
    input: string | BaseMessage[],
    options?: Partial<GenerateOptions>
  ): Promise<GenerateResponse>;

  abstract stream(
    input: string | BaseMessage[],
    options?: Partial<GenerateOptions>
  ): AsyncGenerator<PartialReturn, GenerateResponse, unknown>;

  async *iterate(
    input: string | BaseMessage[],
    options: IterateGenerateOptions
  ): AsyncGenerator<PartialReturn, GenerateResponse, unknown> {
    return yield* iterateFromMethods(this, input, options);
  }
}

Utility functions

String to message conversion

function convertStringToMessages(
  message: string | BaseMessage[]
): BaseMessage[] {
  if (typeof message === 'string') {
    return [
      {
        role: 'user',
        content: [{ type: 'text', text: message }],
      },
    ];
  }
  return message;
}

Tool definition

interface ToolDefinition {
  name: string;
  description: string;
  inputSchema: JSONSchema;
}

interface JSONSchema {
  type: string;
  properties?: Record<string, unknown>;
  required?: string[];
  [key: string]: unknown;
}

Example usage

Simple generation

import { BedrockProvider } from '@agenite/bedrock';

const provider = new BedrockProvider({
  model: 'anthropic.claude-3-5-sonnet-20240620-v1:0',
});

// Simple text input
const response = await provider.generate('Hello, how are you?');
console.log(response.content);

// Structured message input
const response2 = await provider.generate([
  {
    role: 'user',
    content: [{ type: 'text', text: 'Tell me about AI' }],
  },
]);

Streaming responses

const generator = provider.stream('Explain quantum computing');

for await (const chunk of generator) {
  if (chunk.type === 'text') {
    process.stdout.write(chunk.text);
  }
}

Tool usage

const toolDefinition = {
  name: 'calculator',
  description: 'Perform basic math operations',
  inputSchema: {
    type: 'object',
    properties: {
      operation: {
        type: 'string',
        enum: ['add', 'subtract', 'multiply', 'divide'],
      },
      a: { type: 'number' },
      b: { type: 'number' },
    },
    required: ['operation', 'a', 'b'],
  },
};

const response = await provider.generate(
  'What is 25 × 16?',
  {
    tools: [toolDefinition],
    temperature: 0,
  }
);

// Check if the response includes a tool use
const toolUse = response.content.find(
  (block) => block.type === 'toolUse'
);

if (toolUse) {
  console.log('Tool called:', toolUse.name);
  console.log('Input:', toolUse.input);
}

Creating a custom provider

import { BaseLLMProvider, GenerateOptions, GenerateResponse } from '@agenite/llm';

class CustomProvider extends BaseLLMProvider {
  name = 'custom-provider';
  version = '1.0.0';
  
  constructor(private apiKey: string) {
    super();
  }
  
  async generate(
    input: string | BaseMessage[],
    options?: Partial<GenerateOptions>
  ): Promise<GenerateResponse> {
    // Convert input to standard messages format
    const messages = convertStringToMessages(input);
    
    // Make API call to your custom LLM service
    // ...
    
    // Return standardized response
    return {
      content: [{ type: 'text', text: 'Response from custom LLM' }],
      tokens: [{ modelId: 'custom-model', inputTokens: 10, outputTokens: 20 }],
      duration: 500,
    };
  }
  
  async *stream(
    input: string | BaseMessage[],
    options?: Partial<GenerateOptions>
  ): AsyncGenerator<PartialReturn, GenerateResponse, unknown> {
    // Stream implementation
    // ...
    
    yield { type: 'text', text: 'Partial response...' };
    
    return {
      content: [{ type: 'text', text: 'Complete response' }],
      tokens: [{ modelId: 'custom-model', inputTokens: 10, outputTokens: 20 }],
      duration: 500,
    };
  }
}

Best practices

  1. Message formatting: Use the standard message format for consistent behavior
  2. Error handling: Implement robust error handling in custom providers
  3. Streaming: Support streaming for better user experience with long responses
  4. Token tracking: Track token usage for monitoring and rate limiting
  5. Content typing: Use the appropriate content block types for rich responses

Next steps