@yolk-sdk/knowledge 0.0.1-canary.0

package/src/search.ts

@@ -0,0 +1,270 @@
+import { Effect } from 'effect'
+import type { KnowledgeChunk, KnowledgeDocument, KnowledgeSearchScope } from './documents.ts'
+import { KnowledgeEmbedder } from './embeddings.ts'
+import { KnowledgeSearchError } from './errors.ts'
+import { SearchIndexStore } from './search-store.ts'
+
+export type KnowledgeSearchInput = {
+  readonly scope: KnowledgeSearchScope
+  readonly query: string
+  readonly limit?: number
+  readonly minScore?: number
+  readonly contextChunks?: number
+  readonly mode?: KnowledgeSearchMode
+  readonly vectorLimit?: number
+  readonly textLimit?: number
+}
+
+export type KnowledgeSearchMode = 'vector' | 'hybrid'
+
+export type KnowledgeSearchScores = {
+  readonly vector?: number
+  readonly text?: number
+  readonly fused?: number
+}
+
+export type KnowledgeSearchResult = {
+  readonly chunk: KnowledgeChunk
+  readonly score: number
+  readonly document: KnowledgeDocument
+  readonly context?: ReadonlyArray<KnowledgeChunk>
+  readonly scores?: KnowledgeSearchScores
+}
+
+export type KnowledgeSearchContext = {
+  readonly query: string
+  readonly results: ReadonlyArray<KnowledgeSearchResult>
+  readonly text: string
+}
+
+export type KnowledgeSearcher = {
+  readonly search: (
+    input: KnowledgeSearchInput
+  ) => Effect.Effect<ReadonlyArray<KnowledgeSearchResult>, KnowledgeSearchError>
+}
+
+const packedResultText = (result: KnowledgeSearchResult) =>
+  result.context?.map(chunk => chunk.content).join('\n\n') ?? result.chunk.content
+
+export const packKnowledgeSearchContext = (query: string, results: ReadonlyArray<KnowledgeSearchResult>): KnowledgeSearchContext =>
+  ({
+    query,
+    results,
+    text: results.map(packedResultText).join('\n\n')
+  })
+
+const scopeIds = (scope: KnowledgeSearchScope): ReadonlyArray<string> => {
+  switch (scope._tag) {
+    case 'KnowledgeCollection':
+      return [scope.id]
+    case 'KnowledgeCollections':
+      return scope.ids
+  }
+}
+
+const defaultHybridCandidateLimit = (limit: number) => Math.max(limit * 5, 40)
+
+const reciprocalRankFusionScore = (rank: number, rankConstant: number) => 1 / (rankConstant + rank)
+
+export const fuseKnowledgeSearchResults = (input: {
+  readonly vectorResults: ReadonlyArray<KnowledgeSearchResult>
+  readonly textResults: ReadonlyArray<KnowledgeSearchResult>
+  readonly limit: number
+  readonly rankConstant?: number
+}): ReadonlyArray<KnowledgeSearchResult> => {
+  const rankConstant = input.rankConstant ?? 60
+  const fused = new Map<
+    string,
+    {
+      readonly result: KnowledgeSearchResult
+      readonly vector?: number
+      readonly text?: number
+      readonly fused: number
+    }
+  >()
+
+  const addResults = (results: ReadonlyArray<KnowledgeSearchResult>, kind: 'vector' | 'text') => {
+    results.forEach((result, index) => {
+      const rank = index + 1
+      const previous = fused.get(result.chunk.id)
+      const nextFused = (previous?.fused ?? 0) + reciprocalRankFusionScore(rank, rankConstant)
+      fused.set(result.chunk.id, {
+        result: previous?.result ?? result,
+        vector: kind === 'vector' ? result.score : previous?.vector,
+        text: kind === 'text' ? result.score : previous?.text,
+        fused: nextFused
+      })
+    })
+  }
+
+  addResults(input.vectorResults, 'vector')
+  addResults(input.textResults, 'text')
+
+  return Array.from(fused.values())
+    .sort((left, right) => right.fused - left.fused)
+    .slice(0, input.limit)
+    .map(item => ({
+      ...item.result,
+      score: item.fused,
+      scores: { vector: item.vector, text: item.text, fused: item.fused }
+    }))
+}
+
+const validateSearchInput = (input: KnowledgeSearchInput) => {
+  const query = input.query.trim()
+  const ids = scopeIds(input.scope)
+  const limit = input.limit ?? 10
+  const contextChunks = input.contextChunks ?? 0
+  const mode = input.mode ?? 'hybrid'
+  const vectorLimit = input.vectorLimit ?? (mode === 'hybrid' ? defaultHybridCandidateLimit(limit) : limit)
+  const textLimit = input.textLimit ?? defaultHybridCandidateLimit(limit)
+
+  if (query.length === 0) {
+    return Effect.fail(new KnowledgeSearchError({ message: 'Search query is empty', stage: 'embed' }))
+  }
+
+  if (ids.length === 0) {
+    return Effect.fail(new KnowledgeSearchError({ message: 'Search scope is empty', stage: 'store' }))
+  }
+
+  if (!Number.isInteger(limit) || limit < 1) {
+    return Effect.fail(new KnowledgeSearchError({ message: 'Search limit must be positive', stage: 'store' }))
+  }
+
+  if (!Number.isInteger(vectorLimit) || vectorLimit < 1) {
+    return Effect.fail(new KnowledgeSearchError({ message: 'Vector search limit must be positive', stage: 'store' }))
+  }
+
+  if (!Number.isInteger(textLimit) || textLimit < 1) {
+    return Effect.fail(new KnowledgeSearchError({ message: 'Text search limit must be positive', stage: 'store' }))
+  }
+
+  if (mode !== 'vector' && mode !== 'hybrid') {
+    return Effect.fail(new KnowledgeSearchError({ message: 'Search mode is invalid', stage: 'store' }))
+  }
+
+  if (!Number.isInteger(contextChunks) || contextChunks < 0) {
+    return Effect.fail(
+      new KnowledgeSearchError({ message: 'Context chunks must be zero or positive', stage: 'store' })
+    )
+  }
+
+  if (input.minScore !== undefined && !Number.isFinite(input.minScore)) {
+    return Effect.fail(
+      new KnowledgeSearchError({ message: 'Search minScore must be finite', stage: 'store' })
+    )
+  }
+
+  return Effect.succeed({ query, limit, contextChunks, mode, vectorLimit, textLimit })
+}
+
+const searchVectorChunks = (input: {
+  readonly scope: KnowledgeSearchScope
+  readonly query: string
+  readonly limit: number
+  readonly minScore?: number
+}) =>
+  Effect.gen(function* () {
+    const store = yield* SearchIndexStore
+    const embedder = yield* KnowledgeEmbedder
+    const embedding = yield* embedder
+      .embedQuery(input.query)
+      .pipe(
+        Effect.mapError(
+          error => new KnowledgeSearchError({ message: error.message, stage: 'embed', cause: error })
+        )
+      )
+
+    return yield* store
+      .searchChunks({
+        scope: input.scope,
+        embedding,
+        limit: input.limit,
+        minScore: input.minScore
+      })
+      .pipe(
+        Effect.mapError(
+          error => new KnowledgeSearchError({ message: error.message, stage: 'store', cause: error })
+        )
+      )
+  })
+
+const searchTextChunks = (input: {
+  readonly scope: KnowledgeSearchScope
+  readonly query: string
+  readonly limit: number
+}) =>
+  Effect.gen(function* () {
+    const store = yield* SearchIndexStore
+    return yield* store
+      .searchChunksByText({ scope: input.scope, query: input.query, limit: input.limit })
+      .pipe(
+        Effect.mapError(
+          error => new KnowledgeSearchError({ message: error.message, stage: 'store', cause: error })
+        )
+      )
+  })
+
+export const searchKnowledge = (input: KnowledgeSearchInput) =>
+  Effect.gen(function* () {
+    const valid = yield* validateSearchInput(input)
+    yield* Effect.annotateCurrentSpan({
+      'knowledge_search.query_length': valid.query.length,
+      'knowledge_search.limit': valid.limit,
+      'knowledge_search.mode': valid.mode,
+      'knowledge_search.context_chunks': valid.contextChunks,
+      'knowledge_search.scope_count': scopeIds(input.scope).length
+    })
+    const store = yield* SearchIndexStore
+    const results: ReadonlyArray<KnowledgeSearchResult> = valid.mode === 'vector'
+      ? yield* searchVectorChunks({
+        scope: input.scope,
+        query: valid.query,
+        limit: valid.vectorLimit,
+        minScore: input.minScore
+      })
+      : yield* Effect.gen(function* () {
+        const searches = yield* Effect.all(
+          {
+            vectorResults: searchVectorChunks({
+              scope: input.scope,
+              query: valid.query,
+              limit: valid.vectorLimit,
+              minScore: input.minScore
+            }),
+            textResults: searchTextChunks({
+              scope: input.scope,
+              query: valid.query,
+              limit: valid.textLimit
+            })
+          },
+          { concurrency: 'unbounded' }
+        )
+
+        return fuseKnowledgeSearchResults({
+          vectorResults: searches.vectorResults,
+          textResults: searches.textResults,
+          limit: valid.limit
+        })
+      })
+
+    if (valid.contextChunks === 0) {
+      return results
+    }
+
+    return yield* Effect.forEach(results, result =>
+      store
+        .getContextChunks({
+          collectionId: result.chunk.collectionId,
+          documentId: result.chunk.documentId,
+          position: result.chunk.position,
+          contextChunks: valid.contextChunks
+        })
+        .pipe(
+          Effect.map(context => ({ ...result, context })),
+          Effect.mapError(
+            error => new KnowledgeSearchError({ message: error.message, stage: 'store', cause: error })
+          )
+        )
+    )
+  }).pipe(Effect.withSpan('knowledge_search.search'))

