Add WIP of new file select agent. Implement vercel ai for deepseek, o…

…penai and vertex

bin/aider

@@ -2,5 +2,6 @@
 source variables/local.env
 export VERTEXAI_PROJECT=$GCLOUD_PROJECT
 export VERTEXAI_LOCATION=$GCLOUD_CLAUDE_REGION
+export OPENAI_API_KEY=$OPENAI_API_KEY
 echo $VERTEXAI_PROJECT $VERTEXAI_LOCATION
-aider --model vertex_ai/claude-3-5-sonnet@20240620
+aider --o1-preview --editor-model vertex_ai/claude-3-5-sonnet@20240620

src/llm/llm.ts

@@ -71,7 +71,7 @@ export function assistant(text: string): LlmMessage {
 export interface LLM {
 	generateTextFromMessages(messages: LlmMessage[], opts?: GenerateTextOptions): Promise<string>;
 
-	/* Generates a response that is expected to be in JSON format, and returns the object */
+	/* Generates a response that is expected to be in JSON format, or end with a JSON object wrapped in <json> tags or Markdown triple ticks, and returns the object */
 	generateJsonFromMessages<T>(messages: LlmMessage[], opts?: GenerateJsonOptions): Promise<T>;
 
 	/* Generates text from a LLM */
@@ -108,7 +108,7 @@ export interface LLM {
 	getService(): string;
 
 	/**
-	 * The LLM model identifier
+	 * The LLM model identifier. This should match the model ids in the Vercel ai module (https://github.com/vercel/ai)
 	 */
 	getModel(): string;
 

src/llm/models/deepseek.ts

@@ -1,7 +1,8 @@
+import { createOpenAI } from '@ai-sdk/openai';
+import { LanguageModel } from 'ai';
 import axios from 'axios';
 import { addCost, agentContext } from '#agent/agentContextLocalStorage';
 import { LlmCall } from '#llm/llmCallService/llmCall';
-import { CallerId } from '#llm/llmCallService/llmCallService';
 import { withSpan } from '#o11y/trace';
 import { currentUser } from '#user/userService/userContext';
 import { sleep } from '#utils/async-utils';
@@ -35,6 +36,7 @@ export function deepseekChat(): LLM {
  */
 export class DeepseekLLM extends BaseLLM {
 	_client: any;
+	aimodel: LanguageModel;
 
 	client() {
 		if (!this._client) {
@@ -52,6 +54,16 @@ export class DeepseekLLM extends BaseLLM {
 		return Boolean(currentUser().llmConfig.deepseekKey || process.env.DEEPSEEK_API_KEY);
 	}
 
+	aiModel(): LanguageModel {
+		if (!this.aimodel) {
+			this.aimodel = createOpenAI({
+				baseURL: 'https://api.deepseek.com',
+				apiKey: currentUser().llmConfig.deepseekKey || envVar('DEEPSEEK_API_KEY'),
+			})('deepseek-coder');
+		}
+		return this.aimodel;
+	}
+
 	constructor(
 		displayName: string,
 		model: string,

src/llm/models/openai.ts

@@ -1,4 +1,6 @@
+import { createOpenAI, openai } from '@ai-sdk/openai';
 import { OpenAIChatModelId } from '@ai-sdk/openai/internal';
+import { LanguageModel } from 'ai';
 import { OpenAI as OpenAISDK } from 'openai';
 import { addCost, agentContext } from '#agent/agentContextLocalStorage';
 import { LlmCall } from '#llm/llmCallService/llmCall';
@@ -77,11 +79,12 @@ export function GPT4oMini() {
 
 export class OpenAI extends BaseLLM {
 	openAISDK: OpenAISDK | null = null;
+	aimodel: LanguageModel;
 
 	constructor(
 		name: string,
 		model: Model,
-		aiModel: OpenAIChatModelId,
+		private aiModelId: OpenAIChatModelId,
 		maxInputTokens: number,
 		calculateInputCost: (input: string) => number,
 		calculateOutputCost: (output: string) => number,
@@ -102,6 +105,15 @@ export class OpenAI extends BaseLLM {
 		return Boolean(currentUser().llmConfig.openaiKey || process.env.OPENAI_API_KEY);
 	}
 
+	aiModel(): LanguageModel {
+		if (!this.aimodel) {
+			this.aimodel = createOpenAI({
+				apiKey: currentUser().llmConfig.openaiKey || envVar('OPENAI_API_KEY'),
+			})(this.getModel());
+		}
+		return this.aimodel;
+	}
+
 	async generateImage(description: string): Promise<string> {
 		const response = await this.sdk().images.generate({
 			model: 'dall-e-3',

src/llm/models/vertexai.ts

@@ -1,4 +1,6 @@
+import { createVertex, vertex } from '@ai-sdk/google-vertex';
 import { GenerativeModel, HarmBlockThreshold, HarmCategory, SafetySetting, VertexAI } from '@google-cloud/vertexai';
+import { LanguageModel } from 'ai';
 import axios from 'axios';
 import { addCost, agentContext } from '#agent/agentContextLocalStorage';
 import { AgentLLMs } from '#agent/agentContextTypes';
@@ -41,11 +43,11 @@ export function vertexLLMRegistry(): Record<string, () => LLM> {
 // gemini-1.5-pro-latest
 // gemini-1.5-pro-exp-0801
 // exp-0801
-export function Gemini_1_5_Pro(version = '002') {
+export function Gemini_1_5_Pro() {
 	return new VertexLLM(
 		'Gemini 1.5 Pro',
 		VERTEX_SERVICE,
-		`gemini-1.5-pro-${version}`,
+		'gemini-1.5-pro',
 		1_000_000,
 		(input: string) => (input.length * (input.length > 128_000 * 4 ? 0.0003125 : 0.000625)) / 1000,
 		(output: string) => (output.length * (output.length > 128_000 * 4 ? 0.0025 : 0.00125)) / 1000,
@@ -63,11 +65,11 @@ export function Gemini_1_5_Experimental() {
 	);
 }
 
-export function Gemini_1_5_Flash(version = '002') {
+export function Gemini_1_5_Flash() {
 	return new VertexLLM(
 		'Gemini 1.5 Flash',
 		VERTEX_SERVICE,
-		`gemini-1.5-flash-${version}`,
+		'gemini-1.5-flash',
 		1_000_000,
 		(input: string) => (input.length * 0.000125) / 1000,
 		(output: string) => (output.length * 0.000375) / 1000,
@@ -147,6 +149,7 @@ export function Vertex_Llama3_405b() {
  */
 class VertexLLM extends BaseLLM {
 	_vertex: VertexAI;
+	aimodel: LanguageModel;
 
 	vertex(): VertexAI {
 		if (!this._vertex) {
@@ -158,6 +161,14 @@ class VertexLLM extends BaseLLM {
 		return this._vertex;
 	}
 
+	aiModel(): LanguageModel {
+		if (!this.aimodel) {
+			const provider = createVertex({ project: process.env.GCLOUD_PROJECT, location: process.env.GCLOUD_REGION });
+			this.aimodel = provider(this.getModel());
+		}
+		return this.aimodel;
+	}
+
 	async generateText(userPrompt: string, systemPrompt?: string, opts?: GenerateTextOptions): Promise<string> {
 		return withActiveSpan(`generateText ${opts?.id ?? ''}`, async (span) => {
 			if (systemPrompt) span.setAttribute('systemPrompt', systemPrompt);

src/routes/code/code-routes.ts

@@ -9,8 +9,8 @@ import { ClaudeVertexLLMs } from '#llm/models/anthropic-vertex';
 import { Gemini_1_5_Flash } from '#llm/models/vertexai';
 import { logger } from '#o11y/logger';
 import { CodeEditingAgent } from '#swe/codeEditingAgent';
-import { codebaseQuery } from '#swe/codebaseQuery';
-import { SelectFilesResponse, selectFilesToEdit } from '#swe/selectFilesToEdit';
+import { codebaseQuery } from '#swe/discovery/codebaseQuery';
+import { SelectFilesResponse, selectFilesToEdit } from '#swe/discovery/selectFilesToEdit';
 import { AppFastifyInstance } from '../../app';
 import { sophiaDirName, systemDir } from '../../appVars';
 

src/swe/SWEBenchAgent.ts

@@ -5,9 +5,9 @@ import { ClaudeVertexLLMs } from '#llm/models/anthropic-vertex';
 import { countTokens } from '#llm/tokens';
 import { logger } from '#o11y/logger';
 import { CodeEditingAgent } from '#swe/codeEditingAgent';
+import { selectFilesToEdit } from '#swe/discovery/selectFilesToEdit';
 import { PythonTools } from '#swe/lang/python/pythonTools';
 import { ProjectInfo } from '#swe/projectDetection';
-import { selectFilesToEdit } from '#swe/selectFilesToEdit';
 import { MAP_REPO_TO_TEST_FRAMEWORK, MAP_VERSION_TO_INSTALL, VersionInstallation } from '#swe/sweBenchConstant';
 import { execCommand } from '#utils/exec';
 

src/swe/codeEditingAgent.ts

@@ -13,9 +13,9 @@ import { execCommand, runShellCommand } from '#utils/exec';
 import { appContext } from '../app';
 import { cacheRetry } from '../cache/cacheRetry';
 import { AiderCodeEditor } from './aiderCodeEditor';
+import { SelectFilesResponse, selectFilesToEdit } from './discovery/selectFilesToEdit';
 import { ProjectInfo, detectProjectInfo } from './projectDetection';
 import { basePrompt } from './prompt';
-import { SelectFilesResponse, selectFilesToEdit } from './selectFilesToEdit';
 import { summariseRequirements } from './summariseRequirements';
 
 export function buildPrompt(args: {

src/swe/codeFunctions.ts

@@ -0,0 +1,36 @@
+import { func, funcClass } from '#functionSchema/functionDecorators';
+import { codebaseQuery } from '#swe/discovery/codebaseQuery';
+import { SelectFilesResponse, selectFilesToEdit } from '#swe/discovery/selectFilesToEdit';
+import { getProjectInfo } from '#swe/projectDetection';
+import { reviewChanges } from '#swe/reviewChanges';
+
+@funcClass(__filename)
+export class CodeFunctions {
+	/**
+	 * Searches across files under the current working directory to provide an answer to the query
+	 * @param query
+	 */
+	@func()
+	async queryRepository(query: string): Promise<string> {
+		return await codebaseQuery(query);
+	}
+
+	/**
+	 *
+	 * @param requirements
+	 */
+	@func()
+	async selectFilesToEdit(requirements: string): Promise<SelectFilesResponse> {
+		return await selectFilesToEdit(requirements, await getProjectInfo());
+	}
+
+	/**
+	 * Reviews the changes committed to git since a commit or start of a branch
+	 * @param requirements
+	 * @param sourceBranchOrCommit
+	 */
+	@func()
+	async reviewChanges(requirements: string, sourceBranchOrCommit: string) {
+		return await reviewChanges(requirements, sourceBranchOrCommit);
+	}
+}

src/swe/discovery/README.md

@@ -0,0 +1,2 @@
+This folder contains agents which perform the discovery/research 
+phase of a task.

src/swe/discovery.ts → src/swe/discovery/discovery.ts

@@ -1,7 +1,7 @@
 import { getFileSystem, llms } from '#agent/agentContextLocalStorage';
 import { logger } from '#o11y/logger';
+import { SelectedFile, selectFilesToEdit } from '#swe/discovery/selectFilesToEdit';
 import { getProjectInfo } from '#swe/projectDetection';
-import { SelectedFile, selectFilesToEdit } from '#swe/selectFilesToEdit';
 
 interface DiscoveryResult {
 	readyForExecution: boolean;