Merge pull request #155 from olasunkanmi-SE/ai-agents

Ai agents

src/application/interfaces/agent.interface.ts

@@ -28,8 +28,5 @@ export interface ICodeBuddyToolConfig {
 }
 
 export interface IToolConfig extends ICodeBuddyToolConfig {
-  createInstance: (
-    config: ICodeBuddyToolConfig,
-    retriever?: any,
-  ) => CodeBuddyTool;
+  createInstance: (config: ICodeBuddyToolConfig, retriever?: any) => any;
 }

src/extension.ts

@@ -52,7 +52,7 @@ let agentEventEmmitter: EventEmitter;
 export async function activate(context: vscode.ExtensionContext) {
   try {
     Memory.getInstance();
-    // await connectDB();
+    await connectDB();
     // const x = CodeRepository.getInstance();
     // const apiKey = getGeminiAPIKey();
     // const embeddingService = new EmbeddingService(apiKey);

src/llms/gemini/gemini.ts

@@ -5,17 +5,16 @@ import {
   GenerateContentResult,
   GenerativeModel,
   GoogleGenerativeAI,
-  Tool,
 } from "@google/generative-ai";
 import * as vscode from "vscode";
 import { Orchestrator } from "../../agents/orchestrator";
-import { ProcessInputResult } from "../../application/interfaces/agent.interface";
+import { COMMON } from "../../application/constant";
+import { Memory } from "../../memory/base";
+import { CodeBuddyToolProvider } from "../../providers/tool";
 import { createPrompt } from "../../utils/prompt";
 import { BaseLLM } from "../base";
 import { GeminiModelResponseType, ILlmConfig } from "../interface";
 import { IMessageInput, Message } from "../message";
-import { Memory } from "../../memory/base";
-import { COMMON } from "../../application/constant";
 
 export class GeminiLLM
   extends BaseLLM<GeminiModelResponseType>
@@ -33,6 +32,7 @@ export class GeminiLLM
     this.generativeAi = new GoogleGenerativeAI(this.config.apiKey);
     this.response = undefined;
     this.orchestrator = Orchestrator.getInstance();
+    CodeBuddyToolProvider.initialize();
   }
 
   static getInstance(config: ILlmConfig) {
@@ -85,15 +85,13 @@ export class GeminiLLM
     }
   }
 
