wiki-graph-core 0.3.1

package/dist/reader/chunk-batch/parser.js

@@ -0,0 +1,640 @@
+import { z } from "zod";
+import {
+  GuaranteedParseValidationError,
+  ParsedJsonError,
+  requestGuaranteedJson
+} from "../../guaranteed/index.js";
+import {
+  ChunkImportance,
+  ChunkRetention,
+  expectChunkImportance,
+  expectChunkRetention
+} from "../../document/index.js";
+import {
+  resolveEvidenceSelectionList,
+  EvidenceResolver
+} from "../../evidence-selection/index.js";
+const MAX_CHOICE_RETRIES = 3;
+const chunkLinkSchema = z.object({
+  from: z.union([z.number().int(), z.string()]),
+  strength: z.string().optional(),
+  to: z.union([z.number().int(), z.string()])
+});
+const evidenceSelectionItemSchema = z.object({
+  quote: z.string().optional(),
+  sentence_id: z.string().optional()
+}).passthrough();
+const chunkEvidenceSchema = z.union([
+  z.record(z.string(), z.unknown()),
+  z.array(evidenceSelectionItemSchema)
+]);
+const userFocusedChunkSchema = z.object({
+  content: z.string(),
+  evidence: chunkEvidenceSchema.nullish(),
+  label: z.string(),
+  retention: z.enum([
+    ChunkRetention.Verbatim,
+    ChunkRetention.Detailed,
+    ChunkRetention.Focused,
+    ChunkRetention.Relevant
+  ]),
+  temp_id: z.string()
+}).passthrough();
+const userFocusedResponseSchema = z.object({
+  chunks: z.array(userFocusedChunkSchema),
+  fragment_summary: z.string(),
+  links: z.array(chunkLinkSchema)
+});
+const bookCoherenceChunkSchema = z.object({
+  content: z.string(),
+  evidence: chunkEvidenceSchema.nullish(),
+  importance: z.enum([
+    ChunkImportance.Critical,
+    ChunkImportance.Important,
+    ChunkImportance.Helpful
+  ]),
+  label: z.string(),
+  temp_id: z.string()
+}).passthrough();
+const importanceAnnotationSchema = z.object({
+  chunk_id: z.number().int(),
+  importance: z.enum([
+    ChunkImportance.Critical,
+    ChunkImportance.Important,
+    ChunkImportance.Helpful
+  ])
+});
+const bookCoherenceResponseSchema = z.object({
+  chunks: z.array(bookCoherenceChunkSchema),
+  importance_annotations: z.array(importanceAnnotationSchema),
+  links: z.array(chunkLinkSchema)
+});
+const choiceResponseSchema = z.object({
+  choice: z.string()
+});
+var ChunkMetadataField = /* @__PURE__ */ ((ChunkMetadataField2) => {
+  ChunkMetadataField2["Retention"] = "retention";
+  ChunkMetadataField2["Importance"] = "importance";
+  return ChunkMetadataField2;
+})(ChunkMetadataField || {});
+class ChunkBatchParser {
+  #choiceSystemPrompt;
+  #evidenceResolver = new EvidenceResolver();
+  #metadataField;
+  #projection;
+  #responseIntentClassifierPrompt;
+  #requestChoice;
+  #sentenceTextByKey;
+  #sentenceTextSource;
+  #wordsCountByKey;
+  #validImportanceChunkIds;
+  #visibleChunkIds;
+  constructor(input) {
+    this.#choiceSystemPrompt = input.choiceSystemPrompt;
+    this.#metadataField = input.metadataField;
+    this.#projection = input.projection;
+    this.#responseIntentClassifierPrompt = input.responseIntentClassifierPrompt;
+    this.#requestChoice = input.requestChoice;
+    this.#sentenceTextSource = input.sentenceTextSource;
+    this.#visibleChunkIds = input.visibleChunkIds;
+    this.#wordsCountByKey = createWordsCountRecord(
+      input.projection.sentences.map((sentence) => ({
+        sentenceId: sentence.sentenceId,
+        wordsCount: sentence.wordsCount
+      }))
+    );
+    this.#sentenceTextByKey = createSentenceTextRecord(input.projection);
+    this.#validImportanceChunkIds = input.validImportanceChunkIds === void 0 ? void 0 : createMembershipRecord(input.validImportanceChunkIds);
+  }
+  async parse(parsedData, input) {
+    const issues = [];
+    const chunks = [];
+    const tempIds = [];
+    const importanceAnnotations = "importance_annotations" in parsedData ? this.#validateImportanceAnnotations(
+      parsedData.importance_annotations,
+      issues
+    ) : void 0;
+    const fragmentSummary = "fragment_summary" in parsedData ? parsedData.fragment_summary : void 0;
+    const links = normalizeChunkLinks(parsedData.links);
+    for (const [index, data] of parsedData.chunks.entries()) {
+      const chunkIndex = index + 1;
+      const chunkIssues = [];
+      const label = data.label.trim();
+      const content = data.content.trim();
+      if (label === "") {
+        chunkIssues.push(
+          `Chunk #${chunkIndex}: Missing or empty "label" field`
+        );
+      }
+      if (content === "") {
+        chunkIssues.push(
+          `Chunk #${chunkIndex}: Missing or empty "content" field`
+        );
+      }
+      const [matchedSentenceIds, evidenceFailure] = await this.#resolveChunkEvidence({
+        chunkIndex,
+        chunkLabel: label,
+        data,
+        isLastGenerationAttempt: input.isLastGenerationAttempt
+      });
+      if (matchedSentenceIds.length === 0) {
+        if (evidenceFailure !== void 0) {
+          chunkIssues.push(
+            `Chunk #${chunkIndex} ("${label}"): ${evidenceFailure.message}`
+          );
+        } else {
+          chunkIssues.push(
+            `Chunk #${chunkIndex} ("${label}"): Missing evidence`
+          );
+        }
+      }
+      if (chunkIssues.length > 0) {
+        issues.push(...chunkIssues);
+        continue;
+      }
+      const primarySentenceId = matchedSentenceIds[0];
+      if (primarySentenceId === void 0) {
+        issues.push(
+          `Chunk #${chunkIndex} ("${label}"): Unable to resolve any sentence IDs`
+        );
+        continue;
+      }
+      const totalWordsCount = matchedSentenceIds.reduce((sum, sentenceId) => {
+        return sum + (this.#wordsCountByKey[createSentenceKey(sentenceId)] ?? 0);
+      }, 0);
+      if (this.#metadataField === "retention" /* Retention */) {
+        const chunkData = data;
+        chunks.push({
+          content,
+          generation: 0,
+          id: 0,
+          label,
+          links: [],
+          retention: expectChunkRetention(chunkData.retention),
+          sentenceId: primarySentenceId,
+          sentenceIds: [...matchedSentenceIds],
+          wordsCount: totalWordsCount
+        });
+      } else {
+        const chunkData = data;
+        chunks.push({
+          content,
+          generation: 0,
+          id: 0,
+          importance: expectChunkImportance(chunkData.importance),
+          label,
+          links: [],
+          sentenceId: primarySentenceId,
+          sentenceIds: [...matchedSentenceIds],
+          wordsCount: totalWordsCount
+        });
+      }
+      tempIds.push(data.temp_id);
+    }
+    const validLinks = this.#filterAndValidateLinks({
+      issues,
+      links,
+      tempIds,
+      visibleChunkIds: this.#visibleChunkIds
+    });
+    if (issues.length > 0) {
+      throw new ParsedJsonError(issues);
+    }
+    return {
+      chunkBatch: {
+        chunks,
+        links: validLinks,
+        orderCorrect: true,
+        tempIds,
+        ...importanceAnnotations === void 0 ? {} : { importanceAnnotations }
+      },
+      ...fragmentSummary === void 0 ? {} : { fragmentSummary }
+    };
+  }
+  async getChunkSourceSentences(chunk) {
+    const sourceSentences = [];
+    for (const sentenceId of chunk.sentenceIds) {
+      const sentenceKey = createSentenceKey(sentenceId);
+      const sentenceText = this.#sentenceTextByKey[sentenceKey];
+      if (sentenceText !== void 0) {
+        sourceSentences.push(sentenceText);
+        continue;
+      }
+      sourceSentences.push(
+        await this.#sentenceTextSource.getSentence(sentenceId)
+      );
+    }
+    return sourceSentences;
+  }
+  async #resolveChunkEvidence(input) {
+    const evidence = input.data.evidence;
+    if (evidence === void 0 || evidence === null) {
+      return [[], void 0];
+    }
+    if (!isRecord(evidence) && !Array.isArray(evidence)) {
+      return [
+        [],
+        {
+          candidates: [],
+          code: "invalid_evidence",
+          fieldName: "evidence",
+          message: `Chunk #${input.chunkIndex} ("${input.chunkLabel}"): evidence must be an object or array`
+        }
+      ];
+    }
+    const projectedSentences = this.#projection.sentences;
+    const candidateSentenceIds = projectedSentences.map(
+      (sentence) => sentence.sentenceId
+    );
+    const selectionSentences = projectedSentences.map((sentence, index) => ({
+      id: `S${index + 1}`,
+      sentenceId: sentence.sentenceId,
+      text: sentence.projectedText
+    }));
+    const selectionEvidence = createEvidenceSelectionList(evidence);
+    const [selectionResolution, selectionFailure] = selectionEvidence === void 0 ? [void 0, void 0] : resolveEvidenceSelectionList({
+      evidence: selectionEvidence,
+      sentences: selectionSentences
+    });
+    if (selectionResolution !== void 0) {
+      return [selectionResolution.sentenceIds, void 0];
+    }
+    const exactMatchSentenceIds = Array.isArray(evidence) ? [] : this.#resolveExactProjectionEvidence(evidence);
+    if (exactMatchSentenceIds.length > 0) {
+      return [exactMatchSentenceIds, void 0];
+    }
+    const candidateTexts = projectedSentences.map(
+      (sentence) => sentence.projectedText
+    );
+    const [resolution, failure] = Array.isArray(evidence) ? [void 0, void 0] : this.#evidenceResolver.resolve(
+      evidence,
+      candidateSentenceIds,
+      candidateTexts
+    );
+    if (resolution !== void 0) {
+      return [resolution.sentenceIds, void 0];
+    }
+    const fallbackFailure = selectionFailure === void 0 ? failure : toEvidenceResolutionFailure(selectionFailure, "evidence");
+    if (fallbackFailure === void 0) {
+      return [[], void 0];
+    }
+    const shouldUseChoice = fallbackFailure.code.startsWith("ambiguous") || fallbackFailure.code === "low_confidence" && input.isLastGenerationAttempt && fallbackFailure.candidates.length > 0;
+    if (!shouldUseChoice) {
+      return [[], fallbackFailure];
+    }
+    const choiceFieldName = toChoiceFieldName(fallbackFailure.fieldName);
+    if (choiceFieldName === void 0) {
+      return [[], fallbackFailure];
+    }
+    const [choiceCandidate, choiceFailure] = await this.#chooseAmbiguousCandidate({
+      candidates: fallbackFailure.candidates,
+      chunkData: input.data,
+      chunkIndex: input.chunkIndex,
+      chunkLabel: input.chunkLabel,
+      fieldName: choiceFieldName
+    });
+    if (choiceFailure !== void 0) {
+      if (choiceFailure.code === "choice_parse_failed_full_fragment") {
+        return [candidateSentenceIds, void 0];
+      }
+      return [[], choiceFailure];
+    }
+    if (choiceCandidate === void 0) {
+      return [
+        [],
+        {
+          candidates: fallbackFailure.candidates,
+          code: "choice_failed",
+          fieldName: fallbackFailure.fieldName,
+          message: `Second-stage choice failed for ${fallbackFailure.fieldName}: no candidate returned.`
+        }
+      ];
+    }
+    if (choiceFieldName === "evidence") {
+      return [[choiceCandidate.sentenceId], void 0];
+    }
+    if (Array.isArray(evidence)) {
+      return [[], fallbackFailure];
+    }
+    if (input.isLastGenerationAttempt) {
+      const [resolved2, resolveFailure2] = this.#evidenceResolver.resolveWithOverrides({
+        candidateSentenceIds,
+        candidateTexts,
+        evidence,
+        overrides: {
+          [choiceFieldName]: choiceCandidate
+        }
+      });
+      if (resolved2 !== void 0) {
+        return [resolved2.sentenceIds, void 0];
+      }
+      return [[], resolveFailure2];
+    }
+    const repairedEvidence = {
+      ...evidence,
+      [choiceFieldName]: {
+        mode: "full",
+        text: choiceCandidate.text
+      }
+    };
+    const [resolved, resolveFailure] = this.#evidenceResolver.resolve(
+      repairedEvidence,
+      candidateSentenceIds,
+      candidateTexts
+    );
+    if (resolved !== void 0) {
+      return [resolved.sentenceIds, void 0];
+    }
+    return [[], resolveFailure];
+  }
+  #resolveExactProjectionEvidence(evidence) {
+    const startValue = evidence.start_anchor ?? evidence.start;
+    const [startAnchor, startFailure] = this.#evidenceResolver.parseAnchor(
+      startValue,
+      "start_anchor"
+    );
+    if (startFailure !== void 0 || startAnchor?.text === void 0) {
+      return [];
+    }
+    const startMatch = expectSingleSpan(
+      this.#projection.findExactMatches(startAnchor.text)
+    );
+    if (startMatch === void 0) {
+      return [];
+    }
+    const endValue = evidence.end_anchor ?? evidence.end;
+    if (endValue === void 0) {
+      return this.#projection.resolveSentenceIds(startMatch);
+    }
+    const [endAnchor, endFailure] = this.#evidenceResolver.parseAnchor(
+      endValue,
+      "end_anchor"
+    );
+    if (endFailure !== void 0 || endAnchor?.text === void 0) {
+      return [];
+    }
+    const endMatch = expectSingleSpan(
+      this.#projection.findExactMatches(endAnchor.text)
+    );
+    if (endMatch === void 0 || endMatch.offset < startMatch.offset) {
+      return [];
+    }
+    return this.#projection.resolveSentenceIds({
+      length: endMatch.offset + endMatch.length - startMatch.offset,
+      offset: startMatch.offset
+    });
+  }
+  async #chooseAmbiguousCandidate(input) {
+    const messages = this.#buildChoiceMessages(input);
+    const candidateIds = input.candidates.map(
+      (candidate) => candidate.occurrenceId
+    );
+    try {
+      const choice = await requestGuaranteedJson({
+        maxRetries: MAX_CHOICE_RETRIES,
+        messages,
+        parse: (data) => {
+          const candidate = input.candidates.find(
+            (item) => item.occurrenceId === data.choice
+          );
+          if (candidate === void 0) {
+            throw new ParsedJsonError([
+              `Invalid choice "${data.choice}". Expected one of: ${candidateIds.join(", ")}`
+            ]);
+          }
+          return candidate;
+        },
+        responseIntentClassifierPrompt: this.#responseIntentClassifierPrompt,
+        request: this.#requestChoice,
+        schema: choiceResponseSchema
+      });
+      return [choice, void 0];
+    } catch (error) {
+      if (isParsedJsonValidationFailure(error)) {
+        return [
+          void 0,
+          {
+            candidates: input.candidates,
+            code: "choice_parse_failed_full_fragment",
+            fieldName: input.fieldName,
+            message: `Second-stage choice parse validation failed for ${input.fieldName}; falling back to the full fragment span.`
+          }
+        ];
+      }
+      return [
+        void 0,
+        {
+          candidates: input.candidates,
+          code: "choice_failed",
+          fieldName: input.fieldName,
+          message: `Second-stage choice failed for ${input.fieldName}: ${formatError(error)}`
+        }
+      ];
+    }
+  }
+  #buildChoiceMessages(input) {
+    return [
+      {
+        content: this.#choiceSystemPrompt,
+        role: "system"
+      },
+      {
+        content: `Previously generated chunk (do NOT rewrite it):
+\`\`\`json
+${JSON.stringify(input.chunkData, null, 2)}
+\`\`\`
+
+Resolve only this field: "${input.fieldName}" for chunk #${input.chunkIndex} [${input.chunkLabel}].
+Choose exactly one candidate occurrence ID from the list below.
+
+` + input.candidates.map(formatChoiceCandidate).join("\n"),
+        role: "user"
+      }
+    ];
+  }
+  #validateImportanceAnnotations(annotations, issues) {
+    if (annotations.length === 0) {
+      return [];
+    }
+    if (this.#validImportanceChunkIds === void 0) {
+      return annotations.map((annotation) => ({
+        chunkId: annotation.chunk_id,
+        importance: expectChunkImportance(annotation.importance)
+      }));
+    }
+    const result = [];
+    for (const annotation of annotations) {
+      if (!hasMembership(this.#validImportanceChunkIds, annotation.chunk_id)) {
+        issues.push(
+          `importance_annotations references unknown chunk_id ${annotation.chunk_id}`
+        );
+        continue;
+      }
+      result.push({
+        chunkId: annotation.chunk_id,
+        importance: expectChunkImportance(annotation.importance)
+      });
+    }
+    return result;
+  }
+  #filterAndValidateLinks(input) {
+    const validTempIds = createMembershipRecord(input.tempIds);
+    const validChunkIds = createMembershipRecord(input.visibleChunkIds);
+    const retainedLinks = [];
+    for (const [index, link] of input.links.entries()) {
+      const fromValid = this.#validateLinkReference({
+        fieldName: "from",
+        index: index + 1,
+        issues: input.issues,
+        reference: link.from,
+        validChunkIds,
+        validTempIds
+      });
+      const toValid = this.#validateLinkReference({
+        fieldName: "to",
+        index: index + 1,
+        issues: input.issues,
+        reference: link.to,
+        validChunkIds,
+        validTempIds
+      });
+      if (fromValid && toValid) {
+        retainedLinks.push(link);
+      }
+    }
+    return retainedLinks;
+  }
+  #validateLinkReference(input) {
+    if (typeof input.reference === "string") {
+      return hasMembership(input.validTempIds, input.reference);
+    }
+    if (!hasMembership(input.validChunkIds, input.reference)) {
+      input.issues.push(
+        `Link #${input.index}: "${input.fieldName}" chunk_id ${input.reference} does not exist in visible chunks`
+      );
+      return false;
+    }
+    return true;
+  }
+}
+function createWordsCountRecord(sentences) {
+  const wordsCountByKey = createEmptyRecord();
+  for (const sentence of sentences) {
+    const sentenceKey = createSentenceKey(sentence.sentenceId);
+    wordsCountByKey[sentenceKey] = sentence.wordsCount;
+  }
+  return wordsCountByKey;
+}
+function createSentenceTextRecord(projection) {
+  const record = createEmptyRecord();
+  for (const sentence of projection.sentences) {
+    record[createSentenceKey(sentence.sentenceId)] = sentence.rawText;
+  }
+  return record;
+}
+function normalizeChunkLinks(links) {
+  return links.map((link) => {
+    if (link.strength === void 0) {
+      return {
+        from: link.from,
+        to: link.to
+      };
+    }
+    return {
+      from: link.from,
+      strength: link.strength,
+      to: link.to
+    };
+  });
+}
+function formatChoiceCandidate(candidate) {
+  return [
+    candidate.occurrenceId,
+    `prev: ${formatChoiceText(candidate.prevText)}`,
+    `text: ${formatChoiceText(candidate.text)}`,
+    `next: ${formatChoiceText(candidate.nextText)}`
+  ].join("\n");
+}
+function formatChoiceText(text) {
+  const collapsed = text.replace(/\s+/gu, " ").trim();
+  return collapsed === "" ? "(none)" : collapsed;
+}
+function toChoiceFieldName(value) {
+  return value === "evidence" || value === "start_anchor" || value === "end_anchor" ? value : void 0;
+}
+function isRecord(value) {
+  return typeof value === "object" && value !== null;
+}
+function formatError(error) {
+  if (error instanceof Error) {
+    return error.message;
+  }
+  return String(error);
+}
+function isParsedJsonValidationFailure(error) {
+  return error instanceof GuaranteedParseValidationError && error.cause instanceof ParsedJsonError;
+}
+function createSentenceKey(sentenceId) {
+  return sentenceId.join(":");
+}
+function createMembershipRecord(values) {
+  const record = createEmptyRecord();
+  for (const value of values) {
+    record[String(value)] = true;
+  }
+  return record;
+}
+function hasMembership(record, value) {
+  return hasIndexedValue(record, String(value));
+}
+function hasIndexedValue(record, key) {
+  return Object.hasOwn(record, key);
+}
+function createEmptyRecord() {
+  return /* @__PURE__ */ Object.create(null);
+}
+function expectSingleSpan(spans) {
+  return spans.length === 1 ? spans[0] : void 0;
+}
+function createEvidenceSelectionList(evidence) {
+  if (Array.isArray(evidence)) {
+    return evidence.map(createEvidenceSelection);
+  }
+  const hasSelectionEvidence = typeof evidence.quote === "string" || typeof evidence.sentence_id === "string";
+  return hasSelectionEvidence ? createEvidenceSelection(evidence) : void 0;
+}
+function createEvidenceSelection(evidence) {
+  return {
+    ...typeof evidence.quote === "string" ? { quote: evidence.quote } : {},
+    ...typeof evidence.sentence_id === "string" ? { sentence_id: evidence.sentence_id } : {}
+  };
+}
+function toEvidenceResolutionFailure(failure, fieldName) {
+  return {
+    candidates: failure.candidates.map(toRankedSentenceCandidate),
+    code: failure.code,
+    fieldName,
+    message: failure.message
+  };
+}
+function toRankedSentenceCandidate(candidate) {
+  return {
+    exactNormalized: candidate.exactNormalized,
+    exactRaw: candidate.exactRaw,
+    exactSubstring: candidate.exactSubstring,
+    index: candidate.index,
+    nextText: candidate.nextText,
+    occurrenceId: candidate.occurrenceId,
+    prevText: candidate.prevText,
+    score: candidate.score,
+    sentenceId: candidate.sentence.sentenceId,
+    text: candidate.sentence.text
+  };
+}
+export {
+  ChunkBatchParser,
+  ChunkMetadataField,
+  bookCoherenceResponseSchema,
+  userFocusedResponseSchema
+};
+//# sourceMappingURL=parser.js.map