Repository: arabold/docs-mcp-server (ts)

Issue:
Improve Search Results: Merge Hits from Same URL into Single Ordered Result
The current search process, handled by `DocumentStore` and `DocumentRetrieverService`, retrieves individual chunks matching a query. `DocumentStore.findByContent` uses hybrid search (vector + FTS) and RRF ranking to find relevant chunks. `DocumentRetrieverService` then expands the context for *each* matching chunk by fetching its parent, preceding siblings, the chunk itself, child chunks, and subsequent siblings.

**Problem:**

When a search query matches multiple chunks within the *same document* (identified by the same `url` metadata), the `DocumentRetrieverService` currently generates a separate `StoreSearchResult` entry for each initial hit. This leads to several issues:

1.  **Overlapping Content:** Since context (siblings, etc.) is fetched for each hit independently, the final results for the same URL often contain redundant or overlapping content.
2.  **Incorrect Ordering:** The list of `StoreSearchResult` objects is implicitly ordered by the `rrf_score` of the initial hits returned by `findByContent`. This means chunks from the same original document might appear out of their natural sequence, potentially confusing downstream processing (like an LLM trying to understand the flow).
3.  **Inefficiency:** Fetching overlapping context (parents, siblings) multiple times for hits on the same page is inefficient.

**Example:**

Imagine a document (`page.html`) split into Chunks A, B, C, D, E in that order.
*   Search query matches Chunk B (score 0.9) and Chunk D (score 0.8).
*   `findByContent` returns `[Chunk B, Chunk D]`.
*   `DocumentRetrieverService` processes B: Fetches A (prec), B, C (subseq). Result 1: `content: "A\n\nB\n\nC"`, `score: 0.9`, `url: "page.html"`.
*   `DocumentRetrieverService` processes D: Fetches C (prec), D, E (subseq). Result 2: `content: "C\n\nD\n\nE"`, `score: 0.8`, `url: "page.html"`.
*   Final Output: `[Result 1, Result 2]` - Contains duplicate 'C' and is ordered by score, not original sequence.

**Proposed Solution:**

Modify the `DocumentRetrieverService.search` method to consolidate results originating from the same URL.

1.  **Group by URL:** After receiving the initial results from `documentStore.findByContent`, group these results by their `metadata.url`.
2.  **Collect Unique Chunks:** For each URL group:
    *   Identify all unique chunk IDs involved: the initial hits within the group *plus* all their respective parents, preceding siblings, children, and subsequent siblings retrieved via `documentStore` methods.
    *   Fetch the full `Document` objects for these unique chunk IDs from the `DocumentStore` *once*.
3.  **Sort and Merge:**
    *   Sort the fetched `Document` objects based on their original `sort_order` metadata.
    *   Concatenate the `pageContent` of the sorted chunks, using `\n\n` as a separator.
4.  **Apply Code Block Merging:** Apply the logic from Issue #45 (Merge Adjacent Split Code Blocks) to the concatenated content to merge adjacent code blocks.
5.  **Assign Score:** Use the *highest* `rrf_score` among the initial hits belonging to that URL group as the final `score` for the merged `StoreSearchResult`.
6.  **Character Limit (Optional):** Implement a maximum character limit for the merged `content` to prevent excessively long results. If the limit is exceeded, the content could be truncated, or potentially split into multiple results (though simple truncation is likely easier initially).

**Desired Output (Example):**

*   Single Result: `content: "A\n\nB\n\nC\n\nD\n\nE"`, `score: 0.9`, `url: "page.html"` (assuming code block merging is also applied).

**Benefits:**

*   Provides a more coherent and contextually accurate representation of matching sections within a document.
*   Eliminates redundant content.
*   Ensures chunks are presented in their original order.
*   Improves efficiency by reducing redundant database fetches.

**Relevant Code:**

*   `src/store/DocumentRetrieverService.ts`: Needs modification to implement grouping, fetching unique chunks, sorting, and merging.
*   `src/store/DocumentStore.ts`: Provides `sort_order` and methods to fetch related chunks.

