rhombus-node-mcp 0.1.13 → 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.
Files changed (81) hide show
  1. package/README.md +11 -3
  2. package/dist/api/camera-tool-api.js +108 -0
  3. package/dist/api/clips-tool-api.js +36 -0
  4. package/dist/api/create-camera-policy-tool-api.js +9 -0
  5. package/dist/api/create-tool-api.js +53 -0
  6. package/dist/api/entity-lookup-tool-api.js +37 -0
  7. package/dist/api/events-tool-api.js +320 -0
  8. package/dist/api/faces-tool-api.js +110 -0
  9. package/dist/api/get-entity-tool-api.js +176 -0
  10. package/dist/api/get-org-information-tool-api.js +9 -0
  11. package/dist/api/location-tool-api.js +9 -0
  12. package/dist/api/lpr-tool-api.js +68 -0
  13. package/dist/api/policy-alerts-tool-api.js +42 -0
  14. package/dist/api/reboot-cameras-tool-api.js +34 -0
  15. package/dist/api/report-tool-api.js +426 -0
  16. package/dist/api/time-tool-api.js +90 -0
  17. package/dist/api/update-tool-api.js +148 -0
  18. package/dist/createServer.js +50 -0
  19. package/dist/disabled-tools/endpoint-to-keys-tool.js +84 -0
  20. package/dist/disabled-tools/semantic-search-tool.js +90 -0
  21. package/dist/index.js +25 -32
  22. package/dist/logger.js +5 -4
  23. package/dist/network.js +45 -14
  24. package/dist/resources/routes.json.js +1 -1
  25. package/dist/services/embedding-service.js +153 -0
  26. package/dist/services/faiss-search-service.js +261 -0
  27. package/dist/tools/camera-tool.js +100 -0
  28. package/dist/tools/clips-tool.js +23 -38
  29. package/dist/tools/count-tool.js +25 -0
  30. package/dist/tools/create-camera-policy-tool.js +214 -0
  31. package/dist/tools/create-tool.js +25 -75
  32. package/dist/tools/entity-lookup-tool.js +35 -0
  33. package/dist/tools/events-tool.js +188 -93
  34. package/dist/tools/faces-tool.js +59 -133
  35. package/dist/tools/get-entity-tool.js +78 -0
  36. package/dist/tools/get-org-information-tool.js +18 -0
  37. package/dist/tools/location-tool.js +24 -32
  38. package/dist/tools/lpr-tool.js +71 -0
  39. package/dist/tools/policy-alerts-tool.js +32 -43
  40. package/dist/tools/reboot-cameras-tool.js +34 -0
  41. package/dist/tools/report-tool.js +190 -0
  42. package/dist/tools/time-conversion-tool.js +43 -0
  43. package/dist/tools/time-tool.js +17 -55
  44. package/dist/tools/update-tool.js +262 -0
  45. package/dist/transports/stdio.js +10 -0
  46. package/dist/transports/streamable-http.js +201 -0
  47. package/dist/{tools/devices/camera-tool/types.js → types/camera-tool-types.js} +13 -0
  48. package/dist/types/clips-tool-types.js +41 -0
  49. package/dist/types/create-camera-policy-tool-types.js +44 -0
  50. package/dist/types/create-tool-types.js +8 -0
  51. package/dist/types/deviceType.js +1 -0
  52. package/dist/types/endpoint-to-keys-tool-types.js +7 -0
  53. package/dist/types/entity-lookup-tool-types.js +70 -0
  54. package/dist/types/events-tools-types.js +257 -0
  55. package/dist/types/faces-tools-types.js +143 -0
  56. package/dist/types/get-entity-tool-types.js +25 -0
  57. package/dist/types/get-org-information-tool-types.js +3 -0
  58. package/dist/types/location-tool-types.js +11 -0
  59. package/dist/types/lpr-tool-types.js +97 -0
  60. package/dist/types/policy-alerts-tool-types.js +74 -0
  61. package/dist/types/reboot-cameras-tool-types.js +8 -0
  62. package/dist/types/report-tool-types.js +268 -0
  63. package/dist/types/schema-components.js +7093 -0
  64. package/dist/types/schema.js +1 -0
  65. package/dist/types/semantic-search-tool-types.js +5 -0
  66. package/dist/types/time-conversion-tool-types.js +8 -0
  67. package/dist/types/time-tool-types.js +11 -0
  68. package/dist/types/update-tool-types.js +186 -0
  69. package/dist/types/zod-schemas.js +21315 -0
  70. package/dist/types.js +17 -7
  71. package/dist/util.js +96 -14
  72. package/dist/utils/confirmation.js +1 -7
  73. package/dist/utils/reduce-output.js +35 -0
  74. package/dist/utils/remove-nulls.js +28 -0
  75. package/dist/utils/temp.js +8 -0
  76. package/dist/utils/timestampInput.js +12 -0
  77. package/package.json +28 -3
  78. package/dist/tools/devices/camera-tool/camera-tool.js +0 -218
  79. package/dist/tools/devices/get-entity-tool.js +0 -117
  80. package/dist/tools/get-org-information.js +0 -18
  81. package/dist/tools/reboot-cameras.js +0 -63
