@just-every/ensemble
Version:
LLM provider abstraction layer with unified streaming interface
238 lines • 10 kB
JavaScript
;
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
var desc = Object.getOwnPropertyDescriptor(m, k);
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
desc = { enumerable: true, get: function() { return m[k]; } };
}
Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
o["default"] = v;
});
var __importStar = (this && this.__importStar) || (function () {
var ownKeys = function(o) {
ownKeys = Object.getOwnPropertyNames || function (o) {
var ar = [];
for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
return ar;
};
return ownKeys(o);
};
return function (mod) {
if (mod && mod.__esModule) return mod;
var result = {};
if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
__setModuleDefault(result, mod);
return result;
};
})();
Object.defineProperty(exports, "__esModule", { value: true });
exports.testProvider = exports.TestProvider = exports.testProviderConfig = void 0;
exports.resetTestProviderConfig = resetTestProviderConfig;
const base_provider_js_1 = require("./base_provider.cjs");
const uuid_1 = require("uuid");
const index_js_1 = require("../index.cjs");
const event_controller_js_1 = require("../utils/event_controller.cjs");
exports.testProviderConfig = {
streamingDelay: 50,
shouldError: false,
errorMessage: 'Simulated error from test provider',
simulateRateLimit: false,
fixedResponse: undefined,
fixedThinking: undefined,
simulateToolCall: false,
toolName: 'web_search',
toolArguments: { query: 'test query' },
tokenUsage: {
inputTokens: 100,
outputTokens: 200,
},
chunkSize: 5,
};
function resetTestProviderConfig() {
exports.testProviderConfig.streamingDelay = 50;
exports.testProviderConfig.shouldError = false;
exports.testProviderConfig.errorMessage = 'Simulated error from test provider';
exports.testProviderConfig.simulateRateLimit = false;
exports.testProviderConfig.fixedResponse = undefined;
exports.testProviderConfig.fixedThinking = undefined;
exports.testProviderConfig.simulateToolCall = false;
exports.testProviderConfig.toolName = 'web_search';
exports.testProviderConfig.toolArguments = { query: 'test query' };
exports.testProviderConfig.tokenUsage = {
inputTokens: 100,
outputTokens: 200,
};
exports.testProviderConfig.chunkSize = 5;
}
const sleep = (ms) => new Promise(resolve => setTimeout(resolve, ms));
class TestProvider extends base_provider_js_1.BaseModelProvider {
config;
constructor(config = exports.testProviderConfig) {
super('test');
this.config = config;
}
async *createResponseStream(messages, model, agent) {
console.log(`[TestProvider] Creating response stream for model: ${model}`);
const lastUserMessage = messages.filter(m => 'role' in m && m.role === 'user').pop();
const userMessageContent = lastUserMessage && 'content' in lastUserMessage
? typeof lastUserMessage.content === 'string'
? lastUserMessage.content
: JSON.stringify(lastUserMessage.content)
: '';
const inputTokenCount = this.config.tokenUsage?.inputTokens || Math.max(50, Math.ceil(userMessageContent.length / 4));
let response;
if (this.config.simulateRateLimit) {
const rateLimitError = '429 Too Many Requests: The server is currently processing too many requests. Please try again later.';
yield {
type: 'error',
error: rateLimitError,
};
return;
}
if (this.config.shouldError) {
yield {
type: 'error',
error: this.config.errorMessage || 'Simulated error from test provider',
};
return;
}
if (this.config.fixedResponse) {
response = this.config.fixedResponse;
}
else {
response = this.generateResponse(userMessageContent);
}
const messageId = (0, uuid_1.v4)();
yield {
type: 'message_start',
message_id: messageId,
content: '',
};
if (this.config.fixedThinking) {
yield {
type: 'message_delta',
message_id: messageId,
content: '',
thinking_content: this.config.fixedThinking,
thinking_signature: '(Simulated thinking)',
};
await sleep(this.config.streamingDelay || 50);
}
if (this.config.simulateToolCall && agent) {
const { getToolsFromAgent } = await Promise.resolve().then(() => __importStar(require("../utils/agent.cjs")));
const currentTools = getToolsFromAgent(agent);
if (currentTools) {
const toolArray = await currentTools;
if (toolArray.length > 0) {
const availableTool = toolArray.find(tool => this.config.toolName ? tool.definition.function.name === this.config.toolName : true);
if (availableTool) {
const toolCall = {
id: (0, uuid_1.v4)(),
type: 'function',
function: {
name: availableTool.definition.function.name,
arguments: JSON.stringify(this.config.toolArguments || {
query: userMessageContent.slice(0, 50),
}),
},
};
yield {
type: 'tool_start',
tool_call: toolCall,
};
await sleep(this.config.streamingDelay || 50);
response = `I've used the ${toolCall.function.name} tool to help answer your question.\n\n${response}`;
}
}
}
}
const chunkSize = this.config.chunkSize || 5;
let position = 0;
while (position < response.length) {
const chunk = response.slice(position, position + chunkSize);
position += chunkSize;
yield {
type: 'message_delta',
message_id: messageId,
content: chunk,
order: position / chunkSize,
};
await sleep(this.config.streamingDelay || 50);
}
yield {
type: 'message_complete',
message_id: messageId,
content: response,
};
const outputTokenCount = this.config.tokenUsage?.outputTokens || Math.ceil(response.length / 4);
const calculatedUsage = index_js_1.costTracker.addUsage({
model,
input_tokens: inputTokenCount,
output_tokens: outputTokenCount,
});
if (!(0, event_controller_js_1.hasEventHandler)()) {
yield {
type: 'cost_update',
usage: {
...calculatedUsage,
total_tokens: inputTokenCount + outputTokenCount,
},
};
}
}
generateResponse(input) {
const lowercaseInput = input.toLowerCase();
if (lowercaseInput.includes('hello') || lowercaseInput.includes('hi')) {
return "Hello! I'm a test AI model. How can I help you today?";
}
else if (lowercaseInput.includes('help')) {
return "I'm here to help! What do you need assistance with?";
}
else if (lowercaseInput.includes('error') || lowercaseInput.includes('problem')) {
return "I understand you're experiencing an issue. Let me help troubleshoot the problem.";
}
else if (lowercaseInput.includes('json') || lowercaseInput.includes('person')) {
return '{"name": "John Doe", "age": 30}';
}
else if (lowercaseInput.includes('test')) {
return 'This is a test response. The test provider is working correctly!';
}
else if (lowercaseInput.includes('weather')) {
return 'The weather is sunny and 72°F.';
}
else if (lowercaseInput.includes('?')) {
return "That's an interesting question. As a test model, I'm designed to provide simulated responses for testing purposes.";
}
else {
return `I've received your message: "${input.slice(0, 50)}${input.length > 50 ? '...' : ''}". This is a simulated response from the test provider.`;
}
}
async createEmbedding(input, model, opts) {
const generateVector = (text) => {
const dimension = opts?.dimension || 384;
const vector = new Array(dimension);
for (let i = 0; i < dimension; i++) {
const charCode = text.charCodeAt(i % text.length) || 0;
const value = Math.sin(charCode * (i + 1) * 0.01) * 0.5 + 0.5;
vector[i] = value;
}
return vector;
};
if (Array.isArray(input)) {
return input.map(text => generateVector(text));
}
else {
return generateVector(input);
}
}
}
exports.TestProvider = TestProvider;
exports.testProvider = new TestProvider();
//# sourceMappingURL=test_provider.js.map