Personal-Test-Ground/src/FFB/ffb-query-planner.service.ts

import { Injectable, OnModuleInit } from '@nestjs/common';
import * as fs from 'fs';
import * as path from 'path';
import axios from 'axios';
import { AgentQueryPlan } from './ffb-agent.types';

@Injectable()
export class FFBQueryPlannerService implements OnModuleInit {
  private systemPrompt: any;

  async onModuleInit() {
    const filePath = path.join(process.cwd(), 'AgentQueryPlan.json'); // updated file
    const data = fs.readFileSync(filePath, 'utf-8');
    this.systemPrompt = JSON.parse(data);
  }

  private buildPrompt(userMessage: string): string {
    const examplesText = (this.systemPrompt.examples || [])
      .map(
        (ex: any) =>
          `Q: "${ex.question}"\nA: ${JSON.stringify(ex.plan, null, 2)}`
      )
      .join('\n\n');

    return `
${this.systemPrompt.instructions}

Always include the minimal "fields" needed for computation to reduce bandwidth.

${examplesText}

Now, given the following user question, output the JSON only:

Q: "${userMessage}"
`;
  }

  async plan(userMessage: string): Promise<AgentQueryPlan> {
    const promptText = this.buildPrompt(userMessage);
    const responseText = await this.callGemini(promptText);
    const sanitized = this.sanitizeLLMOutput(responseText);

    try {
      return JSON.parse(sanitized);
    } catch (err) {
      console.error('Failed to parse Gemini output:', sanitized);
      throw new Error('LLM returned invalid JSON');
    }
  }

  private async callGemini(prompt: string): Promise<string> {
    const apiKey = process.env.GOOGLE_API_KEY;
    if (!apiKey) throw new Error('Missing GOOGLE_API_KEY');

    const url =
      'https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent';

    const body = {
      contents: [{ role: 'user', parts: [{ text: prompt }] }],
    };

    try {
      const response = await axios.post(url, body, {
        headers: {
          'Content-Type': 'application/json',
          'x-goog-api-key': apiKey,
        },
      });

      const text =
        response.data?.candidates?.[0]?.content?.parts
          ?.map((p: any) => p.text)
          .join(' ') ?? '';

      if (!text) throw new Error('No text generated by Gemini');
      return text;
    } catch (err: any) {
      console.error('Failed to call Gemini:', err.response?.data || err.message);
      throw err;
    }
  }

  private sanitizeLLMOutput(text: string): string {
    return text
      .trim()
      .replace(/^```json\s*/, '') // remove opening ```json
      .replace(/^```\s*/, '')     // remove opening ```
      .replace(/```$/, '')        // remove closing ```
      .trim();
  }
}