@antl3x/toolrag
Version:
Context-aware tool retrieval for language models - unlock the full potential of LLM function calling without context window limitations or constraints.
117 lines • 4.9 kB
JavaScript
;
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
var desc = Object.getOwnPropertyDescriptor(m, k);
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
desc = { enumerable: true, get: function() { return m[k]; } };
}
Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
o["default"] = v;
});
var __importStar = (this && this.__importStar) || (function () {
var ownKeys = function(o) {
ownKeys = Object.getOwnPropertyNames || function (o) {
var ar = [];
for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
return ar;
};
return ownKeys(o);
};
return function (mod) {
if (mod && mod.__esModule) return mod;
var result = {};
if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
__setModuleDefault(result, mod);
return result;
};
})();
Object.defineProperty(exports, "__esModule", { value: true });
exports.EmbeddingProviderGoogle = void 0;
const aiplatform = __importStar(require("@google-cloud/aiplatform"));
const zod_1 = require("zod");
const EmbeddingProviderGoogleConfigSchema = zod_1.z
.object({
model: zod_1.z.string().default('text-embedding-005'),
projectId: zod_1.z
.string()
.optional()
.refine((val) => val !== undefined || process.env.GOOGLE_CLOUD_PROJECT !== undefined, 'Google project ID is required'),
location: zod_1.z.string().default('us-central1'),
dimensions: zod_1.z.number().default(768),
taskType: zod_1.z.string().default('RETRIEVAL_QUERY'),
client: zod_1.z.instanceof(aiplatform.v1.PredictionServiceClient).optional(),
})
.default({});
/**
* Google implementation of the EmbeddingProvider
*/
class EmbeddingProviderGoogle {
_config;
_client;
_endpoint;
constructor(config) {
this._config = EmbeddingProviderGoogleConfigSchema.parse(config);
this._config.projectId = this._config.projectId || process.env.GOOGLE_CLOUD_PROJECT || '';
const clientOptions = {
apiEndpoint: 'us-central1-aiplatform.googleapis.com',
};
this._client = this._config.client || new aiplatform.v1.PredictionServiceClient(clientOptions);
this._endpoint = `projects/${this._config.projectId}/locations/${this._config.location}/publishers/google/models/${this._config.model}`;
}
async getEmbedding(text) {
const toValueHelper = aiplatform.helpers?.toValue;
if (!toValueHelper) {
throw new Error('Helpers not available in Google Cloud AI Platform client');
}
const instances = [
toValueHelper({
content: text,
task_type: this._config.taskType,
}),
];
const parameters = toValueHelper(this._config.dimensions > 0 ? { outputDimensionality: this._config.dimensions } : {});
const request = {
endpoint: this._endpoint,
instances,
parameters,
};
const [response] = await this._client.predict(request);
const predictions = response.predictions;
if (!predictions || predictions.length === 0) {
throw new Error('No embedding returned from Google API');
}
const prediction = predictions[0];
if (!prediction.structValue || !prediction.structValue.fields) {
throw new Error('Invalid embedding response structure');
}
const embeddingsProto = prediction.structValue.fields.embeddings;
if (!embeddingsProto || !embeddingsProto.structValue || !embeddingsProto.structValue.fields) {
throw new Error('Invalid embedding response structure');
}
const valuesProto = embeddingsProto.structValue.fields.values;
if (!valuesProto || !valuesProto.listValue || !valuesProto.listValue.values) {
throw new Error('Invalid embedding response structure');
}
return valuesProto.listValue.values
.map((v) => (typeof v.numberValue === 'number' ? v.numberValue : 0))
.filter((v) => v !== null && v !== undefined);
}
getDimensions() {
return this._config.dimensions;
}
getName() {
return `google`;
}
getModel() {
return this._config.model;
}
}
exports.EmbeddingProviderGoogle = EmbeddingProviderGoogle;
//# sourceMappingURL=EmbeddingProviderGoogle.js.map