Add embedding service with Xenova transformers for vector embeddings

Implement context service for code indexing and function mapping
Configure embedding settings with batch processing and rate limiting
Update code repository to separate table creation and data insertion
Add @xenova/transformers dependency
Refactor TypeScriptAtsMapper to use singleton pattern

package.json

@@ -335,6 +335,7 @@
     "@google/generative-ai": "^0.3.0",
     "@libsql/client": "^0.14.0",
     "@types/node-fetch": "^2.6.11",
+    "@xenova/transformers": "^2.17.2",
     "dotenv": "^16.1.4",
     "groq-sdk": "^0.3.2",
     "markdown-it": "^14.1.0",

src/application/constant.ts

@@ -65,3 +65,12 @@ export enum FSPROPS {
   TS_FILE_PATTERN = "**/*.ts",
   TSCONFIG_FILE = "tsconfig.json",
 }
+
+export const EmbeddingsConfig = {
+  batchSize: 5,
+  maxRetries: 3,
+  retryDelay: 1000,
+  rateLimit: 1500,
+  embeddingModel: "Xenova/all-MiniLM-L6-v2",
+  textModel: "gemini-1.5-flash",
+};

src/extension.ts

@@ -47,7 +47,7 @@ export async function activate(context: vscode.ExtensionContext) {
       const codeMapper = new TypeScriptAtsMapper();
       const mappedCode = await codeMapper.buildCodebaseMap();
       const ats = Object.values(mappedCode).flatMap((repo) =>
-        Object.values(repo.modules),
+        Object.values(repo.modules)
       );
       const mapper = new CodeStructureMapper(ats);
       return mapper.normalizeData();
@@ -70,49 +70,49 @@ export async function activate(context: vscode.ExtensionContext) {
     } = OLA_ACTIONS;
     const getComment = new Comments(
       `${USER_MESSAGE} generates the code comments...`,
-      context,
+      context
     );
     const getInLineChat = new InLineChat(
       `${USER_MESSAGE} generates a response...`,
-      context,
+      context
     );
     const generateOptimizeCode = new OptimizeCode(
       `${USER_MESSAGE} optimizes the code...`,
-      context,
+      context
     );
     const generateRefactoredCode = new RefactorCode(
       `${USER_MESSAGE} refactors the code...`,
-      context,
+      context
     );
     const explainCode = new ExplainCode(
       `${USER_MESSAGE} explains the code...`,
-      context,
+      context
     );
     const generateReview = new ReviewCode(
       `${USER_MESSAGE} reviews the code...`,
-      context,
+      context
     );
     const codeChartGenerator = new CodeChartGenerator(
       `${USER_MESSAGE} creates the code chart...`,
-      context,
+      context
     );
     const codePattern = new FileUploader(context);
     const knowledgeBase = new ReadFromKnowledgeBase(
       `${USER_MESSAGE} generate your code pattern...`,
-      context,
+      context
     );
     const generateCommitMessage = new GenerateCommitMessage(
       `${USER_MESSAGE} generates a commit message...`,
-      context,
+      context
     );
     const generateInterviewQuestions = new InterviewMe(
       `${USER_MESSAGE} generates interview questions...`,
-      context,
+      context
     );
 
     const generateUnitTests = new GenerateUnitTest(
       `${USER_MESSAGE} generates unit tests...`,
-      context,
+      context
     );
 
     const actionMap = {
@@ -126,7 +126,7 @@ export async function activate(context: vscode.ExtensionContext) {
         new FixError(
           `${USER_MESSAGE} finds a solution to the error...`,
           context,
-          errorMessage,
+          errorMessage
         ).execute(errorMessage),
       [explain]: () => explainCode.execute(),
       [pattern]: () => codePattern.uploadFileHandler(),
@@ -137,7 +137,7 @@ export async function activate(context: vscode.ExtensionContext) {
     };
 
     const subscriptions: vscode.Disposable[] = Object.entries(actionMap).map(
-      ([action, handler]) => vscode.commands.registerCommand(action, handler),
+      ([action, handler]) => vscode.commands.registerCommand(action, handler)
     );
 
     const selectedGenerativeAiModel = getConfigValue("generativeAi.option");
@@ -146,7 +146,7 @@ export async function activate(context: vscode.ExtensionContext) {
     const quickFixCodeAction: vscode.Disposable =
       vscode.languages.registerCodeActionsProvider(
         { scheme: "file", language: "*" },
-        quickFix,
+        quickFix
       );
 
     const modelConfigurations: {
@@ -186,13 +186,13 @@ export async function activate(context: vscode.ExtensionContext) {
         key,
         webviewProviderClass,
         subscriptions,
-        quickFixCodeAction,
+        quickFixCodeAction
       );
     }
   } catch (error) {
     Brain.clear();
     vscode.window.showErrorMessage(
-      "An Error occured while setting up generative AI model",
+      "An Error occured while setting up generative AI model"
     );
     console.log(error);
   }

src/infrastructure/repository/code-repository.ts

@@ -1,7 +1,7 @@
 import { ICodeRepository } from "../../application/interfaces/code.repository.interface";
 import { Client, createClient, ResultSet, Row } from "@libsql/client";
 import { Logger } from "../logger/logger";
-import { createTableQuery, selectFunctionProps } from "./sql";
+import { createTableQuery, insertDataQuery, selectFunctionProps } from "./sql";
 
 export class CodeRepository implements ICodeRepository {
   private client: Client | undefined;
@@ -30,9 +30,23 @@ export class CodeRepository implements ICodeRepository {
     return CodeRepository.instance;
   }
 
-  async CreateTable(values: string): Promise<ResultSet[] | undefined> {
+  async CreateTable(): Promise<ResultSet[] | undefined> {
     try {
-      const query = createTableQuery(values);
+      const query = createTableQuery();
+      const table = await this.client?.batch(query, "write");
+      if (table) {
+        this.logger.info("Database initialized successfully");
+      }
+      return table;
+    } catch (error) {
+      this.logger.error("Failed to initialize database", error);
+      throw error;
+    }
+  }
+
+  async InsertData(values: string) {
+    try {
+      const query = insertDataQuery(values);
       const table = await this.client?.batch(query, "write");
       if (table) {
         this.logger.info("Database initialized successfully");
@@ -46,7 +60,7 @@ export class CodeRepository implements ICodeRepository {
 
   async searchSimilarFunctions(
     queryEmbeddings: number[],
-    limit: number,
+    limit: number
   ): Promise<Row[] | undefined> {
     try {
       const query = selectFunctionProps();

src/infrastructure/repository/sql.ts

@@ -1,4 +1,4 @@
-export const createTableQuery = (values: string) => {
+export const createTableQuery = () => {
   return [
     "DROP TABLE IF EXISTS code_functions",
     `CREATE TABLE IF NOT EXISTS code_functions (
@@ -9,8 +9,13 @@ export const createTableQuery = (values: string) => {
       embedding F32_BLOB(768) NOT NULL
     )`,
     "CREATE INDEX IF NOT EXISTS code_functions_idx ON code_functions (libsql_vector_idx(embedding))",
+  ];
+};
+
+export const insertDataQuery = (values: string) => {
+  return [
     `INSERT INTO code_functions (class_name, function_name, file_path, created_at, embedding) 
-     VALUES ${values}`,
+          VALUES ${values}`,
   ];
 };
 

src/services/code-structure.mapper.service.ts

@@ -54,7 +54,7 @@ export class CodeStructureMapper {
    */
   private getMappedFunctions(
     codeEntry: ICodeEntry,
-    primaryClass: ICodeClass | null,
+    primaryClass: ICodeClass | null
   ): IMappedFunction[] {
     try {
       const sourceFunctions = this.getSourceFunctions(codeEntry, primaryClass);
@@ -77,7 +77,7 @@ export class CodeStructureMapper {
    */
   private getSourceFunctions(
     codeEntry: ICodeEntry,
-    primaryClass: ICodeClass | null,
+    primaryClass: ICodeClass | null
   ): IMappedFunction[] | undefined {
     try {
       return (codeEntry.functions?.length ?? 0) > 0
@@ -97,7 +97,7 @@ export class CodeStructureMapper {
    */
   private mapFunctions(
     functions: IMappedFunction[],
-    codeEntry?: ICodeEntry,
+    codeEntry?: ICodeEntry
   ): IMappedFunction[] {
     try {
       return functions.map((func) => ({
@@ -111,7 +111,7 @@ export class CodeStructureMapper {
           func,
           codeEntry?.dependencies && codeEntry?.dependencies.length > 0
             ? codeEntry.dependencies
-            : undefined,
+            : undefined
         ),
       }));
     } catch (error) {
@@ -122,7 +122,7 @@ export class CodeStructureMapper {
 
   private createCompositeText(
     functionProps: IMappedFunction,
-    dependencies: string[] = [],
+    dependencies: string[] = []
   ): string {
     if (!this.createDescriptionText(functionProps)) {
       return "";

src/services/context-service.ts

@@ -0,0 +1,45 @@
+import { EmbeddingService } from "./embedding-service";
+import { APP_CONFIG } from "../application/constant";
+import { IFunctionData } from "../application/interfaces";
+import { getConfigValue } from "../application/utils";
+import { Logger } from "../infrastructure/logger/logger";
+import { CodeStructureMapper } from "./code-structure.mapper.service";
+import { TypeScriptAtsMapper } from "./typescript-ats.service";
+
+export class ContextService {
+  logger: Logger;
+  embeddingService: EmbeddingService;
+  constructor() {
+    this.logger = new Logger("ContextService");
+    const apiKey = this.getAPIKey();
+    this.embeddingService = new EmbeddingService(apiKey);
+  }
+
+  getAPIKey(): string {
+    const { geminiKey } = APP_CONFIG;
+    const apiKey = getConfigValue(geminiKey);
+    if (!apiKey) {
+      this.logger.info("Gemini API Key is required for code indexing");
+      throw new Error("Gemini API Key is required for code indexing");
+    }
+    return apiKey;
+  }
+
+  async buildFunctionStructureMap() {
+    const codeATS = TypeScriptAtsMapper.getInstance();
+    const mappedCode = await codeATS.buildCodebaseMap();
+    const ats = Object.values(mappedCode).flatMap((repo) =>
+      Object.values(repo.modules)
+    );
+    const mapper = new CodeStructureMapper(ats);
+    return mapper.normalizeData();
+  }
+
+  async generateFunctionDescription() {
+    const functions =
+      (await this.buildFunctionStructureMap()) as IFunctionData[];
+    return await this.embeddingService.processFunctions(functions);
+  }
+
+  //   async generateContextEmbeddings
+}