@@ -0,0 +1,50 @@
1
+ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
2
+ import { logger } from "./logger.js";
3
+ import getResources from "./resources/getResources.js";
4
+ import getTools from "./tools/getTools.js";
5
+ let initiated = false;
6
+ let resources;
7
+ let tools;
8
+ export async function serverInit() {
9
+ resources = await getResources();
10
+ logger.info(`📚 Found ${resources.length} resources`);
11
+ for (const resource of resources) {
12
+ logger.debug(`📕 - ${resource.name}`);
13
+ }
14
+ tools = await getTools();
15
+ logger.info(`🛠️ Found ${tools.length} tools`);
16
+ for (const tool of tools) {
17
+ logger.debug(`🔧 - ${tool.name}`);
18
+ }
19
+ initiated = true;
20
+ }
21
+ export default async function createServer() {
22
+ if (!initiated) {
23
+ await serverInit();
24
+ }
25
+ logger.info(`🖥️ Creating Server`);
26
+ const server = new McpServer({
27
+ name: "rhombus-node-mcp",
28
+ version: "1.0.0",
29
+ capabilities: {
30
+ resources: {},
31
+ tools: {},
32
+ },
33
+ });
34
+ for (const resource of resources) {
35
+ resource.create(server);
36
+ }
37
+ logger.info(`🛠️ Registered ${resources.length} resources`);
38
+ for (const tool of tools) {
39
+ try {
40
+ await tool.create(server);
41
+ }
42
+ catch (error) {
43
+ logger.error(`Failed to register tool ${tool.name}:`, error);
44
+ // Continue with other tools instead of failing completely
45
+ }
46
+ }
47
+ logger.info(`🛠️ Registered ${tools.length} tools`);
48
+ logger.info(`✅ Server created`);
49
+ return server;
50
+ }
@@ -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
@@ -1,43 +1,36 @@
1
1
  #!/usr/bin/env node
2
2
  import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
3
- import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
4
3
  import "dotenv/config";
5
- import getTools from "./tools/getTools.js";
4
+ import { serverInit } from "./createServer.js";
6
5
  import { logger } from "./logger.js";
7
- import getResources from "./resources/getResources.js";
6
+ import stdioTransport from "./transports/stdio.js";
7
+ import streamableHttpTransport from "./transports/streamable-http.js";
8
8
  const RHOMBUS_API_KEY = process.env.RHOMBUS_API_KEY;
