License TL;DR Retrieval Benchmark

import { readFile } from 'node:fs/promises'
import { typegraphInit } from '@typegraph-ai/sdk'

type CorpusRow = {
  _id: string
  title?: string
  text: string
  metadata?: Record<string, unknown>
}

const typegraph = await typegraphInit({
  apiKey: process.env.TYPEGRAPH_API_KEY!,
  tenantId: process.env.TYPEGRAPH_TENANT_ID!,
})

const corpus = JSON.parse(
  await readFile('./data/license-tldr-retrieval/corpus.json', 'utf8'),
) as CorpusRow[]

const bucket = await typegraph.bucket.create({
  name: 'license-tldr-retrieval',
  indexDefaults: {
    chunkSize: 2048,
    chunkOverlap: 256,
    graphExtraction: false,
    deduplicateBy: ['content'],
  },
})

const documents = corpus.map((row) => ({
  id: row._id,
  name: row.title ?? row._id,
  content: row.text,
  metadata: {
    ...(row.metadata ?? {}),
    corpusId: row._id,
    documentName: row.title,
  },
}))

await typegraph.document.ingest(documents, { bucketId: bucket.id })

console.log(`Bucket ready: ${bucket.id}`)

import { readFile } from 'node:fs/promises'
import { typegraphInit } from '@typegraph-ai/sdk'

type QueryRow = { _id: string; text: string }

const typegraph = await typegraphInit({
  apiKey: process.env.TYPEGRAPH_API_KEY!,
  tenantId: process.env.TYPEGRAPH_TENANT_ID!,
})

const bucketId = process.env.TYPEGRAPH_BUCKET_ID!
const queries = JSON.parse(
  await readFile('./data/license-tldr-retrieval/queries.json', 'utf8'),
) as QueryRow[]

const run = new Map<string, string[]>()

for (const query of queries) {
  const response = await typegraph.search(query.text, {
    buckets: [bucketId],
    resources: ['documents'],
    weights: { semantic: 1, bm25: false, graph: false, recency: false },
    limit: 12,
  })

  run.set(
    query._id,
    response.results.chunks.map((chunk) =>
      String(chunk.metadata?.corpusId ?? chunk.document.name),
    ),
  )
}

// Score the deduplicated top 10 in `run` against qrels.json with nDCG@10, MAP@10, Recall@10, and P@10.

const response = await typegraph.search(query.text, {
  buckets: [process.env.TYPEGRAPH_BUCKET_ID!],
  resources: ['documents'],
  weights: { semantic: 1, bm25: false, graph: false, recency: false },
  limit: 12,
})

const retrievedCorpusIds = response.results.chunks.map((chunk) =>
  String(chunk.metadata?.corpusId ?? chunk.document.name),
)

const ingestCost =
  (embedIngestTokens / 1_000_000) * 0.12 +
  (ingestComputeMs / 3_600_000) * 0.52

const queryCost =
  (embedQueryTokens / 1_000_000) * 0.04 +
  (queryComputeMs / 3_600_000) * 0.52