@axarai/axar
Version:
TypeScript-based agent framework for building agentic applications powered by LLMs
311 lines (310 loc) • 12.9 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.input = exports.output = void 0;
exports.model = model;
exports.systemPrompt = systemPrompt;
exports.tool = tool;
require("reflect-metadata");
const zod_1 = require("zod");
const meta_keys_1 = require("./meta-keys");
const schema_1 = require("../schema");
/**
* `model` decorator to associate a model identifier and configuration with an agent.
*
* @param modelIdentifier - The model identifier string. List of model identifiers available at https://axar-ai.gitbook.io/axar/basics/model
* @param config - Optional configuration for the model
* @param config.maxTokens - Maximum number of tokens to generate
* @param config.temperature - Sampling temperature between 0 and 1. Use either temperature or topP, not both.
* @param config.topP - Nucleus sampling threshold. Sample from tokens whose cumulative probability exceeds topP. Use either temperature or topP, not both.
* @param config.topK - Only sample from the top K options for each subsequent token. Recommended for advanced use cases only.
* @param config.presencePenalty - Penalize tokens based on their presence in the prompt and generated text. Value between -2.0 and 2.0.
* @param config.frequencyPenalty - Penalize tokens based on their frequency in the generated text. Value between -2.0 and 2.0.
* @param config.maxRetries - Maximum number of retries for failed requests (defaults to 2 in SDK)
* @param config.maxSteps - Maximum number of steps for tool calling (defaults to 3)
* @param config.toolChoice - Tool choice mode - 'auto' or 'none'
* @returns A class decorator function.
*
* @example
* ```typescript
* // Basic usage
* @model('openai:gpt-4-mini')
* class MyAgent extends Agent<string, string> {}
*
* // With configuration
* @model('openai:gpt-4-mini', {
* maxTokens: 100,
* temperature: 0.7,
* maxRetries: 3,
* maxSteps: 5,
* toolChoice: 'auto'
* })
* class MyConfiguredAgent extends Agent<string, string> {}
*
* // With advanced sampling parameters
* @model('openai:gpt-4-mini', {
* topP: 0.9, // nucleus sampling
* topK: 50, // top-k sampling
* presencePenalty: 0.6, // reduce repetition
* frequencyPenalty: 0.5 // reduce common tokens
* })
* class CreativeAgent extends Agent<string, string> {}
* ```
*/
function model(modelIdentifier, config) {
return function (target) {
Reflect.defineMetadata(meta_keys_1.META_KEYS.MODEL, modelIdentifier, target);
if (config) {
Reflect.defineMetadata(meta_keys_1.META_KEYS.MODEL_CONFIG, config, target);
}
return target;
};
}
/**
* Creates a schema from the provided type specification
*
* @param type - The type specification (ZodSchema, class decorated with @schema, or primitive constructor)
* @param decoratorName - The name of the decorator that we're creating the schema for
* @returns The created schema
*/
function createSchema(type, decoratorName) {
if (type instanceof zod_1.ZodSchema) {
return type;
}
const primitiveSchemas = {
[String.name]: zod_1.z.string(),
[Number.name]: zod_1.z.number(),
[Boolean.name]: zod_1.z.boolean(),
};
const primitiveSchema = primitiveSchemas[type.name];
if (primitiveSchema) {
return primitiveSchema;
}
if ((0, schema_1.hasSchemaDef)(type)) {
return (0, schema_1.getSchemaDef)(type);
}
const typeName = typeof type === 'function' && type.name ? type.name : String(type);
throw new Error(`${decoratorName} error: Could not create a schema for "${typeName}". ` +
`Type must be a Zod schema, a class decorated with @schema, or a primitive constructor (String, Number, Boolean).`);
}
/**
*
* Creates a decorator for input/output schema definition
*
* @param metaKey - Input/Output symbol
* @param decoratorName - '@input' or '@output'
* @returns a decorator
*/
function createSchemaDecorator(metaKey, decoratorName) {
return function (type) {
return function (target) {
const schema = createSchema(type, decoratorName);
Reflect.defineMetadata(metaKey, schema, target);
return target;
};
};
}
/**
* Specifies the output schema for a class.
*
* @param type - The output type specification, which can be:
* - A Zod schema
* - A class decorated with @schema
* - A primitive type (String, Number, Boolean)
*
* @example
* ```typescript
* // Using primitive
* @output(String)
* class StringAgent extends Agent<string, string> {}
*
* // Using schema-decorated class
* @output(UserProfile)
* class UserAgent extends Agent<string, UserProfile> {}
*
* // Using Zod schema directly
* @output(z.boolean())
* class BooleanAgent extends Agent<string, boolean> {}
* ```
*/
exports.output = createSchemaDecorator(meta_keys_1.META_KEYS.OUTPUT, '@output');
/**
* Specifies the input schema for a class.
*
* @param type - The input type specification, which can be:
* - A Zod schema
* - A class decorated with @schema
* - A primitive type (String, Number, Boolean)
*
* @example
* ```typescript
* // Using primitive
* @input(Boolean)
* class BooleanAgent extends Agent<boolean, string> {}
*
* // Using schema-decorated class
* @input(UserProfile)
* class UserAgent extends Agent<UserProfile, string> {}
*
* // Using Zod schema directly
* @input(z.boolean())
* class BooleanAgent extends Agent<boolean, string> {}
* ```
*/
exports.input = createSchemaDecorator(meta_keys_1.META_KEYS.INPUT, '@input');
// Implementation
function systemPrompt(prompt) {
// Class Decorator
if (typeof prompt === 'string') {
return systemPromptClass(prompt);
}
else {
// Method Decorator
return systemPromptMethod();
}
}
// Internal helper for class-level system prompts
function systemPromptClass(prompt) {
return function (target) {
const systemPrompts = Reflect.getMetadata(meta_keys_1.META_KEYS.SYSTEM_PROMPTS, target) || [];
// Add class prompt to the beginning
systemPrompts.unshift(async () => prompt);
Reflect.defineMetadata(meta_keys_1.META_KEYS.SYSTEM_PROMPTS, systemPrompts, target);
return target;
};
}
// Internal helper for method-level system prompts
function systemPromptMethod() {
return function (target, propertyKey, descriptor) {
if (typeof descriptor.value !== 'function') {
throw new Error(`@systemPrompt can only be applied to methods, not to property '${String(propertyKey)}'.`);
}
// Retrieve existing system prompts or initialize
const systemPrompts = Reflect.getMetadata(meta_keys_1.META_KEYS.SYSTEM_PROMPTS, target.constructor) || [];
systemPrompts.push(async function () {
const result = await descriptor.value.apply(this); // Use the actual instance's `this`
if (typeof result !== 'string') {
throw new Error(`Method '${String(propertyKey)}' decorated with @systemPrompt must return a string.`);
}
return result;
});
Reflect.defineMetadata(meta_keys_1.META_KEYS.SYSTEM_PROMPTS, systemPrompts, target.constructor);
return descriptor;
};
}
/**
* Maps primitive constructors to their corresponding Zod schemas.
*/
const primitiveToZodSchema = {
[String.name]: zod_1.z.string(),
[Number.name]: zod_1.z.number(),
[Boolean.name]: zod_1.z.boolean(),
};
/**
* `@tool` decorator to mark a method as a tool with a description and schema.
* Supports explicit ZodSchema, class-based schema derivation, and primitive types.
*
* Usage with ZodSchema:
* @tool("Description", z.object({ ... }))
*
* Usage with class-based schema:
* @tool("Description")
* async method(params: ClassBasedParams): Promise<ReturnType> { ... }
*
* Usage with primitive types (String, Number, Boolean):
* @tool("Description")
* async method(query: string): Promise<string> { ... }
*
* When using primitive types, the parameter is automatically wrapped in an object
* with a 'value' property for the LLM, and unwrapped when calling your method.
*
* @param description - Description of the tool's functionality.
* @param schemaOrClass - Optional Zod schema or class constructor.
* @returns A method decorator function.
*/
function tool(description, schemaOrClass) {
return function (target, propertyKey, descriptor) {
if (typeof descriptor.value !== 'function') {
throw new Error(`@tool can only be applied to methods, not to property '${String(propertyKey)}'.`);
}
let schema;
let isPrimitiveParam = false;
if (schemaOrClass) {
if (schemaOrClass instanceof zod_1.z.ZodSchema) {
// Explicit Zod schema provided
schema = schemaOrClass;
}
else if ((0, schema_1.hasSchemaDef)(schemaOrClass)) {
schema = (0, schema_1.getSchemaDef)(schemaOrClass);
}
else {
throw new Error(`${schemaOrClass.name} must be either a Zod schema or a class decorated with @schema`);
}
}
else {
// No schema provided, derive via reflection
const paramTypes = Reflect.getMetadata('design:paramtypes', target, propertyKey);
if (!paramTypes || paramTypes.length === 0) {
// Empty parameter list, so we use an empty schema
schema = zod_1.z.object({});
}
else {
const paramType = paramTypes[0];
// Validate parameter count
if (paramTypes.length > 1) {
throw new Error(`@tool ${String(propertyKey)}: Expected a single parameter, but found ${paramTypes.length}.`);
}
// Check for supported primitive types (String, Number, Boolean)
const primitiveSchema = primitiveToZodSchema[paramType.name];
if (primitiveSchema) {
// Wrap primitive in an object schema with a 'value' property
// This is required because LLM tool parameters must be objects
schema = zod_1.z.object({ value: primitiveSchema });
isPrimitiveParam = true;
}
else {
// Exclude unsupported primitive types
const unsupportedPrimitives = [Symbol, BigInt, Function];
if (unsupportedPrimitives.includes(paramType)) {
throw new Error(`@tool ${String(propertyKey)}: The parameter type (${paramType.name}) is not supported. Use String, Number, Boolean, or an object type.`);
}
// Ensure paramType is constructable and produces an object
if (typeof paramType !== 'function' || !paramType.prototype) {
throw new Error(`@tool ${String(propertyKey)}: The parameter type (${paramType.name}) is not a valid class or constructor.`);
}
// Check if paramType is a class decorated with @schema
if (!(0, schema_1.hasSchemaDef)(paramType)) {
throw new Error(`@tool decorator on ${String(propertyKey)} requires an explicit Zod schema or a parameter class decorated with @schema.`);
}
// Convert the parameter class to Zod schema
schema = (0, schema_1.getSchemaDef)(paramType);
}
}
}
// Retrieve existing tools metadata or initialize an empty array
const tools = Reflect.getMetadata(meta_keys_1.META_KEYS.TOOLS, target.constructor) || [];
// Add the new tool to the metadata
tools.push({
name: String(propertyKey),
description,
method: String(propertyKey),
parameters: schema,
isPrimitiveParam,
});
// Define the updated tools metadata on the constructor
Reflect.defineMetadata(meta_keys_1.META_KEYS.TOOLS, tools, target.constructor);
// Wrap the original method to validate input and unwrap primitives
const originalMethod = descriptor.value;
descriptor.value = function (...args) {
let input = args[0];
if (input !== undefined) {
schema.parse(input); // This will throw if validation fails
// Unwrap primitive parameter from { value: ... } wrapper
if (isPrimitiveParam && typeof input === 'object' && 'value' in input) {
input = input.value;
}
}
return originalMethod.apply(this, [input, ...args.slice(1)]);
};
return descriptor;
};
}