SmithDB của LangChain: Tìm kiếm toàn văn trên Object Storage với độ trễ chỉ 400ms

Tổng quan

SmithDB hỗ trợ tìm kiếm toàn văn (full-text search) và lọc JSON trên các dấu vết (traces) của tác nhân (agent) với độ trễ trung vị (P50) là 400 ms, mặc dù dữ liệu cơ sở bao gồm các tài liệu JSON lớn, lồng nhau sâu được lưu trữ trong object storage.

Tìm kiếm toàn văn là một lĩnh vực đã quá quen thuộc. Lucene đã có tuổi đời hai thập kỷ; Tantivy và Quickwit đã đưa việc tìm kiếm và lập chỉ mục lên object storage. Tuy nhiên, khi xây dựng tính năng tìm kiếm văn bản vào SmithDB, chúng tôi quyết định tiếp cận vấn đề từ những nguyên lý cơ bản vì việc lập chỉ mục các dấu vết tác nhân cho các khối lượng công việc tìm kiếm đặt ra một thách thức độc đáo.

SmithDB đòi hỏi một cách tiếp cận khác đối với việc tìm kiếm

Thách thức 1: Đặc điểm dữ liệu độc đáo của các dấu vết tác nhân

Mọi sự kiện LangSmith đều mã hóa các trường inputs và outputs, chiếm phần lớn tổng số byte của nó. Kích thước payload từ 1 MB trở lên cho inputs và outputs là rất phổ biến, với một số trường hợp lên tới hàng trăm megabyte khi chưa nén. Các cột nội dung này làm lu mờ các cột định danh, dấu thời gian và các cột siêu dữ liệu khác theo nhiều bậc độ lớn.

Ngoài ra, như đã đề cập trong bài blog gốc về SmithDB, các payload liên quan đến dấu vết tác nhân tiếp tục tăng kích thước theo thời gian. Đây là kết quả trực tiếp của việc kích thước cửa sổ ngữ cảnh (context window) của LLM ngày càng lớn và các tác nhân chạy trong thời gian dài hơn, khiến LLM tích lũy nhiều ngữ cảnh hơn.

Những đặc điểm này làm đảo ngược các quy luật kinh tế thông thường của một chỉ mục tìm kiếm. Một công cụ log truyền thống lập chỉ mục hàng tỷ tài liệu nhỏ, vì vậy chỉ mục có kích thước nhỏ so với mỗi tài liệu. Chúng tôi lập chỉ mục hàng tỷ tài liệu khổng lồ, nơi một tài liệu có thể tạo ra nhiều dữ liệu chỉ mục hơn nhiều dòng log nhỏ. Thông thường, tỷ lệ nguồn:chỉ mục cho log là khoảng 1:1,25. Tuy nhiên, đối với các dấu vết tác nhân trong LangSmith, chúng tôi quan sát thấy mức trung bình gần với 1:1,9 hơn. Ba hệ quả rút ra là:

Một chỉ mục đảo ngược (inverted index) là thứ ngăn chặn việc truy vấn nội dung phải quét toàn bộ payload, và nó phải hấp thụ các payload nặng, lệch, bán cấu trúc.

Thách thức 2: Object storage

SmithDB lưu giữ tất cả dữ liệu bền vững trong object storage để việc tính toán tương đối không trạng thái (stateless) và hệ thống có thể mở rộng bằng cách thêm các node mà không cần quản lý đĩa cục bộ.

Chi phí của một truy vấn tỷ lệ thuận với (số yêu cầu gửi đến object storage) × (số byte đọc mỗi yêu cầu). Trên object storage:

Mọi khía cạnh của chỉ mục đảo ngược trong SmithDB, từ bố cục lưu trữ đến thực thi truy vấn, đều được thiết kế với những ràng buộc này trong tâm trí.

Các dạng truy vấn tìm kiếm của SmithDB

Trước khi đi sâu hơn vào bố cục lưu trữ cho chỉ mục đảo ngược của chúng tôi, hãy xem qua các mẫu truy vấn chính mà chỉ mục phải giải quyết. Bề mặt truy vấn của SmithDB quy về ba nhóm vị ngữ (predicate families), và chúng khác nhau ở chỗ chúng khớp với cái gì và chấp nhận cú pháp mẫu nào.

Tóm lại:

Mọi phần sau đều tham chiếu lại bảng này: khi chúng tôi nói "truy vấn chỉ đường dẫn" (path-only query), chúng tôi muốn nói đến json_key, "giá trị có khóa" (keyed value) nghĩa là json_key_search, và "toàn văn" (full-text) nghĩa là search.

Tổng quan về chỉ mục đảo ngược

Chỉ mục đảo ngược là cấu trúc dữ liệu cung cấp sức mạnh cho mọi thư viện tìm kiếm, từ Lucene đến Tantivy. Nó giống như mục lục ở cuối sách giáo khoa: tra cứu một thuật ngữ một lần và nhảy thẳng đến các trang đề cập đến nó thay vì đọc từng trang. SmithDB xây dựng dựa trên ý tưởng này và chuyên biệt hóa bố cục lưu trữ cho các payload dấu vết tác nhân lớn mà nó lưu trữ trong object storage.