Pull request description:
Consolidate search results by URL in DocumentRetrieverService (closes #46)
Refactored DocumentRetrieverService to group initial hits by URL, deduplicate and merge context, and produce a single ordered result per URL. Includes updated and new unit tests for DocumentStore and DocumentRetrieverService. Removes outdated per-hit aggregation tests and replaces with focused, scenario-based tests for the new logic.

Patch:
diff --git a/src/store/DocumentRetrieverService.ts b/src/store/DocumentRetrieverService.ts
index d09994e..9fcb289 100644
--- a/src/store/DocumentRetrieverService.ts
+++ b/src/store/DocumentRetrieverService.ts
@@ -1,3 +1,4 @@
+import type { Document } from "@langchain/core/documents";
 import type { DocumentStore } from "./DocumentStore";
 import type { StoreSearchResult } from "./types";
 
@@ -11,6 +12,120 @@ export class DocumentRetrieverService {
     this.documentStore = documentStore;
   }
 
+  /**
+   * Collects all related chunk IDs for a given initial hit.
+   * Returns an object with url, hitId, relatedIds (Set), and score.
+   */
+  private async getRelatedChunkIds(
+    library: string,
+    version: string,
+    doc: Document,
+    siblingLimit = SIBLING_LIMIT,
+    childLimit = CHILD_LIMIT,
+  ): Promise<{
+    url: string;
+    hitId: string;
+    relatedIds: Set<string>;
+    score: number;
+  }> {
+    const id = doc.id as string;
+    const url = doc.metadata.url as string;
+    const score = doc.metadata.score as number;
+    const relatedIds = new Set<string>();
+    relatedIds.add(id);
+
+    // Parent
+    const parent = await this.documentStore.findParentChunk(library, version, id);
+    if (parent) {
+      relatedIds.add(parent.id as string);
+    }
+
+    // Preceding Siblings
+    const precedingSiblings = await this.documentStore.findPrecedingSiblingChunks(
+      library,
+      version,
+      id,
+      siblingLimit,
+    );
+    for (const sib of precedingSiblings) {
+      relatedIds.add(sib.id as string);
+    }
+
+    // Child Chunks
+    const childChunks = await this.documentStore.findChildChunks(
+      library,
+      version,
+      id,
+      childLimit,
+    );
+    for (const child of childChunks) {
+      relatedIds.add(child.id as string);
+    }
+
+    // Subsequent Siblings
+    const subsequentSiblings = await this.documentStore.findSubsequentSiblingChunks(
+      library,
+      version,
+      id,
+      siblingLimit,
+    );
+    for (const sib of subsequentSiblings) {
+      relatedIds.add(sib.id as string);
+    }
+
+    return { url, hitId: id, relatedIds, score };
+  }
+
+  /**
+   * Groups related chunk info by URL, deduplicates IDs, and finds max score per URL.
+   */
+  private groupAndPrepareFetch(
+    relatedInfos: Array<{
+      url: string;
+      hitId: string;
+      relatedIds: Set<string>;
+      score: number;
+    }>,
+  ): Map<string, { uniqueChunkIds: Set<string>; maxScore: number }> {
+    const urlMap = new Map<string, { uniqueChunkIds: Set<string>; maxScore: number }>();
+    for (const info of relatedInfos) {
+      let entry = urlMap.get(info.url);
+      if (!entry) {
+        entry = { uniqueChunkIds: new Set(), maxScore: info.score };
+        urlMap.set(info.url, entry);
+      }
+      for (const id of info.relatedIds) {
+        entry.uniqueChunkIds.add(id);
+      }
+      if (info.score > entry.maxScore) {
+        entry.maxScore = info.score;
+      }
+    }
+    return urlMap;
+  }
+
+  /**
+   * Finalizes the merged result for a URL group by fetching, sorting, and joining content.
+   */
+  private async finalizeResult(
+    library: string,
+    version: string,
+    url: string,
+    uniqueChunkIds: Set<string>,
+    maxScore: number,
+  ): Promise<StoreSearchResult> {
+    const ids = Array.from(uniqueChunkIds);
+    const docs = await this.documentStore.findChunksByIds(library, version, ids);
+    // Already sorted by sort_order in findChunksByIds
+    const content = docs.map((d) => d.pageContent).join("\n");
+    // TODO: Apply code block merging here if/when implemented
+    return {
+      url,
+      content,
+      score: maxScore,
+    };
+  }
+
   /**
    * Searches for documents and expands the context around the matches.
    * @param library The library name.
@@ -37,63 +152,27 @@ export class DocumentRetrieverService {
       limit ?? 10,
     );
 
-    const results: StoreSearchResult[] = [];
-
-    for (const doc of initialResults) {
-      const id = doc.id as string;
-      let content = "";
-
-      // Parent
-      const parent = await this.documentStore.findParentChunk(
-        library,
-        normalizedVersion,
-        id,
-      );
-      if (parent) {
-        content += `${parent.pageContent}\n\n`;
-      }
-
-      // Preceding Siblings
-      const precedingSiblings = await this.documentStore.findPrecedingSiblingChunks(
-        library,
-        normalizedVersion,
-        id,
-        SIBLING_LIMIT,
-      );
-      if (precedingSiblings.length > 0) {
-        content += `${precedingSiblings.map((d) => d.pageContent).join("\n\n")}\n\n`;
-      }
+    // Step 1: Expand context for each initial hit (collect related chunk IDs)
+    const relatedInfos = await Promise.all(
+      initialResults.map((doc) =>
+        this.getRelatedChunkIds(library, normalizedVersion, doc),
+      ),
+    );
 
-      // Initial Result
-      content += `${doc.pageContent}`;
+    // Step 2: Group by URL, deduplicate, and find max score
+    const urlMap = this.groupAndPrepareFetch(relatedInfos);
 
-      // Child Chunks
-      const childChunks = await this.documentStore.findChildChunks(
-        library,
-        normalizedVersion,
-        id,
-        CHILD_LIMIT,
-      );
-      if (childChunks.length > 0) {
-        content += `\n\n${childChunks.map((d) => d.pageContent).join("\n\n")}`;
-      }
-
-      // Subsequent Siblings
-      const subsequentSiblings = await this.documentStore.findSubsequentSiblingChunks(
+    // Step 3: For each URL group, fetch, sort, and format the merged result
+    const results: StoreSearchResult[] = [];
+    for (const [url, { uniqueChunkIds, maxScore }] of urlMap.entries()) {
+      const result = await this.finalizeResult(
         library,
         normalizedVersion,
-        id,
-        SIBLING_LIMIT,
+        url,
+        uniqueChunkIds,
+        maxScore,
       );
-      if (subsequentSiblings.length > 0) {
-        content += `\n\n${subsequentSiblings.map((d) => d.pageContent).join("\n\n")}`;
-      }
-
-      results.push({
-        url: doc.metadata.url,
-        content,
-        score: doc.metadata.score,
-      });
+      results.push(result);
     }
 
     return results;
diff --git a/src/store/DocumentStore.ts b/src/store/DocumentStore.ts
index 95e356b..b6678c7 100644
--- a/src/store/DocumentStore.ts
+++ b/src/store/DocumentStore.ts
@@ -675,4 +675,31 @@ export class DocumentStore {
       throw new ConnectionError(`Failed to find parent chunk for ID ${id}`, error);
     }
   }
+
+  /**
+   * Fetches multiple documents by their IDs in a single call.
+   * Returns an array of Document objects, sorted by their sort_order.
+   */
+  async findChunksByIds(
+    library: string,
+    version: string,
+    ids: string[],
+  ): Promise<Document[]> {
+    if (!ids.length) return [];
+    try {
+      // Use parameterized query for variable number of IDs
+      const placeholders = ids.map(() => "?").join(",");
+      const stmt = this.db.prepare(
+        `SELECT * FROM documents WHERE library = ? AND version = ? AND id IN (${placeholders}) ORDER BY sort_order`,
+      );
+      const rows = stmt.all(
+        library.toLowerCase(),
+        version.toLowerCase(),
+        ...ids,
+      ) as DbDocument[];
+      return rows.map((row) => mapDbDocumentToDocument(row));
+    } catch (error) {
+      throw new ConnectionError("Failed to fetch documents by IDs", error);
+    }
+  }
 }


