@nixxie-cms/ai-rag 1.0.0

package/package.json

@@ -0,0 +1,37 @@
+{
+  "name": "@nixxie-cms/ai-rag",
+  "version": "1.0.0",
+  "license": "MIT",
+  "main": "dist/nixxie-cms-ai-rag.cjs.js",
+  "module": "dist/nixxie-cms-ai-rag.esm.js",
+  "exports": {
+    ".": {
+      "types": "./dist/nixxie-cms-ai-rag.cjs.js",
+      "module": "./dist/nixxie-cms-ai-rag.esm.js",
+      "default": "./dist/nixxie-cms-ai-rag.cjs.js"
+    },
+    "./package.json": "./package.json"
+  },
+  "dependencies": {
+    "@babel/runtime": "^7.24.7"
+  },
+  "devDependencies": {
+    "@nixxie-cms/core": "^1.1.0"
+  },
+  "peerDependencies": {
+    "@nixxie-cms/core": "^1.0.3"
+  },
+  "optionalDependencies": {
+    "@anthropic-ai/sdk": "^0.69.0",
+    "openai": "^6.0.0"
+  },
+  "preconstruct": {
+    "entrypoints": [
+      "index.ts"
+    ]
+  },
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/nixxiecms/nixxie/tree/main/packages/ai-rag"
+  }
+}

package/src/AiRagService.ts

