openai-code/src/vectordb.test.mjs

An unofficial proxy layer that lets you use Anthropic Claude Code with any OpenAI API backend.

import { describe, it, expect, beforeAll, afterAll } from 'vitest';
import { matmul, addToIndex, removeByPath, search, loadFromDisk, storeToDisk } from './vectordb.mjs';
import { createEmbedding, normalizeEmbedding } from './openai.mjs'
import { copyFileSync, unlinkSync } from 'node:fs';

describe('Vector Database Tests', () => {

  let needle;
  beforeAll(async() => {
    // copy CLAUDE_VECTORDB.json to CLAUDE_VECTORDB.json.bak
    copyFileSync('CLAUDE_VECTORDB.json', 'CLAUDE_VECTORDB.json.bak');
    
    needle = await createEmbedding('We need to calculate a hash for the contents of a file');
  });

  it('should compute the dot product correctly', () => {
    const vectorA = [1, 2, 3];
    const vectorB = [4, 5, 6];
    const result = matmul(vectorA, vectorB);
    expect(result).toBe(32);
  });

  it('should normalize embeddings correctly', () => {
    const embedding = [3, 4];
    const normalized = normalizeEmbedding(embedding);
    expect(normalized).toEqual([0.6, 0.8]);
  });

  it('should add and remove documents from index', async() => {

    const md5Code = `import { createHash } from 'node:crypto';
export const md5 = (fileBuffer) => createHash('md5').update(fileBuffer).digest('hex')`
    const md5CodeEmbedding = await createEmbedding(md5Code)

    const utilsCode = `import { HttpsProxyAgent } from "https-proxy-agent";
import { writeFileSync } from 'node:fs';

export function getOpenAICommonOptions() {
  const options = {};
  if (process.env.PROXY_URL) {
    options.httpAgent = new HttpsProxyAgent(process.env.PROXY_URL);
  }
  return options;
}

export function logToFile(...args) {
  const jsonLog = JSON.stringify(args, null, 2);
  writeFileSync('log.txt', '', { flag: 'a' });
}

import { readFileSync } from 'node:fs';

// right, we don't use dotenv ;)
export const loadDotEnv = (filePath = '.env') => {
  try {
    const envVariables = {};
    const envContent = readFileSync(filePath, 'utf-8');

    for (const line of envContent.split('\n')) {
      // trim spaces and ignore comments
      const trimmedLine = line.trim();
      if (trimmedLine && !trimmedLine.startsWith('#')) {
        const [key, value] = trimmedLine.split('=');
        if (key && value !== undefined) {
          // remove quotes from value if present
          envVariables[key.trim()] = value.trim().replace(/^"|"$/g, '');
        }
      }
    }

    for (const [key, value] of Object.entries(envVariables)) {
      if (!process.env[key]) {
        process.env[key] = value
      }
    }
    return envVariables
  } catch (error) {
    return {}
  }
}

const secrets = loadDotEnv();

/** returns an environment variable, agnostic to the ESM or CJS runtime environment */
/** returns the CI_name invariant when running with CI=true (env) */
export const getEnv = (name) => {
  return secrets[name] || process.env[name] || import.meta?.env?.[name];
}

export const setEnv = (name, value) => {
  secrets[name] = value;
}`
    const utilsLibCodeEmbedding = await createEmbedding(utilsCode)

    addToIndex('src/md5.mjs', md5CodeEmbedding);
    addToIndex('src/fs.mjs', utilsLibCodeEmbedding);

    const searchResult = search({ embedding: needle })

    console.log('searchResult', searchResult);
    expect(searchResult.length).toBe(1);
    expect(searchResult[0].source).toBe(md5Code);
  });

  it('should load and store documents from disk', () => {
    const path = './';
    storeToDisk(path);
    loadFromDisk(path);
    expect(search({ embedding: [1, 2, 3] }).length).toBe(0); // do not find bs
    expect(search({ embedding: needle }).length).toBe(1); // find math.mjs
  });

  it('should remove documents from index and not find those results', () => {

    removeByPath('src/md5.mjs');

    const restSearchResult = search({ embedding: needle });
    console.log('restSearchResult', restSearchResult);
    expect(restSearchResult.length).toBe(0);
  });

  afterAll(() => {
    // restore CLAUDE_VECTORDB.json from CLAUDE_VECTORDB.json.bak
    copyFileSync('CLAUDE_VECTORDB.json.bak', 'CLAUDE_VECTORDB.json')
    unlinkSync('CLAUDE_VECTORDB.json.bak');
  });
});