|
|
@@ -1,44 +1,89 @@
|
|
|
import { Injectable, OnModuleInit } from '@nestjs/common';
|
|
|
-import { pipeline, FeatureExtractionPipeline } from '@xenova/transformers';
|
|
|
+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';
|
|
|
|
|
|
@Injectable()
|
|
|
export class FFBVectorService implements OnModuleInit {
|
|
|
- private embedder: FeatureExtractionPipeline;
|
|
|
private repo: FFBProductionRepository;
|
|
|
- private readonly VECTOR_DIM = 384; // must match your index
|
|
|
+ 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) {}
|
|
|
|
|
|
- /** Initialize model and repository at module startup */
|
|
|
async onModuleInit() {
|
|
|
- const modelName = process.env.EMBEDDING_MODEL || 'Xenova/all-MiniLM-L6-v2';
|
|
|
- console.log(`🔹 Loading embedding model: ${modelName}...`);
|
|
|
-
|
|
|
- this.embedder = (await pipeline('feature-extraction', modelName)) as unknown as FeatureExtractionPipeline;
|
|
|
-
|
|
|
+ // Initialize Mongo repository
|
|
|
const db = await this.mongoCore.getDb();
|
|
|
this.repo = new FFBProductionRepository(db);
|
|
|
await this.repo.init();
|
|
|
|
|
|
- console.log(`✅ Embedding model loaded and repository ready.`);
|
|
|
+ // 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.');
|
|
|
}
|
|
|
|
|
|
- /** Convert an FFBProduction record into text for embedding */
|
|
|
+ /** Convert a record to a string suitable for embedding */
|
|
|
private recordToText(record: FFBProduction): string {
|
|
|
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 vector from text */
|
|
|
+ /** Generate embedding via Google Gemini */
|
|
|
private async embedText(text: string): Promise<number[]> {
|
|
|
- const output = await this.embedder(text, { pooling: 'mean', normalize: true });
|
|
|
- const vector = Array.from(output.data);
|
|
|
- if (vector.length !== this.VECTOR_DIM) {
|
|
|
- throw new Error(`Embedding dimension mismatch. Expected ${this.VECTOR_DIM}, got ${vector.length}`);
|
|
|
+ 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;
|
|
|
}
|
|
|
- return vector;
|
|
|
}
|
|
|
|
|
|
/** Insert a single record with embedding vector */
|
|
|
@@ -46,7 +91,6 @@ export class FFBVectorService implements OnModuleInit {
|
|
|
const text = this.recordToText(record);
|
|
|
const vector = await this.embedText(text);
|
|
|
|
|
|
- // Explicitly tell TypeScript this object matches the repository type
|
|
|
const data: FFBProduction & { vector: number[] } = { ...record, vector };
|
|
|
return this.repo.create(data);
|
|
|
}
|
|
|
@@ -55,17 +99,12 @@ export class FFBVectorService implements OnModuleInit {
|
|
|
async vectorSearch(query: string, k = 5) {
|
|
|
if (!query) throw new Error('Query string cannot be empty');
|
|
|
|
|
|
- // Step 1: Embed the query text
|
|
|
const vector = await this.embedText(query);
|
|
|
-
|
|
|
- // Step 2: Use repository aggregation for vector search
|
|
|
- const results = await this.repo.vectorSearch(vector, k, 50); // numCandidates = 50
|
|
|
- // Step 3: Return results directly (they now include the full document + score)
|
|
|
- return results.map(r => ({
|
|
|
- ...r, // all FFBProduction fields
|
|
|
- _id: r._id.toString(), // convert ObjectId to string if needed
|
|
|
- score: r.score // similarity score
|
|
|
+ const results = await this.repo.vectorSearch(vector, k, 50);
|
|
|
+ return results.map((r) => ({
|
|
|
+ ...r,
|
|
|
+ _id: r._id.toString(),
|
|
|
+ score: r.score,
|
|
|
}));
|
|
|
}
|
|
|
-
|
|
|
}
|
Dr-Swopt