pit-manager
Version:
Centralized prompt management system for Human Behavior AI agents
570 lines (569 loc) • 22.5 kB
JavaScript
"use strict";
/**
* Unified model interface with automatic tracking and chaining for TypeScript
*/
Object.defineProperty(exports, "__esModule", { value: true });
exports.model = exports.ModelResponse = void 0;
const fs = require("fs");
const path = require("path");
const crypto = require("crypto");
const registry_1 = require("../providers/registry");
const operations_1 = require("../versioning/operations");
const structured_output_1 = require("./structured-output");
const execution_tracker_1 = require("../storage/execution-tracker");
const prompts_1 = require("./prompts");
/**
* ModelResponse implementation that enables automatic chaining
*/
class ModelResponse {
constructor(params) {
this._isPitResponse = true;
this.content = params.content;
this.model = params.model;
this.promptHash = params.promptHash;
this.executionId = params.executionId;
this.metadata = params.metadata;
this.rawResponse = params.rawResponse;
}
/**
* String representation returns the content
*/
toString() {
if (typeof this.content === 'object') {
return JSON.stringify(this.content, null, 2);
}
return String(this.content);
}
/**
* Convert response to prompt string for chaining
*/
toPrompt() {
return this.toString();
}
/**
* Concatenate with another ModelResponse or string
*/
plus(other) {
return this.concatenate(other);
}
/**
* Internal concatenation logic
*/
concatenate(other) {
// Convert string to ModelResponse if needed
if (typeof other === 'string') {
other = ModelResponse.fromString(other);
}
// Check for duplicate concatenation
if (this.executionId === other.executionId) {
console.warn(`⚠️ Duplicate concatenation detected: ${this.executionId}`);
return this;
}
// Combine agent names intelligently
const selfTag = this.metadata.tag;
const otherTag = other.metadata.tag;
const combinedTag = this.combineAgentNames(selfTag, otherTag);
// Generate deterministic chain ID based on unique agents
const agentSequence = this.getAgentSequence().concat(other.getAgentSequence());
const newChainId = this.generateDeterministicChainId(agentSequence);
// Combine content
const combinedContent = this.toString() + '\n' + other.toString();
// Create new execution ID
const newExecutionId = crypto
.createHash('sha256')
.update(`${combinedContent}_${Date.now()}`)
.digest('hex');
return new ModelResponse({
content: combinedContent,
model: `${this.model}+${other.model}`,
promptHash: crypto.createHash('sha256').update(combinedContent).digest('hex'),
executionId: newExecutionId,
metadata: {
tag: combinedTag,
provider: `${this.metadata.provider}+${other.metadata.provider}`,
chainId: newChainId,
tokens: {
prompt: this.metadata.tokens.prompt + other.metadata.tokens.prompt,
completion: this.metadata.tokens.completion + other.metadata.tokens.completion,
total: this.metadata.tokens.total + other.metadata.tokens.total
},
latencyMs: Math.max(this.metadata.latencyMs, other.metadata.latencyMs),
structured: this.metadata.structured || other.metadata.structured,
concatenated: true,
originalResponses: [this.executionId, other.executionId],
agentSequence,
timestamp: new Date().toISOString()
}
});
}
/**
* Combine agent names intelligently, avoiding duplicates
*/
combineAgentNames(tag1, tag2) {
// Parse existing combined tags
const agents1 = tag1.includes('+') ? tag1.split('+') : [tag1];
const agents2 = tag2.includes('+') ? tag2.split('+') : [tag2];
// Get unique agents while preserving order
const seen = new Set();
const uniqueAgents = [];
for (const agent of [...agents1, ...agents2]) {
if (!seen.has(agent)) {
seen.add(agent);
uniqueAgents.push(agent);
}
}
return uniqueAgents.join('+');
}
/**
* Get the sequence of agents in this response
*/
getAgentSequence() {
const metadata = this.metadata;
if (metadata.agentSequence) {
return metadata.agentSequence;
}
const tag = this.metadata.tag;
return tag.includes('+') ? tag.split('+') : [tag];
}
/**
* Generate deterministic chain ID from agent sequence
*/
generateDeterministicChainId(agentSequence) {
// Get unique agents and sort for consistency
const uniqueAgents = [...new Set(agentSequence)].sort();
const signature = uniqueAgents.join('->');
// Generate deterministic hash
const hash = crypto.createHash('sha256').update(signature).digest('hex');
return `auto_${hash.substring(0, 12)}`;
}
/**
* Create ModelResponse from string
*/
static fromString(content, tag = 'string-context') {
const executionId = crypto
.createHash('sha256')
.update(`${content}_${Date.now()}`)
.digest('hex');
return new ModelResponse({
content,
model: 'string',
promptHash: crypto.createHash('sha256').update(content).digest('hex'),
executionId,
metadata: {
tag,
provider: 'string',
tokens: { prompt: 0, completion: 0, total: 0 },
latencyMs: 0,
structured: false,
fromString: true,
timestamp: new Date().toISOString()
}
});
}
}
exports.ModelResponse = ModelResponse;
/**
* Unified model interface with automatic tracking
*/
class Model {
// private activeChains: Map<string, string> = new Map();
constructor() {
this.repoPath = this.findPitDirectory();
if (this.repoPath) {
this.versioning = new operations_1.VersioningEngine(this.repoPath);
this.executionTracker = new execution_tracker_1.ExecutionTracker(this.repoPath);
}
else {
this.versioning = null;
this.executionTracker = null;
}
// Initialize provider registry
this.registry = new registry_1.ProviderRegistry();
}
/**
* Execute model with automatic tracking and optional structured output
*/
async complete(modelName, prompt, tag, jsonOutput, options = {}) {
// Detect chaining
let parentExecutionId = options?.parentExecutionId;
let chainId = options?.chainId;
let promptStr;
let multimodalContent;
if (process.env.PIT_DEBUG) {
console.log(`[DEBUG] model.complete input - chainId: ${chainId}, parentExecutionId: ${parentExecutionId}, tag: ${tag}`);
}
if (this.isModelResponse(prompt)) {
// This is a chained call
parentExecutionId = prompt.executionId;
chainId = prompt.metadata.chainId;
promptStr = prompt.toPrompt();
}
else if (typeof prompt === 'string') {
promptStr = prompt;
// Only use detected parent ID if this prompt came from prompts() function
// Check if the prompt contains the special marker that prompts() adds
if (!parentExecutionId && (0, prompts_1.getCurrentPromptFile)()) {
const detectedParentId = (0, prompts_1.getParentExecutionId)();
if (detectedParentId) {
parentExecutionId = detectedParentId;
}
}
}
else if (this.isMultimodalContent(prompt)) {
// Handle multimodal content
multimodalContent = prompt;
// Extract text for prompt hash calculation
promptStr = prompt.parts
.filter(part => part.type === 'text')
.map(part => part.text)
.join('\n');
}
else {
throw new Error('Invalid prompt type');
}
// Capture prompt file before clearing context
const promptFile = (0, prompts_1.getCurrentPromptFile)();
// Clear prompt context to prevent it from affecting subsequent calls
(0, prompts_1.clearPromptContext)();
// Detect provider from model name
const providerName = this.detectProvider(modelName);
// Get provider
const provider = this.registry.getProvider(providerName);
if (!provider) {
throw new Error(`Provider not found: ${providerName}`);
}
// Prepare structured output if needed
let structuredConfig = {};
if (jsonOutput) {
const handler = new structured_output_1.StructuredOutputHandler();
if (providerName === 'openai') {
structuredConfig = handler.prepareOpenAIStructured(jsonOutput);
}
else if (providerName === 'google') {
structuredConfig = handler.prepareGeminiStructured(jsonOutput);
}
else if (providerName === 'anthropic') {
structuredConfig = handler.prepareClaudeStructured(jsonOutput);
}
}
// Execute with tracking
const startTime = Date.now();
try {
// Create messages
const messages = [];
// Add system prompt if provided
if (options.systemPrompt) {
messages.push({ role: 'system', content: options.systemPrompt });
}
// Add user message
if (multimodalContent) {
// Convert multimodal content to message format
messages.push({
role: 'user',
content: {
type: 'multimodal',
parts: multimodalContent.parts.map(part => {
if (part.type === 'text') {
return {
type: 'text',
text: part.text
};
}
else if (part.type === 'image') {
return {
type: 'image',
image: {
data: part.data,
mimeType: part.mimeType
}
};
}
return part;
})
}
});
}
else {
messages.push({ role: 'user', content: promptStr });
}
// Handle structured output for different providers
let response;
if (providerName === 'openai' && jsonOutput && structuredConfig.useBeta) {
// Use OpenAI API for structured output
const OpenAI = require('openai');
const client = new OpenAI({ apiKey: provider.apiKey });
// Filter out PIT-specific options
const { parentExecutionId: _, systemPrompt: __, // Already handled via messages
chainId: ___, // PIT-specific tracking
...cleanOptions } = options;
response = await client.chat.completions.create({
model: modelName,
messages: messages.map(m => ({ role: m.role, content: m.content })),
response_format: structuredConfig.response_format,
temperature: options.temperature,
max_tokens: options.maxTokens,
...cleanOptions
});
}
else {
// Regular API call
// Filter out PIT-specific options from provider options
const { parentExecutionId: _, systemPrompt: __, // systemPrompt is handled differently by each provider
chainId: ___, // chainId is PIT-specific tracking
...cleanOptions } = options;
const providerOptions = {
...cleanOptions,
...structuredConfig
};
response = await provider.chatCompletion(modelName, messages, providerOptions);
}
// Parse structured output if needed
let content;
if (jsonOutput) {
const handler = new structured_output_1.StructuredOutputHandler();
if (providerName === 'openai') {
content = handler.parseOpenAIResponse(response, jsonOutput);
}
else if (providerName === 'google') {
content = handler.parseGeminiResponse(response, jsonOutput);
}
else if (providerName === 'anthropic') {
content = handler.parseClaudeResponse(response, jsonOutput);
}
else {
content = this.extractContent(response);
}
}
else {
// Regular text response
content = this.extractContent(response);
}
// Calculate metrics
const endTime = Date.now();
const latencyMs = endTime - startTime;
// Extract token usage
const tokens = this.extractTokens(response, providerName);
if (process.env.PIT_DEBUG) {
console.debug('Token extraction:', { provider: providerName, response: response.usage, tokens });
}
// Track execution
const promptHash = crypto.createHash('sha256').update(promptStr).digest('hex');
const executionId = crypto.createHash('sha256')
.update(`${promptHash}_${new Date().toISOString()}`)
.digest('hex');
// Create or update chain
if (!chainId && parentExecutionId) {
// Try to inherit chainId from parent execution
if (this.executionTracker) {
try {
const parentExecution = await this.executionTracker.getExecutionById(parentExecutionId);
if (parentExecution && parentExecution.chain_id) {
chainId = parentExecution.chain_id;
}
}
catch (error) {
// If we can't find parent, create new chain
if (process.env.PIT_DEBUG) {
console.debug('Could not inherit chainId from parent:', error);
}
}
}
// If still no chainId, create a new one
if (!chainId) {
chainId = crypto.createHash('sha256')
.update(`${tag}_${new Date().toISOString()}`)
.digest('hex');
}
}
// Track execution if tracker available
if (this.executionTracker) {
// Use the captured prompt file (already captured before clearing context)
const cost = this.calculateCost(tokens, modelName, providerName);
await this.executionTracker.trackExecution(executionId, modelName, promptHash, tag, tokens, latencyMs, cost, chainId || null, parentExecutionId || null, {
provider: providerName,
structured_output: jsonOutput !== undefined,
prompt_file: promptFile
});
// Auto-commit if not in active chain
if (!parentExecutionId && this.versioning) {
try {
await this.versioning.commit(`Model execution: ${tag} with ${modelName}`, 'PIT Auto-Tracker');
}
catch (commitError) {
// Ignore commit errors for executions - they're not critical
if (process.env.PIT_DEBUG) {
console.debug('Auto-commit skipped:', commitError.message);
}
}
}
}
// Create response object
return new ModelResponse({
content,
model: modelName,
promptHash,
executionId,
metadata: {
tag,
provider: providerName,
chainId,
tokens,
latencyMs,
structured: jsonOutput !== undefined
},
rawResponse: response
});
}
catch (error) {
throw new Error(`Model execution failed: ${error.message}`);
}
}
/**
* Check if value is a ModelResponse
*/
isModelResponse(value) {
return value && value._isPitResponse === true;
}
/**
* Check if value is MultimodalContent
*/
isMultimodalContent(value) {
return value && value.parts && Array.isArray(value.parts);
}
/**
* Find the .pit directory by searching up from current directory
*/
findPitDirectory() {
let current = process.cwd();
while (current !== path.dirname(current)) {
const pitDir = path.join(current, '.pit');
if (fs.existsSync(pitDir) && fs.statSync(pitDir).isDirectory()) {
return pitDir;
}
current = path.dirname(current);
}
// Check current directory one more time
const pitDir = path.join(process.cwd(), '.pit');
if (fs.existsSync(pitDir) && fs.statSync(pitDir).isDirectory()) {
return pitDir;
}
return null;
}
/**
* Detect provider from model name
*/
detectProvider(modelName) {
const modelLower = modelName.toLowerCase();
if (['gpt', 'o1', 'davinci', 'curie', 'babbage', 'ada'].some(x => modelLower.includes(x))) {
return 'openai';
}
else if (['gemini', 'palm', 'bison'].some(x => modelLower.includes(x))) {
return 'google';
}
else if (['claude', 'anthropic'].some(x => modelLower.includes(x))) {
return 'anthropic';
}
else {
// Default to OpenAI for unknown models
return 'openai';
}
}
/**
* Extract content from provider response
*/
extractContent(response) {
if (response.content) {
return response.content;
}
else if (response.choices && response.choices[0]?.message?.content) {
return response.choices[0].message.content;
}
else if (response.text) {
return response.text;
}
else {
return String(response);
}
}
/**
* Extract token usage from provider response
*/
extractTokens(response, provider) {
// The response from providers already has the tokens in the correct format
if (response.tokens) {
return {
prompt: response.tokens.input || 0,
completion: response.tokens.output || 0,
total: response.tokens.total || 0
};
}
// Fallback for raw responses
const tokens = { prompt: 0, completion: 0, total: 0 };
if (provider === 'openai') {
if (response.usage) {
tokens.prompt = response.usage.prompt_tokens || 0;
tokens.completion = response.usage.completion_tokens || 0;
tokens.total = response.usage.total_tokens || 0;
}
}
else if (provider === 'google') {
// Gemini doesn't provide token counts in response
// Would need to estimate
}
else if (provider === 'anthropic') {
if (response.usage) {
tokens.prompt = response.usage.input_tokens || 0;
tokens.completion = response.usage.output_tokens || 0;
tokens.total = tokens.prompt + tokens.completion;
}
}
return tokens;
}
/**
* Calculate estimated cost based on tokens and model
*/
calculateCost(tokens, model, provider) {
// Simplified pricing (would need full pricing table)
const pricing = {
openai: {
'gpt-4': { input: 0.03, output: 0.06 },
'gpt-3.5-turbo': { input: 0.001, output: 0.002 }
},
anthropic: {
'claude-3-opus': { input: 0.015, output: 0.075 },
'claude-3-sonnet': { input: 0.003, output: 0.015 }
},
google: {
'gemini-pro': { input: 0.001, output: 0.002 }
}
};
// Get provider pricing
const providerPricing = pricing[provider] || {};
let modelPricing = null;
// Find matching model pricing
for (const modelKey of Object.keys(providerPricing)) {
if (model.toLowerCase().includes(modelKey)) {
modelPricing = providerPricing[modelKey];
break;
}
}
if (!modelPricing) {
return 0.0;
}
// Calculate cost (per 1K tokens)
const inputCost = (tokens.prompt / 1000) * modelPricing.input;
const outputCost = (tokens.completion / 1000) * modelPricing.output;
return inputCost + outputCost;
}
/**
* Get aggregated chain execution summary
*/
async getChainSummary(chainId) {
if (!this.executionTracker) {
throw new Error('Execution tracking not available - no .pit directory found');
}
return await this.executionTracker.getChainSummary(chainId);
}
}
// Create singleton instance
exports.model = new Model();