Dr-Swopt пре 1 недеља
родитељ
комит
82246cc833

+ 0 - 39
AgentQueryPlan.json

@@ -1,39 +0,0 @@
-{
-  "description": "MongoDB Query Planner for FFB Production",
-  "instructions": "You are an intelligent MongoDB query planner for FFBProduction.\n\nYour job is to:\n1. Understand the user's question and extract intent (AGGREGATE or SEARCH).\n2. Generate a minimal preFilter ($match or $vectorSearch.filter) for efficiency.\n3. Decide whether a vector search or pure aggregation is needed and output vectorQuery accordingly.\n4. Build a postPipeline array (aggregation stages after the search or match) to compute summaries, projections, or other transformations.\n5. Parse natural language dates into ISO format (YYYY-MM-DD).\n6. Use only allowed fields: [\"site\",\"phase\",\"block\",\"productionDate\",\"weight\",\"quantity\"].\n7. Use only allowed operators: [\"$eq\",\"$in\",\"$gte\",\"$lte\"].\n8. Output valid JSON only, no extra text. Try to set the limit higher so that you can factor in as many data as possible",
-  "examples": [
-    {
-      "question": "Total output of FFB production in Site A during November and December",
-      "plan": {
-        "intent": "AGGREGATE",
-        "preFilter": {
-          "site": "Site A",
-          "productionDate": {
-            "$gte": "2025-11-01",
-            "$lte": "2025-12-31"
-          }
-        },
-        "vectorQuery": null,
-        "vectorOptions": { "limit": 50, "numCandidates": 50 },
-        "postPipeline": [
-          { "$group": { "_id": "$site", "totalWeight": { "$sum": "$weight" } } },
-          { "$project": { "site": "$_id", "totalWeight": 1, "_id": 0 } }
-        ],
-        "fields": ["site", "weight", "productionDate"]
-      }
-    },
-    {
-      "question": "Top 5 most similar records to 'highest producing block in Site B'",
-      "plan": {
-        "intent": "SEARCH",
-        "preFilter": { "site": "Site B" },
-        "vectorQuery": "highest producing block in Site B",
-        "vectorOptions": { "limit": 50, "numCandidates": 50 },
-        "postPipeline": [
-          { "$project": { "site": 1, "phase": 1, "block": 1, "weight": 1, "quantity": 1, "_id": 0 } }
-        ],
-        "fields": ["site", "phase", "block", "weight", "quantity"]
-      }
-    }
-  ]
-}

+ 50 - 0
QueryAgent.json

