test-ai-code-agents/src/functions/prompts-to-test-spec/src/services/gemini-service.ts

/**
 * Service for handling Gemini API operations
 */
import { VertexAI } from '@google-cloud/vertexai';
import * as fs from 'fs';
import * as path from 'path';
import { Project, Workitem } from '../types';
import { GOOGLE_CLOUD_PROJECT_ID, GOOGLE_CLOUD_LOCATION, GEMINI_MODEL } from '../config';

export class GeminiService {
  private vertexAI: VertexAI;
  private model: string;
  private projectId: string;
  private location: string;

  constructor(projectId?: string, location?: string, model?: string) {
    this.projectId = projectId || GOOGLE_CLOUD_PROJECT_ID;
    this.location = location || GOOGLE_CLOUD_LOCATION;
    this.model = model || GEMINI_MODEL;

    if (!this.projectId) {
      throw new Error('Google Cloud Project ID is required');
    }

    this.vertexAI = new VertexAI({
      project: this.projectId,
      location: this.location,
    });
  }

  /**
   * Apply project guidelines to a workitem
   * @param project Project information
   * @param workitem Workitem to process
   * @param projectRepoPath Path to the cloned project repository
   * @returns Result of the processing
   */
  async processWorkitem(
    project: Project, 
    workitem: Workitem, 
    projectRepoPath: string
  ): Promise<{ success: boolean; error?: string }> {
    try {
      // Skip inactive workitems
      if (!workitem.isActive) {
        console.log(`Skipping inactive workitem: ${workitem.name}`);

        // If the feature file exists, it should be deleted according to guidelines
        const featureFileName = `${workitem.name}.feature`;
        const featurePath = path.join(projectRepoPath, 'nitro-it', 'src', 'test', 'resources', 'workitems', featureFileName);

        if (fs.existsSync(featurePath)) {
          fs.unlinkSync(featurePath);
          console.log(`Deleted feature file for inactive workitem: ${featurePath}`);
        }

        return { success: true };
      }

      // Read project guidelines
      const projectGuidelines = await this.readProjectGuidelines(project.path);

      // Read workitem content
      const workitemContent = fs.readFileSync(workitem.path, 'utf-8');

      // Generate feature file content using Gemini API
      const featureContent = await this.generateFeatureFile(
        projectGuidelines,
        workitemContent,
        workitem.name
      );

      // Ensure the target directory exists
      const targetDir = path.join(projectRepoPath, 'nitro-it', 'src', 'test', 'resources', 'workitems');
      if (!fs.existsSync(targetDir)) {
        fs.mkdirSync(targetDir, { recursive: true });
      }

      // Write the feature file
      const featureFileName = `${workitem.name}.feature`;
      const featurePath = path.join(targetDir, featureFileName);
      fs.writeFileSync(featurePath, featureContent);

      console.log(`Created/updated feature file: ${featurePath}`);

      return { success: true };
    } catch (error) {
      console.error(`Error processing workitem ${workitem.name}:`, error);
      return { 
        success: false, 
        error: error instanceof Error ? error.message : String(error) 
      };
    }
  }

  /**
   * Read AI guidelines for a project
   * @param projectPath Path to the project directory
   * @returns AI guidelines content
   */
  private async readProjectGuidelines(projectPath: string): Promise<string> {
    const aiPath = path.join(projectPath, 'AI.md');

    if (!fs.existsSync(aiPath)) {
      return '';
    }

    return fs.readFileSync(aiPath, 'utf-8');
  }

  /**
   * Generate feature file content using Gemini API
   * @param guidelines Project guidelines
   * @param workitemContent Workitem content
   * @param workitemName Name of the workitem
   * @returns Generated feature file content
   */
  private async generateFeatureFile(
    guidelines: string,
    workitemContent: string,
    workitemName: string
  ): Promise<string> {
    const generativeModel = this.vertexAI.getGenerativeModel({
      model: this.model,
    });

    const currentDate = new Date().toISOString();

    // Send the AI.md file directly to Gemini without hardcoded instructions
    const prompt = `
${guidelines}

Workitem:
${workitemContent}

Include the following comment at the top of the generated file:
# Generated by prompts-to-test-spec on ${currentDate}
# Source: ${workitemName}
`;

    const result = await generativeModel.generateContent({
      contents: [{ role: 'user', parts: [{ text: prompt }] }],
    });

    const response = await result.response;
    const generatedText = response.candidates[0]?.content?.parts[0]?.text || '';

    return generatedText;
  }

  /**
   * Generate a pull request description using Gemini API
   * @param processedWorkitems List of processed workitems
   * @returns Generated pull request description
   */
  async generatePullRequestDescription(
    processedWorkitems: { workitem: Workitem; success: boolean; error?: string }[]
  ): Promise<string> {
    const generativeModel = this.vertexAI.getGenerativeModel({
      model: this.model,
    });

    // Prepare workitem data for the prompt
    const added: string[] = [];
    const updated: string[] = [];
    const deleted: string[] = [];
    const failed: string[] = [];

    for (const item of processedWorkitems) {
      const { workitem, success, error } = item;

      if (!success) {
        failed.push(`- ${workitem.name}: ${error}`);
        continue;
      }

      if (!workitem.isActive) {
        deleted.push(`- ${workitem.name}`);
      } else if (workitem.implementation) {
        updated.push(`- ${workitem.name}`);
      } else {
        added.push(`- ${workitem.name}`);
      }
    }

    // Create a structured summary of changes for Gemini
    let workitemSummary = '';

    if (added.length > 0) {
      workitemSummary += 'Added workitems:\n' + added.join('\n') + '\n\n';
    }

    if (updated.length > 0) {
      workitemSummary += 'Updated workitems:\n' + updated.join('\n') + '\n\n';
    }

    if (deleted.length > 0) {
      workitemSummary += 'Deleted workitems:\n' + deleted.join('\n') + '\n\n';
    }

    if (failed.length > 0) {
      workitemSummary += 'Failed workitems:\n' + failed.join('\n') + '\n\n';
    }

    const prompt = `
You are tasked with creating a pull request description for changes to test specifications.

The following is a summary of the changes made:
${workitemSummary}

Create a clear, professional pull request description that:
1. Explains that this PR was automatically generated by the prompts-to-test-spec function
2. Summarizes the changes (added, updated, deleted, and failed workitems)
3. Uses markdown formatting for better readability
4. Keeps the description concise but informative

The pull request description should be ready to use without further editing.
`;

    const result = await generativeModel.generateContent({
      contents: [{ role: 'user', parts: [{ text: prompt }] }],
    });

    const response = await result.response;
    const generatedText = response.candidates[0]?.content?.parts[0]?.text || '';

    return generatedText;
  }
}