Test patch:
diff --git a/src/store/DocumentRetrieverService.test.ts b/src/store/DocumentRetrieverService.test.ts
index a63f92f..e00e8e9 100644
--- a/src/store/DocumentRetrieverService.test.ts
+++ b/src/store/DocumentRetrieverService.test.ts
@@ -1,384 +1,227 @@
 import { Document } from "@langchain/core/documents";
-import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 import { DocumentRetrieverService } from "./DocumentRetrieverService";
 import { DocumentStore } from "./DocumentStore";
 
 vi.mock("./DocumentStore");
 vi.mock("../utils/logger");
 
-describe("DocumentRetrieverService", () => {
+describe("DocumentRetrieverService (consolidated logic)", () => {
   let retrieverService: DocumentRetrieverService;
   let mockDocumentStore: DocumentStore;
 
   beforeEach(() => {
     vi.clearAllMocks();
-    mockDocumentStore = new DocumentStore("mock_connection_string"); // The constructor argument won't matter
+    mockDocumentStore = new DocumentStore("mock_connection_string");
     retrieverService = new DocumentRetrieverService(mockDocumentStore);
   });
 
-  afterEach(() => {
-    // Cleanup if needed
-  });
-
   it("should return an empty array when no documents are found", async () => {
     vi.spyOn(mockDocumentStore, "findByContent").mockResolvedValue([]);
-
-    const results = await retrieverService.search("test-lib", "1.0.0", "query");
+    const results = await retrieverService.search("lib", "1.0.0", "query");
     expect(results).toEqual([]);
   });
 
-  it("should retrieve and aggregate document content", async () => {
-    const library = "test-lib";
+  it("should consolidate multiple hits from the same URL into a single ordered result", async () => {
+    const library = "lib";
     const version = "1.0.0";
-    const query = "test query";
-    const initialResult = new Document({
+    const query = "test";
+    // Two initial hits from the same URL, with overlapping context
+    const initialResult1 = new Document({
       id: "doc1",
-      pageContent: "Initial content",
-      metadata: { url: "url" },
-    });
-
-    const parent = new Document({
-      id: "parent1",
-      pageContent: "Parent content",
-      metadata: { url: "url" },
+      pageContent: "Chunk A",
+      metadata: { url: "url", score: 0.9 },
     });
-    const precedingSibling1 = new Document({
-      id: "sibling1",
-      pageContent: "Preceding sibling 1",
-      metadata: { url: "url" },
-    });
-    const child1 = new Document({
-      id: "child1",
-      pageContent: "Child 1 content",
-      metadata: { url: "url" },
+    const initialResult2 = new Document({
+      id: "doc3",
+      pageContent: "Chunk C",
+      metadata: { url: "url", score: 0.8 },
     });
-    const subsequentSibling1 = new Document({
-      id: "sibling2",
-      pageContent: "Subsequent sibling 1",
+    const doc2 = new Document({
+      id: "doc2",
+      pageContent: "Chunk B",
       metadata: { url: "url" },
     });
 
-    vi.spyOn(mockDocumentStore, "findByContent").mockResolvedValue([initialResult]);
-    vi.spyOn(mockDocumentStore, "findParentChunk").mockResolvedValue(parent);
-    vi.spyOn(mockDocumentStore, "findPrecedingSiblingChunks").mockResolvedValue([
-      precedingSibling1,
-    ]);
-    vi.spyOn(mockDocumentStore, "findChildChunks").mockResolvedValue([child1]);
-    vi.spyOn(mockDocumentStore, "findSubsequentSiblingChunks").mockResolvedValue([
-      subsequentSibling1,
+    vi.spyOn(mockDocumentStore, "findByContent").mockResolvedValue([
+      initialResult1,
+      initialResult2,
     ]);
-
-    const results = await retrieverService.search(library, version, query);
-
-    expect(mockDocumentStore.findByContent).toHaveBeenCalledWith(
-      library,
-      version,
-      query,
-      10,
-    );
-    expect(mockDocumentStore.findParentChunk).toHaveBeenCalledWith(
-      library,
-      version,
-      "doc1",
+    vi.spyOn(mockDocumentStore, "findParentChunk").mockImplementation(async () => null);
+    vi.spyOn(mockDocumentStore, "findPrecedingSiblingChunks").mockImplementation(
+      async () => [],
     );
-    expect(mockDocumentStore.findPrecedingSiblingChunks).toHaveBeenCalledWith(
-      library,
-      version,
-      "doc1",
-      2,
+    vi.spyOn(mockDocumentStore, "findChildChunks").mockImplementation(
+      async (lib, ver, id) => (id === "doc1" ? [doc2] : []),
     );
-    expect(mockDocumentStore.findChildChunks).toHaveBeenCalledWith(
-      library,
-      version,
-      "doc1",
-      5,
-    );
-    expect(mockDocumentStore.findSubsequentSiblingChunks).toHaveBeenCalledWith(
-      library,
-      version,
-      "doc1",
-      2,
+    vi.spyOn(mockDocumentStore, "findSubsequentSiblingChunks").mockImplementation(
+      async (lib, ver, id) => (id === "doc1" ? [doc2] : []),
     );
-
-    expect(results).toEqual([
-      {
-        content:
-          "Parent content\n\nPreceding sibling 1\n\nInitial content\n\nChild 1 content\n\nSubsequent sibling 1",
-        url: "url",
-        score: undefined,
-      },
-    ]);
-  });
-
-  it("should handle missing parent chunk", async () => {
-    const library = "test-lib";
-    const version = "1.0.0";
-    const query = "test query";
-    const initialResult = new Document({
-      id: "doc1",
-      pageContent: "Initial content",
-      metadata: { url: "url" },
-    });
-
-    const precedingSibling1 = new Document({
-      id: "sibling1",
-      pageContent: "Preceding sibling 1",
-      metadata: { url: "url" },
-    });
-    const child1 = new Document({
-      id: "child1",
-      pageContent: "Child 1 content",
-      metadata: { url: "url" },
-    });
-    const subsequentSibling1 = new Document({
-      id: "sibling2",
-      pageContent: "Subsequent sibling 1",
-      metadata: { url: "url" },
-    });
-
-    vi.spyOn(mockDocumentStore, "findByContent").mockResolvedValue([initialResult]);
-    vi.spyOn(mockDocumentStore, "findParentChunk").mockResolvedValue(null); // Mock missing parent
-    vi.spyOn(mockDocumentStore, "findPrecedingSiblingChunks").mockResolvedValue([
-      precedingSibling1,
-    ]);
-    vi.spyOn(mockDocumentStore, "findChildChunks").mockResolvedValue([child1]);
-    vi.spyOn(mockDocumentStore, "findSubsequentSiblingChunks").mockResolvedValue([
-      subsequentSibling1,
-    ]);
+    const findChunksByIdsSpy = vi
+      .spyOn(mockDocumentStore, "findChunksByIds")
+      .mockResolvedValue([
+        initialResult1, // doc1 (Chunk A)
+        doc2, // doc2 (Chunk B)
+        initialResult2, // doc3 (Chunk C)
+      ]);
 
     const results = await retrieverService.search(library, version, query);
 
+    expect(findChunksByIdsSpy).toHaveBeenCalledWith(
+      library,
+      version,
+      expect.arrayContaining(["doc1", "doc2", "doc3"]),
+    );
     expect(results).toEqual([
       {
-        content:
-          "Preceding sibling 1\n\nInitial content\n\nChild 1 content\n\nSubsequent sibling 1",
+        content: "Chunk A\n\nChunk B\n\nChunk C",
         url: "url",
-        score: undefined,
+        score: 0.9,
       },
     ]);
   });
 
-  it("should handle missing preceding siblings", async () => {
-    const library = "test-lib";
+  it("should return a single result for a single hit with context", async () => {
+    const library = "lib";
     const version = "1.0.0";
-    const query = "test query";
+    const query = "test";
     const initialResult = new Document({
       id: "doc1",
-      pageContent: "Initial content",
-      metadata: { url: "url" },
+      pageContent: "Main chunk",
+      metadata: { url: "url", score: 0.7 },
     });
-
     const parent = new Document({
       id: "parent1",
-      pageContent: "Parent content",
+      pageContent: "Parent",
       metadata: { url: "url" },
     });
-    const child1 = new Document({
+    const child = new Document({
       id: "child1",
-      pageContent: "Child 1 content",
-      metadata: { url: "url" },
-    });
-    const subsequentSibling1 = new Document({
-      id: "sibling2",
-      pageContent: "Subsequent sibling 1",
+      pageContent: "Child",
       metadata: { url: "url" },
     });
 
     vi.spyOn(mockDocumentStore, "findByContent").mockResolvedValue([initialResult]);
     vi.spyOn(mockDocumentStore, "findParentChunk").mockResolvedValue(parent);
-    vi.spyOn(mockDocumentStore, "findPrecedingSiblingChunks").mockResolvedValue([]); // Mock missing preceding siblings
-    vi.spyOn(mockDocumentStore, "findChildChunks").mockResolvedValue([child1]);
-    vi.spyOn(mockDocumentStore, "findSubsequentSiblingChunks").mockResolvedValue([
-      subsequentSibling1,
-    ]);
+    vi.spyOn(mockDocumentStore, "findPrecedingSiblingChunks").mockResolvedValue([]);
+    vi.spyOn(mockDocumentStore, "findChildChunks").mockResolvedValue([child]);
+    vi.spyOn(mockDocumentStore, "findSubsequentSiblingChunks").mockResolvedValue([]);
+    const findChunksByIdsSpy = vi
+      .spyOn(mockDocumentStore, "findChunksByIds")
+      .mockResolvedValue([parent, initialResult, child]);
 
     const results = await retrieverService.search(library, version, query);
 
+    expect(findChunksByIdsSpy).toHaveBeenCalledWith(
+      library,
+      version,
+      expect.arrayContaining(["parent1", "doc1", "child1"]),
+    );
     expect(results).toEqual([
       {
-        content:
-          "Parent content\n\nInitial content\n\nChild 1 content\n\nSubsequent sibling 1",
+        content: "Parent\n\nMain chunk\n\nChild",
         url: "url",
-        score: undefined,
+        score: 0.7,
       },
     ]);
   });
 
-  it("should handle missing child chunks", async () => {
-    const library = "test-lib";
+  it("should return multiple results for hits from different URLs", async () => {
+    const library = "lib";
     const version = "1.0.0";
-    const query = "test query";
-    const initialResult = new Document({
-      id: "doc1",
-      pageContent: "Initial content",
-      metadata: { url: "url" },
-    });
-
-    const parent = new Document({
-      id: "parent1",
-      pageContent: "Parent content",
-      metadata: { url: "url" },
-    });
-    const precedingSibling1 = new Document({
-      id: "sibling1",
-      pageContent: "Preceding sibling 1",
-      metadata: { url: "url" },
+    const query = "test";
+    const docA = new Document({
+      id: "a1",
+      pageContent: "A1",
+      metadata: { url: "urlA", score: 0.8 },
     });
-    const subsequentSibling1 = new Document({
-      id: "sibling2",
-      pageContent: "Subsequent sibling 1",
-      metadata: { url: "url" },
+    const docB = new Document({
+      id: "b1",
+      pageContent: "B1",
+      metadata: { url: "urlB", score: 0.9 },
     });
 
-    vi.spyOn(mockDocumentStore, "findByContent").mockResolvedValue([initialResult]);
-    vi.spyOn(mockDocumentStore, "findParentChunk").mockResolvedValue(parent);
-    vi.spyOn(mockDocumentStore, "findPrecedingSiblingChunks").mockResolvedValue([
-      precedingSibling1,
-    ]);
-    vi.spyOn(mockDocumentStore, "findChildChunks").mockResolvedValue([]); // Mock missing children
-    vi.spyOn(mockDocumentStore, "findSubsequentSiblingChunks").mockResolvedValue([
-      subsequentSibling1,
-    ]);
+    vi.spyOn(mockDocumentStore, "findByContent").mockResolvedValue([docA, docB]);
+    vi.spyOn(mockDocumentStore, "findParentChunk").mockResolvedValue(null);
+    vi.spyOn(mockDocumentStore, "findPrecedingSiblingChunks").mockResolvedValue([]);
+    vi.spyOn(mockDocumentStore, "findChildChunks").mockResolvedValue([]);
+    vi.spyOn(mockDocumentStore, "findSubsequentSiblingChunks").mockResolvedValue([]);
+    vi.spyOn(mockDocumentStore, "findChunksByIds").mockImplementation(
+      async (lib, ver, ids) => {
+        if (ids.includes("a1")) return [docA];
+        if (ids.includes("b1")) return [docB];
+        return [];
+      },
+    );
 
     const results = await retrieverService.search(library, version, query);
 
     expect(results).toEqual([
       {
-        content:
-          "Parent content\n\nPreceding sibling 1\n\nInitial content\n\nSubsequent sibling 1",
-        url: "url",
-        score: undefined,
+        content: "A1",
+        url: "urlA",
+        score: 0.8,
       },
-    ]);
-  });
-
-  it("should handle missing subsequent siblings", async () => {
-    const library = "test-lib";
-    const version = "1.0.0";
-    const query = "test query";
-    const initialResult = new Document({
-      id: "doc1",
-      pageContent: "Initial content",
-      metadata: { url: "url" },
-    });
-
-    const parent = new Document({
-      id: "parent1",
-      pageContent: "Parent content",
-      metadata: { url: "url" },
-    });
-    const precedingSibling1 = new Document({
-      id: "sibling1",
-      pageContent: "Preceding sibling 1",
-      metadata: { url: "url" },
-    });
-    const child1 = new Document({
-      id: "child1",
-      pageContent: "Child 1 content",
-      metadata: { url: "url" },
-    });
-
-    vi.spyOn(mockDocumentStore, "findByContent").mockResolvedValue([initialResult]);
-    vi.spyOn(mockDocumentStore, "findParentChunk").mockResolvedValue(parent);
-    vi.spyOn(mockDocumentStore, "findPrecedingSiblingChunks").mockResolvedValue([
-      precedingSibling1,
-    ]);
-    vi.spyOn(mockDocumentStore, "findChildChunks").mockResolvedValue([child1]);
-    vi.spyOn(mockDocumentStore, "findSubsequentSiblingChunks").mockResolvedValue([]); // Mock missing subsequent siblings
-
-    const results = await retrieverService.search(library, version, query);
-
-    expect(results).toEqual([
       {
-        content:
-          "Parent content\n\nPreceding sibling 1\n\nInitial content\n\nChild 1 content",
-        url: "url",
-        score: undefined,
+        content: "B1",
+        url: "urlB",
+        score: 0.9,
       },
     ]);
   });
 
-  it("should handle multiple initial results", async () => {
-    const library = "test-lib";
+  it("should handle all context lookups returning empty", async () => {
+    const library = "lib";
     const version = "1.0.0";
-    const query = "test query";
-    const initialResult1 = new Document({
+    const query = "test";
+    const initialResult = new Document({
       id: "doc1",
-      pageContent: "Initial content 1",
-      metadata: { url: "url" },
-    });
-    const initialResult2 = new Document({
-      id: "doc2",
-      pageContent: "Initial content 2",
-      metadata: { url: "url" },
+      pageContent: "Main chunk",
+      metadata: { url: "url", score: 0.5 },
     });
 
-    vi.spyOn(mockDocumentStore, "findByContent").mockResolvedValue([
-      initialResult1,
-      initialResult2,
-    ]);
+    vi.spyOn(mockDocumentStore, "findByContent").mockResolvedValue([initialResult]);
     vi.spyOn(mockDocumentStore, "findParentChunk").mockResolvedValue(null);
     vi.spyOn(mockDocumentStore, "findPrecedingSiblingChunks").mockResolvedValue([]);
     vi.spyOn(mockDocumentStore, "findChildChunks").mockResolvedValue([]);
     vi.spyOn(mockDocumentStore, "findSubsequentSiblingChunks").mockResolvedValue([]);
+    const findChunksByIdsSpy = vi
+      .spyOn(mockDocumentStore, "findChunksByIds")
+      .mockResolvedValue([initialResult]);
 
     const results = await retrieverService.search(library, version, query);
 
+    expect(findChunksByIdsSpy).toHaveBeenCalledWith(
+      library,
+      version,
+      expect.arrayContaining(["doc1"]),
+    );
     expect(results).toEqual([
       {
-        content: "Initial content 1",
+        content: "Main chunk",
         url: "url",
-        score: undefined,
-      },
-      {
-        content: "Initial content 2",
-        url: "url",
-        score: undefined,
+        score: 0.5,
       },
     ]);
   });
 
   it("should use the provided limit", async () => {
-    const library = "test-lib";
+    const library = "lib";
     const version = "1.0.0";
-    const query = "test query";
-    const limit = 5;
+    const query = "test";
+    const limit = 3;
     const initialResult = new Document({
       id: "doc1",
-      pageContent: "Initial content",
-      metadata: { url: "url" },
-    });
-
-    const parent = new Document({
-      id: "parent1",
-      pageContent: "Parent content",
-      metadata: { url: "url" },
-    });
-    const precedingSibling1 = new Document({
-      id: "sibling1",
-      pageContent: "Preceding sibling 1",
-      metadata: { url: "url" },
-    });
-    const child1 = new Document({
-      id: "child1",
-      pageContent: "Child 1 content",
-      metadata: { url: "url" },
-    });
-    const subsequentSibling1 = new Document({
-      id: "sibling2",
-      pageContent: "Subsequent sibling 1",
-      metadata: { url: "url" },
+      pageContent: "Main chunk",
+      metadata: { url: "url", score: 0.5 },
     });
 
     vi.spyOn(mockDocumentStore, "findByContent").mockResolvedValue([initialResult]);
-    vi.spyOn(mockDocumentStore, "findParentChunk").mockResolvedValue(parent);
-    vi.spyOn(mockDocumentStore, "findPrecedingSiblingChunks").mockResolvedValue([
-      precedingSibling1,
-    ]);
-    vi.spyOn(mockDocumentStore, "findChildChunks").mockResolvedValue([child1]);
-    vi.spyOn(mockDocumentStore, "findSubsequentSiblingChunks").mockResolvedValue([
-      subsequentSibling1,
-    ]);
+    vi.spyOn(mockDocumentStore, "findParentChunk").mockResolvedValue(null);
+    vi.spyOn(mockDocumentStore, "findPrecedingSiblingChunks").mockResolvedValue([]);
+    vi.spyOn(mockDocumentStore, "findChildChunks").mockResolvedValue([]);
+    vi.spyOn(mockDocumentStore, "findSubsequentSiblingChunks").mockResolvedValue([]);
+    vi.spyOn(mockDocumentStore, "findChunksByIds").mockResolvedValue([initialResult]);
 
     const results = await retrieverService.search(library, version, query, limit);
 
@@ -387,55 +230,13 @@ describe("DocumentRetrieverService", () => {
       version,
       query,
       limit,
-    ); // Verify limit is passed to findByContent
+    );
     expect(results).toEqual([
       {
-        content:
-          "Parent content\n\nPreceding sibling 1\n\nInitial content\n\nChild 1 content\n\nSubsequent sibling 1",
+        content: "Main chunk",
         url: "url",
-        score: undefined,
+        score: 0.5,
       },
     ]);
   });
-
-  // Test for optional version handling
-  describe("Optional Version Handling", () => {
-    const library = "opt-lib";
-    const query = "optional query";
-    const limit = 7;
-
-    it("search should normalize null/undefined/empty version to empty string for store calls", async () => {
-      // Mock store methods to prevent errors and allow checking calls
-      const findByContentSpy = vi
-        .spyOn(mockDocumentStore, "findByContent")
-        .mockResolvedValue([]); // Return empty to simplify test
-      const findParentChunkSpy = vi.spyOn(mockDocumentStore, "findParentChunk");
-      const findPrecedingSiblingsSpy = vi.spyOn(
-        mockDocumentStore,
-        "findPrecedingSiblingChunks",
-      );
-      const findChildChunksSpy = vi.spyOn(mockDocumentStore, "findChildChunks");
-      const findSubsequentSiblingsSpy = vi.spyOn(
-        mockDocumentStore,
-        "findSubsequentSiblingChunks",
-      );
-
-      // Test with null version
-      await retrieverService.search(library, null, query, limit);
-      expect(findByContentSpy).toHaveBeenCalledWith(library, "", query, limit);
-      // We don't need to check other methods if findByContent returns empty,
-      // but if it returned results, we'd check those calls too, e.g.:
-      // expect(findParentChunkSpy).toHaveBeenCalledWith(library, "", expect.any(String));
-
-      // Test with undefined version
-      await retrieverService.search(library, undefined, query, limit);
-      expect(findByContentSpy).toHaveBeenCalledWith(library, "", query, limit);
-
-      // Test with empty string version
-      await retrieverService.search(library, "", query, limit);
-      expect(findByContentSpy).toHaveBeenCalledWith(library, "", query, limit);
-
-      // Restore mocks if necessary, though clearAllMocks in beforeEach handles it
-    });
-  });
 });
diff --git a/src/store/DocumentStore.test.ts b/src/store/DocumentStore.test.ts
index 7dffd2f..043b7d6 100644
--- a/src/store/DocumentStore.test.ts
+++ b/src/store/DocumentStore.test.ts
@@ -80,6 +80,83 @@ describe("DocumentStore", () => {
     vi.restoreAllMocks();
   });
 
+  describe("findChunksByIds", () => {
+    const library = "test-lib";
+    const version = "1.0.0";
+
+    it("should fetch and return documents for given IDs, sorted by sort_order", async () => {
+      const ids = ["id1", "id2", "id3"];
+      const mockRows = [
+        {
+          id: "id2",
+          library,
+          version,
+          url: "url2",
+          content: "content2",
+          metadata: JSON.stringify({ url: "url2", score: 0.5 }),
+          embedding: null,
+          sort_order: 1,
+          score: 0.5,
+        },
+        {
+          id: "id1",
+          library,
+          version,
+          url: "url1",
+          content: "content1",
+          metadata: JSON.stringify({ url: "url1", score: 0.9 }),
+          embedding: null,
+          sort_order: 0,
+          score: 0.9,
+        },
+        {
+          id: "id3",
+          library,
+          version,
+          url: "url3",
+          content: "content3",
+          metadata: JSON.stringify({ url: "url3", score: 0.7 }),
+          embedding: null,
+          sort_order: 2,
+          score: 0.7,
+        },
+      ];
+      // Should be returned sorted by sort_order: id1, id2, id3
+      mockStatementAll.mockReturnValueOnce([mockRows[1], mockRows[0], mockRows[2]]);
+      const result = await documentStore.findChunksByIds(library, version, ids);
+      expect(mockPrepare).toHaveBeenCalledWith(expect.stringContaining("id IN"));
+      expect(mockStatementAll).toHaveBeenCalledWith(
+        library.toLowerCase(),
+        version.toLowerCase(),
+        ...ids,
+      );
+      expect(result.length).toBe(3);
+      expect(result[0].id).toBe("id1");
+      expect(result[1].id).toBe("id2");
+      expect(result[2].id).toBe("id3");
+      expect(result[0].pageContent).toBe("content1");
+      expect(result[1].pageContent).toBe("content2");
+      expect(result[2].pageContent).toBe("content3");
+    });
+
+    it("should return an empty array if no IDs are provided", async () => {
+      const prepareCallsBefore = mockPrepare.mock.calls.length;
+      const allCallsBefore = mockStatementAll.mock.calls.length;
+      const result = await documentStore.findChunksByIds(library, version, []);
+      expect(result).toEqual([]);
+      expect(mockPrepare.mock.calls.length).toBe(prepareCallsBefore);
+      expect(mockStatementAll.mock.calls.length).toBe(allCallsBefore);
+    });
+
+    it("should return an empty array if no documents are found", async () => {
+      mockStatementAll.mockReturnValueOnce([]);
+      const result = await documentStore.findChunksByIds(library, version, ["idX"]);
+      expect(result).toEqual([]);
+      expect(mockPrepare).toHaveBeenCalled();
+      expect(mockStatementAll).toHaveBeenCalled();
+    });
+  });
+
   describe("findByContent", () => {
     const library = "test-lib";
     const version = "1.0.0";