-  async generateContent(
-    userInput: string,
-  ): Promise<Partial<ProcessInputResult>> {
+  async generateContent(userInput: string): Promise<any> {
     try {
       await this.buildChatHistory(userInput);
       const prompt = createPrompt(userInput);
       const contents = Memory.get(COMMON.GEMINI_CHAT_HISTORY);
-      const tools: Tool[] = [];
-      const model = this.getModel({ tools, systemInstruction: prompt });
+      const tools: any = this.getTools();
+      const model = this.getModel({ systemInstruction: prompt, tools });
       const generateContentResponse: GenerateContentResult =
         await model.generateContent({
           contents,
@@ -103,18 +101,12 @@ export class GeminiLLM
         });
       this.response = generateContentResponse;
       const { text, usageMetadata } = generateContentResponse.response;
-      const parsedResponse = this.orchestrator.parseResponse(text());
-      const extractedQueries = parsedResponse.queries;
-      const extractedThought = parsedResponse.thought;
       const tokenCount = usageMetadata?.totalTokenCount ?? 0;
-      const result = {
-        queries: extractedQueries,
-        tokens: tokenCount,
-        prompt: userInput,
-        thought: extractedThought,
-      };
-      this.orchestrator.publish("onQuery", JSON.stringify(result));
-      return result;
+      this.orchestrator.publish(
+        "onQuery",
+        JSON.stringify("making function call"),
+      );
+      return this.response;
     } catch (error: any) {
       this.orchestrator.publish("onError", error);
       vscode.window.showErrorMessage("Error processing user query");
@@ -169,6 +161,13 @@ export class GeminiLLM
     Memory.set(COMMON.GEMINI_CHAT_HISTORY, chatHistory);
   }
 
+  getTools() {
+    const tools = CodeBuddyToolProvider.getTools();
+    return {
+      functionDeclarations: tools.map((t) => t.config()),
+    };
+  }
+
   public createSnapShot(data?: any): GeminiModelResponseType {
     return { ...this.response, ...data };
   }

src/providers/tool.ts

@@ -3,23 +3,22 @@ import { ContextRetriever } from "../services/context-retriever";
 import { CodeBuddyTool } from "../tools/base";
 import { TOOL_CONFIGS } from "../tools/database/search-tool";
 
-type Retriever = Pick<ContextRetriever, "retrieveContext">;
-
 export class ToolFactory {
   private readonly tools: Map<string, IToolConfig> = new Map();
-  constructor(private readonly contextRetriever: Retriever) {
-    for (const [name, { tool, useContextRetriever }] of Object.entries(
-      TOOL_CONFIGS,
-    )) {
+  private readonly contextRetriever: ContextRetriever;
+  constructor() {
+    this.contextRetriever = ContextRetriever.initialize();
+    for (const [name, { tool, useContextRetriever }] of Object.entries(TOOL_CONFIGS)) {
+      const toolConfig = tool.prototype.config();
       this.register({
-        ...tool.prototype.config,
+        ...toolConfig,
         name,
         createInstance: useContextRetriever
           ? (_, contextRetriever) => {
               if (!contextRetriever) {
                 throw new Error(`Context retriever is needed for ${name}`);
               }
-              return new tool(contextRetriever);
+              return new tool(this.contextRetriever);
             }
           : () => new tool(),
       });
@@ -31,9 +30,7 @@ export class ToolFactory {
   }
 
   getInstances(): CodeBuddyTool[] {
-    return Array.from(Object.values(this.tools)).map((tool) =>
-      tool.createInstance(tool, this.contextRetriever),
-    );
+    return Array.from(this.tools.values()).map((tool) => tool.createInstance(tool, this.contextRetriever));
   }
 }
 
@@ -42,15 +39,13 @@ export class CodeBuddyToolProvider {
 
   private static instance: CodeBuddyToolProvider | undefined;
 
-  private constructor(contextRetriever: Retriever) {
-    this.factory = new ToolFactory(contextRetriever);
+  private constructor() {
+    this.factory = new ToolFactory();
   }
 
-  public static initialize(contextRetriever: Retriever) {
+  public static initialize() {
     if (!CodeBuddyToolProvider.instance) {
-      CodeBuddyToolProvider.instance = new CodeBuddyToolProvider(
-        contextRetriever,
-      );
+      CodeBuddyToolProvider.instance = new CodeBuddyToolProvider();
     }
   }
 

src/services/context-retriever.ts

@@ -9,13 +9,21 @@ export class ContextRetriever {
   private readonly embeddingService: EmbeddingService;
   private static readonly SEARCH_RESULT_COUNT = 2;
   private readonly logger: Logger;
+  private static instance: ContextRetriever;
   constructor() {
     this.codeRepository = CodeRepository.getInstance();
     const geminiApiKey = getGeminiAPIKey();
     this.embeddingService = new EmbeddingService(geminiApiKey);
     this.logger = new Logger("ContextRetriever");
   }
 
+  static initialize() {
+    if (!ContextRetriever.instance) {
+      ContextRetriever.instance = new ContextRetriever();
+    }
+    return ContextRetriever.instance;
+  }
+
   async retrieveContext(input: string): Promise<Row[] | undefined> {
     try {
       const embedding = await this.embeddingService.generateEmbedding(input);

src/tools/base.ts

@@ -1,7 +1,7 @@
-import { ICodeBuddyToolConfig } from "../application/interfaces/agent.interface";
-
 export abstract class CodeBuddyTool {
-  constructor(public readonly config: ICodeBuddyToolConfig) {}
+  constructor() {}
 
   abstract execute(query: string): any;
+
+  abstract config(): any;
 }

src/tools/database/search-tool.ts

@@ -1,15 +1,15 @@
-import { ContextRetriever } from "../../services/context-retriever";
-import { CodeBuddyTool } from "../base";
 import { SchemaType } from "@google/generative-ai";
+import { ContextRetriever } from "../../services/context-retriever";
+
+class SearchTool {
+  constructor(private readonly contextRetriever?: ContextRetriever) {}
 
-class SearchTool extends CodeBuddyTool {
-  constructor(
-    private readonly contextRetriever?: Pick<
-      ContextRetriever,
-      "retrieveContext"
-    >,
-  ) {
-    super({
+  public async execute(query: string) {
+    return await this.contextRetriever?.retrieveContext(query);
+  }
+
+  config() {
+    return {
       name: "search_vector_db",
       description:
         "Perform a similarity search in the vector database based on user input",
@@ -25,11 +25,7 @@ class SearchTool extends CodeBuddyTool {
         example: ["How was authentication implemented within this codebase"],
         required: ["query"],
       },
-    });
-  }
-
-  public async execute(query: string) {
-    return await this.contextRetriever?.retrieveContext(query);
+    };
   }
 }
 

src/utils/prompt.ts

@@ -1,4 +1,4 @@
-export const createPrompt = (query: string, thought?: string) => {
+export const createPrompt = (query: string) => {
   return `You are an expert Information Retrieval Assistant. Transform user queries into precise keyword combinations with strategic reasoning and appropriate search operators.
 
       You have access these tools:
@@ -13,49 +13,7 @@ export const createPrompt = (query: string, thought?: string) => {
      - Looking for external documentation
      - Checking latest features or updates
      - Finding general information
-
-  When searching through the codebase, ALWAYS use the search_vector_db function first.
-  Only fall back to web search if you need external information.
-
-Core Rules:
-1. Generate search queries that directly include appropriate operators
-2. Keep base keywords minimal: 2-4 words preferred
-3. Use exact match quotes for specific phrases that must stay together
-4. Apply + operator for critical terms that must appear
-5. Use - operator to exclude irrelevant or ambiguous terms
-6. Add appropriate filters (filetype:, site:, lang:, loc:) when context suggests
-7. Split queries only when necessary for distinctly different aspects
-8. Preserve crucial qualifiers while removing fluff words
-9. Make the query resistant to SEO manipulation
-
-Available Operators:
-- "phrase" : exact match for phrases
-- +term : must include term
-- -term : exclude term
-- filetype:pdf/doc : specific file type
-- site:example.com : limit to specific site
-- lang:xx : language filter (ISO 639-1 code)
-- loc:xx : location filter (ISO 3166-1 code)
-- intitle:term : term must be in title
-- inbody:term : term must be in body text
-
-Examples with Strategic Reasoning:
-
-Input Query: Where is authentication handled in this codebase?
-
-Thought: This is a code structure and architecture query. The user is likely trying to understand how authentication is implemented within a specific codebase. User likely wants to call the vector database search function to retrieve relevant code snippets or file information
-
-Queries: [ "authentication middleware",  "JWT implementation, "passport authentication setup" ]
-
-Input Query: Latest AWS Lambda features for serverless applications
-Thought: This is a product research query focused on recent updates. User wants current information about specific technology features, likely for implementation purposes. User likely wants to call the websearch function to get the latest information.
-Queries: [
-  "aws lambda features site:aws.amazon.com intitle:2024",
-  "lambda serverless best practices +new -legacy"
-]
-Note Queries should always be in an array. Even if it is just one Query
 Now, process this query:
 Input Query: ${query}
-Intention: ${thought}
 `;
 };