Thuật ngữ, danh sách đăng bài (postings), vị trí

Cấu trúc chỉ mục đảo ngược dựa trên ba khái niệm:

Lấy năm dấu vết được lập chỉ mục trên văn bản của chúng:

Chỉ mục giữ một mục cho mỗi thuật ngữ, trỏ đến các tài liệu đề cập đến nó:

Mỗi thuật ngữ là một mục từ điển: tra cứu giá trị, đọc danh sách đăng bài của nó, và bạn biết chính xác tài liệu nào cần lấy. Một truy vấn như search("deep agents") giao các danh sách đăng bài cho deep ([1, 2, 3, 4]) và agents ([0, 1, 2, 3]) để có được [1, 2, 3] mà không cần quét payload.

Cột vị trí ghi lại, cho mỗi tài liệu, (các) offset token nơi thuật ngữ xuất hiện, ví dụ 1:[0] nghĩa là tài liệu 1, vị trí 0. Đó là điều làm cho tìm kiếm cụm từ trở nên khả thi: search("langsmith engine") khớp với tài liệu 1 vì langsmith ở offset 0 và engine ở offset 1 (0 + 1 == 1), nhưng không khớp với tài liệu 4, nơi powers và the nằm giữa chúng (langsmith ở 1, engine ở 4).

Tại sao chúng tôi tận dụng Vortex thay vì Tantivy

Tantivy là một thư viện lập chỉ mục tìm kiếm xuất sắc và là điểm tham chiếu rõ ràng cho tìm kiếm kiểu Lucene trong Rust. Chúng tôi bắt đầu bằng việc tự hỏi liệu có thể áp dụng nó trực tiếp hay không. Thiết kế mà chúng tôi đạt được lấy cảm hứng rất nhiều từ Tantivy, nhưng một vài ràng buộc khiến nó không phù hợp trực tiếp với trường hợp sử dụng của chúng tôi:

Hành trình phát triển chỉ mục đảo ngược của SmithDB

Giới thiệu nhanh về Vortex

Vortex là một định dạng tệp cột và có thể mở rộng mà SmithDB sử dụng cho object storage. Không giống như các định dạng cố định như Parquet, Vortex cho phép các mã hóa có thể cắm vào (pluggable) và bố cục tệp tùy chỉnh, cho phép chúng tôi điều chỉnh nén và các mẫu truy cập I/O cho khối lượng công việc của mình mà không cần phân nhánh định dạng tệp.

Mỗi lần đọc sẽ loại bỏ toàn bộ các nhóm hàng bằng cách sử dụng thống kê, lọc các hàng còn lại thành một mặt nạ (mask), và chỉ chiếu các cột mà truy vấn thực sự cần.

Đơn vị I/O trong tệp Vortex là một phân đoạn (segment): một phạm vi byte vật lý liên tục. Trên object storage, một vòng lặp (round-trip) tốn khoảng 100 ms, vì vậy đòn bẩy chính cho độ trễ truy vấn là giảm thiểu số lượng yêu cầu. Bộ lập lịch I/O của Vortex gộp các lần đọc phân đoạn gần nhau thành một yêu cầu duy nhất, hợp nhất các lần đọc trong khoảng cách 1 MB thành một, lên đến cửa sổ 16 MB, vì vậy các mẫu truy cập tuần tự trong chỉ mục ánh xạ tới rất ít lệnh GET của object store.

Nỗ lực đầu tiên (không thành công) của chúng tôi

Phiên bản đầu tiên là bản dịch gần như nguyên văn của chỉ mục đảo ngược trong sách giáo khoa. Hai cột (term_key cho các đường dẫn và term_value cho các token) cho phép một bố cục phục vụ cả ba dạng truy vấn: sự tồn tại của đường dẫn đọc term_key, tìm kiếm có khóa giao các danh sách đăng bài trên cả hai cột, và toàn văn giao trên term_value đơn thuần. Các danh sách đăng bài được lưu trữ dưới dạng các ô List<u32>, các vị trí dưới dạng List<List<u32>>.

Chúng tôi dựa vào các mặc định của Vortex: mã hóa FSST cho các cột thuật ngữ, mã hóa bitpacked cho các danh sách đăng bài và vị trí, và bố cục lưu trữ theo vùng cho phép loại bỏ tại thời điểm truy vấn. Chỉ riêng các vị trí (cần thiết cho tìm kiếm cụm từ) đã lớn hơn một bậc so với mọi cột khác, vì vậy chúng tôi giữ chỉ mục trong một tệp riêng biệt với dữ liệu chạy cốt lõi. Điều này cho phép chúng tôi tách biệt việc xây dựng và hợp nhất chỉ mục khỏi đường dẫn ghi cốt lõi. Các API của Vortex hoạt động trên các chỉ số hàng và mặt nạ, vì vậy việc ủy quyền lọc chỉ mục cho một tệp anh em được thực hiện một cách tự nhiên.

Ba vấn đề đã xuất hiện ở quy mô lớn: