rhombus-node-mcp 0.1.15 → 0.1.16
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +9 -1
- package/dist/api/camera-tool-api.js +108 -0
- package/dist/api/clips-tool-api.js +36 -0
- package/dist/api/create-camera-policy-tool-api.js +9 -0
- package/dist/api/create-tool-api.js +53 -0
- package/dist/api/entity-lookup-tool-api.js +37 -0
- package/dist/api/events-tool-api.js +320 -0
- package/dist/api/faces-tool-api.js +110 -0
- package/dist/api/get-entity-tool-api.js +176 -0
- package/dist/api/get-org-information-tool-api.js +9 -0
- package/dist/api/location-tool-api.js +9 -0
- package/dist/api/lpr-tool-api.js +68 -0
- package/dist/api/policy-alerts-tool-api.js +42 -0
- package/dist/api/reboot-cameras-tool-api.js +34 -0
- package/dist/api/report-tool-api.js +426 -0
- package/dist/api/time-tool-api.js +90 -0
- package/dist/api/update-tool-api.js +148 -0
- package/dist/createServer.js +7 -1
- package/dist/disabled-tools/endpoint-to-keys-tool.js +84 -0
- package/dist/disabled-tools/semantic-search-tool.js +90 -0
- package/dist/index.js +3 -5
- package/dist/logger.js +4 -3
- package/dist/network.js +42 -15
- package/dist/resources/routes.json.js +1 -1
- package/dist/services/embedding-service.js +153 -0
- package/dist/services/faiss-search-service.js +261 -0
- package/dist/tools/camera-tool.js +100 -0
- package/dist/tools/clips-tool.js +23 -37
- package/dist/tools/count-tool.js +25 -0
- package/dist/tools/create-camera-policy-tool.js +214 -0
- package/dist/tools/create-tool.js +25 -74
- package/dist/tools/entity-lookup-tool.js +35 -0
- package/dist/tools/events-tool.js +188 -92
- package/dist/tools/faces-tool.js +59 -132
- package/dist/tools/get-entity-tool.js +78 -0
- package/dist/tools/get-org-information-tool.js +18 -0
- package/dist/tools/location-tool.js +24 -31
- package/dist/tools/lpr-tool.js +71 -0
- package/dist/tools/policy-alerts-tool.js +32 -42
- package/dist/tools/reboot-cameras-tool.js +34 -0
- package/dist/tools/report-tool.js +190 -0
- package/dist/tools/time-conversion-tool.js +43 -0
- package/dist/tools/time-tool.js +17 -57
- package/dist/tools/update-tool.js +262 -0
- package/dist/transports/streamable-http.js +160 -57
- package/dist/{tools/devices/camera-tool/types.js → types/camera-tool-types.js} +13 -0
- package/dist/types/clips-tool-types.js +41 -0
- package/dist/types/create-camera-policy-tool-types.js +44 -0
- package/dist/types/create-tool-types.js +8 -0
- package/dist/types/deviceType.js +1 -0
- package/dist/types/endpoint-to-keys-tool-types.js +7 -0
- package/dist/types/entity-lookup-tool-types.js +70 -0
- package/dist/types/events-tools-types.js +257 -0
- package/dist/types/faces-tools-types.js +143 -0
- package/dist/types/get-entity-tool-types.js +25 -0
- package/dist/types/get-org-information-tool-types.js +3 -0
- package/dist/types/location-tool-types.js +11 -0
- package/dist/types/lpr-tool-types.js +97 -0
- package/dist/types/policy-alerts-tool-types.js +74 -0
- package/dist/types/reboot-cameras-tool-types.js +8 -0
- package/dist/types/report-tool-types.js +268 -0
- package/dist/types/schema-components.js +7093 -0
- package/dist/types/schema.js +1 -0
- package/dist/types/semantic-search-tool-types.js +5 -0
- package/dist/types/time-conversion-tool-types.js +8 -0
- package/dist/types/time-tool-types.js +11 -0
- package/dist/types/update-tool-types.js +186 -0
- package/dist/types/zod-schemas.js +21315 -0
- package/dist/types.js +17 -7
- package/dist/util.js +94 -2
- package/dist/utils/confirmation.js +1 -1
- package/dist/utils/reduce-output.js +35 -0
- package/dist/utils/remove-nulls.js +28 -0
- package/dist/utils/temp.js +8 -0
- package/dist/utils/timestampInput.js +12 -0
- package/package.json +23 -3
- package/dist/tools/devices/camera-tool/camera-tool.js +0 -218
- package/dist/tools/devices/get-entity-tool.js +0 -118
- package/dist/tools/get-org-information.js +0 -17
- package/dist/tools/reboot-cameras.js +0 -62
|
@@ -0,0 +1,84 @@
|
|
|
1
|
+
import { createToolTextContent } from "../util.js";
|
|
2
|
+
import { TOOL_ARGS } from "../types/endpoint-to-keys-tool-types.js";
|
|
3
|
+
import * as fs from 'fs';
|
|
4
|
+
import * as path from 'path';
|
|
5
|
+
import { fileURLToPath } from 'url';
|
|
6
|
+
const __filename = fileURLToPath(import.meta.url);
|
|
7
|
+
const __dirname = path.dirname(__filename);
|
|
8
|
+
const TOOL_NAME = "endpoint-to-keys-tool";
|
|
9
|
+
const TOOL_DESCRIPTION = `
|
|
10
|
+
Returns the output keys for a given API endpoint.
|
|
11
|
+
This tool helps you understand what fields are available in the response of any API endpoint.
|
|
12
|
+
Use this to determine which fields to include when using other tools that accept includeFields parameters.
|
|
13
|
+
|
|
14
|
+
Example usage:
|
|
15
|
+
- Input: "POST /api/camera/getMinimalCameraStateList"
|
|
16
|
+
- Output: ["wifiSignalStrength", "uuid", "name", "locationUuid", ...]
|
|
17
|
+
|
|
18
|
+
Always use this tool first when you need to know what fields are available for any API endpoint.
|
|
19
|
+
`;
|
|
20
|
+
// Load the routes-to-output-keys.json data
|
|
21
|
+
let routesData = {};
|
|
22
|
+
function loadRoutesData() {
|
|
23
|
+
try {
|
|
24
|
+
// Try multiple possible locations for the routes data
|
|
25
|
+
const possiblePaths = [
|
|
26
|
+
path.join(process.cwd(), 'assets', 'routes-to-output-keys.json'),
|
|
27
|
+
path.join(process.cwd(), 'routes-to-output-keys.json'),
|
|
28
|
+
path.join(__dirname, '..', '..', 'assets', 'routes-to-output-keys.json'),
|
|
29
|
+
];
|
|
30
|
+
let fileContent = '';
|
|
31
|
+
let usedPath = '';
|
|
32
|
+
for (const routesPath of possiblePaths) {
|
|
33
|
+
try {
|
|
34
|
+
fileContent = fs.readFileSync(routesPath, 'utf8');
|
|
35
|
+
usedPath = routesPath;
|
|
36
|
+
break;
|
|
37
|
+
}
|
|
38
|
+
catch (e) {
|
|
39
|
+
// Continue to next path
|
|
40
|
+
}
|
|
41
|
+
}
|
|
42
|
+
if (!fileContent) {
|
|
43
|
+
throw new Error(`Could not find routes-to-output-keys.json in any of: ${possiblePaths.join(', ')}`);
|
|
44
|
+
}
|
|
45
|
+
const data = JSON.parse(fileContent);
|
|
46
|
+
routesData = data.routes || {};
|
|
47
|
+
console.log(`Loaded ${Object.keys(routesData).length} route endpoints from ${usedPath}`);
|
|
48
|
+
}
|
|
49
|
+
catch (error) {
|
|
50
|
+
console.error('Failed to load routes-to-output-keys.json:', error);
|
|
51
|
+
routesData = {};
|
|
52
|
+
}
|
|
53
|
+
}
|
|
54
|
+
// Load data on module initialization
|
|
55
|
+
loadRoutesData();
|
|
56
|
+
const TOOL_HANDLER = async (args, extra) => {
|
|
57
|
+
const { endpoint } = args;
|
|
58
|
+
// Normalize the endpoint format
|
|
59
|
+
const normalizedEndpoint = endpoint.trim();
|
|
60
|
+
// Find the matching endpoint
|
|
61
|
+
const outputKeys = routesData[normalizedEndpoint];
|
|
62
|
+
if (!outputKeys) {
|
|
63
|
+
// Try to find a partial match or suggest similar endpoints
|
|
64
|
+
const availableEndpoints = Object.keys(routesData);
|
|
65
|
+
const similarEndpoints = availableEndpoints.filter(ep => ep.toLowerCase().includes(endpoint.toLowerCase()) ||
|
|
66
|
+
endpoint.toLowerCase().includes(ep.toLowerCase()));
|
|
67
|
+
return createToolTextContent(JSON.stringify({
|
|
68
|
+
error: "Endpoint not found",
|
|
69
|
+
message: `No output keys found for endpoint: ${endpoint}`,
|
|
70
|
+
suggestion: similarEndpoints.length > 0
|
|
71
|
+
? `Similar endpoints found: ${similarEndpoints.slice(0, 5).join(', ')}`
|
|
72
|
+
: "Check the endpoint format. It should be like 'POST /api/camera/getMinimalCameraStateList'",
|
|
73
|
+
availableEndpointsCount: availableEndpoints.length
|
|
74
|
+
}, null, 2));
|
|
75
|
+
}
|
|
76
|
+
return createToolTextContent(JSON.stringify({
|
|
77
|
+
endpoint: normalizedEndpoint,
|
|
78
|
+
outputKeys: outputKeys,
|
|
79
|
+
totalKeys: outputKeys.length
|
|
80
|
+
}, null, 2));
|
|
81
|
+
};
|
|
82
|
+
export function createTool(server) {
|
|
83
|
+
server.tool(TOOL_NAME, TOOL_DESCRIPTION, TOOL_ARGS, TOOL_HANDLER);
|
|
84
|
+
}
|
|
@@ -0,0 +1,90 @@
|
|
|
1
|
+
import { FaissSearchService } from '../services/faiss-search-service.js';
|
|
2
|
+
import { createToolTextContent } from '../util.js';
|
|
3
|
+
import { logger } from '../logger.js';
|
|
4
|
+
import { TOOL_ARGS } from '../types/semantic-search-tool-types.js';
|
|
5
|
+
// Constants following coding style guidelines
|
|
6
|
+
const DEFAULT_SEARCH_LIMIT = 3;
|
|
7
|
+
const DEFAULT_MIN_SIMILARITY = 0.3;
|
|
8
|
+
const TOOL_NAME = "semantic-search";
|
|
9
|
+
const TOOL_DESCRIPTION = `
|
|
10
|
+
This tool performs semantic search across the Rhombus knowledge base using AI embeddings and vector similarity.
|
|
11
|
+
|
|
12
|
+
It searches through documentation, blog posts, support articles, and other content to find the most relevant
|
|
13
|
+
information based on the meaning of your query, not just keyword matching.
|
|
14
|
+
|
|
15
|
+
The tool returns an array of search results, each containing:
|
|
16
|
+
- sourceTitle: The title of the source document
|
|
17
|
+
- sourceUrl: The URL where the content can be found
|
|
18
|
+
- chunkText: The relevant text snippet that matches your query
|
|
19
|
+
|
|
20
|
+
Use this tool when you need to find information about:
|
|
21
|
+
- Camera features, specifications, or comparisons
|
|
22
|
+
- Installation and setup procedures
|
|
23
|
+
- Technical troubleshooting
|
|
24
|
+
- Product capabilities and benefits
|
|
25
|
+
- Access control and security features
|
|
26
|
+
- Cloud storage and management
|
|
27
|
+
- Any other Rhombus-related topics
|
|
28
|
+
`;
|
|
29
|
+
// Singleton instance for reuse across searches
|
|
30
|
+
let globalSearchService = null;
|
|
31
|
+
let isGlobalServiceInitialized = false;
|
|
32
|
+
async function getSearchService() {
|
|
33
|
+
if (!globalSearchService) {
|
|
34
|
+
globalSearchService = new FaissSearchService();
|
|
35
|
+
}
|
|
36
|
+
if (!isGlobalServiceInitialized) {
|
|
37
|
+
await globalSearchService.loadEmbeddings();
|
|
38
|
+
await globalSearchService.loadIndex();
|
|
39
|
+
isGlobalServiceInitialized = true;
|
|
40
|
+
}
|
|
41
|
+
return globalSearchService;
|
|
42
|
+
}
|
|
43
|
+
const TOOL_HANDLER = async (args, extra) => {
|
|
44
|
+
const { query } = args;
|
|
45
|
+
try {
|
|
46
|
+
if (!query || query.trim().length === 0) {
|
|
47
|
+
return createToolTextContent(JSON.stringify({ error: "Query cannot be empty" }));
|
|
48
|
+
}
|
|
49
|
+
const searchService = await getSearchService();
|
|
50
|
+
const searchOptions = {
|
|
51
|
+
query: query.trim(),
|
|
52
|
+
limit: DEFAULT_SEARCH_LIMIT,
|
|
53
|
+
minSimilarity: DEFAULT_MIN_SIMILARITY,
|
|
54
|
+
urlPattern: undefined,
|
|
55
|
+
};
|
|
56
|
+
const { results, stats } = await searchService.search(searchOptions);
|
|
57
|
+
const searchResults = results.map(result => ({
|
|
58
|
+
sourceTitle: result.sourceTitle,
|
|
59
|
+
sourceUrl: result.sourceUrl,
|
|
60
|
+
chunkText: result.chunkText,
|
|
61
|
+
}));
|
|
62
|
+
logger.info('Semantic search results:', results);
|
|
63
|
+
// Log search statistics for debugging
|
|
64
|
+
logger.info('Semantic search stats:', {
|
|
65
|
+
searchTime: stats.searchTime,
|
|
66
|
+
totalVectors: stats.totalVectors,
|
|
67
|
+
resultsReturned: stats.resultsReturned,
|
|
68
|
+
resultsFiltered: stats.resultsFiltered,
|
|
69
|
+
query: query.trim()
|
|
70
|
+
});
|
|
71
|
+
return createToolTextContent(JSON.stringify(searchResults));
|
|
72
|
+
}
|
|
73
|
+
catch (error) {
|
|
74
|
+
return createToolTextContent(JSON.stringify({
|
|
75
|
+
error: `Search failed: ${error instanceof Error ? error.message : 'Unknown error'}`
|
|
76
|
+
}));
|
|
77
|
+
}
|
|
78
|
+
};
|
|
79
|
+
export async function createTool(server) {
|
|
80
|
+
try {
|
|
81
|
+
// Verify search service can be initialized and also save it into the cached variable
|
|
82
|
+
await getSearchService();
|
|
83
|
+
server.tool(TOOL_NAME, TOOL_DESCRIPTION, TOOL_ARGS, TOOL_HANDLER);
|
|
84
|
+
}
|
|
85
|
+
catch (error) {
|
|
86
|
+
logger.error('Failed to initialize semantic search tool:', error);
|
|
87
|
+
logger.error('Make sure you have the index and embeddings files in rhombus-node-mcp/data/embeddings/');
|
|
88
|
+
throw error;
|
|
89
|
+
}
|
|
90
|
+
}
|
package/dist/index.js
CHANGED
|
@@ -8,12 +8,10 @@ import streamableHttpTransport from "./transports/streamable-http.js";
|
|
|
8
8
|
const RHOMBUS_API_KEY = process.env.RHOMBUS_API_KEY;
|
|
9
9
|
const TRANSPORT_TYPE = process.env.TRANSPORT_TYPE || "stdio";
|
|
10
10
|
async function main() {
|
|
11
|
-
if (!RHOMBUS_API_KEY) {
|
|
12
|
-
logger.info("Missing RHOMBUS_API_KEY");
|
|
13
|
-
}
|
|
14
11
|
const serverUrl = process.env.RHOMBUS_API_SERVER || "api2.rhombussystems.com";
|
|
15
|
-
|
|
16
|
-
|
|
12
|
+
if (RHOMBUS_API_KEY) {
|
|
13
|
+
logger.info(`🔑 Using API_KEY: ${RHOMBUS_API_KEY}`);
|
|
14
|
+
}
|
|
17
15
|
logger.info("🌐 Using server url", serverUrl);
|
|
18
16
|
await serverInit();
|
|
19
17
|
const server = new McpServer({
|
package/dist/logger.js
CHANGED
|
@@ -4,21 +4,22 @@ log4js.configure({
|
|
|
4
4
|
appenders: {
|
|
5
5
|
mcp: {
|
|
6
6
|
type: "file",
|
|
7
|
-
filename: path.resolve(process.env.LOG_FOLDER ?? process.cwd(), "rhombus-node-mcp.log"),
|
|
7
|
+
filename: path.resolve(process.env.LOG_FOLDER ?? process.cwd(), "./logs/rhombus-node-mcp.log"),
|
|
8
8
|
maxLogSize: "1M",
|
|
9
9
|
layout: {
|
|
10
10
|
type: "basic",
|
|
11
11
|
},
|
|
12
12
|
},
|
|
13
13
|
stderr: { type: "stderr" },
|
|
14
|
+
stdout: { type: "stdout" },
|
|
14
15
|
},
|
|
15
16
|
categories: {
|
|
16
17
|
default: {
|
|
17
|
-
appenders: ["mcp", "
|
|
18
|
+
appenders: ["mcp", "stdout"],
|
|
18
19
|
level: "trace",
|
|
19
20
|
},
|
|
20
21
|
mcp: {
|
|
21
|
-
appenders: ["mcp", "
|
|
22
|
+
appenders: ["mcp", "stdout"],
|
|
22
23
|
level: "trace",
|
|
23
24
|
},
|
|
24
25
|
},
|
package/dist/network.js
CHANGED
|
@@ -1,8 +1,6 @@
|
|
|
1
1
|
import { logger } from "./logger.js";
|
|
2
|
+
import { authStore } from "./transports/streamable-http.js";
|
|
2
3
|
export const RHOMBUS_API_KEY = process.env.RHOMBUS_API_KEY;
|
|
3
|
-
if (!RHOMBUS_API_KEY) {
|
|
4
|
-
console.error("Missing RHOMBUS_API_KEY");
|
|
5
|
-
}
|
|
6
4
|
export const serverUrl = process.env.RHOMBUS_API_SERVER || "api2.rhombussystems.com";
|
|
7
5
|
export const BASE_URL = `https://${serverUrl}/api`;
|
|
8
6
|
export const STATIC_HEADERS = {
|
|
@@ -14,12 +12,6 @@ export const AUTH_HEADERS = {
|
|
|
14
12
|
"x-auth-apikey": RHOMBUS_API_KEY,
|
|
15
13
|
"x-auth-scheme": "api-token",
|
|
16
14
|
};
|
|
17
|
-
const enableLogs = process.env.ENABLE_LOGS;
|
|
18
|
-
const log = (msg) => {
|
|
19
|
-
if (!enableLogs)
|
|
20
|
-
return;
|
|
21
|
-
console.error(msg);
|
|
22
|
-
};
|
|
23
15
|
export const appendQueryParams = (url, params) => {
|
|
24
16
|
if (!params || typeof params !== "object")
|
|
25
17
|
return url;
|
|
@@ -34,15 +26,48 @@ export const appendQueryParams = (url, params) => {
|
|
|
34
26
|
const queryString = existingSearchParams.toString();
|
|
35
27
|
return queryString ? `${baseUrl}?${queryString}` : baseUrl;
|
|
36
28
|
};
|
|
37
|
-
export async function postApi(route, body, modifiers) {
|
|
29
|
+
export async function postApi({ route, body, modifiers, sessionId, }) {
|
|
30
|
+
let url = BASE_URL + route;
|
|
31
|
+
// construct auth headers
|
|
32
|
+
let authHeaders = {};
|
|
33
|
+
if (!sessionId) {
|
|
34
|
+
// if no sessionId, we fall back to the api key in our environment variables
|
|
35
|
+
authHeaders = AUTH_HEADERS;
|
|
36
|
+
}
|
|
37
|
+
else {
|
|
38
|
+
// use sessionId to get auth
|
|
39
|
+
const auth = authStore.get(sessionId);
|
|
40
|
+
if (!auth) {
|
|
41
|
+
logger.error(`No auth found for sessionId: ${sessionId}`);
|
|
42
|
+
throw new Error(`No auth found for sessionId: ${sessionId}`);
|
|
43
|
+
}
|
|
44
|
+
if ("apiKey" in auth) {
|
|
45
|
+
authHeaders = {
|
|
46
|
+
"x-auth-apikey": auth.apiKey,
|
|
47
|
+
"x-auth-scheme": "api-token",
|
|
48
|
+
};
|
|
49
|
+
}
|
|
50
|
+
else if ("sessionId" in auth) {
|
|
51
|
+
authHeaders = {
|
|
52
|
+
"x-auth-session": auth.sessionId,
|
|
53
|
+
"x-auth-chat": auth.latestRecordUuid,
|
|
54
|
+
"x-auth-scheme": "chatbot",
|
|
55
|
+
};
|
|
56
|
+
url = appendQueryParams(url, { _rs: auth.sessionId });
|
|
57
|
+
}
|
|
58
|
+
else if ("cookie" in auth) {
|
|
59
|
+
authHeaders = {
|
|
60
|
+
"x-auth-scheme": "web2",
|
|
61
|
+
cookie: auth.cookie,
|
|
62
|
+
};
|
|
63
|
+
}
|
|
64
|
+
}
|
|
38
65
|
// merge headers
|
|
39
|
-
|
|
66
|
+
const requestHeaders = {
|
|
40
67
|
...STATIC_HEADERS,
|
|
41
|
-
...
|
|
68
|
+
...authHeaders,
|
|
42
69
|
...(modifiers?.headers ?? {}),
|
|
43
70
|
};
|
|
44
|
-
let url = BASE_URL + route;
|
|
45
|
-
// attach query params if provided
|
|
46
71
|
if (modifiers?.query) {
|
|
47
72
|
url = appendQueryParams(url, modifiers.query);
|
|
48
73
|
}
|
|
@@ -72,7 +97,9 @@ export async function postApi(route, body, modifiers) {
|
|
|
72
97
|
// throw new Error(`HTTP error! status: ${response.status}`);
|
|
73
98
|
}
|
|
74
99
|
const ret = await response.json();
|
|
75
|
-
|
|
100
|
+
const jsonStr = JSON.stringify(ret);
|
|
101
|
+
const truncatedJson = jsonStr.length > 150 ? jsonStr.substring(0, 150) + "..." : jsonStr;
|
|
102
|
+
logger.debug(`✅ RESPONSE - ${response.ok} - ${truncatedJson}`);
|
|
76
103
|
return ret;
|
|
77
104
|
}
|
|
78
105
|
catch (error) {
|
|
@@ -13,7 +13,7 @@ The paths are categorized into "common_paths", "parameterized_paths", and "exter
|
|
|
13
13
|
|
|
14
14
|
It is in the form of a JSON file.
|
|
15
15
|
|
|
16
|
-
Use these paths to help the user navigate to the correct that they may want to go to. However, try not directly reference the path itself, rather describe it, and make sure to show the user a button to help navigate with.
|
|
16
|
+
Use these paths to help the user navigate to the correct page that they may want to go to. However, try not directly reference the path itself, rather describe it, and make sure to show the user a button to help navigate with.
|
|
17
17
|
When providing a path, make sure to be very exact. You're only allowed to substitute in path segments that begin with : or are surrounded by brackets []. For example, /locations/:locationUuid, you need to get a location's UUID and replace :locationUuid with the actual UUID`,
|
|
18
18
|
}, async (uri) => {
|
|
19
19
|
const currentDir = path.dirname(fileURLToPath(import.meta.url));
|
|
@@ -0,0 +1,153 @@
|
|
|
1
|
+
import OpenAI from "openai";
|
|
2
|
+
// Constants following coding style guidelines
|
|
3
|
+
const EMBEDDING_MODEL = "text-embedding-3-small";
|
|
4
|
+
const EMBEDDING_BATCH_SIZE = 100;
|
|
5
|
+
const MAX_TOKENS_PER_REQUEST = 8000;
|
|
6
|
+
const DELAY_BETWEEN_BATCHES_MS = 1000;
|
|
7
|
+
const MAX_RETRY_ATTEMPTS = 3;
|
|
8
|
+
const RETRY_DELAY_MS = 2000;
|
|
9
|
+
export class EmbeddingService {
|
|
10
|
+
openai;
|
|
11
|
+
model;
|
|
12
|
+
batchSize;
|
|
13
|
+
constructor(apiKey, model = EMBEDDING_MODEL) {
|
|
14
|
+
this.openai = new OpenAI({
|
|
15
|
+
apiKey: apiKey || process.env.OPENAI_API_KEY,
|
|
16
|
+
});
|
|
17
|
+
this.model = model;
|
|
18
|
+
this.batchSize = EMBEDDING_BATCH_SIZE;
|
|
19
|
+
}
|
|
20
|
+
/**
|
|
21
|
+
* Generate embeddings for a batch of text chunks
|
|
22
|
+
* @param requests - Array of embedding requests
|
|
23
|
+
* @returns Batch result with embeddings and cost information
|
|
24
|
+
*/
|
|
25
|
+
async generateEmbeddings(requests) {
|
|
26
|
+
if (!requests || requests.length === 0) {
|
|
27
|
+
return {
|
|
28
|
+
results: [],
|
|
29
|
+
totalTokens: 0,
|
|
30
|
+
totalCost: 0,
|
|
31
|
+
batchCount: 0,
|
|
32
|
+
};
|
|
33
|
+
}
|
|
34
|
+
const results = [];
|
|
35
|
+
let totalTokens = 0;
|
|
36
|
+
let totalCost = 0;
|
|
37
|
+
const batches = this.createBatches(requests);
|
|
38
|
+
console.log(`Processing ${requests.length} embeddings in ${batches.length} batches`);
|
|
39
|
+
for (let batchIndex = 0; batchIndex < batches.length; batchIndex++) {
|
|
40
|
+
const batch = batches[batchIndex];
|
|
41
|
+
console.log(`Processing batch ${batchIndex + 1}/${batches.length} (${batch.length} items)`);
|
|
42
|
+
try {
|
|
43
|
+
const batchResult = await this.processBatch(batch);
|
|
44
|
+
results.push(...batchResult.results);
|
|
45
|
+
totalTokens += batchResult.totalTokens;
|
|
46
|
+
totalCost += batchResult.totalCost;
|
|
47
|
+
// Add delay between batches to respect rate limits
|
|
48
|
+
if (batchIndex < batches.length - 1) {
|
|
49
|
+
await this.delay(DELAY_BETWEEN_BATCHES_MS);
|
|
50
|
+
}
|
|
51
|
+
}
|
|
52
|
+
catch (error) {
|
|
53
|
+
console.error(`Error processing batch ${batchIndex + 1}:`, error);
|
|
54
|
+
throw error;
|
|
55
|
+
}
|
|
56
|
+
}
|
|
57
|
+
return {
|
|
58
|
+
results,
|
|
59
|
+
totalTokens,
|
|
60
|
+
totalCost,
|
|
61
|
+
batchCount: batches.length,
|
|
62
|
+
};
|
|
63
|
+
}
|
|
64
|
+
/**
|
|
65
|
+
* Generate embedding for a single text chunk
|
|
66
|
+
* @param text - Text to embed
|
|
67
|
+
* @param chunkId - Unique identifier for the chunk
|
|
68
|
+
* @returns Single embedding result
|
|
69
|
+
*/
|
|
70
|
+
async generateSingleEmbedding(text, chunkId) {
|
|
71
|
+
const result = await this.generateEmbeddings([{ text, chunkId }]);
|
|
72
|
+
if (result.results.length === 0) {
|
|
73
|
+
throw new Error(`Failed to generate embedding for chunk ${chunkId}`);
|
|
74
|
+
}
|
|
75
|
+
return result.results[0];
|
|
76
|
+
}
|
|
77
|
+
/**
|
|
78
|
+
* Estimate cost for embedding generation
|
|
79
|
+
* @param tokenCount - Total number of tokens to embed
|
|
80
|
+
* @returns Estimated cost in USD
|
|
81
|
+
*/
|
|
82
|
+
estimateCost(tokenCount) {
|
|
83
|
+
// text-embedding-3-small costs $0.02 per 1M tokens
|
|
84
|
+
const costPerMilTokens = 0.02;
|
|
85
|
+
return (tokenCount / 1_000_000) * costPerMilTokens;
|
|
86
|
+
}
|
|
87
|
+
async processBatch(batch) {
|
|
88
|
+
const texts = batch.map(req => req.text);
|
|
89
|
+
let attempt = 0;
|
|
90
|
+
while (attempt < MAX_RETRY_ATTEMPTS) {
|
|
91
|
+
try {
|
|
92
|
+
const response = await this.openai.embeddings.create({
|
|
93
|
+
model: this.model,
|
|
94
|
+
input: texts,
|
|
95
|
+
});
|
|
96
|
+
const results = batch.map((request, index) => {
|
|
97
|
+
const embeddingData = response.data[index];
|
|
98
|
+
const tokenCount = response.usage?.total_tokens
|
|
99
|
+
? Math.round(response.usage.total_tokens / batch.length)
|
|
100
|
+
: 0;
|
|
101
|
+
return {
|
|
102
|
+
chunkId: request.chunkId,
|
|
103
|
+
embedding: embeddingData.embedding,
|
|
104
|
+
tokenCount,
|
|
105
|
+
cost: this.estimateCost(tokenCount),
|
|
106
|
+
model: this.model,
|
|
107
|
+
};
|
|
108
|
+
});
|
|
109
|
+
return {
|
|
110
|
+
results,
|
|
111
|
+
totalTokens: response.usage?.total_tokens || 0,
|
|
112
|
+
totalCost: this.estimateCost(response.usage?.total_tokens || 0),
|
|
113
|
+
batchCount: 1,
|
|
114
|
+
};
|
|
115
|
+
}
|
|
116
|
+
catch (error) {
|
|
117
|
+
attempt++;
|
|
118
|
+
console.error(`Attempt ${attempt} failed for batch:`, error);
|
|
119
|
+
if (attempt >= MAX_RETRY_ATTEMPTS) {
|
|
120
|
+
throw new Error(`Failed to process batch after ${MAX_RETRY_ATTEMPTS} attempts: ${error instanceof Error ? error.message : "Unknown error"}`);
|
|
121
|
+
}
|
|
122
|
+
// Exponential backoff delay
|
|
123
|
+
const delayMs = RETRY_DELAY_MS * Math.pow(2, attempt - 1);
|
|
124
|
+
console.log(`Retrying in ${delayMs}ms...`);
|
|
125
|
+
await this.delay(delayMs);
|
|
126
|
+
}
|
|
127
|
+
}
|
|
128
|
+
throw new Error("Unexpected end of retry attempts");
|
|
129
|
+
}
|
|
130
|
+
createBatches(requests) {
|
|
131
|
+
const batches = [];
|
|
132
|
+
for (let i = 0; i < requests.length; i += this.batchSize) {
|
|
133
|
+
const batch = requests.slice(i, i + this.batchSize);
|
|
134
|
+
// Check if batch exceeds token limit (rough estimate)
|
|
135
|
+
const estimatedTokens = batch.reduce((sum, req) => sum + req.text.length / 4, 0);
|
|
136
|
+
if (estimatedTokens > MAX_TOKENS_PER_REQUEST) {
|
|
137
|
+
// Split large batch into smaller ones
|
|
138
|
+
const halfSize = Math.floor(batch.length / 2);
|
|
139
|
+
batches.push(batch.slice(0, halfSize));
|
|
140
|
+
if (halfSize < batch.length) {
|
|
141
|
+
batches.push(batch.slice(halfSize));
|
|
142
|
+
}
|
|
143
|
+
}
|
|
144
|
+
else {
|
|
145
|
+
batches.push(batch);
|
|
146
|
+
}
|
|
147
|
+
}
|
|
148
|
+
return batches;
|
|
149
|
+
}
|
|
150
|
+
async delay(ms) {
|
|
151
|
+
return new Promise(resolve => setTimeout(resolve, ms));
|
|
152
|
+
}
|
|
153
|
+
}
|