package/src/store.ts

@@ -0,0 +1,53 @@
+import { Context } from 'effect'
+import type { Effect } from 'effect'
+import type { KnowledgeArtifact } from './artifacts.ts'
+import type { KnowledgeStoreError } from './errors.ts'
+import type { KnowledgeLink } from './links.ts'
+import type {
+  CreateKnowledgeRecordInput,
+  KnowledgeRecord,
+  KnowledgeScope,
+  UpdateKnowledgeRecordInput
+} from './records.ts'
+import type { KnowledgeProvenance } from './provenance.ts'
+import type { KnowledgeRepresentation } from './representations.ts'
+
+export type GetKnowledgeRecordInput = {
+  readonly scope: KnowledgeScope
+  readonly id: string
+}
+
+export type ListKnowledgeRecordsInput = {
+  readonly scope: KnowledgeScope
+  readonly limit: number
+}
+
+export type ListKnowledgeRecordsResult = {
+  readonly records: ReadonlyArray<KnowledgeRecord>
+}
+
+export type ListPinnedKnowledgeInput = {
+  readonly scope: KnowledgeScope
+  readonly limit: number
+}
+
+export type ListPinnedKnowledgeResult = {
+  readonly records: ReadonlyArray<KnowledgeRecord>
+  readonly representations: ReadonlyArray<KnowledgeRepresentation>
+}
+
+export type KnowledgeStoreApi = {
+  readonly createRecord: (input: CreateKnowledgeRecordInput) => Effect.Effect<KnowledgeRecord, KnowledgeStoreError>
+  readonly updateRecord: (input: UpdateKnowledgeRecordInput) => Effect.Effect<KnowledgeRecord, KnowledgeStoreError>
+  readonly getRecord: (input: GetKnowledgeRecordInput) => Effect.Effect<KnowledgeRecord, KnowledgeStoreError>
+  readonly listRecords: (input: ListKnowledgeRecordsInput) => Effect.Effect<ListKnowledgeRecordsResult, KnowledgeStoreError>
+  readonly listPinned: (input: ListPinnedKnowledgeInput) => Effect.Effect<ListPinnedKnowledgeResult, KnowledgeStoreError>
+  readonly deleteRecord: (input: GetKnowledgeRecordInput) => Effect.Effect<void, KnowledgeStoreError>
+  readonly listArtifacts: (input: GetKnowledgeRecordInput) => Effect.Effect<ReadonlyArray<KnowledgeArtifact>, KnowledgeStoreError>
+  readonly listProvenance: (input: GetKnowledgeRecordInput) => Effect.Effect<ReadonlyArray<KnowledgeProvenance>, KnowledgeStoreError>
+  readonly listLinks: (input: GetKnowledgeRecordInput) => Effect.Effect<ReadonlyArray<KnowledgeLink>, KnowledgeStoreError>
+}
+
+export class KnowledgeStore extends Context.Service<KnowledgeStore, KnowledgeStoreApi>()(
+  '@yolk-sdk/knowledge/KnowledgeStore'
+) {}

