# Directory Structure
```
├── .env.example
├── .eslintrc.js
├── .gitignore
├── .prettierrc
├── LICENSE
├── nodemon.json
├── package-lock.json
├── package.json
├── README.md
├── requirements.txt
├── scripts
│ ├── setup.js
│ └── start.js
├── src
│ ├── index.ts
│ ├── lib
│ │ ├── chunking-util.ts
│ │ ├── context-manager.ts
│ │ ├── embedding-service.ts
│ │ ├── errors.ts
│ │ ├── http-server.ts
│ │ ├── logger.ts
│ │ ├── namespace-manager.ts
│ │ ├── search-module.ts
│ │ ├── sqlite-store.ts
│ │ ├── tool-definitions.ts
│ │ └── tool-implementations.ts
│ └── python
│ └── embedding_service.py
└── tsconfig.json
```
# Files
--------------------------------------------------------------------------------
/.env.example:
--------------------------------------------------------------------------------
```
1 | # Path to SQLite database file
2 | DB_PATH=./data/context.db
3 |
4 | PORT=3000
5 |
6 | # Use HTTP SSE or Stdio
7 | USE_HTTP_SSE=true
8 |
9 | # Logging Configuration: debug, info, warn, error
10 | LOG_LEVEL=info
```
--------------------------------------------------------------------------------
/.prettierrc:
--------------------------------------------------------------------------------
```
1 | {
2 | "singleQuote": true,
3 | "trailingComma": "es5",
4 | "printWidth": 100,
5 | "tabWidth": 2,
6 | "bracketSpacing": true,
7 | "jsxBracketSameLine": false,
8 | "semi": true,
9 | "useTabs": false,
10 | "proseWrap": "always"
11 | }
```
--------------------------------------------------------------------------------
/.eslintrc.js:
--------------------------------------------------------------------------------
```javascript
1 | module.exports = {
2 | "env": {
3 | "browser": true,
4 | "es2021": true
5 | },
6 | "plugins": ["simple-import-sort", "prettier", "n", "promise"],
7 | "extends": "standard-with-typescript",
8 | "overrides": [
9 | {
10 | "env": {
11 | "node": true
12 | },
13 | "files": [
14 | ".eslintrc.{js,cjs}"
15 | ],
16 | "parserOptions": {
17 | "sourceType": "script"
18 | }
19 | }
20 | ],
21 | "parserOptions": {
22 | "ecmaVersion": "latest",
23 | "sourceType": "module"
24 | },
25 | "rules": {
26 | '@typescript-eslint/no-explicit-any': 'off',
27 | '@typescript-eslint/explicit-module-boundary-types': 'off',
28 | '@typescript-eslint/no-unused-vars': 'error',
29 | 'simple-import-sort/imports': 'warn',
30 | 'simple-import-sort/exports': 'warn',
31 | 'no-async-promise-executor': 'off',
32 | 'prefer-arrow-callback': 'error',
33 | 'no-prototype-builtins': 'off',
34 | 'prefer-const': 'error',
35 | 'no-var': 'error',
36 | 'prefer-template': 'error',
37 | 'no-useless-escape': 'off',
38 | "indent": "off",
39 | "no-use-before-define": "off",
40 | '@typescript-eslint/indent': 'off',
41 | '@typescript-eslint/restrict-template-expressions': 'off',
42 | '@typescript-eslint/prefer-nullish-coalescing': 'off',
43 | '@typescript-eslint/semi': 'off',
44 | '@typescript-eslint/explicit-function-return-type': 'off',
45 | '@typescript-eslint/space-before-function-paren': 'off',
46 | '@typescript-eslint/no-floating-promises': 'off',
47 | '@typescript-eslint/strict-boolean-expressions': 'off',
48 | '@typescript-eslint/comma-dangle': 'off',
49 | "@typescript-eslint/member-delimiter-style": "off"
50 | }
51 | }
52 |
```
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
```
1 | # Logs
2 | logs
3 | *.log
4 | npm-debug.log*
5 | yarn-debug.log*
6 | yarn-error.log*
7 | lerna-debug.log*
8 | .pnpm-debug.log*
9 |
10 | # Diagnostic reports (https://nodejs.org/api/report.html)
11 | report.[0-9]*.[0-9]*.[0-9]*.[0-9]*.json
12 |
13 | # Runtime data
14 | pids
15 | *.pid
16 | *.seed
17 | *.pid.lock
18 |
19 | # Directory for instrumented libs generated by jscoverage/JSCover
20 | lib-cov
21 |
22 | # Coverage directory used by tools like istanbul
23 | coverage
24 | *.lcov
25 |
26 | # nyc test coverage
27 | .nyc_output
28 |
29 | # Grunt intermediate storage (https://gruntjs.com/creating-plugins#storing-task-files)
30 | .grunt
31 |
32 | # Bower dependency directory (https://bower.io/)
33 | bower_components
34 |
35 | # node-waf configuration
36 | .lock-wscript
37 |
38 | # Compiled binary addons (https://nodejs.org/api/addons.html)
39 | build/Release
40 |
41 | # Dependency directories
42 | node_modules/
43 | jspm_packages/
44 |
45 | # Snowpack dependency directory (https://snowpack.dev/)
46 | web_modules/
47 |
48 | # TypeScript cache
49 | *.tsbuildinfo
50 |
51 | # Optional npm cache directory
52 | .npm
53 |
54 | # Optional eslint cache
55 | .eslintcache
56 |
57 | # Optional stylelint cache
58 | .stylelintcache
59 |
60 | # Microbundle cache
61 | .rpt2_cache/
62 | .rts2_cache_cjs/
63 | .rts2_cache_es/
64 | .rts2_cache_umd/
65 |
66 | # Optional REPL history
67 | .node_repl_history
68 |
69 | # Output of 'npm pack'
70 | *.tgz
71 |
72 | # Yarn Integrity file
73 | .yarn-integrity
74 |
75 | # dotenv environment variable files
76 | .env
77 | .env.development.local
78 | .env.test.local
79 | .env.production.local
80 | .env.local
81 |
82 | # parcel-bundler cache (https://parceljs.org/)
83 | .cache
84 | .parcel-cache
85 |
86 | # Next.js build output
87 | .next
88 | out
89 |
90 | # Nuxt.js build / generate output
91 | .nuxt
92 | dist
93 |
94 | # Gatsby files
95 | .cache/
96 | # Comment in the public line in if your project uses Gatsby and not Next.js
97 | # https://nextjs.org/blog/next-9-1#public-directory-support
98 | # public
99 |
100 | # vuepress build output
101 | .vuepress/dist
102 |
103 | # vuepress v2.x temp and cache directory
104 | .temp
105 | .cache
106 |
107 | # Docusaurus cache and generated files
108 | .docusaurus
109 |
110 | # Serverless directories
111 | .serverless/
112 |
113 | # FuseBox cache
114 | .fusebox/
115 |
116 | # DynamoDB Local files
117 | .dynamodb/
118 |
119 | # TernJS port file
120 | .tern-port
121 |
122 | # Stores VSCode versions used for testing VSCode extensions
123 | .vscode-test
124 |
125 | # yarn v2
126 | .yarn/cache
127 | .yarn/unplugged
128 | .yarn/build-state.yml
129 | .yarn/install-state.gz
130 | .pnp.*
131 |
132 | # misc
133 | .idea
134 | venv
135 | /data/
136 | src/python/download_model.py
137 |
```
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
```markdown
1 | # Simple Memory Extension MCP Server
2 |
3 | An MCP server to extend the context window / memory of agents. Useful when coding big features or vibe coding and need to store/recall progress, key moments or changes or anything worth remembering. Simply ask the agent to store memories and recall whenever you need or ask the agent to fully manage its memory (through cursor rules for example) however it sees fit.
4 |
5 | ## Usage
6 |
7 | ### Starting the Server
8 |
9 | ```bash
10 | npm install
11 | npm start
12 | ```
13 |
14 | ### Available Tools
15 |
16 | #### Context Item Management
17 | - `store_context_item` - Store a value with key in namespace
18 | - `retrieve_context_item_by_key` - Get value by key
19 | - `delete_context_item` - Delete key-value pair
20 |
21 | #### Namespace Management
22 | - `create_namespace` - Create new namespace
23 | - `delete_namespace` - Delete namespace and all contents
24 | - `list_namespaces` - List all namespaces
25 | - `list_context_item_keys` - List keys in a namespace
26 |
27 | #### Semantic Search
28 | - `retrieve_context_items_by_semantic_search` - Find items by meaning
29 |
30 | ### Semantic Search Implementation
31 |
32 | 1. Query converted to vector using E5 model
33 | 2. Text automatically split into chunks for better matching
34 | 3. Cosine similarity calculated between query and stored chunks
35 | 4. Results filtered by threshold and sorted by similarity
36 | 5. Top matches returned with full item values
37 |
38 | ## Development
39 |
40 | ```bash
41 | # Dev server
42 | npm run dev
43 |
44 | # Format code
45 | npm run format
46 | ```
47 |
48 | ## .env
49 |
50 | ```
51 | # Path to SQLite database file
52 | DB_PATH=./data/context.db
53 |
54 | PORT=3000
55 |
56 | # Use HTTP SSE or Stdio
57 | USE_HTTP_SSE=true
58 |
59 | # Logging Configuration: debug, info, warn, error
60 | LOG_LEVEL=info
61 | ```
62 |
63 | ## Semantic Search
64 |
65 | This project includes semantic search capabilities using the E5 embedding model from Hugging Face. This allows you to find context items based on their meaning rather than just exact key matches.
66 |
67 | ### Setup
68 |
69 | The semantic search feature requires Python dependencies, but these *should be* automatically installed when you run: `npm run start`
70 |
71 | ### Embedding Model
72 |
73 | We use the [intfloat/multilingual-e5-large-instruct](https://huggingface.co/intfloat/multilingual-e5-large-instruct)
74 |
75 |
76 | ### Notes
77 |
78 | Developed mostly while vibe coding, so don't expect much :D. But it works, and I found it helpful so w/e. Feel free to contribute or suggest improvements.
79 |
```
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
```
1 | torch>=2.0.0
2 | transformers>=4.30.0
3 | sentencepiece>=0.1.99
```
--------------------------------------------------------------------------------
/nodemon.json:
--------------------------------------------------------------------------------
```json
1 | {
2 | "watch": ["src"],
3 | "ext": "ts,json",
4 | "ignore": ["src/**/*.spec.ts"],
5 | "exec": "node --loader ts-node/esm src/index.ts"
6 | }
7 |
```
--------------------------------------------------------------------------------
/package.json:
--------------------------------------------------------------------------------
```json
1 | {
2 | "name": "simple-memory-extension-mcp-server",
3 | "version": "1.0.0",
4 | "description": "An MCP server providing a persistent key-value memory store with semantic search for your agents",
5 | "main": "dist/index.js",
6 | "start": "node dist/index.js",
7 | "bin": {
8 | "enhanced-kv-mcp-server": "./dist/index.js"
9 | },
10 | "type": "module",
11 | "scripts": {
12 | "build": "tsc",
13 | "start": "node scripts/start.js",
14 | "setup": "node scripts/setup.js",
15 | "dev": "nodemon",
16 | "format": "prettier --write \"src/**/*.ts\""
17 | },
18 | "keywords": [
19 | "mcp",
20 | "key-value",
21 | "sqlite",
22 | "persistence",
23 | "namespace",
24 | "agents",
25 | "llm",
26 | "model-context-protocol",
27 | "semantic-search"
28 | ],
29 | "author": "gmacev",
30 | "dependencies": {
31 | "@modelcontextprotocol/sdk": "^1.4.1",
32 | "dotenv": "^16.4.5",
33 | "sqlite3": "^5.1.7",
34 | "zod": "^3.22.4"
35 | },
36 | "devDependencies": {
37 | "@types/node": "^22.13.9",
38 | "@types/sqlite3": "^3.1.11",
39 | "@typescript-eslint/eslint-plugin": "^6.21.0",
40 | "@typescript-eslint/parser": "^6.21.0",
41 | "eslint": "^8.57.1",
42 | "eslint-config-prettier": "^9.1.0",
43 | "eslint-config-standard-with-typescript": "^39.1.1",
44 | "eslint-plugin-import": "^2.31.0",
45 | "eslint-plugin-n": "^16.6.2",
46 | "eslint-plugin-prettier": "^5.2.3",
47 | "eslint-plugin-promise": "^6.6.0",
48 | "eslint-plugin-simple-import-sort": "^10.0.0",
49 | "nodemon": "^3.0.1",
50 | "prettier": "^3.5.3",
51 | "ts-node": "^10.9.2",
52 | "typescript": "^5.8.2"
53 | },
54 | "engines": {
55 | "node": ">=18.0.0"
56 | },
57 | "license": "MIT"
58 | }
59 |
```
--------------------------------------------------------------------------------
/src/lib/errors.ts:
--------------------------------------------------------------------------------
```typescript
1 | /**
2 | * Custom error classes for better error classification and handling
3 | */
4 |
5 | /**
6 | * Base error class for all MCP server errors
7 | */
8 | export class McpServerError extends Error {
9 | constructor(message: string) {
10 | super(message);
11 | this.name = 'McpServerError';
12 | // Ensure proper prototype chain for instanceof checks
13 | Object.setPrototypeOf(this, McpServerError.prototype);
14 | }
15 | }
16 |
17 | /**
18 | * Database-related errors
19 | */
20 | export class DatabaseError extends McpServerError {
21 | constructor(
22 | message: string,
23 | public readonly cause?: Error
24 | ) {
25 | super(message);
26 | this.name = 'DatabaseError';
27 | Object.setPrototypeOf(this, DatabaseError.prototype);
28 | }
29 | }
30 |
31 | /**
32 | * Validation-related errors
33 | */
34 | export class ValidationError extends McpServerError {
35 | constructor(message: string) {
36 | super(message);
37 | this.name = 'ValidationError';
38 | Object.setPrototypeOf(this, ValidationError.prototype);
39 | }
40 | }
41 |
42 | /**
43 | * Error for invalid namespace names
44 | */
45 | export class InvalidNamespaceError extends ValidationError {
46 | constructor(message: string) {
47 | super(message);
48 | this.name = 'InvalidNamespaceError';
49 | Object.setPrototypeOf(this, InvalidNamespaceError.prototype);
50 | }
51 | }
52 |
53 | /**
54 | * Error for invalid key names
55 | */
56 | export class InvalidKeyError extends ValidationError {
57 | constructor(message: string) {
58 | super(message);
59 | this.name = 'InvalidKeyError';
60 | Object.setPrototypeOf(this, InvalidKeyError.prototype);
61 | }
62 | }
63 |
64 | /**
65 | * Helper function to convert unknown errors to typed errors
66 | * @param error The original error
67 | * @param defaultMessage Default message if error is not an Error object
68 | * @returns A properly typed error
69 | */
70 | export function normalizeError(error: unknown, defaultMessage = 'Unknown error'): Error {
71 | if (error instanceof Error) {
72 | return error;
73 | }
74 |
75 | if (typeof error === 'string') {
76 | return new Error(error);
77 | }
78 |
79 | return new Error(defaultMessage);
80 | }
81 |
```
--------------------------------------------------------------------------------
/src/lib/namespace-manager.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { SQLiteDataStore } from './sqlite-store.js';
2 | import { InvalidNamespaceError } from './errors.js';
3 | import { logger } from './logger.js';
4 |
5 | /**
6 | * Creates a new namespace
7 | * @param db Database connection
8 | * @param namespace Name of the namespace to create
9 | * @returns Promise that resolves to true if a new namespace was created, false if it already existed
10 | */
11 | export async function createNamespace(db: SQLiteDataStore, namespace: string): Promise<boolean> {
12 | if (!namespace) {
13 | logger.error('Cannot create namespace: namespace is required');
14 | throw new InvalidNamespaceError('Namespace is required');
15 | }
16 |
17 | logger.debug('Creating namespace', { namespace });
18 | const created = await db.createNamespace(namespace);
19 |
20 | if (created) {
21 | logger.info('Created new namespace', { namespace });
22 | } else {
23 | logger.debug('Namespace already exists', { namespace });
24 | }
25 |
26 | return created;
27 | }
28 |
29 | /**
30 | * Deletes a namespace and all items within it
31 | * @param db Database connection
32 | * @param namespace Name of the namespace to delete
33 | * @returns Promise that resolves to true if the namespace was deleted, false otherwise
34 | */
35 | export async function deleteNamespace(db: SQLiteDataStore, namespace: string): Promise<boolean> {
36 | if (!namespace) {
37 | logger.error('Cannot delete namespace: namespace is required');
38 | throw new InvalidNamespaceError('Namespace is required');
39 | }
40 |
41 | logger.debug('Deleting namespace', { namespace });
42 | const deleted = await db.deleteNamespace(namespace);
43 |
44 | if (deleted) {
45 | logger.info('Deleted namespace', { namespace });
46 | } else {
47 | logger.debug('Namespace does not exist, nothing to delete', { namespace });
48 | }
49 |
50 | return deleted;
51 | }
52 |
53 | /**
54 | * Lists all namespaces
55 | * @param db Database connection
56 | * @returns Array of namespace names
57 | */
58 | export async function listNamespaces(db: SQLiteDataStore): Promise<string[]> {
59 | logger.debug('Listing namespaces');
60 | return await db.listAllNamespaces();
61 | }
62 |
```
--------------------------------------------------------------------------------
/src/lib/context-manager.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { SQLiteDataStore } from './sqlite-store.js';
2 | import { InvalidNamespaceError, InvalidKeyError } from './errors.js';
3 | import { logger } from './logger.js';
4 |
5 | /**
6 | * Stores a context item
7 | * @param db Database connection
8 | * @param namespace Namespace for the context item
9 | * @param key Key for the context item
10 | * @param value Value to store
11 | * @returns Promise that resolves when the operation is complete
12 | */
13 | export async function storeContextItem(
14 | db: SQLiteDataStore,
15 | namespace: string,
16 | key: string,
17 | value: any
18 | ): Promise<void> {
19 | if (!namespace || namespace.trim() === '') {
20 | logger.warn('Attempted to store context item with empty namespace', { key });
21 | throw new InvalidNamespaceError('Namespace cannot be empty');
22 | }
23 |
24 | if (!key || key.trim() === '') {
25 | logger.warn('Attempted to store context item with empty key', { namespace });
26 | throw new InvalidKeyError('Key cannot be empty');
27 | }
28 |
29 | logger.debug('Storing context item', { namespace, key });
30 | // Store the item in the database
31 | await db.storeDataItem(namespace, key, value);
32 | }
33 |
34 | /**
35 | * Retrieves a context item by its key
36 | * @param db Database connection
37 | * @param namespace Namespace for the context item
38 | * @param key Key for the context item
39 | * @returns The context item with timestamps or null if not found
40 | */
41 | export async function retrieveContextItemByKey(
42 | db: SQLiteDataStore,
43 | namespace: string,
44 | key: string
45 | ): Promise<{ value: any; created_at: string; updated_at: string } | null> {
46 | if (!namespace || namespace.trim() === '') {
47 | logger.warn('Attempted to retrieve context item with empty namespace', { key });
48 | throw new InvalidNamespaceError('Namespace cannot be empty');
49 | }
50 |
51 | if (!key || key.trim() === '') {
52 | logger.warn('Attempted to retrieve context item with empty key', { namespace });
53 | throw new InvalidKeyError('Key cannot be empty');
54 | }
55 |
56 | logger.debug('Retrieving context item by key', { namespace, key });
57 | const item = await db.retrieveDataItem(namespace, key);
58 |
59 | if (!item) {
60 | logger.debug('Context item not found', { namespace, key });
61 | return null;
62 | }
63 |
64 | logger.debug('Context item retrieved successfully', { namespace, key });
65 | return {
66 | value: item.value,
67 | created_at: item.created_at,
68 | updated_at: item.updated_at,
69 | };
70 | }
71 |
72 | /**
73 | * Delete a context item
74 | * @param db Database connection
75 | * @param namespace Namespace to delete from
76 | * @param key Key to delete
77 | * @returns True if item was deleted, false if it did not exist
78 | */
79 | export async function deleteContextItem(
80 | db: SQLiteDataStore,
81 | namespace: string,
82 | key: string
83 | ): Promise<boolean> {
84 | if (!namespace || namespace.trim() === '') {
85 | logger.warn('Attempted to delete context item with empty namespace', { key });
86 | throw new InvalidNamespaceError('Namespace cannot be empty');
87 | }
88 |
89 | if (!key || key.trim() === '') {
90 | logger.warn('Attempted to delete context item with empty key', { namespace });
91 | throw new InvalidKeyError('Key cannot be empty');
92 | }
93 |
94 | logger.debug('Deleting context item', { namespace, key });
95 | const deleted = await db.deleteDataItem(namespace, key);
96 | logger.debug('Context item deletion result', { namespace, key, deleted });
97 | return deleted;
98 | }
99 |
100 | /**
101 | * List context item keys in a namespace
102 | * @param db Database instance
103 | * @param namespace Namespace to list items from
104 | * @returns Array of key objects for all items in the namespace
105 | */
106 | export async function listContextItemKeys(
107 | db: SQLiteDataStore,
108 | namespace: string
109 | ): Promise<
110 | Array<{
111 | key: string;
112 | created_at: string;
113 | updated_at: string;
114 | }>
115 | > {
116 | if (!namespace || namespace.trim() === '') {
117 | logger.warn('Attempted to list context item keys with empty namespace');
118 | throw new InvalidNamespaceError('Namespace cannot be empty');
119 | }
120 |
121 | logger.debug('Listing context item keys', { namespace });
122 | const keys = await db.listContextItemKeys(namespace);
123 | logger.debug('Listed context item keys', { namespace, count: keys.length });
124 | return keys;
125 | }
126 |
```
--------------------------------------------------------------------------------
/src/lib/logger.ts:
--------------------------------------------------------------------------------
```typescript
1 | /**
2 | * Structured logging utility for the MCP server
3 | * Provides consistent logging format across the application
4 | */
5 |
6 | // Log levels
7 | export enum LogLevel {
8 | DEBUG = 'debug',
9 | INFO = 'info',
10 | WARN = 'warn',
11 | ERROR = 'error',
12 | }
13 |
14 | // Log entry structure
15 | export interface LogEntry {
16 | timestamp: string;
17 | level: LogLevel;
18 | message: string;
19 | context?: Record<string, any>;
20 | error?: Error | unknown;
21 | }
22 |
23 | /**
24 | * Logger class for structured logging
25 | */
26 | export class Logger {
27 | private static instance: Logger;
28 | private logLevel: LogLevel = LogLevel.INFO;
29 |
30 | private constructor() {}
31 |
32 | /**
33 | * Get the singleton logger instance
34 | */
35 | public static getInstance(): Logger {
36 | if (!Logger.instance) {
37 | Logger.instance = new Logger();
38 | }
39 | return Logger.instance;
40 | }
41 |
42 | /**
43 | * Set the minimum log level
44 | * @param level Minimum log level to display
45 | */
46 | public setLogLevel(level: LogLevel): void {
47 | this.logLevel = level;
48 | }
49 |
50 | /**
51 | * Log a debug message
52 | * @param message Log message
53 | * @param context Optional context object
54 | */
55 | public debug(message: string, context?: Record<string, any>): void {
56 | this.log(LogLevel.DEBUG, message, context);
57 | }
58 |
59 | /**
60 | * Log an info message
61 | * @param message Log message
62 | * @param context Optional context object
63 | */
64 | public info(message: string, context?: Record<string, any>): void {
65 | this.log(LogLevel.INFO, message, context);
66 | }
67 |
68 | /**
69 | * Log a warning message
70 | * @param message Log message
71 | * @param context Optional context object
72 | */
73 | public warn(message: string, context?: Record<string, any>): void {
74 | this.log(LogLevel.WARN, message, context);
75 | }
76 |
77 | /**
78 | * Log an error message
79 | * @param message Log message
80 | * @param error Error object
81 | * @param context Optional context object
82 | */
83 | public error(message: string, error?: Error | unknown, context?: Record<string, any>): void {
84 | this.log(LogLevel.ERROR, message, context, error);
85 | }
86 |
87 | /**
88 | * Internal logging method
89 | * @param level Log level
90 | * @param message Log message
91 | * @param context Optional context object
92 | * @param error Optional error object
93 | */
94 | private log(
95 | level: LogLevel,
96 | message: string,
97 | context?: Record<string, any>,
98 | error?: Error | unknown
99 | ): void {
100 | // Skip logging if level is below configured level
101 | if (!this.shouldLog(level)) {
102 | return;
103 | }
104 |
105 | const entry: LogEntry = {
106 | timestamp: new Date().toISOString(),
107 | level,
108 | message,
109 | context,
110 | error,
111 | };
112 |
113 | // Format and output the log entry
114 | this.output(entry);
115 | }
116 |
117 | /**
118 | * Check if a log level should be displayed
119 | * @param level Log level to check
120 | * @returns True if the log should be displayed
121 | */
122 | private shouldLog(level: LogLevel): boolean {
123 | const levels = [LogLevel.DEBUG, LogLevel.INFO, LogLevel.WARN, LogLevel.ERROR];
124 | const configuredIndex = levels.indexOf(this.logLevel);
125 | const messageIndex = levels.indexOf(level);
126 |
127 | return messageIndex >= configuredIndex;
128 | }
129 |
130 | /**
131 | * Output a log entry
132 | * @param entry Log entry to output
133 | */
134 | private output(entry: LogEntry): void {
135 | // Format the log entry
136 | let output = `[${entry.timestamp}] [${entry.level.toUpperCase()}] ${entry.message}`;
137 |
138 | // Add context if available
139 | if (entry.context && Object.keys(entry.context).length > 0) {
140 | output += ` | Context: ${JSON.stringify(entry.context)}`;
141 | }
142 |
143 | // Add error details if available
144 | if (entry.error) {
145 | if (entry.error instanceof Error) {
146 | output += ` | Error: ${entry.error.message}`;
147 | if (entry.error.stack) {
148 | output += `\n${entry.error.stack}`;
149 | }
150 | } else {
151 | output += ` | Error: ${String(entry.error)}`;
152 | }
153 | }
154 |
155 | // Output to console based on log level
156 | switch (entry.level) {
157 | case LogLevel.DEBUG:
158 | case LogLevel.INFO:
159 | console.log(output);
160 | break;
161 | case LogLevel.WARN:
162 | console.warn(output);
163 | break;
164 | case LogLevel.ERROR:
165 | console.error(output);
166 | break;
167 | }
168 | }
169 | }
170 |
171 | // Export a singleton instance
172 | export const logger = Logger.getInstance();
173 |
```
--------------------------------------------------------------------------------
/scripts/start.js:
--------------------------------------------------------------------------------
```javascript
1 | #!/usr/bin/env node
2 |
3 | /**
4 | * This script orchestrates the complete startup process:
5 | * 1. Checks if Python setup is needed (first run)
6 | * 2. Builds the TypeScript code
7 | * 3. Sets up Python environment and pre-downloads model if needed
8 | * 4. Starts the server
9 | */
10 |
11 | import { spawn } from 'child_process';
12 | import path from 'path';
13 | import fs from 'fs';
14 | import os from 'os';
15 | import { fileURLToPath } from 'url';
16 |
17 | // Get the directory name in ES modules
18 | const __filename = fileURLToPath(import.meta.url);
19 | const __dirname = path.dirname(__filename);
20 |
21 | // Important paths
22 | const VENV_DIR = path.join(__dirname, '..', 'venv');
23 | const PYTHON_SCRIPT_DIR = path.join(__dirname, '..', 'src', 'python');
24 | const EMBEDDING_SCRIPT = path.join(PYTHON_SCRIPT_DIR, 'embedding_service.py');
25 | const SERVER_PATH = path.join(__dirname, '..', 'dist', 'index.js');
26 |
27 | // Colors for console output
28 | const colors = {
29 | reset: '\x1b[0m',
30 | green: '\x1b[32m',
31 | yellow: '\x1b[33m',
32 | blue: '\x1b[34m',
33 | red: '\x1b[31m',
34 | };
35 |
36 | /**
37 | * Logs a message with color
38 | */
39 | function log(message, color = colors.reset) {
40 | console.log(`${color}${message}${colors.reset}`);
41 | }
42 |
43 | /**
44 | * Runs a command and returns a promise
45 | */
46 | function runCommand(command, args, options = {}) {
47 | return new Promise((resolve, reject) => {
48 | log(`Running: ${command} ${args.join(' ')}`, colors.blue);
49 |
50 | const proc = spawn(command, args, {
51 | stdio: 'inherit',
52 | shell: true,
53 | ...options
54 | });
55 |
56 | proc.on('close', (code) => {
57 | if (code === 0) {
58 | resolve();
59 | } else {
60 | reject(new Error(`Command failed with exit code ${code}`));
61 | }
62 | });
63 |
64 | proc.on('error', (err) => {
65 | reject(err);
66 | });
67 | });
68 | }
69 |
70 | /**
71 | * Checks if setup is needed
72 | */
73 | function isSetupNeeded() {
74 | // Check if Python virtual environment exists
75 | const venvPython = os.platform() === 'win32'
76 | ? path.join(VENV_DIR, 'Scripts', 'python.exe')
77 | : path.join(VENV_DIR, 'bin', 'python');
78 |
79 | const venvExists = fs.existsSync(venvPython);
80 |
81 | // Also make sure the embedding script exists
82 | const scriptExists = fs.existsSync(EMBEDDING_SCRIPT);
83 |
84 | // If either doesn't exist, setup is needed
85 | const setupNeeded = !venvExists || !scriptExists;
86 |
87 | if (setupNeeded) {
88 | log('First-time setup is needed', colors.yellow);
89 | } else {
90 | log('Setup already completed, skipping setup phase', colors.green);
91 | }
92 |
93 | return setupNeeded;
94 | }
95 |
96 | /**
97 | * Builds the TypeScript code
98 | */
99 | async function buildTypeScript() {
100 | log('Building TypeScript code...', colors.blue);
101 | try {
102 | // On Windows use npm directly with shell: true to handle the npm batch file
103 | await runCommand('npm', ['run', 'build']);
104 | log('TypeScript build completed successfully.', colors.green);
105 | } catch (error) {
106 | log(`TypeScript build failed: ${error.message}`, colors.red);
107 | throw error;
108 | }
109 | }
110 |
111 | /**
112 | * Runs the Python setup if needed
113 | */
114 | async function setupIfNeeded() {
115 | // Check if setup is needed
116 | if (!isSetupNeeded()) {
117 | return;
118 | }
119 |
120 | log('Setting up Python environment...', colors.blue);
121 | try {
122 | const setupScript = path.join(__dirname, 'setup.js');
123 |
124 | // Make setup script executable on Unix systems
125 | if (process.platform !== 'win32') {
126 | try {
127 | fs.chmodSync(setupScript, '755');
128 | } catch (error) {
129 | // Ignore errors here, the script will still run with Node
130 | }
131 | }
132 |
133 | // Use node with shell: true to handle scripts on Windows
134 | await runCommand('node', [setupScript]);
135 | log('Python setup completed successfully.', colors.green);
136 | } catch (error) {
137 | log(`Python setup failed: ${error.message}`, colors.red);
138 | throw error;
139 | }
140 | }
141 |
142 | /**
143 | * Starts the server
144 | */
145 | async function startServer() {
146 | log('Starting server...', colors.blue);
147 | try {
148 | // Use node with shell: true to run the built server
149 | await runCommand('node', [SERVER_PATH]);
150 | } catch (error) {
151 | log(`Server failed to start: ${error.message}`, colors.red);
152 | throw error;
153 | }
154 | }
155 |
156 | /**
157 | * Main function
158 | */
159 | async function main() {
160 | log('Starting launch sequence...', colors.yellow);
161 |
162 | try {
163 | // Step 1: Build TypeScript
164 | await buildTypeScript();
165 |
166 | // Step 2: Setup Python if needed
167 | await setupIfNeeded();
168 |
169 | // Step 3: Start server
170 | await startServer();
171 | } catch (error) {
172 | log(`Launch sequence failed: ${error.message}`, colors.red);
173 | process.exit(1);
174 | }
175 | }
176 |
177 | // Run the main function
178 | main();
```
--------------------------------------------------------------------------------
/src/index.ts:
--------------------------------------------------------------------------------
```typescript
1 | #!/usr/bin/env node
2 |
3 | import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
4 | import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
5 | import * as dotenv from 'dotenv';
6 | import { SQLiteDataStore } from './lib/sqlite-store.js';
7 | import { registerTools } from './lib/tool-implementations.js';
8 | import { logger, LogLevel } from './lib/logger.js';
9 | import { setupHttpServer, closeAllTransports } from './lib/http-server.js';
10 | import http from 'http';
11 | import path from 'path';
12 |
13 | // Load environment variables from .env file
14 | dotenv.config();
15 |
16 | // Configure logging level from environment variable
17 | const configuredLogLevel = process.env.LOG_LEVEL || 'info';
18 | if (configuredLogLevel) {
19 | switch (configuredLogLevel.toLowerCase()) {
20 | case 'debug':
21 | logger.setLogLevel(LogLevel.DEBUG);
22 | break;
23 | case 'info':
24 | logger.setLogLevel(LogLevel.INFO);
25 | break;
26 | case 'warn':
27 | logger.setLogLevel(LogLevel.WARN);
28 | break;
29 | case 'error':
30 | logger.setLogLevel(LogLevel.ERROR);
31 | break;
32 | default:
33 | logger.warn(`Unknown log level: ${configuredLogLevel}, using default`);
34 | }
35 | }
36 |
37 | /**
38 | * Main entry point for the server
39 | */
40 | async function main() {
41 | let db: SQLiteDataStore;
42 | let server: McpServer;
43 |
44 | try {
45 | logger.info('Initializing simple-memory-extension-mcp-server...');
46 |
47 | // Get database path from environment variable or use default path
48 | const dbPath = process.env.DB_PATH || path.join(process.cwd(), 'data', 'context.db');
49 | logger.info(`Using database at: ${dbPath}`);
50 |
51 | // Initialize the SQLite database
52 | try {
53 | db = new SQLiteDataStore({
54 | dbPath: dbPath,
55 | enableChunking: true,
56 | });
57 | logger.info('Database initialized successfully');
58 | } catch (dbError) {
59 | logger.error('Failed to initialize database', dbError);
60 | process.exit(1);
61 | }
62 |
63 | // Create an MCP server
64 | server = new McpServer({
65 | name: 'simple-memory-extension-mcp-server',
66 | version: '1.0.0',
67 | });
68 |
69 | // Register all tools with the server using schemas from tool-definitions.ts
70 | registerTools(server, db);
71 | logger.info('All tools registered with the server');
72 |
73 | // Check if we should use HTTP server
74 | const useHttpSSE = process.env.USE_HTTP_SSE === 'true';
75 | const port = parseInt(process.env.PORT || '3000', 10);
76 |
77 | if (useHttpSSE) {
78 | // Set up HTTP server SSE transport
79 | const httpServer = setupHttpServer(server, port);
80 |
81 | // Setup graceful shutdown for HTTP server
82 | setupGracefulShutdown(db, httpServer);
83 | } else {
84 | // Use standard stdio transport
85 | const transport = new StdioServerTransport();
86 | await server.connect(transport);
87 | logger.info('Server ready to receive requests via stdio');
88 |
89 | // Setup graceful shutdown
90 | setupGracefulShutdown(db);
91 | }
92 | } catch (error) {
93 | logger.error('Error initializing server', error);
94 | process.exit(1);
95 | }
96 | }
97 |
98 | /**
99 | * Set up graceful shutdown to properly close database connections
100 | */
101 | function setupGracefulShutdown(db: SQLiteDataStore, httpServer?: http.Server) {
102 | const shutdown = async () => {
103 | logger.info('Shutting down server...');
104 |
105 | // Set a timeout to force exit if graceful shutdown takes too long
106 | const forceExitTimeout = setTimeout(() => {
107 | logger.error('Forced shutdown after timeout');
108 | process.exit(1);
109 | }, 5000); // Force exit after 5 seconds
110 |
111 | try {
112 | // Close all active SSE transports
113 | await closeAllTransports();
114 |
115 | if (httpServer) {
116 | await new Promise<void>((resolve) => {
117 | httpServer.close(() => {
118 | logger.info('HTTP server closed');
119 | resolve();
120 | });
121 | });
122 | }
123 |
124 | if (db) {
125 | // Get underlying DB and close it
126 | const sqliteDb = await db.getDb();
127 | if (sqliteDb) {
128 | await new Promise<void>((resolve) => {
129 | sqliteDb.close(() => {
130 | logger.info('Database connection closed');
131 | resolve();
132 | });
133 | });
134 | }
135 | }
136 |
137 | logger.info('Server shutdown complete');
138 | clearTimeout(forceExitTimeout);
139 | process.exit(0);
140 | } catch (error) {
141 | logger.error('Error during shutdown', error);
142 | clearTimeout(forceExitTimeout);
143 | process.exit(1);
144 | }
145 | };
146 |
147 | // Register shutdown handlers
148 | process.on('SIGINT', shutdown);
149 | process.on('SIGTERM', shutdown);
150 | process.on('unhandledRejection', (reason, promise) => {
151 | logger.error('Unhandled Promise rejection', { reason, promise });
152 | });
153 | }
154 |
155 | // Start the server
156 | main().catch((error) => {
157 | logger.error('Unhandled error in main function', error);
158 | process.exit(1);
159 | });
160 |
```
--------------------------------------------------------------------------------
/src/python/embedding_service.py:
--------------------------------------------------------------------------------
```python
1 | #!/usr/bin/env python3
2 | """
3 | Embedding service using the E5 model from Hugging Face.
4 | This script provides a bridge between Node.js and the E5 model.
5 | """
6 |
7 | import sys
8 | import json
9 | import torch
10 | import torch.nn.functional as F
11 | from transformers import AutoTokenizer, AutoModel
12 |
13 | # Initialize model
14 | MODEL_NAME = "intfloat/multilingual-e5-large-instruct"
15 | tokenizer = None
16 | model = None
17 |
18 | def initialize_model():
19 | """Initialize the model and tokenizer."""
20 | global tokenizer, model
21 |
22 | if tokenizer is None or model is None:
23 | print("Initializing E5 model...", file=sys.stderr)
24 | tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
25 | model = AutoModel.from_pretrained(MODEL_NAME)
26 | print(f"Model initialized: {MODEL_NAME}", file=sys.stderr)
27 |
28 | def average_pool(last_hidden_states, attention_mask):
29 | """Average pooling function for the model output."""
30 | last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
31 | return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
32 |
33 | def format_text_for_embedding(text, is_query=False):
34 | """
35 | Format text for embedding based on whether it's a query or passage.
36 | For queries, we add an instruction prefix.
37 | For passages, we use the text as is.
38 | """
39 | if is_query:
40 | # For queries, add instruction
41 | task_description = "Given a web search query, retrieve relevant passages that answer the query"
42 | return f'Instruct: {task_description}\nQuery: {text}'
43 | else:
44 | # For passages, use as is
45 | return text
46 |
47 | def generate_embedding(text, is_query=False):
48 | """Generate embedding for a single text."""
49 | initialize_model()
50 |
51 | # Format text based on type
52 | input_text = format_text_for_embedding(text, is_query)
53 |
54 | # Tokenize
55 | encoded_input = tokenizer(
56 | input_text,
57 | max_length=512,
58 | padding=True,
59 | truncation=True,
60 | return_tensors='pt'
61 | )
62 |
63 | # Generate embedding
64 | with torch.no_grad():
65 | model_output = model(**encoded_input)
66 |
67 | # Pool and normalize
68 | embedding = average_pool(model_output.last_hidden_state, encoded_input['attention_mask'])
69 | embedding = F.normalize(embedding, p=2, dim=1)
70 |
71 | # Convert to list
72 | return embedding[0].tolist()
73 |
74 | def generate_embeddings(texts, is_query=False):
75 | """Generate embeddings for multiple texts."""
76 | initialize_model()
77 |
78 | if not texts or not isinstance(texts, list):
79 | raise ValueError("Texts must be a non-empty list of strings")
80 |
81 | # Format each text based on type
82 | input_texts = [format_text_for_embedding(text, is_query) for text in texts]
83 |
84 | # Tokenize
85 | encoded_input = tokenizer(
86 | input_texts,
87 | max_length=512,
88 | padding=True,
89 | truncation=True,
90 | return_tensors='pt'
91 | )
92 |
93 | # Generate embeddings
94 | with torch.no_grad():
95 | model_output = model(**encoded_input)
96 |
97 | # Pool and normalize
98 | embeddings = average_pool(model_output.last_hidden_state, encoded_input['attention_mask'])
99 | embeddings = F.normalize(embeddings, p=2, dim=1)
100 |
101 | # Convert to list
102 | return embeddings.tolist()
103 |
104 | def process_command(command_json):
105 | """Process a command from Node.js."""
106 | try:
107 | command = command_json.get("command")
108 |
109 | if command == "initialize":
110 | initialize_model()
111 | return {"status": "initialized"}
112 |
113 | elif command == "generate_embedding":
114 | text = command_json.get("text")
115 | is_query = command_json.get("is_query", False)
116 |
117 | if not text:
118 | return {"error": "No text provided"}
119 |
120 | embedding = generate_embedding(text, is_query)
121 | return {"embedding": embedding}
122 |
123 | elif command == "generate_embeddings":
124 | texts = command_json.get("texts")
125 | is_query = command_json.get("is_query", False)
126 |
127 | if not texts or not isinstance(texts, list):
128 | return {"error": "No texts provided or invalid format"}
129 |
130 | embeddings = generate_embeddings(texts, is_query)
131 | return {"embeddings": embeddings}
132 |
133 | else:
134 | return {"error": f"Unknown command: {command}"}
135 |
136 | except Exception as e:
137 | return {"error": str(e)}
138 |
139 | def main():
140 | """Main function to process commands from stdin."""
141 | print("E5 Embedding Service started", file=sys.stderr)
142 | initialize_model()
143 |
144 | for line in sys.stdin:
145 | try:
146 | command_json = json.loads(line)
147 | result = process_command(command_json)
148 | print(json.dumps(result), flush=True)
149 | except json.JSONDecodeError:
150 | print(json.dumps({"error": "Invalid JSON"}), flush=True)
151 | except Exception as e:
152 | print(json.dumps({"error": str(e)}), flush=True)
153 |
154 | if __name__ == "__main__":
155 | main()
```
--------------------------------------------------------------------------------
/src/lib/tool-definitions.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { z } from 'zod';
2 |
3 | // Define Zod schemas for tool parameters
4 | const namespaceSchema = z
5 | .string()
6 | .min(1, 'Namespace must not be empty')
7 | .describe(
8 | 'A unique identifier for a collection of related key-value pairs. Use logical naming patterns like "user_preferences", "conversation_context", or "task_data" to organize information.'
9 | );
10 | const keySchema = z
11 | .string()
12 | .min(1, 'Key must not be empty')
13 | .describe(
14 | 'A unique identifier for a value within a namespace. Use descriptive keys that indicate the data\'s purpose, like "last_search_query", "user_location", or "current_task_id".'
15 | );
16 | const nSchema = z
17 | .number()
18 | .int('Number of turns must be an integer')
19 | .positive('Number of turns must be greater than 0')
20 | .describe(
21 | 'A positive integer specifying the quantity of items to retrieve. Use smaller values for recent context, larger values for more comprehensive history.'
22 | );
23 |
24 | // Define tool schemas
25 | export const toolSchemas = {
26 | // Context Item Management
27 | retrieve_context_item_by_key: {
28 | name: 'retrieve_context_item_by_key',
29 | description:
30 | 'Retrieves a stored value by its key from a specific namespace. Although values are stored as strings, they are automatically parsed and returned with their original data types. For example, JSON objects/arrays will be returned as structured data, numbers as numeric values, booleans as true/false, and null as null.',
31 | schema: z.object({
32 | namespace: namespaceSchema,
33 | key: keySchema,
34 | }),
35 | },
36 |
37 | store_context_item: {
38 | name: 'store_context_item',
39 | description:
40 | 'Stores a NEW value associated with a key in a specified namespace. IMPORTANT: All values must be passed as strings, but the system will intelligently parse them to preserve their data types. This tool will fail if the key already exists - use update_context_item for modifying existing items.',
41 | schema: z.object({
42 | namespace: namespaceSchema,
43 | key: keySchema,
44 | value: z
45 | .string()
46 | .describe(
47 | 'The value to store as a string. Examples: "{\"name\":\"John\"}" for objects, "[1,2,3]" for arrays, "42.5" for numbers, "true"/"false" for booleans, "null" for null. The system will automatically detect and preserve the appropriate data type when retrieving.'
48 | ),
49 | }),
50 | },
51 |
52 | update_context_item: {
53 | name: 'update_context_item',
54 | description:
55 | 'Updates an EXISTING value associated with a key in a specified namespace. This tool will fail if the key does not exist - use store_context_item for creating new items. All values must be passed as strings but will be intelligently parsed to preserve data types.',
56 | schema: z.object({
57 | namespace: namespaceSchema,
58 | key: keySchema,
59 | value: z
60 | .string()
61 | .describe(
62 | 'The new value to store as a string. Examples: "{\"name\":\"John\"}" for objects, "[1,2,3]" for arrays, "42.5" for numbers, "true"/"false" for booleans, "null" for null. The system will automatically detect and preserve the appropriate data type when retrieving.'
63 | ),
64 | }),
65 | },
66 |
67 | delete_context_item: {
68 | name: 'delete_context_item',
69 | description:
70 | 'Deletes a key-value pair from a namespace. Use this to remove data that is no longer needed or to clean up temporary storage.',
71 | schema: z.object({
72 | namespace: namespaceSchema,
73 | key: keySchema,
74 | }),
75 | },
76 |
77 | // Namespace Management
78 | create_namespace: {
79 | name: 'create_namespace',
80 | description:
81 | 'Creates a new namespace for storing key-value pairs. Use this before storing items in a new namespace. If the namespace already exists, this is a no-op and returns success.',
82 | schema: z.object({
83 | namespace: namespaceSchema,
84 | }),
85 | },
86 |
87 | delete_namespace: {
88 | name: 'delete_namespace',
89 | description:
90 | 'Deletes an entire namespace and all key-value pairs within it. Use this for cleanup when all data in a namespace is no longer needed.',
91 | schema: z.object({
92 | namespace: namespaceSchema,
93 | }),
94 | },
95 |
96 | list_namespaces: {
97 | name: 'list_namespaces',
98 | description:
99 | 'Lists all available namespaces. Use this to discover what namespaces exist before retrieving or storing data.',
100 | schema: z.object({}),
101 | },
102 |
103 | // Semantic search
104 | retrieve_context_items_by_semantic_search: {
105 | name: 'retrieve_context_items_by_semantic_search',
106 | description: 'Retrieves context items using semantic search based on query relevance',
107 | schema: z.object({
108 | namespace: namespaceSchema,
109 | query: z.string().describe('The semantic query to search for'),
110 | similarity_threshold: z
111 | .number()
112 | .optional()
113 | .describe('Minimum similarity score (0-1) to include in results'),
114 | limit: z.number().optional().describe('Maximum number of results to return'),
115 | }),
116 | },
117 |
118 | // Context Item Keys Listing
119 | list_context_item_keys: {
120 | name: 'list_context_item_keys',
121 | description:
122 | 'Lists keys and timestamps for all items in a namespace without retrieving their values. Use this to discover what data is available before retrieving specific items.',
123 | schema: z.object({
124 | namespace: namespaceSchema,
125 | }),
126 | },
127 | };
128 |
```
--------------------------------------------------------------------------------
/src/lib/search-module.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { SQLiteDataStore } from './sqlite-store.js';
2 | import { ContextItem } from './sqlite-store.js';
3 | import { logger } from './logger.js';
4 | import { EmbeddingService } from './embedding-service.js';
5 |
6 | // Singleton embedding service instance
7 | let embeddingService: EmbeddingService | null = null;
8 |
9 | /**
10 | * Gets or creates the embedding service
11 | */
12 | async function getEmbeddingService(): Promise<EmbeddingService> {
13 | if (!embeddingService) {
14 | embeddingService = new EmbeddingService();
15 | await embeddingService.initialize();
16 | }
17 | return embeddingService;
18 | }
19 |
20 | /**
21 | * Retrieves context items using semantic search
22 | * @param db Database connection
23 | * @param namespace Namespace to search in
24 | * @param query The semantic query text
25 | * @param options Additional search options
26 | * @returns Array of matching context items with similarity scores
27 | */
28 | export async function retrieveBySemantic(
29 | db: SQLiteDataStore,
30 | namespace: string,
31 | query: string,
32 | options: {
33 | limit?: number;
34 | threshold?: number;
35 | tags?: string[];
36 | } = {}
37 | ): Promise<Array<ItemWithEmbedding & { similarity: number }>> {
38 | if (!namespace || namespace.trim() === '') {
39 | throw new Error('Namespace cannot be empty');
40 | }
41 |
42 | if (!query || query.trim() === '') {
43 | throw new Error('Query cannot be empty');
44 | }
45 |
46 | logger.debug('Performing semantic search', { namespace, query, options });
47 |
48 | try {
49 | // Generate embedding for the query
50 | const queryEmbedding = await generateEmbedding(query);
51 |
52 | // Get all items from the namespace with optional tag filtering
53 | const items = await getNamespaceItems(db, namespace);
54 |
55 | // For items with embeddings, calculate similarity and sort
56 | const results = await calculateSimilarities(items, queryEmbedding);
57 |
58 | // Filter by threshold and limit results
59 | return results
60 | .filter((item) => item.similarity >= (options.threshold || 0.7))
61 | .sort((a, b) => b.similarity - a.similarity)
62 | .slice(0, options.limit || 10);
63 | } catch (error: any) {
64 | logger.error('Error in semantic search', error, { namespace, query });
65 | throw new Error(`Failed to perform semantic search: ${error.message}`);
66 | }
67 | }
68 |
69 | /**
70 | * Generates embedding for text using the E5 model
71 | */
72 | async function generateEmbedding(text: string): Promise<number[]> {
73 | try {
74 | logger.debug('Generating embedding for text', { textLength: text.length });
75 | const service = await getEmbeddingService();
76 | return service.generateEmbedding(text);
77 | } catch (error: any) {
78 | logger.error('Error generating embedding', error);
79 | // Fallback to random vector if embedding generation fails
80 | logger.debug('Using fallback random embedding');
81 | return Array.from({ length: 1024 }, () => Math.random() - 0.5);
82 | }
83 | }
84 |
85 | /**
86 | * Retrieve items from namespace with optional tag filtering
87 | */
88 | async function getNamespaceItems(db: SQLiteDataStore, namespace: string): Promise<ContextItem[]> {
89 | logger.debug(`Getting all items in namespace: ${namespace}`);
90 | return await db.retrieveAllItemsInNamespace(namespace);
91 | }
92 |
93 | // Define a type for items with parsed embeddings
94 | interface ItemWithEmbedding extends Omit<ContextItem, 'embedding'> {
95 | embedding: number[];
96 | }
97 |
98 | /**
99 | * Calculate similarity between query embedding and items
100 | */
101 | async function calculateSimilarities(
102 | items: ContextItem[],
103 | queryEmbedding: number[]
104 | ): Promise<Array<ItemWithEmbedding & { similarity: number }>> {
105 | logger.debug('Calculating similarities', { itemCount: items.length });
106 |
107 | // Generate embeddings for all items
108 | const itemsWithEmbeddings = await generateEmbeddingsForItems(items);
109 |
110 | // Calculate cosine similarity for each item
111 | return itemsWithEmbeddings.map((item) => ({
112 | ...item,
113 | similarity: calculateCosineSimilarity(queryEmbedding, item.embedding),
114 | }));
115 | }
116 |
117 | /**
118 | * Generates embeddings for a list of items
119 | */
120 | async function generateEmbeddingsForItems(items: ContextItem[]): Promise<ItemWithEmbedding[]> {
121 | logger.debug(`Generating embeddings for ${items.length} items`);
122 |
123 | const service = await getEmbeddingService();
124 | const results: ItemWithEmbedding[] = [];
125 |
126 | for (const item of items) {
127 | try {
128 | const text = extractTextFromItem(item);
129 | const embedding = await service.generateEmbedding(text);
130 |
131 | results.push({
132 | ...item,
133 | embedding,
134 | });
135 | } catch (error: any) {
136 | logger.error('Failed to generate embedding for item', {
137 | namespace: item.namespace,
138 | key: item.key,
139 | error: error.message,
140 | });
141 | // Skip items that fail embedding generation
142 | continue;
143 | }
144 | }
145 |
146 | return results;
147 | }
148 |
149 | /**
150 | * Extract text representation from an item
151 | */
152 | function extractTextFromItem(item: ContextItem): string {
153 | // Convert item value to text for embedding
154 | if (typeof item.value === 'string') {
155 | return item.value;
156 | }
157 |
158 | if (typeof item.value === 'object') {
159 | // For objects, concatenate string values
160 | return Object.entries(item.value)
161 | .filter(([_, v]) => typeof v === 'string')
162 | .map(([k, v]) => `${k}: ${v}`)
163 | .join('\n');
164 | }
165 |
166 | return `${item.key}: ${JSON.stringify(item.value)}`;
167 | }
168 |
169 | /**
170 | * Calculate cosine similarity between two embeddings
171 | */
172 | function calculateCosineSimilarity(vec1: number[], vec2: number[]): number {
173 | if (vec1.length !== vec2.length) {
174 | throw new Error('Vectors must have the same dimensions');
175 | }
176 |
177 | const dotProduct = vec1.reduce((sum, val, i) => sum + val * vec2[i], 0);
178 | const norm1 = Math.sqrt(vec1.reduce((sum, val) => sum + val * val, 0));
179 | const norm2 = Math.sqrt(vec2.reduce((sum, val) => sum + val * val, 0));
180 |
181 | if (norm1 === 0 || norm2 === 0) {
182 | return 0; // Avoid division by zero
183 | }
184 |
185 | return dotProduct / (norm1 * norm2);
186 | }
187 |
```
--------------------------------------------------------------------------------
/src/lib/chunking-util.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { logger } from './logger.js';
2 |
3 | /**
4 | * Configuration for text chunking
5 | */
6 | export interface ChunkingConfig {
7 | // Maximum tokens per chunk (default matches E5 model constraints)
8 | maxTokens: number;
9 | // Tokens to overlap between chunks for context preservation
10 | overlapTokens: number;
11 | // Rough estimation of characters per token (varies by language)
12 | charsPerToken: number;
13 | }
14 |
15 | // Default configuration optimized for the E5 model
16 | export const DEFAULT_CHUNKING_CONFIG: ChunkingConfig = {
17 | maxTokens: 400, // Target 400 tokens to stay safely below 512 limit
18 | overlapTokens: 50, // ~12.5% overlap to maintain context
19 | charsPerToken: 4, // Rough estimate - varies by language
20 | };
21 |
22 | /**
23 | * Estimates token count from character count
24 | * This is a rough approximation - actual tokenization depends on model and language
25 | */
26 | export function estimateTokenCount(
27 | text: string,
28 | charsPerToken: number = DEFAULT_CHUNKING_CONFIG.charsPerToken
29 | ): number {
30 | return Math.ceil(text.length / charsPerToken);
31 | }
32 |
33 | /**
34 | * Chunk text by semantic boundaries like paragraphs and sentences
35 | * Tries to respect natural text boundaries while staying within token limits
36 | */
37 | export function chunkTextBySemanticBoundaries(
38 | text: string,
39 | config: ChunkingConfig = DEFAULT_CHUNKING_CONFIG
40 | ): string[] {
41 | logger.debug(`Chunking text of length ${text.length} characters`);
42 |
43 | // If text is already small enough, return as-is
44 | if (estimateTokenCount(text, config.charsPerToken) <= config.maxTokens) {
45 | logger.debug('Text fits in a single chunk, no chunking needed');
46 | return [text];
47 | }
48 |
49 | const chunks: string[] = [];
50 |
51 | // First split by double newlines (paragraphs)
52 | const paragraphs = text.split(/\n\s*\n/);
53 | logger.debug(`Split into ${paragraphs.length} paragraphs`);
54 |
55 | let currentChunk = '';
56 | let currentTokenCount = 0;
57 |
58 | // Process paragraph by paragraph
59 | for (let i = 0; i < paragraphs.length; i++) {
60 | const para = paragraphs[i];
61 | const paraTokens = estimateTokenCount(para, config.charsPerToken);
62 |
63 | // If this paragraph alone exceeds max tokens, split it into sentences
64 | if (paraTokens > config.maxTokens) {
65 | logger.debug(`Large paragraph found (est. ${paraTokens} tokens), splitting into sentences`);
66 |
67 | // If we have accumulated content in current chunk, save it first
68 | if (currentChunk) {
69 | chunks.push(currentChunk);
70 | currentChunk = '';
71 | currentTokenCount = 0;
72 | }
73 |
74 | // Split large paragraph into sentences and process them
75 | const sentences = para.split(/(?<=[.!?])\s+/);
76 | let sentenceChunk = '';
77 | let sentenceTokenCount = 0;
78 |
79 | for (const sentence of sentences) {
80 | const sentenceTokens = estimateTokenCount(sentence, config.charsPerToken);
81 |
82 | // If single sentence exceeds limit, we have to split it by character count
83 | if (sentenceTokens > config.maxTokens) {
84 | logger.debug(
85 | `Very long sentence found (est. ${sentenceTokens} tokens), splitting by character count`
86 | );
87 |
88 | // Save any accumulated content first
89 | if (sentenceChunk) {
90 | chunks.push(sentenceChunk);
91 | sentenceChunk = '';
92 | sentenceTokenCount = 0;
93 | }
94 |
95 | // Force split the long sentence into multiple chunks
96 | const maxChars = config.maxTokens * config.charsPerToken;
97 | for (let j = 0; j < sentence.length; j += maxChars) {
98 | const subChunk = sentence.substring(j, j + maxChars);
99 | chunks.push(subChunk);
100 | }
101 | }
102 | // If adding this sentence exceeds limit, save current and start new
103 | else if (sentenceTokenCount + sentenceTokens > config.maxTokens) {
104 | chunks.push(sentenceChunk);
105 |
106 | // Start new chunk with overlap if possible
107 | if (sentenceChunk && config.overlapTokens > 0) {
108 | // Extract last N tokens worth of text as overlap
109 | const overlapChars = config.overlapTokens * config.charsPerToken;
110 | const overlapText = sentenceChunk.substring(
111 | Math.max(0, sentenceChunk.length - overlapChars)
112 | );
113 | sentenceChunk = overlapText + ' ' + sentence;
114 | sentenceTokenCount = estimateTokenCount(sentenceChunk, config.charsPerToken);
115 | } else {
116 | sentenceChunk = sentence;
117 | sentenceTokenCount = sentenceTokens;
118 | }
119 | }
120 | // Otherwise add to current sentence chunk
121 | else {
122 | sentenceChunk = sentenceChunk ? `${sentenceChunk} ${sentence}` : sentence;
123 | sentenceTokenCount += sentenceTokens;
124 | }
125 | }
126 |
127 | // Add the last sentence chunk if not empty
128 | if (sentenceChunk) {
129 | chunks.push(sentenceChunk);
130 | }
131 | }
132 | // If adding this paragraph would exceed the token limit
133 | else if (currentTokenCount + paraTokens > config.maxTokens) {
134 | // Save current chunk
135 | chunks.push(currentChunk);
136 |
137 | // Start new chunk with overlap if possible
138 | if (currentChunk && config.overlapTokens > 0) {
139 | // Extract last N tokens worth of text as overlap
140 | const overlapChars = config.overlapTokens * config.charsPerToken;
141 | const overlapText = currentChunk.substring(Math.max(0, currentChunk.length - overlapChars));
142 | currentChunk = overlapText + '\n\n' + para;
143 | currentTokenCount = estimateTokenCount(currentChunk, config.charsPerToken);
144 | } else {
145 | currentChunk = para;
146 | currentTokenCount = paraTokens;
147 | }
148 | }
149 | // Otherwise add to current chunk
150 | else {
151 | if (currentChunk) {
152 | currentChunk += '\n\n' + para;
153 | } else {
154 | currentChunk = para;
155 | }
156 | currentTokenCount += paraTokens;
157 | }
158 | }
159 |
160 | // Add the last chunk if not empty
161 | if (currentChunk) {
162 | chunks.push(currentChunk);
163 | }
164 |
165 | logger.debug(`Text chunked into ${chunks.length} semantic chunks`);
166 | return chunks;
167 | }
168 |
```
--------------------------------------------------------------------------------
/scripts/setup.js:
--------------------------------------------------------------------------------
```javascript
1 | #!/usr/bin/env node
2 |
3 | /**
4 | * This script automates the setup process for the semantic search feature:
5 | * - Creates a Python virtual environment (if needed)
6 | * - Installs Python dependencies
7 | * - Pre-downloads the E5 model to avoid delays on first use
8 | */
9 |
10 | import { spawn, execSync } from 'child_process';
11 | import fs from 'fs';
12 | import path from 'path';
13 | import os from 'os';
14 | import { fileURLToPath } from 'url';
15 |
16 | // Get the directory name in ES modules
17 | const __filename = fileURLToPath(import.meta.url);
18 | const __dirname = path.dirname(__filename);
19 |
20 | const PYTHON_COMMAND = os.platform() === 'win32' ? 'python' : 'python3';
21 | const VENV_DIR = path.join(__dirname, '..', 'venv');
22 | const REQUIREMENTS_FILE = path.join(__dirname, '..', 'requirements.txt');
23 | const PYTHON_SCRIPT_DIR = path.join(__dirname, '..', 'src', 'python');
24 | const MODEL_DIR = path.join(os.homedir(), '.cache', 'huggingface', 'transformers');
25 |
26 | // Ensure the Python script directory exists
27 | if (!fs.existsSync(PYTHON_SCRIPT_DIR)) {
28 | fs.mkdirSync(PYTHON_SCRIPT_DIR, { recursive: true });
29 | }
30 |
31 | // Colors for console output
32 | const colors = {
33 | reset: '\x1b[0m',
34 | green: '\x1b[32m',
35 | yellow: '\x1b[33m',
36 | blue: '\x1b[34m',
37 | red: '\x1b[31m',
38 | };
39 |
40 | /**
41 | * Logs a message with color
42 | */
43 | function log(message, color = colors.reset) {
44 | console.log(`${color}${message}${colors.reset}`);
45 | }
46 |
47 | /**
48 | * Checks if a command exists in PATH
49 | */
50 | function commandExists(command) {
51 | try {
52 | const devNull = os.platform() === 'win32' ? 'NUL' : '/dev/null';
53 | execSync(`${command} --version`, { stdio: 'ignore' });
54 | return true;
55 | } catch (e) {
56 | return false;
57 | }
58 | }
59 |
60 | /**
61 | * Checks if virtual environment exists
62 | */
63 | function venvExists() {
64 | const venvPython = os.platform() === 'win32'
65 | ? path.join(VENV_DIR, 'Scripts', 'python.exe')
66 | : path.join(VENV_DIR, 'bin', 'python');
67 | return fs.existsSync(venvPython);
68 | }
69 |
70 | /**
71 | * Creates a Python virtual environment
72 | */
73 | async function createVirtualEnv() {
74 | if (venvExists()) {
75 | log('Python virtual environment already exists.', colors.green);
76 | return;
77 | }
78 |
79 | log('Creating Python virtual environment...', colors.blue);
80 | return new Promise((resolve, reject) => {
81 | const proc = spawn(PYTHON_COMMAND, ['-m', 'venv', VENV_DIR]);
82 |
83 | proc.on('close', (code) => {
84 | if (code === 0) {
85 | log('Python virtual environment created successfully.', colors.green);
86 | resolve();
87 | } else {
88 | log(`Failed to create virtual environment (exit code ${code}).`, colors.red);
89 | reject(new Error(`Failed to create virtual environment (exit code ${code})`));
90 | }
91 | });
92 |
93 | proc.on('error', (err) => {
94 | log(`Error creating virtual environment: ${err.message}`, colors.red);
95 | reject(err);
96 | });
97 | });
98 | }
99 |
100 | /**
101 | * Installs Python dependencies
102 | */
103 | async function installDependencies() {
104 | const pipCmd = os.platform() === 'win32'
105 | ? path.join(VENV_DIR, 'Scripts', 'pip.exe')
106 | : path.join(VENV_DIR, 'bin', 'pip');
107 |
108 | log('Installing Python dependencies...', colors.blue);
109 | return new Promise((resolve, reject) => {
110 | const proc = spawn(pipCmd, ['install', '-r', REQUIREMENTS_FILE]);
111 |
112 | proc.stdout.on('data', (data) => {
113 | process.stdout.write(data.toString());
114 | });
115 |
116 | proc.stderr.on('data', (data) => {
117 | process.stderr.write(data.toString());
118 | });
119 |
120 | proc.on('close', (code) => {
121 | if (code === 0) {
122 | log('Python dependencies installed successfully.', colors.green);
123 | resolve();
124 | } else {
125 | log(`Failed to install dependencies (exit code ${code}).`, colors.red);
126 | reject(new Error(`Failed to install dependencies (exit code ${code})`));
127 | }
128 | });
129 |
130 | proc.on('error', (err) => {
131 | log(`Error installing dependencies: ${err.message}`, colors.red);
132 | reject(err);
133 | });
134 | });
135 | }
136 |
137 | /**
138 | * Pre-downloads the E5 model
139 | */
140 | async function preDownloadModel() {
141 | const pythonCmd = os.platform() === 'win32'
142 | ? path.join(VENV_DIR, 'Scripts', 'python.exe')
143 | : path.join(VENV_DIR, 'bin', 'python');
144 |
145 | const downloadScriptPath = path.join(PYTHON_SCRIPT_DIR, 'download_model.py');
146 |
147 | // Check if the download script already exists
148 | if (!fs.existsSync(downloadScriptPath)) {
149 | // Create a simple script to download the model
150 | const downloadScript = `
151 | import os
152 | from transformers import AutoTokenizer, AutoModel
153 |
154 | # Set the model name
155 | MODEL_NAME = "intfloat/multilingual-e5-large-instruct"
156 |
157 | print(f"Pre-downloading model: {MODEL_NAME}")
158 | print("This might take a few minutes...")
159 |
160 | # Download the model and tokenizer
161 | tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
162 | model = AutoModel.from_pretrained(MODEL_NAME)
163 |
164 | print(f"Model downloaded and cached at: {os.path.expanduser('~/.cache/huggingface/transformers')}")
165 | print("Setup completed successfully!")
166 | `;
167 |
168 | fs.writeFileSync(downloadScriptPath, downloadScript);
169 | }
170 |
171 | log('Pre-downloading the E5 model (this may take a few minutes)...', colors.blue);
172 | return new Promise((resolve, reject) => {
173 | const proc = spawn(pythonCmd, [downloadScriptPath]);
174 |
175 | proc.stdout.on('data', (data) => {
176 | process.stdout.write(data.toString());
177 | });
178 |
179 | proc.stderr.on('data', (data) => {
180 | process.stderr.write(data.toString());
181 | });
182 |
183 | proc.on('close', (code) => {
184 | if (code === 0) {
185 | log('E5 model downloaded successfully.', colors.green);
186 | resolve();
187 | } else {
188 | log(`Failed to download model (exit code ${code}).`, colors.red);
189 | // Don't reject since this is not critical - the model will be downloaded on first use
190 | resolve();
191 | }
192 | });
193 |
194 | proc.on('error', (err) => {
195 | log(`Error downloading model: ${err.message}`, colors.red);
196 | // Don't reject since this is not critical - the model will be downloaded on first use
197 | resolve();
198 | });
199 | });
200 | }
201 |
202 | /**
203 | * Main function
204 | */
205 | async function main() {
206 | log('Starting setup for semantic search...', colors.blue);
207 |
208 | // Check if Python is installed
209 | if (!commandExists(PYTHON_COMMAND)) {
210 | log(`${PYTHON_COMMAND} not found. Please install Python 3.8 or later.`, colors.red);
211 | process.exit(1);
212 | }
213 |
214 | try {
215 | // Create virtual environment
216 | await createVirtualEnv();
217 |
218 | // Install dependencies
219 | await installDependencies();
220 |
221 | // Pre-download model (optional)
222 | await preDownloadModel();
223 |
224 | log('Setup completed successfully!', colors.green);
225 | } catch (error) {
226 | log(`Setup failed: ${error.message}`, colors.red);
227 | process.exit(1);
228 | }
229 | }
230 |
231 | // Run the main function
232 | main();
```
--------------------------------------------------------------------------------
/src/lib/embedding-service.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { spawn } from 'child_process';
2 | import path from 'path';
3 | import os from 'os';
4 | import { logger } from './logger.js';
5 |
6 | /**
7 | * Service for generating embeddings using the E5 model
8 | */
9 | export class EmbeddingService {
10 | private pythonProcess: any;
11 | private initialized: boolean = false;
12 | private initializationPromise: Promise<void> | null = null;
13 | private readonly pythonScriptPath: string;
14 | private readonly pythonCommand: string;
15 | private readonly taskDescription: string;
16 |
17 | /**
18 | * Creates a new embedding service
19 | * @param options Configuration options
20 | */
21 | constructor(
22 | options: {
23 | pythonScriptPath?: string;
24 | taskDescription?: string;
25 | useVenv?: boolean;
26 | } = {}
27 | ) {
28 | this.pythonScriptPath =
29 | options.pythonScriptPath || path.join(process.cwd(), 'src', 'python', 'embedding_service.py');
30 | this.taskDescription =
31 | options.taskDescription || 'Given a document, find semantically similar documents';
32 |
33 | // Determine which Python executable to use based on the useVenv option
34 | const useVenv = options.useVenv !== undefined ? options.useVenv : true;
35 | if (useVenv) {
36 | const venvDir = path.join(process.cwd(), 'venv');
37 | this.pythonCommand =
38 | os.platform() === 'win32'
39 | ? path.join(venvDir, 'Scripts', 'python.exe')
40 | : path.join(venvDir, 'bin', 'python');
41 | } else {
42 | this.pythonCommand = os.platform() === 'win32' ? 'python' : 'python3';
43 | }
44 | }
45 |
46 | /**
47 | * Initializes the embedding service
48 | */
49 | async initialize(): Promise<void> {
50 | if (this.initialized) {
51 | return;
52 | }
53 |
54 | if (this.initializationPromise) {
55 | return this.initializationPromise;
56 | }
57 |
58 | this.initializationPromise = new Promise<void>((resolve, reject) => {
59 | try {
60 | logger.debug('Initializing embedding service', {
61 | scriptPath: this.pythonScriptPath,
62 | pythonCommand: this.pythonCommand,
63 | });
64 |
65 | // Spawn Python process
66 | this.pythonProcess = spawn(this.pythonCommand, [this.pythonScriptPath]);
67 |
68 | // Handle process exit
69 | this.pythonProcess.on('exit', (code: number) => {
70 | logger.error('Embedding service process exited', { code });
71 | this.initialized = false;
72 | this.pythonProcess = null;
73 | });
74 |
75 | // Handle process errors
76 | this.pythonProcess.on('error', (error: Error) => {
77 | logger.error('Error in embedding service process', error);
78 | reject(error);
79 | });
80 |
81 | // Log stderr output
82 | this.pythonProcess.stderr.on('data', (data: Buffer) => {
83 | logger.debug('Embedding service stderr:', { message: data.toString().trim() });
84 | });
85 |
86 | // Initialize the model
87 | this.sendCommand({ command: 'initialize' })
88 | .then(() => {
89 | this.initialized = true;
90 | logger.debug('Embedding service initialized successfully');
91 | resolve();
92 | })
93 | .catch((error) => {
94 | logger.error('Failed to initialize embedding service', error);
95 | reject(error);
96 | });
97 | } catch (error) {
98 | logger.error('Error initializing embedding service', error);
99 | reject(error);
100 | }
101 | });
102 |
103 | return this.initializationPromise;
104 | }
105 |
106 | /**
107 | * Generates an embedding for a single text
108 | * @param text The text to generate an embedding for
109 | * @returns The embedding vector
110 | */
111 | async generateEmbedding(text: string): Promise<number[]> {
112 | if (!this.initialized) {
113 | await this.initialize();
114 | }
115 |
116 | const result = await this.sendCommand({
117 | command: 'generate_embedding',
118 | text,
119 | task: this.taskDescription,
120 | });
121 |
122 | if (result.error) {
123 | throw new Error(`Embedding generation failed: ${result.error}`);
124 | }
125 |
126 | return result.embedding;
127 | }
128 |
129 | /**
130 | * Generate embeddings for multiple texts
131 | * @param texts Array of texts to generate embeddings for
132 | * @param options Options for embedding generation
133 | * @returns Promise that resolves to an array of embedding vectors
134 | */
135 | async generateEmbeddings(
136 | texts: string[],
137 | options?: { input_type?: 'query' | 'passage' }
138 | ): Promise<number[][]> {
139 | if (!this.initialized) {
140 | await this.initialize();
141 | }
142 |
143 | if (!texts || texts.length === 0) {
144 | throw new Error('No texts provided for embedding generation');
145 | }
146 |
147 | logger.debug(`Generating embeddings for ${texts.length} texts`);
148 |
149 | try {
150 | const command = {
151 | command: 'generate_embeddings',
152 | texts: texts,
153 | is_query: options?.input_type === 'query',
154 | };
155 |
156 | const result = await this.sendCommand(command);
157 |
158 | if (result.error) {
159 | throw new Error(`Batch embedding generation failed: ${result.error}`);
160 | }
161 |
162 | if (!result.embeddings || !Array.isArray(result.embeddings)) {
163 | throw new Error('Invalid response from embedding service');
164 | }
165 |
166 | return result.embeddings;
167 | } catch (error: any) {
168 | logger.error('Error generating embeddings', error);
169 | throw new Error(`Batch embedding generation failed: ${error.message}`);
170 | }
171 | }
172 |
173 | /**
174 | * Sends a command to the Python process
175 | * @param command The command to send
176 | * @returns The result from the Python process
177 | */
178 | private sendCommand(command: any): Promise<any> {
179 | return new Promise((resolve, reject) => {
180 | if (!this.pythonProcess) {
181 | reject(new Error('Python process not initialized'));
182 | return;
183 | }
184 |
185 | // Buffer to accumulate response chunks
186 | let responseBuffer = '';
187 |
188 | // Set up data handler that accumulates chunks
189 | const responseHandler = (data: Buffer) => {
190 | const chunk = data.toString();
191 | responseBuffer += chunk;
192 |
193 | try {
194 | // Try to parse the accumulated buffer
195 | const response = JSON.parse(responseBuffer);
196 |
197 | // If successful parsing, clean up and resolve
198 | this.pythonProcess.stdout.removeListener('data', responseHandler);
199 | resolve(response);
200 | } catch (error) {
201 | // Incomplete JSON, continue accumulating
202 | // This is expected for large responses split across chunks
203 | // We'll keep collecting chunks until we get valid JSON
204 | }
205 | };
206 |
207 | // Handle error and end events
208 | const errorHandler = (error: Error) => {
209 | this.pythonProcess.stdout.removeListener('data', responseHandler);
210 | reject(error);
211 | };
212 |
213 | // Set up event listeners
214 | this.pythonProcess.stdout.on('data', responseHandler);
215 | this.pythonProcess.stdout.once('error', errorHandler);
216 |
217 | // Send command to Python process
218 | this.pythonProcess.stdin.write(JSON.stringify(command) + '\n');
219 |
220 | // Set a reasonable timeout (30 seconds)
221 | setTimeout(() => {
222 | this.pythonProcess.stdout.removeListener('data', responseHandler);
223 | reject(new Error('Timeout waiting for embedding service response'));
224 | }, 30000);
225 | });
226 | }
227 |
228 | /**
229 | * Closes the embedding service
230 | */
231 | async close(): Promise<void> {
232 | if (this.pythonProcess) {
233 | this.pythonProcess.kill();
234 | this.pythonProcess = null;
235 | this.initialized = false;
236 | this.initializationPromise = null;
237 | logger.debug('Embedding service closed');
238 | }
239 | }
240 | }
241 |
```
--------------------------------------------------------------------------------
/src/lib/http-server.ts:
--------------------------------------------------------------------------------
```typescript
1 | import http from 'http';
2 | import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
3 | import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js';
4 | import { logger } from './logger.js';
5 | import { URL } from 'url';
6 |
7 | // Store active SSE transports by session ID
8 | const activeTransports = new Map<string, SSEServerTransport>();
9 |
10 | // Store heartbeat intervals by session ID
11 | const heartbeatIntervals = new Map<string, NodeJS.Timeout>();
12 |
13 | /**
14 | * Get transport by session ID - useful for reconnection
15 | */
16 | export function getTransport(sessionId: string): SSEServerTransport | undefined {
17 | return activeTransports.get(sessionId);
18 | }
19 |
20 | /**
21 | * Set up an HTTP server with SSE transport for the MCP server
22 | */
23 | export function setupHttpServer(server: McpServer, port: number): http.Server {
24 | // Create HTTP server
25 | const httpServer = http.createServer((req, res) => {
26 | res.setHeader('Access-Control-Allow-Origin', '*');
27 | res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
28 | res.setHeader('Access-Control-Allow-Headers', 'Content-Type, Authorization');
29 |
30 | // Handle preflight requests
31 | if (req.method === 'OPTIONS') {
32 | res.writeHead(204);
33 | res.end();
34 | return;
35 | }
36 |
37 | // Parse URL
38 | const reqUrl = new URL(req.url || '/', `http://localhost:${port}`);
39 | const path = reqUrl.pathname;
40 |
41 | // Health check endpoint
42 | if (path === '/health' && req.method === 'GET') {
43 | res.writeHead(200, { 'Content-Type': 'application/json' });
44 | res.end(JSON.stringify({ status: 'ok' }));
45 | return;
46 | }
47 |
48 | // SSE endpoint
49 | if (path === '/sse' && req.method === 'GET') {
50 | logger.info('New SSE connection established');
51 |
52 | // Prevent timeout for long-running connections
53 | req.socket.setTimeout(0);
54 | req.socket.setKeepAlive(true);
55 | req.socket.setMaxListeners(20);
56 |
57 | // Check if there's a session ID for reconnection
58 | const urlParams = new URLSearchParams(reqUrl.search);
59 | const existingSessionId = urlParams.get('sessionId');
60 |
61 | // Create SSE transport with a fixed endpoint path
62 | const sseTransport = new SSEServerTransport('message', res);
63 |
64 | // If reconnecting with an existing session ID, store it
65 | if (existingSessionId) {
66 | logger.debug(`Attempting to reconnect with existing session ID: ${existingSessionId}`);
67 | }
68 |
69 | // Store the transport in our map
70 | const sessionId = sseTransport.sessionId;
71 | activeTransports.set(sessionId, sseTransport);
72 |
73 | logger.debug(`Created new SSE transport with session ID: ${sessionId}`);
74 |
75 | // Set up heartbeat to keep connection alive
76 | // This is critical for preventing connection timeouts and ensuring we don't lose message correlation
77 | // The heartbeat sends an empty comment every 30 seconds to keep the connection open
78 | const heartbeatInterval = setInterval(() => {
79 | if (!res.writableEnded) {
80 | try {
81 | // Send a comment as heartbeat
82 | res.write(`:heartbeat\n\n`);
83 | logger.debug(`Sent heartbeat to session ${sessionId}`);
84 | } catch (err) {
85 | logger.error(`Failed to send heartbeat to session ${sessionId}`, err);
86 | clearInterval(heartbeatInterval);
87 | activeTransports.delete(sessionId);
88 | }
89 | } else {
90 | // Connection already ended
91 | clearInterval(heartbeatInterval);
92 | activeTransports.delete(sessionId);
93 | }
94 | }, 30000);
95 |
96 | // Store the heartbeat interval
97 | heartbeatIntervals.set(sessionId, heartbeatInterval);
98 |
99 | // Connect server to transport
100 | server.connect(sseTransport).catch((error) => {
101 | logger.error('Failed to connect to SSE transport', error);
102 | clearInterval(heartbeatIntervals.get(sessionId)!);
103 | heartbeatIntervals.delete(sessionId);
104 | activeTransports.delete(sessionId);
105 | });
106 |
107 | // Handle connection close
108 | req.on('close', () => {
109 | logger.info(`SSE connection closed for session: ${sessionId}`);
110 |
111 | // Clean up resources
112 | const interval = heartbeatIntervals.get(sessionId);
113 | if (interval) {
114 | clearInterval(interval);
115 | heartbeatIntervals.delete(sessionId);
116 | }
117 |
118 | activeTransports.delete(sessionId);
119 | sseTransport.close().catch((error) => {
120 | logger.error('Error closing SSE transport', error);
121 | });
122 | });
123 |
124 | return;
125 | }
126 |
127 | // Message endpoint
128 | if (path === '/message' && req.method === 'POST') {
129 | const urlParams = new URLSearchParams(reqUrl.search);
130 | const sessionId = urlParams.get('sessionId');
131 |
132 | if (!sessionId) {
133 | res.writeHead(400, { 'Content-Type': 'application/json' });
134 | res.end(JSON.stringify({ error: 'Missing sessionId parameter' }));
135 | return;
136 | }
137 |
138 | const transport = activeTransports.get(sessionId);
139 | if (!transport) {
140 | res.writeHead(404, { 'Content-Type': 'application/json' });
141 | res.end(JSON.stringify({
142 | error: 'Session not found or expired',
143 | reconnect: true
144 | }));
145 | return;
146 | }
147 |
148 | // Pass the request to the SSE transport's handlePostMessage method
149 | try {
150 | transport.handlePostMessage(req, res);
151 | } catch (error) {
152 | logger.error(`Error handling message for session ${sessionId}`, error);
153 | if (!res.headersSent) {
154 | res.writeHead(500, { 'Content-Type': 'application/json' });
155 | res.end(JSON.stringify({
156 | error: 'Internal server error',
157 | details: error instanceof Error ? error.message : String(error)
158 | }));
159 | }
160 | }
161 |
162 | return;
163 | }
164 |
165 | // Not found
166 | res.writeHead(404, { 'Content-Type': 'application/json' });
167 | res.end(JSON.stringify({ error: 'Not found' }));
168 | });
169 |
170 | // Start HTTP server
171 | httpServer.listen(port, () => {
172 | logger.info(`HTTP server listening on port ${port}`);
173 | logger.info('Server ready to receive requests via HTTP');
174 | logger.info(`SSE endpoint: http://localhost:${port}/sse`);
175 | logger.info(`Message endpoint: http://localhost:${port}/message?sessionId=<SESSION_ID>`);
176 | });
177 |
178 | return httpServer;
179 | }
180 |
181 | /**
182 | * Close all active SSE transports
183 | */
184 | export async function closeAllTransports(): Promise<void> {
185 | logger.info(`Closing ${activeTransports.size} active SSE transports`);
186 |
187 | // Clear all heartbeat intervals
188 | for (const [sessionId, interval] of heartbeatIntervals.entries()) {
189 | clearInterval(interval);
190 | heartbeatIntervals.delete(sessionId);
191 | }
192 |
193 | // Use Promise.allSettled to ensure we attempt to close all transports
194 | // even if some fail
195 | const closePromises = Array.from(activeTransports.entries()).map(
196 | async ([sessionId, transport]) => {
197 | try {
198 | logger.debug(`Closing SSE transport for session: ${sessionId}`);
199 | await Promise.race([
200 | transport.close(),
201 | // Add a timeout to prevent hanging
202 | new Promise((_, reject) =>
203 | setTimeout(() => reject(new Error('Transport close timeout')), 1000)
204 | ),
205 | ]);
206 | return { sessionId, success: true };
207 | } catch (error) {
208 | logger.error(`Error closing SSE transport for session ${sessionId}`, error);
209 | return { sessionId, success: false };
210 | }
211 | }
212 | );
213 |
214 | await Promise.allSettled(closePromises);
215 |
216 | // Clear the map regardless of success/failure
217 | activeTransports.clear();
218 | logger.info('All SSE transports closed or timed out');
219 | }
220 |
```
--------------------------------------------------------------------------------
/tsconfig.json:
--------------------------------------------------------------------------------
```json
1 | {
2 | "compilerOptions": {
3 | /* Visit https://aka.ms/tsconfig to read more about this file */
4 |
5 | /* Projects */
6 | // "incremental": true, /* Save .tsbuildinfo files to allow for incremental compilation of projects. */
7 | // "composite": true, /* Enable constraints that allow a TypeScript project to be used with project references. */
8 | // "tsBuildInfoFile": "./.tsbuildinfo", /* Specify the path to .tsbuildinfo incremental compilation file. */
9 | // "disableSourceOfProjectReferenceRedirect": true, /* Disable preferring source files instead of declaration files when referencing composite projects. */
10 | // "disableSolutionSearching": true, /* Opt a project out of multi-project reference checking when editing. */
11 | // "disableReferencedProjectLoad": true, /* Reduce the number of projects loaded automatically by TypeScript. */
12 |
13 | /* Language and Environment */
14 | "target": "ES2020",
15 | // "lib": [], /* Specify a set of bundled library declaration files that describe the target runtime environment. */
16 | // "jsx": "preserve", /* Specify what JSX code is generated. */
17 | // "libReplacement": true, /* Enable lib replacement. */
18 | // "experimentalDecorators": true, /* Enable experimental support for legacy experimental decorators. */
19 | // "emitDecoratorMetadata": true, /* Emit design-type metadata for decorated declarations in source files. */
20 | // "jsxFactory": "", /* Specify the JSX factory function used when targeting React JSX emit, e.g. 'React.createElement' or 'h'. */
21 | // "jsxFragmentFactory": "", /* Specify the JSX Fragment reference used for fragments when targeting React JSX emit e.g. 'React.Fragment' or 'Fragment'. */
22 | // "jsxImportSource": "", /* Specify module specifier used to import the JSX factory functions when using 'jsx: react-jsx*'. */
23 | // "reactNamespace": "", /* Specify the object invoked for 'createElement'. This only applies when targeting 'react' JSX emit. */
24 | // "noLib": true, /* Disable including any library files, including the default lib.d.ts. */
25 | // "useDefineForClassFields": true, /* Emit ECMAScript-standard-compliant class fields. */
26 | // "moduleDetection": "auto", /* Control what method is used to detect module-format JS files. */
27 |
28 | /* Modules */
29 | "module": "NodeNext",
30 | "moduleResolution": "NodeNext",
31 | // "rootDir": "./", /* Specify the root folder within your source files. */
32 | // "baseUrl": "./", /* Specify the base directory to resolve non-relative module names. */
33 | // "paths": {}, /* Specify a set of entries that re-map imports to additional lookup locations. */
34 | // "rootDirs": [], /* Allow multiple folders to be treated as one when resolving modules. */
35 | // "typeRoots": [], /* Specify multiple folders that act like './node_modules/@types'. */
36 | // "types": [], /* Specify type package names to be included without being referenced in a source file. */
37 | // "allowUmdGlobalAccess": true, /* Allow accessing UMD globals from modules. */
38 | // "moduleSuffixes": [], /* List of file name suffixes to search when resolving a module. */
39 | // "allowImportingTsExtensions": true, /* Allow imports to include TypeScript file extensions. Requires '--moduleResolution bundler' and either '--noEmit' or '--emitDeclarationOnly' to be set. */
40 | // "rewriteRelativeImportExtensions": true, /* Rewrite '.ts', '.tsx', '.mts', and '.cts' file extensions in relative import paths to their JavaScript equivalent in output files. */
41 | // "resolvePackageJsonExports": true, /* Use the package.json 'exports' field when resolving package imports. */
42 | // "resolvePackageJsonImports": true, /* Use the package.json 'imports' field when resolving imports. */
43 | // "customConditions": [], /* Conditions to set in addition to the resolver-specific defaults when resolving imports. */
44 | // "noUncheckedSideEffectImports": true, /* Check side effect imports. */
45 | "resolveJsonModule": true,
46 | // "allowArbitraryExtensions": true, /* Enable importing files with any extension, provided a declaration file is present. */
47 | // "noResolve": true, /* Disallow 'import's, 'require's or '<reference>'s from expanding the number of files TypeScript should add to a project. */
48 |
49 | /* JavaScript Support */
50 | // "allowJs": true, /* Allow JavaScript files to be a part of your program. Use the 'checkJS' option to get errors from these files. */
51 | // "checkJs": true, /* Enable error reporting in type-checked JavaScript files. */
52 | // "maxNodeModuleJsDepth": 1, /* Specify the maximum folder depth used for checking JavaScript files from 'node_modules'. Only applicable with 'allowJs'. */
53 |
54 | /* Emit */
55 | "declaration": true,
56 | // "declarationMap": true, /* Create sourcemaps for d.ts files. */
57 | // "emitDeclarationOnly": true, /* Only output d.ts files and not JavaScript files. */
58 | "sourceMap": true,
59 | // "inlineSourceMap": true, /* Include sourcemap files inside the emitted JavaScript. */
60 | // "noEmit": true, /* Disable emitting files from a compilation. */
61 | // "outFile": "./", /* Specify a file that bundles all outputs into one JavaScript file. If 'declaration' is true, also designates a file that bundles all .d.ts output. */
62 | "outDir": "./dist",
63 | // "removeComments": true, /* Disable emitting comments. */
64 | // "importHelpers": true, /* Allow importing helper functions from tslib once per project, instead of including them per-file. */
65 | // "downlevelIteration": true, /* Emit more compliant, but verbose and less performant JavaScript for iteration. */
66 | // "sourceRoot": "", /* Specify the root path for debuggers to find the reference source code. */
67 | // "mapRoot": "", /* Specify the location where debugger should locate map files instead of generated locations. */
68 | // "inlineSources": true, /* Include source code in the sourcemaps inside the emitted JavaScript. */
69 | // "emitBOM": true, /* Emit a UTF-8 Byte Order Mark (BOM) in the beginning of output files. */
70 | // "newLine": "crlf", /* Set the newline character for emitting files. */
71 | // "stripInternal": true, /* Disable emitting declarations that have '@internal' in their JSDoc comments. */
72 | // "noEmitHelpers": true, /* Disable generating custom helper functions like '__extends' in compiled output. */
73 | // "noEmitOnError": true, /* Disable emitting files if any type checking errors are reported. */
74 | // "preserveConstEnums": true, /* Disable erasing 'const enum' declarations in generated code. */
75 | // "declarationDir": "./", /* Specify the output directory for generated declaration files. */
76 |
77 | /* Interop Constraints */
78 | // "isolatedModules": true, /* Ensure that each file can be safely transpiled without relying on other imports. */
79 | // "verbatimModuleSyntax": true, /* Do not transform or elide any imports or exports not marked as type-only, ensuring they are written in the output file's format based on the 'module' setting. */
80 | // "isolatedDeclarations": true, /* Require sufficient annotation on exports so other tools can trivially generate declaration files. */
81 | // "erasableSyntaxOnly": true, /* Do not allow runtime constructs that are not part of ECMAScript. */
82 | // "allowSyntheticDefaultImports": true, /* Allow 'import x from y' when a module doesn't have a default export. */
83 | "esModuleInterop": true,
84 | // "preserveSymlinks": true, /* Disable resolving symlinks to their realpath. This correlates to the same flag in node. */
85 | "forceConsistentCasingInFileNames": true,
86 |
87 | /* Type Checking */
88 | "strict": true,
89 | // "noImplicitAny": true, /* Enable error reporting for expressions and declarations with an implied 'any' type. */
90 | // "strictNullChecks": true, /* When type checking, take into account 'null' and 'undefined'. */
91 | // "strictFunctionTypes": true, /* When assigning functions, check to ensure parameters and the return values are subtype-compatible. */
92 | // "strictBindCallApply": true, /* Check that the arguments for 'bind', 'call', and 'apply' methods match the original function. */
93 | // "strictPropertyInitialization": true, /* Check for class properties that are declared but not set in the constructor. */
94 | // "strictBuiltinIteratorReturn": true, /* Built-in iterators are instantiated with a 'TReturn' type of 'undefined' instead of 'any'. */
95 | // "noImplicitThis": true, /* Enable error reporting when 'this' is given the type 'any'. */
96 | // "useUnknownInCatchVariables": true, /* Default catch clause variables as 'unknown' instead of 'any'. */
97 | // "alwaysStrict": true, /* Ensure 'use strict' is always emitted. */
98 | // "noUnusedLocals": true, /* Enable error reporting when local variables aren't read. */
99 | // "noUnusedParameters": true, /* Raise an error when a function parameter isn't read. */
100 | // "exactOptionalPropertyTypes": true, /* Interpret optional property types as written, rather than adding 'undefined'. */
101 | // "noImplicitReturns": true, /* Enable error reporting for codepaths that do not explicitly return in a function. */
102 | // "noFallthroughCasesInSwitch": true, /* Enable error reporting for fallthrough cases in switch statements. */
103 | // "noUncheckedIndexedAccess": true, /* Add 'undefined' to a type when accessed using an index. */
104 | // "noImplicitOverride": true, /* Ensure overriding members in derived classes are marked with an override modifier. */
105 | // "noPropertyAccessFromIndexSignature": true, /* Enforces using indexed accessors for keys declared using an indexed type. */
106 | // "allowUnusedLabels": true, /* Disable error reporting for unused labels. */
107 | // "allowUnreachableCode": true, /* Disable error reporting for unreachable code. */
108 |
109 | /* Completeness */
110 | // "skipDefaultLibCheck": true, /* Skip type checking .d.ts files that are included with TypeScript. */
111 | "skipLibCheck": true
112 | },
113 | "include": ["src/**/*"],
114 | "exclude": ["node_modules", "dist"]
115 | }
116 |
```
--------------------------------------------------------------------------------
/src/lib/tool-implementations.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
2 | import { SQLiteDataStore } from './sqlite-store.js';
3 | import { createNamespace, deleteNamespace, listNamespaces } from './namespace-manager.js';
4 | import {
5 | storeContextItem,
6 | retrieveContextItemByKey,
7 | deleteContextItem,
8 | listContextItemKeys,
9 | } from './context-manager.js';
10 | import { toolSchemas } from './tool-definitions.js';
11 | import { logger } from './logger.js';
12 | import { normalizeError } from './errors.js';
13 | import { retrieveBySemantic } from './search-module.js';
14 |
15 | /**
16 | * Registers all tool implementations with the MCP server
17 | * @param server MCP server instance
18 | * @param db Database connection
19 | */
20 | export function registerTools(server: McpServer, db: SQLiteDataStore): void {
21 | server.tool(
22 | toolSchemas.retrieve_context_item_by_key.name,
23 | toolSchemas.retrieve_context_item_by_key.description,
24 | toolSchemas.retrieve_context_item_by_key.schema.shape,
25 | async ({ namespace, key }) => {
26 | try {
27 | logger.debug('Executing retrieve_context_item_by_key tool', { namespace, key });
28 | const result = await retrieveContextItemByKey(db, namespace, key);
29 |
30 | return {
31 | content: [{ type: 'text', text: JSON.stringify({ result }) }],
32 | };
33 | } catch (error: any) {
34 | const normalizedError = normalizeError(error);
35 | logger.error('Error retrieving context item', normalizedError, { namespace, key });
36 | return {
37 | isError: true,
38 | content: [
39 | { type: 'text', text: `Error retrieving context item: ${normalizedError.message}` },
40 | ],
41 | };
42 | }
43 | }
44 | );
45 |
46 | server.tool(
47 | toolSchemas.store_context_item.name,
48 | toolSchemas.store_context_item.description,
49 | toolSchemas.store_context_item.schema.shape,
50 | async ({ namespace, key, value }) => {
51 | try {
52 | logger.debug('Executing store_context_item tool', { namespace, key });
53 |
54 | // Try to parse the value as JSON or number before storing
55 | let parsedValue: any = value;
56 |
57 | // First try to parse as JSON
58 | try {
59 | // Only attempt to parse if the value looks like JSON (starts with { or [)
60 | if (value.trim().startsWith('{') || value.trim().startsWith('[')) {
61 | parsedValue = JSON.parse(value);
62 | logger.debug('Parsed value as JSON object', { namespace, key });
63 | }
64 | } catch (parseError: any) {
65 | // If JSON parsing fails, keep the original string value
66 | logger.debug('Failed to parse value as JSON, keeping as string', {
67 | namespace,
68 | key,
69 | error: parseError.message,
70 | });
71 | }
72 |
73 | // If it's still a string, try to parse as other types
74 | if (typeof parsedValue === 'string') {
75 | const trimmedValue = parsedValue.trim();
76 |
77 | // Check for number
78 | if (/^-?\d+(\.\d+)?$/.test(trimmedValue)) {
79 | const numValue = Number(trimmedValue);
80 | if (!isNaN(numValue)) {
81 | parsedValue = numValue;
82 | logger.debug('Parsed value as number', { namespace, key });
83 | }
84 | }
85 | // Check for boolean
86 | else if (trimmedValue === 'true' || trimmedValue === 'false') {
87 | parsedValue = trimmedValue === 'true';
88 | logger.debug('Parsed value as boolean', { namespace, key });
89 | }
90 | // Check for null
91 | else if (trimmedValue === 'null') {
92 | parsedValue = null;
93 | logger.debug('Parsed value as null', { namespace, key });
94 | }
95 | }
96 |
97 | // Check if the item already exists
98 | const existingItem = await retrieveContextItemByKey(db, namespace, key);
99 | if (existingItem) {
100 | // Item already exists, return as an error
101 | return {
102 | isError: true,
103 | content: [
104 | {
105 | type: 'text',
106 | text: `Error storing context item: An item with the key "${key}" already exists in namespace "${namespace}". Use update_context_item to modify existing items.`,
107 | },
108 | ],
109 | };
110 | }
111 |
112 | // Store the parsed value
113 | await storeContextItem(db, namespace, key, parsedValue);
114 | return {
115 | content: [{ type: 'text', text: JSON.stringify({ result: true }) }],
116 | };
117 | } catch (error: any) {
118 | const normalizedError = normalizeError(error);
119 | logger.error('Error storing context item', normalizedError, { namespace, key });
120 | return {
121 | isError: true,
122 | content: [
123 | { type: 'text', text: `Error storing context item: ${normalizedError.message}` },
124 | ],
125 | };
126 | }
127 | }
128 | );
129 |
130 | server.tool(
131 | toolSchemas.update_context_item.name,
132 | toolSchemas.update_context_item.description,
133 | toolSchemas.update_context_item.schema.shape,
134 | async ({ namespace, key, value }) => {
135 | try {
136 | logger.debug('Executing update_context_item tool', { namespace, key });
137 |
138 | // Check if the item exists
139 | const existingItem = await retrieveContextItemByKey(db, namespace, key);
140 | if (!existingItem) {
141 | // Item doesn't exist, return as an error
142 | return {
143 | isError: true,
144 | content: [
145 | {
146 | type: 'text',
147 | text: `Error updating context item: No item with the key "${key}" exists in namespace "${namespace}". Use store_context_item to create new items.`,
148 | },
149 | ],
150 | };
151 | }
152 |
153 | // Try to parse the value as JSON or number before storing
154 | let parsedValue: any = value;
155 |
156 | // First try to parse as JSON
157 | try {
158 | // Only attempt to parse if the value looks like JSON (starts with { or [)
159 | if (value.trim().startsWith('{') || value.trim().startsWith('[')) {
160 | parsedValue = JSON.parse(value);
161 | logger.debug('Parsed value as JSON object', { namespace, key });
162 | }
163 | } catch (parseError: any) {
164 | // If JSON parsing fails, keep the original string value
165 | logger.debug('Failed to parse value as JSON, keeping as string', {
166 | namespace,
167 | key,
168 | error: parseError.message,
169 | });
170 | }
171 |
172 | // If it's still a string, try to parse as other types
173 | if (typeof parsedValue === 'string') {
174 | const trimmedValue = parsedValue.trim();
175 |
176 | // Check for number
177 | if (/^-?\d+(\.\d+)?$/.test(trimmedValue)) {
178 | const numValue = Number(trimmedValue);
179 | if (!isNaN(numValue)) {
180 | parsedValue = numValue;
181 | logger.debug('Parsed value as number', { namespace, key });
182 | }
183 | }
184 | // Check for boolean
185 | else if (trimmedValue === 'true' || trimmedValue === 'false') {
186 | parsedValue = trimmedValue === 'true';
187 | logger.debug('Parsed value as boolean', { namespace, key });
188 | }
189 | // Check for null
190 | else if (trimmedValue === 'null') {
191 | parsedValue = null;
192 | logger.debug('Parsed value as null', { namespace, key });
193 | }
194 | }
195 |
196 | // Update the item
197 | await storeContextItem(db, namespace, key, parsedValue);
198 | return {
199 | content: [{ type: 'text', text: JSON.stringify({ result: true }) }],
200 | };
201 | } catch (error: any) {
202 | const normalizedError = normalizeError(error);
203 | logger.error('Error updating context item', normalizedError, { namespace, key });
204 | return {
205 | isError: true,
206 | content: [
207 | { type: 'text', text: `Error updating context item: ${normalizedError.message}` },
208 | ],
209 | };
210 | }
211 | }
212 | );
213 |
214 | server.tool(
215 | toolSchemas.delete_context_item.name,
216 | toolSchemas.delete_context_item.description,
217 | toolSchemas.delete_context_item.schema.shape,
218 | async ({ namespace, key }) => {
219 | try {
220 | logger.debug('Executing delete_context_item tool', { namespace, key });
221 | const deleted = await deleteContextItem(db, namespace, key);
222 |
223 | return {
224 | content: [
225 | {
226 | type: 'text',
227 | text: JSON.stringify({
228 | result: deleted, // Will be true if item was deleted, false if it didn't exist
229 | }),
230 | },
231 | ],
232 | };
233 | } catch (error: any) {
234 | const normalizedError = normalizeError(error);
235 | logger.error('Error deleting context item', normalizedError, { namespace, key });
236 | return {
237 | isError: true,
238 | content: [
239 | { type: 'text', text: `Error deleting context item: ${normalizedError.message}` },
240 | ],
241 | };
242 | }
243 | }
244 | );
245 |
246 | // Namespace Management
247 | server.tool(
248 | toolSchemas.create_namespace.name,
249 | toolSchemas.create_namespace.description,
250 | toolSchemas.create_namespace.schema.shape,
251 | async ({ namespace }) => {
252 | try {
253 | logger.debug('Executing create_namespace tool', { namespace });
254 | const created = await createNamespace(db, namespace);
255 |
256 | if (!created) {
257 | // Namespace already exists, return as an error
258 | return {
259 | isError: true,
260 | content: [
261 | {
262 | type: 'text',
263 | text: `Error creating namespace: A namespace with the name "${namespace}" already exists.`,
264 | },
265 | ],
266 | };
267 | }
268 |
269 | return {
270 | content: [
271 | {
272 | type: 'text',
273 | text: JSON.stringify({
274 | result: true,
275 | }),
276 | },
277 | ],
278 | };
279 | } catch (error: any) {
280 | const normalizedError = normalizeError(error);
281 | logger.error('Error creating namespace', normalizedError, { namespace });
282 | return {
283 | isError: true,
284 | content: [{ type: 'text', text: `Error creating namespace: ${normalizedError.message}` }],
285 | };
286 | }
287 | }
288 | );
289 |
290 | server.tool(
291 | toolSchemas.delete_namespace.name,
292 | toolSchemas.delete_namespace.description,
293 | toolSchemas.delete_namespace.schema.shape,
294 | async ({ namespace }) => {
295 | try {
296 | logger.debug('Executing delete_namespace tool', { namespace });
297 | const deleted = await deleteNamespace(db, namespace);
298 |
299 | return {
300 | content: [
301 | {
302 | type: 'text',
303 | text: JSON.stringify({
304 | result: deleted, // Will be true if namespace was deleted, false if it didn't exist
305 | }),
306 | },
307 | ],
308 | };
309 | } catch (error: any) {
310 | const normalizedError = normalizeError(error);
311 | logger.error('Error deleting namespace', normalizedError, { namespace });
312 | return {
313 | isError: true,
314 | content: [{ type: 'text', text: `Error deleting namespace: ${normalizedError.message}` }],
315 | };
316 | }
317 | }
318 | );
319 |
320 | server.tool(
321 | toolSchemas.list_namespaces.name,
322 | toolSchemas.list_namespaces.description,
323 | toolSchemas.list_namespaces.schema.shape,
324 | async () => {
325 | try {
326 | logger.debug('Executing list_namespaces tool');
327 | const namespaces = await listNamespaces(db);
328 |
329 | return {
330 | content: [{ type: 'text', text: JSON.stringify({ result: namespaces }) }],
331 | };
332 | } catch (error: any) {
333 | const normalizedError = normalizeError(error);
334 | logger.error('Error listing namespaces', normalizedError);
335 | return {
336 | isError: true,
337 | content: [{ type: 'text', text: `Error listing namespaces: ${normalizedError.message}` }],
338 | };
339 | }
340 | }
341 | );
342 |
343 | // Keys Listing
344 | server.tool(
345 | toolSchemas.list_context_item_keys.name,
346 | toolSchemas.list_context_item_keys.description,
347 | toolSchemas.list_context_item_keys.schema.shape,
348 | async ({ namespace }) => {
349 | try {
350 | logger.debug('Executing list_context_item_keys tool', { namespace });
351 | const keys = await listContextItemKeys(db, namespace);
352 |
353 | return {
354 | content: [{ type: 'text', text: JSON.stringify({ result: keys }) }],
355 | };
356 | } catch (error: any) {
357 | const normalizedError = normalizeError(error);
358 | logger.error('Error listing context item keys', normalizedError, { namespace });
359 | return {
360 | isError: true,
361 | content: [
362 | {
363 | type: 'text',
364 | text: `Error listing context item keys: ${normalizedError.message}`,
365 | },
366 | ],
367 | };
368 | }
369 | }
370 | );
371 |
372 | // Semantic search
373 | server.tool(
374 | 'retrieve_context_items_by_semantic_search',
375 | toolSchemas.retrieve_context_items_by_semantic_search.schema.shape,
376 | async ({ namespace, query, limit, similarity_threshold }) => {
377 | try {
378 | logger.debug('Executing retrieve_context_items_by_semantic_search tool', {
379 | namespace,
380 | query,
381 | limit,
382 | similarity_threshold,
383 | });
384 |
385 | const items = await retrieveBySemantic(db, namespace, query, {
386 | limit: limit,
387 | threshold: similarity_threshold,
388 | });
389 |
390 | logger.debug('Retrieved items by semantic search', {
391 | namespace,
392 | query,
393 | count: items.length,
394 | });
395 |
396 | return {
397 | content: [
398 | {
399 | type: 'text',
400 | text: JSON.stringify({
401 | result: items.map((item) => ({
402 | value: item.value,
403 | similarity: item.similarity,
404 | created_at: item.created_at,
405 | updated_at: item.updated_at,
406 | })),
407 | }),
408 | },
409 | ],
410 | };
411 | } catch (error: any) {
412 | const normalizedError = normalizeError(error);
413 | logger.error('Error retrieving items by semantic search', normalizedError, {
414 | namespace,
415 | query,
416 | });
417 | return {
418 | isError: true,
419 | content: [
420 | {
421 | type: 'text',
422 | text: `Error retrieving items by semantic search: ${normalizedError.message}`,
423 | },
424 | ],
425 | };
426 | }
427 | }
428 | );
429 | }
430 |
```
--------------------------------------------------------------------------------
/src/lib/sqlite-store.ts:
--------------------------------------------------------------------------------
```typescript
1 | import sqlite3 from 'sqlite3';
2 | import { promisify } from 'util';
3 | import fs from 'fs';
4 | import path from 'path';
5 | import { logger } from './logger.js';
6 | import { DatabaseError, normalizeError } from './errors.js';
7 | import { chunkTextBySemanticBoundaries } from './chunking-util.js';
8 |
9 | // Define types for database operations
10 | export interface ContextItem {
11 | namespace: string;
12 | key: string;
13 | value: any;
14 | created_at: string;
15 | updated_at: string;
16 | }
17 |
18 | export interface Database {
19 | run: (sql: string, ...params: any[]) => Promise<void>;
20 | get: (sql: string, ...params: any[]) => Promise<any>;
21 | all: (sql: string, ...params: any[]) => Promise<any[]>;
22 | close: () => Promise<void>;
23 | }
24 |
25 | export interface SQLiteDataStoreOptions {
26 | dbPath: string;
27 | // Add optional chunking flag
28 | enableChunking?: boolean;
29 | }
30 |
31 | // Define the DataStore interface
32 | interface DataStore {
33 | storeDataItem(namespace: string, key: string, value: any): Promise<void>;
34 | retrieveDataItem(namespace: string, key: string): Promise<any | null>;
35 | deleteDataItem(namespace: string, key: string): Promise<boolean>;
36 | dataItemExists(namespace: string, key: string): Promise<boolean>;
37 | listAllNamespaces(): Promise<string[]>;
38 | listContextItemKeys(namespace: string): Promise<any[]>;
39 | createNamespace(namespace: string): Promise<boolean>;
40 | deleteNamespace(namespace: string): Promise<boolean>;
41 | // Optional method for non-chunked embeddings
42 | retrieveDataItemsByEmbeddingSimilarity?(
43 | namespace: string,
44 | queryText: string,
45 | embeddingService: any,
46 | options: any
47 | ): Promise<any[]>;
48 | retrieveAllItemsInNamespace(namespace: string): Promise<any[]>;
49 | }
50 |
51 | export class SQLiteDataStore implements DataStore {
52 | private dbPath: string;
53 | private db: sqlite3.Database | null = null;
54 | // Add chunking option
55 | private enableChunking: boolean;
56 |
57 | constructor(options: SQLiteDataStoreOptions) {
58 | this.dbPath = options.dbPath;
59 | // Default to false for backward compatibility
60 | this.enableChunking = options.enableChunking ?? false;
61 | logger.debug(
62 | `SQLite data store initialized, chunking ${this.enableChunking ? 'enabled' : 'disabled'}`
63 | );
64 | }
65 |
66 | /**
67 | * Get or create a database connection
68 | */
69 | async getDb(): Promise<sqlite3.Database> {
70 | if (!this.db) {
71 | logger.debug('Creating new database connection');
72 | this.db = new sqlite3.Database(this.dbPath);
73 |
74 | // Enable foreign key constraints
75 | const run = promisify(this.db.run.bind(this.db)) as (
76 | sql: string,
77 | ...params: any[]
78 | ) => Promise<void>;
79 | await run('PRAGMA foreign_keys = ON');
80 | logger.debug('Foreign key constraints enabled');
81 |
82 | await this.initializeDatabase();
83 | }
84 | return this.db;
85 | }
86 |
87 | /**
88 | * Check if a data item exists
89 | */
90 | async dataItemExists(namespace: string, key: string): Promise<boolean> {
91 | const db = await this.getDb();
92 | const get = promisify(db.get.bind(db)) as (sql: string, ...params: any[]) => Promise<any>;
93 |
94 | try {
95 | const result = await get(
96 | 'SELECT 1 FROM context_items WHERE namespace = ? AND key = ?',
97 | namespace,
98 | key
99 | );
100 | return !!result;
101 | } catch (error) {
102 | logger.error(
103 | `Error checking if data item exists | Context: ${JSON.stringify({ namespace, key, error })}`
104 | );
105 | throw error;
106 | }
107 | }
108 |
109 | /**
110 | * Initializes the SQLite database and creates necessary tables if they don't exist
111 | * @param dbPath Path to the SQLite database file
112 | * @returns A database connection object
113 | */
114 | async initializeDatabase(): Promise<Database> {
115 | try {
116 | logger.info('Initializing database', { dbPath: this.dbPath });
117 |
118 | // Ensure the directory exists
119 | const dbDir = path.dirname(this.dbPath);
120 | if (!fs.existsSync(dbDir)) {
121 | logger.debug('Creating database directory', { dbDir });
122 | fs.mkdirSync(dbDir, { recursive: true });
123 | }
124 |
125 | // Create a new database connection
126 | const db = new sqlite3.Database(this.dbPath);
127 | logger.debug('Database connection established');
128 |
129 | // Promisify database methods with proper type assertions
130 | const run = promisify(db.run.bind(db)) as (sql: string, ...params: any[]) => Promise<void>;
131 | const get = promisify(db.get.bind(db)) as (sql: string, ...params: any[]) => Promise<any>;
132 | const all = promisify(db.all.bind(db)) as (sql: string, ...params: any[]) => Promise<any[]>;
133 | const close = promisify(db.close.bind(db)) as () => Promise<void>;
134 |
135 | // Create tables if they don't exist
136 | logger.debug("Creating tables if they don't exist");
137 | await run(`
138 | CREATE TABLE IF NOT EXISTS namespaces (
139 | namespace TEXT PRIMARY KEY,
140 | created_at TEXT NOT NULL
141 | )
142 | `);
143 |
144 | await run(`
145 | CREATE TABLE IF NOT EXISTS context_items (
146 | namespace TEXT NOT NULL,
147 | key TEXT NOT NULL,
148 | value TEXT NOT NULL,
149 | created_at TEXT NOT NULL,
150 | updated_at TEXT NOT NULL,
151 | PRIMARY KEY (namespace, key),
152 | FOREIGN KEY (namespace) REFERENCES namespaces(namespace) ON DELETE CASCADE
153 | )
154 | `);
155 |
156 | // Create embeddings table for chunked items
157 | await run(`
158 | CREATE TABLE IF NOT EXISTS embeddings (
159 | id INTEGER PRIMARY KEY AUTOINCREMENT,
160 | namespace TEXT NOT NULL,
161 | item_key TEXT NOT NULL,
162 | chunk_index INTEGER NOT NULL,
163 | chunk_text TEXT NOT NULL,
164 | embedding TEXT,
165 | created_at TEXT NOT NULL,
166 | FOREIGN KEY (namespace, item_key) REFERENCES context_items(namespace, key) ON DELETE CASCADE
167 | )
168 | `);
169 |
170 | // Create indexes for embeddings table
171 | await run(
172 | 'CREATE INDEX IF NOT EXISTS idx_embeddings_namespace_item_key ON embeddings(namespace, item_key)'
173 | );
174 | await run(
175 | 'CREATE INDEX IF NOT EXISTS idx_embeddings_has_embedding ON embeddings(namespace, embedding IS NOT NULL)'
176 | );
177 |
178 | // Create indexes for performance
179 | logger.debug("Creating indexes if they don't exist");
180 | await run(
181 | 'CREATE INDEX IF NOT EXISTS idx_context_items_namespace ON context_items(namespace)'
182 | );
183 |
184 | logger.info('Database initialization completed successfully');
185 | return {
186 | run,
187 | get,
188 | all,
189 | close,
190 | };
191 | } catch (error) {
192 | const normalizedError = normalizeError(error);
193 | logger.error('Failed to initialize database', normalizedError, { dbPath: this.dbPath });
194 | throw new DatabaseError(`Failed to initialize database: ${normalizedError.message}`);
195 | }
196 | }
197 |
198 | /**
199 | * Store a data item with optional chunking support
200 | */
201 | async storeDataItem(namespace: string, key: string, value: any): Promise<void> {
202 | const db = await this.getDb();
203 | const run = promisify(db.run.bind(db)) as (sql: string, ...params: any[]) => Promise<void>;
204 | const all = promisify(db.all.bind(db)) as (sql: string, ...params: any[]) => Promise<any[]>;
205 |
206 | logger.debug(`Storing data item | Context: ${JSON.stringify({ namespace, key })}`);
207 |
208 | // Check if the item already exists
209 | const exists = await this.dataItemExists(namespace, key);
210 |
211 | const timestamp = new Date().toISOString();
212 |
213 | try {
214 | logger.debug(
215 | `Beginning transaction for storing data item | Context: ${JSON.stringify({ namespace, key })}`
216 | );
217 | await run('BEGIN TRANSACTION');
218 |
219 | if (exists) {
220 | logger.debug(`Updating existing item | Context: ${JSON.stringify({ namespace, key })}`);
221 | await run(
222 | `UPDATE context_items
223 | SET value = ?, updated_at = ?
224 | WHERE namespace = ? AND key = ?`,
225 | JSON.stringify(value),
226 | timestamp,
227 | namespace,
228 | key
229 | );
230 |
231 | // If chunking is enabled, delete existing chunks before adding new ones
232 | if (this.enableChunking) {
233 | logger.debug(
234 | `Deleting existing chunks for updated item | Context: ${JSON.stringify({ namespace, key })}`
235 | );
236 | await run('DELETE FROM embeddings WHERE namespace = ? AND item_key = ?', namespace, key);
237 | }
238 | } else {
239 | logger.debug(`Inserting new item | Context: ${JSON.stringify({ namespace, key })}`);
240 | await run(
241 | `INSERT INTO context_items
242 | (namespace, key, value, created_at, updated_at)
243 | VALUES (?, ?, ?, ?, ?)`,
244 | namespace,
245 | key,
246 | JSON.stringify(value),
247 | timestamp,
248 | timestamp
249 | );
250 | }
251 |
252 | // Process chunking if enabled
253 | if (this.enableChunking) {
254 | // Extract the text value for chunking - stringify objects
255 | const textValue = typeof value === 'string' ? value : JSON.stringify(value);
256 |
257 | // Create chunks based on semantic boundaries
258 | const chunks = chunkTextBySemanticBoundaries(textValue);
259 |
260 | logger.debug(
261 | `Storing ${chunks.length} chunks for item | Context: ${JSON.stringify({ namespace, key })}`
262 | );
263 |
264 | // Store each chunk with its index
265 | for (let i = 0; i < chunks.length; i++) {
266 | await run(
267 | `INSERT INTO embeddings
268 | (namespace, item_key, chunk_index, chunk_text, created_at)
269 | VALUES (?, ?, ?, ?, ?)`,
270 | namespace,
271 | key,
272 | i,
273 | chunks[i],
274 | timestamp
275 | );
276 | }
277 |
278 | // Immediately generate embeddings for the newly created chunks
279 | try {
280 | // Import the embedding service dynamically to avoid circular dependencies
281 | const { EmbeddingService } = await import('./embedding-service.js');
282 | const embeddingService = new EmbeddingService();
283 | await embeddingService.initialize();
284 |
285 | logger.debug(
286 | `Generating embeddings for chunks | Context: ${JSON.stringify({ namespace, key, count: chunks.length })}`
287 | );
288 |
289 | // Get the chunk IDs for the newly inserted chunks
290 | const chunkIds = await all(
291 | `SELECT id, chunk_text FROM embeddings
292 | WHERE namespace = ? AND item_key = ? AND embedding IS NULL
293 | ORDER BY chunk_index`,
294 | namespace,
295 | key
296 | );
297 |
298 | if (chunkIds.length > 0) {
299 | // Extract just the text for embedding generation
300 | const chunkTexts = chunkIds.map((chunk) => chunk.chunk_text);
301 |
302 | // Generate embeddings
303 | const embeddings = await embeddingService.generateEmbeddings(chunkTexts);
304 |
305 | // Update each chunk with its embedding
306 | for (let i = 0; i < chunkIds.length; i++) {
307 | const embedding = JSON.stringify(embeddings[i]);
308 | await run(
309 | `UPDATE embeddings SET embedding = ? WHERE id = ?`,
310 | embedding,
311 | chunkIds[i].id
312 | );
313 | }
314 |
315 | logger.debug(
316 | `Embeddings generated and stored for ${chunkIds.length} chunks | Context: ${JSON.stringify({ namespace, key })}`
317 | );
318 | }
319 | } catch (error) {
320 | // Log the error but don't fail the store operation
321 | logger.error(
322 | `Error generating embeddings for chunks | Context: ${JSON.stringify({ namespace, key, error })}`
323 | );
324 | }
325 | }
326 |
327 | await run('COMMIT');
328 | logger.debug(
329 | `Transaction committed successfully | Context: ${JSON.stringify({ namespace, key })}`
330 | );
331 | } catch (error) {
332 | logger.error(
333 | `Error storing data item | Context: ${JSON.stringify({ namespace, key, error })}`
334 | );
335 | await run('ROLLBACK');
336 | throw error;
337 | }
338 | }
339 |
340 | /**
341 | * Generate embeddings for all chunks that don't have them yet
342 | */
343 | async generateEmbeddingsForChunks(namespace: string, embeddingService: any): Promise<number> {
344 | if (!this.enableChunking) {
345 | logger.debug(`Chunking is disabled, skipping embedding generation`);
346 | return 0;
347 | }
348 |
349 | const db = await this.getDb();
350 | const all = promisify(db.all.bind(db)) as (sql: string, ...params: any[]) => Promise<any[]>;
351 | const run = promisify(db.run.bind(db)) as (sql: string, ...params: any[]) => Promise<void>;
352 |
353 | // Find chunks without embeddings
354 | const chunksWithoutEmbeddings = await all(
355 | `SELECT id, namespace, item_key, chunk_index, chunk_text
356 | FROM embeddings
357 | WHERE namespace = ? AND embedding IS NULL
358 | ORDER BY item_key, chunk_index`,
359 | namespace
360 | );
361 |
362 | if (chunksWithoutEmbeddings.length === 0) {
363 | logger.debug(`No chunks found without embeddings for namespace: ${namespace}`);
364 | return 0;
365 | }
366 |
367 | logger.debug(
368 | `Generating embeddings for ${chunksWithoutEmbeddings.length} chunks | Context: { namespace: ${namespace} }`
369 | );
370 |
371 | let processedCount = 0;
372 |
373 | // Process chunks in batches to avoid memory issues
374 | const BATCH_SIZE = 10;
375 | for (let i = 0; i < chunksWithoutEmbeddings.length; i += BATCH_SIZE) {
376 | const batch = chunksWithoutEmbeddings.slice(i, i + BATCH_SIZE);
377 | const chunkTexts = batch.map((chunk) => chunk.chunk_text);
378 |
379 | try {
380 | // Generate embeddings for the batch
381 | const embeddings = await embeddingService.generateEmbeddings(chunkTexts);
382 |
383 | // Update each chunk with its embedding
384 | await run('BEGIN TRANSACTION');
385 | for (let j = 0; j < batch.length; j++) {
386 | const chunk = batch[j];
387 | const embedding = JSON.stringify(embeddings[j]);
388 |
389 | await run(
390 | `UPDATE embeddings
391 | SET embedding = ?
392 | WHERE id = ?`,
393 | embedding,
394 | chunk.id
395 | );
396 | processedCount++;
397 | }
398 | await run('COMMIT');
399 |
400 | logger.debug(
401 | `Generated embeddings for batch ${i / BATCH_SIZE + 1} | Context: { count: ${batch.length} }`
402 | );
403 | } catch (error) {
404 | logger.error(
405 | `Error generating embeddings for batch ${i / BATCH_SIZE + 1} | Context: ${JSON.stringify(error)}`
406 | );
407 | await run('ROLLBACK');
408 | throw error;
409 | }
410 | }
411 |
412 | logger.debug(`Completed embedding generation | Context: { processed: ${processedCount} }`);
413 | return processedCount;
414 | }
415 |
416 | /**
417 | * Retrieve items by semantic search, using the chunked embeddings if enabled
418 | */
419 | async retrieveDataItemsBySemanticSearch(
420 | namespace: string,
421 | queryText: string,
422 | embeddingService: any,
423 | options: {
424 | limit?: number;
425 | similarityThreshold?: number;
426 | } = {}
427 | ): Promise<any[]> {
428 | logger.debug(
429 | `Performing semantic search | Context: ${JSON.stringify({ namespace, queryText, options })}`
430 | );
431 |
432 | const { limit = 10, similarityThreshold = 0.5 } = options;
433 |
434 | // If chunking is not enabled, fall back to the legacy method if it exists
435 | if (!this.enableChunking) {
436 | logger.debug(`Chunking disabled, using legacy semantic search method`);
437 | return this.retrieveDataItemsByEmbeddingSimilarity(
438 | namespace,
439 | queryText,
440 | embeddingService,
441 | options
442 | );
443 | }
444 |
445 | const db = await this.getDb();
446 | const all = promisify(db.all.bind(db)) as (sql: string, ...params: any[]) => Promise<any[]>;
447 |
448 | try {
449 | // Generate embedding for the query
450 | logger.debug(
451 | `Generating embedding for text | Context: ${JSON.stringify({ textLength: queryText.length })}`
452 | );
453 | const queryEmbedding = await embeddingService.generateEmbeddings([queryText], {
454 | input_type: 'query',
455 | });
456 |
457 | if (!queryEmbedding || !queryEmbedding[0] || queryEmbedding[0].length === 0) {
458 | logger.error(`Failed to generate query embedding`);
459 | return [];
460 | }
461 |
462 | // Get all chunks with embeddings
463 | const chunksWithEmbeddings = await all(
464 | `SELECT c.namespace, c.item_key, c.chunk_index, c.chunk_text, c.embedding,
465 | i.value, i.tags, i.description, i.created_at, i.updated_at
466 | FROM embeddings c
467 | JOIN context_items i ON c.namespace = i.namespace AND c.item_key = i.key
468 | WHERE c.namespace = ? AND c.embedding IS NOT NULL`,
469 | namespace
470 | );
471 |
472 | if (chunksWithEmbeddings.length === 0) {
473 | logger.debug(`No chunks with embeddings found for namespace: ${namespace}`);
474 | return [];
475 | }
476 |
477 | logger.debug(
478 | `Calculating similarities | Context: ${JSON.stringify({ itemCount: chunksWithEmbeddings.length })}`
479 | );
480 |
481 | // Calculate similarity scores
482 | const similarities = chunksWithEmbeddings.map((chunk) => {
483 | const chunkEmbedding = JSON.parse(chunk.embedding);
484 | const similarity = this.calculateCosineSimilarity(queryEmbedding[0], chunkEmbedding);
485 | return {
486 | ...chunk,
487 | similarity,
488 | };
489 | });
490 |
491 | // Filter by similarity threshold and sort by similarity
492 | const filteredResults = similarities
493 | .filter((item) => item.similarity >= similarityThreshold)
494 | .sort((a, b) => b.similarity - a.similarity);
495 |
496 | // Group by item to avoid duplicates, taking the highest similarity score
497 | const itemMap = new Map();
498 | for (const result of filteredResults) {
499 | const itemKey = `${result.namespace}:${result.item_key}`;
500 |
501 | if (!itemMap.has(itemKey) || itemMap.get(itemKey).similarity < result.similarity) {
502 | // Parse the value field
503 | try {
504 | result.value = JSON.parse(result.value);
505 | } catch (e) {
506 | // If parsing fails, keep the original value
507 | }
508 |
509 | // Parse tags
510 | try {
511 | result.tags = JSON.parse(result.tags);
512 | } catch (e) {
513 | result.tags = [];
514 | }
515 |
516 | itemMap.set(itemKey, result);
517 | }
518 | }
519 |
520 | // Convert back to array and take top results
521 | const results = Array.from(itemMap.values())
522 | .slice(0, limit)
523 | .map((item) => ({
524 | namespace: item.namespace,
525 | key: item.item_key,
526 | value: item.value,
527 | tags: item.tags,
528 | description: item.description,
529 | matchedChunk: item.chunk_text,
530 | similarity: item.similarity,
531 | created_at: item.created_at,
532 | updated_at: item.updated_at,
533 | }));
534 |
535 | logger.debug(
536 | `Retrieved items by semantic search | Context: ${JSON.stringify({ namespace, queryText, count: results.length })}`
537 | );
538 |
539 | return results;
540 | } catch (error) {
541 | logger.error(
542 | `Error retrieving items by semantic search | Context: ${JSON.stringify({ namespace, error })}`
543 | );
544 | throw error;
545 | }
546 | }
547 |
548 | /**
549 | * Legacy method for retrieving items by embedding similarity (without chunking)
550 | */
551 | async retrieveDataItemsByEmbeddingSimilarity(
552 | namespace: string,
553 | queryText: string,
554 | embeddingService: any,
555 | options: {
556 | limit?: number;
557 | similarityThreshold?: number;
558 | } = {}
559 | ): Promise<any[]> {
560 | logger.warn(`Legacy embedding similarity search called, but not implemented`);
561 | return []; // Empty implementation - this would be implemented in your system if needed
562 | }
563 |
564 | /**
565 | * Retrieve a data item by namespace and key
566 | */
567 | async retrieveDataItem(namespace: string, key: string): Promise<any | null> {
568 | const db = await this.getDb();
569 | const get = promisify(db.get.bind(db)) as (sql: string, ...params: any[]) => Promise<any>;
570 |
571 | try {
572 | logger.debug(`Retrieving data item | Context: ${JSON.stringify({ namespace, key })}`);
573 |
574 | const result = await get(
575 | 'SELECT * FROM context_items WHERE namespace = ? AND key = ?',
576 | namespace,
577 | key
578 | );
579 |
580 | if (!result) {
581 | logger.debug(`Data item not found | Context: ${JSON.stringify({ namespace, key })}`);
582 | return null;
583 | }
584 |
585 | try {
586 | // Parse JSON value
587 | result.value = JSON.parse(result.value);
588 | } catch (e) {
589 | // If parsing fails, keep the original value
590 | logger.warn(
591 | `Failed to parse value as JSON | Context: ${JSON.stringify({ namespace, key })}`
592 | );
593 | }
594 |
595 | logger.debug(
596 | `Data item retrieved successfully | Context: ${JSON.stringify({ namespace, key })}`
597 | );
598 | return result;
599 | } catch (error) {
600 | logger.error(
601 | `Error retrieving data item | Context: ${JSON.stringify({ namespace, key, error })}`
602 | );
603 | throw error;
604 | }
605 | }
606 |
607 | /**
608 | * Delete a data item
609 | */
610 | async deleteDataItem(namespace: string, key: string): Promise<boolean> {
611 | const db = await this.getDb();
612 | const run = promisify(db.run.bind(db)) as (sql: string, ...params: any[]) => Promise<void>;
613 |
614 | try {
615 | logger.debug(`Deleting data item | Context: ${JSON.stringify({ namespace, key })}`);
616 |
617 | await run('BEGIN TRANSACTION');
618 |
619 | // First delete any associated chunks if chunking is enabled
620 | if (this.enableChunking) {
621 | await run('DELETE FROM embeddings WHERE namespace = ? AND item_key = ?', namespace, key);
622 | }
623 |
624 | // Then delete the main item
625 | const result = await run(
626 | 'DELETE FROM context_items WHERE namespace = ? AND key = ?',
627 | namespace,
628 | key
629 | );
630 |
631 | await run('COMMIT');
632 |
633 | logger.debug(`Data item deleted | Context: ${JSON.stringify({ namespace, key, result })}`);
634 | return true; // SQLite doesn't return affected rows in the same way as other DBs
635 | } catch (error) {
636 | await run('ROLLBACK');
637 | logger.error(
638 | `Error deleting data item | Context: ${JSON.stringify({ namespace, key, error })}`
639 | );
640 | throw error;
641 | }
642 | }
643 |
644 | /**
645 | * List all namespaces
646 | */
647 | async listAllNamespaces(): Promise<string[]> {
648 | const db = await this.getDb();
649 | const all = promisify(db.all.bind(db)) as (sql: string, ...params: any[]) => Promise<any[]>;
650 |
651 | try {
652 | logger.debug('Listing all namespaces');
653 |
654 | const results = await all('SELECT namespace FROM namespaces ORDER BY namespace');
655 | const namespaces = results.map((row) => row.namespace);
656 |
657 | logger.debug(`Retrieved ${namespaces.length} namespaces`);
658 | return namespaces;
659 | } catch (error) {
660 | logger.error(`Error listing namespaces | Context: ${JSON.stringify(error)}`);
661 | throw error;
662 | }
663 | }
664 |
665 | /**
666 | * List context item keys
667 | */
668 | async listContextItemKeys(namespace: string): Promise<any[]> {
669 | const db = await this.getDb();
670 | const all = promisify(db.all.bind(db)) as (sql: string, ...params: any[]) => Promise<any[]>;
671 |
672 | try {
673 | logger.debug(`Listing context item keys | Context: ${JSON.stringify({ namespace })}`);
674 |
675 | const results = await all(
676 | 'SELECT key, created_at, updated_at FROM context_items WHERE namespace = ? ORDER BY key',
677 | namespace
678 | );
679 |
680 | logger.debug(
681 | `Retrieved keys for ${results.length} items | Context: ${JSON.stringify({ namespace })}`
682 | );
683 | return results;
684 | } catch (error) {
685 | logger.error(
686 | `Error listing context item keys | Context: ${JSON.stringify({ namespace, error })}`
687 | );
688 | throw error;
689 | }
690 | }
691 |
692 | /**
693 | * Create a namespace
694 | */
695 | async createNamespace(namespace: string): Promise<boolean> {
696 | const db = await this.getDb();
697 | const get = promisify(db.get.bind(db)) as (sql: string, ...params: any[]) => Promise<any>;
698 | const run = promisify(db.run.bind(db)) as (sql: string, ...params: any[]) => Promise<void>;
699 |
700 | try {
701 | logger.debug(`Creating namespace | Context: ${JSON.stringify({ namespace })}`);
702 |
703 | // Check if namespace already exists
704 | const existingNamespace = await get(
705 | 'SELECT 1 FROM namespaces WHERE namespace = ?',
706 | namespace
707 | );
708 |
709 | if (existingNamespace) {
710 | logger.debug(`Namespace already exists | Context: ${JSON.stringify({ namespace })}`);
711 | return false;
712 | }
713 |
714 | // Create the namespace
715 | await run(
716 | 'INSERT INTO namespaces (namespace, created_at) VALUES (?, ?)',
717 | namespace,
718 | new Date().toISOString()
719 | );
720 |
721 | logger.debug(`Namespace created successfully | Context: ${JSON.stringify({ namespace })}`);
722 | return true;
723 | } catch (error) {
724 | logger.error(`Error creating namespace | Context: ${JSON.stringify({ namespace, error })}`);
725 | throw error;
726 | }
727 | }
728 |
729 | /**
730 | * Delete a namespace
731 | */
732 | async deleteNamespace(namespace: string): Promise<boolean> {
733 | const db = await this.getDb();
734 | const get = promisify(db.get.bind(db)) as (sql: string, ...params: any[]) => Promise<any>;
735 | const run = promisify(db.run.bind(db)) as (sql: string, ...params: any[]) => Promise<void>;
736 |
737 | try {
738 | logger.debug(`Attempting to delete namespace | Context: ${JSON.stringify({ namespace })}`);
739 |
740 | // Check if namespace exists
741 | const existingNamespace = await get(
742 | 'SELECT 1 FROM namespaces WHERE namespace = ?',
743 | namespace
744 | );
745 |
746 | if (!existingNamespace) {
747 | logger.debug(
748 | `Namespace does not exist, nothing to delete | Context: ${JSON.stringify({ namespace })}`
749 | );
750 | return false;
751 | }
752 |
753 | // Delete using a transaction to ensure atomicity
754 | logger.debug(
755 | `Beginning transaction for namespace deletion | Context: ${JSON.stringify({ namespace })}`
756 | );
757 | await run('BEGIN TRANSACTION');
758 |
759 | // Delete all context items in namespace (cascades to embeddings due to foreign key)
760 | logger.debug(
761 | `Deleting all context items in namespace | Context: ${JSON.stringify({ namespace })}`
762 | );
763 | await run('DELETE FROM context_items WHERE namespace = ?', namespace);
764 |
765 | // Delete the namespace
766 | logger.debug(`Deleting namespace | Context: ${JSON.stringify({ namespace })}`);
767 | await run('DELETE FROM namespaces WHERE namespace = ?', namespace);
768 |
769 | await run('COMMIT');
770 |
771 | logger.debug(`Namespace deleted successfully | Context: ${JSON.stringify({ namespace })}`);
772 | return true;
773 | } catch (error) {
774 | await run('ROLLBACK');
775 | logger.error(`Error deleting namespace | Context: ${JSON.stringify({ namespace, error })}`);
776 | throw error;
777 | }
778 | }
779 |
780 | /**
781 | * Calculate cosine similarity between two embedding vectors
782 | */
783 | private calculateCosineSimilarity(vec1: number[], vec2: number[]): number {
784 | if (vec1.length !== vec2.length) {
785 | throw new Error(`Vector dimensions don't match: ${vec1.length} vs ${vec2.length}`);
786 | }
787 |
788 | let dotProduct = 0;
789 | let mag1 = 0;
790 | let mag2 = 0;
791 |
792 | for (let i = 0; i < vec1.length; i++) {
793 | dotProduct += vec1[i] * vec2[i];
794 | mag1 += vec1[i] * vec1[i];
795 | mag2 += vec2[i] * vec2[i];
796 | }
797 |
798 | mag1 = Math.sqrt(mag1);
799 | mag2 = Math.sqrt(mag2);
800 |
801 | if (mag1 === 0 || mag2 === 0) {
802 | return 0;
803 | }
804 |
805 | return dotProduct / (mag1 * mag2);
806 | }
807 |
808 | /**
809 | * Execute a SQL query with parameters and return a single row
810 | * (For compatibility with old Database interface)
811 | */
812 | async get(sql: string, ...params: any[]): Promise<any> {
813 | const db = await this.getDb();
814 | const get = promisify(db.get.bind(db)) as (sql: string, ...params: any[]) => Promise<any>;
815 |
816 | try {
817 | logger.debug(`Executing SQL get query: ${sql}`);
818 | const result = await get(sql, ...params);
819 | return result;
820 | } catch (error) {
821 | logger.error(`Error executing SQL get query: ${error}`);
822 | throw error;
823 | }
824 | }
825 |
826 | /**
827 | * Execute a SQL query with parameters and return all rows
828 | * (For compatibility with old Database interface)
829 | */
830 | async all(sql: string, ...params: any[]): Promise<any[]> {
831 | const db = await this.getDb();
832 | const all = promisify(db.all.bind(db)) as (sql: string, ...params: any[]) => Promise<any[]>;
833 |
834 | try {
835 | logger.debug(`Executing SQL all query: ${sql}`);
836 | const results = await all(sql, ...params);
837 | return results;
838 | } catch (error) {
839 | logger.error(`Error executing SQL all query: ${error}`);
840 | throw error;
841 | }
842 | }
843 |
844 | /**
845 | * Execute a SQL query with parameters (no return value)
846 | * (For compatibility with old Database interface)
847 | */
848 | async run(sql: string, ...params: any[]): Promise<void> {
849 | const db = await this.getDb();
850 | const run = promisify(db.run.bind(db)) as (sql: string, ...params: any[]) => Promise<void>;
851 |
852 | try {
853 | logger.debug(`Executing SQL run query: ${sql}`);
854 | await run(sql, ...params);
855 | } catch (error) {
856 | logger.error(`Error executing SQL run query: ${error}`);
857 | throw error;
858 | }
859 | }
860 |
861 | /**
862 | * Close the database connection
863 | * (For compatibility with old Database interface)
864 | */
865 | async close(): Promise<void> {
866 | if (this.db) {
867 | const db = this.db;
868 | await new Promise<void>((resolve, reject) => {
869 | db.close((err) => {
870 | if (err) {
871 | reject(err);
872 | } else {
873 | this.db = null;
874 | resolve();
875 | }
876 | });
877 | });
878 | }
879 | }
880 |
881 | /**
882 | * Retrieve all items in a namespace
883 | */
884 | async retrieveAllItemsInNamespace(namespace: string): Promise<any[]> {
885 | const db = await this.getDb();
886 | const all = promisify(db.all.bind(db)) as (sql: string, ...params: any[]) => Promise<any[]>;
887 |
888 | const results = await all(
889 | 'SELECT * FROM context_items WHERE namespace = ? ORDER BY key',
890 | namespace
891 | );
892 |
893 | return results.map((result) => {
894 | try {
895 | result.value = JSON.parse(result.value);
896 | } catch (e) {
897 | // If parsing fails, keep the original value
898 | }
899 |
900 | return result;
901 | });
902 | }
903 | }
904 |
```