UNPKG

sourcesailor

Version:

A CLI tool for analyzing and documenting codebases

195 lines (194 loc) 9.87 kB
import Anthropic from "@anthropic-ai/sdk"; import { readConfig } from "./utils.mjs"; import { prompts } from "./prompts.mjs"; import { maskSensitiveInfo } from "./utils.mjs"; // Assuming this function exists in utils.mjs const modelMapping = { "claude-3-haiku": 'claude-3-haiku-20240307', "claude-3.5-haiku": 'claude-3-5-haiku-20241022', "claude-3.5-haiku-latest": 'claude-3-5-haiku-latest', 'claude-3-sonnet': 'claude-3-sonnet-20240229', 'claude-3-opus': 'claude-3-opus-20240229', 'claude-3.5-sonnet': 'claude-3-5-sonnet-20241022', 'claude-3.5-sonnet-legacy': 'claude-3-5-sonnet-20240620', 'claude-3.5-sonnet-latest': 'claude-3-5-sonnet-latest', 'haiku-3': 'claude-3-haiku-20240307', 'haiku-3.5': 'claude-3-5-haiku-20241022', 'haiku-3.5-latest': 'claude-3-5-haiku-latest', 'sonnet-3': 'claude-3-sonnet-20240229', 'opus-3': 'claude-3-opus-20240229', 'sonnet-3.5': 'claude-3-5-sonnet-20241022', 'sonnet-3.5-legacy': 'claude-3-5-sonnet-20240620', 'sonnet-3.5-latest': 'claude-3-5-sonnet-latest', }; export class AnthropicInterface { getModelId(model) { return modelMapping[model]; } getClient(isVerbose) { const config = readConfig(); const client = new Anthropic({ apiKey: config.ANTHROPIC_API_KEY || process.env.ANTHROPIC_API_KEY, }); if (isVerbose) { console.log("Anthropic client initialized"); } return client; } getName() { return "Anthropic"; } async listModels(isVerbose) { if (isVerbose) { console.log("Available Anthropic models:", maskSensitiveInfo(JSON.stringify(modelMapping, null, 2))); } return Object.keys(modelMapping); } async inferProjectDirectory(directoryStructure, allowStreaming, isVerbose, userExpertise, modelName) { const model = this.getModel(modelName); const systemPrompt = `${prompts.commonSystemPrompt.prompt}\n${prompts.rootUnderstanding.prompt}`; const userPrompt = `<FileStructure>${JSON.stringify(directoryStructure)}</FileStructure>`; const tools = prompts.rootUnderstanding.params ? { name: prompts.rootUnderstanding.params.name, input_schema: { type: "object", properties: { isMonorepo: { type: "boolean", description: prompts.rootUnderstanding.params.parameters.properties['isMonorepo'].description }, directories: { type: "array", items: { type: "string" }, description: prompts.rootUnderstanding.params.parameters.properties['directories'].description }, programmingLanguage: { type: "string", description: prompts.rootUnderstanding.params.parameters.properties['programmingLanguage'].description }, framework: { type: "string", description: prompts.rootUnderstanding.params.parameters.properties['framework'].description, }, dependenciesFile: { type: "string", description: prompts.rootUnderstanding.params.parameters.properties['dependenciesFile'].description }, lockFile: { type: "string", description: prompts.rootUnderstanding.params.parameters.properties['lockFile'].description }, entryPointFile: { type: "string", description: prompts.rootUnderstanding.params.parameters.properties['entryPointFile'].description }, workflow: { type: "string", description: prompts.rootUnderstanding.params.parameters.properties['workflow'].description }, }, }, description: prompts.rootUnderstanding.params.description, } : undefined; if (isVerbose) { console.log("Inferring project directory structure"); console.log("Model:", model); console.log("System prompt:", maskSensitiveInfo(systemPrompt)); console.log("User prompt:", maskSensitiveInfo(userPrompt)); console.log("Tools:", maskSensitiveInfo(JSON.stringify(tools, null, 2))); } return this.callApiAndReturnString(model, systemPrompt, userPrompt, tools, allowStreaming, isVerbose, userExpertise); } async inferDependency(dependencyFile, workflow, allowStreaming, isVerbose, userExpertise, modelName) { const model = this.getModel(modelName); const systemPrompt = `${prompts.commonSystemPrompt.prompt}\n${prompts.dependencyUnderstanding.prompt}`; const userPrompt = `<DependencyFile>${JSON.stringify(dependencyFile)}</DependencyFile>\n<Workflow>${workflow}</Workflow> ${prompts.commonMarkdownPrompt.prompt}`; return this.callApiAndReturnString(model, systemPrompt, userPrompt, undefined, allowStreaming, isVerbose, userExpertise); } async inferCode(directoryStructure, allowStreaming, isVerbose, userExpertise, modelName) { const model = this.getModel(modelName); const systemPrompt = `${prompts.commonSystemPrompt.prompt}\n${prompts.codeUnderstanding.prompt}`; const userPrompt = `<Code>${JSON.stringify(directoryStructure)}</Code> ${prompts.commonMarkdownPrompt.prompt}`; return this.callApiAndReturnString(model, systemPrompt, userPrompt, undefined, allowStreaming, isVerbose, userExpertise); } async inferInterestingCode(directoryStructure, allowStreaming, isVerbose, userExpertise, modelName) { const model = this.getModel(modelName); const systemPrompt = prompts.interestingCodeParts.prompt; const userPrompt = `<Code>${JSON.stringify(directoryStructure)}</Code> ${prompts.commonMarkdownPrompt.prompt}`; return this.callApiAndReturnString(model, systemPrompt, userPrompt, undefined, allowStreaming, isVerbose, userExpertise); } async generateReadme(directoryStructure, dependencyInference, codeInference, allowStreaming, isVerbose, userExpertise, modelName) { const model = this.getModel(modelName); const systemPrompt = prompts.readmePrompt.prompt; const userPrompt = `<DirectoryStructure>${JSON.stringify(directoryStructure)}</DirectoryStructure>\n<DependencyInference>${JSON.stringify(dependencyInference)}</DependencyInference>\n<CodeInference>${JSON.stringify(codeInference)}</CodeInference> ${prompts.commonMarkdownPrompt.prompt}`; return this.callApiAndReturnString(model, systemPrompt, userPrompt, undefined, allowStreaming, isVerbose, userExpertise); } getModel(modelName) { const config = readConfig(); const selectedModel = modelName || config.DEFAULT_ANTHROPIC_MODEL || process.env.DEFAULT_ANTHROPIC_MODEL || "claude-3-opus-20240229"; return this.getModelId(selectedModel); } async callApiAndReturnString(modelId, systemPrompt, userPrompt, tools, allowStreaming = false, isVerbose = false, userExpertise) { const client = this.getClient(isVerbose); let finalSystemPrompt = systemPrompt; if (userExpertise) { finalSystemPrompt += `\n<Expertise>${JSON.stringify(userExpertise)}</Expertise>`; } const messageConfig = { max_tokens: 8192, model: modelId, system: finalSystemPrompt, messages: [ { role: 'user', content: userPrompt } ], stream: allowStreaming }; if (tools) { messageConfig.tools = [tools]; messageConfig.tool_choice = { type: 'tool', name: tools.name }; delete messageConfig.stream; allowStreaming = false; } if (isVerbose) { console.log("Sending request to Anthropic API"); console.log("Message config:", maskSensitiveInfo(JSON.stringify(messageConfig, null, 2))); } const response = await client.messages.create(messageConfig); if (allowStreaming) { const stream = response; const streamedYield = this.convertStreamToStringStream(stream, isVerbose); return streamedYield; } else { const message = response; if (isVerbose) { console.log("Received response from Anthropic API"); console.log("Response:", maskSensitiveInfo(JSON.stringify(message, null, 2))); } if (message.stop_reason === 'tool_use') { const toolContent = message.content.find(contentData => contentData.type === 'tool_use'); if (toolContent) { return JSON.stringify(toolContent.input); } else { return ""; } } const data = message.content.filter(content => content.type === 'text').map(content => content.text).join("\n"); return data; } } async *convertStreamToStringStream(response, isVerbose) { for await (const chunk of response) { if (chunk.type === 'content_block_delta') { if (chunk.delta.type === 'text_delta') { if (isVerbose) { console.log("Received chunk:", maskSensitiveInfo(chunk.delta.text)); } yield chunk.delta.text; } } yield ""; } } } export default AnthropicInterface;