package/src/summarization.ts

@@ -0,0 +1,28 @@
+import { Context, Effect, Layer } from 'effect'
+import type { KnowledgeMetadata } from './documents.ts'
+import type { KnowledgeSummarizationError } from './errors.ts'
+
+export type KnowledgeDocumentSummary = {
+  readonly title?: string
+  readonly summary?: string
+}
+
+export type SummarizeKnowledgeDocumentInput = {
+  readonly content: string
+  readonly sourceTitle?: string
+  readonly metadata?: KnowledgeMetadata
+}
+
+export type KnowledgeSummarizerApi = {
+  readonly summarize: (
+    input: SummarizeKnowledgeDocumentInput
+  ) => Effect.Effect<KnowledgeDocumentSummary, KnowledgeSummarizationError>
+}
+
+export class KnowledgeSummarizer extends Context.Service<KnowledgeSummarizer, KnowledgeSummarizerApi>()(
+  '@yolk-sdk/knowledge/KnowledgeSummarizer'
+) {}
+
+export const NoopKnowledgeSummarizerLive = Layer.succeed(KnowledgeSummarizer, {
+  summarize: input => Effect.succeed({ title: input.sourceTitle })
+})

package/src/vector-store.ts

@@ -0,0 +1,6 @@
+export { SearchIndexStoreError as VectorStoreError } from './errors.ts'
+export type {
+  KnowledgeChunkSearchInput as VectorSearchQuery,
+  KnowledgeChunkSearchResult as VectorSearchResult,
+  ReplaceKnowledgeDocumentChunksInput as VectorRecord
+} from './search-store.ts'