9
- if (!RHOMBUS_API_KEY) {
10
- logger.info("Missing RHOMBUS_API_KEY");
11
- }
12
- const serverUrl = process.env.RHOMBUS_API_SERVER || "api2.rhombussystems.com";
13
- logger.info(`Using API_KEY: ${RHOMBUS_API_KEY}`);
14
- logger.info(`To hit API server: ${serverUrl}`);
15
- logger.info("🌐 Using server url", serverUrl);
16
- export const server = new McpServer({
17
- name: "rhombus",
18
- version: "1.0.0",
19
- capabilities: {
20
- resources: {},
21
- tools: {},
22
- },
23
- });
9
+ const TRANSPORT_TYPE = process.env.TRANSPORT_TYPE || "stdio";
24
10
  async function main() {
25
- const resources = await getResources();
26
- logger.info(`🛠️ Registering ${resources.length} resources`);
27
- for (const resource of resources) {
28
- resource.create(server);
29
- logger.debug(`🔧 Registered resource ${resource.name}`);
11
+ const serverUrl = process.env.RHOMBUS_API_SERVER || "api2.rhombussystems.com";
12
+ if (RHOMBUS_API_KEY) {
13
+ logger.info(`🔑 Using API_KEY: ${RHOMBUS_API_KEY}`);
14
+ }
15
+ logger.info("🌐 Using server url", serverUrl);
16
+ await serverInit();
17
+ const server = new McpServer({
18
+ name: "rhombus",
19
+ version: "1.0.0",
20
+ capabilities: {
21
+ resources: {},
22
+ tools: {},
23
+ },
24
+ });
25
+ if (TRANSPORT_TYPE === "stdio") {
26
+ await stdioTransport();
27
+ }
28
+ else if (TRANSPORT_TYPE === "streamable-http") {
29
+ await streamableHttpTransport();
30
30
  }
31
- const tools = await getTools();
32
- logger.info(`🛠️ Registering ${tools.length} tools`);
33
- for (const tool of tools) {
34
- tool.create(server);
35
- logger.debug(`🔧 Registered tool ${tool.name}`);
31
+ else {
32
+ throw new Error(`Invalid transport type: ${TRANSPORT_TYPE}`);
36
33
  }
37
- const transport = new StdioServerTransport();
38
- logger.info(`🚙 Starting stdio transport`);
39
- await server.connect(transport);
40
- logger.info(`🚙🌬️ Connected.`);
41
34
  }
42
35
  main().catch(error => {
43
36
  console.error("Fatal error in main():", error);
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"),
8
- maxLogSize: "1K",
7
+ filename: path.resolve(process.env.LOG_FOLDER ?? process.cwd(), "./logs/rhombus-node-mcp.log"),
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", "stderr"],
18
+ appenders: ["mcp", "stdout"],
18
19
  level: "trace",
19
20
  },
20
21
  mcp: {
21
- appenders: ["mcp", "stderr"],
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,52 @@ 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 = undefined) {
38
- let requestHeaders = {
39
- ...(modifiers?.headers ?? AUTH_HEADERS),
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
+ }
65
+ // merge headers
66
+ const requestHeaders = {
40
67
  ...STATIC_HEADERS,
68
+ ...authHeaders,
69
+ ...(modifiers?.headers ?? {}),
41
70
  };
42
- let url = BASE_URL + route;
43
71
  if (modifiers?.query) {
44
72
  url = appendQueryParams(url, modifiers.query);
45
73
  }
74
+ // stringify body if it's not already a string
46
75
  if (typeof body === "object") {
47
76
  body = JSON.stringify(body);
48
77
  }
@@ -68,7 +97,9 @@ export async function postApi(route, body, modifiers = undefined) {
68
97
  // throw new Error(`HTTP error! status: ${response.status}`);
69
98
  }
70
99
  const ret = await response.json();
71
- logger.debug(`✅ RESPONSE - ${response.ok} - ${JSON.stringify(ret)}`);
100
+ const jsonStr = JSON.stringify(ret);
101
+ const truncatedJson = jsonStr.length > 150 ? jsonStr.substring(0, 150) + "..." : jsonStr;
102
+ logger.debug(`✅ RESPONSE - ${response.ok} - ${truncatedJson}`);
72
103
  return ret;
73
104
  }
74
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
+ }