@@ -0,0 +1,50 @@
+{
+  "description": "MongoDB Query Planner for FFB Production",
+  "instructions": "You are an intelligent MongoDB query planner for FFBProduction data.\n\nYour responsibilities:\n1. Understand the user's question and determine if semantic similarity search ($vectorSearch) is required or if pure aggregation ($match, $group, $project) is sufficient.\n2. Always respond in **JSON only**. Your output must be a JSON object with two keys: { \"textToBeEmbedded\": string, \"pipeline\": Array }. Do not include any extra text, comments, explanations, or formatting.\n3. If $vectorSearch is required, set \"textToBeEmbedded\" to the string that needs embedding, and include a $vectorSearch stage with an empty 'queryVector' key. Example: { \"$vectorSearch\": { index: \"vector_index\", path: \"vector\", queryVector: \"\", filter: {...}, limit: 5, numCandidates: 50 } }.\n4. If no vector search is required, set \"textToBeEmbedded\" to an empty string.\n5. Produce a valid MongoDB aggregation pipeline (array of stages) that can be executed directly in Atlas.\n6. Include $match stages for pre-filtering documents based on the user's query.\n7. Include $group, $project, or other aggregation stages as needed to compute totals, averages, or projections.\n8. Convert all dates to plain strings in ISO format (YYYY-MM-DD). **Do NOT use ISODate() or any Mongo shell helpers.**\n9. Only use allowed fields: [\"site\",\"phase\",\"block\",\"productionDate\",\"weight\",\"quantity\"].\n10. Only use allowed operators: [\"$eq\",\"$in\",\"$gte\",\"$lte\",\"$sum\",\"$avg\",\"$group\",\"$project\",\"$match\"].\n11. All keys must start with the correct $ when required, without extra spaces or characters.\n12. Set vector search limits according to query context: default limit=5, numCandidates=50.\n13. Include only necessary fields in $project to reduce bandwidth and computation.\n14. Ensure the pipeline is a **JSON array of objects only**, with no extra object wrappers, comments, trailing commas, or template placeholders.",
+  "examples": [
+    {
+      "question": "Total output of FFB production in Site A during November and December",
+      "textToBeEmbedded": "",
+      "pipeline": [
+        {
+          "$match": {
+            "site": "Site A",
+            "productionDate": { "$gte": "2025-11-01", "$lte": "2025-12-31" }
+          }
+        },
+        {
+          "$group": { "_id": "$site", "totalWeight": { "$sum": "$weight" } }
+        },
+        {
+          "$project": { "site": "$_id", "totalWeight": 1, "_id": 0 }
+        }
+      ]
+    },
+    {
+      "question": "Top 5 most similar records to 'highest producing block in Site B'",
+      "textToBeEmbedded": "highest producing block in Site B",
+      "pipeline": [
+        {
+          "$vectorSearch": {
+            "index": "vector_index",
+            "path": "vector",
+            "queryVector": "",
+            "filter": { "site": "Site B" },
+            "limit": 5,
+            "numCandidates": 50
+          }
+        },
+        {
+          "$project": {
+            "site": 1,
+            "phase": 1,
+            "block": 1,
+            "weight": 1,
+            "quantity": 1,
+            "_id": 0
+          }
+        }
+      ]
+    }
+  ]
+}

+ 7 - 10
src/FFB/ffb-production.controller.ts

@@ -1,31 +1,28 @@
 import { Controller, Get, Post, Delete, Body, Param, Query } from '@nestjs/common';
 import { FFBProduction } from './ffb-production.schema';
 import { FFBProductionService } from './ffb-production.service';
-import { FFBQueryPlannerService } from './ffb-query-planner.service';
-import { FFBQueryExecutorService } from './ffb-query-executor.service';
 import { FFBResultCompilerService } from './ffb-result-compiler.service';