@@ -0,0 +1,640 @@
+import type {
+  NixxieAiMessage,
+  NixxieAiRagService,
+  NixxieContext,
+  NixxieRagAnswer,
+  NixxieRagAskOptions,
+  NixxieRagChunk,
+  NixxieRagCitation,
+  NixxieRagDocument,
+  NixxieRagDocumentInput,
+  NixxieRagDocumentQuery,
+  NixxieRagIndexStats,
+  NixxieRagRetrieveOptions,
+  NixxieRagStreamEvent,
+} from '@nixxie-cms/core'
+import { chunkText } from './chunking'
+import { checkGrounding, resolveGuard, shouldRefuseForNoContext, type ResolvedGuard } from './guard'
+import {
+  resolveEmbeddingProvider,
+  resolveGenerationProvider,
+  type EmbeddingProvider,
+  type GenerationProvider,
+  type RagGenerateOptions,
+} from './providers'
+import { ALLOW_KNOWLEDGE_SYSTEM_PROMPT, buildRagPrompt, DEFAULT_SYSTEM_PROMPT } from './prompt'
+import { Semaphore } from './semaphore'
+import type {
+  AiRagConfig,
+  RagGenerationConfig,
+  VectorRecord,
+  VectorStore,
+} from './types'
+import { SqlVectorStore } from './vector-store'
+
+type Resolved = {
+  documentsCollection: string
+  chunksCollection: string
+  generation: Required<Pick<RagGenerationConfig, 'temperature' | 'maxTokens'>> & {
+    systemPrompt: string
+    topP?: number
+    model?: string
+    buildPrompt?: RagGenerationConfig['buildPrompt']
+  }
+  embeddingBatchSize: number
+  embeddingModel?: string
+  topK: number
+  minScore: number
+  maxContextChars: number
+  candidateMultiplier: number
+  chunkSize: number
+  chunkOverlap: number
+  chunkStrategy: 'recursive' | 'sentence' | 'fixed'
+  historyLimit: number
+  maxQueryChars: number
+  guard: ResolvedGuard
+  indexing: {
+    auto: boolean
+    onConnect: boolean
+    schedule?: string | number
+    concurrency: number
+  }
+}
+
+function snippet(text: string, max = 240): string {
+  const t = text.trim().replace(/\s+/g, ' ')
+  return t.length > max ? `${t.slice(0, max)}…` : t
+}
+
+/** The main RAG assistant. Create via `createAiRag()`; register via `ragPlugin()`. */
+export class AiRagService implements NixxieAiRagService {
+  private documents: any | null = null
+  private generation: GenerationProvider
+  private embedder: EmbeddingProvider
+  private vectorStore: VectorStore
+  private rerank?: (query: string, chunks: NixxieRagChunk[]) => Promise<NixxieRagChunk[]>
+  private chatGate: Semaphore
+  private r: Resolved
+
+  constructor(config: AiRagConfig = {}) {
+    const documentsCollection = config.collections?.documents ?? 'KnowledgeBase'
+    const chunksCollection = config.collections?.chunks ?? 'KnowledgeChunk'
+
+    this.generation = resolveGenerationProvider(config.generation)
+    this.embedder = resolveEmbeddingProvider(config.embedding)
+    this.vectorStore = config.retrieval?.vectorStore ?? new SqlVectorStore(chunksCollection)
+    this.rerank = config.retrieval?.rerank
+
+    const guard = resolveGuard(config.guard)
+    const defaultSystem = guard.allowModelKnowledge
+      ? ALLOW_KNOWLEDGE_SYSTEM_PROMPT
+      : DEFAULT_SYSTEM_PROMPT
+
+    this.r = {
+      documentsCollection,
+      chunksCollection,
+      generation: {
+        temperature: config.generation?.temperature ?? 0.2,
+        maxTokens: config.generation?.maxTokens ?? 1024,
+        topP: config.generation?.topP,
+        model: config.generation?.model,
+        systemPrompt: config.generation?.systemPrompt ?? defaultSystem,
+        buildPrompt: config.generation?.buildPrompt,
+      },
+      embeddingBatchSize: config.embedding?.batchSize ?? 64,
+      embeddingModel: config.embedding?.model,
+      topK: config.retrieval?.topK ?? 5,
+      minScore: config.retrieval?.minScore ?? 0.2,
+      maxContextChars: config.retrieval?.maxContextChars ?? 6000,
+      candidateMultiplier: config.retrieval?.candidateMultiplier ?? 4,
+      chunkSize: config.chunking?.chunkSize ?? 1200,
+      chunkOverlap: config.chunking?.chunkOverlap ?? 200,
+      chunkStrategy: config.chunking?.strategy ?? 'recursive',
+      historyLimit: config.chat?.historyLimit ?? 10,
+      maxQueryChars: config.limits?.maxQueryChars ?? 8000,
+      guard,
+      indexing: {
+        auto: config.indexing?.auto ?? true,
+        onConnect: config.indexing?.onConnect ?? true,
+        schedule: config.indexing?.schedule,
+        concurrency: config.indexing?.concurrency ?? 4,
+      },
+    }
+    this.chatGate = new Semaphore(config.limits?.maxConcurrentChats ?? 8)
+  }
+
+  /** Names of the collections this service reads/writes (used by the plugin). */
+  get collections(): { documents: string; chunks: string } {
+    return { documents: this.r.documentsCollection, chunks: this.r.chunksCollection }
+  }
+
+  get indexingSchedule(): string | number | undefined {
+    return this.r.indexing.schedule
+  }
+
+  // ── Lifecycle ──
+
+  async init(context: NixxieContext): Promise<void> {
+    this.documents = this.requireDelegate(context, this.r.documentsCollection)
+    await this.vectorStore.init?.(context)
+    if (this.r.indexing.onConnect) {
+      // Index anything still pending without blocking boot on failures.
+      this.indexPending().catch(err =>
+        console.error('[@nixxie-cms/ai-rag] Initial indexing failed:', err)
+      )
+    }
+  }
+
+  private requireDelegate(context: NixxieContext, listKey: string): any {
+    const delegate = (context.prisma as any)?.[listKey[0].toLowerCase() + listKey.slice(1)]
+    if (!delegate) {
+      throw new Error(
+        `[@nixxie-cms/ai-rag] Collection "${listKey}" was not found in the Prisma client. ` +
+          `Register it via \`ragPlugin()\` (or add \`${listKey}: knowledgeBaseCollection()\`) and run a migration.`
+      )
+    }
+    return delegate
+  }
+
+  private requireDocuments(): any {
+    if (!this.documents) {
+      throw new Error(
+        '[@nixxie-cms/ai-rag] Not ready yet — the knowledge base is available once the database has connected.'
+      )
+    }
+    return this.documents
+  }
+
+  // ── Knowledge-base content ──
+
+  async addDocument(doc: NixxieRagDocumentInput): Promise<NixxieRagDocument> {
+    const [created] = await this.addDocuments([doc])
+    return created!
+  }
+
+  async addDocuments(docs: NixxieRagDocumentInput[]): Promise<NixxieRagDocument[]> {
+    const model = this.requireDocuments()
+    const out: NixxieRagDocument[] = []
+    for (const doc of docs) {
+      const row = await model.create({
+        data: {
+          title: doc.title ?? null,
+          content: doc.content,
+          source: doc.source ?? null,
+          tags: doc.tags ?? [],
+          metadata: doc.metadata ?? null,
+          status: 'pending',
+          chunkCount: 0,
+        },
+      })
+      out.push(rowToDocument(row))
+    }
+    if (this.r.indexing.auto) {
+      for (const d of out) {
+        try {
+          await this.index(d.id)
+        } catch (err) {
+          console.error(`[@nixxie-cms/ai-rag] Failed to index document ${d.id}:`, err)
+        }
+      }
+      // Re-read to reflect post-index status.
+      const refreshed = await Promise.all(out.map(d => this.getDocument(d.id)))
+      return refreshed.map((d, i) => d ?? out[i]!)
+    }
+    return out
+  }
+
+  async getDocument(id: string): Promise<NixxieRagDocument | undefined> {
+    const row = await this.requireDocuments().findUnique({ where: { id } })
+    return row ? rowToDocument(row) : undefined
+  }
+
+  async listDocuments(query: NixxieRagDocumentQuery = {}): Promise<NixxieRagDocument[]> {
+    const where: any = {}
+    if (query.status) where.status = { equals: query.status }
+    if (query.search) {
+      where.OR = [
+        { title: { contains: query.search } },
+        { content: { contains: query.search } },
+        { source: { contains: query.search } },
+      ]
+    }
+    const rows: any[] = await this.requireDocuments().findMany({
+      where,
+      orderBy: { createdAt: 'desc' },
+      skip: query.skip ?? 0,
+      ...(query.take !== undefined ? { take: query.take } : {}),
+    })
+    let docs = rows.map(rowToDocument)
+    // Tag filtering is done in Node for cross-database portability (tags are stored as JSON).
+    if (query.tags && query.tags.length) {
+      docs = docs.filter(d => query.tags!.every(t => (d.tags ?? []).includes(t)))
+    }
+    return docs
+  }
+
+  async updateDocument(
+    id: string,
+    patch: Partial<NixxieRagDocumentInput>
+  ): Promise<NixxieRagDocument> {
+    const model = this.requireDocuments()
+    const data: any = { status: 'pending' }
+    if (patch.title !== undefined) data.title = patch.title
+    if (patch.content !== undefined) data.content = patch.content
+    if (patch.source !== undefined) data.source = patch.source
+    if (patch.tags !== undefined) data.tags = patch.tags
+    if (patch.metadata !== undefined) data.metadata = patch.metadata
+    await model.update({ where: { id }, data })
+    if (this.r.indexing.auto) {
+      try {
+        await this.index(id)
+      } catch (err) {
+        console.error(`[@nixxie-cms/ai-rag] Failed to re-index document ${id}:`, err)
+      }
+    }
+    const doc = await this.getDocument(id)
+    if (!doc) throw new Error(`[@nixxie-cms/ai-rag] Document not found after update: ${id}`)
+    return doc
+  }
+
+  async removeDocument(id: string): Promise<void> {
+    await this.vectorStore.deleteByDocument(id)
+    await this.requireDocuments().delete({ where: { id } })
+  }
+
+  /**
+   * Delete a document's indexed chunks without touching the document row. Used by the
+   * auto-index delete hook, where the KB row has already been removed by the CMS.
+   */
+  async purgeChunks(documentId: string): Promise<void> {
+    await this.vectorStore.deleteByDocument(documentId)
+  }
+
+  // ── Indexing ──
+
+  async index(documentId: string): Promise<void> {
+    const model = this.requireDocuments()
+    const row = await model.findUnique({ where: { id: documentId } })
+    if (!row) throw new Error(`[@nixxie-cms/ai-rag] Document not found: ${documentId}`)
+    const doc = rowToDocument(row)
+    if (doc.status === 'disabled') {
+      await this.vectorStore.deleteByDocument(documentId)
+      await model.update({ where: { id: documentId }, data: { chunkCount: 0 } })
+      return
+    }
+
+    await model.update({ where: { id: documentId }, data: { status: 'indexing', error: null } })
+    try {
+      const text = [doc.title, doc.content].filter(Boolean).join('\n\n')
+      const pieces = chunkText(text, {
+        strategy: this.r.chunkStrategy,
+        chunkSize: this.r.chunkSize,
+        chunkOverlap: this.r.chunkOverlap,
+      })
+      const embeddings = await this.embedBatched(pieces)
+      const records: VectorRecord[] = pieces.map((content, i) => ({
+        id: `${documentId}:${i}`,
+        documentId,
+        content,
+        embedding: embeddings[i]!,
+        title: doc.title,
+        source: doc.source,
+        tags: doc.tags,
+        metadata: doc.metadata,
+      }))
+      await this.vectorStore.upsert(documentId, records)
+      await model.update({
+        where: { id: documentId },
+        data: {
+          status: 'indexed',
+          chunkCount: records.length,
+          error: null,
+          indexedAt: new Date(),
+        },
+      })
+    } catch (err) {
+      await model.update({
+        where: { id: documentId },
+        data: { status: 'error', error: err instanceof Error ? err.message : String(err) },
+      })
+      throw err
+    }
+  }
+
+  async reindex(options: { force?: boolean; tags?: string[] } = {}): Promise<NixxieRagIndexStats> {
+    const start = Date.now()
+    const docs = await this.listDocuments({ tags: options.tags })
+    const targets = options.force ? docs : docs.filter(d => d.status !== 'indexed')
+    let chunks = 0
+    let errors = 0
+    await mapWithConcurrency(targets, this.r.indexing.concurrency, async d => {
+      try {
+        await this.index(d.id)
+        const refreshed = await this.getDocument(d.id)
+        chunks += refreshed?.chunkCount ?? 0
+      } catch {
+        errors++
+      }
+    })
+    return { documents: targets.length, chunks, errors, durationMs: Date.now() - start }
+  }
+
+  async indexPending(): Promise<NixxieRagIndexStats> {
+    const start = Date.now()
+    const pending = (await this.listDocuments()).filter(
+      d => d.status === 'pending' || d.status === 'error'
+    )
+    let chunks = 0
+    let errors = 0
+    await mapWithConcurrency(pending, this.r.indexing.concurrency, async d => {
+      try {
+        await this.index(d.id)
+        const refreshed = await this.getDocument(d.id)
+        chunks += refreshed?.chunkCount ?? 0
+      } catch {
+        errors++
+      }
+    })
+    return { documents: pending.length, chunks, errors, durationMs: Date.now() - start }
+  }
+
+  private async embedBatched(texts: string[]): Promise<number[][]> {
+    if (texts.length === 0) return []
+    const out: number[][] = []
+    const size = this.r.embeddingBatchSize
+    for (let i = 0; i < texts.length; i += size) {
+      const batch = texts.slice(i, i + size)
+      out.push(...(await this.embedder.embed(batch, this.r.embeddingModel)))
+    }
+    return out
+  }
+
+  // ── Retrieval ──
+
+  async retrieve(query: string, options: NixxieRagRetrieveOptions = {}): Promise<NixxieRagChunk[]> {
+    const topK = options.topK ?? this.r.topK
+    const minScore = options.minScore ?? this.r.minScore
+    const [embedding] = await this.embedder.embed([query], this.r.embeddingModel)
+    if (!embedding) return []
+    const candidates = await this.vectorStore.query({
+      embedding,
+      topK: Math.max(topK, topK * this.r.candidateMultiplier),
+      tags: options.tags,
+      minScore,
+    })
+    let chunks: NixxieRagChunk[] = candidates.map(c => ({
+      id: c.id,
+      documentId: c.documentId,
+      title: c.title,
+      source: c.source,
+      content: c.content,
+      score: c.score,
+      tags: c.tags,
+      metadata: c.metadata,
+    }))
+    if (this.rerank) chunks = await this.rerank(query, chunks)
+    return chunks.slice(0, topK)
+  }
+
+  // ── Chat ──
+
+  async ask(question: string, options: NixxieRagAskOptions = {}): Promise<NixxieRagAnswer> {
+    return this.chat([{ role: 'user', content: question }], options)
+  }
+
+  async chat(messages: NixxieAiMessage[], options: NixxieRagAskOptions = {}): Promise<NixxieRagAnswer> {
+    const release = await this.chatGate.acquire()
+    try {
+      const prepared = await this.prepare(messages, options)
+      if (prepared.refusal) {
+        return {
+          text: prepared.refusal,
+          sources: [],
+          grounded: true,
+          refused: true,
+          model: this.generation.defaultModel,
+        }
+      }
+      const result = await this.generation.generate(prepared.messages, prepared.genOptions)
+      return this.finalize(result.text, result.model, result.usage, prepared.context, options)
+    } finally {
+      release()
+    }
+  }
+
+  async *stream(
+    messages: NixxieAiMessage[],
+    options: NixxieRagAskOptions = {}
+  ): AsyncIterable<NixxieRagStreamEvent> {
+    const release = await this.chatGate.acquire()
+    try {
+      const prepared = await this.prepare(messages, options)
+      if (prepared.refusal) {
+        const answer: NixxieRagAnswer = {
+          text: prepared.refusal,
+          sources: [],
+          grounded: true,
+          refused: true,
+          model: this.generation.defaultModel,
+        }
+        yield { type: 'token', token: prepared.refusal }
+        yield { type: 'done', answer }
+        return
+      }
+
+      yield { type: 'sources', sources: prepared.citations }
+
+      let text = ''
+      let model = this.generation.defaultModel
+      let usage: { inputTokens?: number; outputTokens?: number } | undefined
+      if (this.generation.stream) {
+        for await (const part of this.generation.stream(prepared.messages, prepared.genOptions)) {
+          if (part.delta) {
+            text += part.delta
+            yield { type: 'token', token: part.delta }
+          }
+          if (part.done) {
+            model = part.model ?? model
+            usage = part.usage
+          }
+        }
+      } else {
+        const result = await this.generation.generate(prepared.messages, prepared.genOptions)
+        text = result.text
+        model = result.model
+        usage = result.usage
+        yield { type: 'token', token: text }
+      }
+
+      const answer = await this.finalize(text, model, usage, prepared.context, options)
+      // If the grounding check rewrote the answer to a refusal, surface that token too.
+      if (answer.refused && answer.text !== text) yield { type: 'token', token: answer.text }
+      yield { type: 'done', answer }
+    } catch (err) {
+      yield { type: 'error', error: err instanceof Error ? err.message : String(err) }
+    } finally {
+      release()
+    }
+  }
+
+  /** Shared retrieve → guard → prompt pipeline for chat/stream. */
+  private async prepare(
+    messages: NixxieAiMessage[],
+    options: NixxieRagAskOptions
+  ): Promise<{
+    refusal?: string
+    context: NixxieRagChunk[]
+    citations: NixxieRagCitation[]
+    messages: { role: 'user' | 'assistant'; content: string }[]
+    genOptions: RagGenerateOptions
+  }> {
+    const turns = messages.filter(m => m.role !== 'system')
+    const lastUser = [...turns].reverse().find(m => m.role === 'user')
+    const question = lastUser?.content ?? ''
+    if (question.length > this.r.maxQueryChars) {
+      throw new Error(
+        `[@nixxie-cms/ai-rag] Query exceeds the ${this.r.maxQueryChars}-character limit.`
+      )
+    }
+
+    const guardOn = options.guard ?? this.r.guard.enabled
+    const minScore = options.minScore ?? this.r.minScore
+
+    const context =
+      options.context ??
+      (question ? await this.retrieve(question, options) : [])
+
+    if (guardOn && shouldRefuseForNoContext(context, this.r.guard, minScore)) {
+      return {
+        refusal: this.r.guard.refusal,
+        context: [],
+        citations: [],
+        messages: [],
+        genOptions: {},
+      }
+    }
+
+    const history = turns
+      .slice(0, -1)
+      .slice(-this.r.historyLimit)
+      .map(m => ({ role: m.role as 'user' | 'assistant', content: m.content }))
+
+    const systemBase = options.systemSuffix
+      ? `${this.r.generation.systemPrompt}\n\n${options.systemSuffix}`
+      : this.r.generation.systemPrompt
+
+    const built = this.r.generation.buildPrompt
+      ? this.r.generation.buildPrompt({
+          question,
+          history,
+          context,
+          systemPrompt: systemBase,
+          requireCitations: guardOn && this.r.guard.requireCitations,
+        })
+      : buildRagPrompt(
+          {
+            question,
+            history,
+            context,
+            systemPrompt: systemBase,
+            requireCitations: guardOn && this.r.guard.requireCitations,
+          },
+          { maxContextChars: this.r.maxContextChars }
+        )
+
+    const genOptions: RagGenerateOptions = {
+      system: built.system,
+      model: options.model ?? this.r.generation.model,
+      temperature: options.temperature ?? this.r.generation.temperature,
+      maxTokens: options.maxTokens ?? this.r.generation.maxTokens,
+      topP: this.r.generation.topP,
+    }
+
+    return {
+      context,
+      citations: toCitations(context),
+      messages: built.messages,
+      genOptions,
+    }
+  }
+
+  /** Apply the grounding check and assemble the final answer with citations. */
+  private async finalize(
+    text: string,
+    model: string,
+    usage: { inputTokens?: number; outputTokens?: number } | undefined,
+    context: NixxieRagChunk[],
+    options: NixxieRagAskOptions
+  ): Promise<NixxieRagAnswer> {
+    const guardOn = options.guard ?? this.r.guard.enabled
+    let grounded = true
+    if (guardOn && this.r.guard.groundingCheck) {
+      const check = await checkGrounding(
+        this.generation,
+        text,
+        context,
+        this.r.guard.groundingModel ?? this.r.generation.model
+      )
+      grounded = check.grounded
+      if (!grounded && this.r.guard.refuseWhenUngrounded) {
+        return {
+          text: this.r.guard.refusal,
+          sources: [],
+          grounded: false,
+          refused: true,
+          model,
+          usage,
+        }
+      }
+    }
+    return { text, sources: toCitations(context), grounded, refused: false, model, usage }
+  }
+
+  async close(): Promise<void> {
+    // No long-lived resources to release; indexing runs are awaited by their callers.
+  }
+}
+
+function toCitations(chunks: NixxieRagChunk[]): NixxieRagCitation[] {
+  return chunks.map(c => ({
+    documentId: c.documentId,
+    chunkId: c.id,
+    title: c.title,
+    source: c.source,
+    score: c.score,
+    snippet: snippet(c.content),
+  }))
+}
+
+function rowToDocument(row: any): NixxieRagDocument {
+  return {
+    id: String(row.id),
+    title: row.title ?? undefined,
+    content: row.content ?? '',
+    source: row.source ?? undefined,
+    tags: Array.isArray(row.tags) ? row.tags : (row.tags ?? undefined),
+    metadata: (row.metadata as Record<string, unknown>) ?? undefined,
+    status: row.status ?? 'pending',
+    chunkCount: row.chunkCount ?? 0,
+    error: row.error ?? undefined,
+    createdAt: row.createdAt ?? new Date(),
+    updatedAt: row.updatedAt ?? row.createdAt ?? new Date(),
+    indexedAt: row.indexedAt ?? undefined,
+  }
+}
+
+/** Run `fn` over items with at most `limit` in flight at once. */
+async function mapWithConcurrency<T>(
+  items: T[],
+  limit: number,
+  fn: (item: T) => Promise<void>
+): Promise<void> {
+  const queue = [...items]
+  const workers = Array.from({ length: Math.max(1, Math.min(limit, items.length)) }, async () => {
+    while (queue.length) {
+      const item = queue.shift()!
+      await fn(item)
+    }
+  })
+  await Promise.all(workers)
+}