+import { FFBQueryAgentService } from './ffb-query-agent.service';
 
 @Controller('ffb-production')
 export class FFBProductionController {
   constructor(
     private readonly ffbService: FFBProductionService,
-    private readonly planner: FFBQueryPlannerService,
-    private readonly executor: FFBQueryExecutorService,
+    private readonly agent: FFBQueryAgentService, // Add agent
     private readonly compiler: FFBResultCompilerService, // Add compiler
   ) { }
 
   @Post('query')
   async query(@Body('message') message: string) {
-    // 1. Planner generates AgentQueryPlan
-    const plan = await this.planner.plan(message);
-    // 2. Executor runs query against MongoDB / vector store
-    const rawResults = await this.executor.execute(plan);
-    // 3. Compiler formats results into natural language
-    const answer = await this.compiler.compile(plan, rawResults);
+    // 1. Agent generates pipeline and executes it directly
+    const rawResults = await this.agent.query(message);
 
+    // 2. Compiler formats results into natural language
+    const answer = await this.compiler.compile(rawResults);
     return { answer }; // Return final human-readable response
   }
 
+
   /** Vector search endpoint */
   @Get('search')
   async search(@Query('q') q: string, @Query('k') k?: string) {

+ 3 - 3
src/FFB/ffb-production.module.ts

@@ -3,17 +3,17 @@ import { FFBProductionController } from './ffb-production.controller';
 import { FFBProductionService } from './ffb-production.service';
 import { MongoModule } from 'src/mongo/mongo.module';
 import { FFBVectorService } from './ffb-vector.service';
-import { FFBQueryExecutorService } from './ffb-query-executor.service';
-import { FFBQueryPlannerService } from './ffb-query-planner.service';
 import { FFBResultCompilerService } from './ffb-result-compiler.service';
 import { FFBGateway } from './ffb.gateway';
+import { FFBQueryAgentService } from './ffb-query-agent.service';
+import { GeminiEmbeddingService } from './gemini-embedding.service';
 
 @Module({
   imports: [
     MongoModule
   ],
   controllers: [FFBProductionController],
-  providers: [FFBProductionService, FFBVectorService, FFBQueryExecutorService, FFBQueryPlannerService, FFBResultCompilerService, FFBGateway],
+  providers: [FFBProductionService, FFBVectorService, FFBQueryAgentService, FFBResultCompilerService, GeminiEmbeddingService, FFBGateway],
   exports: [FFBGateway],
 })
 export class FFBProductionModule { }

+ 160 - 0
src/FFB/ffb-query-agent.service.ts

@@ -0,0 +1,160 @@
+import { Injectable, OnModuleInit } from "@nestjs/common";
+import { MongoCoreService } from "src/mongo/mongo-core.service";
+import { FFBProductionRepository } from "src/mongo/mongo-ffb-production.repository";
+import { FFBGateway } from "./ffb.gateway";
+import path from "path";
+import fs from "fs";
+import axios from "axios";
+import jwt from "jsonwebtoken";
+import { GeminiEmbeddingService } from "./gemini-embedding.service";
+
+@Injectable()
+export class FFBQueryAgentService implements OnModuleInit {
+  private systemPrompt: any;
+  private repo: FFBProductionRepository;
+  private readonly VECTOR_DIM = parseInt(process.env.VECTOR_DIM || '3072'); // Gemini default
+
+  private serviceAccount: any; // parsed service account JSON
+  private tokenExpiry: number = 0;
+  private accessToken: string = '';
+
+  constructor(
+    private readonly mongoCore: MongoCoreService,
+    private readonly gateway: FFBGateway,
+    private readonly embeddingService: GeminiEmbeddingService
+  ) { }
+
+  async onModuleInit() {
+    const filePath = path.join(process.cwd(), 'QueryAgent.json');
+    const data = fs.readFileSync(filePath, 'utf-8');
+    this.systemPrompt = JSON.parse(data);
+
+    // Load service account for OAuth
+    const credsPath = process.env.GOOGLE_APPLICATION_CREDENTIALS;
+    if (!credsPath) throw new Error('Missing GOOGLE_APPLICATION_CREDENTIALS');
+    const credsData = fs.readFileSync(credsPath, 'utf-8');
+    this.serviceAccount = JSON.parse(credsData);
+  }
+
+
+
+  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 JSON only in the format:
+{ "textToBeEmbedded": string, "pipeline": [ ... ] }
+
+Q: "${userMessage}"
+`;
+  }
+
+  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 {
+    let sanitized = text
+      .trim()
+      .replace(/^```json\s*/, '')
+      .replace(/^```\s*/, '')
+      .replace(/```$/, '')
+      .trim();
+    sanitized = sanitized.replace(/ISODate\(["'](.+?)["']\)/g, '"$1"');
+    sanitized = sanitized.replace(/"\s*\$/g, '"$');
+    return sanitized;
+  }
+
+  private parseLLMOutput(sanitized: string): { textToBeEmbedded: string; pipeline: any[] } {
+    try {
+      const parsed = JSON.parse(sanitized);
+      if (!('pipeline' in parsed) || !Array.isArray(parsed.pipeline)) {
+        throw new Error('LLM output missing pipeline array');
+      }
+      return parsed;
+    } catch {
+      try {
+        return eval(`(${sanitized})`);
+      } catch (err2) {
+        console.error('Failed to parse LLM output even with fallback:', sanitized);
+        throw new Error('LLM returned invalid JSON');
+      }
+    }
+  }
+
+  private async getRepo(): Promise<FFBProductionRepository> {
+    if (!this.repo) {
+      const db = await this.mongoCore.getDb();
+      this.repo = new FFBProductionRepository(db);
+      await this.repo.init();
+    }
+    return this.repo;
+  }
+
+  /** Main entry point: plan + conditional vector search execution */
+  async query(userMessage: string): Promise<any[]> {
+    const promptText = this.buildPrompt(userMessage);
+    const llmResponse = await this.callGemini(promptText);
+
+    const sanitized = this.sanitizeLLMOutput(llmResponse);
+    const { textToBeEmbedded, pipeline } = this.parseLLMOutput(sanitized);
+
+    if (textToBeEmbedded && textToBeEmbedded.trim()) {
+      const embedding = await this.embeddingService.embedText(textToBeEmbedded.trim());
+      for (const stage of pipeline) {
+        if ('$vectorSearch' in stage) stage.$vectorSearch.queryVector = embedding;
+      }
+    }
+
+    const pipelineForSocket = pipeline.map(stage => ('$vectorSearch' in stage ? { ...stage, $vectorSearch: { ...stage.$vectorSearch, queryVector: '[VECTOR]' } } : stage));
+
+    this.gateway.emitAgentOutput({
+      stage: 'pipeline_generated',
+      rawLLMOutput: llmResponse,
+      pipeline: pipelineForSocket,
+      prettyPipeline: JSON.stringify(pipelineForSocket, null, 2),
+      textToBeEmbedded,
+    });
+
+    const repo = await this.getRepo();
+    const results = await repo.aggregate(pipeline);
+
+    this.gateway.emitAgentOutput({
+      stage: 'pipeline_executed',
+      pipeline: pipelineForSocket,
+      count: results.length,
+      results,
+    });
+
+    return results;
+  }
+
+}

+ 0 - 93
src/FFB/ffb-query-executor.service.ts

@@ -1,93 +0,0 @@
-import { Injectable } from '@nestjs/common';
-import { AgentQueryPlan } from './ffb-agent.types';
-import { FFBVectorService } from './ffb-vector.service';
-import { MongoCoreService } from 'src/mongo/mongo-core.service';
-import { FFBProductionRepository } from 'src/mongo/mongo-ffb-production.repository';
-import { FFBGateway } from './ffb.gateway';
-
-@Injectable()
-export class FFBQueryExecutorService {
-    private repo: FFBProductionRepository;
-
-    constructor(
-        private readonly mongoCore: MongoCoreService,
-        private readonly vectorService: FFBVectorService,
-        private readonly gateway: FFBGateway,
-    ) { }
-
-    private async getRepo() {
-        if (!this.repo) {
-            const db = await this.mongoCore.getDb();
-            this.repo = new FFBProductionRepository(db);
-            await this.repo.init();
-        }
-        return this.repo;
-    }
-
-    /** Convert any string dates in preFilter to Date objects */
-    private convertDates(preFilter?: Record<string, any>) {
-        if (!preFilter) return;
-        const pd = preFilter.productionDate;
-        if (pd) {
-            if (pd.$gte) pd.$gte = new Date(pd.$gte);
-            if (pd.$lte) pd.$lte = new Date(pd.$lte);
-        }
-    }
-
-    async execute(plan: AgentQueryPlan) {
-        const repo = await this.getRepo();
-
-        // Convert string dates to Date objects
-        this.convertDates(plan.preFilter);
-
-        // Build $project if fields are specified
-        let postPipeline: Array<Record<string, any>> = plan.postPipeline ?? [];
-        if (plan.fields && plan.fields.length > 0) {
-            postPipeline = [
-                { $project: plan.fields.reduce((acc, f) => ({ ...acc, [f]: 1 }), { _id: 0 }) },
-                ...postPipeline,
-            ];
-        }
-
-        // Construct the aggregation pipeline
-        let pipeline: Array<Record<string, any>> = [];
-
-        if (plan.vectorQuery) {
-            const vector = await this.vectorService['embedText'](plan.vectorQuery);
-
-            pipeline = [
-                {
-                    $vectorSearch: {
-                        index: 'vector_index',
-                        path: 'vector',
-                        queryVector: vector,
-                        filter: plan.preFilter,
-                        limit: plan.vectorOptions?.limit ?? 5,
-                        numCandidates: plan.vectorOptions?.numCandidates ?? 50,
-                    },
-                },
-                ...postPipeline,
-            ];
-        } else {
-            // Pure aggregation
-            pipeline = [
-                ...(plan.preFilter ? [{ $match: plan.preFilter }] : []),
-                ...postPipeline,
-            ];
-        }
-
-        // console.log('--- Aggregation pipeline ---\n', JSON.stringify(pipeline, null, 2));
-
-        const results = await repo.aggregate(pipeline);
-        console.log('--- Raw results ---\n', results);
-
-        this.gateway.emitExecutorResult({
-            receivedPlan: plan,
-            pipeline: pipeline,
-            count: results.length,
-            results,
-        });
-
-        return results;
-    }
-}

+ 0 - 113
src/FFB/ffb-query-planner.service.ts

@@ -1,113 +0,0 @@
-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';
-import { FFBGateway } from './ffb.gateway';
-
-@Injectable()
-export class FFBQueryPlannerService implements OnModuleInit {
-  private systemPrompt: any;
-
-  constructor(private readonly gateway: FFBGateway,) {
-    // Logic here 
-  }
-
-  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);
-
-    this.gateway.emitPlannerOutput({
-      stage: 'raw_llm_output',
-      content: responseText,
-    });
-
-    try {
-      const plan = JSON.parse(sanitized);
-
-      this.gateway.emitPlannerOutput({
-        stage: 'parsed_plan',
-        plan,
-      });
-
-      return plan;
-    } catch (err) {
-      console.error('Failed to parse Gemini output:', sanitized);
-      this.gateway.emitError({
-        source: 'planner',
-        error: 'Invalid JSON returned by LLM',
-        raw: 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();
-  }
-}

+ 6 - 6
src/FFB/ffb-result-compiler.service.ts

@@ -14,18 +14,18 @@ export class FFBResultCompilerService implements OnModuleInit {
     this.systemPrompt = JSON.parse(data);
   }
 
-  async compile(plan: AgentQueryPlan, rawResults: any[]): Promise<string> {
-    const promptText = this.buildPrompt(plan, rawResults);
+  async compile(rawResults: any[]): Promise<string> {
+    const promptText = this.buildPrompt(rawResults);
     const response = await this.callGemini(promptText);
     return response.trim();
   }
 
-  private buildPrompt(plan: AgentQueryPlan, results: any[]): string {
+  private buildPrompt(results: any[]): string {
     // Build example text for LLM
     const examplesText = (this.systemPrompt.examples || [])
       .map(
         (ex: any) =>
-          `Plan: ${JSON.stringify(ex.plan, null, 2)}\nResults: ${JSON.stringify(ex.results, null, 2)}\nAnswer: ${ex.answer}`
+          `Results: ${JSON.stringify(ex.results, null, 2)}\nAnswer: ${ex.answer}`
       )
       .join('\n\n');
 
@@ -34,8 +34,8 @@ ${this.systemPrompt.instructions}
 
 ${examplesText}
 
-Here is the actual AgentQueryPlan: ${JSON.stringify(plan, null, 2)}
 Here are the raw results from MongoDB: ${JSON.stringify(results, null, 2)}
+Please produce a clear, concise natural language answer for the user.
 `;
   }
 
@@ -69,4 +69,4 @@ Here are the raw results from MongoDB: ${JSON.stringify(results, null, 2)}
       throw err;
     }
   }
-}
+}

+ 8 - 66
src/FFB/ffb-vector.service.ts

@@ -1,17 +1,17 @@
 import { Injectable, OnModuleInit } from '@nestjs/common';
-import axios from 'axios';
-import { GoogleAuth } from 'google-auth-library';
 import { MongoCoreService } from 'src/mongo/mongo-core.service';
 import { FFBProductionRepository } from 'src/mongo/mongo-ffb-production.repository';
 import { FFBProduction } from './ffb-production.schema';
+import { GeminiEmbeddingService } from './gemini-embedding.service';
 
 @Injectable()
 export class FFBVectorService implements OnModuleInit {
   private repo: FFBProductionRepository;
-  private readonly VECTOR_DIM = parseInt(process.env.VECTOR_DIM || '3072'); // Gemini default
-  private accessToken: string; // OAuth token
 
-  constructor(private readonly mongoCore: MongoCoreService) { }
+  constructor(
+    private readonly mongoCore: MongoCoreService,
+    private readonly embeddingService: GeminiEmbeddingService
+  ) {}
 
   async onModuleInit() {
     // Initialize Mongo repository
@@ -19,19 +19,6 @@ export class FFBVectorService implements OnModuleInit {
     this.repo = new FFBProductionRepository(db);
     await this.repo.init();
 
-    // Authenticate with Google service account
-    const auth = new GoogleAuth({
-      keyFile: process.env.GOOGLE_APPLICATION_CREDENTIALS,
-      scopes: ['https://www.googleapis.com/auth/cloud-platform'],
-    });
-
-    const client = await auth.getClient();
-    const tokenResponse = await client.getAccessToken();
-    if (!tokenResponse.token) {
-      throw new Error('Failed to obtain OAuth 2 access token for Gemini embeddings');
-    }
-    this.accessToken = tokenResponse.token;
-
     console.log('✅ Gemini embedding service ready. Repository initialized.');
   }
 
@@ -40,56 +27,10 @@ export class FFBVectorService implements OnModuleInit {
     return `Production on ${new Date(record.productionDate).toISOString()} at ${record.site} in ${record.phase} ${record.block} produced ${record.quantity} ${record.quantityUom} with a total weight of ${record.weight} ${record.weightUom}.`;
   }
 
-  /** Generate embedding via Google Gemini */
-  private async embedText(text: string): Promise<number[]> {
-    const project = process.env.GOOGLE_PROJECT_ID!;
-    const model = process.env.EMBEDDING_MODEL || 'gemini-embedding-001';
-    const location = 'us-central1';
-
-    const url = `https://${location}-aiplatform.googleapis.com/v1/projects/${project}/locations/${location}/publishers/google/models/${model}:predict`;
-
-    try {
-      const response = await axios.post(
-        url,
-        {
-          instances: [{ content: text }],
-          parameters: { autoTruncate: true, outputDimensionality: this.VECTOR_DIM },
-        },
-        {
-          headers: {
-            Authorization: `Bearer ${this.accessToken}`,
-            'Content-Type': 'application/json',
-          },
-        },
-      );
-
-      const predictions = response.data?.predictions;
-      if (!predictions || predictions.length === 0) {
-        throw new Error(`No predictions returned from Gemini API: ${JSON.stringify(response.data)}`);
-      }
-
-      const embedding = predictions[0]?.embeddings?.values;
-      if (!embedding || !Array.isArray(embedding)) {
-        throw new Error(`Invalid embedding format returned: ${JSON.stringify(predictions[0])}`);
-      }
-
-      if (embedding.length !== this.VECTOR_DIM) {
-        console.warn(
-          `⚠️ Warning: embedding dimension mismatch. Expected ${this.VECTOR_DIM}, got ${embedding.length}`,
-        );
-      }
-
-      return embedding;
-    } catch (err: any) {
-      console.error('Failed to generate embedding:', err.response?.data || err.message);
-      throw err;
-    }
-  }
-
   /** Insert a single record with embedding vector */
   async insertWithVector(record: FFBProduction) {
     const text = this.recordToText(record);
-    const vector = await this.embedText(text);
+    const vector = await this.embeddingService.embedText(text);
 
     const data: FFBProduction & { vector: number[] } = { ...record, vector };
     return this.repo.create(data);
@@ -99,8 +40,9 @@ export class FFBVectorService implements OnModuleInit {
   async vectorSearch(query: string, k = 5) {
     if (!query) throw new Error('Query string cannot be empty');
 
-    const vector = await this.embedText(query);
+    const vector = await this.embeddingService.embedText(query);
     const results = await this.repo.vectorSearch(vector, k, 50);
+
     return results.map((r) => ({
       ...r,
       _id: r._id.toString(),

+ 6 - 13
src/FFB/ffb.gateway.ts

@@ -1,27 +1,20 @@
-import {
-  WebSocketGateway,
-  WebSocketServer,
-} from '@nestjs/websockets';
+import { WebSocketGateway, WebSocketServer } from '@nestjs/websockets';
 import { Server } from 'socket.io';
 
 @WebSocketGateway({
   namespace: '/ffb',
-  cors: {
-    origin: '*',
-  },
+  cors: { origin: '*' },
 })
 export class FFBGateway {
   @WebSocketServer()
   server: Server;
 
-  emitPlannerOutput(payload: any) {
-    this.server.emit('planner.output', payload);
-  }
-
-  emitExecutorResult(payload: any) {
-    this.server.emit('executor.result', payload);
+  /** Emits any intermediate output from the Agent (planning + execution) */
+  emitAgentOutput(payload: any) {
+    this.server.emit('agent.output', payload);
   }
 
+  /** Emits any errors that occur during agent execution */
   emitError(payload: any) {
     this.server.emit('error', payload);
   }

+ 76 - 0
src/FFB/gemini-embedding.service.ts

@@ -0,0 +1,76 @@
+import { Injectable, OnModuleInit } from '@nestjs/common';
+import fs from 'fs';
+import path from 'path';
+import axios from 'axios';
+import jwt from 'jsonwebtoken';
+
+@Injectable()
+export class GeminiEmbeddingService implements OnModuleInit {
+  private serviceAccount: any;
+  private tokenExpiry: number = 0;
+  private accessToken: string = '';
+  private readonly VECTOR_DIM = parseInt(process.env.VECTOR_DIM || '3072');
+
+  async onModuleInit() {
+    const credsPath = process.env.GOOGLE_APPLICATION_CREDENTIALS;
+    if (!credsPath) throw new Error('Missing GOOGLE_APPLICATION_CREDENTIALS');
+    const credsData = fs.readFileSync(credsPath, 'utf-8');
+    this.serviceAccount = JSON.parse(credsData);
+  }
+
+  /** Get OAuth2 access token dynamically, refresh if expired */
+  async getAccessToken(): Promise<string> {
+    const now = Math.floor(Date.now() / 1000);
+    if (this.accessToken && now < this.tokenExpiry - 60) return this.accessToken;
+
+    const iat = now;
+    const exp = iat + 3600;
+    const payload = {
+      iss: this.serviceAccount.client_email,
+      sub: this.serviceAccount.client_email,
+      aud: 'https://oauth2.googleapis.com/token',
+      iat,
+      exp,
+      scope: 'https://www.googleapis.com/auth/cloud-platform',
+    };
+
+    const signedJWT = jwt.sign(payload, this.serviceAccount.private_key, { algorithm: 'RS256' });
+
+    const resp = await axios.post(
+      'https://oauth2.googleapis.com/token',
+      new URLSearchParams({
+        grant_type: 'urn:ietf:params:oauth:grant-type:jwt-bearer',
+        assertion: signedJWT,
+      }).toString(),
+      { headers: { 'Content-Type': 'application/x-www-form-urlencoded' } }
+    );
+
+    this.accessToken = resp.data.access_token;
+    this.tokenExpiry = now + resp.data.expires_in;
+    return this.accessToken;
+  }
+
+  /** Embed a text string using Gemini embeddings */
+  async embedText(text: string): Promise<number[]> {
+    if (!text) return [];
+
+    const project = process.env.GOOGLE_PROJECT_ID!;
+    const model = process.env.EMBEDDING_MODEL || 'gemini-embedding-001';
+    const location = 'us-central1';
+
+    const token = await this.getAccessToken();
+
+    const url = `https://${location}-aiplatform.googleapis.com/v1/projects/${project}/locations/${location}/publishers/google/models/${model}:predict`;
+
+    const response = await axios.post(
+      url,
+      { instances: [{ content: text }], parameters: { autoTruncate: true, outputDimensionality: this.VECTOR_DIM } },
+      { headers: { Authorization: `Bearer ${token}`, 'Content-Type': 'application/json' } }
+    );
+
+    const predictions = response.data?.predictions;
+    const embedding = predictions?.[0]?.embeddings?.values;
+    if (!embedding || !Array.isArray(embedding)) throw new Error('Invalid embedding returned');
+    return embedding;
+  }
+}