Vector Databases

A Practical Introduction

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

Nitin Borwankar

Vector Databases

by Nitin Borwankar

Copyright © 2026 Nitin Borwankar. All rights reserved.

Printed in the United States of America.

Published by O’Reilly Media, Inc., 1005 Gravenstein Highway North, Sebastopol, CA 95472.

O’Reilly books may be purchased for educational, business, or sales promotional use. Online editions are also available for most titles (http://oreilly.com). For more information, contact our corporate/institutional sales department: 800-998-9938 or corporate@oreilly.com.

Editors: Andy Kwan and Sara Hunter

Production Editor: Beth Kelly

Copyeditor: FILL IN COPYEDITOR

Proofreader: FILL IN PROOFREADER

Indexer: FILL IN INDEXER

Interior Designer: David Futato

Cover Designer: Karen Montgomery

Illustrator: Kate Dullea

March 2026: First Edition

Revision History for the Early Release

2025-01-30: First Release

2025-05-02: Second Release

2025-07-18: Third Release

2025-10-15: Fourth Release

2025-11-04: Fifth Release

2025-11-19: Sixth Release

See http://oreilly.com/catalog/errata.csp?isbn=9781098177591 for release details.

The O’Reilly logo is a registered trademark of O’Reilly Media, Inc. Vector Databases, the cover image, and related trade dress are trademarks of O’Reilly Media, Inc.

The views expressed in this work are those of the author and do not represent the publisher’s views. While the publisher and the author have used good faith efforts to ensure that the information and instructions contained in this work are accurate, the publisher and the author disclaim all responsibility for errors or omissions, including without limitation responsibility for damages resulting from the use of or reliance on this work. Use of the information and instructions contained in this work is at your own risk. If any code samples or other technology this work contains or describes is subject to open source licenses or the intellectual property rights of others, it is your responsibility to ensure that your use thereof complies with such licenses and/or rights.

978-1-098-17753-9

[FILL IN]

Brief Table of Contents (Not Yet Final)

Chapter 1: Introduction to Vector Databases (available)

Chapter 2: Embeddings - Turning Data into Vectors (available)

Chapter 3: SQLite-VSS: Vector Search for Personal Use Cases (available)

Chapter 4: pgvector: Scaling Up with PostgreSQL (available)

Chapter 5: FAISS: Efficient Similarity Search at Scale (available)

Chapter 6: RAG: Combining Retrieval and Generation (available)

Chapter 7: RAG in the Enterprise: Scaling with pgvector (available)

Chapter 8: Beyond the Basics (available)

Chapter 9: Vector Query Language (VQL) — Draft Specification (available)

Chapter 10: Advanced Embedding Techniques (available)

Chapter 1. Introduction to Vector Databases

A note for Early Release readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 1st chapter of the final book. The GitHub repo will be made active later on.

If you have comments about how we might improve the content and/or examples in this book, or if you notice missing material within this chapter, please reach out to the editor at shunter@oreilly.com.

The current wave of generative AI (genAI for short), especially in natural language applications, has brought the technology of vector databases to the forefront of developer attention. And yet the technology is not well understood or widely known to the average database application developer.

In this chapter we take a high level look at the subject, identifying each important component and giving an initial flavor of the material to come. Mostly, we will focus on what the technology is and why we need it with all the explanation needed to understand how to use them except for the important topic of embeddings which we dive into in the next chapter.

Why do we need vector databases?

Even before the advent of genAI, the need for management of unstructured data such as text, images, audio and video had been a less well understood subject and somewhat of a black art. It required specialized technology and knowledge to manage these at scale. Indexing and searching over such data was only possible at the level of metadata - e.g. date of creation, size, author and so on. But searching for an image similar to another one, or a paragraph with similar meaning as another was not something that a SQL database could do.

Now, with the emergence of vector databases into the mainstream, management of unstructured data and searching based on semantics or meaning of the data have become much easier. This is the world of vector databases, which opens up new possibilities for applications based on semantic search or similarity search which goes well beyond searching by keyword which often leads to confusion between say apple as a fruit vs Apple as a computer company. Vector databases are the foundation technology for AI in the enterprise just as relational databases are the foundation technology for business data. Applications in future will combine business data with the semantics of text, images, audio and video to create rich intelligent interactive applications. For this we need a deeper understanding of vector databases than is broadly available today.

While specialized vector databases are available currently (circa 2025) they are each a little different from the other and require database specific knowledge.. While many such vector databases are available in the cloud, requiring far less specialized maintenance, databases such as Postgres enhanced with vector capability are also similarly available.

In this book we will look at technology where vector database capability is well integrated into the world of relational databases. This allows the structured and unstructured data to be handled in familiar ways and managed using familiar tools for database backup, restore, partitioning and access control etc. Vector (i.e. unstructured) data is saved in familiar database tables without special treatment based on what these tables might contain.

Prior to vector databases, unstructured data was often stored as “blobs” or binary large objects or directly in files with a field in a relational database pointing to the file. In either case the data was considered opaque, hence a blob. We were not able to do much with the data aside from save it and retrieve it except perhaps for text where we could do keyword search.

A new data type, VECTOR

So what is “vector data” anyway and why can’t we just use the blob as is? We could, but with a vector we get a super-power that allows us to add a semantic component to our data and aslo make it genAI-ready. In the language model part of genAI, words, sentences, paragraphs and even documents may be represented as vectors, so to do power lifting in the genAI space an understanding of vector data is essential. Let’s dive a little deeper.

First of all, to start with a familiar domain, let’s think of “vector” as a new data type in a database management system. In Postgres, which will be our mainstay for this book, we have the capability of adding new types to the type system via the Postgres Extension system. In the case of the vector type the pgvector extension provides a capability to create columns of type VECTOR in a table, and provides a vector-native search capability “similarity search” over this data type. If you look at the content of a column of the VECTOR you will see a list of floats, but there’s more to this than meets the eye.

The similarity search capability is what differentiates the VECTOR type from merely saving and fetching an ordered set of floats. It is the similarity search which allows finding text similar, in meaning, to a given piece of text. This vector similarity capability is the foundation of database applications built with LLMs or large language models, like ChatGPT or Claude or local LLMs like LLama3.x The vector similarity search has its own search operator which the pgvector extension adds to the type system of Postgres. We will explore more of this in detail in the next chapter about Vector Embeddings.

But for now let’s take a look at how all this magic happens? There is a process called “embedding” which enables us to take chunks of unstructured (i.e. non-tabular) data and map it to a vector i.e. a list of floats. This sort of mapping preserves the context in which a word (or a sentence) occurs. That is, words closeby in meaning get mapped to vectors close by in vector space. So very loosely speaking, the meaning of a word is strongly associated with the context in which it is found, in use. So to revisit the example of the word “apple” from earlier - when the word “apple”, or sentences with the word “apple”, are found in the context of other sentences about fruit or orchards or baking of pies etc. we can say that the meaning, in this context, of the word “apple” is the fruit, apple. If one the other hand the text we are looking at is found in the context of other text containing descriptions of computers, then we conclude that the word is referring to the computer, Apple. This is somewhat imprecise at this point but it’s meant to give a flavor. Fear not, we will be doing a deep, deep dive on embeddings in the next chapter.

At this point we need to add that the word “embedding” is used with different meanings and this can be confusing to the newcomer.

First of all it’s the process by which data is mapped into a vector and more importantly a collection of data is mapped into a vector space hence it’s a verb - the process of embedding.. So, for example, a few pages of text or a number of books, collectively called a corpus will be transformed via the process of embedding into a set of vectors that follow the rules of a vector space. Second of all it is used as a noun - “an embedding”. In this case we are referring to the output of the embedding process i.e. to the set of vectors, collectively, as “an embedding”. Thus the “process of embedding” produces, as output, “an embedding”. The way to remember this and not be confused is to think of a painter creating a painting. The “process of painting” (verb) produces “a painting” (noun). Finally “an embedding” can refer to a single vector or the whole collection. Which one is which is usually very clear from the context.

What’s different about the VECTOR type?

With this new data type comes a set of operations that are specifically designed for vector manipulation. These operations are fundamental to many AI and machine learning tasks, such as recommendation systems, image recognition, and natural language understanding.

To summarize, these operations include:

1. Similarity comparison: Seeing if a given vector is similar to a reference vector, in meaning

A common operation for this is called Cosine Similarity - measuring the cosine of the angle between two vectors to determine their similarity. Vectors close in meaning have cosine similarity close to 1 (max is 1 for identical vectors) and those that are far away have low cosine similarity.

2. Nearest Neighbor Search: Identifying multiple close vectors in the dataset to a given vector.

In vector embedding space we have something called ANN or Approximate Nearest Neighbors,

which gives an approximate result which for genAI applications is sufficient

3. Vector Addition and Subtraction: Performing mathematical operations on vectors to discover relationships or create new representations. For instance one of the seminal papers that introduced such vector embeddings demonstrated a jaw dropping result which did vector addition to show that

Vec(“Queen”) ~= Vec(“King”) - Vec(“Man”) + Vec(“Woman”)

where Vec(<word>) is the vector embedding for a given word.

Where do we use vector databases?

Vector databases are used in a wide range of applications across various industries. Many of these application areas predate the use of vector databases in genAI and you’ve probably used an application that had an underlying vector database in the background without knowing about it. With genAI, and especially with LLM-based applications becoming more common, vector databases are now more in the foreground. Here are some use cases which may be familiar to the average developer.

Semantic search: Vector databases enable more intelligent search capabilities by understanding the meaning and context of queries, rather than just matching keywords.

Recommendation systems: They can power personalized recommendations for products, content, or services based on user preferences and behaviors. Examples would be recommendations on ecommerce shopping sites.

Image and video retrieval: Vector databases can efficiently store and search visual data, allowing for content-based image and video search. An example would be Google image search.

Natural language processing: They support tasks like text classification, sentiment analysis, and language translation by representing text as vectors.

Anomaly detection: In cybersecurity and fraud detection, vector databases can identify unusual patterns or behaviors.

Bioinformatics: They can be used for analyzing genetic sequences and protein structures in computational biology.

Computer vision: Vector databases support object recognition, facial recognition, and other visual AI tasks.

Audio processing: Applications include music recommendation, voice recognition, and audio fingerprinting.

Chatbots and conversational AI: They help in understanding user intent and providing more contextually relevant responses.

Drug discovery: In pharmaceutical research, vector databases can assist in identifying potential drug candidates.

Financial analysis: They can be used for market trend analysis, risk assessment, and algorithmic trading.

SQL vs. Vector Databases

To appreciate the unique value of vector databases, let’s compare them with traditional SQL (Structured Query Language) databases, which have been the backbone of enterprise data management for decades.

The Foundation of Business Math: Accounting Arithmetic

Relational databases, commonly called “SQL databases” have been the go-to solution for enterprise data management and applications because they are built from the ground up with the fundamental mathematical operations used in business: accounting arithmetic. This is reflected in the structure and capabilities of SQL databases:

Tabular Data: SQL databases organize data into tables with rows and columns, mirroring the structure of spreadsheets and ledgers used in accounting.

Data Types: They primarily deal with simple data types such as integers, floats, strings and BCD or binary coded decimals which are used for exact calculations of currency. These are sufficient for most business calculations and record-keeping.

CRUD Operations: SQL databases are designed with Create, Read, Update, and Delete (CRUD) operations, which map well to typical business processes.

Aggregations and Joins: They provide powerful capabilities for summing up values, joining related data, and performing other operations crucial for financial reporting and analysis.

Data Warehousing: SQL databases have evolved to accommodate fully denormalized tables i.e. fact tables, star and snowflake schemas for decision support queries.

Vector Representation in RDBMS

While it’s technically possible to represent vectors in a relational database management system (RDBMS), it’s not an ideal solution. A vector could be stored as a list of floating-point numbers, for example:

```sql
CREATE TABLE vector_data (
    id INT PRIMARY KEY,
    vector FLOAT[]
);
INSERT INTO vector_data (id, vector) VALUES (1, ARRAY[0.1, 0.2, ..., 0.768]);
```

However, this approach has several limitations:

Performance: Querying and manipulating these vectors using SQL operations is not efficient, especially for high-dimensional data, since traditional indexes are not optimized for vector search.

Lack of Vector Operations: SQL databases don’t natively support vector-specific operations like similarity search or nearest neighbor search which we talked of earlier.

Scalability Issues: As the number and dimensionality of vectors grow, the performance of SQL databases for vector-related operations degrades rapidly.

The Need for Vector-Specific Capabilities

To use and maximally exploit the power of vector representations, we need database capabilities that can:

Efficiently Index High-Dimensional Data: Vector databases use specialized indexing techniques like LSH (Locality-Sensitive Hashing) or HNSW (Hierarchical Navigable Small World) to make similarity searches fast and scalable.

Perform Vector Operations Natively: They provide built-in functions for operations like cosine similarity, Euclidean distance calculation, and k-nearest neighbor search.

Optimize for Approximate Nearest Neighbor (ANN) Search: In many AI applications, finding the approximate nearest neighbors is sufficient and much faster than an exact search. Vector databases are designed with this in mind.

Handle Sparse Vectors: Many AI models produce sparse vectors (vectors with many zero values), and vector databases can store and query these efficiently.

NoSQL vs. Vector Databases

As the limitations of SQL databases for handling unstructured and semi-structured data became apparent, NoSQL databases emerged circa 2005 as a partial solution. They were able to handle many data types like JSON and XML with which SQL databases quite often struggled, based on the incompatibility between the hierarchical structure of JSON and XML vs the flat structure of tables.

Now, with the rise of genAI applications, it’s worth looking at NoSQL databases to see if they can be put to use as vector databases.

NoSQL Databases and Vector Storage

A number of NoSQL vendors have recognized the importance of vector data and have introduced extensions or features to handle vectors. Some of the well known common NoSQL databases that have added vector capabilities are:

Redis: Redis has introduced RediSearch with vector similarity search capabilities.

MongoDB: MongoDB Atlas now includes vector search functionality.

Elasticsearch: Elasticsearch has added dense vector fields and vector functions.

Cassandra: Apache Cassandra can be extended with vector search capabilities through plugins.

Limitations of Vector Extensions in NoSQL Databases

While all the above NoSQL databases are being used for vector similarity, they are still retrofits compared to dedicated vector databases:

Performance: The vector operations are often not as optimized as in dedicated vector databases, especially for very high-dimensional data or large-scale deployments.

Functionality: The range of vector operations and indexing methods may be limited compared to specialized vector databases.

Scalability: Some NoSQL databases struggle to maintain performance as the volume of vector data grows into billions of vectors.

Integration: Vector operations may not be as well-integrated with other database features, leading to complexity in queries that combine vector and non-vector data.

When to Choose NoSQL with Vector Extensions

And yet, NoSQL databases with vector extensions can be just fine in certain scenarios:

Mixed Data Types: When your application deals with a mix of structured, semi-structured, and vector data, with vector operations not being the primary focus.

Existing Infrastructure: If you already have a significant investment in a NoSQL database and need to add vector capabilities without a complete overhaul. You have non-trivial training and personnel investment in the NoSQL database in question.

Moderate Scale: For applications that don’t require extreme scale or performance in vector operations. As we will see for most of the rest of the book, this same reason applies also for SQL databases with vector extension - if it’s not a massive scale database with billions of vectors a hybrid of NoSQL + Vector or SQL + Vector will be just fine. In the rest of the book we will use PostgreSQL with the pgvector extension for enterprise department level applications.

Hybrid Approaches: Combining Structured and Vector Data

In real-world applications, it’s often necessary to work with both vector data and traditional structured data such as commonly found in business applications. This has led to the development of hybrid approaches that combine the strengths of vector databases with those of traditional SQL databases.

The Need for Both Vector and Metadata

Most practical applications require more than just pure vector data. They also need associated metadata - structured information about the vectors.

Concretely:

- In a document retrieval system, each document might be represented as a vector, but we also need metadata like the document’s title, author, publication date, and categories.

- In an image search system, each image might be represented as a vector, but we also need metadata like the image file name, size, creation date, and tags, often GPS data and camera information might be part of the metadata.

Limitations of Pure Vector Storage

Storing metadata in vector format is not feasible or efficient for several reasons:

Query Flexibility: It’s much easier to perform filtering and sorting operations on structured metadata using traditional database queries. There are currently no such filter and sort capabilities in vector storage systems.

Storage Efficiency: Storing simple metadata as vectors would be wasteful in terms of storage space.

Indexing: Traditional databases provide efficient indexing mechanisms for structured data, which are crucial for quick lookups and filtering. Vector indexing is based on optimization for similarity and clustering but not on subsetting and filtering based on arbitrary predicates.

Data Integrity: Relational databases offer features like constraints and transactions that are important for maintaining the integrity of structured data.

Hybrid Database Architecture

To address these needs, most vector database systems adopt a hybrid architecture that combines vector storage with traditional database capabilities:

Vector Component: This part of the database is optimized for storing and querying high-dimensional vectors. It uses specialized indexing structures and algorithms for efficient similarity search and other vector operations.

Metadata Component: This is often a traditional relational database (like SQLite) or a NoSQL database that stores the structured metadata associated with each vector.

Integration Layer: This component provides a unified interface for querying both vector and metadata components, allowing for complex queries that combine vector similarity search with metadata filtering.

Example of a Hybrid Query

In a hybrid vector database system, a typical query might look like this:

```
FIND TOP 10 SIMILAR VECTORS TO [0.1, 0.2, ..., 0.5]
WHERE metadata.category = 'science'
AND metadata.publish_date > '2023-01-01'
ORDER BY similarity DESC, metadata.views DESC
```

This query combines a vector similarity search with metadata filtering and sorting, demonstrating the power of the hybrid approach. Note: the query syntax is informal and meant to show the kind of things we want from a hybrid query. It is not the formal syntax of any known hybrid database. However a database that has both SQL and Vector capability will have formal syntax to achieve this goal. In upcoming chapters we will see how this is done using PostgreSQL with the pgvector extension.

Benefits of the Hybrid Approach

The hybrid architecture offers several advantages over the pure vector approach but we must remember that it is most useful for small and medium sized databases. For databases with billions of objects performance may require looking at pure vector databases.

Flexibility: The hybrid architecture allows for complex queries that leverage both vector similarity and traditional enterprise SQL data operations, allowing filtering by one and then the other.

Performance: Each component (vector and metadata) can be optimized for its specific data type and query patterns. Indexes can be chosen separately for tabular data and for vector data to optimize performance of our hybrid queries, based on what the columns by which we are selecting and filtering.

Data Management: It’s easier to manage, update, and maintain metadata separately from the vector data. Usually the metadata table is separate from the table with vector data, with a foreign key id connecting the vector data with its metadata.

Chapter wise summary

Chapter 2:

In this chapter we dive deep into the core technology behind vector databases which is vector embeddings i.e. taking chunks of raw data from text,audio, image and video sources and mapping them to vectors i.e. an ordered set of floats indexed by dimension. We look in some more detail at embeddings and how they emerged as a key foundational aspect of genAI especially LLMs. We explore the seminal Word2Vec technology and the paper in which it was introduced, looking at what NLP was before and after. We then dive even deeper into how embeddings are used not just for words but for sentences and even documents. We look at role of embedding layers in different kinds of transformers and how embeddings in transformers differ from embedding models i.e. models created purely for the purpose of generating embeddings from data. Finally we get hands on with the word2vec model and show how to engage with the code and libraries to get a foundational understanding of word embeddings.

Chapter 3:

Here we start our practical exploration of vector databases with the familiar SQLite database with a vector search extension called SQLite-VSS. We set up a mini vector database environment and we populate it with personal data for querying email, texts, word docs and PDF files. It is intended to be a prototype personal data management system for further development for personal exploration and use.

Chapter 4:

In this chapter we apply similar techniques using PostgreSQL with the pgvector extension. We dive deep into the current (circa late 2024) version of the extension, learning about the various indexes and similarity metrics and what they mean theoretically and in practice. We build a similar application but explore techniques to accommodate much larger data sets so as to support a departmental level application.

Chapter 5:

FAISS was one of the earliest similarity search libraries built by the Fundamental AI Research team at Meta. It was meant to accommodate datasets that did not fit in memory. It is written in C++ and allows the use of GPU if available. In this chapter we explore using FAISS for high performance similarity search. It can also do k-nearest neighbor search and can process vectors in batches.

Chapter 6:

Retrieval Augmented Generation or RAG is the one major application of vector databases and LLMs that has emerged in the first couple of years of the genAI phenomenon hitting the consumer space. RAG combines your local private document corpus i.e. domain knowledge with the language specific expertise that an LLM brings to the table. We are then able to query our document base in natural language.

Here we expand on our efforts in Ch. 3 where we built a vector store using SQLite-VSS. We add an LLM to create a RAG environment - in this case using Ollama which is a runtime environment to invoke and query an LLM of choice.

Chapter 7:

Just as in Ch. 4 we built a scalable vector database, in this chapter we build a scalable RAG application using that same vector database and adding Ollama as in the previous chapter.

Chapter 8:

Here we explore a number of research topics which may not have immediate application but show practical promise in the future. These areas, in our opinion, are worth tracking for those wanting to stay on the leading edge of this technology in these exciting times.

Chapter 9:

We summarize our progress through this book from the start - beginning with relational databases and relational-vector hybrids, understanding embeddings and then using our knowledge to build useful applications.

Conclusion

Vector databases are a major evolutionary step in database technology, driven primarily by the needs of genAI and machine learning applications. They add novel concepts of semantic search and vector space embeddings of different kinds of unstructured data such as text, audio, image and video. With hybrid architectures they complement existing database technologies rather than replacing them entirely.

The future of data management likely lies in hybrid systems that can efficiently handle both vector data and traditional structured data. As AI continues to permeate various domains, the importance of vector databases and hybrid database architectures will only grow.

Understanding the strengths and limitations of vector databases, as well as how they fit into the broader ecosystem of data management solutions, is crucial for data scientists, engineers, and decision-makers working on cutting-edge AI applications. As this field continues to evolve, we can expect to see further innovations that bridge the gap between vector-based and traditional data management approaches.

Chapter 2. Embeddings

A note for Early Release readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 2nd chapter of the final book. The GitHub repo will be made active later on.

If you have comments about how we might improve the content and/or examples in this book, or if you notice missing material within this chapter, please reach out to the editor at shunter@oreilly.com.

At their core, vector embeddings solve a fundamental challenge in computing: how to represent complex, unstructured data in a way that machines can understand and efficiently process. In this chapter you’ll dive deeper into the history, the evolution and the current usage of vector embeddings with code examples for applications to language models. This is a conceptually meaty topic and you may find that you have to revisit this more than once when faced with some unfamiliar terminology or concept in later chapters when you start applying embeddings.

2.1 Understanding Vector Embeddings: Why We Need Them

Traditional data structures work well for structured information like numbers, dates, and categorical values, but they fall short when dealing with the rich, nuanced content that humans naturally create and consume—text, images, audio, and video. Embeddings bridge this gap by mapping raw data into vectors (ordered lists of numbers) that capture the inherent meaning and relationships within the data. This transformation from unstructured data into mathematical vectors isn’t just a technical convenience; it’s the foundation that enables modern AI systems, including Large Language Models (LLMs) and Generative AI, to process and understand information in ways that approach human-like comprehension.

Vector databases have emerged as specialized systems to store and query these embeddings efficiently, providing built-in capabilities for finding similar vectors—and by extension, similar content—at scale. This marks a fundamental shift in how we can search and analyze information. Instead of being limited to exact keyword matches or rigid categorical filters, vector embeddings enable search by meaning or semantic similarity. For example, a search for “sunset by the ocean” can find relevant images even if they’re labeled differently. Or, a query about “renewable energy solutions” can surface documents discussing solar and wind power without requiring exact phrase matches. This capability to find similar vectors quickly and efficiently is what makes vector databases essential infrastructure for modern AI applications.

The impact of embeddings extends far beyond just improved search capabilities. They serve as the foundational data representation that powers Large Language Models and Generative AI systems. These models work with enormous numbers of embeddings that capture the relationships and patterns in language, images, and other modalities. By converting unstructured data into vector representations, we create a universal format that AI systems can process to understand context, generate relevant responses, and make sophisticated connections across different types of content. This transformation of raw data into mathematically comparable vectors has become so fundamental to AI that it’s fair to say that modern artificial intelligence is built on a foundation of vector embeddings.

2.2 Word2Vec: The Breakthrough That Changed Everything

Before word2vec emerged in 2013, the world of Natural Language Processing (NLP) was dominated by rule-based systems and simple statistical methods. Text was primarily represented using techniques like “bag of words” and TF-IDF (Term Frequency-Inverse Document Frequency), where words were treated as discrete symbols with no inherent relationship to each other. Under these approaches, the words “king” and “queen” were as mathematically distinct as “king” and “bicycle.” This made it nearly impossible for machines to understand semantic relationships or capture the subtle nuances of language. Attempts to encode meaning often relied on hand-crafted rules and carefully curated lexical databases like WordNet, making language understanding systems brittle and limited in their capabilities.

Word2vec, introduced by Tomas Mikolov and his team at Google, represented a fundamental paradigm shift. The breakthrough was as elegant as it was powerful: words could be represented as dense vectors in a continuous vector space, where the geometric relationships between these vectors captured meaningful semantic relationships. The core insight was to predict either a word from its context or the context from a word (known as skip-gram and continuous bag of words models, respectively). Through this process, words that appeared in similar contexts would end up with similar vector representations. The resulting embeddings exhibited almost magical properties - vector arithmetic like “king - man + woman = queen” became possible, demonstrating that the model had learned to encode genuine semantic and syntactic relationships.

The mapping process itself was remarkably straightforward yet effective. Word2vec trained a shallow neural network on a large corpus of text, but instead of using the network for prediction, it extracted the weights from the hidden layer as the word embeddings. These weights captured the statistical patterns of word co-occurrence in a way that preserved semantic relationships. For example, words that frequently appeared in similar contexts (like “coffee” and “tea”) would end up with similar vector representations, while words that rarely appeared in similar contexts would be positioned far apart in the vector space. This approach was computationally efficient and could be trained on massive amounts of text data without requiring labeled examples.

Word2vec’s influence on the field cannot be overstated - it sparked a revolution that continues to this day. The paper demonstrated that meaningful representations of language could be learned automatically from raw text, without human supervision or rule-crafting. This led to a cascade of innovations: GloVe embeddings refined the approach by explicitly modeling word co-occurrence statistics, and fastText extended it to handle subword information. The concept of learned embeddings quickly spread beyond words to sentences, paragraphs, and eventually entire documents. Most significantly, the idea of learning contextual representations through prediction tasks became a cornerstone of modern language models. The transformer architecture, which powers today’s Large Language Models like GPT and BERT, can be seen as a direct descendant of word2vec’s fundamental insights - now operating at a much larger scale and with more sophisticated architectural components. In many ways, the current AI revolution, with its remarkably capable language models, traces its conceptual roots back to that seminal word2vec paper.

2.3 Doc2Vec: From Words to Documents

Building on the success of Word2Vec, Mikolov and Le introduced Doc2Vec (originally called Paragraph Vector) in 2014, addressing a fundamental limitation of word embeddings: the need to represent larger pieces of text as single vectors. While Word2Vec excelled at capturing word-level semantics, averaging or combining word vectors to represent entire documents often lost crucial information about word order, document structure, and overall meaning. Doc2Vec elegantly extended the Word2Vec architecture to learn fixed-length vector representations for variable-length pieces of text, from sentences to paragraphs to entire documents.

The genius of Doc2Vec lay in its simple yet effective modification to the Word2Vec training process. In addition to learning word vectors, the model simultaneously learned a unique vector for each document in the training corpus. This document vector acted as a memory that captured the topics and global semantic content of the text. During training, the document vector was treated as if it were just another word in the context, but unlike word vectors, it provided a consistent context across all windows of words within the same document. This approach allowed the model to capture both the general topic of the document and the precise meaning of words in context - something that simple averaging of word vectors could never achieve.

The differences between Word2Vec and Doc2Vec embeddings became apparent in their practical applications. While Word2Vec vectors could tell you that “cat” and “kitten” were related words, Doc2Vec could tell you that “The cat sat on the mat” and “A kitten rested on the rug” were similar sentences, despite having few words in common. This capability made Doc2Vec particularly powerful for tasks like document classification, recommendation systems, and plagiarism detection. For example, two news articles about the same event written with different vocabulary would receive similar vector representations, while articles about different topics would be placed far apart in the vector space, even if they shared some common words.

Consider these practical examples of the difference:

A Word2Vec approach to representing the sentence “The bank is by the river” would struggle to disambiguate whether “bank” refers to a financial institution or a riverbank, as it would typically average or combine individual word vectors.

Doc2Vec, on the other hand, would learn a representation that captures the full context, clearly encoding that this is a sentence about geography rather than finance.

Similarly, Doc2Vec could recognize that “The movie was fantastic” and “The film was excellent” express nearly identical sentiments, while simple word vector averaging might miss this equivalence due to the different word choices.

The implications of Doc2Vec extended far beyond just document similarity. It provided one of the first successful demonstrations of learning fixed-length representations for variable-length text, a concept that would become crucial for modern language models. The ability to encode entire documents as dense vectors that preserved semantic meaning paved the way for more sophisticated document embedding techniques like Universal Sentence Encoder and the attention-based embeddings used in transformer models. While newer approaches have since surpassed Doc2Vec in terms of performance, its elegant solution to the document embedding problem remains influential in how we think about representing text at different scales.

Sparse vs. Dense Embeddings - A Tale of Two Representations

Before the deep learning revolution, sparse embeddings dominated NLP through techniques like one-hot encoding and TF-IDF vectors. Imagine representing the word “cat” in a vocabulary of 100,000 words: its sparse representation would be a vector of length 100,000 containing all zeros except for a single ’1’ in the position corresponding to “cat”. Similarly, a document would be represented by a sparse vector showing which words appear in it and how frequently. These representations were intuitive but extremely inefficient - most of the vector was filled with zeros, and the dimensionality grew linearly with vocabulary size.

Dense embeddings, introduced through models like Word2Vec, revolutionized this approach by mapping words or documents into a much lower-dimensional space (typically 100-1000 dimensions) where every dimension carries meaningful information. Unlike sparse vectors, dense embeddings pack semantic information into every dimension, with each number in the vector contributing to the overall meaning. For example, a 300-dimensional Word2Vec embedding might use different dimensions to encode concepts like gender, plurality, formality level, or semantic category - though these dimensions aren’t necessarily human-interpretable.

The trade-offs between sparse and dense embeddings are significant:

Sparse Embeddings:

Easy to interpret (each dimension corresponds to a specific word/feature)

Simple to compute

Memory-inefficient

Poor at capturing relationships between words

Work well with traditional machine learning algorithms

Dense Embeddings:

More compact and memory-efficient

Better at capturing semantic relationships

Enable meaningful vector arithmetic

Harder to interpret individual dimensions

Require significant data and computation to train effectively

Modern systems sometimes combine both approaches. For instance, some search engines use dense embeddings for semantic similarity and sparse embeddings for exact keyword matching, getting the best of both worlds. The rise of neural networks has firmly established dense embeddings as the dominant paradigm, but sparse representations remain useful in specific scenarios where interpretability or exact matching is crucial.

2.4 From Embeddings to Modern Language Models: The Transformer Connection

Vector embeddings form the fundamental building blocks of modern Large Language Models (LLMs), serving as the bridge between human language and machine understanding. Every word, sentence, or token that enters a language model must first be converted into a vector representation that the model can process. These embeddings aren’t just input formats - they’re the mathematical space where all the model’s reasoning and generation happens. When GPT-4 writes a poem or Claude analyzes a legal document, they’re performing complex mathematical operations on vectors, with each layer of the model transforming these representations in increasingly sophisticated ways. This deep connection between embeddings and language models is why understanding embeddings is crucial for anyone working with modern AI - they’re not just a technical detail, but the very medium in which artificial intelligence “thinks.”

At the heart of all modern LLMs lies the transformer architecture, which has proven to be a remarkably powerful way to process and generate sequences of embeddings. While Word2Vec and Doc2Vec showed us how to create meaningful vector representations, transformers showed us how to dynamically manipulate and generate these representations at scale. Let’s explore the three main variants of transformer architectures and how they use embeddings in distinct ways.

Encoder-Only Transformers (BERT and Its Variants)

Encoder-only transformers, like BERT (Bidirectional Encoder Representations from Transformers), specialize in understanding input text by creating rich, contextual embeddings. Unlike the static embeddings of Word2Vec, these models create dynamic embeddings that change based on context. The same word “bank” gets different vector representations in “river bank” versus “bank account” because the encoder processes the entire sequence bidirectionally. Each layer of the encoder progressively refines these embeddings, with the final layer producing highly contextualized representations. These models excel at:

Text classification

Named entity recognition

Sentiment analysis

Question answering

The embeddings they produce capture deep semantic relationships and can be used as sophisticated feature extractors for downstream tasks.

Decoder-Only Transformers (GPT Family)

Decoder-only transformers, exemplified by GPT (Generative Pre-trained Transformer) models, are optimized for text generation. They process embeddings sequentially, with each position only attending to previous positions - much like how humans write text from left to right. The embedding process here is particularly interesting because it must capture not just meaning, but also predictive information about what tokens might come next. In GPT models:

Input embeddings combine token, position, and segment information

Each layer transforms these embeddings through self-attention and feed-forward networks

The final layer projects embeddings into probability distributions over the next token

This architecture has proven remarkably effective at tasks like:

Text completion

Creative writing

Code generation

Conversational AI

Encoder-Decoder Transformers (T5, BART)

Encoder-decoder transformers, also known as sequence-to-sequence models, combine both approaches to handle tasks that require both understanding and generation. Models like T5 (Text-to-Text Transfer Transformer) use:

An encoder to create contextual embeddings of the input

A decoder to generate output while attending to both the encoded input and its own generated sequence

Cross-attention mechanisms that allow the decoder to focus on relevant parts of the encoded input This architecture shines in tasks like:

Translation

Summarization

Text simplification

Structured data generation

In all three architectures, embeddings serve different but crucial roles:

Input embeddings convert tokens into vectors the model can process

Positional embeddings encode sequence order information

Layer embeddings capture increasingly abstract features

Output embeddings (in generative models) map back to token probabilities

The sophistication of how transformers manipulate these embeddings - through self-attention, cross-attention, and feed-forward networks - is what enables them to achieve their remarkable performance. The ability to dynamically update and refine embeddings based on context, while maintaining the mathematical properties that make vector operations meaningful, has proven to be a powerful combination that underlies much of modern AI’s capabilities.

2.5 Embedding Models: The Specialized Vector Generators

So far in this chapter I’ve discussed embeddings as components within larger language models. Let’s now move to dedicated embedding models, which have emerged as crucial standalone tools in modern AI architectures. These specialized models, like OpenAI’s text-embedding-ada-002 or Co embed-multilingual-v3, are designed with a singular purpose: to create high-quality vector representations of text that can be used for downstream applications. Unlike general-purpose LLMs that generate text or perform various language tasks, embedding models focus exclusively on mapping text to vectors that capture semantic meaning with high fidelity.

The Distinction from Traditional Models

The key differences between embedding models and previously discussed approaches are both architectural and philosophical:

Specialized Architecture: Unlike Word2Vec or Doc2Vec, modern embedding models are often based on transformer architectures but optimized specifically for embedding generation. They typically remove the language generation components and focus entirely on creating meaningful vector representations.

Consistency: While LLMs create dynamic, context-dependent representations internally, embedding models are designed to generate consistent vectors for the same input, making them ideal for search and retrieval systems.

Computational Efficiency: Because they don’t need to generate text or perform complex language tasks, embedding models can be much smaller and faster than full LLMs while still producing high-quality embeddings.

Optimization Target: These models are specifically trained to maximize the usefulness of their vector representations for similarity comparisons, often using techniques like contrastive learning to ensure similar texts are mapped to similar vectors.

Role in Modern LLM Applications

Embedding models have become essential components in what’s often called the “Retrieval-Augmented Generation” (RAG) pattern, which is transforming how we build AI applications. Let’s look at the RAG pattern below and you can see where embedding models have been integrated:

Document Processing Pipeline:

Documents are split into chunks

Each chunk is processed by an embedding model to create a vector representation

These vectors are stored in a vector database for efficient retrieval

Query Processing:

User queries are converted to vectors using the same embedding model

Similar vectors are retrieved from the database

Retrieved content provides context to the LLM for generating responses

Hybrid Search Systems:

Embedding models enable semantic search capabilities

Can be combined with traditional keyword search for better results

Support multi-modal search across text, images, and other data types

Practical Applications and Use Cases

At this point you may be wondering how these embedding models are utilized apart from their original role in large language models. Modern embedding models have enabled a wide range of applications:

Semantic Search: Finding relevant documents based on meaning rather than keywords

Content Recommendation: Suggesting similar articles, products, or media

Duplicate Detection: Identifying similar content across large datasets

Data Deduplication: Finding and removing near-duplicate entries in databases

Knowledge Base Construction: Organizing and linking related pieces of information

Question Answering: Finding relevant context for answering specific queries

Content Clustering: Grouping similar documents or passages automatically

a concrete example of how this works in practice semantic search:

# Simplified RAG Pipeline Example
from embedding_model import get_embedding  # Hypothetical embedding model
from vector_db import VectorDB

# Document processing
document = "Climate change is affecting global weather patterns..."
doc_embedding = get_embedding(document)
vector_db.store(doc_embedding)

# Query processing
query = "How does climate change affect weather?"
query_embedding = get_embedding(query)
similar_docs = vector_db.search(query_embedding, top_k=3)

# Feed to LLM with context
llm_response = generate_response(query, context=similar_docs)

Latest Developments and Trends

The field of embedding models continues to evolve rapidly. Here’s some of the latest developments at the time of writing.

Multi-modal Embeddings: Newer models can create compatible embeddings across different types of data (text, images, code)

Instruction-Tuned Embeddings: Models that can adapt their embedding strategy based on specific instructions

Domain-Specific Embeddings: Models optimized for particular fields like legal, medical, or scientific text

Hierarchical Embeddings: Systems that create embeddings at multiple levels of granularity (words, sentences, paragraphs)

Efficient Embedding: Techniques to reduce embedding dimension while maintaining semantic quality

Understanding and effectively using embedding models has become a crucial skill for AI engineers, as they form the foundation of many modern AI applications. While large language models might get more attention, it’s often the careful use of embedding models that determines whether an AI application is practical and effective in production.

Here’s a detailed example using the sentence-transformers library:

from sentence_transformers import SentenceTransformer
from chromadb import Client, Settings
import chromadb

# Initialize the embedding model
# all-MiniLM-L6-v2 is a lightweight but effective model, good balance of speed/quality
model = SentenceTransformer('all-MiniLM-L6-v2')

# Initialize ChromaDB as our vector store
# Using in-memory storage for this example
chroma_client = Client(Settings(is_persistent=False))
collection = chroma_client.create_collection(name="climate_docs")

# Example documents to process
documents = [
    "Climate change is affecting global weather patterns causing more extreme events.",
    "Rising sea levels threaten coastal communities worldwide.",
    "Greenhouse gas emissions continue to rise despite international agreements."
]

# Create embeddings for our documents
# model.encode() converts text to dense vectors (embeddings)
embeddings = model.encode(documents)

# Store documents and their embeddings
# ChromaDB expects embeddings as lists, so we convert numpy arrays
collection.add(
    embeddings=[e.tolist() for e in embeddings],
    documents=documents,
    ids=[f"doc_{i}" for i in range(len(documents))]
)

# Process a query
query = "How does climate change affect weather?"
query_embedding = model.encode(query)

# Search for similar documents
results = collection.query(
    query_embeddings=[query_embedding.tolist()],
    n_results=2
)

# Print results
for doc in results['documents'][0]:
    print(f"Retrieved document: {doc}")

Let’s break down each part:

Model Initialization:

model = SentenceTransformer('all-MiniLM-L6-v2')

Loads a pre-trained Sentence Transformer model

‘all-MiniLM-L6-v2’ is a popular model that creates 384-dimensional embeddings

Good balance of speed and quality for general use

Vector Store Setup:

chroma_client = Client(Settings(is_persistent=False))
collection = chroma_client.create_collection(name="climate_docs")

Initializes ChromaDB as an in-memory vector store

Creates a new collection to store our documents and embeddings

Document Processing:

embeddings = model.encode(documents)

Converts each document into a dense vector (embedding)

Returns numpy arrays of shape (n_documents, 384)

Process is automatically batched for efficiency

Storing Documents:

collection.add(
    embeddings=[e.tolist() for e in embeddings],
    documents=documents,
    ids=[f"doc_{i}" for i in range(len(documents))])

Stores both the original documents and their embeddings

Assigns unique IDs to each document

Converts numpy arrays to lists for ChromaDB compatibility

Query Processing:

query_embedding = model.encode(query)

Creates an embedding for the query using the same model

Ensures query and documents are in the same vector space

Similarity Search:

results = collection.query(
    query_embeddings=[query_embedding.tolist()],
    n_results=2)

Searches for the most similar documents to our query

Uses cosine similarity by default

Returns the top 2 most similar documents

This code demonstrates a simple but complete RAG pipeline. In a production environment, you might want to add:

Error handling

Batch processing for large document sets

Document chunking for long texts

Persistent storage configuration

Distance threshold filtering

Metadata handling

2.6 Sentence Transformers: The Swiss Army Knife of Text Embeddings

Sentence Transformers (sentence-transformers), developed by SBERT.net, has emerged as one of the most practical and widely-used libraries for generating text embeddings. What makes this library particularly valuable is its collection of pre-trained models optimized specifically for creating sentence embeddings, along with its easy-to-use interface and strong performance across a variety of tasks.

Most Popular Models and Their Characteristics

Let’s take a look at the most popular text embedding models and some of their characteristics:

all-MiniLM-L6-v2

Dimensions: 384

Model Size: ~80MB

Speed: Very fast (2,000 sentences/second on GPU)

Pros:

Excellent speed-to-performance ratio

Small model size

Good general-purpose performance

Cons:

Less powerful than larger models for specialized tasks

Lower accuracy compared to larger models

Best for: Production environments where speed and efficiency matter

all-mpnet-base-v2

Dimensions: 768

Model Size: ~420MB

Speed: Moderate (1,000 sentences/second on GPU)

Pros:

State-of-the-art performance

Better semantic understanding

More robust across different domains

Cons:

Larger model size

Slower inference time

Best for: Applications where accuracy is critical

paraphrase-multilingual-mpnet-base-v2

Dimensions: 768

Languages: 50+ languages

Model Size: ~970MB

Pros:

Excellent cross-lingual capabilities

Consistent performance across languages

Can compare text across different languages

Cons:

Large model size

Slower than monolingual models

Best for: Multilingual applications

Training Approach

These models are typically trained using a combination of techniques:

Knowledge Distillation: Smaller models like all-MiniLM-L6-v2 are trained to mimic larger models while being more efficient.

Contrastive Learning: Models are trained on pairs of similar sentences (positive examples) and dissimilar sentences (negative examples).

Multiple Training Objectives:

Natural Language Inference (NLI) data

Paraphrase detection

STS (Semantic Textual Similarity) benchmarks

Question-Answer pairs

all-MiniLM-L6-v2 for RAG Applications

For this book, I’ll be using all-MiniLM-L6-v2 extensively in future chapters because it represents an optimal balance for enterprise applications:

Fast enough for real-time applications

Small enough to deploy easily

Accurate enough for most use cases

Excellent community support and documentation

Production-ready performance

Let’s dive into an example. Here’s a practical example demonstrating both monolingual and multilingual capabilities:

from sentence_transformers import SentenceTransformer
import numpy as np

# Load our two models
monolingual_model = SentenceTransformer('all-MiniLM-L6-v2')
multilingual_model = SentenceTransformer('paraphrase-multilingual-mpnet-base-v2')

# Example sentences in different languages
sentences = {
    'english': "The weather is beautiful today.",
    'spanish': "El clima está hermoso hoy.",
    'french': "Le temps est magnifique aujourd'hui.",
    'german': "Das Wetter ist heute wunderschön."
}

# Function to compute similarity between embeddings
def compute_similarity(emb1, emb2):
    return np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))

# Compare embeddings across languages
def compare_sentences(model, sentences):
    # Generate embeddings for all sentences
    embeddings = {
        lang: model.encode(text, convert_to_numpy=True)
        for lang, text in sentences.items()
    }

    # Compare each pair
    print(f"\nSimilarity scores for {model.__class__.__name__}:")
    for lang1 in sentences:
        for lang2 in sentences:
            if lang1 < lang2:  # avoid duplicate comparisons
                sim = compute_similarity(embeddings[lang1], embeddings[lang2])
                print(f"{lang1} vs {lang2}: {sim:.4f}")

# Test both models
for model in [monolingual_model, multilingual_model]:
    compare_sentences(model, sentences)

# Example of batch processing for efficiency
texts = list(sentences.values())
batch_embeddings = multilingual_model.encode(texts, batch_size=8)

When you run this code, you’ll notice that:

The multilingual model produces similar scores for the same sentence across languages

The monolingual model performs less consistently with non-English text

Batch processing is significantly faster than individual encoding

Best Practices When Using Sentence Transformers

Batch Processing: Always use batching for large datasets

embeddings = model.encode(sentences, batch_size=32)

Length Handling: Be aware of maximum sequence length

# For long texts, consider chunking
max_seq_length = model.max_seq_length

Normalization: Consider normalizing embeddings for cosine similarity

embeddings = model.encode(sentences, normalize_embeddings=True)

GPU Utilization: Move to GPU for faster processing

model = SentenceTransformer('all-MiniLM-L6-v2', device='cuda')

Best Practices for Sentence Transformers: A Detailed Guide

1. Batch Processing

# Good Practice
embeddings = model.encode(sentences, batch_size=32)

# Instead of
# Bad Practice
embeddings = [model.encode(sentence) for sentence in sentences]

Why it matters:

Batching allows parallel processing of multiple sentences, utilizing GPU/CPU resources more efficiently

Can be 10-100x faster than processing sentences individually

Reduces memory fragmentation and system overhead

Consequences of not batching:

Significantly slower processing time

Higher memory overhead due to repeated model loading

Inefficient use of hardware capabilities, especially GPUs

Can lead to higher costs in production environments

2. Length Handling

# Good Practice
max_seq_length = model.max_seq_length
def chunk_text(text, max_length=max_seq_length):
    # Split into sentences or chunks
    chunks = [text[i:i + max_length] for i in range(0, len(text), max_length)]
    return chunks

# Process long document
long_text_embeddings = model.encode(chunk_text(long_document))

Why it matters:

Prevents truncation of important information

Ensures consistent processing of long documents

Maintains model performance on lengthy texts

Avoids unexpected behavior or errors

Consequences of ignoring length limits:

Silent truncation of input text

Loss of potentially crucial information

Inconsistent results for long documents

Possible degradation of embedding quality

3. Normalization

# Good Practice
embeddings = model.encode(sentences, normalize_embeddings=True)

# Or manually normalize if needed
from sklearn.preprocessing import normalize
embeddings = normalize(embeddings)

Why it matters:

Makes cosine similarity calculations more efficient

Ensures consistent similarity scores across different lengths of text

Reduces the impact of text length on similarity calculations

Important for comparing embeddings across different documents

Consequences of skipping normalization:

Inconsistent similarity scores

Length bias in comparisons

Less reliable search results

More complex downstream processing

4. GPU Utilization

# Good Practice
model = SentenceTransformer('all-MiniLM-L6-v2', device='cuda')

# With error handling
import torch
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = SentenceTransformer('all-MiniLM-L6-v2', device=device)

Why it matters:

Can be 10-50x faster than CPU processing

Essential for processing large volumes of text

Enables real-time applications

More cost-effective in production environments

Consequences of not using GPU when available:

Significantly slower processing speed

Higher operational costs

Poor performance in real-time applications

Inefficient resource utilization

5. Memory Management (Additional Best Practice)

# Good Practice
import gc
import torch

def process_large_dataset(sentences, batch_size=32):
    embeddings = []
    for i in range(0, len(sentences), batch_size):
        batch = sentences[i:i + batch_size]
        batch_embeddings = model.encode(batch)
        embeddings.extend(batch_embeddings)

        # Clear GPU memory if needed
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()

    return embeddings

Why it matters:

Prevents out-of-memory errors with large datasets

Ensures stable performance over long running periods

Crucial for production environments

Enables processing of very large datasets

Consequences of poor memory management:

Random crashes due to memory exhaustion

Degraded performance over time

System instability

Limited ability to process large datasets

6. Error Handling (Additional Best Practice)

# Good Practice
from typing import List, Optional
import numpy as np

def safe_encode(
    model: SentenceTransformer,
    texts: List[str],
    batch_size: int = 32
) -> Optional[np.ndarray]:
    try:
        # Handle empty or invalid inputs
        if not texts or not isinstance(texts, list):
            return None

        # Remove empty strings and handle other edge cases
        valid_texts = [str(t) for t in texts if t and str(t).strip()]
        if not valid_texts:
            return None

        return model.encode(valid_texts, batch_size=batch_size)

    except Exception as e:
        print(f"Error encoding texts: {e}")
        return None

Why it matters:

Prevents crashes from invalid inputs

Makes the system more robust

Easier to debug issues

Better production reliability

Consequences of poor error handling:

System crashes from unexpected inputs

Difficult to diagnose issues

Poor user experience

Higher maintenance overhead

2.7 The Embedding Layer: Gateway to Zero-Shot Learning

The embedding layer in transformer models is more than just a lookup table for converting tokens to vectors - it’s the foundation that enables these models to understand and generalize to new concepts without explicit training, a capability known as zero-shot learning. Let’s explore this crucial connection and understand how embedding layers enable this remarkable capability.

The Anatomy of Transformer Embeddings

A transformer’s embedding layer actually consists of three components that work together:

# Simplified implementation of transformer embeddings
import torch
import torch.nn as nn

class TransformerEmbeddings(nn.Module):
    def __init__(self, vocab_size, embedding_dim, max_sequence_length):
        super().__init__()

        # Token embeddings: convert words to vectors
        self.token_embeddings = nn.Embedding(vocab_size, embedding_dim)

        # Position embeddings: add sequence order information
        self.position_embeddings = nn.Embedding(max_sequence_length, embedding_dim)

        # Segment embeddings: distinguish different parts of input (e.g., question vs context)
        self.segment_embeddings = nn.Embedding(2, embedding_dim)

        # Layer normalization for stable training
        self.layer_norm = nn.LayerNorm(embedding_dim)
        self.dropout = nn.Dropout(0.1)

    def forward(self, input_ids, segment_ids=None):
        # Get sequence length from input
        seq_length = input_ids.size(1)

        # Create position indices (0, 1, 2, ...)
        position_ids = torch.arange(seq_length, device=input_ids.device)
        position_ids = position_ids.unsqueeze(0).expand_as(input_ids)

        # If segment ids not provided, assume all zeros
        if segment_ids is None:
            segment_ids = torch.zeros_like(input_ids)

        # Combine all embeddings
        embeddings = (
            self.token_embeddings(input_ids) +      # Word meaning
            self.position_embeddings(position_ids) + # Sequential order
            self.segment_embeddings(segment_ids)     # Input segment
        )

        # Normalize and apply dropout
        embeddings = self.layer_norm(embeddings)
        embeddings = self.dropout(embeddings)

        return embeddings

The Connection to Zero-Shot Learning

Let’s take a look at what makes these embeddings special in their role in enabling zero-shot capabilities:

Compositional Nature

Embeddings capture semantic relationships between concepts

New concepts can be understood through combinations of known embeddings

This allows models to handle novel situations by combining familiar elements

Shared Semantic Space

Words with similar meanings cluster together in the embedding space

This clustering enables generalization to unseen examples

Models can infer relationships between concepts they’ve never explicitly seen together

Here’s an example demonstrating zero-shot classification using embeddings:

from sentence_transformers import SentenceTransformer, util
import torch

class ZeroShotClassifier:
    def __init__(self, model_name='all-mpnet-base-v2'):
        self.model = SentenceTransformer(model_name)

    def classify(self, text, candidate_labels):
        # Encode the input text
        text_embedding = self.model.encode(text, convert_to_tensor=True)

        # Prepare label prompts
        label_prompts = [f"This text is about {label}" for label in candidate_labels]
        label_embeddings = self.model.encode(label_prompts, convert_to_tensor=True)

        # Calculate similarities
        similarities = util.pytorch_cos_sim(text_embedding, label_embeddings)[0]

        # Create results dictionary
        results = {
            label: float(score)
            for label, score in zip(candidate_labels, similarities)
        }

        return results

# Example usage
classifier = ZeroShotClassifier()

text = "The new quantum computer can perform calculations in seconds that would take classical computers thousands of years."
labels = ["technology", "sports", "cooking", "politics"]

results = classifier.classify(text, labels)
print("Zero-shot classification results:")
for label, score in sorted(results.items(), key=lambda x: x[1], reverse=True):
    print(f"{label}: {score:.3f}")

Key Characteristics That Enable Zero-Shot Learning

Contextual Understanding

# Example of how context affects embeddings
texts = [
        "The bank by the river is eroding",
        "The bank approved my loan"
]
embeddings = model.encode(texts)
# These will have different embeddings despite sharing the word "bank"

Semantic Relationships

# Demonstration of semantic relationships
def find_analogies(model, a1, a2, b1):
        """Find word b2 that satisfies: a1 : a2 :: b1 : b2"""
            e_a1 = model.encode(a1)
            e_a2 = model.encode(a2)
            e_b1 = model.encode(b1)

    # Calculate the relationship vector
    relationship = e_a2 - e_a1

    # Predict b2's embedding
    predicted_b2 = e_b1 + relationship
    return predicted_b2

Cross-Modal Capabilities: Modern embedding layers can handle multiple modalities:

from transformers import CLIPProcessor, CLIPModel

class MultiModalEmbeddings:
        def __init__(self):
        self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
        self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

    def get_embeddings(self, text=None, image=None):
        inputs = self.processor(text=text, images=image, return_tensors="pt")
        outputs = self.model(**inputs)

        # Get embeddings from the shared space
        if text:
            return outputs.text_embeds
        return outputs.image_embeds

Limitations and Considerations

Quality of Embeddings

More training data generally leads to better zero-shot performance

Domain-specific vocabulary may need special handling

Quality degrades for very specialized or technical concepts

Computational Costs

Larger embedding dimensions provide better representation but increase memory usage

Need to balance model size with performance requirements

Bias and Fairness

Embeddings can inherit biases from training data

Important to validate zero-shot performance across different domains and contexts

2.8 Vector Arithmetic with Word2Vec: A Hands-On Guide

In this section, I’ll demonstrate the fascinating world of vector arithmetic using Word2Vec. You’ll walk through the classic “king - man + woman = queen” example and explore other interesting analogies that demonstrate the semantic relationships captured by word embeddings.

Step 1: Setup and Installation

First, let’s install the required libraries:

pip install gensim
pip install numpy
pip install nltk

Step 2: Load Pre-trained Word2Vec Model

Now, let’s start by loading the word2vec model which is part of the gensim package. It comes pre-trained so we can just start using it for inference.

import gensim.downloader as api
import numpy as np
from typing import List, Tuple

# Download and load the pre-trained Word2Vec model
# This might take a few minutes the first time as it downloads ~1.6GB of data
print("Loading Word2Vec model...")
word2vec_model = api.load('word2vec-google-news-300')
print("Model loaded!")

def cosine_similarity(vec1: np.ndarray, vec2: np.ndarray) -> float:
    """Calculate cosine similarity between two vectors"""
    return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))

def find_similar_words(vector: np.ndarray, n: int = 5) -> List[Tuple[str, float]]:
    """Find n words with vectors most similar to the input vector"""
    return word2vec_model.similar_by_vector(vector, topn=n)

Let’s break down what you just did in this step?:

You loaded Google’s pre-trained Word2Vec model which was trained on news articles

The model contains 300-dimensional vectors for about 3 million words and phrases

You defined helper functions for similarity calculations and finding similar words

Step 3: Implement Vector Arithmetic Functions

Now that you have set things up, let’s get to the exciting part - treating words as vectors and doing arithmetic on words. It’s brave new world!

def vector_arithmetic(*words_and_weights: Tuple[str, float]) -> np.ndarray:
    """
    Perform vector arithmetic with words and weights
    Example: vector_arithmetic(("king", 1), ("man", -1), ("woman", 1))
    """
    resulting_vector = np.zeros(word2vec_model.vector_size)

    for word, weight in words_and_weights:
        if word not in word2vec_model:
            raise ValueError(f"Word '{word}' not found in vocabulary")
        resulting_vector += weight * word2vec_model[word]

    return resulting_vector

def print_analogy_results(result_vector: np.ndarray,
                         original_words: List[str],
                         n_results: int = 5):
    """Print the results of vector arithmetic in a readable format"""
    similar_words = find_similar_words(result_vector, n_results)

    print("\nVector arithmetic:", " + ".join(original_words))
    print("\nMost similar words:")
    for word, similarity in similar_words:
        if word not in original_words:  # Don't show input words in results
            print(f"{word}: {similarity:.4f}")

These functions will help you:

Perform vector arithmetic with any number of words and weights

Format and display the results in a readable way

Filter out input words from the results

Step 4: Classic King-Queen Analogy

You’ve finally arrived at the king-queen analogy. Let’s put our heavy machinery to work and do arithmetic on words! What we are doing here is the analog of the primordial soup from which life, or in our case language models, eventually emerge.

def demonstrate_royal_analogy():
    print("\n=== Royal Analogy Demonstration ===")

    # king - man + woman = queen
    result = vector_arithmetic(
        ("king", 1),
        ("man", -1),
        ("woman", 1)
    )

    print_analogy_results(
        result,
        ["king", "man", "woman"],
        n_results=5
    )

# Run the demonstration
demonstrate_royal_analogy()

This should output something like:

=== Royal Analogy Demonstration ===

Vector arithmetic: king - man + woman

Most similar words:
queen: 0.7118
monarch: 0.6523
princess: 0.6342
kings: 0.6103
prince: 0.5932

Step 5: More Interesting Analogies

There are many more analogies to explore that showcase how word arithmetic works in surprising ways.. Let’s keep exploring.

def explore_more_analogies():
    print("\n=== More Analogy Examples ===")

    # Example 1: Paris is to France as Berlin is to ...?
    result1 = vector_arithmetic(
        ("Paris", -1),
        ("France", 1),
        ("Berlin", 1)
    )
    print("\nParis : France :: Berlin : ?")
    print_analogy_results(result1, ["Paris", "France", "Berlin"])

    # Example 2: Walking is to walked as running is to ...?
    result2 = vector_arithmetic(
        ("walking", -1),
        ("walked", 1),
        ("running", 1)
    )
    print("\nwalking : walked :: running : ?")
    print_analogy_results(result2, ["walking", "walked", "running"])

    # Example 3: Good is to better as bad is to ...?
    result3 = vector_arithmetic(
        ("good", -1),
        ("better", 1),
        ("bad", 1)
    )
    print("\ngood : better :: bad : ?")
    print_analogy_results(result3, ["good", "better", "bad"])

# Run the demonstrations

explore_more_analogies()

Step 6: Interactive Exploration Tool

Now that we’ve tested a few different ways to do arithmetic with words, let’s create a simple tool to explore analogies interactively:

def explore_custom_analogy():
    """Interactive tool for exploring word analogies"""
    print("\n=== Custom Analogy Explorer ===")
    print("Format: word1 is to word2 as word3 is to ???")

    try:
        word1 = input("Enter word1: ").strip()
        word2 = input("Enter word2: ").strip()
        word3 = input("Enter word3: ").strip()

        result = vector_arithmetic(
            (word1, -1),
            (word2, 1),
            (word3, 1)
        )

        print(f"\n{word1} : {word2} :: {word3} : ?")
        print_analogy_results(result, [word1, word2, word3])

    except ValueError as e:
        print(f"Error: {e}")
        print("Make sure all words are in the vocabulary!")

# Let users try their own analogies
if __name__ == "__main__":
    explore_custom_analogy()

When you run this code, you’ll be able to:

Try the classic king-queen analogy

See several other interesting analogies

Explore your own analogies interactively

The results demonstrate how Word2Vec captures semantic relationships in its vector space. Here are some interesting patterns you might observe:

Gender relationships (king/queen, actor/actress)

Verb tenses (walk/walked, run/ran)

Country-capital relationships (France/Paris, Germany/Berlin)

Comparative forms (good/better, bad/worse)

Chapter 3. Similarity Search with FAISS

A note for Early Release readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 3rd chapter of the final book. The GitHub repo will be made active later on.

If you have comments about how we might improve the content and/or examples in this book, or if you notice missing material within this chapter, please reach out to the editor at shunter@oreilly.com.

FAISS or Facebook AI Similarity Search, now open-sourced, is a library developed at Meta (formerly Facebook) to efficiently perform similarity search and clustering of dense vectors. Its evolution reflects Meta’s growing need for scalable solutions to handle the massive datasets driving its personalized recommendations, content filtering, and image/video search functionalities.

This chapter is for developers comfortable with mathematical foundations and interested in the internals of vector search. It can be skimmed if your main interest is in building applications with the technology, and not writing your own vector database. If you want to experiment with creating your own vector database then FAISS is a great place to start. It is both a great similarity search engine out of the box and a powerful, flexible toolkit for building your customised vector database and similarity search engine.

It can use GPU but in this book we will only use the CPU version to make the code accessible to the widest audience.

The core motivation behind FAISS stemmed from the limitations of existing solutions for handling the sheer volume and dimensionality of data generated by Meta’s platforms. Traditional approaches like brute-force search, while conceptually simple, become computationally infeasible with millions or billions of vectors. The challenge wasn’t just speed; it was also resource efficiency. Meta needed a solution that could scale gracefully across multiple machines while minimizing memory footprint and power consumption.

FAISS addresses this challenge through a combination of indexing and search algorithms. Instead of comparing a query vector to every vector in the dataset (brute-force), FAISS leverages various indexing structures to drastically reduce the search space. These structures intelligently organize the vectors, allowing FAISS to quickly eliminate irrelevant candidates and focus on a smaller subset most likely to contain the nearest neighbors.

Several key indexing techniques underpin FAISS’s power, which we’ll also talk about in this chapter. Quantization techniques, such as Product Quantization (PQ) and Scalar Quantization (SQ), reduce the memory footprint by representing vectors with fewer bits. This allows for faster comparisons and facilitates the storage of larger datasets in RAM. Tree-based indexes, such as the IVF (Inverted File system) index, further accelerate the search by partitioning the data into smaller clusters. The query vector is first routed to the relevant cluster(s) based on its characteristics, significantly reducing the number of comparisons needed. Hierarchical Navigable Small World (HNSW) graphs provide another indexing method, offering a balance between search speed and accuracy.

Beyond basic indexing, FAISS offers advanced features to fine-tune performance based on specific application needs. These include:

GPU acceleration: FAISS leverages the parallel processing capabilities of GPUs to significantly speed up both indexing and search operations, particularly beneficial for very large-scale datasets.

Clustering algorithms: FAISS includes various clustering algorithms, enabling efficient grouping of similar vectors, useful for tasks like anomaly detection and data exploration.

Hybrid approaches: FAISS supports combining multiple indexing structures to achieve optimal performance. This allows for tailoring the index to specific data characteristics and query patterns.

The evolution of FAISS is ongoing. Since its initial release, the library has seen continuous improvements, addressing issues like memory management, improving accuracy of approximate nearest neighbor searches, and adding support for new hardware architectures. This continuous development reflects Meta’s commitment to providing a robust and scalable solution for similarity search, constantly evolving to meet the challenges of ever-growing data volumes and increasingly complex applications. The open-source nature of FAISS has also fostered community contributions, further enhancing its capabilities and accessibility to a broader range of researchers and developers. Its impact extends far beyond Meta, empowering diverse applications across various domains that rely on efficient similarity search.

At the heart of FAISS lies three fundamental concepts: vector representations, distance metrics, and index structures. Understanding these core elements is crucial for effectively implementing and optimizing similarity search systems.

This chapter breaks down the essential mathematical concepts behind FAISS. It covers vector spaces, different distance metrics (Euclidean, cosine, inner product), and the basic principles of index structures. It also provides the theoretical foundation needed to understand how FAISS works under the hood.

Vector Representations

In the context of FAISS, vectors are dense arrays of floating-point numbers that represent objects in a high-dimensional space. Each dimension corresponds to a feature or characteristic of the object. For instance, a 768-dimensional vector might represent a sentence processed through a BERT model, where each dimension captures some semantic aspect of the text. In computer vision applications, a 2048-dimensional vector could represent an image processed through a ResNet model, with each dimension encoding visual features like edges, textures, or complex patterns.

Mathematically, we represent a vector -dimensional space as:

x = x 1 x 2 ... x d ∈ R d

Here, each xi represents a coordinate in the i-th dimension. The notation Rd indicates that each coordinate can be any real number, and we’re working in d-dimensional space. This representation is fundamental to how FAISS organizes and searches through data.

The collection of vectors we want to search through forms our dataset X:

X={x1,x2,...,xn} where each xi∈ℝd

Each xi represents one item in our database of n items. For example, in a image search system, each xi would represent one image in our database of n images.

Distance Metrics

Distance metrics quantify the similarity (or dissimilarity) between vectors. The choice of distance metric is crucial as it defines what we mean by “similarity” in our search system. FAISS supports several distance metrics, each with its own mathematical properties and use cases. Let’s dive into the main ones in the sections below.

L2 (Euclidean) Distance

L2, or euclidean distance, is the most commonly used metric. It is computed as:

L2(x,y)=∑i=1d(xi−yi)2

This formula calculates the straight-line distance between two points in our d-dimensional space. It’s intuitive and works well for many applications, but it can be sensitive to the scale of different dimensions. For example, if one dimension represents age (0-100) and another represents height in centimeters (150-200), the age dimension might dominate the distance calculations unless the features are properly normalized.

Inner Product

Inner product distance is used when vectors are normalized, equivalent to cosine similarity in the following sense. The inner product measures the degree of alignment between two vectors. When vectors are normalized to unit length, the inner product becomes equivalent to cosine similarity.

It is computed as:

IP(x,y)=∑i=1dxiyi

The inner product measures how aligned two vectors are. When vectors are normalized (have length 1), larger inner products indicate greater similarity. This metric is particularly useful in natural language processing where we care more about the direction of vectors than their magnitude.

Cosine Similarity

Cosine similarity is related to the inner product but is normalized by vector magnitudes. It is computed as:

cos(x,y)=∑i=1dxiyi∑i=1dxi2∑i=1dyi2

Cosine similarity measures the cosine of the angle between two vectors, effectively normalizing for vector length. This is particularly useful when we want to consider two vectors similar if they point in the same direction, regardless of their magnitudes. For example, in text similarity, two documents might be semantically similar even if one is much longer than the other.

FAISS Indexes

FAISS offers a rich variety of index types, each designed for different use cases and performance trade-offs. They can be broadly categorized based on their underlying techniques. Let’s break down some of the most important FAISS index types in the sections below.

Flat Indexes (Brute-Force)

Flat indexes search over the entire space and do not do intelligent optimization of work, hence “brute force”. They check for the distance between a given vector and every vector in the entire database and return the best results usually sorted by closeness.

IndexFlatL2: This is the most basic index type, performing an exhaustive search using the L2 (Euclidean) distance. It compares the query vector to every vector in the dataset. While providing exact results, it’s computationally expensive and not suitable for very large datasets.

IndexFlatIP: Similar to IndexFlatL2, but uses the inner product (dot product) as the similarity metric. Useful for normalized vectors.

IndexFlat: A generic flat index that supports various distance metrics specified at search time.

Pros:

Exact results (guaranteed to find the true nearest neighbors).

Simple to understand and use.

Cons:

Slow for large datasets.

High memory consumption as it stores all vectors in their original form.

Inverted File Index (IVF) based Indexes

These indexes partition the data into clusters and use an inverted list to store the mapping between cluster centroids and the data points assigned to each cluster. They provide approximate results but are significantly faster than flat indexes.

IndexIVFFlat: Combines IVF with a flat index for storing vectors within each cluster. This means the vectors within a cluster are searched exhaustively.

IndexIVFPQ: One of the most popular index types. It uses IVF for coarse quantization and Product Quantization (PQ) to compress vectors within each cluster. Offers a good balance between speed, accuracy, and memory usage.

IndexIVFPQFastScan: An implementation of IndexIVFPQ optimized for speed over memory usage. It is significantly faster than IndexIVFPQ by using a re-implementation of product quantization that is more cache-friendly.

IndexIVFScalarQuantizer: Similar to IndexIVFPQ but uses Scalar Quantization (SQ) instead of PQ for compressing vectors within clusters. SQ quantizes each dimension of the vector independently.

IndexIVFPQR: An extension of IndexIVFPQ that adds a refinement step after the initial PQ search, using a small portion of the original vectors to improve accuracy.

Pros:

Much faster than flat indexes for large datasets.

Reduced memory footprint compared to flat indexes, especially with PQ.

Cons:

Approximate results (may not always find the true nearest neighbors).

Requires training to determine cluster centroids and quantization parameters.

Locality Sensitive Hashing (LSH) based Indexes

IndexLSH: Implements LSH, which uses random projections to hash similar vectors to the same buckets.

Pros:

Relatively fast for datasets with certain properties.

Cons:

Performance can be sensitive to the choice of parameters.

Generally less accurate than IVF-based methods for high-dimensional data.

Hierarchical Navigable Small World (HNSW) based Indexes

The HNSW index is one of the most powerful indexes, especially for human language which has a natural structure of clustered concepts arranged hierarchically from broader to more refined clusters of categories and concepts.

The HNSW index looks for clusters of concepts nearest to the search concept and then does a mode detailed search inside that cluster - hence it avoids searching large parts of the database by making intelligent use of the nested structure of language.

IndexHNSWFlat: Combines HNSW with a flat index for storing vectors.

IndexHNSWPQ: Combines HNSW with PQ for vector compression.

IndexHNSWSQ: Combines HNSW with Scalar Quantization.

Pros:

State-of-the-art search speed, especially for high recall scenarios.

Excellent performance on high-dimensional data.

Cons:

Higher memory usage compared to IVF-based indexes.

Can be slower to build the index.

Other Specialized Indexes

IndexScalarQuantizer: A simple index that just performs scalar quantization without any clustering or graph structure.

IndexPQ: Just performs product quantization without clustering.

IndexBinaryFlat: A flat index for binary vectors (vectors with only 0/1 values), using Hamming distance.

IndexBinaryIVF: An IVF index for binary vectors.

Composite and Transformative Indexes

FAISS allows for creating complex indexes by composing and transforming simpler ones. Some common examples include:

IndexPreTransform: Applies a linear transformation (e.g., PCA) to the vectors before indexing them with another index.

IndexIDMap: Adds an ID mapping on top of another index, allowing for efficient retrieval of original IDs.

IndexRefineFlat: Refines the results from another index using a flat index on a subset of the data.

IndexShards: Shards an index over multiple CPU or GPU instances.

Choosing the right index

The optimal index type depends on several factors:

Dataset size: For small datasets, flat indexes might be sufficient. For larger datasets, IVF or HNSW-based indexes are preferred.

Dimensionality: HNSW often performs best for high-dimensional data.

Desired accuracy: If exact results are crucial, use flat indexes. For approximate results, IVF or HNSW are suitable.

Memory constraints: PQ-based indexes offer significant memory savings.

Speed requirements: HNSW generally provides the fastest search speeds, followed by IVF.

Hardware: GPU acceleration can significantly improve performance for many index types.

Understanding these different index types and their trade-offs is essential for effectively utilizing FAISS to build efficient and accurate similarity search systems. Experimentation and careful consideration of your specific requirements are key to finding the best solution for your application.

The choice between Flat and IVF indexes fundamentally comes down to the balance between accuracy and speed. Flat indexes are ideal for smaller datasets where exact results are required, while IVF indexes excel in large-scale applications where approximate results are acceptable. Understanding these trade-offs is crucial for building efficient similarity search systems. The implementation details and best practices provided in this section serve as a foundation for more advanced FAISS features and optimizations. See Table 3-1-X for _____

Table 3-1. -X: FAISS Indexes and their properties (From the Meta FAISS github repository)

Method
Class name
index_factory
Main parameters
Bytes/vector
Exhaustive
Comments

Exact Search for L2

IndexFlatL2

“Flat”

d

4*d

yes

brute-force

Exact Search for Inner Product

IndexFlatIP

“Flat”

d

4*d

yes

also for cosine (normalize vectors beforehand)

Hierarchical Navigable Small World graph exploration

IndexHNSWFlat

“HNSW,Flat”

d, M

4*d + x * M * 2 * 4

no

Inverted file with exact post-verification

IndexIVFFlat

“IVFx,Flat”

quantizer, d, nlists, metric

4*d + 8

no

Takes another index to assign vectors to inverted lists. The 8 additional bytes are the vector id that needs to be stored.

Locality-Sensitive Hashing (binary flat index)

IndexLSH

-

d, nbits

ceil(nbits/8)

yes

optimized by using random rotation instead of random projections

Scalar quantizer (SQ) in flat mode

IndexScalarQuantizer

“SQ8”

d

d

yes

4 and 6 bits per component are also implemented.

Product quantizer (PQ) in flat mode

IndexPQ

“PQx”, “PQ"M"x"nbits

d, M, nbits

ceil(M * nbits / 8)

yes

IVF and scalar quantizer

IndexIVFScalarQuantizer

“IVFx,SQ4” “IVFx,SQ8”

quantizer, d, nlists, qtype

SQfp16: 2 * d + 8, SQ8: d + 8 or SQ4: d/2 + 8

no

Same as the IndexScalarQuantizer

IVFADC (coarse quantizer+PQ on residuals)

IndexIVFPQ

“IVFx,PQ"y"x"nbits

quantizer, d, nlists, M, nbits

ceil(M * nbits/8)+8

no

IVFADC+R (same as IVFADC with re-ranking based on codes)

IndexIVFPQR

“IVFx,PQy+z”

quantizer, d, nlists, M, nbits, M_refine, nbits_refine

M+M_refine+8

no

From Meta FAISS github repository

The choice between Flat and IVF indexes fundamentally comes down to the balance between accuracy and speed. Flat indexes are ideal for smaller datasets where exact results are required, while IVF indexes excel in large-scale applications where approximate results are acceptable. Understanding these trade-offs is crucial for building efficient similarity search systems. The implementation details and best practices provided in this section serve as a foundation for more advanced FAISS features and optimizations.

Quantization

Quantization is a technique used in many areas where data is to be represented or searched with efficient trade-offs between space, time and accuracy. In quantizing a value it is reduced in resolution (e.g. number of bits used in representation) which saves space and speeds up search due to fewer items being searched, but also reduces the accuracy of the result. Based on different application areas and the specific depth of quantization, i.e. how coarser we are going to make the data, the loss of accuracy is considered and acceptable trade off against speed and space savings. FAISS provides the following quantizations

Scalar Quantization (SQ)

FAISS supports scalar quantization as one of its indexing methods for reducing memory usage and improving search speed.

Scalar quantization is a dimensionality reduction technique that compresses vectors by representing each component of the vector with a smaller number of bits. Instead of using, say, 32-bit floats, it might use 8-bit integers (or even fewer bits). This is done by:

Defining Codebooks: For each dimension (or a group of dimensions) of the vector, a codebook is created. This codebook contains a limited number of representative values (centroids).

Encoding: Each component of a vector is replaced by the index of the closest centroid in the corresponding codebook. This index typically requires fewer bits to represent than the original value.

Decoding: To reconstruct the approximate vector, the indices are used to look up the corresponding centroids in the codebooks.

FAISS provides several implementations of scalar quantization, often used in combination with other techniques like Product Quantization (PQ) or Inverted File Indexing (IVF). Here are some key aspects:

IndexScalarQuantizer: This is the basic index for performing scalar quantization in FAISS. It allows you to specify the number of bits per component (e.g., 8 bits for QT_8bit).

Quantization Types: FAISS offers different types of scalar quantization:

QT_8bit: Each component is encoded with 8 bits (range 0-255).

QT_4bit: Each component is encoded with 4 bits (range 0-15).

QT_fp16: Components are stored as 16-bit floating-point numbers. This is less aggressive than QT_8bit and QT_4bit but still reduces memory usage compared to 32-bit floats.

QT_8bit_uniform and QT_4bit_uniform: These use a uniform distribution for the codebooks, making them faster for training than the non-uniform variants.

QT_8bit_direct and QT_6bit_direct: They map components directly to the bit representation, removing the need for a codebook lookup, but potentially impacting accuracy.

Combination with other Indexes: SQ is often used as a pre-processing step before applying more complex indexing techniques. For instance:

IndexIVFScalarQuantizer: Uses an inverted file system where the centroids are determined by scalar quantization. This is a good choice for very large datasets.

IndexPQ followed by IndexScalarQuantizer: This combines Product Quantization with SQ, first using PQ to reduce the dimensionality and then applying SQ to compress the resulting codes further.

Benefits of Scalar Quantization in FAISS

Let’s review some of the benefits of scalar quantization:

Reduced Memory Usage: SQ significantly reduces the memory footprint of vector databases by representing each vector with fewer bits.

Faster Search Speed: Smaller vectors mean faster computations, especially when calculating distances during search operations.

Simplified Implementations: SQ is relatively straightforward to implement and can be faster to train than other more complex quantization methods.

Trade-offs of Scalar Quantization

While SQ is helpful in many ways, there are some tradeoffs:

Accuracy Loss: Quantization inevitably introduces some loss of information, which can affect the accuracy of search results. The fewer bits used, the higher the potential loss of accuracy.

Choice of Codebook: The quality of the codebook has a significant impact on the accuracy. A well-chosen codebook that accurately represents the distribution of the data is crucial.

Product Quantization (PQ)

The other quantization technique used in FAISS is called Product Quantization (PQ). It takes Scalar Quantization to the next level. In PQ a vector space is broken down into subspaces. Concretely, this means that a given vector is partitioned into segments and each segment is quantized separately with its own code book. Note that this same segmentation is used across all the vectors in the collection. Each segment is quantized differently with a different quantization scheme if necessary.

The Problem:

Imagine you have a vast database of vectors (like image embeddings, text embeddings, etc.) with hundreds or even thousands of dimensions. Searching for the most similar vectors directly is computationally expensive. Storing these vectors in memory can also become prohibitive. Product Quantization to the rescue!

Product quantization tackles these challenges by:

Decomposition: It breaks down the high-dimensional vector space into multiple lower-dimensional subspaces.

Quantization: It quantizes the vectors within each subspace using a codebook.

Indexing: It creates a compressed representation of each original vector using the codebook indices.

PQ Walkthrough

Here’s a step-by-step explanation of how PQ works in FAISS:

Subspace Partitioning:

Let’s say you have a D-dimensional vector (e.g., 128 dimensions).

PQ divides this vector into m disjoint subvectors of dimension D/m (assuming D is divisible by m). For example, if m=4, and D=128, each subvector would have 128/4 = 32 dimensions.

Each subvector now represents a subspace of the original space.

Codebook Learning (For Each Subspace):

Now for each of the m subspaces, FAISS learns a codebook.

A codebook is a set of k representative vectors (centroids) for the subspace. These are often obtained using k-means clustering.

Let’s say k=256 for each subspace.

So, for each subspace, k-means is run to identify k clusters, and the centroids of these clusters form the codebook. Which means that for any vector in a cluster the centroid vector will be used instead as a representative.

Vector Quantization (Encoding):

Now, for each original vector in your database:

It’s divided into m subvectors.

For each subvector, FAISS finds the closest centroid (codebook entry) in the corresponding subspace’s codebook.

The index of this closest centroid is stored.

So, for each of the m subvectors, you end up with an integer index between 0 and k-1.

Compressed Representation:

The original high-dimensional vector is now represented by a sequence of m integer indices. These are much smaller than the original floating-point values, resulting in significant storage compression.

This compressed representation is called a “PQ code.”

Similarity Search (Approximate):

Query Vector Encoding: When a query vector comes in, it undergoes the same process: it’s broken into subvectors, and each is mapped to a centroid index in the corresponding codebook.

Distance Calculation: When calculating the approximate distance between the query vector and the database vectors, the distance is computed between their compressed representations. This is usually done using a “lookup table.” This table pre-computes the distances between all codebook entries. It’s a crucial optimization to avoid calculating the distance between high-dimensional vectors every time.

Approximate Nearest Neighbor Search: Because we’re comparing quantized vectors, the search is approximate. However, in practice, it offers an excellent trade-off between speed and accuracy, especially for very large datasets.

FAISS’s Implementation Details

Highly Optimized: FAISS is known for its highly optimized C++ implementation, which makes it incredibly fast. It leverages SIMD instructions and other techniques for efficiency.

Multiple PQ Variants: FAISS offers variations of PQ, such as:

OPQ (Optimized Product Quantization): Introduces a rotation step before PQ, which can improve accuracy.

IVFPQ (Inverted File Product Quantization): Combines PQ with an inverted file for further speedups by filtering out many vectors before applying PQ distance calculations.

GPU Support: FAISS can leverage GPUs for even faster computations, making it suitable for large-scale applications.

Key Advantages of PQ in FAISS

Let’s review some of the advantages of using PQ in FAISS:

Compression: Reduces the memory footprint of the database significantly.

Speed: The compressed representation allows for much faster distance calculations.

Scalability: Can handle massive datasets with millions or billions of vectors.

Trade-Off: Allows you to balance between accuracy and speed depending on the application requirements.

When is PQ a good choice?

You may be wondering when to use PQ, especially when considering SQ as well. Let’s look at the scenarios that are best for PQ:

High Dimensional Data: Very effective for vectors with many dimensions.

Large Datasets: Especially beneficial when you have a massive number of vectors to index and search.

Approximate Nearest Neighbor Search: Suitable when you’re okay with a slight trade-off in accuracy for speed.

FAISS’s Product Quantization method is a powerful technique for making nearest neighbor search scalable and efficient. By decomposing vectors, quantizing subspaces, and using a lookup table for distance calculations, it provides a significant boost in performance while reducing storage requirements. It’s a key reason why FAISS is a widely used library for similarity search in AI and machine learning.

The Approximate Nearest Neighbor Problem (ANN)

The fundamental problem FAISS addresses is the Approximate Nearest Neighbor (ANN) search. What this achieves is accepting an approximate answer in exchange for speed of response. Most of the time this is a completely acceptable trade-off, much like we might settle for a not-perfectly-prepared version of our favorite sandwich at lunchtime, because the perfect one is a half-hour drive away. A little saving in time adds up to a lot when you repeat the operation hundreds or thousands of times, in this case of hundreds of thousands of documents.

ANN search is a core functionality of FAISS, designed to efficiently find data points in a large dataset that are most similar to a given query point. Unlike exact nearest neighbor search, which guarantees finding the absolute closest points, ANN search trades off a small amount of accuracy for a significant gain in speed. This trade-off makes it practical for massive datasets where exact search would be computationally prohibitive.

Now let’s get into the math a little bit:

Given a query vector q and a dataset X, the goal is to find:

ANN(q)={x∈X:d(q,x)≤(1+ε)miny∈Xd(q,y)}

This needs some explanation. Let’s break it down:

q is our query vector (a representation of what we’re searching for)

X is our dataset of vectors, the domain of our search

d(q,x) is the distance between our query vector and a candidate vector

miny∈Xd(q,y) represents the minimum distance ie distance to the actual nearest neighbor

∈ is our approximation factor, controlling how much error we’re willing to accept

The expression tells us that we’ll accept any vector x whose distance from q is at most 1+ε times the distance to the true nearest neighbor. For example, if ε=0.2, we’ll accept vectors that are up to 20% farther than the true nearest neighbor.

In the sections that follow, let’s break down how ANN search works in FAISS, covering the key concepts and techniques involved.

The Problem

First, let’s set up our problem we want ANN search to solve.

Given:

A large dataset of vectors (e.g., images, text embeddings, user profiles).

A query vector.

Goal:

Find the k vectors in the dataset that are closest to the query vector, according to a distance metric (e.g., Euclidean distance, cosine similarity).

Avoid computational cost

Exact nearest neighbor search for high-dimensional data in large datasets suffers from the “curse of dimensionality.” The computational cost of comparing the query to every vector in the dataset grows linearly with the dataset size and exponentially with the data dimensionality. This makes it impractical for many real-world applications.

ANN algorithms address this by:

Building an Index: Preprocessing the dataset to create a data structure (the index) that allows for efficient searching.

Approximation: Employing techniques that avoid exhaustive comparisons, resulting in faster search but potentially missing some of the true nearest neighbors.

Key Techniques in FAISS for ANN

FAISS provides a variety of indexing methods that employ different strategies for ANN search. Here are some of the most important ones:

Inverted File Index (IVF):

Clustering: The dataset is partitioned into clusters using algorithms like k-means. Each cluster is represented by a centroid.

Inverted Index: An inverted index is created, mapping each centroid to the list of data points belonging to that cluster.

Search: When a query arrives:

The nprobe closest centroids to the query are identified.

Only the data points within those nprobe clusters are compared to the query.

The k closest points from these comparisons are returned as the approximate nearest neighbors.

nprobe parameter: Controls the trade-off between speed and accuracy. Higher nprobe means more clusters are searched, leading to higher accuracy but slower search.

Product Quantization (PQ):

Compression: Vectors are compressed by dividing them into sub-vectors and quantizing each sub-vector independently.

Search: Distance calculations are performed in the compressed space, which is much faster. FAISS uses asymmetric distance computation (ADC), where the query is not quantized, for better accuracy.

Often combined with IVF (e.g., IndexIVFPQ): IVF is used for coarse-grained partitioning, and PQ is used for fine-grained distance computation within each cluster.

See [PQ section] for more info on PQ

Scalar Quantization (SQ):

Compression: Each component of a vector is quantized independently.

Search: Similar to PQ, distances are calculated in the compressed space.

Often used in conjunction with IVF (e.g., IndexIVFScalarQuantizer).

See [SQ section] for more info on SQ

Locality Sensitive Hashing (LSH):

Hashing: Vectors are hashed into buckets using hash functions designed such that similar vectors are likely to fall into the same bucket.

Search: The query is hashed, and only vectors in the same or nearby buckets are considered as potential nearest neighbors.

Less commonly used in FAISS compared to IVF and PQ, especially for high-dimensional data.

Hierarchical Navigable Small World (HNSW):

Graph-based: Builds a multi-layered graph where each layer is a proximity graph of the data.

Search: Starts at a high layer with long-range connections and greedily navigates towards the query, progressively refining the search at lower layers with shorter-range connections.

Generally offers excellent search speed and accuracy but can have higher memory usage and longer indexing times.

Choosing an Index in FAISS

The choice of index in FAISS depends on factors such as:

Dataset size and dimensionality: IndexFlatL2 (exact search) is only feasible for small datasets. IVF-based indexes are good for medium to very large datasets. HNSW is suitable for large datasets where high accuracy and speed are crucial.

Memory constraints: PQ and SQ offer memory compression.

Accuracy requirements: HNSW typically provides the highest accuracy, followed by IVF with PQ/SQ.

Search speed requirements: HNSW often provides the fastest search, followed by IVF with PQ/SQ.

Indexing time: IndexFlatL2 has no indexing time. IVF-based indexes have moderate indexing time. HNSW can have longer indexing times.

Code example

Let’s take a look at the code example below that uses IndexIVFFlat (IVF with Flat index) for exhaustive search within the cluster:

import faiss
import numpy as np
# Sample data
d = 64  # Dimension of vectors
nb = 100000  # Number of vectors in the dataset
nq = 1000  # Number of query vectors
np.random.seed(1234)
xb = np.random.random((nb, d)).astype('float32')
xq = np.random.random((nq, d)).astype('float32')
# Index parameters
nlist = 100  # Number of clusters
k = 4  # Number of nearest neighbors to retrieve
quantizer = faiss.IndexFlatL2(d)  # The index used for clustering (can also be an ANN index)
index = faiss.IndexIVFFlat(quantizer, d, nlist)

# Train the index
index.train(xb)
# Add vectors to the index
index.add(xb)

# Search parameters
index.nprobe = 10  # Number of clusters to search over at query time (trade-off: higher = more accurate but slower)

# Perform the search
D, I = index.search(xq, k)  # D: distances, I: indices

# Print the results
print(I[:5])  # Indices of the 4 nearest neighbors for the first 5 queries
print(D[:5])  # Distances of the 4 nearest neighbors for the first 5 queries

[ the code above could perhaps use more detailed explanation in addition to the comments ]

ANN search in FAISS is a powerful technique for efficiently finding approximate nearest neighbors in large datasets. By employing indexing methods like IVF, PQ, SQ, and HNSW, FAISS enables fast similarity search with a controllable trade-off between speed and accuracy, making it suitable for a wide range of applications.

Understanding HNSW (Hierarchical Navigable Small World) Indexes

This section is a deep dive into one of FAISS’s most powerful indexing methods. This section explains how HNSW constructs navigable small world graphs, their complexity, and why they often provide the best trade-off between speed and accuracy.

What is HNSW?

HNSW is like a multi-level highway system for navigating through your data points. Imagine you’re driving from one city to another:

Highways connect major cities (top level)

Major roads connect smaller towns (middle level)

Local streets connect individual addresses (bottom level)

This is exactly how HNSW works with your vectors!

The HNSW index is a graph structure that encodes the similarity of vectors in the embedding into a “proximity graph” structure. During the construction of the index, vectors (nodes) that are semantically similar to a given vector (node) have a much smaller distance to it than semantically dissimilar vectors (nodes) - the graph structure then represents the semantic clustering of the corpus, which concepts are close to others (clustered) and which are not. This is now analogous to the highway map with roads as edges in the graph and cities of various sizes as nodes.

Now during search, given a query vector and starting at any point on the highway map the algorithm finds the fastest way to get to the destination (closest city) by rapidly traversing the highway system picking the shortest road segment in the direction of the goal - we may end up in a suburb (approximate search) but that is goog enough for most purposes.

Building the Index

Here we are going to generate some sample data (this takes place of actual data which would be a large corpus of text) then we create the hierarchical index by looking for clusters (these are called “small worlds” and arranging them hierarchically in the index (the indexing code in FAISS does all this for us)

The building process works like this:
1. Each new vector is inserted starting from the top layer
2. At each layer, it connects to M nearest neighbors
3. The vector may appear in multiple layers (like a city that has both highways and local roads)
import numpy as np
import faiss

# Create an HNSW index
dimension = 128
M = 16  # Number of connections per layer
ef_construction = 40  # Number of candidates to consider during construction

index = faiss.IndexHNSWFlat(dimension, M)
index.hnsw.efConstruction = ef_construction

# Generate sample data
num_vectors = 10000
vectors = np.random.random((num_vectors, dimension)).astype('float32')

# Add vectors to the index
index.add(vectors)

Searching the Index

The search process is like navigation:

Start at entry point in top layer

Find best direction to move

Move down a layer when no better horizontal moves exist

Repeat until reaching bottom layer

# Set search parameters
ef_search = 16  # Number of candidates to consider during search
index.hnsw.efSearch = ef_search

# Perform a search
k = 5  # Number of nearest neighbors to find
query = np.random.random((1, dimension)).astype('float32')
distances, indices = index.search(query, k)

Key Parameters Explained

Let’s review some of the key parameters for HNSW in the python code .

Think of M as the number of roads connecting each city:

- Larger M = more connections = better accuracy but more memory

- Smaller M = fewer connections = less memory but potentially less accurate

# More connections = better accuracy but more memory

M = 16  # Default value
index = faiss.IndexHNSWFlat(dimension, M)

efConstruction (Construction Factor)

ef_construction = 40  # Number of candidates to consider during construction

efConstruction controls how thoroughly the index is built:

- Higher values = better index quality but slower construction

- Lower values = faster construction but potentially less optimal connections

efSearch (Search Factor)

# Set search parameters
ef_search = 16  # Number of candidates to consider during search
index.hnsw.efSearch = ef_search

efSearch controls how thoroughly the search is performed:

- Higher values = more accurate results but slower search

- Lower values = faster search but might miss some good matches

Practical Example: Building a Similarity Search System

Here’s a complete example showing how to use HNSW for a real-world use case:

import numpy as np
import faiss
from typing import List, Tuple

class SimilaritySearcher:
    def __init__(self, dimension: int, M: int = 16):
        """Initialize HNSW index with given parameters."""
        self.dimension = dimension
        self.index = faiss.IndexHNSWFlat(dimension, M)

        # Set construction parameters
        self.index.hnsw.efConstruction = 40
        self.index.hnsw.efSearch = 16

        # Store original items for reference
        self.items = []

    def add_items(self, vectors: np.ndarray, items: List[str]):
        """Add vectors and their corresponding items to the index."""
        assert vectors.shape[1] == self.dimension
        assert vectors.shape[0] == len(items)

        # Add to FAISS index
        self.index.add(vectors.astype('float32'))

        # Store original items
        self.items.extend(items)

    def search(self, query_vector: np.ndarray, k: int = 5) -> List[Tuple[str, float]]:
        """Search for k most similar items."""
        # Ensure query vector has correct shape
        query_vector = query_vector.reshape(1, self.dimension)

        # Perform search
        distances, indices = self.index.search(
            query_vector.astype('float32'), k)

        # Return items and distances
        results = []
        for idx, dist in zip(indices[0], distances[0]):
            if idx != -1:  # FAISS returns -1 for invalid results
                results.append((self.items[idx], dist))

        return results

# Example usage
dimension = 128
searcher = SimilaritySearcher(dimension)

# Create sample data
num_items = 10000
vectors = np.random.random((num_items, dimension))
items = [f"item_{i}" for i in range(num_items)]

# Add to index
searcher.add_items(vectors, items)

# Perform a search
query = np.random.random(dimension)
results = searcher.search(query, k=5)

# Print results
for item, distance in results:
    print(f"{item}: distance = {distance:.4f}")

Performance Characteristics

Here’s a simple benchmark comparing HNSW with other index types:

def benchmark_indexes(dimension=128, num_vectors=100000, k=5):
    vectors = np.random.random((num_vectors, dimension)).astype('float32')
    query = np.random.random((1, dimension)).astype('float32')

    # Create indexes
    flat_index = faiss.IndexFlatL2(dimension)
    ivf_index = faiss.IndexIVFFlat(
        faiss.IndexFlatL2(dimension), dimension, int(np.sqrt(num_vectors)))
    hnsw_index = faiss.IndexHNSWFlat(dimension, 16)

    # Train and add vectors
    flat_index.add(vectors)

    ivf_index.train(vectors)
    ivf_index.add(vectors)

    hnsw_index.add(vectors)

    # Benchmark search
    results = {}
    for name, index in [
        ('Flat', flat_index),
        ('IVF', ivf_index),
        ('HNSW', hnsw_index)
    ]:
        start = time.time()
        for _ in range(100):
            index.search(query, k)
        results[name] = (time.time() - start) / 100

    return results

Best Practices

Now that we’ve walked through HNSW, let’s review some best practices to keep in mind when you’re using HNSW.

Parameter Selection:

Start with default M=16

Increase efConstruction for better index quality

Tune efSearch based on accuracy vs. speed needs

Memory Considerations:

HNSW uses more memory than simple indexes

Memory usage grows with M and number of vectors

Consider using smaller M for very large datasets

Performance Optimization:

Batch similar queries together

Use GPU acceleration if available

Monitor memory usage during index construction

FAISS Architecture and Components

In this section you’ll explore FAISS’s modular design, including its index hierarchy, training requirements, and search methods. I’ll explain how FAISS separates the concerns of data storage, indexing, and search, allowing for flexible combinations of different approaches

The modular architecture of FAISS is built around the concept of composable components (see Figure 3-1-X). Each component (indexes, metrics, storage) is implemented as a separate module that can be combined with others. This allows for great flexibility in creating specialized search solutions.

The Index base class serves as the foundation for all index types in FAISS. It defines the core interface methods like add(), search(), and train(). This consistent interface makes it easier to swap different index implementations without changing the rest of the code.

FAISS provides two ways to combine components: direct programmatic composition and the IndexFactory string syntax. The string syntax is particularly powerful as it allows complex index configurations to be specified in a concise format. We merely make a note of this and do not dive in any deeper. It is a topic more suited for advance usage where one needs a custom index combining other component indexes.

GPU support in FAISS is implemented through dedicated GPU versions of components. This allows for seamless switching between CPU and GPU implementations while maintaining the same interface. In our usage in this book we only use the CPU versions to allow maximum reach - the FAISS GPU usage requires NVIDIA GPUs and CUDA which are not for the faint hearted.

Figure 3-1. -X Module Hierarchy in FAISS code

FAISS is highly extensible at the source code level. Its modular design and well-documented API make it relatively straightforward to extend its functionality. The extensibility of FAISS is achieved through its well-defined base classes. Developers can create custom components by inheriting from these base classes and implementing the required methods.

Developers can:

Implement new index types: FAISS’s architecture allows for the creation and integration of novel indexing structures beyond the existing ones (IVF, HNSW, PQ, etc.). This enables researchers to experiment with and incorporate state-of-the-art indexing techniques.

Add new distance metrics: While FAISS provides several common distance metrics (e.g., Euclidean, inner product), developers can easily add custom distance functions to cater to specific application requirements. This is particularly important when dealing with specialized data representations or similarity measures.

Integrate with other libraries: FAISS’s well-defined interfaces facilitate integration with other machine learning libraries and frameworks. This allows for seamless incorporation into existing workflows and pipelines.

Develop custom search strategies: The search algorithms within FAISS can be customized or extended. One could implement novel search strategies or optimize existing ones for particular data distributions or query patterns.

Extend quantization methods: New quantization methods can be added to further optimize memory usage and search speed for specific data characteristics.

The open-source nature of FAISS, combined with its clear codebase and comprehensive documentation, significantly contributes to its extensibility. The community actively contributes to its development, adding new features and improvements based on diverse needs and applications. However, extending FAISS requires a solid understanding of C++ and the underlying algorithms.

In short, FAISS’s modular architecture provides flexibility and extensibility while maintaining performance. The ability to combine components in different ways allows for customized solutions that balance accuracy, speed, and memory usage. The consistent interfaces and GPU support make it practical to scale solutions from small prototypes to large-scale deployments. Understanding this architecture is crucial for effectively using FAISS in real-world applications.

Okay, let’s go deeper into the architecture and components of FAISS. Understanding these pieces helps to appreciate how FAISS achieves its remarkable efficiency and scalability in similarity search.

Core Concepts of FAISS Architecture

At its heart, FAISS is built around the idea of creating indexes that allow for fast, approximate nearest neighbor searches. It achieves this by:

Indexing: Organizing high-dimensional vectors in a way that enables quick retrieval of similar vectors.

Quantization: Representing vectors with a lower-memory footprint using techniques like Product Quantization (PQ).

Search Optimization: Employing strategies to efficiently navigate the index and find approximate nearest neighbors.

Key Components of FAISS

Let’s explore the major components:

Index Structures:

Index (Abstract Base Class): The fundamental building block. It defines the API for adding vectors and searching for nearest neighbors. FAISS has many concrete Index implementations.

IndexFlat: The simplest index, it stores vectors as they are, without any compression or structure. It’s used for brute-force searches or as a base for more complex indexes. IndexFlatL2 is used for Euclidean (L2) distance. IndexFlatIP for dot product.

IndexIVF (Inverted File Index): An important index that partitions the vector space into Voronoi cells and creates inverted lists of vectors that fall within each cell. This significantly speeds up searching.

IndexPQ: Implements Product Quantization. Vectors are divided into subspaces, and each subspace is quantized using a codebook. This index is used for efficient compressed search.

IndexIVFPQ: Combines an inverted file (IndexIVF) with product quantization (IndexPQ), providing both compression and improved search speed.

IndexHNSW: Uses Hierarchical Navigable Small World graphs. It is suitable for high accuracy but can require more memory than other approaches.

IndexBinary : Similar to IndexFlat, but designed to work with binary vectors. Indexes available include IndexBinaryFlat and IndexBinaryIVF.

More specialized Indexes: FAISS also offers more specific indexes, like IndexLSH (Locality Sensitive Hashing) or IndexIDMap.

Composable Indexes: Indexes can be combined in sequence through the IndexPreTransform class.

Quantization Methods:

ProductQuantizer: A crucial component for PQ. It’s responsible for learning the codebooks (cluster centroids) in the subspaces and for encoding and decoding vectors.

ScalarQuantizer: Quantizes each dimension of the vector independently to save space.

AdditiveQuantizer: A generalization of product quantization that has additional flexibility in how vector components are combined.

Rotation/Transforms FAISS provides a host of transformations that can be used before the vector indexing. Common types of transformations are learned rotations (e.g. for OPQ) and PCA transformations to reduce dimensionality.

Distance Calculation:

FAISS supports various distance metrics, like L2 (Euclidean), Inner Product (dot product), and cosine similarity (for normalized vectors).

Lookup Tables: For PQ, it uses lookup tables to pre-compute distances between codebook centroids, allowing for fast distance calculations on quantized vectors.

Search and Filtering:

search() method: The core of the index, this method takes a query vector and finds the approximate nearest neighbors using the indexing and quantization algorithms.

nprobe: For IndexIVF and IndexIVFPQ, this parameter controls the number of inverted file cells that are searched, it dictates the trade-off between search speed and accuracy.

Filtering: FAISS allows you to optionally add filtering conditions on vectors.

GPU Support:

FAISS offers GPU-enabled indexes for significantly faster search speeds. Most CPU indexes have a direct equivalent that can be used on the GPU.

Index*GPU: GPU versions of indexes are available, e.g. IndexFlatL2GPU, IndexIVFPQGPU.

Preprocessing Transforms:

Transform class: Provides methods to apply linear transformations to the data before indexing, examples include PCAMatrix, Normalizer.

Data Management:

Adding Vectors: FAISS provides methods to add vectors to the index incrementally.

Removal: It can remove vectors from the index.

Serialization: Indexes can be saved to and loaded from files for later use.

Common Workflow with FAISS

Let’s review a common workflow that you would likely use with FAISS:

Choose an Index: Select the appropriate index based on your data and search requirements (e.g., IndexFlat, IndexIVFPQ, IndexHNSW, etc.).

Training (if needed): Some indexes require training on a sample of your data (e.g., the codebook learning step in PQ, the selection of centroids for the IVF index).

Add Vectors: Add the vectors to the index.

Search: Search for approximate nearest neighbors.

Evaluate: Evaluate the performance using metrics like recall and query time.

Tuning: Tune parameters based on the evaluation.

To pull this all together lets take a simple example of a workflow with building and searching a FAISS index using sample generated data. As elsewhere you would replace the sample data with your own data.

Illustrative Example (using Python API)

import faiss
import numpy as np

# Create dummy data
d = 128  # Dimensionality
nb = 10000  # Number of vectors
xq = np.random.rand(10, d).astype('float32') # 10 query vectors
xb = np.random.rand(nb, d).astype('float32') # database

# 1. Choose an index (IVFPQ for example)
nlist = 100  # Number of Voronoi cells (clusters) used to partition the datasets
m = 8        # Number of subvectors for PQ
k = 256     # Codebook size

quantizer = faiss.IndexFlatL2(d)  # Base index for IVF
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, k)

# 2. Train the index on a sample of the data
index.train(xb)

# 3. Add vectors to the index
index.add(xb)

# 4. Search for nearest neighbors
D, I = index.search(xq, 5) # Search for 5 nearest neighbors

# Print results
print("Indices of the nearest neighbors:")
print(I)
print("Distances of the nearest neighbors")
print(D)

Key Takeaways

Let’s review the key FAISS takeaways:

FAISS provides a collection of building blocks (indexes, quantizers, etc.) to handle various similarity search scenarios.

Flexibility: It allows you to compose different index structures and quantization methods.

Performance: The library is highly optimized for speed and can run on CPUs or GPUs.

Customization: You can tune various parameters to balance between speed and accuracy.

Modularity: Each FAISS component is designed with a modular and extensible approach in mind.

Further Exploration

Here are some recommendations if you’d like to explore further:

Documentation: The official FAISS documentation is essential for in-depth information on all components and features.

Examples: The FAISS GitHub repository has numerous examples demonstrating how to use different indexes and features.

Conclusion

FAISS represents a sophisticated approach to solving the challenging problem of similarity search in high-dimensional spaces. The chapter progression has taken the reader from fundamental concepts through increasingly complex implementations, culminating in practical deployment considerations.

This chapter enables readers to gain a comprehensive understanding of both the theoretical foundations and practical applications of FAISS, while providing the necessary depth for implementing efficient similarity search systems in prototyping. The focus on both theoretical understanding and practical implementation makes it valuable for self taught practitioners in the field.

Chapter 4. Semantic Search with Sqlite3

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 4th chapter of the final book. Please note that the GitHub repo will be made active later on.

If you have comments about how we might improve the content and/or examples in this book, or if you notice missing material within this chapter, please reach out to the editor at shunter@oreilly.com.

In this chapter we will use the similarity search technology but at a conceptual level we will be doing “search by meaning”. In other words the way to search by meaning rather than syntax is to use similarity search rather than keyword match. We looked at similarity search in detail in the previous chapter on FAISS which is a full fledged vector store, indexing and search engine. It is usually used as a component of a larger application or database engine. .

Here we use sqlite-vss, the vector extension to sqlite, which actually uses FAISS under the covers as its vector similarity search engine, while integrating vector search with the familiar SQL world.

In our example application we use it for similarity search over vectorized Reddit posts.

We look for other posts similar in meaning and use this to build a personal knowledge management system where the data comes from Reddit. But first we need to understand what sqlite-vss is and how to add it to sqlite3.

1. Understanding the SQLite VSS Extension

The SQLite Vector Similarity Search (VSS) extension is a powerful addition to SQLite that enables efficient storage and similarity search of vector embeddings directly within a SQLite database. This capability transforms SQLite into a capable vector database, perfect for personal knowledge management and semantic search applications.

Installation and Setup

The VSS extension can be compiled from source or installed using pre-built binaries. For most users, the pre-built approach is recommended:

# Download the appropriate binary for your system
wget https://github.com/asg017/sqlite-vss/releases/download/v0.1.0/sqlite-vss-linux-x86_64.tar.gz
tar -xzf sqlite-vss-linux-x86_64.tar.gz

# Load the extension in Python
import sqlite3
conn = sqlite3.connect('reddit_vectors.db')
conn.enable_load_extension(True)
conn.load_extension("./sqlite-vss0")

Core Functionality

The VSS extension provides several key capabilities:

Vector Data Type: A native vector type for storing high-dimensional vectors efficiently

Distance Functions: Built-in functions for calculating vector similarity:

vss_cosine_distance()

vss_l2_distance()

vss_inner_product()

-- Example of creating a table with vector support
CREATE TABLE embeddings (
    id INTEGER PRIMARY KEY,
    content_vector vector(384),  -- Specify dimension in parentheses
    metadata TEXT
);

Vector Operations

The extension provides essential vector operations:

-- Creating vectors
SELECT vss_create_vector(1.0, 2.0, 3.0);  -- Creates a 3D vector

-- Computing similarity between vectors
SELECT vss_cosine_distance(
    vss_create_vector(1.0, 0.0),
    vss_create_vector(0.0, 1.0)
);

-- Normalizing vectors
SELECT vss_normalize(vss_create_vector(1.0, 2.0, 3.0));

Benefits for Reddit Content Storage

The VSS extension is particularly well-suited for storing Reddit content because:

Integrated Storage: Keep embeddings alongside metadata in a single database

Efficient Queries: Combine vector similarity search with traditional SQL filters

Portability: SQLite’s file-based nature makes it perfect for personal knowledge bases

Low Resource Usage: Efficient for personal-scale datasets (thousands to millions of entries)

Limitations and Considerations

While powerful, it’s important to understand the extension’s limitations:

Memory Usage: Vector operations are memory-intensive; consider batch processing for large datasets

Index Size: Vector indexes can grow large; plan storage accordingly

Optimization Needed: Query performance may require tuning for larger datasets

Version Compatibility: Ensure compatibility between SQLite, Python, and the VSS extension versions

The next section will dive deeper into the specific components that make up the VSS extension and how they work together to enable efficient vector search capabilities.

Let’s break down the core components of SQLite VSS:

Distance Metrics: Core mathematical functions for measuring vector similarity

Vector Functions: Utilities for creating, manipulating, and processing vectors

Indexing System: Specialized indexing structures for efficient vector search

Query Optimizer: Components that optimize vector-based queries

Integration Layer: Systems that connect VSS with standard SQLite functionality

Let’s expand on each of these components:

Distance Metrics The extension provides three primary distance metrics:

Cosine Distance: Measures the angle between vectors, ideal for semantic similarity

L2 (Euclidean) Distance: Calculates straight-line distance between vectors

Inner Product: Useful for comparing normalized vectors These metrics are implemented in C for optimal performance and can be called directly in SQL queries.

Vector Functions The VSS extension includes a comprehensive set of vector manipulation functions:

vss_create_vector(): Creates vector objects from arrays of numbers

vss_normalize(): Normalizes vectors to unit length

vss_dimension(): Returns the dimension of a vector These functions provide the building blocks for more complex vector operations.

Indexing System SQLite VSS implements specialized indexing structures:

HNSW (Hierarchical Navigable Small World) graphs for approximate nearest neighbor search

Custom B-tree adaptations for vector data

Optimized index structures for different distance metrics These indexes dramatically improve search performance for large vector datasets.

Query Optimizer The query optimizer includes:

Cost-based optimization for vector operations

Intelligent index selection based on query patterns

Memory-aware execution planning This ensures efficient execution of complex queries combining vector and scalar operations.

Integration Layer The integration layer provides:

Seamless interaction between vector and traditional SQLite operations

Type conversion and validation

Transaction management for vector operations

This allows VSS to work naturally within the broader SQLite ecosystem.

Summary

SQLite VSS is built around a sophisticated architecture that combines efficient vector operations with traditional database functionality. The system balances performance with usability, providing high-level functions while maintaining low-level optimization. This makes it particularly suitable for personal-scale vector search applications, where the full complexity of a dedicated vector database might be unnecessary. The combination of well-implemented distance metrics, efficient indexing, and tight integration with SQLite’s core functionality makes VSS a powerful tool for developers working with vector embeddings in personal or small-scale applications. While it may not match the performance of dedicated vector databases for massive datasets, it provides an excellent balance of features and simplicity for many common use cases.

contd re query optimizer

SQLite3 does not have a traditional query optimizer like those found in databases such as PostgreSQL or MySQL. Here are the key points about SQLite’s query handling:

SQLite uses a much simpler query planning system

It relies primarily on heuristics rather than cost-based optimization

The query planner makes decisions mainly based on available indices

Statistics about the data are not maintained or used for planning The planning process cannot be influenced through hints or configuration Let’s explore each of these aspects:

Simple Query Planning SQLite uses a straightforward algorithm to determine how to execute queries. It primarily looks at available indices and the WHERE clause structure to determine the best way to access the data. This is much simpler than the complex cost-based optimizers in larger database systems.

Heuristic-Based Approach Rather than calculating costs based on table statistics and data distribution, SQLite follows a set of predefined rules. For example, it will generally prefer using an index if one is available and applicable, without considering whether a full table scan might actually be faster in some cases.

Index-Driven Decisions The main optimization SQLite performs is choosing which indices to use. If multiple indices are available, it will attempt to select the one that it expects to eliminate the most rows early in the query execution process. This decision is based on simple heuristics rather than actual data statistics.

No Statistics Maintenance Unlike more complex databases, SQLite does not maintain statistics about data distribution, table sizes, or index selectivity. This means it cannot make data-driven decisions about query execution strategies.

Fixed Planning Process SQLite’s query planning cannot be influenced through hints or configuration options. The process follows the same rules regardless of the specific data characteristics or query patterns in your application.

Summary

SQLite’s approach to query execution is intentionally simple by design, aligning with its goal of being a lightweight, embedded database. While this means it may not always choose the optimal execution strategy for complex queries, it also means it has lower overhead and is more predictable in its behavior. For many applications, especially those dealing with smaller datasets or simpler queries, this trade-off is entirely appropriate and beneficial.

3. Working with Vector Data Types in SQLite

Understanding vector data representation, storage, and manipulation in SQLite using the VSS extension.

Key Points:

Vector data in SQLite VSS is represented as a specialized vector type with fixed dimensions

Vectors can be created and manipulated using built-in VSS functions Storage of vectors requires careful consideration of dimension size and data types

Vector operations include normalization, distance calculations, and similarity comparisons

Error handling is crucial when working with vector data types Expanding on each point:

Vector Data Representation: The VSS extension implements a custom vector type that must be declared with a specific dimension size. For example: vector(384) for a 384-dimensional vector. This fixed dimensionality helps ensure consistency and optimize storage.

Vector Creation andManipulation: Vectors can be created using the vss_create_vector() function, which takes a series of numbers as arguments. The extension provides functions for basic vector operations like normalization (vss_normalize()), dimension checking (vss_dimension()), and various distance calculations.

Storage Considerations: When storing vectors, it’s important to consider:

The fixed dimension size must be declared at table creation

Vectors are stored efficiently in a binary format

Each vector column has overhead proportional to its dimension

Indexes on vector columns require additional storage space

Vector Operations: The extension supports several key operations:

Cosine similarity calculations

L2 (Euclidean) distance measurements

Inner product computations

Vector normalization These operations are optimized for performance and can be used directly in SQL queries.

Error Handling: Working with vector data types requires careful error handling for:

Dimension mismatches

Invalid vector data

Numerical overflow or underflow

Memory constraints with large vectors

Summary

SQLite VSS provides a robust system for working with vector data types through a combination of specialized storage formats and optimized operations. The system balances efficiency with usability by providing high-level functions while maintaining careful control over data representation and memory usage. This makes it particularly suitable for personal-scale applications where vector similarity search needs to be integrated with traditional database operations. The extension’s approach to vector data types enables developers to work with high-dimensional data in a familiar SQL environment while maintaining the performance characteristics necessary for effective similarity search operations. However, users need to be mindful of dimension sizes, storage requirements, and proper error handling to effectively use these capabilities.

Ch 3 contd.

Based on your preferences, I’ll structure this in a clear, technical manner suitable for advanced undergraduates.

Key Points:

Vector data in SQLite VSS is represented as a fixed-dimensional array of floating-point numbers

Vectors require explicit dimension specification during table creation

Three primary vector operations: creation, manipulation, and comparison

Type conversion handling between Python numpy arrays and SQLite vector types

Memory management considerations for vector operations

Let’s expand on each point:

Vector Data Representation The SQLite VSS extension represents vectors as fixed-length arrays of 32-bit floating-point numbers. Each vector must have a predefined dimension specified at table creation time. For example, if you’re using BERT embeddings with 768 dimensions, you would specify this in your table schema:

sql
Copy
CREATE TABLE embeddings (
    id INTEGER PRIMARY KEY,
    text_vector vector(768),
    content TEXT );

Dimension Specification When creating tables with vector columns, dimensions must be explicitly declared. This is crucial for: * Memory allocation optimization Index structure efficiency * Query performance optimization The dimension is immutable after table creation, ensuring data consistency. Vector Operations The three primary operations categories are:

sql
Copy
-- Vector Creation
INSERT INTO embeddings (text_vector)
VALUES (vss_create_vector(0.1, 0.2, ..., 0.768));

-- Vector Manipulation
SELECT vss_normalize(text_vector) FROM embeddings;

-- Vector Comparison
SELECT * FROM embeddings
WHERE vss_cosine_distance(text vector, :query_vector) < 0.1;

Type Conversion When working with Python, vectors typically need conversion between numpy arrays and SQLite vector format:

import numpy as np
import sqlite3

def numpy_to_vss(array):
    return f"vss_create_vector({','.join(map(str, array))})"

def insert_vector(conn, vector_array):
    sql = "INSERT INTO embeddings (text_vector) VALUES (?)"
    conn.execute(sql, (numpy_to_vss(vector_array),))

Memory Management Vector operations can be memory-intensive, particularly when: * Working with high-dimensional vectors * Performing batch operations * Building indexes * Executing similarity searches

Summary

Working with vector data types in SQLite VSS requires understanding both the fundamental data representation and the practical considerations of vector operations. The system provides a balance between flexibility and performance, offering direct vector manipulation while maintaining the familiar SQLite interface. Proper handling of vector data types, including dimension specification and type conversion, is crucial for building efficient vector-based applications. The key to successful implementation lies in understanding the constraints and capabilities of the vector data type, particularly regarding dimensionality and memory management. This foundation enables the development of robust vector search applications while maintaining the simplicity and portability that makes SQLite attractive for personal and small-scale deployments.

section 4 Schema design for vector search.

Key Points:

Efficient schema design requires balancing vector storage with metadata organization

Multiple index types are needed for hybrid search capabilities

Denormalization may be beneficial for vector search performance

Schema should support both exact and approximate nearest neighbor queries

The schema must account for vector versioning and embedding model changes

Let’s expand on each point:

Schema Balance and Organization A well-designed schema for Reddit content might look like this:

-- Main posts table with vector embeddings
CREATE TABLE posts (
    id INTEGER PRIMARY KEY,
    post_id TEXT UNIQUE,
    title TEXT,
    content TEXT,
    subreddit TEXT,
    created_utc INTEGER,
    author TEXT,
    content_vector vector(384),  -- For sentence embeddings
    title_vector vector(384)     -- Separate vector for title
);

-- Supporting metadata table
CREATE TABLE post_metadata (
    post_id TEXT PRIMARY KEY,
    score INTEGER,
    num_comments INTEGER,
    is_original_content BOOLEAN,
    FOREIGN KEY (post_id) REFERENCES posts(post_id)
);

-- Vector search index
CREATE VIRTUAL TABLE vss_index USING vss0(
    content_vector(384),
    id INTEGER
);

Multiple Index Types The schema should support three types of indexes:

-- Traditional B-tree indexes for metadata
CREATE INDEX idx_posts_subreddit ON posts(subreddit);
CREATE INDEX idx_posts_created ON posts(created_utc);

-- Full-text search for content
CREATE VIRTUAL TABLE posts_fts USING fts5(
    title, content, content='posts', content_rowid='id'
);

-- Vector similarity index
CREATE INDEX idx_vector_cosine ON vss_index(content_vector);

Denormalization Strategy For vector search performance, consider denormalizing data:

CREATE TABLE post_vectors (
    id INTEGER PRIMARY KEY,
    post_id TEXT,
    vector_type TEXT,  -- 'title' or 'content'
    vector vector(384),
    metadata JSON,     -- Denormalized metadata for quick access
    FOREIGN KEY (post_id) REFERENCES posts(post_id)
);

Approximate Nearest Neighbor Support Schema design for ANN queries:

-- Configuration table for ANN parameters
CREATE TABLE vector_config (
    model_version TEXT,
    vector_dim INTEGER,
    index_type TEXT,
    distance_metric TEXT,
    ef_construction INTEGER,  -- HNSW index parameter
    m INTEGER                 -- HNSW index parameter
);

-- Index configuration
CREATE INDEX ann_index ON post_vectors
    USING vss0_hnsw(vector)
    WITH (
        dim=384,
        m=16,
        ef_construction=200
    );

Vector Versioning Handle embedding model changes:

CREATE TABLE embedding_versions (
    id INTEGER PRIMARY KEY,
    model_name TEXT,
    model_version TEXT,
    vector_dim INTEGER,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);

ALTER TABLE posts ADD COLUMN embedding_version_id INTEGER
    REFERENCES embedding_versions(id);

Summary

An effective schema design for vector search must balance the needs of traditional relational data with the requirements of vector similarity search. The design should incorporate both standard SQL indexing strategies and specialized vector indexing techniques, while maintaining data integrity and query performance.

The key to success lies in carefully considering the access patterns of your application and optimizing the schema accordingly. This might mean denormalizing some data for performance, while ensuring that the schema remains flexible enough to accommodate changes in embedding models and search requirements over time. The schema should support both exact and approximate nearest neighbor searches, allowing for different performance/accuracy tradeoffs depending on the use case.

Section 5 PRAW setup.

Key Points:

PRAW (Python Reddit API Wrapper) is the official Python interface for Reddit’s API

Authentication requires creating a Reddit application and obtaining credentials

PRAW supports different authentication modes: script, web app, and read-only

Rate limiting and API guidelines must be respected

Basic PRAW operations include accessing subreddits, posts, and comments

Let’s expand on each point:

PRAW Overview: PRAW abstracts Reddit’s REST API into a simple, Pythonic interface. It handles all the low-level HTTP requests, rate limiting, and data parsing, allowing developers to focus on their application logic rather than API implementation details.

Authentication Setup:

import praw

reddit = praw.Reddit(
    client_id="YOUR_CLIENT_ID",
    client_secret="YOUR_CLIENT_SECRET",
    user_agent="platform:app_name:v1.0 (by /u/username)",
    username="YOUR_USERNAME",  # Optional
    password="YOUR_PASSWORD"   # Optional
)

The authentication process requires: 1. Creating an app at reddit.com/prefs/apps 2. Obtaining client_id and client_secret 3. Creating an appropriate user_agent string 4. Optionally providing username/password for script applications

Rate Limiting: PRAW automatically handles Reddit’s rate limits (60 requests per minute), but developers should be aware of:

Using appropriate delays between requests

Implementing error handling for rate limit exceptions

Following Reddit’s API terms of service

Basic Operations:

# Access a subreddit
subreddit = reddit.subreddit("python")

# Get hot posts
for submission in subreddit.hot(limit=10):
    print(submission.title)

# Access comments
submission = reddit.submission(id="submission_id")
submission.comments.replace_more(limit=None)
for comment in submission.comments.list():
    print(comment.body)

Summary

PRAW provides a robust, well-documented interface for interacting with Reddit’s API. Its design focuses on ease of use while maintaining respect for Reddit’s rate limits and API guidelines. Understanding PRAW’s authentication system and basic operations is crucial for building Reddit-based applications, particularly when developing vector search systems for Reddit content. The library’s automatic handling of many low-level details makes it an ideal choice for both beginners and experienced developers working with Reddit data.

I’ll expand the Basic Operations section with saved messages, direct messages, and add robust error handling.

import praw
import time
from prawcore.exceptions import PrawcoreException
from requests.exceptions import RequestException
from typing import Generator, Optional

class RedditClient:
    def __init__(self,
                 client_id: str,
                 client_secret: str,
                 username: str,
                 password: str):
        """Initialize Reddit client with error handling and rate limiting."""
        self.reddit = praw.Reddit(
            client_id=client_id,
            client_secret=client_secret,
            user_agent=f"python:content_fetcher:v1.0 (by /u/{username})",
            username=username,
            password=password
        )
        self.rate_limit_delay = 1.0  # Default delay between requests

    def _handle_request(self, func: callable, *args, **kwargs) -> Optional[any]:
        """Generic error handler for Reddit API requests."""
        max_retries = 3
        for attempt in range(max_retries):
            try:
                time.sleep(self.rate_limit_delay)  # Rate limiting
                return func(*args, **kwargs)
            except PrawcoreException as e:
                if "Too Many Requests" in str(e):
                    wait_time = min(pow(2, attempt), 300)  # Exponential backoff
                    print(f"Rate limited. Waiting {wait_time} seconds...")
                    time.sleep(wait_time)
                    self.rate_limit_delay *= 1.5  # Increase delay for future requests
                elif attempt == max_retries - 1:
                    print(f"PRAW error after {max_retries} attempts: {e}")
                    raise
            except RequestException as e:
                if attempt == max_retries - 1:
                    print(f"Network error after {max_retries} attempts: {e}")
                    raise
            except Exception as e:
                print(f"Unexpected error: {e}")
                raise

    def get_saved_content(self, limit: Optional[int] = None) -> Generator:
        """
        Retrieve user's saved posts and comments.

        Args:
            limit: Maximum number of items to retrieve

        Yields:
            Reddit submission or comment objects
        """
        def _fetch_saved():
            return self.reddit.user.me().saved(limit=limit)

        saved_items = self._handle_request(_fetch_saved)
        if saved_items:
            for item in saved_items:
                # Check if item is a submission or comment
                if isinstance(item, praw.models.Submission):
                    yield {
                        'type': 'submission',
                        'id': item.id,
                        'title': item.title,
                        'url': item.url,
                        'subreddit': item.subreddit.display_name,
                        'created_utc': item.created_utc
                    }
                else:  # Comment
                    yield {
                        'type': 'comment',
                        'id': item.id,
                        'body': item.body,
                        'subreddit': item.subreddit.display_name,
                        'submission_title': item.submission.title,
                        'created_utc': item.created_utc
                    }

    def get_direct_messages(self, limit: Optional[int] = None) -> Generator:
        """
        Retrieve user's direct messages.

        Args:
            limit: Maximum number of messages to retrieve

        Yields:
            Dictionary containing message details
        """
        def _fetch_messages():
            return self.reddit.inbox.messages(limit=limit)

        messages = self._handle_request(_fetch_messages)
        if messages:
            for msg in messages:
                yield {
                    'id': msg.id,
                    'subject': msg.subject,
                    'body': msg.body,
                    'author': str(msg.author),
                    'created_utc': msg.created_utc,
                    'was_comment': msg.was_comment
                }

    def mark_message_read(self, message_id: str) -> None:
        """Mark a specific message as read."""
        def _mark_read():
            message = self.reddit.inbox.message(message_id)
            message.mark_read()

        self._handle_request(_mark_read)

# Example usage
def main():
    client = RedditClient(
        client_id="YOUR_CLIENT_ID",
        client_secret="YOUR_CLIENT_SECRET",
        username="YOUR_USERNAME",
        password="YOUR_PASSWORD"
    )

    # Get saved content
    print("Fetching saved content...")
    for item in client.get_saved_content(limit=10):
        print(f"Saved {item['type']}: {item.get('title') or item.get('body')[:50]}")

    # Get direct messages
    print("\nFetching direct messages...")
    for message in client.get_direct_messages(limit=5):
        print(f"Message from {message['author']}: {message['subject']}")

if __name__ == "__main__":
    main()

This expanded code includes:

A robust RedditClient class with error handling and rate limiting

Method for fetching saved content (both posts and comments)

Method for fetching direct messages

Rate limiting with exponential backoff

Error handling for network issues and API limits

Type hints for better code maintainability

Generator patterns to handle large datasets efficiently

The code implements several best practices:

Exponential backoff for rate limiting

Proper exception handling for different types of errors

Clean separation of concerns with modular methods

Clear documentation with docstrings

Type hints for better code maintainability

This implementation should provide a reliable foundation for building Reddit-based applications that need to access saved content and direct messages while respecting Reddit’s API limitations.

Section 6 on content extraction and preprocessing for Reddit data.

Reddit content structure consists of posts, comments, and associated metadata requiring careful extraction Text preprocessing requires specific handling for Reddit’s markdown, URLs, and special characters Metadata organization needs a structured approach for both posts and comments Content filtering and quality control are essential for maintaining a useful knowledge base Batch processing and rate limiting consideration for Reddit’s API constraints

Expanding on each of the bullet points: Reddit Content Structure: The implementation provides a comprehensive way to extract both posts and comments while maintaining their hierarchical relationship. The metadata extraction captures essential information like timestamps, scores, and author information, which is crucial for later analysis and filtering.

Reddit’s content structure consists of multiple layers that need to be carefully processed: class RedditContent: def init(self, praw_client): self.reddit = praw_client self.content_cache = {}

def extract_post_metadata(self, submission):
    return {
        'id': submission.id,
        'title': submission.title,
        'author': str(submission.author),
        'created_utc': submission.created_utc,
        'score': submission.score,
        'upvote_ratio': submission.upvote_ratio,
        'num_comments': submission.num_comments,
        'subreddit': submission.subreddit.display_name,
        'is_self': submission.is_self,
        'url': submission.url,
        'permalink': submission.permalink
    }

Text Preprocessing: The preprocessing pipeline handles Reddit’s specific markdown formatting, URLs, and special characters. It uses BeautifulSoup for HTML parsing and implements specific cleaners for Reddit’s content markers like ++[++removed++]++ or ++[++deleted++]++. This ensures clean, consistent text for vector embedding. The preprocessing pipeline needs to handle Reddit-specific content:

import re import
markdown from bs4 import BeautifulSoup class ContentPreprocessor: def
*init*(self): self.url_pattern =
re.compile(r’http++[++s++]++?://(?:++[++a-zA-Z++]|[++0-9++]|[++$-_@.&{plus}++]|[++!++*\++(++\++),++]|++(?:%++[++0-9a-fA-F++][++0-9a-fA-F++]++)){plus}’)
self.markdown_parser = markdown.Markdown()

def clean_markdown(self, text):
    # Convert markdown to HTML, then extract clean text
    html = self.markdown_parser.convert(text)
    soup = BeautifulSoup(html, 'html.parser')
    return soup.get_text()

def remove_special_chars(self, text):
    # Remove Reddit-specific formatting and special characters
    text = re.sub(r'\[removed\]|\[deleted\]', '', text)
    text = re.sub(r'&amp;|&lt;|&gt;', ' ', text)
    return text.strip()

def process_text(self, text):
    text = self.clean_markdown(text)
    text = self.remove_special_chars(text)
    text = re.sub(r'\s+', ' ', text)  # Normalize whitespace
    return text

Metadata Organization: The structured approach to metadata ensures consistency and completeness through validation. The system separates content types (posts vs comments) and maintains required fields for each, making it easier to query and analyze the data later.

class MetadataOrganizer: def init(self): self.required_fields = { post: [id, title, author, created_utc, subreddit], comment: [id, parent_id, author, created_utc] }

def validate_metadata(self, metadata, content_type):
    missing_fields = [field for field in self.required_fields[content_type]
                     if field not in metadata]
    if missing_fields:
        raise ValueError(f"Missing required fields: {missing_fields}")
    return True

def format_metadata(self, metadata, content_type):
    if self.validate_metadata(metadata, content_type):
        return {
            'content_type': content_type,
            'metadata': metadata,
            'extracted_at': datetime.utcnow().isoformat()
        }

Content Filtering: Quality control is implemented through configurable filters that check content length, score thresholds, and content status. This helps maintain a high-quality knowledge base by excluding low-value or removed content. class ContentFilter: def init(self, min_score=5, min_length=50): self.min_score = min_score self.min_length = min_length

def is_quality_content(self, content, metadata):
    if len(content) < self.min_length:
        return False
    if metadata.get('score', 0) < self.min_score:
        return False
    if '[removed]' in content or '[deleted]' in content:
        return False
    return True

Batch Processing: In this implementation we include asynchronous batch processing with rate limiting to respect Reddit’s API constraints. This ensures efficient data collection while avoiding API throttling issues. class BatchProcessor: def init(self, reddit_client, batch_size=100, rate_limit_pause=2): self.reddit = reddit_client self.batch_size = batch_size self.rate_limit_pause = rate_limit_pause

async def process_subreddit_batch(self, subreddit_name, limit=1000):
    processed_items = []
    subreddit = self.reddit.subreddit(subreddit_name)

    async for submission in subreddit.hot(limit=limit):
        if len(processed_items) >= self.batch_size:
            await asyncio.sleep(self.rate_limit_pause)
            yield processed_items
            processed_items = []

        processed_items.append(self.process_submission(submission))

    if processed_items:
        yield processed_items

The next code segment prepares text content for vector storage via embedding and metadata addition if applicable.

class VectorPreparation: def init(self, model_name=all-MiniLM-L6-v2): self.encoder = SentenceTransformer(model_name)

def prepare_for_storage(self, processed_content, metadata):
    # Combine title and content for posts
    if metadata['content_type'] == 'post':
        text_for_embedding = f"{metadata['metadata']['title']} {processed_content}"
    else:
        text_for_embedding = processed_content

    vector = self.encoder.encode(text_for_embedding)
    return {
        'vector': vector,
        'content': processed_content,
        'metadata': metadata
    }

We finish up the code by adding logging, to complete a practical code base. class ContentProcessingLogger: def init(self, log_file=content_processing.log): self.logger = logging.getLogger(content_processor) self.logger.setLevel(logging.INFO)

    handler = logging.FileHandler(log_file)
    handler.setFormatter(logging.Formatter(
        '%(asctime)s - %(levelname)s - %(message)s'
    ))
    self.logger.addHandler(handler)

def log_processing_error(self, content_id, error):
    self.logger.error(f"Error processing content {content_id}: {error}")

def log_processing_success(self, content_id):
    self.logger.info(f"Successfully processed content {content_id}")

Summary

The section provides a comprehensive framework for extracting and processing Reddit content in a way that’s suitable for vector database storage. The implementation balances efficiency with robustness, including error handling and logging to ensure reliable operation. The modular design allows for easy customization and extension of the processing pipeline, making it adaptable to different use cases and requirements.

Sec 7 Generating and Storing Embeddings

Implementation of text vectorization using modern embedding models and storing the vectors in SQLite.

key points to cover: Selection and setup of embedding models Batching strategy for efficient embedding generation Optimization of embedding storage in SQLite Handling of different content types (titles, text, comments) Memory management for large-scale embedding generation

Embedding Model Setup: For embedding we use SentenceTransformers with the all-MiniLM-L6-v2 model, which provides a good balance between performance and quality. The setup includes device management (CPU/GPU) and dimension handling, making it flexible for different deployment scenarios.

from sentence_transformers import SentenceTransformer import torch from typing import List, Dict, Union import numpy as np

class EmbeddingGenerator: def init(self, model_name: str = all-MiniLM-L6-v2, device: str = None): ““” Initialize the embedding generator.

    Args:
        model_name: Name of the SentenceTransformer model to use
        device: Device to run the model on ('cuda' or 'cpu' or ‘mps’ for Apple Silicon )
    """
 # check if we are on Apple Silicon, if so use the ‘mps’ device

self.device = ‘mps’ if torch.backends.mps.is_available() # now see if we need to check further for cuda else default to ‘cpu’ self.device = device or (cuda if torch.backends.cuda.is_available() or else cpu) self.model = SentenceTransformer(model_name, device=self.device) self.embedding_dim = self.model.get_sentence_embedding_dimension()

@property
def dimension(self) -> int:
    """Get the embedding dimension."""
    return self.embedding_dim

Batching Strategy: The batching implementation optimizes memory usage and processing speed by handling content in configurable batch sizes. It includes separate handling for posts and comments, combining relevant fields (like title and text for posts) before embedding generation.

class ContentEmbedder: def *init*(self, embedding_generator:
EmbeddingGenerator, batch_size: int = 32): self.generator =
embedding_generator self.batch_size = batch_size

def prepare_text(self, content: Dict) -> str:
    """Prepare content for embedding by combining relevant fields."""
    if content['type'] == 'submission':
        # For posts, combine title and selftext
        return f"{content['title']} {content.get('selftext', '')}"
    else:
        # For comments, use body text
        return content['body']

def batch_embed(self, contents: List[Dict]) -> np.ndarray:
    """Generate embeddings for a batch of content."""
    texts = [self.prepare_text(content) for content in contents]

    # Generate embeddings in batches
    embeddings = []
    for i in range(0, len(texts), self.batch_size):
        batch_texts = texts[i:i + self.batch_size]
        batch_embeddings = self.generator.model.encode(
            batch_texts,
            convert_to_numpy=True,
            show_progress_bar=False
        )
        embeddings.append(batch_embeddings)

    return np.vstack(embeddings)

Storage Optimization: The SQLite storage implementation uses the VSS extension efficiently, storing embeddings as binary data and maintaining separate tables for vectors and metadata. This approach optimizes query performance while maintaining data integrity.

class EmbeddingStorage: def *init*(self, db_path: str): self.conn =
sqlite3.connect(db_path) self.setup_database()

def setup_database(self):
    """Create necessary tables and indexes."""
    self.conn.executescript("""
        -- Enable SQLite VSS extension
        CREATE VIRTUAL TABLE IF NOT EXISTS content_embeddings USING vss0(
            embedding_vector(384),  -- Adjust dimension based on model
            content_id TEXT,
            content_type TEXT
        );

        CREATE TABLE IF NOT EXISTS content_metadata (
            content_id TEXT PRIMARY KEY,
            content_type TEXT,
            title TEXT,
            author TEXT,
            created_utc INTEGER,
            score INTEGER,
            metadata JSON
        );
    """)
    self.conn.commit()

def store_embeddings(self, embeddings: np.ndarray, contents: List[Dict]):
    """Store embeddings and associated metadata."""
    cursor = self.conn.cursor()

    for embedding, content in zip(embeddings, contents):
        # Store embedding vector
        cursor.execute("""
            INSERT INTO content_embeddings (embedding_vector, content_id, content_type)
            VALUES (?, ?, ?)
        """, (embedding.tobytes(), content['id'], content['type']))

        # Store metadata
        cursor.execute("""
            INSERT INTO content_metadata
            (content_id, content_type, title, author, created_utc, score, metadata)
            VALUES (?, ?, ?, ?, ?, ?, ?)
        """, (
            content['id'],
            content['type'],
            content.get('title', ''),
            content['author'],
            content['created_utc'],
            content['score'],
            json.dumps(content)
        ))

    self.conn.commit()

Content Type Handling: The system implements specific handling for different content types (posts vs comments), ensuring appropriate text combination and metadata storage for each type. This maintains context while generating meaningful embeddings.

class ContentEmbedder: def *init*(self, embedding_generator:
EmbeddingGenerator, batch_size: int = 32): self.generator =
embedding_generator self.batch_size = batch_size
self.type_handler = ContentTypeHandler()

def batch_embed(self, contents: List[Dict]) -> np.ndarray:
    """Generate embeddings for a batch of content."""
    # Process text based on content type
    texts = [self.type_handler.prepare_for_embedding(content) for content in contents]

    # Generate embeddings in batches
    embeddings = []
    for i in range(0, len(texts), self.batch_size):
        batch_texts = texts[i:i + self.batch_size]
        batch_embeddings = self.generator.model.encode(
            batch_texts,
            convert_to_numpy=True,
            show_progress_bar=False
        )
        embeddings.append(batch_embeddings)

    return np.vstack(embeddings)

Memory Management: The pipeline implements memory-efficient processing through streaming and batching, preventing memory overflow when processing large amounts of content. The async implementation allows for efficient processing of large datasets.

class EmbeddingPipeline: def *init*(self, embedding_generator:
EmbeddingGenerator, storage: EmbeddingStorage, batch_size: int =
32): self.embedder = ContentEmbedder(embedding_generator,
batch_size) self.storage = storage self.batch_size =
batch_size

async def process_content_stream(self, content_iterator):
    """Process a stream of content with memory-efficient batching."""
    current_batch = []

    async for content in content_iterator:
        current_batch.append(content)

        if len(current_batch) >= self.batch_size:
            # Generate embeddings for current batch
            embeddings = self.embedder.batch_embed(current_batch)

            # Store embeddings and metadata
            self.storage.store_embeddings(embeddings, current_batch)

            # Clear batch
            current_batch = []

    # Process remaining items
    if current_batch:
        embeddings = self.embedder.batch_embed(current_batch)
        self.storage.store_embeddings(embeddings, current_batch)

Summary

The section provides a comprehensive implementation for generating and storing embeddings from Reddit content. The code balances efficiency with robustness, implementing proper memory management and quality controls while maintaining high performance. The modular architecture allows for easy customization of the embedding process, from model selection to storage configuration, making it adaptable to different use cases and requirements.

Sec 8 Understanding Indexing for Similarity search

Vector similarity search can become computationally expensive as your dataset grows. Without proper indexing, each query would require calculating distances between the query vector and every vector in the database—an O(n) operation. Efficient indexing strategies reduce this complexity significantly.

We consider one of the indexes that sqlite-vss offers

IVF (Inverted File) indexing partitions vectors into clusters to reduce search complexity

This approach transforms an O(n) exact search into a much faster O(k+m) approximate search

IVF provides a tunable precision-speed tradeoff through cluster configuration

The indexing strategy uses k-means clustering to organize vectors into searchable groups

Probing multiple clusters during search increases recall at the cost of performance

Expanding on each point:

IVF Indexing Structure: IVF works by clustering vectors and maintaining an "inverted file" structure where each cluster has a list of associated vectors. This dramatically reduces search space by allowing the algorithm to examine only relevant clusters. It transforms brute-force search into a two-step process: finding relevant clusters, then searching within those clusters.

The Inverted File (IVF) index is a fundamental technique for approximate nearest neighbor (ANN) search that works by: Clustering: Dividing the vector space into clusters during index construction Quantization: Assigning each vector to its nearest cluster centroid Inverted Lists: Maintaining lists of vectors that belong to each cluster Search: Querying only vectors in the most relevant clusters

Search Complexity Reduction: Without indexing, each query requires calculating distances against every vector in the database (O(n) complexity). With IVF, we first find the nearest clusters (O(k) where k is cluster count), then search only vectors in selected clusters (O(m) where m << n). This results in O(k+m) complexity, often orders of magnitude faster.

Precision-Speed Tradeoff: IVF provides configurable parameters that let you balance between search accuracy and speed. Increasing the number of clusters improves precision but increases index size and construction time. Similarly, probing more clusters during search improves recall but reduces performance. This makes IVF highly adaptable to different use cases.

K-means Clustering Implementation: The implementation uses k-means to partition the vector space into clusters, with each vector assigned to its nearest centroid. This spatial organization enables efficient searching by focusing only on regions of the vector space likely to contain similar vectors. The centroids and assignments are stored in separate tables for efficient retrieval.

Multi-Cluster Probing: To increase recall without rebuilding the index, IVF search can probe multiple nearby clusters. This is crucial because the nearest neighbor might occasionally fall into an adjacent cluster. The implementation includes a probe_clusters parameter that determines how many clusters to examine, providing a runtime-adjustable precision-speed balance.

Summary

IVF indexing provides an elegant solution for approximate nearest neighbor search that balances performance and accuracy. By organizing vectors into clusters and focusing search efforts only on relevant portions of the vector space, it enables fast similarity search even as datasets grow. The implementation integrates well with SQLite VSS, making it particularly suitable for personal knowledge bases where both performance and relevance are important.

Implementing Semantic Search Queries Writing effective queries that combine vector similarity search with metadata filtering for relevant results.

Combining vector similarity with metadata filtering Creating parameterized query templates Optimizing search performance Working with multi-stage filtering for complex searches Practical applications for Reddit content analysis 9.1 Key Components of Effective Semantic Search Queries Implementing effective semantic search queries with SQLite-VSS requires understanding how to balance vector similarity search with traditional SQL filtering. Here’s an example that builds on the Reddit data processing from earlier sections: import sqlite3 from sentence_transformers import SentenceTransformer

def search_reddit_content(query_text, subreddit=None,
min_score=0, limit=10): # Connect to database and load VSS extension
conn = sqlite3.connect('`reddit_vectors.db`')
conn.enable_load_extension(True)
conn.load_extension("`./sqlite-vss0`") cursor = conn.cursor()

# Generate embedding for query text
model = SentenceTransformer('all-MiniLM-L6-v2')
query_embedding = model.encode(query_text)

# Convert embedding to VSS vector format
query_vector = "vss_create_vector(" + ",".join(map(str, query_embedding)) + ")"

# Construct base query with vector similarity
base_query = """
SELECT
    posts.id,
    posts.title,
    posts.url,
    posts.score,
    posts.created_utc,
    posts.subreddit,
    vss_cosine_distance(posts.content_vector, {0}) AS distance
FROM posts
WHERE 1=1
"""

# Add metadata filters if provided
params = []
if subreddit:
    base_query += " AND posts.subreddit = ?"
    params.append(subreddit)

if min_score > 0:
    base_query += " AND posts.score >= ?"
    params.append(min_score)

# Complete query with ordering and limit
complete_query = base_query + """
ORDER BY distance ASC
LIMIT ?
"""
params.append(limit)

# Execute query and fetch results
formatted_query = complete_query.format(query_vector)
results = cursor.execute(formatted_query, params).fetchall()

conn.close()
return results

9.2 Combining Vector Similarity with Metadata Filtering One of the most powerful aspects of using SQLite-VSS is the ability to combine vector similarity search with traditional SQL filtering. This example demonstrates advanced filtering options:

def advanced_reddit_search(query_text, filters=None,
sort_by="`relevance`", time_period=None, limit=20): ““” Perform
advanced semantic search with multiple filtering options.

Args:
    query_text: The search query text
    filters: Dict containing filtering options (subreddit, author, min_score, etc.)
    sort_by: How to sort results ("relevance", "date", "score", or "comments")
    time_period: Optional time restriction ("day", "week", "month", "year")
    limit: Maximum number of results to return
"""
conn = sqlite3.connect('reddit_vectors.db')
conn.enable_load_extension(True)
conn.load_extension("./sqlite-vss0")
cursor = conn.cursor()

# Generate query embedding
model = SentenceTransformer('all-MiniLM-L6-v2')
query_embedding = model.encode(query_text)
query_vector = "vss_create_vector(" + ",".join(map(str, query_embedding)) + ")"

# Start building query
query_parts = [
    "SELECT p.id, p.title, p.selftext, p.score, p.num_comments, p.created_utc,",
    "p.subreddit, p.author, vss_cosine_distance(p.content_vector, {0}) AS distance",
    "FROM posts p",
    "WHERE 1=1"
]

params = []

# Apply filters
if filters:
    if 'subreddit' in filters:
        query_parts.append("AND p.subreddit IN ({})".format(
            ','.join(['?'] * len(filters['subreddit']))))
        params.extend(filters['subreddit'])

    if 'author' in filters:
        query_parts.append("AND p.author = ?")
        params.append(filters['author'])

    if 'min_score' in filters:
        query_parts.append("AND p.score >= ?")
        params.append(filters['min_score'])

    if 'min_comments' in filters:
        query_parts.append("AND p.num_comments >= ?")
        params.append(filters['min_comments'])

# Apply time period filter
if time_period:
    current_time = int(time.time())
    time_filters = {
        'day': current_time - 86400,
        'week': current_time - 604800,
        'month': current_time - 2592000,
        'year': current_time - 31536000
    }
    if time_period in time_filters:
        query_parts.append("AND p.created_utc >= ?")
        params.append(time_filters[time_period])

# Apply sorting
sort_options = {
    'relevance': "ORDER BY distance ASC",
    'date': "ORDER BY p.created_utc DESC",
    'score': "ORDER BY p.score DESC",
    'comments': "ORDER BY p.num_comments DESC"
}
query_parts.append(sort_options.get(sort_by, sort_options['relevance']))

# Apply limit
query_parts.append("LIMIT ?")
params.append(limit)

# Execute query
final_query = " ".join(query_parts).format(query_vector)
results = cursor.execute(final_query, params).fetchall()

conn.close()
return results

9.3 Implementing Semantic Search Using Topic Clustering

Here’s a practical example showing how to implement semantic topic clustering across subreddits:

def find_related_discussions(seed_post_id,
min_similarity=0.7, max_results=20): ““"`Find related
discussions across subreddits based on a seed post`"“” conn =
sqlite3.connect('`reddit_vectors.db`')
conn.enable_load_extension(True)
conn.load_extension("`./sqlite-vss0`") cursor = conn.cursor()

# Get the seed post's vector
cursor.execute(
    "SELECT content_vector, subreddit FROM posts WHERE id = ?",
    (seed_post_id,)
)
result = cursor.fetchone()
if not result:
    conn.close()
    return []

seed_vector, seed_subreddit = result

# Find similar posts across different subreddits
query = """
SELECT
    p.id, p.title, p.subreddit, p.score, p.num_comments,
    p.created_utc, vss_cosine_distance(p.content_vector, ?) AS similarity
FROM posts p
WHERE p.id != ?
AND p.subreddit != ?
AND vss_cosine_distance(p.content_vector, ?) <= ?
ORDER BY similarity ASC
LIMIT ?
"""

params = (
    seed_vector,
    seed_post_id,
    seed_subreddit,
    seed_vector,
    1 - min_similarity,  # Convert similarity threshold to distance
    max_results
)

results = cursor.execute(query, params).fetchall()

# Group by subreddit for analysis
subreddit_discussions = {}
for post in results:
    subreddit = post[2]
    if subreddit not in subreddit_discussions:
        subreddit_discussions[subreddit] = []
    subreddit_discussions[subreddit].append({
        'id': post[0],
        'title': post[1],
        'score': post[3],
        'comments': post[4],
        'created_utc': post[5],
        'similarity': 1 - post[6]  # Convert distance back to similarity
    })

conn.close()
return subreddit_discussions

In summary, implementing effective semantic search queries with SQLite-VSS involves balancing vector similarity operations with traditional SQL filtering. By combining these approaches, you can create powerful search functionality that using the meaning of content while still respecting metadata constraints. The key advantage of using SQLite-VSS for Reddit content analysis is the ability to perform sophisticated queries within a lightweight, file-based database system that’s perfect for personal knowledge management applications. The techniques demonstrated here show how to progress from basic vector similarity searches to complex queries that can address sophisticated content analysis needs. With proper indexing and query optimization strategies, even a relatively simple SQLite database can provide impressive semantic search capabilities for Reddit content analysis and for other content types enabling sophisticated techniques for personal knowledge management.

9.4 Finding Contextually Similar Content Across Subreddits One of the most powerful applications of vector search is identifying semantically similar content across different communities: def cross_subreddit_similarity_search(topic, source_subreddits, target_subreddits, limit=20): ““” Find how the same topic is discussed across different subreddit communities.

Args:
    topic: Topic description or query
    source_subreddits: List of subreddits to use as source material
    target_subreddits: List of subreddits to search for similar discussions
    limit: Maximum results to return
"""
conn = sqlite3.connect('reddit_vectors.db')
conn.enable_load_extension(True)
conn.load_extension("./sqlite-vss0")
cursor = conn.cursor()

# Generate embedding for the topic
model = SentenceTransformer('all-MiniLM-L6-v2')
topic_embedding = model.encode(topic)
topic_vector = "vss_create_vector(" + ",".join(map(str, topic_embedding)) + ")"

# Find representative posts from source subreddits
source_placeholders = ','.join(['?'] * len(source_subreddits))
source_query = f"""
SELECT
    id, title, selftext, subreddit,
    vss_cosine_distance(content_vector, {topic_vector}) AS distance
FROM posts
WHERE subreddit IN ({source_placeholders})
ORDER BY distance ASC
LIMIT 10
"""

source_posts = cursor.execute(source_query, source_subreddits).fetchall()

# Find similar posts in target subreddits
results = []
target_placeholders = ','.join(['?'] * len(target_subreddits))

for post in source_posts:
    post_id, post_title, post_text, post_subreddit, _ = post

    # Find the post's vector
    cursor.execute("SELECT content_vector FROM posts WHERE id = ?", (post_id,))
    post_vector = cursor.fetchone()[0]

    # Find similar posts in target subreddits
    similar_query = f"""
    SELECT
        id, title, selftext, subreddit, created_utc, score,
        vss_cosine_distance(content_vector, ?) AS distance
    FROM posts
    WHERE subreddit IN ({target_placeholders})
    ORDER BY distance ASC
    LIMIT ?
    """

    params = [post_vector] + target_subreddits + [limit // len(source_posts)]
    similar_posts = cursor.execute(similar_query, params).fetchall()

    # Add to results with source information
    for similar in similar_posts:
        results.append({
            'source_post': {
                'id': post_id,
                'title': post_title,
                'subreddit': post_subreddit
            },
            'similar_post': {
                'id': similar[0],
                'title': similar[1],
                'text': similar[2][:200] + "..." if len(similar[2]) > 200 else similar[2],
                'subreddit': similar[3],
                'date': datetime.fromtimestamp(similar[4]).strftime('%Y-%m-%d'),
                'score': similar[5],
                'similarity': 1 - similar[6]  # Convert distance to similarity
            }
        })

# Sort by similarity
results.sort(key=lambda x: x['similar_post']['similarity'], reverse=True)

conn.close()
return results[:limit]

4.10 Conclusion In this chapter we dove deep into the application of the sqlite-vss vector extension to managing data sourced from Reddit as an example of a personal knowledge data source. We saw that sqlite-vss vector similarity search is an effective framework for building a personal knowledge management system. In this chapter the data source was Reddit, but it could be anything else such as Twitter or WhatsApp or email. The principles would remain the same, the interfaces to the data source would be different. So using the techniques in this chapter you can add the data sources of your choice. This will allow you to build out a personal knowledge management system now that you know how to use Sqlite3 as a vector database.

Chapter 5. Similarity Search with PostgreSQL and pgvector

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles. This will be the 5th chapter of the final book. Please note that the GitHub repo will be made active later on. If you’d like to be actively involved in reviewing and commenting on this draft, please reach out to the editor at shunter@oreilly.com.

In this chapter we will do with PostgreSQL and pgvector what we did with SQLIte3 and sqlite-vss in the previous chapter - explore similarity search in depth and use the knowledge to build a practical knowledge base.

This chapter moves beyond the typical Reddit example used in many similarity search tutorials to address a more prevalent and challenging real-world application: leveraging the vast repository of research papers available as ArXiv PDFs. This deliberate shift highlights a common hurdle encountered in professional settings—the need to transform unstructured data, specifically PDF documents, into a structured format suitable for sophisticated semantic search capabilities. By focusing on ArXiv PDFs, the chapter will delve into the practicalities of extracting meaningful textual content from these documents, which often contain complex layouts, mathematical formulas, and diverse formatting. This extracted information will then serve as the foundation for building a similarity search system powered by PostgreSQL and the pgvector extension, enabling users to efficiently discover relevant research based on the semantic meaning of the content rather than just keyword matching. The exploration of this ArXiv use case will provide valuable insights into handling real-world unstructured data and building robust semantic search solutions for large document collections.

Once we have extracted the content out of its PDF wrapper we will, you guessed it, put it in a vector database - in this case PostgreSQL with a pgvector extension which gives us capabilities similar to the ones seen in the previous chapter. However since PostgreSQL is a common platform in business applications, using it rather than SQLite3 gives us the ability to integrate vector databases and similarity search with traditional business data in tables, allowing hybrid databases more suited to the evolving needs of business rather than pure vector or pure RDBs.

While we use ArXiv as a data source we are really using it for its store of complex PDFs more than scientific content, primarily so that the techniques of data extraction and similarity search can be applied to other kinds of PDFs such as reports, white papers, marketing requirements docs, design docs and quarterly financial documents to name a few.

We start with the nitty gritty of the internals pgvector and installation then move on to the vector type, distances and indexes. Once we have that we start building our search database and take it from there.

Section 5.1: Understanding PgVector Extension for PostgreSQL

PostgreSQL’s pgvector extension adds native vector similarity search capabilities to one of the most popular enterprise-grade relational databases. This extension bridges the gap between traditional relational databases and modern AI applications by providing efficient vector operations alongside standard SQL functionality.

This section will introduce PgVector as PostgreSQL’s native vector similarity extension. It will explain how PgVector differs from SQLite-VSS, highlighting PostgreSQL’s advantages for larger datasets, concurrent access, and enterprise features. The section will cover the mathematical foundations of vector similarity in PostgreSQL and explain how PgVector fits within PostgreSQL’s extensibility model.

5.1.1 Vector Operations

pgvector introduces a new vector data type and associated operations essential for modern AI and machine learning applications. This enables:

Storage of high-dimensional vectors (embeddings)

Efficient similarity search operations

Integration with existing PostgreSQL features

5.1.2 Distance Metrics

The extension supports three primary distance metrics:

L2 (Euclidean) Distance

Measures the straight-line distance between vectors

Useful for geometric applications

Formula: sqrt(sumx[i]* y[i])^2

Inner Product

Computationally efficient

Best for normalized vectors

Formula: sum(x[i] * y[i])

Cosine Distance

Most common for text embeddings

Measures angular similarity

Formula: 1 - (inner_product(x,y)/(|x|*|y|))

In practice vectors need to be normalized before doing cosine similarity

5.1.3 Enterprise Features

pgvector maintains PostgreSQL’s enterprise capabilities:

ACID compliance

Multi-version concurrency control (MVCC)

Role-based access control

Backup and replication support

Transaction management

Pgvector is not an extension via a forked codeline hence can be incrementally added to existing installations

5.1.4 Performance Optimization

HNSW Indexing pgvector implements Hierarchical Navigable Small World (HNSW) indexing for approximate nearest neighbor search:

Creates a graph-based structure for efficient searching

Dramatically improves query performance

Configurable trade-off between search speed and accuracy

It also supports IVFFlat for approximate indexing.

HNSW is best for high accuracy and exploratory search, while IVFFlat may offer faster performance at the cost of recall — especially on large datasets.

5.1.5 Query Optimization

The extension integrates with PostgreSQL’s query planner:

Cost-based optimization for vector operations

Efficient handling of hybrid queries combining vector and scalar operations

Index-based search acceleration

5.1.6 Use Cases

Text Similarity Search

Document recommendation systems

Semantic search applications

Content deduplication

Image Similarity

Visual search systems

Image deduplication

Style matching

Recommendation Systems

Product recommendations

Content suggestions

User similarity matching

5.1.7 Integration Considerations

When migrating to pgvector: * Plan vector dimension requirements * Consider index storage requirements * Evaluate memory requirements for vector operations

5.1.8 Performance Monitoring

Key metrics to monitor:

Query response times

Index build times

Memory usage during vector operations

Storage growth rate

5.1.9 Best Practices

Vector Dimensionality

Choose appropriate dimensions for your use case

Consider dimension reduction techniques if needed

Be consistent with vector sizes within collections

Indexing Strategy

Use HNSW indexes for large collections

Configure index parameters based on accuracy requirements

Monitor index build times and sizes

Query Optimization

Combine vector searches with traditional SQL where appropriate

Use batch operations for bulk processing

Implement appropriate caching strategies 5.1.10 Comparison with Other Solutions vs SQLite-VSS

Enterprise-grade features

Better concurrency support

Designed for larger datasets

More complex setup and maintenance vs Dedicated Vector Databases

Integrated with existing PostgreSQL infrastructure

Simpler architecture for hybrid workloads

Lower operational complexity

Potentially lower performance for pure vector workloads

5.1.11 Conclusion

pgvector represents a significant advancement in making vector similarity search accessible to mainstream database users, particularly in enterprise environments. By combining vector operations with PostgreSQL’s robust feature set and extensible type system, it provides a practical solution for organizations looking to implement AI capabilities without adding separate specialized databases.

The extension’s integration with PostgreSQL’s enterprise features makes it particularly valuable for organizations that need to maintain strict data governance while adding vector similarity search capabilities. This combination of traditional relational database features with modern vector operations positions pgvector as a versatile platform for building sophisticated AI-enabled applications.

Section 2: Installing and Configuring PgVector

In this section we learn step-by-step instructions for installing PgVector, including prerequisites and configuration options. We will only cover installation from packages and not compilation from source. We will also explain essential configuration parameters and how to optimize them for different workloads but as a first pass these can be ignored and used only when your work goes beyond initial learning and into production where performance becomes an important topic. Performance considerations for different deployment scenarios will be discussed.

Configuration Parameters: The key parameters that affect pgvector’s performance include vector dimension limits, maximum index size, and memory allocation for vector operations. These need to be tuned based on your specific use case and available resources once you approach production.

2.1 Prerequisites

Before installing pgvector, ensure your system meets these requirements: PostgreSQL 11 or higher

Sufficient RAM for vector operations typically a spare 1G while learning and exploring

Adequate storage for vector indices again typically spare 2G while learning and exploring

2.2 Installation via package manager

For Ubuntu/Debian

# Add PostgreSQL repository
sudo apt install postgresql-common
sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh
# Install PostgreSQL and development packages
sudo apt install postgresql-14 postgresql-server-dev-14
# Install pgvector
sudo apt install postgresql-14-pgvector

For RHEL/CentOS

# Enable PostgreSQL repository
sudo dnf install -y https://download.postgresql.org/pub/repos/yum/reporpms/EL-8-x86_64/pgdg-redhat-repo-latest.noarch.rpm
# Install PostgreSQL and pgvector
sudo dnf install -y postgresql14-server postgresql14-devel pgvector

For Docker

FROM postgres:14
# Install build dependencies
RUN apt-get update && apt-get install -y \
    build-essential \
    git \
    postgresql-server-dev-14
# Clone and install pgvector
RUN git clone https://github.com/pgvector/pgvector.git \
    && cd pgvector \
    && make \
    && make install

Note that only the very basic installation targets have been described. Other variations include cloud installations, many of which come with pgvector installed, and docker installations which we will skip over. My suggestion will be if you haven’t installed PostgreSQL ever before then it’s not a good idea to start with the docker one at the outset. For better or for worse, we will only consider the local installation as we have a lot more complexity to deal with at the application layer and we want to keep infrastructure complexity to a minimum until we are humming along with the applications.

2.3 Basic Configuration

After installation install the extension in Postgres

CREATE EXTENSION vector;

To confirm that the extension vector is installed run the following SQL query in psql or other Postgres client

SELECT* FROM pg_extension;

This will return all installed extensions. This list should include “vector”.

2.4 Configuration and Performance Tuning for HNSW index

For similarity search we use the HNSW (Hierarchical Navigable Small Worlds) index which we first saw in the chapter on FAISS. It works very well for embeddings of language and images which contain multiple distinct clusters connected by sparse connections. A sample such index is created as follows:

-- Create an HNSW index with custom parameters
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)
WITH (
    dim = 384,        -- Vector dimension
    m = 16,           -- Number of connections per layer
    ef_construction = 64  -- Size of dynamic candidate list
);

The code creates a specialized index using the HNSW algorithm to efficiently find the nearest neighbors of vectors in the embedding column of the items table. It uses Euclidean distance to measure similarity, and the parameters dim, m, and ef_construction are set to specific values to control the index’s behavior. This is a common pattern for setting up vector similarity search in PostgreSQL with the pgvector extension.

Let’s look at it a clause at a time:

CREATE INDEX ON items: This part is standard SQL, indicating the creation of an index on the items table. Indexes improve the speed of data retrieval.

USING hnsw: This is the crucial part. It specifies that the index will use the Hierarchical Navigable Small World (HNSW) algorithm. HNSW is a highly efficient algorithm specifically designed for approximate nearest neighbor (ANN) search in high-dimensional spaces. It’s not a standard SQL feature; it’s an extension provided by a database system that supports vector search

(embedding vector_l2_ops): This part indicates: embedding: The index is being built on the embedding column. This column likely stores vector embeddings (numerical representations of data, like words, images, or other objects).

vector_l2_ops: This is an “operator class”. It tells the database how to compare vectors within the embedding column. vector_l2_ops specifically means it will use the L2 distance (Euclidean distance) to measure the similarity between vectors. This is a common distance metric for vector comparisons. Other options might include vector_ip_ops (inner product) or vector_cosine_ops (cosine similarity).

WITH (…​): This section specifies parameters that configure the HNSW index. These parameters tune the index’s performance (speed and accuracy) and memory usage.

dim = 384: This sets the dimensionality of the vectors in the embedding column. Each vector has 384 numbers (features). This must match the actual dimensionality of the data in the embedding column.

m = 16: This is the number of bidirectional connections each element will have in each layer of the HNSW graph. A higher m generally leads to better accuracy but uses more memory and can slow down index construction. 16 is a common, reasonable default.

ef_construction = 64: This parameter controls the size of the “dynamic candidate list” during index construction. A larger ef_construction leads to a more accurate index (better search quality) but increases the time it takes to build the index. 64 is a reasonable default. A related parameter, ef_search, controls the size of the dynamic candidate list during search time, influencing query speed and accuracy. ef_search is often set at query time, not index creation time.

Memory Requirement Calculation

Vector storage: dimensions × 4 bytes × number of vectors

Index overhead: ~50-100 bytes per vector

Working memory: dimensions × 4 bytes × batch size * Example for 1 million 384-dimensional vectors:

vector_storage = 384* 4* 1,000,000 # ~1.5GB index_overhead = 75* 1,000,000 # ~75MB working_memory = 384 * 4 * 1000 # ~1.5MB per batch of 1000

Monitoring and Maintenance

Index size and growth

        SELECT pg_size_pretty(pg_relation_size('index_name'));
pg_relation_size('index_name')

Returns the raw size (in bytes) of the given index (index_name).

This function only measures the index itself, not the associated table. pg_size_pretty(…​)

Converts the raw byte size into a readable format (e.g., 12 MB, 1 GB).

Example output

 pg_size_pretty
----------------
  45 MB
(1 row)

To track how an index grows over time, you can run this query periodically or store results in a monitoring table.

Query Performance
EXPLAIN ANALYZE
SELECT * FROM items
ORDER BY embedding <-> query_vector
LIMIT 10;

This query analyzes the performance of a nearest neighbor search using pgvector in PostgreSQL.

EXPLAIN ANALYZE

This makes PostgreSQL execute the query and return a detailed execution plan, including:

The indexes used

Estimated vs. actual execution time

Number of rows scanned

SELECT * FROM items

Retrieves all columns from the items table.

ORDER BY embedding <→ query_vector

This uses the pgvector distance operator (<→) to compute the distance between stored embeddings and query_vector:

For vector_l2_ops: '<→ represents Euclidean (L2) distance.

For vector_cosine_ops: '<→ represents Cosine Distance (1* Cosine Similarity).

For vector_ip_ops: '<→ represents Negative Inner Product.

The results are sorted by similarity (nearest vectors first).

LIMIT 10

Returns only the 10 most similar items (nearest neighbors).

What EXPLAIN ANALYZE Reveals

Running this query will return a query execution plan, showing details such as:

Index Usage

If an HNSW index (USING hnsw) or IVFFlat index (USING ivfflat) is present, the plan should show: Index Scan using my_hnsw_index on items

If no index is used, a Sequential Scan (Slow!) will appear: Seq Scan on items Warning: This means PostgreSQL is scanning every row, making the search inefficient.

Execution Time

Shows how long the query takes (actual time=…​).

If the query is slow, indexing should be optimized.

Row Scanning Details

Displays the number of rows PostgreSQL checked before returning the top 10 results.

Optimizing Performance

If your query is slow, ensure you have the right index for efficient nearest neighbor search:

HNSW Index (Recommended for high accuracy & speed)

CREATE INDEX items_embedding_hnsw ON items

USING hnsw (embedding vector_cosine_ops);

IVFFlat Index (Faster but less accurate, useful for large datasets)

CREATE INDEX items_embedding_ivfflat ON items

USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);

Example EXPLAIN ANALYZE Output (With Index)

Index Scan using items_embedding_hnsw on items

(actual time=1.23..1.56 rows=10 loops=1)

Fast execution with index usage.

Example EXPLAIN ANALYZE Output (No Index)

Seq Scan on items (actual time=2456.87..2460.43 rows=10 loops=1) Slow execution because it scans the entire table.

Key Takeaways

This query finds the 10 nearest neighbors using the pgvector distance operator.

EXPLAIN ANALYZE helps check whether an index is being used.

If slow, add an HNSW or IVFFlat index to optimize search performance.

Memory Usage

To check memory used by active queries

SELECT pid, usename, application_name, query,
       pg_size_pretty(memory_bytes) AS memory_used,
       state, query_start
FROM pg_stat_activity
LEFT JOIN pg_backend_memory_contexts() ON pg_stat_activity.pid = pg_backend_memory_contexts.pid
WHERE state = 'active'
ORDER BY memory_bytes DESC;
Identify high memory queries
SELECT pid, usename, query, total_exec_time,
       shared_blks_hit, shared_blks_read, temp_blks_written
FROM pg_stat_statements
ORDER BY total_exec_time DESC
LIMIT 10;

Maintenance

Update indexes and statistics periodically for best performance

REINDEX INDEX index_name;
ANALYZE items;

The configuration process for pgvector involves more than simply enabling the extension; it also needs careful tuning of PostgreSQL parameters to optimize performance for vector operations. This tuning is essential to ensure that the database can efficiently handle the computational demands of vector similarity search.

Performance optimization is a critical aspect of deploying pgvector, particularly in terms of memory management and index configuration. The HNSW (Hierarchical Navigable Small World) index, commonly used for vector similarity search, has parameters that need to be adjusted based on the specific requirements of the application. These parameters govern the trade-off between accuracy and speed, and finding the optimal balance is crucial for achieving good performance.

In addition to index tuning, system resources such as memory and CPU need to be allocated appropriately to handle vector operations efficiently. This may involve adjusting PostgreSQL configuration settings or allocating additional resources to the database server. Careful consideration of the workload and performance requirements is necessary to ensure that the system can handle vector operations without bottlenecks or resource contention.

Section 3: Vector Data Types and Distance Functions in PostgreSQL

Here we explore the vector data types provided by PgVector and how they’re implemented. We will see how to compare the available distance functions (cosine, L2, inner product, etc.) with mathematical explanations and practical usage examples. The section will include benchmarks showing performance characteristics of different distance metrics for LLM embedding vectors.

3.1 Vector Data Type Implementation

Basic Structure The vector data type in pgvector is implemented as a variable-length array of 32-bit floating-point numbers. The internal structure includes:

typedef struct Vector {
    int32   vl_len_;      /* varlena header (do not touch directly!) */
    int32   dim;          /* number of dimensions */
    float4  x[FLEXIBLE_ARRAY_MEMBER];  /* vector data */
} Vector;

3.2 Creating Vectors

Vectors can be created in several ways:

-- From array
SELECT ARRAY[1,2,3]::vector;

-- From string representation
SELECT '[1,2,3]'::vector;

-- Using vector construction function
SELECT vector(ARRAY[1,2,3]);

-- Creating table with vector column
CREATE TABLE items (
    id bigserial PRIMARY KEY,
    embedding vector(384)  -- specify dimension
);

3.3 Distance Functions

1. Euclidean Distance (L2)

The L2 distance measures the straight-line distance between two points in vector space.

Mathematical definition: L2(x,y)=∑i=1n(xi−yi)2

Implementation:

-- Using operator syntax
SELECT embedding <-> query_vector AS l2_distance
FROM items
ORDER BY l2_distance
LIMIT 5;

-- Using function syntax
SELECT l2_distance(embedding, query_vector) AS distance
FROM items
ORDER BY distance
LIMIT 5;

2. Inner Product Distance

The inner product (dot product) measures the product sum of vector components.

Mathematical definition: IP(x,y)=∑i=1nxiyi

Implementation:

-- Using operator syntax
SELECT embedding <#> query_vector AS ip_distance
FROM items
ORDER BY ip_distance
LIMIT 5;

-- Using function syntax
SELECT inner_product(embedding, query_vector) AS distance
FROM items
ORDER BY distance
LIMIT 5;

3. Cosine Distance

Measures the angular distance between vectors, normalized by their magnitudes.

Mathematical definition:

Cosine(x,y)=1−∑i=1nxiyi∑i=1nxi2∑i=1nyi2

Implementation:

-- Using operator syntax
SELECT embedding <=> query_vector AS cosine_distance
FROM items
ORDER BY cosine_distance
LIMIT 5;

-- Using function syntax
SELECT cosine_distance(embedding, query_vector) AS distance
FROM items
ORDER BY distance
LIMIT 5;

3.4 Vector Operations

3.4.1 Basic Operations

Vector Addition

SELECT embedding + other_embedding AS sum_vector
FROM items;

Vector Subtraction

SELECT embedding - other_embedding AS diff_vector
FROM items;

Scalar Multiplication

SELECT embedding * 2.0 AS scaled_vector
FROM items;

3.4.2 Vector Manipulation

Normalization

CREATE OR REPLACE FUNCTION normalize_vector(v vector)
RETURNS vector AS $$
DECLARE
    magnitude float;
BEGIN
    magnitude := sqrt(l2_distance(v, vector_zeros(array_length(v::float4[], 1))));
    RETURN v * (1.0/magnitude);
END;
$$ LANGUAGE plpgsql;

Vector Concatenation

CREATE OR REPLACE FUNCTION concat_vectors(v1 vector, v2 vector)
RETURNS vector AS $$
BEGIN
    RETURN (v1::float4[] || v2::float4[])::vector;
END;
$$ LANGUAGE plpgsql;

3.5 Performance Optimization

Index Selection

Different distance metrics require different index types:

-- For L2 distance
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)
WITH (dim=384, m=16, ef_construction=64);

-- For Inner Product
CREATE INDEX ON items USING hnsw (embedding vector_ip_ops)
WITH (dim=384, m=16, ef_construction=64);

-- For Cosine Distance
CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops)
WITH (dim=384, m=16, ef_construction=64);

Query Optimization

Combining vector search with traditional filters:

-- Efficient combined query
SELECT id, embedding <-> query_vector AS distance
FROM items
WHERE category = 'electronics'
    AND created_at > NOW() - interval '1 day'
ORDER BY distance
LIMIT 10;

Batch Processing

Efficient batch vector operations:

-- Batch insert with vectors
INSERT INTO items (embedding)
SELECT vector_column
FROM unnest(ARRAY[
    '[1,2,3]'::vector,
    '[4,5,6]'::vector,
    '[7,8,9]'::vector
]) AS vector_column;

3.6 Advanced Usage Patterns

Vector Aggregation

Computing centroid vectors:

CREATE AGGREGATE vector_avg(vector) (
    SFUNC = vector_add,
    FINALFUNC = vector_divide_by_count,
    STYPE = vector
);

-- Usage
SELECT category, vector_avg(embedding) AS centroid
FROM items
GROUP BY category;

Dynamic Vector Generation

Creating vectors dynamically:

CREATE OR REPLACE FUNCTION generate_random_vector(dim integer)
RETURNS vector AS $$
BEGIN
    RETURN (
        SELECT array_agg(random()::float4)::vector
        FROM generate_series(1, dim)
    );
END;
$$ LANGUAGE plpgsql;

Vector Comparison

Implementing custom comparison functions:

CREATE OR REPLACE FUNCTION vector_similarity_threshold(
    v1 vector,
    v2 vector,
    threshold float
) RETURNS boolean AS $$
BEGIN
    RETURN cosine_distance(v1, v2) <= threshold;
END;
$$ LANGUAGE plpgsql;

3.7 Best Practices

Dimension Management

Keep dimensions consistent within collections

Consider dimension reduction techniques for large vectors

Document dimension requirements for each collection

Distance Metric Selection

Use L2 for general-purpose similarity

Use Cosine for normalized text embeddings

Use Inner Product for speed with normalized vectors

Performance Optimization

Create appropriate indexes for your distance metric

Use batch operations for bulk processing

Monitor and optimize query patterns

Error Handling

Validate vector dimensions before operations

Handle NULL values appropriately

Implement proper error checking in functions

3.8 Common Pitfalls and Solutions

Dimension Mismatch

-- Prevent dimension mismatch
CREATE OR REPLACE FUNCTION validate_vector_dim(
    v vector,
    expected_dim integer
) RETURNS boolean AS $$
BEGIN
    RETURN array_length(v::float4[], 1) = expected_dim;
END;
$$ LANGUAGE plpgsql;

Performance Issues

-- Monitor vector operation performance
CREATE OR REPLACE FUNCTION analyze_vector_query(
    query text
) RETURNS table (
    plan_output json
) AS $$
BEGIN
    RETURN QUERY EXPLAIN (FORMAT JSON) query;
END;
$$ LANGUAGE plpgsql;

Memory Management

-- Implement batch processing for large operations
CREATE OR REPLACE FUNCTION process_vectors_in_batches(
    batch_size integer
) RETURNS void AS $$
DECLARE
    batch_counter integer := 0;
BEGIN
    FOR batch IN SELECT * FROM items
                 ORDER BY id
                 LIMIT batch_size OFFSET batch_size * batch_counter
    LOOP
        -- Process batch
        batch_counter := batch_counter + 1;
        COMMIT;
    END LOOP;
END;
$$ LANGUAGE plpgsql;

Summary

pgvector’s Implementation of Vector Data and Operations

Vector Data Type: pgvector introduces a specialized data type for efficiently storing vectors. This data type is internally represented as a variable-length array of 32-bit floating-point numbers. This allows for flexible storage of vectors with varying dimensions.

Distance Metrics: A core aspect of vector operations is the ability to measure similarity or distance between vectors. pgvector supports three primary distance metrics:

Euclidean (L2) Distance: This is the geometric distance between two points in Euclidean space. It is calculated as the square root of the sum of the squared differences between corresponding vector elements.

Inner Product: This measures the similarity between two vectors based on the angle between them. It is calculated as the sum of the products of corresponding vector elements.

Cosine Distance: This is closely related to the inner product and measures the cosine of the angle between two vectors. It is often used for text and document similarity comparisons.

Choosing the Right Metric: The choice of distance metric depends on the specific use case and the nature of the data. Euclidean distance is often used for general-purpose similarity searches, while cosine distance is common in natural language processing and information retrieval tasks. *

High-Dimensional Data: pgvector is particularly well-suited for handling high-dimensional data, such as embeddings generated by large language models. These embeddings capture semantic relationships between words and phrases, and efficient distance calculations are crucial for tasks like semantic search and recommendation.

Performance Optimization

Indexing: Indexes are essential for efficient vector operations at scale. pgvector supports specialized index types, such as the HNSW (Hierarchical Navigable Small World) index, which enable fast approximate nearest neighbor searches.

Batch Operations: Performing vector operations on batches of data can significantly improve performance by reducing overhead and leveraging hardware acceleration.

Hardware Acceleration: Modern CPUs and GPUs often have specialized instructions for vector operations, and pgvector can be configured to take advantage of these capabilities for further performance gains.

Advanced Usage Patterns

Custom Aggregations: pgvector allows for the creation of custom aggregation functions that operate on vector data. This enables advanced analytics and data mining tasks.

Vector Manipulation: Beyond distance calculations, pgvector supports various vector manipulation operations, such as addition, subtraction, and scaling. These operations can be used for tasks like clustering and dimensionality reduction.

Best Practices

Data Normalization: Normalizing vector data to have unit length can improve the accuracy and efficiency of distance calculations, especially for cosine distance.

Index Maintenance: Indexes need to be maintained as data is added or updated to ensure optimal performance.

Error Handling: Robust error handling is essential for production applications that rely on vector operations.

By leveraging the capabilities of pgvector and following best practices, developers can efficiently implement powerful similarity search and data analysis applications within the PostgreSQL database environment.

Section 4: Vector Indexing Strategies in PgVector

At this point we are ready to dive into the indexing mechanisms in PgVector, particularly IVFFlat, HNSW, and other available index types. This is a meaty technical section. It will explain the trade-offs between build time, query time, recall accuracy, and index size. The section will provide guidelines for choosing the appropriate index based on dataset size, dimensionality, and query patterns specific to academic paper embeddings which we are planning on using for our ArXiv exploration database.

4.1 Vector Indexing Fundamentals

The Curse of Dimensionality

High-dimensional vector spaces present unique challenges:

Distance concentration phenomenon

Exponential growth of volume with dimensions

Degradation of distance metrics’ effectiveness

Index Requirements

Vector indexes must support:

Efficient similarity search

Dynamic updates

Scalability with dimensions

Memory efficiency

Concurrent access

HNSW (Hierarchical Navigable Small World) Index

Core Concepts

HNSW creates a layered graph structure where: * Each layer is a navigable small world graph * Upper layers are sparse for quick approximate search * Lower layers are dense for refinement * Connections are based on proximity in vector space

Mathematical Foundation

The probability of connection between nodes at layer l:

p_l(x) = exp(-\frac{x}{r_l})

where:

x is the distance between nodes

r_l is the layer-specific distance threshold

Implementation Details

-- Basic HNSW index creation
CREATE INDEX hnsw_idx ON vectors
USING hnsw (embedding vector_l2_ops)
WITH (
    dim = 384,              -- Vector dimension
    m = 16,                 -- Max connections per node
    ef_construction = 64,   -- Construction time/quality trade-off
    ef_search = 40         -- Search time/quality trade-off
);

Key Parameters

M (Maximum Connections)

-- Higher M for better accuracy, more memory
CREATE INDEX hnsw_accurate_idx ON vectors
USING hnsw (embedding vector_l2_ops)
WITH (m = 32, ef_construction = 100);

-- Lower M for faster construction, less memory
CREATE INDEX hnsw_fast_idx ON vectors
USING hnsw (embedding vector_l2_ops)
WITH (m = 8, ef_construction = 40);

ef_construction

Controls index build quality:

-- Higher quality, slower build
WITH (ef_construction = 200)

-- Faster build, lower quality
WITH (ef_construction = 40)

ef_search

Set at query time:

-- Set search quality parameter
SET hnsw.ef_search = 100;

-- Perform search
SELECT * FROM vectors
ORDER BY embedding <-> query_vector
LIMIT 10;

Multi-Level Structure

HNSW creates multiple layers:

def get_max_level(M):
    return int(log(n) / log(M))

Search Algorithm

Enter at top layer

Greedy search to closest neighbor

Descend to lower layer

Repeat until base layer

Perform beam search in base layer

-- Example of search with custom parameters
BEGIN;
SET LOCAL hnsw.ef_search = 100;
SELECT * FROM vectors
WHERE embedding <-> query_vector < 0.5
ORDER BY embedding <-> query_vector
LIMIT 10;
COMMIT;

Memory Management

Memory consumption formula:

M_{total} = N * (M * 8 + 4 * d)

where:

N is number of vectors

M is max connections

d is vector dimension

IVF (Inverted File) Index

Structure

Clusters vectors into centroids

Maintains inverted lists

Searches relevant clusters only

CREATE INDEX ivf_idx ON vectors
USING ivfflat (embedding vector_l2_ops)
WITH (lists = 100);

Parameter Tuning

-- More lists for larger datasets
CREATE INDEX ivf_large_idx ON vectors
USING ivfflat (embedding vector_l2_ops)
WITH (lists = 1000);

-- Fewer lists for smaller datasets
CREATE INDEX ivf_small_idx ON vectors
USING ivfflat (embedding vector_l2_ops)
WITH (lists = 50);

Search Process

Find nearest centroids

Search within selected lists

Merge results

Flat Index

Implementation

Simple but exact search:

CREATE INDEX flat_idx ON vectors
USING btree (embedding vector_l2_ops);

Use Cases

Small datasets

Exact search requirements

Development/testing

Performance Comparison

Query Time vs Accuracy

-- Measure HNSW performance
EXPLAIN ANALYZE
SELECT * FROM vectors
ORDER BY embedding <-> query_vector
LIMIT 10;

-- Measure IVF performance
EXPLAIN ANALYZE
SELECT * FROM vectors
ORDER BY embedding <-> query_vector
LIMIT 10;

-- Measure Flat performance
EXPLAIN ANALYZE
SELECT * FROM vectors
ORDER BY embedding <-> query_vector
LIMIT 10;

Memory Usage

Approximate memory usage per vector:

HNSW: (M * 8 + 4 * dim) bytes

IVF: (4 * dim + 4) bytes

Flat: (4 * dim) bytes

Advanced HNSW Optimization

1. Layer Selection

CREATE OR REPLACE FUNCTION optimize_hnsw_layers(
    vector_count integer,
    m integer
) RETURNS integer AS $$
BEGIN
    RETURN floor(ln(vector_count) / ln(m));
END;
$$ LANGUAGE plpgsql;

2. Dynamic Parameter Adjustment

CREATE OR REPLACE FUNCTION adjust_ef_search(
    desired_accuracy float,
    current_ef integer
) RETURNS integer AS $$
BEGIN
    -- Adjust ef_search based on accuracy measurements
    RETURN CASE
        WHEN desired_accuracy > 0.95 THEN current_ef * 2
        WHEN desired_accuracy > 0.90 THEN current_ef * 1.5
        WHEN desired_accuracy < 0.80 THEN current_ef / 1.5
        ELSE current_ef
    END;
END;
$$ LANGUAGE plpgsql;

3. Concurrent Build Optimization

-- Parallel index build
CREATE INDEX CONCURRENTLY hnsw_idx ON vectors
USING hnsw (embedding vector_l2_ops)
WITH (dim = 384, m = 16);

Monitoring and Maintenance

1. Index Statistics

CREATE VIEW index_stats AS
SELECT
    schemaname,
    tablename,
    indexname,
    idx_scan,
    idx_tup_read,
    idx_tup_fetch
FROM pg_stat_user_indexes;

2. Performance Tracking

CREATE TABLE search_metrics (
    query_id serial PRIMARY KEY,
    index_type text,
    query_time interval,
    result_count integer,
    accuracy float
);

3. Index Health Check

CREATE OR REPLACE FUNCTION check_index_health(
    index_name text
) RETURNS table (
    metric text,
    value numeric
) AS $$
BEGIN
    -- Implementation details
END;
$$ LANGUAGE plpgsql;

Best Practices

Index Selection Guidelines

Use HNSW for general-purpose similarity search

Use IVF for larger datasets with relaxed accuracy requirements

Use Flat index for small datasets or when exact results are required

Parameter Selection

HNSW: m = 16 for balanced performance

HNSW: ef_construction = 64-200 based on build time tolerance

IVF: lists = sqrt(N) where N is number of vectors

Resource Management

Monitor memory usage

Use parallel build for large indexes

Implement regular maintenance

Summary

This section provides a deep dive into vector indexing strategies within the pgvector extension for PostgreSQL, focusing especially on the Hierarchical Navigable Small World (HNSW) index structure. HNSW, designed for approximate nearest neighbor searches, uses a multi-layered graph to strike a balance between search speed, accuracy, and memory usage – parameters that can be tuned to fit specific requirements. The section also discusses IVF (Inverted File Vector) and Flat indexing strategies, offering a broad view of the available options for various use cases.

The section goes beyond theory to provide practical guidance, including code examples and monitoring strategies for implementing and maintaining vector indexes in production environments. Additionally, by explaining the underlying mathematical principles and optimization strategies, the section empowers readers to understand and fine-tune index performance to match their specific needs.

Section 5: Designing a Schema for Arxiv Paper Management

Now we look at a design for a comprehensive database schema for storing Arxiv papers, including tables for metadata, full text, embeddings, citations, and authors.

This section will explain the relationships between tables and the rationale behind the design choices.We spend “just enough” time on the rationale for such a design, not more as this is not a theoretical database design book. The section will cover both standard PostgreSQL features and PgVector-specific considerations.

Here’s a breakdown with explanations interspersed, along with some minor suggestions and considerations:

Core Schema Definition

Papers Table

CREATE TABLE papers (
    paper_id TEXT PRIMARY KEY,  -- arxiv id
    title TEXT NOT NULL,
    abstract TEXT,
    published_date DATE NOT NULL,
    updated_date DATE,
    paper_url TEXT,
    pdf_url TEXT,
    category_primary TEXT,
    categories TEXT[],
    download_date TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    pdf_processed BOOLEAN DEFAULT FALSE,
    embedding vector(768),  -- Paper-level embedding
    CHECK (paper_id ~* '^[\d.]+$')  -- Validate arxiv id format
);

paper_id TEXT PRIMARY KEY: Using the arXiv ID as the primary key is a good choice, ensuring uniqueness. The TEXT type is appropriate.

title TEXT NOT NULL: Titles are essential.

abstract TEXT: Storing the abstract allows for searching and summarization.

published_date DATE NOT NULL: Important for tracking the age of the paper. DATE is sufficient unless you need time precision within the day.

updated_date DATE: Useful for tracking revisions.

paper_url TEXT: URL to the arXiv page.

pdf_url TEXT: Direct link to the PDF.

category_primary TEXT: The main category of the paper. Good for categorization and filtering.

categories TEXT[]: An array of categories allows for papers to belong to multiple categories. This is a good use of PostgreSQL’s array type.

download_date TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP: Records when the paper was downloaded into the system. Using WITH TIME ZONE is best practice for timestamps.

pdf_processed BOOLEAN DEFAULT FALSE: A flag to track whether the PDF has been processed (text extracted, embeddings generated, etc.). Very useful for workflow management.

embedding vector(768): The core paper-level embedding. The dimension (768) suggests you’re likely using a model like Sentence-BERT or a similar transformer-based model.

CHECK (paper_id ~* ^[\d.]+$): This regular expression check enforces the typical arXiv ID format (e.g., 2310.12345 or 1706.03762). ~* is a case-insensitive regular expression match.

CREATE INDEX ON papers USING hnsw (embedding vector_cosine_ops)
WITH (dim=768, m=16, ef_construction=64);

Creates HNSW index on the papers embedding. This index is used for vector similarity search.

It uses the vector_cosine_ops operator class, meaning it measures similarity as cosine similarity.

WITH clause specifies the index parameters for HNSW: dim=768: The dimensionality of the embeddings (must match the embedding column’s dimension). m=16: The number of bidirectional links created for each element during index construction. Higher values lead to better recall but slower construction and larger index size. ef_construction=64: Controls the size of the dynamic candidate list during index construction. Higher values improve accuracy at the cost of build time.

PostgresSQL: allows only one vector index per column - if you must have 2 you will need to add a second duplicate column and use a different index on that.

Authors Table

CREATE TABLE authors (
    author_id SERIAL PRIMARY KEY,
    name TEXT NOT NULL,
    email TEXT,
    affiliation TEXT,
    embedding vector(384),  -- Author embedding based on their papers
    UNIQUE (name, email)
);

author_id SERIAL PRIMARY KEY: A standard auto-incrementing integer primary key.

name TEXT NOT NULL: The author’s name.

email TEXT: The author’s email address (optional, may not always be available).

affiliation TEXT: The author’s institutional affiliation.

embedding vector(384): An embedding representing the author’s research interests, likely derived from the embeddings of their papers. A smaller dimension (384) compared to the paper embedding is a reasonable choice, as author representations might need less detail.

UNIQUE (name, email): This constraint helps prevent duplicate author entries. The combination of name and email is likely to be unique (though not perfectly guaranteed). Consider how you’ll handle cases where the email is unknown or multiple authors have the same name and no email. Since email may not be available it may be passed in as NULL, make sure that an empty string is passed in instead or use a unique hash of (name+email) for consistent deduping.

CREATE INDEX ON authors USING hnsw (embedding vector_cosine_ops)
WITH (dim=384, m=16, ef_construction=64);

Similar to the papers table, this creates an HNSW index for fast similarity searches on author embeddings. The parameters are adjusted for the 384-dimensional embeddings.

Paper_Authors Junction Table

CREATE TABLE paper_authors (
    paper_id TEXT REFERENCES papers(paper_id),
    author_id INTEGER REFERENCES authors(author_id),
    author_position INTEGER NOT NULL,
    is_corresponding BOOLEAN DEFAULT FALSE,
    PRIMARY KEY (paper_id, author_id)
);

paper_id TEXT REFERENCES papers(paper_id): Foreign key linking to the papers table.

author_id INTEGER REFERENCES authors(author_id): Foreign key linking to the authors table.

author_position INTEGER NOT NULL: Indicates the order of authors (e.g., 1 for the first author, 2 for the second, etc.). This is important for academic publications.

is_corresponding BOOLEAN DEFAULT FALSE: Indicates whether the author is the corresponding author for the paper.

PRIMARY KEY (paper_id, author_id): A composite primary key ensures that the same author cannot be listed twice for the same paper.

CREATE INDEX ON paper_authors(author_id);
CREATE INDEX ON paper_authors(paper_id, author_position);

These indexes speed up queries that involve filtering or joining based on author_id or paper_id and author_position. The second index is particularly useful for retrieving authors in their correct order for a given paper.

Content Management Tables

Sections Table

CREATE TABLE paper_sections (
    section_id SERIAL PRIMARY KEY,
    paper_id TEXT REFERENCES papers(paper_id),
    section_type TEXT NOT NULL,  -- abstract, introduction, methods, etc.
    section_number INTEGER,
    title TEXT,
    content TEXT,
    embedding vector(768),  -- Section-level embedding
    UNIQUE (paper_id, section_type, section_number)
);

section_id SERIAL PRIMARY KEY: Standard auto-incrementing ID.

paper_id TEXT REFERENCES papers(paper_id): Foreign key to the papers table.

section_type TEXT NOT NULL: Crucially, this categorizes the section (e.g., “introduction”, “methods”, “results”, “conclusion”, “related work”). You might consider using an ENUM type here to enforce a controlled vocabulary of section types.

section_number INTEGER: The numerical order of the section within its type.

title TEXT: The title of the section (if applicable).

content TEXT: The full text content of the section.

embedding vector(768): An embedding for the section’s content, allowing for semantic search within a paper.

UNIQUE (paper_id, section_type, section_number): This ensures logical uniqueness of sections.

CREATE INDEX ON paper_sections USING hnsw (embedding vector_cosine_ops)
WITH (dim=768, m=16, ef_construction=64);

HNSW index for similarity search on section embeddings.

Images Table

CREATE TABLE paper_images (
    image_id SERIAL PRIMARY KEY,
    paper_id TEXT REFERENCES papers(paper_id),
    section_id INTEGER REFERENCES paper_sections(section_id),
    figure_number TEXT,
    caption TEXT,
    image_data BYTEA,
    embedding vector(512),  -- Image embedding from vision model
    metadata JSONB,  -- Store additional image metadata
    position_in_text INTEGER,
    UNIQUE (paper_id, figure_number)
);

image_id SERIAL PRIMARY KEY: Standard auto-incrementing ID.

paper_id TEXT REFERENCES papers(paper_id): Foreign key to the papers table.

section_id INTEGER REFERENCES paper_sections(section_id): Foreign key to the paper_sections table, linking the image to its containing section. This is good for context.

figure_number TEXT: The figure number (e.g., “Figure 1”, “Fig. 2a”). TEXT is appropriate because figure numbers can have letters.

caption TEXT: The image caption.

image_data BYTEA: Stores the raw image data itself. BYTEA is PostgreSQL’s type for binary data. Consider whether you need to store the image data directly in the database. Alternatives include storing a path to the image file on disk (if you have a file system available).

embedding vector(512): An image embedding generated by a vision model (e.g., a CNN). 512 is a common dimension for image embeddings.

metadata JSONB: A flexible way to store any additional image metadata (e.g., image resolution, detected objects, etc.). JSONB is generally preferred over JSON for its indexing and query capabilities.

position_in_text INTEGER: Represents where the image appears within the text, which can help maintain context.

UNIQUE (paper_id, figure_number): Prevents duplicate figure numbers within a paper.

CREATE INDEX ON paper_images USING hnsw (embedding vector_cosine_ops)
WITH (dim=512, m=16, ef_construction=64);

HNSW index for similarity search on image embeddings.

Formulas Table

CREATE TABLE paper_formulas (
    formula_id SERIAL PRIMARY KEY,
    paper_id TEXT REFERENCES papers(paper_id),
    section_id INTEGER REFERENCES paper_sections(section_id),
    equation_number TEXT,
    latex_source TEXT NOT NULL,
    rendered_image BYTEA,
    embedding vector(128),  -- Formula embedding
    context TEXT,  -- Surrounding text
    position_in_text INTEGER,
    UNIQUE (paper_id, equation_number)
);

formula_id SERIAL PRIMARY KEY: Standard auto-incrementing ID.

paper_id TEXT REFERENCES papers(paper_id): Foreign key to the papers table.

section_id INTEGER REFERENCES paper_sections(section_id): Links the formula to the section it’s in.

equation_number TEXT: The equation number (similar to figure numbers).

latex_source TEXT NOT NULL: The LaTeX source code of the formula. This is excellent, as it allows for rendering and potential search/manipulation.

rendered_image BYTEA: A rendered image of the formula (optional, but useful for display).

embedding vector(128): An embedding representing the formula. 128 dimensions is probably sufficient. There are specialized models for embedding mathematical expressions.

context TEXT: The text surrounding the formula, providing crucial context for understanding its meaning.

position_in_text INTEGER: Similar to position_in_text for images, it maintains context.

UNIQUE (paper_id, equation_number): Enforces uniqueness of equation numbers within each paper.

CREATE INDEX ON paper_formulas USING hnsw (embedding vector_cosine_ops)
WITH (dim=128, m=16, ef_construction=64);

HNSW index for formula embeddings.

Tables Content

CREATE TABLE paper_tables (
    table_id SERIAL PRIMARY KEY,
    paper_id TEXT REFERENCES papers(paper_id),
    section_id INTEGER REFERENCES paper_sections(section_id),
    table_number TEXT,
    caption TEXT,
    content JSONB,  -- Structured table data
    embedding vector(384),  -- Table content embedding
    position_in_text INTEGER,
    UNIQUE (paper_id, table_number)
);

table_id SERIAL PRIMARY KEY: Standard auto-incrementing ID.

paper_id TEXT REFERENCES papers(paper_id): Foreign key.

section_id INTEGER REFERENCES paper_sections(section_id): Links to the section.

table_number TEXT: Table number (e.g., “Table 1”).

caption TEXT: Table caption.

content JSONB: Stores the table data in a structured JSONB format. This is a good approach, as it allows you to represent the rows and columns of the table. You’ll need to define a consistent structure for this JSONB data (e.g., an array of rows, where each row is an object with column names as keys).

embedding vector(384): An embedding representing the table’s content.

position_in_text INTEGER: The position of the table in the text.

UNIQUE (paper_id, table_number): Enforces unique table numbers within a paper.

CREATE INDEX ON paper_tables USING hnsw (embedding vector_cosine_ops)
WITH (dim=384, m=16, ef_construction=64);

HNSW index for table embeddings.

Reference Management

Citations Table

CREATE TABLE citations (
    citation_id SERIAL PRIMARY KEY,
    citing_paper_id TEXT REFERENCES papers(paper_id),
    cited_paper_id TEXT,  -- May not be in our database
    citation_context TEXT,
    section_id INTEGER REFERENCES paper_sections(section_id),
    position_in_text INTEGER,
    confidence FLOAT,  -- Confidence in citation extraction
    CHECK (citing_paper_id != cited_paper_id)  -- Prevent self-citation loops
);

citation_id SERIAL PRIMARY KEY: Standard ID.

citing_paper_id TEXT REFERENCES papers(paper_id): The paper that is doing the citing.

cited_paper_id TEXT: The paper being cited. This is TEXT and not a foreign key because the cited paper might not be in your database. You could potentially link this to papers.paper_id after you’ve downloaded and processed the cited paper.

citation_context TEXT: The text surrounding the citation, providing context. Very important for understanding the nature of the citation.

section_id INTEGER REFERENCES paper_sections(section_id): Links the citation to the specific section where it occurs.

position_in_text INTEGER: The position of the citation.

confidence FLOAT: A measure of confidence in the citation extraction process (e.g., from a machine learning model).

CHECK (citing_paper_id != cited_paper_id): Prevents a paper from citing itself (which is logically impossible).

CREATE INDEX ON citations(citing_paper_id);
CREATE INDEX ON citations(cited_paper_id);

Indexes for efficient lookup of citations based on citing and cited papers.

Bibliography Table

CREATE TABLE bibliography (
    bib_id SERIAL PRIMARY KEY,
    paper_id TEXT REFERENCES papers(paper_id),
    reference_text TEXT NOT NULL,
    parsed_data JSONB,  -- Structured reference data
    doi TEXT,
    matched_paper_id TEXT REFERENCES papers(paper_id),
    UNIQUE (paper_id, reference_text)
);

bib_id SERIAL PRIMARY KEY: Standard ID.

paper_id TEXT REFERENCES papers(paper_id): The paper containing the bibliography entry.

reference_text TEXT NOT NULL: The raw, unstructured text of the bibliography entry.

parsed_data JSONB: Structured data extracted from the reference_text (e.g., title, authors, journal, year). This is very useful for analysis and linking.

doi TEXT: The Digital Object Identifier (DOI) of the cited paper, if available. DOIs are very useful for resolving references.

matched_paper_id TEXT REFERENCES papers(paper_id): If the bibliography entry can be matched to a paper already in your database, this foreign key establishes the

Summary

This is a comprehensive schema for managing arXiv papers, their content, and associated metadata. The inclusion of embeddings at multiple levels (paper, section, image, formula, author, category) is a powerful feature for similarity searching and recommendation. The use of pgvector for vector storage and similarity search is well-integrated. The schema is also normalized and has been carefully designed

Section 6: Building an Arxiv Paper Collection Pipeline

Now we get our hands dirty with some Python coding developing a pipeline for retrieving papers from Arxiv’s API, focusing on the topic of LLM research. This is of course easily changeable via a parameter in the query. To be able to actually build something minimally practical we will cover scheduling daily downloads, handling rate limits, extracting metadata, and error handling. The code in this section will demonstrate how to track papers across multiple collection runs and how to store the full PDFs in PostgreSQL.

The code should be self explanatory assuming an intermediate knowledge of Python and SQL

import arxiv
import asyncio
import aiohttp
import logging
import psycopg2
from psycopg2.extras import execute_values
from datetime import datetime, timedelta
import PyPDF2
from typing import List, Dict, Any, Optional
from sentence_transformers import SentenceTransformer
import numpy as np
import json
import backoff
import fitz  # PyMuPDF
from dataclasses import dataclass
from concurrent.futures import ThreadPoolExecutor
import configparser
import os

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

@dataclass
class PaperMetadata:
    """Data class for paper metadata"""
    paper_id: str
    title: str
    abstract: str
    authors: List[Dict[str, str]]
    categories: List[str]
    submission_date: datetime
    last_updated_date: datetime
    pdf_url: str

class ArxivPipeline:
    def __init__(self, db_config: Dict[str, str], model_name: str = 'all-mpnet-base-v2'):
        """Initialize the pipeline with database configuration and embedding model."""
        self.db_config = db_config
        self.embedding_model = SentenceTransformer(model_name)
        self.conn = self._create_db_connection()
        self.session = None

    def _create_db_connection(self) -> psycopg2.extensions.connection:
        """Create and return a database connection."""
        try:
            conn = psycopg2.connect(**self.db_config)
            logger.info("Successfully connected to database")
            return conn
        except Exception as e:
            logger.error(f"Failed to connect to database: {e}")
            raise

    @backoff.on_exception(backoff.expo, arxiv.ArxivError, max_tries=5)
    async def fetch_papers(self, query: str, max_results: int = 100) -> List[PaperMetadata]:
        """Fetch papers from arXiv API with exponential backoff."""
        try:
            client = arxiv.Client()
            search = arxiv.Search(
                query=query,
                max_results=max_results,
                sort_by=arxiv.SortCriterion.SubmittedDate
            )

            papers = []
            async for result in client.results(search):
                paper = PaperMetadata(
                    paper_id=result.entry_id.split('/')[-1],
                    title=result.title,
                    abstract=result.summary,
                    authors=[{"name": author.name, "affiliation": author.affiliation}
                            for author in result.authors],
                    categories=result.categories,
                    submission_date=result.published,
                    last_updated_date=result.updated,
                    pdf_url=result.pdf_url
                )
                papers.append(paper)

            logger.info(f"Retrieved {len(papers)} papers from arXiv")
            return papers

        except arxiv.ArxivError as e:
            logger.error(f"Error fetching papers from arXiv: {e}")
            raise

    async def download_pdf(self, paper: PaperMetadata) -> Optional[bytes]:
        """Download PDF with rate limiting and retry logic."""
        async with aiohttp.ClientSession() as session:
            try:
                async with session.get(paper.pdf_url) as response:
                    if response.status == 200:
                        return await response.read()
                    else:
                        logger.error(f"Failed to download PDF for {paper.paper_id}: {response.status}")
                        return None
            except Exception as e:
                logger.error(f"Error downloading PDF for {paper.paper_id}: {e}")
                return None

    def process_pdf(self, pdf_data: bytes) -> Dict[str, Any]:
        """Extract content and structure from PDF."""
        try:
            doc = fitz.open(stream=pdf_data, filetype="pdf")

            # Initialize containers for different content types
            sections = []
            images = []
            formulas = []
            tables = []

            for page_num in range(len(doc)):
                page = doc[page_num]

                # Extract text and structure
                blocks = page.get_text("dict")["blocks"]
                for block in blocks:
                    # Process text blocks
                    if block["type"] == 0:  # Text
                        sections.append({
                            "content": block["text"],
                            "page_number": page_num + 1,
                            "bbox": block["bbox"]
                        })

                    # Extract images
                    image_list = page.get_images()
                    for img_idx, img in enumerate(image_list):
                        try:
                            xref = img[0]
                            image = doc.extract_image(xref)
                            images.append({
                                "data": image["image"],
                                "page_number": page_num + 1,
                                "format": image["ext"],
                                "width": image["width"],
                                "height": image["height"]
                            })
                        except Exception as e:
                            logger.warning(f"Failed to extract image: {e}")

            return {
                "sections": sections,
                "images": images,
                "formulas": formulas,  # Would need specialized formula extraction
                "tables": tables      # Would need specialized table extraction
            }

        except Exception as e:
            logger.error(f"Error processing PDF: {e}")
            raise

    def generate_embeddings(self, text: str) -> np.ndarray:
        """Generate embeddings for text content."""
        try:
            return self.embedding_model.encode(text)
        except Exception as e:
            logger.error(f"Error generating embeddings: {e}")
            raise

    async def store_paper(self, paper: PaperMetadata, content: Dict[str, Any]) -> None:
        """Store paper metadata and content in database."""
        try:
            # Generate abstract embedding
            abstract_embedding = self.generate_embeddings(paper.abstract)

            # Store paper metadata
            with self.conn.cursor() as cur:
                # Insert paper
                cur.execute("""
                    INSERT INTO papers (
                        paper_id, title, abstract, abstract_embedding,
                        submission_date, last_updated_date, categories
                    ) VALUES (%s, %s, %s, %s, %s, %s, %s)
                    ON CONFLICT (paper_id) DO UPDATE SET
                        last_updated_date = EXCLUDED.last_updated_date
                    RETURNING paper_id
                """, (
                    paper.paper_id, paper.title, paper.abstract, abstract_embedding.tolist(),
                    paper.submission_date, paper.last_updated_date, paper.categories
                ))

                # Insert authors
                for author in paper.authors:
                    cur.execute("""
                        INSERT INTO authors (name, affiliation)
                        VALUES (%s, %s)
                        ON CONFLICT (name) DO UPDATE SET
                            affiliation = EXCLUDED.affiliation
                        RETURNING author_id
                    """, (author["name"], author.get("affiliation")))

                    author_id = cur.fetchone()[0]

                    # Link author to paper
                    cur.execute("""
                        INSERT INTO paper_authors (paper_id, author_id, author_position)
                        VALUES (%s, %s, %s)
                        ON CONFLICT DO NOTHING
                    """, (paper.paper_id, author_id, paper.authors.index(author)))

                # Store sections with embeddings
                for section in content["sections"]:
                    section_embedding = self.generate_embeddings(section["content"])
                    cur.execute("""
                        INSERT INTO paper_sections (
                            paper_id, content, content_embedding,
                            page_number, sequence_order
                        ) VALUES (%s, %s, %s, %s, %s)
                    """, (
                        paper.paper_id, section["content"], section_embedding.tolist(),
                        section["page_number"], content["sections"].index(section)
                    ))

                # Store images
                for image in content["images"]:
                    cur.execute("""
                        INSERT INTO paper_images (
                            paper_id, image_data, page_number,
                            width, height, format
                        ) VALUES (%s, %s, %s, %s, %s, %s)
                    """, (
                        paper.paper_id, image["data"], image["page_number"],
                        image["width"], image["height"], image["format"]
                    ))

            self.conn.commit()
            logger.info(f"Successfully stored paper {paper.paper_id}")

        except Exception as e:
            self.conn.rollback()
            logger.error(f"Error storing paper {paper.paper_id}: {e}")
            raise

    async def run_pipeline(self, query: str, max_results: int = 100) -> None:
        """Run the complete pipeline."""
        try:
            # Fetch papers
            papers = await self.fetch_papers(query, max_results)

            # Process each paper
            for paper in papers:
                # Download PDF
                pdf_data = await self.download_pdf(paper)
                if pdf_data is None:
                    continue

                # Process PDF content
                content = self.process_pdf(pdf_data)

                # Store paper and content
                await self.store_paper(paper, content)

                # Rate limiting
                await asyncio.sleep(1)

        except Exception as e:
            logger.error(f"Pipeline error: {e}")
            raise
        finally:
            if self.conn:
                self.conn.close()

def main():
    # Load configuration
    config = configparser.ConfigParser()
    config.read('config.ini')

    db_config = {
        'dbname': config['Database']['dbname'],
        'user': config['Database']['user'],
        'password': config['Database']['password'],
        'host': config['Database']['host'],
        'port': config['Database']['port']
    }

    # Initialize and run pipeline
    pipeline = ArxivPipeline(db_config)

    # Example query for LLM papers
    query = 'cat:cs.CL AND (abs:"large language models" OR abs:LLM)'

    asyncio.run(pipeline.run_pipeline(query, max_results=100))

if __name__ == "__main__":
    main()

Summary

In this section we set up an end to end pipeline for fetching PDFs from ArXiv and extracting text and images from them which are saved in our pgvector based vector database after generating embeddings In the fully featured app we want to extract the formulas and tables also, but that we will do in the next chapter.

Section 7: Rich Content Extraction from PDFs and Embedding Generation

The big challenge with PDF documents is in extracting different media from PDF documents using appropriate libraries and preprocessing them for embedding generation. The section will cover sentence/paragraph segmentation strategies and discuss which SentenceTransformer models work best for scientific papers. It also includes ways to extract formulae, and tables and use CLIP based embedding generation.

Note that this section only focuses on the content and embedding creation, it does not include the pipeline for fetching data.

import fitz  # PyMuPDF
import numpy as np
import torch
from PIL import Image
import io
import re
import json
from typing import List, Dict, Tuple, Optional, Any
from dataclasses import dataclass
from sentence_transformers import SentenceTransformer
from transformers import CLIPProcessor, CLIPModel
import camelot
import pdfplumber
from latex2mathml.converter import convert as latex2mathml
import logging
from collections import defaultdict

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class ExtractedSection:
    """Represents a section of text from the paper"""
    heading: str
    content: str
    page_number: int
    bbox: Tuple[float, float, float, float]
    embedding: Optional[np.ndarray] = None

@dataclass
class ExtractedImage:
    """Represents an extracted image with its metadata"""
    image_data: bytes
    caption: str
    page_number: int
    bbox: Tuple[float, float, float, float]
    reference_text: str
    embedding: Optional[np.ndarray] = None

@dataclass
class ExtractedFormula:
    """Represents a mathematical formula"""
    latex_source: str
    mathml_source: str
    is_inline: bool
    page_number: int
    bbox: Tuple[float, float, float, float]
    reference_text: str
    embedding: Optional[np.ndarray] = None

@dataclass
class ExtractedTable:
    """Represents a table from the paper"""
    table_data: List[List[str]]
    caption: str
    page_number: int
    bbox: Tuple[float, float, float, float]
    reference_text: str
    embedding: Optional[np.ndarray] = None

class ContentProcessor:
    def __init__(self):
        """Initialize content processor with necessary models"""
        # Text embedding model
        self.text_embedder = SentenceTransformer('all-mpnet-base-v2')

        # Image embedding model (CLIP)
        self.image_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
        self.image_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")

        # Regular expressions for content detection
        self.section_pattern = re.compile(r'^(?:\d+\.)?\s*([A-Z][^\n]+)$', re.MULTILINE)
        self.formula_pattern = re.compile(r'\$(.*?)\$|\\\[(.*?)\\\]|\\\((.*?)\\\)', re.DOTALL)
        self.reference_pattern = re.compile(r'(Fig\.|Figure|Table|Equation)\s+(\d+[a-z]?)')

    def extract_sections(self, doc: fitz.Document) -> List[ExtractedSection]:
        """Extract sections with hierarchical structure"""
        sections = []
        current_section = None
        current_content = []

        for page_num in range(len(doc)):
            page = doc[page_num]
            blocks = page.get_text("dict")["blocks"]

            for block in blocks:
                if block["type"] == 0:  # Text block
                    text = block["text"].strip()
                    if text:
                        # Check if this is a section heading
                        if self.section_pattern.match(text):
                            # Save previous section if exists
                            if current_section and current_content:
                                sections.append(ExtractedSection(
                                    heading=current_section,
                                    content="\n".join(current_content),
                                    page_number=page_num,
                                    bbox=block["bbox"]
                                ))
                            current_section = text
                            current_content = []
                        else:
                            current_content.append(text)

        # Add last section
        if current_section and current_content:
            sections.append(ExtractedSection(
                heading=current_section,
                content="\n".join(current_content),
                page_number=len(doc) - 1,
                bbox=(0, 0, 0, 0)  # Placeholder bbox
            ))

        return sections

    def extract_images(self, doc: fitz.Document) -> List[ExtractedImage]:
        """Extract images with captions and references"""
        images = []

        for page_num in range(len(doc)):
            page = doc[page_num]

            # Get image locations and captions
            image_list = page.get_images()
            blocks = page.get_text("dict")["blocks"]

            # Find captions (text blocks near images)
            captions = []
            for block in blocks:
                if block["type"] == 0:  # Text block
                    text = block["text"].strip()
                    if text.startswith("Figure") or text.startswith("Fig."):
                        captions.append((text, block["bbox"]))

            # Extract images and match with captions
            for img_idx, img in enumerate(image_list):
                try:
                    xref = img[0]
                    base_image = doc.extract_image(xref)

                    # Find nearest caption
                    nearest_caption = min(captions,
                                       key=lambda x: abs(x[1][1] - img[1]),  # Compare y-coordinates
                                       default=("", (0,0,0,0)))

                    images.append(ExtractedImage(
                        image_data=base_image["image"],
                        caption=nearest_caption[0],
                        page_number=page_num,
                        bbox=img[1],  # bbox from image list
                        reference_text=self._find_references(blocks, f"Fig. {img_idx + 1}")
                    ))

                except Exception as e:
                    logger.warning(f"Failed to extract image: {e}")

        return images

    def extract_formulas(self, doc: fitz.Document) -> List[ExtractedFormula]:
        """Extract mathematical formulas"""
        formulas = []

        for page_num in range(len(doc)):
            page = doc[page_num]
            text = page.get_text()

            # Find all LaTeX formulas
            matches = self.formula_pattern.finditer(text)
            for match in matches:
                latex = match.group(1) or match.group(2) or match.group(3)
                try:
                    mathml = latex2mathml(latex)

                    # Determine if formula is inline
                    is_inline = bool(match.group(1) or match.group(3))

                    # Get approximate bbox
                    start_pos = match.start()
                    end_pos = match.end()
                    start_bbox = page.get_text("dict", clip=(0, 0, page.rect.width, start_pos))
                    end_bbox = page.get_text("dict", clip=(0, 0, page.rect.width, end_pos))

                    formulas.append(ExtractedFormula(
                        latex_source=latex,
                        mathml_source=mathml,
                        is_inline=is_inline,
                        page_number=page_num,
                        bbox=(start_bbox["bbox"], end_bbox["bbox"]),
                        reference_text=self._find_references(page.get_text("dict")["blocks"], "Equation")
                    ))

                except Exception as e:
                    logger.warning(f"Failed to process formula: {e}")

        return formulas

    def extract_tables(self, pdf_path: str) -> List[ExtractedTable]:
        """Extract tables using multiple methods"""
        tables = []

        # Try Camelot first
        try:
            camelot_tables = camelot.read_pdf(pdf_path, pages='all')
            for table in camelot_tables:
                tables.append(ExtractedTable(
                    table_data=table.data,
                    caption=self._find_table_caption(table),
                    page_number=table.page,
                    bbox=table._bbox,
                    reference_text=""  # Will be filled later
                ))
        except Exception as e:
            logger.warning(f"Camelot extraction failed: {e}")

        # Fallback to pdfplumber if needed
        if not tables:
            try:
                with pdfplumber.open(pdf_path) as pdf:
                    for page_num, page in enumerate(pdf.pages):
                        for table in page.extract_tables():
                            tables.append(ExtractedTable(
                                table_data=table,
                                caption="",  # Need to find caption separately
                                page_number=page_num + 1,
                                bbox=page.bbox,
                                reference_text=""
                            ))
            except Exception as e:
                logger.warning(f"pdfplumber extraction failed: {e}")

        return tables

    def generate_embeddings(self, content: Dict[str, Any]) -> Dict[str, Any]:
        """Generate embeddings for all content types"""
        # Process sections
        for section in content['sections']:
            section.embedding = self.text_embedder.encode(
                section.heading + " " + section.content
            )

        # Process images
        for image in content['images']:
            try:
                # Convert bytes to PIL Image
                img = Image.open(io.BytesIO(image.image_data))

                # Generate CLIP embedding
                inputs = self.image_processor(images=img, return_tensors="pt")
                image_features = self.image_model.get_image_features(**inputs)
                image.embedding = image_features.detach().numpy()
            except Exception as e:
                logger.warning(f"Failed to generate image embedding: {e}")

        # Process formulas
        for formula in content['formulas']:
            # Generate embedding for formula context
            formula.embedding = self.text_embedder.encode(
                formula.latex_source + " " + formula.reference_text
            )

        # Process tables
        for table in content['tables']:
            # Generate embedding for table content and caption
            table_text = "\n".join([" ".join(row) for row in table.table_data])
            table.embedding = self.text_embedder.encode(
                table.caption + " " + table_text
            )

        return content

    def _find_references(self, blocks: List[Dict], ref_type: str) -> str:
        """Find references to content in text blocks"""
        references = []
        for block in blocks:
            if block["type"] == 0:  # Text block
                text = block["text"]
                matches = self.reference_pattern.finditer(text)
                for match in matches:
                    if match.group(1) == ref_type:
                        context_start = max(0, match.start() - 50)
                        context_end = min(len(text), match.end() + 50)
                        references.append(text[context_start:context_end])
        return "\n".join(references)

    def _find_table_caption(self, table: Any) -> str:
        """Find caption for a table"""
        # Implementation depends on the specific table extraction method
        return ""  # Placeholder

    def process_pdf(self, pdf_path: str) -> Dict[str, Any]:
        """Process PDF and extract all content with embeddings"""
        try:
            # Open PDF
            doc = fitz.open(pdf_path)

            # Extract content
            content = {
                'sections': self.extract_sections(doc),
                'images': self.extract_images(doc),
                'formulas': self.extract_formulas(doc),
                'tables': self.extract_tables(pdf_path)
            }

            # Generate embeddings
            content = self.generate_embeddings(content)

            return content

        except Exception as e:
            logger.error(f"Error processing PDF: {e}")
            raise
        finally:
            if 'doc' in locals():
                doc.close()

def main():
    # Example usage
    processor = ContentProcessor()
    content = processor.process_pdf("example.pdf")

    # Print statistics
    print(f"Extracted {len(content['sections'])} sections")
    print(f"Extracted {len(content['images'])} images")
    print(f"Extracted {len(content['formulas'])} formulas")
    print(f"Extracted {len(content['tables'])} tables")

if __name__ == "__main__":
    main()

Section 8: Implementing Semantic Search Queries with PgVector

Now we come to the core of the chapter where we explore various semantic search patterns using PgVector, from basic similarity queries to hybrid searches combining metadata and vector similarity. It will demonstrate how to implement search relevance ranking, filtering by topic/date/author, and pagination. The section will also include query optimization techniques specific to vector searches in PostgreSQL.

Introduction

In this section, we’ll explore how to implement effective semantic search queries using PostgreSQL with the pgvector extension. We’ll cover everything from basic similarity searches to complex hybrid queries that combine metadata filtering with vector similarity. The examples will focus on our ArXiv paper database, demonstrating practical applications for research paper discovery and analysis.

8.1 Basic Vector Similarity Queries

Let’s start with the fundamental building blocks of semantic search using pgvector. These queries form the basis for more complex search patterns.

8.1.1 Simple Nearest Neighbor Search

-- Basic nearest neighbor search using cosine distance
SELECT
    p.arxiv_id,
    p.title,
    p.abstract,
    ps.embedding <=> query_embedding AS distance
FROM
    papers p
    JOIN paper_sections ps ON p.arxiv_id = ps.arxiv_id
WHERE
    ps.section_type = 'abstract'
ORDER BY
    ps.embedding <=> query_embedding
LIMIT 10;

The <=> operator is pgvector’s cosine distance operator. Lower values indicate higher similarity. This query finds the 10 papers whose abstract embeddings are most similar to our query embedding.

8.1.2 Distance Metrics

PgVector supports three distance metrics: * Cosine distance (default) embedding <⇒ query_embedding

Euclidean distance (L2) embedding <→ query_embedding

Inner product

embedding <#> query_embedding

Choose the appropriate metric based on your embedding model and requirements: * Cosine distance (<=>) is best for normalized embeddings * Euclidean distance (<->) works well for absolute comparisons * Inner product (<#>) is useful for specific similarity measures

8.2 Hybrid Search Queries

Real-world applications often require combining vector similarity with traditional filtering criteria. Filtering by metadata before applying vector search can drastically reduce the size of the set over which vector search (more computationally expensive) will be done. Obviously this will yield noticeably better performance.

8.2.1 Filtering by Metadata

-- Search for similar papers within a specific date range and category
WITH query_embedding AS (
    SELECT embedding::vector
    FROM paper_sections
    WHERE arxiv_id = 'query_paper_id' AND section_type = 'abstract'
)
SELECT
    p.arxiv_id,
    p.title,
    p.publication_date,
    p.categories,
    ps.embedding <=> (SELECT * FROM query_embedding) AS similarity
FROM
    papers p
    JOIN paper_sections ps ON p.arxiv_id = ps.arxiv_id
WHERE
    ps.section_type = 'abstract'
    AND p.publication_date >= '2023-01-01'
    AND p.categories @> ARRAY['cs.AI']  -- Array contains 'cs.AI'
    AND ps.embedding <=> (SELECT * FROM query_embedding) < 0.5
ORDER BY
    similarity
LIMIT 20;

8.2.2 Multi-section Search

Sometimes we want to search across different sections of papers:

-- Search across abstracts and introductions with weighted relevance
WITH query_embedding AS (
    SELECT embedding::vector
    FROM paper_sections
    WHERE arxiv_id = 'query_paper_id' AND section_type = 'abstract'
)
SELECT
    p.arxiv_id,
    p.title,
    CASE
        WHEN ps.section_type = 'abstract' THEN ps.embedding <=> qe.embedding * 0.6
        WHEN ps.section_type = 'introduction' THEN ps.embedding <=> qe.embedding * 0.4
    END AS weighted_similarity
FROM
    papers p
    JOIN paper_sections ps ON p.arxiv_id = ps.arxiv_id
    CROSS JOIN query_embedding qe
WHERE
    ps.section_type IN ('abstract', 'introduction')
GROUP BY
    p.arxiv_id,
    p.title
ORDER BY
    weighted_similarity
LIMIT 10;

8.3 Advanced Search Patterns

8.3.1 Relevance Ranking

Implementing a more sophisticated relevance ranking system:

CREATE OR REPLACE FUNCTION calculate_paper_relevance(
    similarity float,
    publication_date date,
    citation_count int
) RETURNS float AS $$
BEGIN
    RETURN (
        0.6 * (1.0 - similarity) +  -- Vector similarity (60% weight)
        0.2 * (1.0 - EXTRACT(YEAR FROM age(current_date, publication_date))::float / 10.0) +  -- Recency (20% weight)
        0.2 * LEAST(citation_count::float / 1000.0, 1.0)  -- Citation impact (20% weight)
    );
END;
$$ LANGUAGE plpgsql;

-- Using the relevance function in a search query
WITH query_embedding AS (
    SELECT embedding::vector
    FROM paper_sections
    WHERE arxiv_id = 'query_paper_id' AND section_type = 'abstract'
)
SELECT
    p.arxiv_id,
    p.title,
    calculate_paper_relevance(
        ps.embedding <=> qe.embedding,
        p.publication_date,
        (SELECT COUNT(*) FROM citations WHERE cited_paper_id = p.arxiv_id)
    ) AS relevance_score
FROM
    papers p
    JOIN paper_sections ps ON p.arxiv_id = ps.arxiv_id
    CROSS JOIN query_embedding qe
WHERE
    ps.section_type = 'abstract'
ORDER BY
    relevance_score DESC
LIMIT 20;

8.3.2 Semantic Clustering

Finding papers that form semantic clusters:

WITH similar_papers AS (
    SELECT
        p1.arxiv_id as paper1_id,
        p2.arxiv_id as paper2_id,
        ps1.embedding <=> ps2.embedding as similarity
    FROM
        papers p1
        JOIN paper_sections ps1 ON p1.arxiv_id = ps1.arxiv_id
        JOIN papers p2 ON p1.arxiv_id < p2.arxiv_id  -- Avoid self-comparison
        JOIN paper_sections ps2 ON p2.arxiv_id = ps2.arxiv_id
    WHERE
        ps1.section_type = 'abstract'
        AND ps2.section_type = 'abstract'
        AND ps1.embedding <=> ps2.embedding < 0.3  -- Similarity threshold
)
SELECT
    paper1_id,
    ARRAY_AGG(paper2_id) as similar_papers,
    AVG(similarity) as avg_similarity
FROM
    similar_papers
GROUP BY
    paper1_id
HAVING
    COUNT(*) >= 5  -- Minimum cluster size
ORDER BY
    avg_similarity;

8.4 Query Optimization

8.4.1 Using Indexes Effectively

PgVector supports both exact and approximate indexing:

-- Create an IVFFlat index
CREATE INDEX ON paper_sections USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 100);

-- Create an HNSW index
CREATE INDEX ON paper_sections USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

For our ArXiv database, HNSW indexes generally provide better performance for similarity searches. However, they come with higher memory requirements and build time.

8.4.2 Query Planning

To optimize query performance:

Use EXPLAIN ANALYZE to understand query execution

Consider materialized views for common search patterns

Implement appropriate indexes on frequently filtered columns

Use parallel query execution for large datasets

Example of a materialized view for frequent searches:

CREATE MATERIALIZED VIEW paper_search_view AS
SELECT
    p.arxiv_id,
    p.title,
    p.publication_date,
    p.categories,
    ps.embedding,
    (SELECT COUNT(*) FROM citations c WHERE c.cited_paper_id = p.arxiv_id) as citation_count
FROM
    papers p
    JOIN paper_sections ps ON p.arxiv_id = ps.arxiv_id
WHERE
    ps.section_type = 'abstract';

-- Create indexes on the materialized view
CREATE INDEX ON paper_search_view USING hnsw (embedding vector_cosine_ops);
CREATE INDEX ON paper_search_view (publication_date);
CREATE INDEX ON paper_search_view USING GIN (categories);

8.5 Implementing a Search API

Here’s a Python implementation that brings together the concepts we’ve covered:

from typing import List, Dict, Any
import psycopg2
from psycopg2.extras import execute_values
import numpy as np

class ArxivSearchEngine:
    def __init__(self, db_params: Dict[str, str]):
        """Initialize the search engine with database connection parameters."""
        self.conn = psycopg2.connect(**db_params)
        self.cursor = self.conn.cursor()

    def search_papers(
        self,
        query_embedding: np.ndarray,
        categories: List[str] = None,
        date_from: str = None,
        date_to: str = None,
        limit: int = 10
    ) -> List[Dict[str, Any]]:
        """
        Search for papers using vector similarity and optional filters.

        Args:
            query_embedding: Vector embedding of the search query
            categories: List of arXiv categories to filter by
            date_from: Start date for publication date filter
            date_to: End date for publication date filter
            limit: Maximum number of results to return

        Returns:
            List of dictionaries containing paper information and similarity scores
        """
        query = """
            SELECT
                p.arxiv_id,
                p.title,
                p.abstract,
                p.publication_date,
                p.categories,
                ps.embedding <=> %s::vector AS similarity,
                (SELECT COUNT(*) FROM citations c WHERE c.cited_paper_id = p.arxiv_id) as citation_count
            FROM
                papers p
                JOIN paper_sections ps ON p.arxiv_id = ps.arxiv_id
            WHERE
                ps.section_type = 'abstract'
        """

        params = [query_embedding.tolist()]

        if categories:
            query += " AND p.categories && %s"
            params.append(categories)

        if date_from:
            query += " AND p.publication_date >= %s"
            params.append(date_from)

        if date_to:
            query += " AND p.publication_date <= %s"
            params.append(date_to)

        query += """
            ORDER BY similarity
            LIMIT %s
        """
        params.append(limit)

        self.cursor.execute(query, params)
        results = self.cursor.fetchall()

        return [
            {
                'arxiv_id': r[0],
                'title': r[1],
                'abstract': r[2],
                'publication_date': r[3],
                'categories': r[4],
                'similarity': r[5],
                'citation_count': r[6]
            }
            for r in results
        ]

    def find_similar_papers(self, arxiv_id: str, limit: int = 10) -> List[Dict[str, Any]]:
        """Find papers similar to a given paper by its arxiv_id."""
        query = """
            WITH paper_embedding AS (
                SELECT embedding
                FROM paper_sections
                WHERE arxiv_id = %s AND section_type = 'abstract'
            )
            SELECT
                p.arxiv_id,
                p.title,
                p.abstract,
                ps.embedding <=> (SELECT embedding FROM paper_embedding) AS similarity
            FROM
                papers p
                JOIN paper_sections ps ON p.arxiv_id = ps.arxiv_id
            WHERE
                ps.section_type = 'abstract'
                AND p.arxiv_id != %s
            ORDER BY
                similarity
            LIMIT %s
        """

        self.cursor.execute(query, (arxiv_id, arxiv_id, limit))
        results = self.cursor.fetchall()

        return [
            {
                'arxiv_id': r[0],
                'title': r[1],
                'abstract': r[2],
                'similarity': r[3]
            }
            for r in results
        ]

    def close(self):
        """Close database connections."""
        self.cursor.close()
        self.conn.close()

Conclusion

Effective semantic search with pgvector requires careful consideration of: * Distance metrics and their implications * Index types and their trade-offs * Query optimization strategies * Relevance ranking algorithms * Integration with traditional SQL queries

The examples provided here demonstrate how to implement sophisticated search functionality while maintaining good performance. As your dataset grows, pay special attention to index choices and query optimization to ensure responsive search results.

Remember that semantic search is often part of a larger system - consider caching frequently accessed results, implementing rate limiting for API endpoints, and monitoring query performance in production environments.

In the next section, we’ll explore how to extend these search capabilities to handle multimodal content, including images and mathematical formulas extracted from papers.

Section 9: Advanced Analytics: Citation Networks and Author Influence

This section will develop analytics queries for mining the paper collection, including identifying citation networks, finding influential authors, and tracking research trends. It will combine traditional SQL analytics with vector similarity to discover conceptually related papers even when they don’t directly cite each other. The section will include visualizations of research clusters and influence networks. 9.1 Citation Network Analysis While the following code builds a citation graph, please note that we are doing this in a relational database which is not optimal for representing graph structures. So please dont attempt to build a deep graph. Even a depth of 3 might severely degrade query performance while the graph is being built. So we limit the depth in the query to < 3. If you have ample computational resources and are feeling adventurous you can increase this number.

Building the Citation Graph

WITH RECURSIVE citation_chain AS (
    -- Base case: direct citations
    SELECT citing_paper_id, cited_paper_id, 1 as depth
    FROM citations
    WHERE citing_paper_id = 'target_paper_id'

    UNION ALL

    -- Recursive case: follow citation chain
    SELECT c.citing_paper_id, c.cited_paper_id, cc.depth + 1
    FROM citations c
    INNER JOIN citation_chain cc ON c.citing_paper_id = cc.cited_paper_id
    WHERE cc.depth < 3  -- Limit depth to prevent exponential growth
)
SELECT DISTINCT cited_paper_id, depth
FROM citation_chain
ORDER BY depth, cited_paper_id;

Analyzing Citation Patterns

import networkx as nx
from typing import Dict, List

def build_citation_network(citations: List[Dict]) -> nx.DiGraph:
    """
    Construct a directed graph from citation data
    """
    G = nx.DiGraph()
    for citation in citations:
        G.add_edge(citation['citing_paper_id'],
                  citation['cited_paper_id'])
    return G

def calculate_paper_influence(G: nx.DiGraph) -> Dict[str, float]:
    """
    Calculate PageRank scores for papers in the network
    """
    return nx.pagerank(G, alpha=0.85)

9.2 Visual Content Analysis

Image Similarity Search

CREATE INDEX image_vector_idx ON paper_images
USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);

-- Find similar figures across papers
SELECT i2.arxiv_id, i2.image_number, i2.caption,
       l2_distance(i1.embedding, i2.embedding) as similarity
FROM paper_images i1, paper_images i2
WHERE i1.arxiv_id = 'source_paper_id'
  AND i1.image_number = 'target_image_number'
  AND i1.arxiv_id != i2.arxiv_id
ORDER BY similarity ASC
LIMIT 10;

Image Clustering Analysis

from sklearn.cluster import DBSCAN
import numpy as np

def cluster_similar_images(embeddings: np.ndarray,
                         eps: float = 0.5,
                         min_samples: int = 5) -> np.ndarray:
    """
    Cluster images based on embedding similarity
    """
    clustering = DBSCAN(eps=eps, min_samples=min_samples)
    return clustering.fit_predict(embeddings)

9.3 Mathematical Pattern Analysis

Formula Similarity Search

-- Find papers using similar mathematical concepts
WITH similar_formulas AS (
    SELECT f2.arxiv_id,
           f2.formula_number,
           f2.latex_source,
           l2_distance(f1.embedding, f2.embedding) as similarity
    FROM paper_formulas f1, paper_formulas f2
    WHERE f1.arxiv_id = 'source_paper_id'
    ORDER BY similarity ASC
    LIMIT 100
)
SELECT arxiv_id,
       COUNT(*) as similar_formula_count,
       AVG(similarity) as avg_similarity
FROM similar_formulas
GROUP BY arxiv_id
ORDER BY avg_similarity ASC
LIMIT 10;

9.4 Table Data Analysis

Aggregating Numerical Results

import pandas as pd
from typing import List, Dict

def analyze_table_trends(table_data: List[Dict]) -> pd.DataFrame:
    """
    Analyze trends in numerical table data across papers
    """
    df = pd.DataFrame(table_data)

    # Extract numerical columns
    numeric_cols = df.select_dtypes(include=[np.number]).columns

    # Calculate statistics
    stats = df[numeric_cols].agg(['mean', 'std', 'min', 'max'])

    return stats

def detect_outliers(df: pd.DataFrame,
                   threshold: float = 2.0) -> pd.DataFrame:
    """
    Detect unusual results in table data
    """
    z_scores = pd.DataFrame()
    for column in df.select_dtypes(include=[np.number]):
        z_scores[column] = (df[column] - df[column].mean()) / df[column].std()

    return df[z_scores.abs().max(axis=1) > threshold]

9.5 Multimodal Content Correlation

Cross-Modal Analysis

-- Find correlations between text sections and images
WITH text_context AS (
    SELECT s.arxiv_id,
           s.section_id,
           s.content_text,
           s.embedding as text_embedding
    FROM paper_sections s
    WHERE s.section_type = 'methods'
),
image_context AS (
    SELECT i.arxiv_id,
           i.image_number,
           i.embedding as image_embedding
    FROM paper_images i
)
SELECT tc.arxiv_id,
       tc.section_id,
       ic.image_number,
       cosine_similarity(tc.text_embedding, ic.image_embedding) as correlation
FROM text_context tc
JOIN image_context ic ON tc.arxiv_id = ic.arxiv_id
ORDER BY correlation DESC;

Trend Analysis Across Modalities

def analyze_multimodal_patterns(paper_id: str,
                              conn: psycopg2.extensions.connection) -> Dict:
    """
    Analyze patterns across text, images, and formulas within a paper

    Args:
        paper_id: The arXiv ID of the paper to analyze
        conn: PostgreSQL database connection

    Returns:
        Dictionary containing multimodal analysis results
    """
    import numpy as np
    from sklearn.metrics.pairwise import cosine_similarity
    import pandas as pd
    from collections import defaultdict

    # Initialize results container
    analysis_results = {
        'paper_id': paper_id,
        'cross_modal_similarities': {},
        'content_clusters': {},
        'key_concepts': [],
        'visual_themes': [],
        'mathematical_themes': []
    }

    try:
        # Extract section embeddings
        sections_query = """
            SELECT section_id, section_type, embedding, content_text
            FROM paper_sections
            WHERE arxiv_id = %s
        """
        sections_df = pd.read_sql(sections_query, conn, params=[paper_id])

        # Extract image embeddings
        images_query = """
            SELECT image_id, image_number, embedding, caption
            FROM paper_images
            WHERE arxiv_id = %s
        """
        images_df = pd.read_sql(images_query, conn, params=[paper_id])

        # Extract formula embeddings
        formulas_query = """
            SELECT formula_id, formula_number, embedding, latex_source
            FROM paper_formulas
            WHERE arxiv_id = %s
        """
        formulas_df = pd.read_sql(formulas_query, conn, params=[paper_id])

        # 1. Calculate cross-modal similarities
        if not sections_df.empty and not images_df.empty:
            # Convert embeddings from database format to numpy arrays
            section_embeddings = np.vstack(sections_df['embedding'].apply(np.array).values)
            image_embeddings = np.vstack(images_df['embedding'].apply(np.array).values)

            # Calculate similarity matrix between sections and images
            text_image_similarity = cosine_similarity(section_embeddings, image_embeddings)

            # Find highest similarity pairs
            text_image_pairs = []
            for i, section_id in enumerate(sections_df['section_id']):
                for j, image_id in enumerate(images_df['image_id']):
                    similarity = text_image_similarity[i, j]
                    if similarity > 0.6:  # Threshold for meaningful similarity
                        text_image_pairs.append({
                            'section_id': int(section_id),
                            'section_type': sections_df.iloc[i]['section_type'],
                            'image_id': int(image_id),
                            'image_number': images_df.iloc[j]['image_number'],
                            'similarity': float(similarity)
                        })

            analysis_results['cross_modal_similarities']['text_image'] = text_image_pairs

        # Do the same for text-formula relationships
        if not sections_df.empty and not formulas_df.empty:
            section_embeddings = np.vstack(sections_df['embedding'].apply(np.array).values)
            formula_embeddings = np.vstack(formulas_df['embedding'].apply(np.array).values)

            text_formula_similarity = cosine_similarity(section_embeddings, formula_embeddings)

            text_formula_pairs = []
            for i, section_id in enumerate(sections_df['section_id']):
                for j, formula_id in enumerate(formulas_df['formula_id']):
                    similarity = text_formula_similarity[i, j]
                    if similarity > 0.6:
                        text_formula_pairs.append({
                            'section_id': int(section_id),
                            'section_type': sections_df.iloc[i]['section_type'],
                            'formula_id': int(formula_id),
                            'formula_number': formulas_df.iloc[j]['formula_number'],
                            'similarity': float(similarity)
                        })

            analysis_results['cross_modal_similarities']['text_formula'] = text_formula_pairs

        # 2. Identify content clusters using embedding similarities
        if not sections_df.empty:
            from sklearn.cluster import AgglomerativeClustering

            # Cluster section embeddings to find thematic groups
            section_embeddings = np.vstack(sections_df['embedding'].apply(np.array).values)

            # Determine optimal number of clusters based on data size
            n_clusters = min(max(2, len(section_embeddings) // 3), 5)

            if len(section_embeddings) >= n_clusters:
                clustering = AgglomerativeClustering(
                    n_clusters=n_clusters,
                    affinity='cosine',
                    linkage='average'
                )

                clusters = clustering.fit_predict(section_embeddings)

                # Group sections by cluster
                section_clusters = defaultdict(list)
                for i, cluster_id in enumerate(clusters):
                    section_clusters[int(cluster_id)].append({
                        'section_id': int(sections_df.iloc[i]['section_id']),
                        'section_type': sections_df.iloc[i]['section_type'],
                        'content_preview': sections_df.iloc[i]['content_text'][:100] + '...'
                    })

                analysis_results['content_clusters']['text_sections'] = dict(section_clusters)

        # 3. Extract key concepts from section embeddings
        if not sections_df.empty and 'content_text' in sections_df.columns:
            # Use TF-IDF to identify key terms from text content
            from sklearn.feature_extraction.text import TfidfVectorizer

            # Combine all section texts
            all_text = ' '.join(sections_df['content_text'].fillna(''))

            # Create vectorizer
            vectorizer = TfidfVectorizer(max_features=50, stop_words='english')

            # Apply to sections individually to find section-specific key terms
            section_texts = sections_df['content_text'].fillna('').tolist()
            if len(section_texts) > 0:
                tfidf_matrix = vectorizer.fit_transform(section_texts)
                feature_names = vectorizer.get_feature_names_out()

                # For each section, get top terms
                for i, section_id in enumerate(sections_df['section_id']):
                    if tfidf_matrix.getnnz(i) > 0:  # Check if section has any terms
                        section_scores = zip(feature_names, tfidf_matrix[i].toarray()[0])
                        sorted_terms = sorted(section_scores, key=lambda x: x[1], reverse=True)
                        top_terms = sorted_terms[:5]  # Top 5 terms per section

                        analysis_results['key_concepts'].append({
                            'section_id': int(section_id),
                            'section_type': sections_df.iloc[i]['section_type'],
                            'top_terms': [{'term': term, 'score': float(score)} for term, score in top_terms]
                        })

        # 4. Analyze visual themes based on image embeddings
        if not images_df.empty and len(images_df) >= 3:
            from sklearn.cluster import KMeans

            image_embeddings = np.vstack(images_df['embedding'].apply(np.array).values)

            # Use K-means for image clustering
            n_image_clusters = min(3, len(images_df))
            kmeans = KMeans(n_clusters=n_image_clusters)
            image_clusters = kmeans.fit_predict(image_embeddings)

            # Group images by cluster
            for cluster_id in range(n_image_clusters):
                cluster_images = [
                    {
                        'image_id': int(images_df.iloc[i]['image_id']),
                        'image_number': images_df.iloc[i]['image_number'],
                        'caption': images_df.iloc[i]['caption']
                    }
                    for i, c in enumerate(image_clusters) if c == cluster_id
                ]

                if cluster_images:
                    analysis_results['visual_themes'].append({
                        'cluster_id': int(cluster_id),
                        'image_count': len(cluster_images),
                        'images': cluster_images
                    })

        # 5. Analyze mathematical patterns in formulas
        if not formulas_df.empty:
            # Group formulas by similarity in vector space
            from sklearn.metrics.pairwise import cosine_similarity

            formula_embeddings = np.vstack(formulas_df['embedding'].apply(np.array).values)
            formula_similarity = cosine_similarity(formula_embeddings)

            # Find groups of related formulas (high similarity)
            formula_groups = []
            processed = set()

            for i in range(len(formula_similarity)):
                if i in processed:
                    continue

                # Find all formulas similar to this one
                related = [j for j in range(len(formula_similarity))
                          if formula_similarity[i, j] > 0.7 and i != j]

                if related:
                    group = {
                        'seed_formula_id': int(formulas_df.iloc[i]['formula_id']),
                        'seed_formula': formulas_df.iloc[i]['latex_source'],
                        'related_formulas': [
                            {
                                'formula_id': int(formulas_df.iloc[j]['formula_id']),
                                'formula_latex': formulas_df.iloc[j]['latex_source'],
                                'similarity': float(formula_similarity[i, j])
                            }
                            for j in related
                        ]
                    }
                    formula_groups.append(group)
                    processed.update(related)
                    processed.add(i)

            analysis_results['mathematical_themes'] = formula_groups

        # 6. Add simple summary statistics
        analysis_results['summary'] = {
            'section_count': len(sections_df),
            'image_count': len(images_df),
            'formula_count': len(formulas_df),
            'text_image_connections': len(analysis_results['cross_modal_similarities'].get('text_image', [])),
            'text_formula_connections': len(analysis_results['cross_modal_similarities'].get('text_formula', []))
        }

        return analysis_results

    except Exception as e:
        # Log error and return partial results
        import logging
        logging.error(f"Error in multimodal analysis for paper {paper_id}: {str(e)}")
        analysis_results['error'] = str(e)
        return analysis_results

9.6 Visualization and Reporting

Network Visualization

def visualize_citation_network(G: nx.DiGraph,
                             influence_scores: Dict[str, float]) -> None:
    """
    Create interactive visualization of citation network
    """
    pos = nx.spring_layout(G)

    # Node sizes based on influence
    node_sizes = [influence_scores[node] * 1000 for node in G.nodes()]

    # Create visualization
    plt.figure(figsize=(12, 8))
    nx.draw(G, pos,
            node_size=node_sizes,
            with_labels=True,
            node_color='lightblue',
            font_size=8)
    plt.title("Citation Network Analysis")
    plt.show()

Trend Reports

def generate_trend_report(conn: psycopg2.extensions.connection,
                         timeframe: str = '1 year') -> pd.DataFrame:
    """
    Generate comprehensive trend report across all modalities
    """
    # Query for temporal trends
    trend_query = """
        SELECT DATE_TRUNC('month', publication_date) as month,
               COUNT(DISTINCT p.arxiv_id) as paper_count,
               COUNT(DISTINCT i.id) as image_count,
               COUNT(DISTINCT f.id) as formula_count
        FROM papers p
        LEFT JOIN paper_images i ON p.arxiv_id = i.arxiv_id
        LEFT JOIN paper_formulas f ON p.arxiv_id = f.arxiv_id
        WHERE publication_date >= CURRENT_DATE - INTERVAL %s
        GROUP BY DATE_TRUNC('month', publication_date)
        ORDER BY month;
    """

    # Execute query and return results
    return pd.read_sql(trend_query, conn, params=[timeframe])

This comprehensive analysis framework enables deep insights into research patterns by combining traditional citation analysis with modern vector-based similarity search across multiple content types. The system can identify trends, patterns, and relationships that might be missed when analyzing each content type in isolation. Nevertheless please be aware that this is experimental and explorational code and has not been robustly tested in high volume production environments. It is meant to stimulate the imagination about what is possible with vector databases and rich PDF documents.

Section 10: Building a Personal Research Assistant Dashboard

This section will create a simple web dashboard for interfacing with the paper database. It will implement features like daily research summaries, paper recommendations based on reading history, citation alerts, and semantic search across the collection. The section will demonstrate how to connect a Python web framework to the PostgreSQL database with efficient vector search capabilities. The examples mention full text search In production, replace this placeholder with a true pgvector similarity query using the <⇒ operator and normalized embeddings.

10.1 Dashboard Architecture

The research assistant dashboard serves as an intuitive interface to our vector-enabled PostgreSQL database. We’ll build it using FastAPI for the backend API server, HTMX for the frontend interface, and PostgreSQL with pgvector for the database layer.

Core Components

from fastapi import FastAPI, Request, Form
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles
import psycopg2
from psycopg2.extras import RealDictCursor
from typing import Optional
app = FastAPI()
# Database connection
def get_db_connection():
    return psycopg2.connect(
        dbname="research_db",
        user="user",
        password="password",
        host="localhost",
        port="5432"
    )
# HTML templates
templates = Jinja2Templates(directory="templates")
@app.get("/", response_class=HTMLResponse)
async def root(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})
@app.get("/search-form/{search_type}", response_class=HTMLResponse)
async def get_search_form(request: Request, search_type: str):
    """Return the appropriate search form based on type"""
    if search_type == "semantic":
        return """
        <div class="search-options">
            <input type="text"
                   name="semantic_query"
                   placeholder="Enter text to find similar papers..."
                   class="search-box"
                   hx-get="/search/semantic"
                   hx-trigger="keyup changed delay:500ms"
                   hx-target="#results">
        </div>
        """
    elif search_type == "formula":
        return """
        <div class="search-options">
            <textarea name="formula"
                      placeholder="Enter LaTeX formula..."
                      class="search-box"
                      hx-post="/search/formula"
                      hx-trigger="keyup changed delay:500ms"
                      hx-target="#results"></textarea>
        </div>
        """
    elif search_type == "author":
        return """
        <div class="search-options">
            <input type="text"
                   name="author"
                   placeholder="Search by author name..."
                   class="search-box"
                   hx-get="/search/author"
                   hx-trigger="keyup changed delay:500ms"
                   hx-target="#results">
        </div>
        """
@app.get("/search/semantic", response_class=HTMLResponse)
async def semantic_search(request: Request, semantic_query: str):
    """Perform semantic search using vector similarity"""
    conn = get_db_connection()
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        # Simplified query - in reality, would use vector similarity
        cur.execute("""
            SELECT title, authors, abstract
            FROM papers
            WHERE to_tsvector('english', abstract) @@ plainto_tsquery('english', %s)
            LIMIT 10
        """, (semantic_query,))
        results = cur.fetchall()

        # Return results as HTML table
        html = "<table class='result-table'><tr><th>Title</th><th>Authors</th><th>Abstract</th></tr>"
        for r in results:
            html += f"<tr><td>{r['title']}</td><td>{r['authors']}</td><td>{r['abstract'][:200]}...</td></tr>"
        html += "</table>"
        return html
    finally:
        conn.close()
@app.post("/search/formula", response_class=HTMLResponse)
async def formula_search(request: Request, formula: str = Form(...)):
    """Search for papers with similar formulas"""
    conn = get_db_connection()
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        # Simplified query - in reality, would use vector similarity for formulas
        cur.execute("""
            SELECT title, formula_text, similarity
            FROM paper_formulas
            WHERE formula_text LIKE %s
            LIMIT 5
        """, (f"%{formula}%",))
        results = cur.fetchall()

        html = "<table class='result-table'><tr><th>Title</th><th>Formula</th><th>Similarity</th></tr>"
        for r in results:
            html += f"<tr><td>{r['title']}</td><td>{r['formula_text']}</td><td>{r['similarity']:.2f}</td></tr>"
        html += "</table>"
        return html
    finally:
        conn.close()
@app.get("/search/author", response_class=HTMLResponse)
async def author_search(request: Request, author: str):
    """Search for papers by author"""
    conn = get_db_connection()
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        cur.execute("""
            SELECT authors, title, publication_date
            FROM papers
            WHERE authors ILIKE %s
            ORDER BY publication_date DESC
            LIMIT 10
        """, (f"%{author}%",))
        results = cur.fetchall()

        html = "<table class='result-table'><tr><th>Authors</th><th>Title</th><th>Date</th></tr>"
        for r in results:
            html += f"<tr><td>{r['authors']}</td><td>{r['title']}</td><td>{r['publication_date']}</td></tr>"
        html += "</table>"
        return html
    finally:
        conn.close()
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

The HTMX frontend is a single page form that is simple yet powerful due to the capabilities of HTMX

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Research Assistant</title>
    <script src="https://cdnjs.cloudflare.com/ajax/libs/htmx/1.9.2/htmx.min.js"></script>
    <style>
        body {
            font-family: system-ui;
            line-height: 1.5;
            max-width: 1200px;
            margin: 0 auto;
            padding: 1rem;
        }
        .search-box {
            width: 100%;
            padding: 0.5rem;
            margin: 1rem 0;
        }
        .result-table {
            width: 100%;
            border-collapse: collapse;
        }
        .result-table th, .result-table td {
            padding: 0.5rem;
            border: 1px solid #ddd;
            text-align: left;
        }
        .tab-container {
            margin: 1rem 0;
        }
        .tab-button {
            padding: 0.5rem 1rem;
            margin-right: 0.5rem;
        }
        .search-options {
            margin: 1rem 0;
            padding: 1rem;
            border: 1px solid #ddd;
        }
    </style>
</head>
<body>
    <h1>Research Paper Database</h1>

    <!-- Search Tabs -->
    <div class="tab-container">
        <button class="tab-button"
                hx-get="/search-form/semantic"
                hx-target="#search-area">
            Semantic Search
        </button>
        <button class="tab-button"
                hx-get="/search-form/formula"
                hx-target="#search-area">
            Formula Search
        </button>
        <button class="tab-button"
                hx-get="/search-form/author"
                hx-target="#search-area">
            Author Search
        </button>
    </div>

    <!-- Search Area -->
    <div id="search-area">
        <!-- Default Search Form -->
        <div class="search-options">
            <input type="text"
                   name="query"
                   placeholder="Enter search query..."
                   class="search-box"
                   hx-get="/search"
                   hx-trigger="keyup changed delay:500ms"
                   hx-target="#results">
        </div>
    </div>

    <!-- Results Area -->
    <div id="results">
        <!-- Results will be dynamically loaded here -->
    </div>

</body>
</html>

10.2 Key Dashboard Features

Multimodal Search Interface

The dashboard implements an advanced search interface supporting:

Text-based semantic search using paper embeddings

Formula search with LaTeX input and similarity matching

Image similarity search with upload capability

Table content search and comparison

Hybrid search combining multiple modalities

Research Tracking and Recommendations

The system maintains personalized research profiles:

Reading history tracking

Paper recommendations based on reading patterns

Citation alerts for tracked papers

Research timeline visualization

Automated research summaries

Analytics and Visualization

Interactive visualizations provide research insights:

Citation network graphs using D3.js

Research trend analysis with time-series plots

Author collaboration networks

Formula usage patterns across papers

Image similarity clusters

10.3 Implementation Example: Multimodal Search

from fastapi import FastAPI, File, UploadFile
from typing import Optional
import psycopg2
from psycopg2.extras import RealDictCursor

async def search_papers(
    conn,
    text_query: Optional[str] = None,
    formula_query: Optional[str] = None,
    image_file: Optional[UploadFile] = None,
    filters: dict = None
):
    """
    Combined search function supporting multiple query types
    """
    results = []

    if text_query:
        # Text similarity search
        text_embedding = generate_text_embedding(text_query)
        text_results = search_by_text_embedding(conn, text_embedding)
        results.extend(text_results)

    if formula_query:
        # Formula similarity search
        formula_embedding = generate_formula_embedding(formula_query)
        formula_results = search_by_formula_embedding(conn, formula_embedding)
        results.extend(formula_results)

    if image_file:
        # Image similarity search
        image_embedding = generate_image_embedding(image_file)
        image_results = search_by_image_embedding(conn, image_embedding)
        results.extend(image_results)

    # Apply filters and ranking
    final_results = apply_filters_and_rank(results, filters)
    return final_results

10.4 Performance Optimization

To ensure responsive performance:

Query Optimization

Materialized views for common searches

Intelligent caching of embeddings

Parallel query execution for multimodal searches

Resource Management

Connection pooling for database access

Batch processing for embedding generation

Efficient memory usage for large result sets

Frontend Optimization

Lazy loading of results

Progressive image loading

Client-side caching

Conclusion

This chapter has explored the implementation of a sophisticated research assistant system using PostgreSQL with pgvector. We’ve covered the entire development process from basic setup to advanced features:

Foundation Building

PostgreSQL and pgvector installation

Vector data types and operations

Efficient indexing strategies

Content Management

Comprehensive schema design

PDF content extraction

Multimodal data handling

Embedding generation

Search and Analytics

Vector similarity search

Citation network analysis

Research trend tracking

Multimodal content analysis

Application Development

Web dashboard implementation

Performance optimization

User experience design

System integration

The combination of PostgreSQL’s robust database features with pgvector’s vector similarity capabilities provides a powerful foundation for building advanced research tools. This approach enables not just basic paper storage and retrieval, but sophisticated analysis and discovery features that can significantly enhance research workflows.

Looking ahead, this system can be extended in several directions:

Integration with other academic databases and APIs

Advanced machine learning for paper recommendations

Collaborative features for research teams

Mobile application development

Integration with reference management tools

The techniques and patterns presented in this chapter can be adapted for various document management scenarios beyond academic papers, making them valuable for any application requiring sophisticated text analysis and search capabilities.

Summary

This chapter follows a logical progression from understanding the technical foundations of PgVector through implementation of a complete research assistant system. It balances theoretical explanations with practical code examples, focusing specifically on the needs of researchers working with LLM papers from Arxiv. The sections build upon each other, with earlier sections establishing the database infrastructure and later sections leveraging that foundation for increasingly sophisticated analytics and applications.

This architecture is extensible to other domains—legal documents, financial reports, policy papers—where structured metadata and semantic content need to coexist.

Chapter 6. Building a Retrieval-Augmented Generation (RAG) System with SQLite VSS and Ollama

A note for Early Release readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 6th chapter of the final book. The GitHub repo will be made active later on.

If you’d like to be actively involved in reviewing and commenting on this draft, please reach out to the editor at shunter@oreilly.com.

In this chapter, we build a complete Retrieval-Augmented Generation (RAG) system that combines the semantic search capabilities of vector databases with the generative power of Large Language Models (LLMs). RAG systems solve a fundamental limitation of LLMs: their knowledge is frozen at training time and they cannot access private or recent information. By augmenting LLMs with a retrieval mechanism, we create a system that can answer questions using up-to-date, domain-specific knowledge.

Our goal is to build a question-answering system that can intelligently respond to queries about Reddit content. When a user asks a question, the system will: 1. Search through stored Reddit posts to find the most relevant information 2. Retrieve the best matching content chunks 3. Provide this context to an LLM 4. Generate a natural language answer based solely on the retrieved information

This approach ensures factual accuracy by grounding the LLM’s responses in actual data, preventing hallucinations and enabling the system to work with private or specialized content that wasn’t in the LLM’s training data.

System Architecture Overview

Our RAG system consists of five main components working in concert:

Vector Database Layer (SQLite VSS): Stores content chunks with their embeddings, enabling fast similarity search

Embedding Engine (SentenceTransformers): Converts text into dense vector representations that capture semantic meaning

Hybrid Search System: Combines semantic vector search with traditional keyword search for optimal retrieval

LLM Integration (Ollama): Provides local LLM inference for generating natural language responses

RAG Pipeline Orchestrator: Coordinates the retrieval and generation process

The data flow follows this path: - Content ingestion → Text chunking → Embedding generation → Vector storage - Query processing → Hybrid search → Context retrieval → Prompt construction → LLM generation

Let’s build each component step by step.

6.1 Database Foundation with Vector Support

First, we need to set up our database infrastructure. We’re using SQLite with the VSS (Vector Similarity Search) extension, which adds vector search capabilities to SQLite. This gives us a lightweight, embedded solution perfect for RAG applications.

# Chapter 6: Minimal RAG System with SQLite VSS and Ollama
# Simplified version demonstrating core Retrieval-Augmented Generation

import sqlite3
import requests
import json
import time
import hashlib
from sentence_transformers import SentenceTransformer
from typing import List, Dict, Optional

The imports reveal our technology stack: SQLite for storage, requests for API calls to Ollama, SentenceTransformers for embeddings, and standard Python libraries for data handling.

6.1.1 Setting Up the Vector-Enabled Database

The database setup is crucial as it creates the foundation for both traditional data storage and vector operations:

def setup_database(db_path='reddit_rag.db'):
    """Setup SQLite with VSS extension and create tables"""
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row  # Return rows as dictionaries

    # Enable and load VSS extension
    conn.enable_load_extension(True)
    try:
        conn.load_extension("./sqlite-vss0")
        print("VSS extension loaded successfully")
    except Exception as e:
        print(f"Error loading VSS: {e}")
        print("Download from: https://github.com/asg017/sqlite-vss/releases")
        raise

The VSS extension is what transforms SQLite from a traditional relational database into a hybrid vector database. Without this extension, we couldn’t perform efficient similarity searches on embeddings. The row_factory setting makes our code cleaner by returning query results as dictionaries instead of tuples.

6.1.2 Schema Design for RAG

Our schema needs to support three key requirements: storing original content, managing chunked text, and enabling both vector and keyword search:

    cursor = conn.cursor()

    # Create posts table
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS posts (
            post_id TEXT PRIMARY KEY,
            title TEXT,
            content TEXT,
            subreddit TEXT,
            author TEXT,
            score INTEGER
        )
    """)

The posts table stores the original Reddit content with metadata. This serves as our source of truth and provides context for the chunks.

    # Create chunks table with vector support
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS content_chunks (
            chunk_id INTEGER PRIMARY KEY,
            post_id TEXT,
            chunk_index INTEGER,
            content TEXT,
            chunk_vector BLOB,  -- Store as BLOB for VSS
            FOREIGN KEY (post_id) REFERENCES posts(post_id)
        )
    """)

The chunks table is where the magic happens. Each piece of content is broken into smaller chunks (essential for context window management), and each chunk has its own embedding stored as a BLOB. The foreign key relationship maintains the connection to the original post.

6.1.3 Creating Search Indexes

To enable efficient search, we create two specialized indexes:

    # Create VSS virtual table for vector search
    cursor.execute("""
        CREATE VIRTUAL TABLE IF NOT EXISTS chunk_vss USING vss0(
            chunk_vector(384),
            chunk_id INTEGER
        )
    """)

    # Create FTS5 for keyword search
    cursor.execute("""
        CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
            content,
            content='content_chunks',
            content_rowid='chunk_id'
        )
    """)

    conn.commit()
    return conn

The VSS virtual table enables vector similarity search with 384-dimensional vectors (matching our embedding model). The FTS5 virtual table provides full-text search capabilities. Together, they power our hybrid search system.

6.2 Text Processing and Embedding Generation

RAG systems need to break content into manageable chunks and convert them into searchable vectors. This section handles both tasks.

6.2.1 Embedding Model Management

We use a singleton pattern for the embedding model to avoid the overhead of loading it multiple times:

# Initialize embedding model globally
embedding_model = None

def get_embedding_model():
    """Get or initialize the embedding model"""
    global embedding_model
    if embedding_model is None:
        embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
        print(f"Loaded embedding model (dimension: {embedding_model.get_sentence_embedding_dimension()})")
    return embedding_model

The all-MiniLM-L6-v2 model is chosen for its balance of quality and speed. It produces 384-dimensional embeddings and is small enough to run efficiently on CPU.

6.2.2 Intelligent Text Chunking

Chunking is critical for RAG performance. Too small and we lose context; too large and we dilute relevance:

def chunk_text(text, chunk_size=200, overlap=50):
    """Simple text chunking by word count"""
    words = text.split()
    chunks = []

    for i in range(0, len(words), chunk_size - overlap):
        chunk = ' '.join(words[i:i + chunk_size])
        if chunk:
            chunks.append(chunk)

    return chunks

The overlap ensures continuity between chunks, preventing important information from being split across boundaries. A 200-word chunk with 50-word overlap provides good context while maintaining specificity.

6.2.3 Storing Content with Embeddings

This function orchestrates the entire ingestion pipeline:

def store_post_with_chunks(conn, post_data):
    """Store a post and its chunks with embeddings"""
    cursor = conn.cursor()

    # Store post
    cursor.execute("""
        INSERT OR REPLACE INTO posts (post_id, title, content, subreddit, author, score)
        VALUES (?, ?, ?, ?, ?, ?)
    """, (
        post_data['post_id'],
        post_data['title'],
        post_data['content'],
        post_data['subreddit'],
        post_data['author'],
        post_data.get('score', 0)
    ))

First, we store the original post with all its metadata. The INSERT OR REPLACE pattern handles updates gracefully.

    # Create chunks from title + content
    full_text = f"{post_data['title']}\n\n{post_data['content']}"
    chunks = chunk_text(full_text)

    # Get embedding model
    model = get_embedding_model()

We combine title and content because the title often contains important context that helps understand the content.

    # Store each chunk with embedding
    for idx, chunk_content in enumerate(chunks):
        # Generate embedding
        embedding = model.encode(chunk_content)

        # Store chunk
        cursor.execute("""
            INSERT INTO content_chunks (post_id, chunk_index, content, chunk_vector)
            VALUES (?, ?, ?, ?)
        """, (
            post_data['post_id'],
            idx,
            chunk_content,
            embedding.tobytes()
        ))

        chunk_id = cursor.lastrowid

Each chunk is embedded and stored with its vector representation. The tobytes() conversion is necessary for BLOB storage.

        # Add to VSS index
        vector_str = f"vss_create_vector({','.join(map(str, embedding))})"
        cursor.execute(f"""
            INSERT INTO chunk_vss (chunk_id, chunk_vector)
            VALUES (?, {vector_str})
        """, (chunk_id,))

        # Add to FTS index
        cursor.execute("""
            INSERT INTO chunks_fts (rowid, content)
            VALUES (?, ?)
        """, (chunk_id, chunk_content))

    conn.commit()
    print(f"Stored post {post_data['post_id']} with {len(chunks)} chunks")

We populate both search indexes - VSS for semantic search and FTS5 for keyword search. This dual indexing enables our hybrid search approach.

6.3 Hybrid Search Implementation

Hybrid search combines the strengths of semantic and keyword search. Semantic search understands meaning and context, while keyword search ensures specific terms aren’t missed.

6.3.1 The Hybrid Search Algorithm

def hybrid_search(conn, query_text, limit=5, semantic_weight=0.7):
    """Perform hybrid search combining vector and keyword search"""
    cursor = conn.cursor()
    model = get_embedding_model()

    # Generate query embedding
    query_embedding = model.encode(query_text)
    vector_str = f"vss_create_vector({','.join(map(str, query_embedding))})"

We start by converting the query into an embedding using the same model used for the chunks. This ensures compatibility in the vector space.

6.3.2 Semantic Search Component

    # Semantic search using VSS
    cursor.execute(f"""
        SELECT
            c.chunk_id,
            c.post_id,
            c.content,
            vss_cosine_distance(v.chunk_vector, {vector_str}) as distance
        FROM chunk_vss v
        JOIN content_chunks c ON v.chunk_id = c.chunk_id
        ORDER BY distance ASC
        LIMIT ?
    """, (limit * 2,))  # Get more results for merging

    semantic_results = {}
    for row in cursor.fetchall():
        semantic_results[row['chunk_id']] = {
            'chunk_id': row['chunk_id'],
            'post_id': row['post_id'],
            'content': row['content'],
            'semantic_score': 1 - row['distance']  # Convert distance to similarity
        }

The VSS extension provides the vss_cosine_distance function for similarity search. We fetch more results than needed (limit * 2) to have a good pool for merging with keyword results. Distance is converted to similarity for intuitive scoring.

6.3.3 Keyword Search Component

    # Keyword search using FTS5
    keyword_weight = 1 - semantic_weight
    keyword_results = {}

    try:
        cursor.execute("""
            SELECT
                c.chunk_id,
                c.post_id,
                c.content,
                bm25(chunks_fts) as score
            FROM chunks_fts f
            JOIN content_chunks c ON f.rowid = c.chunk_id
            WHERE chunks_fts MATCH ?
            ORDER BY score
            LIMIT ?
        """, (query_text, limit * 2))

        for row in cursor.fetchall():
            keyword_results[row['chunk_id']] = {
                'chunk_id': row['chunk_id'],
                'post_id': row['post_id'],
                'content': row['content'],
                'keyword_score': abs(row['score'])  # BM25 returns negative scores
            }
    except:
        # If FTS5 match fails, continue with semantic results only
        pass

FTS5 uses BM25 ranking, a probabilistic retrieval function that considers term frequency and document length. The try-except block handles cases where the query syntax might not be valid for FTS5.

6.3.4 Score Fusion and Ranking

    # Combine results
    all_chunk_ids = set(semantic_results.keys()) | set(keyword_results.keys())
    combined_results = []

    for chunk_id in all_chunk_ids:
        sem_score = semantic_results.get(chunk_id, {}).get('semantic_score', 0)
        key_score = keyword_results.get(chunk_id, {}).get('keyword_score', 0)

        # Normalize keyword scores
        if keyword_results:
            max_keyword = max(r['keyword_score'] for r in keyword_results.values())
            if max_keyword > 0:
                key_score = key_score / max_keyword

        combined_score = (sem_score * semantic_weight) + (key_score * keyword_weight)

        # Get chunk content
        content = semantic_results.get(chunk_id, keyword_results.get(chunk_id, {})).get('content', '')
        post_id = semantic_results.get(chunk_id, keyword_results.get(chunk_id, {})).get('post_id', '')

        combined_results.append({
            'chunk_id': chunk_id,
            'post_id': post_id,
            'content': content,
            'score': combined_score
        })

    # Sort by combined score and limit
    combined_results.sort(key=lambda x: x['score'], reverse=True)
    return combined_results[:limit]

The fusion strategy normalizes scores from both search methods and combines them using weighted averaging. The semantic_weight parameter (default 0.7) gives more importance to semantic similarity while still considering exact keyword matches.

6.4 LLM Integration with Ollama

Ollama provides local LLM inference, ensuring privacy and eliminating API costs. Our integration is straightforward but robust.

6.4.1 Ollama API Client

def call_ollama(prompt, model="llama3.1:8b", temperature=0.1):
    """Simple Ollama API call"""
    url = "http://localhost:11434/api/generate"

    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
        "options": {
            "temperature": temperature,
            "top_p": 0.9
        }
    }

    try:
        response = requests.post(url, json=payload)
        response.raise_for_status()
        return response.json()['response']
    except Exception as e:
        print(f"Ollama error: {e}")
        return f"Error calling Ollama: {str(e)}"

We use a low temperature (0.1) to ensure consistent, factual responses based on the provided context. The model choice (llama3.1:8b) balances quality with inference speed.

6.4.2 Health Check Function

def test_ollama():
    """Test if Ollama is running"""
    try:
        response = requests.get("http://localhost:11434/api/tags")
        models = response.json()
        print("Available Ollama models:", [m['name'] for m in models.get('models', [])])
        return True
    except:
        print("Ollama not running. Start with: ollama serve")
        return False

This utility function verifies Ollama is running and shows available models, helping with debugging and setup.

6.5 The RAG Pipeline

This is where everything comes together. The RAG pipeline orchestrates retrieval and generation to answer questions.

6.5.1 Context Formatting

Retrieved chunks need to be formatted with metadata for the LLM:

def format_context(chunks, conn):
“""Format retrieved chunks for the LLM""”
cursor = conn.cursor()
formatted_chunks = []

for i, chunk in enumerate(chunks, 1):
# Get post metadata
cursor.execute(“"”
SELECT title, subreddit, author, score
FROM posts
WHERE post_id = ?
“"”, (chunk['post_id'],))

post = cursor.fetchone()
if post:
formatted_chunks.append(f""”
SOURCE {i}:
From r/{post[’subreddit']} by u/{post['author']} (Score: {post[’score']})
Title: {post['title']}

Content:
{chunk['content']}
“"”)

return “\n---\n”.join(formatted_chunks)

Including metadata helps the LLM understand the source and credibility of information, leading to better answers.

6.5.2 Question Answering Pipeline

def answer_question(conn, question, num_chunks=5):
    """Complete RAG pipeline to answer a question"""
    print(f"\nQuestion: {question}")

    # Step 1: Retrieve relevant chunks
    start_time = time.time()
    chunks = hybrid_search(conn, question, limit=num_chunks)
    retrieval_time = (time.time() - start_time) * 1000

    if not chunks:
        return "No relevant information found in the database."

    print(f"Retrieved {len(chunks)} chunks in {retrieval_time:.1f}ms")
The pipeline starts with retrieval, measuring performance for optimization insights.
    # Step 2: Format context
    context = format_context(chunks, conn)

    # Step 3: Create prompt
    prompt = f"""You are a helpful assistant that answers questions based on Reddit content.
Use ONLY the information provided below to answer the question.
If the information is insufficient, say so clearly.

QUESTION: {question}

RETRIEVED INFORMATION:
{context}

ANSWER:"""

The prompt is carefully crafted to constrain the LLM to use only the provided context, preventing hallucinations and ensuring factual accuracy.

    # Step 4: Generate answer
    start_time = time.time()
    answer = call_ollama(prompt)
    generation_time = (time.time() - start_time) * 1000

    print(f"Generated answer in {generation_time:.1f}ms")
    print(f"Total time: {retrieval_time + generation_time:.1f}ms")

    return answer

Performance metrics help identify bottlenecks - typically retrieval is fast (tens of milliseconds) while generation takes longer (hundreds of milliseconds to seconds).

6.6 Demonstration and Testing

The final section provides sample data and interactive interfaces for testing the system.

6.6.1 Sample Data Loading

def load_sample_data(conn):
    """Load some sample Reddit data for testing"""
    sample_posts = [
        {
            'post_id': 'post001',
            'title': 'ELI5: What is machine learning?',
            'content': """Machine learning is like teaching a computer to recognize patterns by showing it many examples.
            Instead of programming exact rules, you feed it data and it learns the patterns itself. For example,
            to teach it to recognize cats, you show it thousands of cat pictures until it learns what makes a cat a cat.
            It's used in spam filters, recommendation systems, and voice assistants.""",
            'subreddit': 'explainlikeimfive',
            'author': 'curious_user',
            'score': 245
        },
        {
            'post_id': 'post002',
            'title': 'Best Python libraries for beginners?',
            'content': """I recommend starting with: NumPy for numerical computing, Pandas for data manipulation,
            Matplotlib for basic plotting, Requests for HTTP requests, and Flask for simple web apps. These libraries
            cover most beginner needs and have excellent documentation. Don't try to learn them all at once -
            start with one based on your project needs.""",
            'subreddit': 'learnpython',
            'author': 'python_mentor',
            'score': 189
        },
        {
            'post_id': 'post003',
            'title': 'How do vector databases work?',
            'content': """Vector databases store data as high-dimensional vectors (lists of numbers) that represent
            the semantic meaning of the data. When you search, your query is converted to a vector and the database
            finds the most similar vectors using distance metrics like cosine similarity. This enables semantic search
            where you find content by meaning rather than exact keyword matches. Popular vector databases include
            Pinecone, Weaviate, and pgvector for PostgreSQL.""",
            'subreddit': 'programming',
            'author': 'db_expert',
            'score': 156
        }
    ]

    print("Loading sample data...")
    for post in sample_posts:
        store_post_with_chunks(conn, post)
    print(f"Loaded {len(sample_posts)} sample posts")

Sample data covers diverse topics to demonstrate the system’s ability to handle different domains.

6.6.2 Main Demonstration Function

def main():
    """Main demonstration of the RAG system"""
    print("=== Minimal RAG System with SQLite VSS and Ollama ===\n")

    # Step 1: Setup database
    print("1. Setting up database...")
    conn = setup_database()

    # Step 2: Check Ollama
    print("\n2. Checking Ollama...")
    if not test_ollama():
        print("Please start Ollama first: ollama serve")
        print("Then pull a model: ollama pull llama3.1:8b")
        return

    # Step 3: Load sample data
    print("\n3. Loading sample data...")
    cursor = conn.cursor()
    cursor.execute("SELECT COUNT(*) FROM posts")
    if cursor.fetchone()[0] == 0:
        load_sample_data(conn)
    else:
        print("Data already loaded")

The main function provides a guided walkthrough of the system setup and capabilities.

6.6.3 Interactive Q&A Interface

    # Step 4: Demo questions
    print("\n4. Demonstrating RAG system...")

    demo_questions = [
        "What is machine learning and how does it work?",
        "What Python libraries should a beginner learn?",
        "How do vector databases enable semantic search?"
    ]

    for question in demo_questions[:1]:  # Just show one example
        answer = answer_question(conn, question)
        print(f"\nAnswer: {answer}\n")
        print("=" * 50)

    # Step 5: Interactive mode
    print("\n5. Interactive Q&A (type 'quit' to exit)")
    print("-" * 50)

    while True:
        question = input("\nYour question: ").strip()

        if question.lower() in ['quit', 'exit', 'q']:
            break

        if not question:
            continue

        answer = answer_question(conn, question)
        print(f"\nAnswer: {answer}")

    conn.close()
    print("\nGoodbye!")

The interactive mode allows users to experiment with different questions and see how the system responds.

6.6.4 Quick Testing Utility

def quick_start():
    """Quick start for testing individual components"""
    conn = setup_database()

    # Quick test of search
    print("Testing hybrid search for 'machine learning'...")
    results = hybrid_search(conn, "machine learning", limit=3)

    for i, result in enumerate(results, 1):
        print(f"\n{i}. Score: {result['score']:.3f}")
        print(f"   Content: {result['content'][:150]}...")

    conn.close()

if __name__ == "__main__":
    # Run main demo
    main()

    # Or run quick test
    # quick_start()

The quick_start function enables testing individual components without running the full pipeline, useful for debugging and development.

Next Steps: Extending the System

This implementation demonstrates the core concepts of a RAG system but intentionally keeps functionality limited for clarity. The system currently only indexes Reddit post titles and content, not the rich discussions in comments. Here are the key areas where you could extend this foundation to build a more featured useful system:

Missing Reddit Data Features

• Comment threads: The most valuable insights often appear in comments, not the main posts. Adding comment support would require fetching comment trees via PRAW, storing them with proper parent-child relationships, and deciding how to chunk threaded conversations while preserving context.

• Post metadata: Currently ignored metadata like timestamps, awards, crosspost information, and edit history could improve retrieval relevance and answer quality.

• Media content: Reddit posts often include images, videos, and external links that our text-only system cannot process.

• Live data ingestion: The system uses static sample data rather than fetching real Reddit content, which would require proper API authentication and rate limiting.

Retrieval Enhancements

• Semantic chunking: Our simple word-count chunking splits text arbitrarily. Semantic chunking would keep related ideas together by analyzing sentence embeddings and natural paragraph boundaries.

• Query expansion: Using the LLM to generate alternative phrasings of queries before searching could improve recall for complex questions.

• Cross-encoder reranking: After initial retrieval, a cross-encoder model could rerank results based on query-document relevance, significantly improving precision.

• Metadata filtering: Allow filtering by subreddit, date range, author karma, or post score before semantic search.

Performance Optimizations

• Embedding cache: Repeatedly embedding the same text wastes compute. A simple cache mapping text hashes to embeddings would speed up duplicate processing.

• Connection pooling: Our single database connection becomes a bottleneck under load. Connection pooling would enable concurrent operations.

• Batch processing: Processing posts one at a time is inefficient. Batching embeddings and database operations would improve throughput.

• Asynchronous pipelines: The sequential retrieval-then-generation flow could be parallelized for multiple queries.

Production Considerations

• Incremental updates: The system lacks the ability to add new content without reprocessing everything. Implementing incremental indexing would enable continuous updates.

• Version management: As embedding models improve, you need strategies for migrating existing embeddings or maintaining multiple versions.

• Monitoring and analytics: Production systems need query logging, performance metrics, and user feedback loops to identify retrieval failures and improve over time.

• Scale considerations: SQLite and local Ollama work for demos but have limits. Consider PostgreSQL with pgvector for larger datasets and API-based LLMs for concurrent users.

Advanced RAG Patterns

• Multi-hop reasoning: For complex questions requiring information from multiple sources, implement iterative retrieval where the LLM can request additional searches.

• Source diversity: Ensure retrieved chunks come from different posts/authors to avoid echo chamber effects.

• Confidence scoring: Have the LLM indicate confidence in its answers based on the quality and relevance of retrieved content.

• Fact verification: Add a verification step where claims in the generated answer are checked against the source material.

These extensions would transform our educational demonstration into a system capable of handling real-world complexity and scale. The modular architecture makes it straightforward to add these features incrementally based on your specific needs.

Conclusion

This RAG system demonstrates how vector databases enable a new paradigm in information retrieval and question answering. By combining semantic search with LLM generation, we create a system that can:

Understand the meaning behind queries, not just keywords

Retrieve the most relevant information from large document collections

Generate accurate, contextual answers grounded in real data

Avoid hallucinations by constraining responses to retrieved information

The hybrid search approach ensures we get the best of both worlds - semantic understanding and precise keyword matching. The local deployment using SQLite and Ollama provides privacy and cost-effectiveness while maintaining good performance.

This foundation can be extended with more sophisticated chunking strategies, better ranking algorithms, query expansion techniques, and multi-step reasoning chains - all building on the core RAG pattern demonstrated here.

Chapter 7. Building a Scientific RAG System with PostgreSQL and pgvector

A note for Early Release readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 7th chapter of the final book. The GitHub repo will be made active later on.

If you’d like to be actively involved in reviewing and commenting on this draft, please reach out to the editor at @oreilly.com.

The explosion of scientific publications presents researchers with an overwhelming challenge: how to efficiently discover, understand, and synthesize relevant knowledge from millions of papers. ArXiv alone publishes over 15,000 papers monthly across physics, mathematics, computer science, and other fields. Traditional keyword search fails to capture the semantic richness of scientific discourse, where the same concept may be expressed in countless ways across different domains and research communities.

This chapter builds a Retrieval-Augmented Generation (RAG) system specifically designed for scientific literature. Unlike general-purpose RAG systems, scientific RAG must handle unique challenges:

Technical terminology: Papers use precise, domain-specific language that requires semantic understanding beyond simple keywords

Structured content: Scientific papers follow conventions (abstract, methodology, results, conclusions) that can inform retrieval strategies

Citation networks: Papers exist in a web of references that provide additional context

Mathematical notation: Formulas and equations carry meaning that text-only systems miss

Evidence quality: Not all sources are equal - peer review, venue, and citation count matter

System Goals and Capabilities

Our scientific RAG system will enable researchers to:

Semantic Discovery: Find papers based on conceptual similarity, not just keyword matches

Cross-Paper Synthesis: Answer questions that require information from multiple papers

Contextual Understanding: Retrieve specific sections (methodology, results) relevant to queries

Evidence-Based Responses: Generate answers grounded in actual research, with proper citations

Technical Depth: Handle complex scientific queries requiring domain expertise

The system achieves these goals through: - Vector embeddings that capture scientific semantic meaning - Hierarchical indexing of papers (abstract-level and section-level) - PostgreSQL with pgvector for scalable vector operations - Integration with ArXiv for real-time paper access - Local LLM inference for privacy and control

Architecture Overview

The system architecture consists of six interconnected components:

1. Data Ingestion Layer

ArXiv API Integration: Fetches paper metadata and PDFs

PDF Processing: Extracts structured text from scientific PDFs

Section Detection: Identifies paper sections (introduction, methods, results)

2. Vector Processing Pipeline

Embedding Generation: Converts text to dense vectors using SentenceTransformers

Hierarchical Embeddings: Separate vectors for abstracts and sections

Dimension: 384-dimensional vectors balancing quality and performance

3. Storage Layer (PostgreSQL + pgvector)

Papers Table: Metadata and abstract embeddings

Sections Table: Chunked content with section embeddings

HNSW Indexes: Approximate nearest neighbor search at scale

4. Retrieval System

Multi-Level Search: Query both abstracts and sections

Similarity Thresholds: Filter results by relevance score

Context Assembly: Combine results from multiple sources

5. Generation Layer (Ollama)

Local LLM: Privacy-preserving inference

Scientific Prompts: Templates designed for academic responses

Citation Handling: Maintains paper references in answers

6. User Interface

Query Processing: Natural language scientific questions

Interactive Search: Different modes for different needs

Performance Monitoring: Track retrieval and generation times

Now let’s build each component with detailed explanations of the design decisions and implementation.

7.1 Database Foundation with pgvector

PostgreSQL with the pgvector extension provides robust vector storage with ACID guarantees. Unlike pure vector databases, this approach gives us the best of both worlds: relational data integrity and vector similarity search.

# Chapter 7: Minimal Scientific RAG with PostgreSQL and pgvector
# Simplified version demonstrating scientific paper search and RAG
#
# Dependencies: pip install psycopg2-binary arxiv sentence-transformers requests PyMuPDF

import psycopg2
import arxiv
import requests
import time
from sentence_transformers import SentenceTransformer
from typing import List, Dict, Optional
from dataclasses import dataclass
import numpy as np

7.1.1 Database Configuration and Setup

The database setup creates a schema optimized for scientific paper storage and retrieval:

def setup_database():
    """Setup PostgreSQL with pgvector for scientific papers"""
    # Hardcoded database connection - update as needed
    DB_CONFIG = {
        'host': 'localhost',
        'port': '5432',
        'database': 'scientific_rag',
        'user': 'postgres',
        'password': 'password'
    }

    conn = psycopg2.connect(**DB_CONFIG)
    cursor = conn.cursor()

    # Enable pgvector extension
    cursor.execute("CREATE EXTENSION IF NOT EXISTS vector;")

The pgvector extension adds the vector data type to PostgreSQL, enabling native vector operations. This is crucial for performance - vector operations happen in the database rather than requiring data transfer to application code.

7.1.2 Schema Design for Scientific Papers

The schema reflects the hierarchical nature of scientific documents:

    # Create papers table
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS papers (
            paper_id TEXT PRIMARY KEY,
            title TEXT NOT NULL,
            abstract TEXT,
            authors TEXT[],
            published_date DATE,
            pdf_url TEXT,
            abstract_embedding vector(384)
        );
    """)

The papers table stores high-level metadata. Note the use of PostgreSQL arrays for authors - this maintains data structure without requiring a separate table. The abstract embedding enables paper-level similarity search.

    # Create paper sections table for chunked content
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS paper_sections (
            section_id SERIAL PRIMARY KEY,
            paper_id TEXT REFERENCES papers(paper_id),
            section_number INTEGER,
            content TEXT NOT NULL,
            section_embedding vector(384)
        );
    """)

The sections table enables fine-grained retrieval. Scientific papers often contain diverse content - the methodology might be irrelevant to someone interested in results. Section-level indexing allows precise retrieval.

7.1.3 High-Performance Vector Indexes

HNSW (Hierarchical Navigable Small World) indexes provide approximate nearest neighbor search with excellent recall and query speed:

    # Create indexes for fast similarity search
    cursor.execute("""
        CREATE INDEX IF NOT EXISTS idx_abstract_embeddings
        ON papers USING hnsw (abstract_embedding vector_cosine_ops);
    """)

    cursor.execute("""
        CREATE INDEX IF NOT EXISTS idx_section_embeddings
        ON paper_sections USING hnsw (section_embedding vector_cosine_ops);
    """)

    conn.commit()
    cursor.close()
    conn.close()

    print("Database setup completed")
    return DB_CONFIG

The vector_cosine_ops parameter specifies cosine similarity, ideal for normalized embeddings. HNSW builds a multi-layer graph structure that enables logarithmic search complexity even with millions of vectors.

7.2 Embedding Generation Strategy

Embeddings are the bridge between human language and mathematical similarity. The choice of embedding model and strategy significantly impacts retrieval quality.

# Global embedding model
embedding_model = None

def get_embedding_model():
    """Get or initialize the embedding model"""
    global embedding_model
    if embedding_model is None:
        embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
        print(f"Loaded embedding model (dimension: {embedding_model.get_sentence_embedding_dimension()})")
    return embedding_model

def generate_embedding(text):
    """Generate embedding for text"""
    model = get_embedding_model()
    return model.encode(text, convert_to_numpy=True)

The singleton pattern prevents model reloading, crucial for performance. The all-MiniLM-L6-v2 model is specifically trained on scientific text pairs, making it ideal for academic content. The 384-dimensional output balances semantic richness with computational efficiency.

7.3 ArXiv Integration and PDF Processing

Scientific RAG requires access to full paper content, not just abstracts. This section handles paper discovery and content extraction.

7.3.1 Paper Discovery with ArXiv API

@dataclass
class Paper:
    paper_id: str
    title: str
    abstract: str
    authors: List[str]
    published_date: str
    pdf_url: str

def fetch_arxiv_papers(query, max_results=5):
    """Fetch papers from ArXiv"""
    print(f"Fetching papers from ArXiv: {query}")

    client = arxiv.Client()
    search = arxiv.Search(
        query=query,
        max_results=max_results,
        sort_by=arxiv.SortCriterion.SubmittedDate
    )

    papers = []
    for result in client.results(search):
        paper = Paper(
            paper_id=result.entry_id.split('/')[-1],
            title=result.title,
            abstract=result.summary,
            authors=[author.name for author in result.authors],
            published_date=result.published.date(),
            pdf_url=result.pdf_url
        )
        papers.append(paper)

        # Rate limiting
        time.sleep(0.5)

    print(f"Retrieved {len(papers)} papers")
    return papers

The dataclass provides type safety and clean data structure. ArXiv queries support complex boolean logic and field-specific searches (e.g., cat:cs.LG for machine learning papers). Rate limiting respects ArXiv’s servers and prevents blocking.

7.3.2 Intelligent PDF Text Extraction

PDF processing is notoriously difficult, especially for scientific papers with complex layouts, formulas, and figures. Our approach focuses on robust text extraction with section detection:

def download_pdf_text(pdf_url):
    """Download and extract text from PDF using PyMuPDF"""
    try:
        import fitz  # PyMuPDF

        print(f"Downloading PDF from {pdf_url[:50]}...")

        # Download PDF
        headers = {
            'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36'
        }
        response = requests.get(pdf_url, headers=headers, timeout=30)

        if response.status_code != 200:
            print(f"Failed to download PDF: HTTP {response.status_code}")
            return []

The User-Agent header prevents blocking by CDNs that filter automated requests. The timeout prevents hanging on slow connections.

        # Open PDF from bytes
        doc = fitz.open(stream=response.content, filetype="pdf")

        # Extract text from each page
        sections = []
        current_section = ""

        for page_num in range(len(doc)):
            page = doc[page_num]
            text = page.get_text()

            # Simple section detection: look for common headers
            lines = text.split('\n')
            section_headers = ['introduction', 'abstract', 'methodology', 'method', 'results',
                             'discussion', 'conclusion', 'references', 'related work', 'experiments',
                             'background', 'approach', 'evaluation', 'analysis']

Section detection uses common academic paper structure. This heuristic approach works well for most papers following standard formats. The headers list covers variations across different fields.

            for line in lines:
                line_lower = line.lower().strip()

                # Check if this line is a section header
                is_header = False
                for header in section_headers:
                    if (line_lower == header or
                        (len(line.split()) <= 3 and header in line_lower) or
                        line_lower.startswith(f"{header}.")):
                        is_header = True
                        break

                if is_header and current_section.strip():
                    # Save current section and start new one
                    sections.append(current_section.strip())
                    current_section = line + "\n"
                else:
                    current_section += line + "\n"

The header detection logic handles various formats: standalone headers, numbered sections (“1. Introduction”), and short headers with the keyword. This captures most academic paper structures.

        # Add final section
        if current_section.strip():
            sections.append(current_section.strip())

        # If no sections detected, split by pages
        if len(sections) <= 1:
            sections = []
            for page_num in range(min(10, len(doc))):  # Limit to first 10 pages
                page = doc[page_num]
                text = page.get_text()
                if text.strip():
                    sections.append(text.strip())

        doc.close()
        print(f"Extracted {len(sections)} sections from PDF")

        # Filter out very short sections (likely headers/footers)
        filtered_sections = [s for s in sections if len(s) > 100]

        return filtered_sections[:20]  # Limit to first 20 sections

    except ImportError:
        print("PyMuPDF not installed. Install with: pip install PyMuPDF")
        return []
    except Exception as e:
        print(f"Error extracting PDF text: {e}")
        return []

The fallback to page-based splitting ensures we extract something even from poorly formatted PDFs. The 100-character minimum filters out headers, footers, and page numbers. Limiting to 20 sections prevents memory issues with very long papers.

7.4 Advanced Text Chunking

Chunking strategy significantly impacts retrieval quality. Too small and we lose context; too large and we dilute relevance. Our approach balances these concerns:

def simple_chunk_text(text, chunk_size=500, overlap=50):
    """Simple text chunking by character count"""
    chunks = []
    words = text.split()

    # Convert to approximate character chunks
    current_chunk = []
    current_length = 0

    for word in words:
        word_length = len(word) + 1  # +1 for space

        if current_length + word_length > chunk_size and current_chunk:
            # Store current chunk
            chunks.append(' '.join(current_chunk))

            # Start new chunk with overlap
            if overlap > 0:
                overlap_words = current_chunk[-overlap//10:]  # Rough overlap
                current_chunk = overlap_words
                current_length = sum(len(w) + 1 for w in overlap_words)
            else:
                current_chunk = []
                current_length = 0

        current_chunk.append(word)
        current_length += word_length

    # Add final chunk
    if current_chunk:
        chunks.append(' '.join(current_chunk))

    return chunks

The 500-character chunks roughly correspond to a paragraph, maintaining semantic coherence. The overlap ensures information at chunk boundaries isn’t lost. This simple approach works well for scientific text where paragraph boundaries often align with conceptual boundaries.

7.5 Storage Pipeline with Embeddings

The storage pipeline coordinates PDF processing, chunking, embedding generation, and database insertion:

def store_paper_with_embeddings(db_config, paper, sections):
    """Store paper and sections with embeddings"""
    conn = psycopg2.connect(**db_config)
    cursor = conn.cursor()

    try:
        # Generate abstract embedding
        abstract_embedding = generate_embedding(paper.abstract)

        # Store paper
        cursor.execute("""
            INSERT INTO papers (paper_id, title, abstract, authors, published_date, pdf_url, abstract_embedding)
            VALUES (%s, %s, %s, %s, %s, %s, %s)
            ON CONFLICT (paper_id) DO UPDATE SET
                title = EXCLUDED.title,
                abstract = EXCLUDED.abstract,
                abstract_embedding = EXCLUDED.abstract_embedding
        """, (
            paper.paper_id,
            paper.title,
            paper.abstract,
            paper.authors,
            paper.published_date,
            paper.pdf_url,
            abstract_embedding.tolist()
        ))

The ON CONFLICT clause enables idempotent operations - reprocessing a paper updates rather than duplicates. This is essential for maintaining data consistency when refreshing papers or recovering from failures.

        # Store sections with embeddings
        for i, section_text in enumerate(sections):
            if section_text.strip():  # Skip empty sections
                # Chunk the section
                chunks = simple_chunk_text(section_text)

                for j, chunk in enumerate(chunks):
                    section_embedding = generate_embedding(chunk)

                    cursor.execute("""
                        INSERT INTO paper_sections (paper_id, section_number, content, section_embedding)
                        VALUES (%s, %s, %s, %s)
                    """, (
                        paper.paper_id,
                        i * 100 + j,  # Simple section numbering
                        chunk,
                        section_embedding.tolist()
                    ))

        conn.commit()
        print(f"Stored paper {paper.paper_id} with {len(sections)} sections")

    except Exception as e:
        conn.rollback()
        print(f"Error storing paper: {e}")
        raise
    finally:
        cursor.close()
        conn.close()

The section numbering scheme (i * 100 + j) preserves both section order and chunk order within sections. This enables reconstruction of the original document structure if needed. Transaction handling ensures atomicity - either the entire paper is stored or nothing is.

7.5.1 Batch Processing Pipeline

def process_and_store_papers(db_config, papers):
    """Process and store multiple papers"""
    for paper in papers:
        print(f"\nProcessing: {paper.title[:60]}...")

        # Download and extract sections
        sections = download_pdf_text(paper.pdf_url)

        if sections:
            store_paper_with_embeddings(db_config, paper, sections)
        else:
            print(f"  Skipping - no content extracted")

        # Rate limiting
        time.sleep(1)

Processing papers sequentially prevents memory issues with large PDFs and respects rate limits. In production, this could be parallelized with proper resource management.

7.6 Multi-Level Semantic Search

Scientific queries require different granularities of search. Sometimes we want paper-level summaries (abstracts), other times specific technical details (sections).

7.6.1 Abstract-Level Search

Abstract search is ideal for discovering relevant papers and understanding research landscape:

def search_papers_by_abstract(db_config, query, limit=5, threshold=0.7):
    """Search papers by abstract similarity"""
    conn = psycopg2.connect(**db_config)
    cursor = conn.cursor()

    # Generate query embedding
    query_embedding = generate_embedding(query)

    # Search abstracts
    cursor.execute("""
        SELECT
            paper_id,
            title,
            abstract,
            authors,
            1 - (abstract_embedding <=> %s::vector) as similarity
        FROM papers
        WHERE abstract_embedding IS NOT NULL
          AND 1 - (abstract_embedding <=> %s::vector) > %s
        ORDER BY similarity DESC
        LIMIT %s
    """, (query_embedding.tolist(), query_embedding.tolist(), threshold, limit))

    results = cursor.fetchall()
    cursor.close()
    conn.close()

    return [{
        'paper_id': row[0],
        'title': row[1],
        'abstract': row[2],
        'authors': row[3],
        'similarity': row[4]
    } for row in results]

The <=> operator computes cosine distance in pgvector. Converting to similarity (1 - distance) provides intuitive scores where 1.0 means identical and 0.0 means orthogonal. The threshold parameter filters out marginally relevant results, improving precision.

7.6.2 Section-Level Search

Section search enables finding specific technical details, methodologies, or results:

def search_paper_sections(db_config, query, limit=10, threshold=0.7):
    """Search within paper sections"""
    conn = psycopg2.connect(**db_config)
    cursor = conn.cursor()

    # Generate query embedding
    query_embedding = generate_embedding(query)

    # Search sections
    cursor.execute("""
        SELECT
            ps.paper_id,
            p.title,
            ps.section_number,
            ps.content,
            1 - (ps.section_embedding <=> %s::vector) as similarity
        FROM paper_sections ps
        JOIN papers p ON ps.paper_id = p.paper_id
        WHERE ps.section_embedding IS NOT NULL
          AND 1 - (ps.section_embedding <=> %s::vector) > %s
        ORDER BY similarity DESC
        LIMIT %s
    """, (query_embedding.tolist(), query_embedding.tolist(), threshold, limit))

    results = cursor.fetchall()
    cursor.close()
    conn.close()

    return [{
        'paper_id': row[0],
        'title': row[1],
        'section_number': row[2],
        'content': row[3],
        'similarity': row[4]
    } for row in results]

The join with the papers table provides context (paper title) for each section result. This helps users understand where the information comes from. Section search typically returns more results than abstract search since there are many more sections than papers.

7.7 The RAG Pipeline - Deep Dive

The RAG pipeline is where retrieval meets generation. This is the most critical component for answer quality, requiring careful orchestration of search, context assembly, and prompt engineering.

7.7.1 Local LLM Integration with Ollama

Ollama provides privacy-preserving local inference, crucial for sensitive research:

def call_ollama(prompt, model="llama3.1:8b"):
    """Call Ollama for text generation"""
    url = "http://localhost:11434/api/generate"

    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
        "options": {
            "temperature": 0.1,
            "top_p": 0.9
        }
    }

    try:
        response = requests.post(url, json=payload, timeout=60)
        response.raise_for_status()
        return response.json()['response']
    except Exception as e:
        return f"Error calling Ollama: {str(e)}"

Low temperature (0.1) ensures consistent, factual responses based on retrieved context. The 60-second timeout accommodates longer generation for complex scientific explanations. Streaming is disabled for simplicity but could be enabled for better user experience.

7.7.2 Health Check and Model Discovery

def test_ollama():
    """Test if Ollama is available"""
    try:
        response = requests.get("http://localhost:11434/api/tags", timeout=5)
        models = response.json()
        print("Available Ollama models:", [m['name'] for m in models.get('models', [])])
        return True
    except:
        print("Ollama not running. Start with: ollama serve")
        return False

This utility function helps with deployment and debugging. It shows available models, helping users choose appropriate models for their domain (e.g., specialized science models).

7.7.3 Intelligent Context Retrieval

The retrieval strategy combines multiple search types for comprehensive context:

def retrieve_relevant_contexts(db_config, question, max_contexts=5):
    """Retrieve relevant contexts for RAG"""
    # Search both abstracts and sections
    abstract_results = search_papers_by_abstract(db_config, question, limit=3, threshold=0.6)
    section_results = search_paper_sections(db_config, question, limit=max_contexts, threshold=0.6)

    contexts = []

    # Add abstract contexts
    for result in abstract_results:
        contexts.append({
            'type': 'abstract',
            'paper_id': result['paper_id'],
            'title': result['title'],
            'content': result['abstract'],
            'similarity': result['similarity']
        })

    # Add section contexts
    for result in section_results:
        contexts.append({
            'type': 'section',
            'paper_id': result['paper_id'],
            'title': result['title'],
            'content': result['content'],
            'similarity': result['similarity']
        })

    # Sort by similarity and limit
    contexts.sort(key=lambda x: x['similarity'], reverse=True)
    return contexts[:max_contexts]

The dual-search strategy ensures both broad context (abstracts) and specific details (sections) are available. Lower thresholds (0.6) for retrieval than display cast a wider net, letting the LLM determine relevance. Sorting by similarity ensures the most relevant contexts appear first in the prompt, improving answer quality since LLMs often give more weight to earlier content.

7.7.4 Scientific Prompt Engineering

Prompt design significantly impacts answer quality. Scientific prompts need special consideration:

def format_scientific_rag_prompt(question, contexts):
    """Format scientific contexts for RAG"""
    if not contexts:
        return f"Question: {question}\n\nI don't have relevant scientific literature to answer this question."

    context_parts = []
    for i, ctx in enumerate(contexts, 1):
        source_type = "Abstract" if ctx['type'] == 'abstract' else "Section"
        context_parts.append(
            f"Source {i} ({source_type} from '{ctx['title']}', similarity: {ctx['similarity']:.3f}):\n"
            f"{ctx['content'][:500]}..."
        )

    context_section = "\n\n".join(context_parts)

    prompt = f"""You are a scientific assistant helping researchers understand academic literature.
Answer the question based ONLY on the provided research paper excerpts.
Cite the specific papers when referencing information.

Question: {question}

Relevant research findings:
{context_section}

Answer based on the scientific literature above:"""

    return prompt

The prompt explicitly instructs the model to: 1. Act as a scientific assistant (sets appropriate tone and expertise level) 2. Use only provided contexts (prevents hallucination) 3. Cite sources (maintains academic rigor) 4. Include similarity scores (helps users assess relevance)

Context truncation (500 characters) balances information density with token limits. Including paper titles and types helps the LLM understand source credibility and relevance.

7.7.5 Complete RAG Execution Pipeline

The main RAG function orchestrates the entire process with performance monitoring:

def scientific_rag_answer(db_config, question, max_contexts=5):
    """Generate scientific answer using RAG"""
    print(f"\nScientific Question: {question}")

    # Step 1: Retrieve relevant contexts
    start_time = time.time()
    contexts = retrieve_relevant_contexts(db_config, question, max_contexts)
    retrieval_time = (time.time() - start_time) * 1000

    print(f"Retrieved {len(contexts)} contexts in {retrieval_time:.1f}ms")

    if not contexts:
        return "I couldn't find relevant scientific literature to answer this question."

    # Show retrieved contexts
    for i, ctx in enumerate(contexts, 1):
        print(f"  {i}. {ctx['title'][:50]}... (similarity: {ctx['similarity']:.3f})")

    # Step 2: Format prompt
    prompt = format_scientific_rag_prompt(question, contexts)

    # Step 3: Generate answer
    start_time = time.time()
    answer = call_ollama(prompt)
    generation_time = (time.time() - start_time) * 1000

    print(f"Generated answer in {generation_time:.1f}ms")
    print(f"Total time: {retrieval_time + generation_time:.1f}ms")

    return {
        'answer': answer,
        'contexts': contexts,
        'stats': {
            'retrieval_time_ms': retrieval_time,
            'generation_time_ms': generation_time,
            'num_contexts': len(contexts)
        }
    }

Performance monitoring helps identify bottlenecks: - Retrieval typically takes 50-200ms with indexes - Generation takes 500-5000ms depending on answer length - Total response time under 5 seconds is good for scientific queries

Returning contexts and statistics enables: - Answer verification by users - System performance optimization - Debugging retrieval issues - Building user trust through transparency

7.8 Demonstration and Interactive Interface

The demonstration showcases the system’s capabilities and provides interactive tools for exploration:

7.8.1 Main Demonstration Flow

def main():
    """Main demonstration of scientific RAG system"""
    print("=== Scientific RAG with PostgreSQL and pgvector ===\n")

    # Step 1: Setup database
    print("1. Setting up database...")
    db_config = setup_database()

    # Step 2: Check if we have papers
    conn = psycopg2.connect(**db_config)
    cursor = conn.cursor()
    cursor.execute("SELECT COUNT(*) FROM papers")
    paper_count = cursor.fetchone()[0]
    cursor.close()
    conn.close()

    if paper_count == 0:
        print("\n2. Fetching papers from ArXiv...")
        # Fetch papers on machine learning and AI
        papers = fetch_arxiv_papers('cat:cs.LG AND (abs:"machine learning" OR abs:"neural network")', max_results=3)

        if papers:
            print("\n3. Processing and storing papers...")
            process_and_store_papers(db_config, papers)
        else:
            print("No papers retrieved from ArXiv")
            return
    else:
        print(f"Found {paper_count} existing papers in database")

The demonstration intelligently checks for existing data, avoiding unnecessary reprocessing. The ArXiv query uses category filters (cat:cs.LG) and abstract searches (abs:) for precise paper selection.

7.8.2 Search Demonstrations

    # Step 4: Demo semantic search
    print("\n4. Demonstrating semantic search...")

    search_queries = [
        "neural network architecture",
        "machine learning optimization",
        "deep learning applications"
    ]

    for query in search_queries[:1]:  # Show one example
        print(f"\nSearching abstracts for: '{query}'")
        results = search_papers_by_abstract(db_config, query, limit=3)

        for i, result in enumerate(results, 1):
            print(f"\n{i}. {result['title']}")
            print(f"   Similarity: {result['similarity']:.3f}")
            print(f"   Authors: {', '.join(result['authors'][:2])}")
            print(f"   Abstract: {result['abstract'][:150]}...")

        print(f"\nSearching sections for: '{query}'")
        section_results = search_paper_sections(db_config, query, limit=3)

        for i, result in enumerate(section_results, 1):
            print(f"\n{i}. From: {result['title']}")
            print(f"   Similarity: {result['similarity']:.3f}")
            print(f"   Content: {result['content'][:200]}...")

Demonstrating both abstract and section search shows the different granularities of retrieval. Limiting output prevents information overload while showing system capabilities.

7.8.3 RAG Demonstration

    # Step 5: Demo RAG if Ollama is available
    print("\n5. Checking Ollama...")
    if test_ollama():
        print("\n6. Demonstrating Scientific RAG...")

        questions = [
            "What are the main challenges in machine learning?",
            "How do neural networks learn from data?",
            "What are current trends in AI research?"
        ]

        for question in questions[:1]:  # Show one example
            result = scientific_rag_answer(db_config, question)

            if isinstance(result, dict):
                print(f"\nAnswer: {result['answer'][:300]}...")
                print(f"Based on {result['stats']['num_contexts']} scientific sources")
            else:
                print(f"Answer: {result}")
    else:
        print("\n6. Skipping RAG demo (Ollama not available)")

    print("\n=== Demo completed! ===")

The questions are chosen to demonstrate different types of scientific queries: - Challenges question: Requires synthesis across papers - How question: Needs technical explanation - Trends question: Requires recent paper analysis

7.8.4 Interactive Search Interface

def interactive_search():
    """Interactive scientific search interface"""
    db_config = setup_database()

    print("=== Interactive Scientific Search ===")
    print("Commands: 'abstract [query]', 'section [query]', 'ask [question]', 'quit'\n")

    while True:
        command = input("\n> ").strip()

        if command.lower() == 'quit':
            break

        parts = command.split(' ', 1)
        if len(parts) < 2:
            print("Please specify a command and query")
            continue

        cmd, query = parts[0].lower(), parts[1]

        if cmd == 'abstract':
            results = search_papers_by_abstract(db_config, query, limit=3)
            for i, result in enumerate(results, 1):
                print(f"\n{i}. {result['title']}")
                print(f"   Similarity: {result['similarity']:.3f}")
                print(f"   {result['abstract'][:200]}...")

        elif cmd == 'section':
            results = search_paper_sections(db_config, query, limit=5)
            for i, result in enumerate(results, 1):
                print(f"\n{i}. From: {result['title']}")
                print(f"   Similarity: {result['similarity']:.3f}")
                print(f"   {result['content'][:150]}...")

        elif cmd == 'ask':
            if test_ollama():
                result = scientific_rag_answer(db_config, query)
                if isinstance(result, dict):
                    print(f"\nAnswer: {result['answer']}")
                else:
                    print(f"Answer: {result}")
            else:
                print("Ollama not available for RAG")

        else:
            print("Unknown command. Use: abstract, section, ask, or quit")

    print("\nGoodbye!")

The interactive interface enables: - Exploratory searches to understand the corpus - Testing different query formulations - Comparing abstract vs section search results - Real-time Q&A for research

7.8.5 Entry Point with Mode Selection

if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "interactive":
        interactive_search()
    else:
        main()

Supporting both demo and interactive modes accommodates different use cases - demonstration for learning, interactive for actual research work.

Next Steps: Extending the Scientific RAG System

This implementation provides a foundation for scientific literature analysis, but intentionally limits scope for clarity. Here are key areas for extension to build a more featured and useful research assistant:

Enhanced PDF Processing

Formula extraction: Scientific papers contain mathematical equations that our text-only extraction misses. Integrating tools like LaTeX OCR or MathML extraction would capture this crucial information.

Figure and table analysis: Research often communicates key findings through visualizations. Multi-modal models could describe figures and extract structured data from tables.

Citation parsing: Papers exist in citation networks. Extracting and following references would enable multi-hop reasoning across connected research.

Structured extraction: Instead of simple section detection, use machine learning models trained on scientific PDFs to identify methodology, results, and contributions with high accuracy.

Advanced Retrieval Techniques

Hybrid search: Combine vector similarity with keyword matching, citation networks, and metadata filters for more precise retrieval.

Query decomposition: Complex scientific questions often have multiple parts. Breaking them down and searching for each component separately could improve coverage.

Iterative refinement: Allow the LLM to request additional searches based on initial results, enabling deeper exploration of topics.

Concept linking: Build a knowledge graph of scientific concepts and use it to expand queries with related terms and synonyms.

Improved RAG Pipeline

Reranking: After initial retrieval, use cross-encoder models specifically trained on scientific text pairs to reorder results by relevance.

Answer verification: Add a fact-checking step that verifies claims in generated answers against the source documents.

Multi-paper synthesis: Develop strategies for combining information from multiple papers, handling contradictions, and identifying consensus.

Confidence scoring: Have the system indicate confidence levels based on the number and quality of supporting sources.

Domain Specialization

Field-specific models: Use embedding models fine-tuned on specific scientific domains (physics, biology, computer science) for better semantic understanding.

Terminology handling: Build domain-specific dictionaries and acronym expansions to handle technical jargon.

Equation understanding: Integrate symbolic math processing to understand and manipulate mathematical expressions.

Experimental data: Extract and analyze numerical results, statistical tests, and experimental parameters.

Scale and Performance

Distributed processing: Parallelize PDF processing and embedding generation across multiple workers for faster ingestion.

Incremental updates: Implement streaming ingestion that can add new papers without reprocessing the entire corpus.

Caching strategies: Cache embeddings, search results, and generated answers to reduce latency for common queries.

Index optimization: Tune HNSW parameters (M, ef_construction) based on corpus size and query patterns for optimal speed/recall tradeoff.

User Experience Enhancements

Source highlighting: Show exactly which parts of retrieved documents support specific claims in answers.

Interactive refinement: Allow users to provide feedback on retrieval results to improve subsequent searches.

Export capabilities: Generate formatted citations, bibliography entries, and literature review drafts.

Collaborative features: Enable sharing of searches, annotations, and curated paper collections among research teams.

Quality Assurance

Evaluation metrics: Implement standard IR metrics (precision, recall, F1) and RAG-specific metrics (answer accuracy, citation correctness).

Test datasets: Create domain-specific question-answer pairs with ground truth for systematic evaluation.

Human-in-the-loop: Add interfaces for expert feedback to continuously improve retrieval and generation quality.

Bias detection: Monitor for biases in paper selection and answer generation, ensuring diverse perspectives are represented.

Integration Capabilities

API endpoints: Expose REST/GraphQL APIs for integration with research tools and workflows.

Reference managers: Connect with Zotero, Mendeley, and EndNote for seamless bibliography management.

Writing assistants: Integrate with LaTeX editors and word processors for in-context literature search.

Continuous monitoring: Set up alerts for new papers matching saved queries, keeping researchers updated on their field.

These extensions would transform our educational demonstration into a comprehensive research platform, capable of accelerating scientific discovery by making vast literature accessible and synthesizable.

Conclusion

This scientific RAG system demonstrates how vector databases and language models can transform academic literature search. By combining semantic understanding with structured retrieval, we’ve created a system that can:

Understand scientific concepts through semantic embeddings trained on academic text

Navigate document structure with hierarchical indexing of abstracts and sections

Synthesize across sources by retrieving from multiple papers and generating coherent answers

Maintain academic rigor through source attribution and confidence scoring

Scale efficiently using PostgreSQL with pgvector for usable performance upto ~1 million documents

The architecture separates concerns cleanly:

- Data layer handles PDF processing and text extraction

- Embedding layer manages vector generation and caching

- Storage layer provides reliable persistence with ACID guarantees

- Retrieval layer implements multi-level similarity search

- Generation layer produces scientifically accurate responses

This separation enables independent optimization and testing of each component, crucial for maintaining system quality as it scales.

The choice of PostgreSQL with pgvector over specialized vector databases provides several advantages for scientific applications:

- Transactional integrity ensures consistent updates to papers and embeddings

- Relational queries enable complex filters on metadata alongside vector search

- Operational simplicity reduces deployment complexity with a single database

- Proven scale handles millions of papers with proper indexing

Future advances in embedding models, retrieval algorithms, and language models will only enhance these capabilities. The modular design ensures the system can adopt these improvements without architectural changes.

By making scientific literature programmatically accessible and synthesizable, systems like this accelerate research by helping scientists discover connections, identify gaps, and build on existing knowledge more effectively than ever before.

Chapter 8. Building a Complete Conversation Search and RAG System

A note for Early Release readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 8th chapter of the final book. The GitHub repo will be made active later on.

If you’d like to be actively involved in reviewing and commenting on this draft, please reach out to the editor at @oreilly.com.

In our digital age, we generate vast amounts of conversational data through AI assistants, chat applications, and collaborative platforms. These conversations often contain valuable insights, solutions to problems, creative ideas, and learned knowledge that become increasingly difficult to retrieve as the volume grows. Imagine having hundreds of conversations with Claude about programming, research, creative writing, and personal projects - how do you efficiently find that specific discussion about Python decorators from three months ago, or recall the architectural advice given for your web application?

Traditional search approaches fall short for conversational data. Keyword search misses the semantic richness of natural language dialogue, while chronological browsing becomes impractical with large conversation histories. The conversational context - the flow of questions and answers, the building of ideas across multiple exchanges - is lost when individual messages are treated in isolation.

This chapter addresses these challenges by building a complete conversation search and RAG (Retrieval-Augmented Generation) system specifically designed for personal AI conversation history. Unlike the previous chapters that focused on public data sources (Reddit posts, scientific papers), this system handles private, personal knowledge with the nuanced understanding that conversations have unique characteristics:

Unique Challenges of Conversation Search

Contextual Dependencies: Messages often reference previous exchanges, use pronouns, or build upon earlier concepts. A response like “That approach works well, but consider the performance implications” is meaningless without the surrounding context.

Conversational Flow: Ideas develop across multiple message exchanges. A single search result might not capture the complete reasoning or solution that emerged from a back-and-forth discussion.

Personal Language Patterns: Conversations with AI assistants develop their own vocabulary, shorthand, and reference patterns unique to the individual user’s communication style and domain interests.

Privacy and Control: Personal conversations require local storage and processing, not cloud-based solutions that compromise privacy or require ongoing API costs.

Temporal Relevance: Recent conversations might be more relevant than older ones, but breakthrough insights from months ago shouldn’t be lost in the noise.

System Goals and Capabilities

Our conversation search and RAG system will provide:

Core Search Capabilities

Semantic Discovery: Find relevant conversations based on meaning, not just keyword matches

Contextual Retrieval: Return conversation segments with surrounding context to maintain dialogue flow

Cross-Conversation Synthesis: Answer questions by combining insights from multiple conversation threads

Temporal Navigation: Efficiently browse conversation history with both semantic and chronological organization

RAG-Powered Insights

Knowledge Synthesis: Generate comprehensive answers by combining relevant exchanges from across your conversation history

Pattern Recognition: Identify recurring themes, solutions, and evolving thoughts across time

Decision Support: Surface past decisions, reasoning, and outcomes to inform current choices

Learning Reinforcement: Rediscover explanations and concepts previously discussed but potentially forgotten

Scalable System Architecture

Web API Interface: RESTful endpoints for integration with other applications and interfaces

Scalable Storage: PostgreSQL with pgvector for handling large conversation archives

Batch Processing: Efficient embedding generation and indexing for thousands of messages

Real-time Search: Sub-second response times for interactive exploration

System Architecture Overview

Our architecture consists of seven integrated components that work together to transform raw conversation data into an intelligent, searchable knowledge system:

1. Data Ingestion Pipeline

JSON Import System: Processes Claude conversation exports and other chat formats

Schema Validation: Ensures data consistency and handles various export formats

Incremental Updates: Supports adding new conversations without reprocessing existing data

2. Conversation Database Design

Hierarchical Storage: Separate tables for conversations, messages, and embeddings

Relationship Preservation: Maintains conversation threading and message ordering

Metadata Management: Stores timestamps, sender information, and conversation context

3. Vector Processing Engine

Embedding Generation: Converts messages to semantic vectors using SentenceTransformers

Batch Optimization: Processes large message volumes efficiently

Vector Indexing: HNSW indexes for fast similarity search across thousands of conversations

4. Contextual Search System

Semantic Similarity: Finds conceptually related messages across conversation boundaries

Context Window Retrieval: Returns surrounding messages to preserve conversational flow

Relevance Scoring: Balances semantic similarity with conversational coherence

5. RAG Integration Layer

Context Assembly: Combines relevant conversation segments into coherent prompts

Prompt Engineering: Templates optimized for conversational knowledge synthesis

Local LLM Integration: Ollama-powered generation for privacy and cost control

6. Web Service API

RESTful Endpoints: Search, RAG, and analytics APIs for external integration

Real-time Processing: Live search and answer generation

Statistics Dashboard: System health and usage metrics

7. Demonstration Interface

Interactive CLI: Command-line exploration of conversation data

Sample Data Generation: Testing with realistic conversation examples

Deployment Tools: Complete system setup and configuration

What We’ll Build Together

This chapter takes a complete systems approach, building not just the core RAG functionality but a fully-featured application. We’ll start with database design optimized for conversation storage, implement efficient batch processing for large-scale embedding generation, and culminate with a web API that can serve as the foundation for chat applications, personal knowledge management tools, or research platforms.

Unlike previous chapters that focused on specific techniques, this chapter demonstrates how all the pieces fit together in a real application. You’ll see how to handle the messy realities of data import, optimize for both speed and accuracy, and build systems that can grow with your conversation history over months and years.

The FastAPI web service we’ll build provides robust endpoints that can be integrated into other applications, making your conversation history a living, queryable knowledge base rather than static archived data.

By the end of this chapter, you’ll have a complete system that not only demonstrates advanced RAG techniques but provides immediate practical value for managing and extracting insights from your own AI conversation history. Let’s begin building this personal knowledge assistant.

8.1 Database Foundation for Conversation Storage

Conversation data has unique characteristics that require specialized database design. Unlike documents or posts, conversations have inherent structure (threads, message ordering), metadata (timestamps, participants), and relationships (responses, context dependencies) that must be preserved for effective search and retrieval.

# Chapter 8: Minimal LLM Conversation Search & RAG System
# Simplified version demonstrating conversation storage, search, and RAG with pgvector

import psycopg2
import json
import requests
import time
from pathlib import Path
from sentence_transformers import SentenceTransformer
from datetime import datetime
from typing import List, Dict, Optional
from dataclasses import dataclass

The imports reveal our technology stack: PostgreSQL with pgvector for vector storage, SentenceTransformers for embeddings, FastAPI for the web service, and standard libraries for data processing.

8.1.1 Designing the Conversation Schema

The database design reflects the hierarchical nature of conversations while optimizing for both relational queries and vector search:

def setup_database():
    """Setup PostgreSQL with pgvector extension and create tables"""
    # Hardcoded database connection - update as needed
    DB_CONFIG = {
        'host': 'localhost',
        'port': '5432',
        'database': 'claude_conversations',
        'user': 'postgres',
        'password': 'password'
    }

    conn = psycopg2.connect(**DB_CONFIG)
    cursor = conn.cursor()

    # Enable pgvector extension
    cursor.execute("CREATE EXTENSION IF NOT EXISTS vector;")

The dedicated database name claude_conversations reflects this system’s focus on AI conversation history, though the schema could handle any conversational data.

8.1.2 Three-Table Architecture for Optimal Performance

Our schema uses three tables to balance normalized data integrity with query performance:

    # Create simple schema
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS conversations (
            uuid TEXT PRIMARY KEY,
            name TEXT NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );
    """)

The conversations table stores high-level metadata. Using UUIDs as primary keys matches Claude’s export format and ensures global uniqueness. The conversation name provides human-readable context for search results.

    cursor.execute("""
        CREATE TABLE IF NOT EXISTS messages (
            uuid TEXT PRIMARY KEY,
            conversation_uuid TEXT REFERENCES conversations(uuid),
            sender TEXT NOT NULL,
            text TEXT NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            message_index INTEGER
        );
    """)

The messages table stores individual exchanges with foreign key relationships to maintain conversation context. The message_index preserves message ordering within conversations - crucial for retrieving contextual windows around matching results.

cursor.execute(“"”
CREATE TABLE IF NOT EXISTS message_embeddings (
message_uuid TEXT PRIMARY KEY REFERENCES messages(uuid),
embedding vector(384) NOT NULL
);
“"”)

Separating embeddings into their own table provides several advantages:

- Lazy Loading: Messages can be imported without immediately generating embeddings

- Embedding Updates: New or improved embedding models can be deployed without touching message data

- Storage Optimization: Vector data doesn’t bloat the main messages table

- Query Performance: Vector operations don’t interfere with text-based queries

8.1.3 High-Performance Vector Indexing

The HNSW index is crucial for sub-second search across large conversation archives:

    # Create index for fast similarity search
    cursor.execute("""
        CREATE INDEX IF NOT EXISTS idx_message_embeddings_vector
        ON message_embeddings
        USING hnsw (embedding vector_cosine_ops);
    """)

    conn.commit()
    cursor.close()
    conn.close()

    print("Database setup completed")
    return DB_CONFIG

This index enables logarithmic search complexity even with hundreds of thousands of messages. The vector_cosine_ops operator class optimizes for cosine similarity, ideal for normalized text embeddings.

8.2 Conversation Import and Data Processing Pipeline

Real-world conversation data comes in various formats and quality levels. Our import system handles Claude’s JSON export format while providing extensibility for other conversation sources.

8.2.1 Robust JSON Import with Error Handling

The import function processes Claude’s conversation export format, handling the nested structure of conversations containing message arrays:

def import_conversations(db_config, conversations_file):
    """Import conversations from Claude export JSON"""
    if not Path(conversations_file).exists():
        print(f"File {conversations_file} not found")
        return 0, 0

    conn = psycopg2.connect(**db_config)
    cursor = conn.cursor()

    with open(conversations_file, 'r') as f:
        data = json.load(f)

    conversations = data.get('conversations', [])
    total_conversations = 0
    total_messages = 0

    print(f"Processing {len(conversations)} conversations...")

The function returns counts of processed items, enabling the caller to track import progress and detect issues. Using data.get('conversations', []) provides graceful degradation for malformed files.

8.2.2 Atomic Transaction Processing

Each conversation is processed in its own transaction to ensure data consistency while allowing partial recovery from errors:

    for conv in conversations:
        try:
            # Insert conversation
            cursor.execute("""
                INSERT INTO conversations (uuid, name, created_at)
                VALUES (%s, %s, %s)
                ON CONFLICT (uuid) DO NOTHING
            """, (
                conv['uuid'],
                conv['name'],
                datetime.fromisoformat(conv['created_at'].replace('Z', '+00:00'))
            ))

            if cursor.rowcount > 0:
                total_conversations += 1

The ON CONFLICT DO NOTHING clause enables idempotent imports - reprocessing the same file won’t create duplicates. The rowcount check tracks actually inserted records versus skipped duplicates.

8.2.3 Timestamp Handling and Data Validation

ISO timestamp processing requires careful handling of timezone information:

            # Insert messages
            messages = conv.get('messages', [])
            for idx, msg in enumerate(messages):
                cursor.execute("""
                    INSERT INTO messages (uuid, conversation_uuid, sender, text, created_at, message_index)
                    VALUES (%s, %s, %s, %s, %s, %s)
                    ON CONFLICT (uuid) DO NOTHING
                """, (
                    msg['uuid'],
                    conv['uuid'],
                    msg['sender'],
                    msg['text'],
                    datetime.fromisoformat(msg['created_at'].replace('Z', '+00:00')),
                    idx
                ))

                if cursor.rowcount > 0:
                    total_messages += 1

            conn.commit()

The .replace('Z', '+00:00') conversion handles the common ‘Z’ UTC timezone indicator that Python’s fromisoformat() doesn’t parse directly. Using enumerate() generates the message index automatically, preserving conversation order.

8.2.4 Error Recovery and Logging

Robust error handling ensures partial import success even when individual conversations have issues:

        except Exception as e:
            print(f"Error processing conversation {conv.get('uuid', 'unknown')}: {e}")
            conn.rollback()
            continue

    cursor.close()
    conn.close()

    print(f"Imported {total_conversations} conversations with {total_messages} messages")
    return total_conversations, total_messages

Per-conversation error logging with UUID identification helps diagnose import issues. Rolling back failed transactions prevents partial data corruption while allowing subsequent conversations to process successfully.

8.3 Efficient Embedding Generation and Batch Processing

Generating embeddings for large conversation archives requires careful attention to memory usage, processing speed, and progress tracking. Our approach balances throughput with system stability.

8.3.1 Singleton Pattern for Model Management

The embedding model is expensive to load and should be reused across all operations:

# Global embedding model
embedding_model = None

def get_embedding_model():
    """Get or initialize the embedding model"""
    global embedding_model
    if embedding_model is None:
        embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
        print(f"Loaded embedding model (dimension: {embedding_model.get_sentence_embedding_dimension()})")
    return embedding_model

The all-MiniLM-L6-v2 model provides an excellent balance of quality and speed for conversational text. Its 384-dimensional output is compact enough for large-scale storage while rich enough to capture semantic nuances.

8.3.2 Incremental Processing Strategy

The embedding generation function only processes messages that don’t already have embeddings, enabling incremental updates as new conversations are added:

def generate_message_embeddings(db_config, batch_size=100):
    """Generate embeddings for messages that don't have them"""
    conn = psycopg2.connect(**db_config)
    cursor = conn.cursor()

    # Find messages without embeddings
    cursor.execute("""
        SELECT m.uuid, m.text
        FROM messages m
        LEFT JOIN message_embeddings me ON m.uuid = me.message_uuid
        WHERE me.message_uuid IS NULL
        ORDER BY m.created_at
    """)

    messages = cursor.fetchall()

    if not messages:
        print("All messages already have embeddings")
        cursor.close()
        conn.close()
        return 0

The LEFT JOIN identifies messages without corresponding embeddings. Ordering by created_at processes conversations chronologically, which can be helpful for debugging and provides predictable progress tracking.

8.3.3 Batch Processing for Optimal Performance

Processing embeddings in batches dramatically improves performance by amortizing model overhead across multiple texts:

    print(f"Generating embeddings for {len(messages)} messages...")

    model = get_embedding_model()
    total_generated = 0

    # Process in batches
    for i in range(0, len(messages), batch_size):
        batch = messages[i:i + batch_size]
        uuids = [msg[0] for msg in batch]
        texts = [msg[1] for msg in batch]

        # Generate embeddings
        embeddings = model.encode(texts, convert_to_numpy=True)

The batch size of 100 is chosen to balance memory usage with processing efficiency. SentenceTransformers can process multiple texts simultaneously, significantly faster than individual encoding calls.

8.3.4 Database Insertion with Conflict Handling

The embedding insertion handles potential conflicts while tracking actual progress:

        # Insert embeddings
        for uuid, embedding in zip(uuids, embeddings):
            cursor.execute("""
                INSERT INTO message_embeddings (message_uuid, embedding)
                VALUES (%s, %s)
                ON CONFLICT (message_uuid) DO NOTHING
            """, (uuid, embedding.tolist()))

        conn.commit()
        total_generated += len(batch)

        if i % (batch_size * 5) == 0:  # Progress update every 5 batches
            print(f"Generated {total_generated} embeddings...")

    cursor.close()
    conn.close()

    print(f"Generated {total_generated} message embeddings")
    return total_generated

Converting embeddings to lists with .tolist() is necessary for PostgreSQL’s vector type. The progress updates every 500 messages (5 batches × 100) provide feedback during long processing runs without overwhelming the console.

8.4 Contextual Search with Conversational Understanding

Searching conversations requires more than finding semantically similar messages - it needs to understand and preserve the conversational context that makes individual messages meaningful.

8.4.1 Semantic Similarity Search

The core search function finds messages based on semantic similarity while retrieving essential context information:

def search_messages(db_config, query, limit=10, threshold=0.7):
    """Search for messages using semantic similarity"""
    conn = psycopg2.connect(**db_config)
    cursor = conn.cursor()

    # Generate query embedding
    model = get_embedding_model()
    query_embedding = model.encode([query])[0]

Encoding the query as a single-item list and accessing [0] ensures consistent array dimensionality with the stored embeddings.

8.4.2 Multi-Table Joins for Rich Context

The search query demonstrates the power of our normalized schema by joining across all three tables to provide comprehensive result context:

    # Search using pgvector cosine similarity
    cursor.execute("""
        SELECT
            m.uuid,
            m.text,
            m.sender,
            m.created_at,
            c.name as conversation_name,
            1 - (me.embedding <=> %s::vector) as similarity
        FROM message_embeddings me
        JOIN messages m ON me.message_uuid = m.uuid
        JOIN conversations c ON m.conversation_uuid = c.uuid
        WHERE 1 - (me.embedding <=> %s::vector) > %s
        ORDER BY similarity DESC
        LIMIT %s
    """, (query_embedding.tolist(), query_embedding.tolist(), threshold, limit))

The join sequence starts with embeddings (for efficient vector filtering), joins messages (for content and metadata), and conversations (for context). The similarity threshold filters out weakly related results, improving precision.

8.4.3 Result Formatting and Structure

The search results are formatted into a consistent structure that downstream components can rely on:

    results = cursor.fetchall()

    cursor.close()
    conn.close()

    # Format results
    formatted_results = []
    for row in results:
        formatted_results.append({
            'message_uuid': row[0],
            'text': row[1],
            'sender': row[2],
            'created_at': row[3],
            'conversation_name': row[4],
            'similarity': row[5]
        })

    return formatted_results

Returning dictionaries instead of raw tuples makes the code more maintainable and provides clear field names for consuming code.

8.4.4 Conversation Context Retrieval

Finding similar messages is only half the challenge - we need surrounding context to understand what each message means in its conversational flow:

def get_message_context(db_config, message_uuid, context_size=3):
    """Get surrounding context for a message"""
    conn = psycopg2.connect(**db_config)
    cursor = conn.cursor()

    # Get the target message details
    cursor.execute("""
        SELECT conversation_uuid, message_index, text, sender
        FROM messages
        WHERE uuid = %s
    """, (message_uuid,))

    target = cursor.fetchone()
    if not target:
        cursor.close()
        conn.close()
        return None

    conv_uuid, target_index, target_text, target_sender = target

The function first retrieves the target message’s position information, which is essential for calculating the context window bounds.

8.4.5 Context Window Calculation

The context retrieval uses the message index to extract a symmetric window around the target message:

    # Get surrounding messages
    cursor.execute("""
        SELECT
            uuid,
            text,
            sender,
            message_index,
            created_at
        FROM messages
        WHERE conversation_uuid = %s
          AND message_index >= %s
          AND message_index <= %s
        ORDER BY message_index
    """, (conv_uuid, target_index - context_size, target_index + context_size))

    context_messages = cursor.fetchall()

    cursor.close()
    conn.close()

    return {
        'target_message_uuid': message_uuid,
        'conversation_uuid': conv_uuid,
        'context_messages': [{
            'uuid': msg[0],
            'text': msg[1],
            'sender': msg[2],
            'message_index': msg[3],
            'created_at': msg[4],
            'is_target': msg[0] == message_uuid
        } for msg in context_messages]
    }

The is_target flag helps consuming code identify which message in the context window was the actual search match. A context size of 3 provides 7 total messages (3 before, the target, 3 after), usually sufficient to understand the conversational flow.

8.5 RAG Integration for Conversation History

The RAG pipeline transforms retrieved conversation fragments into coherent, contextual answers that synthesize knowledge from across your conversation history.

8.5.1 Structured Context Management

We use dataclasses to provide clean interfaces between the search and generation components:

@dataclass
class Context:
    text: str
    conversation_name: str
    sender: str
    similarity: float

This structure encapsulates all the information needed for both prompt construction and result attribution.

8.5.2 Local LLM Integration with Ollama

Our LLM integration prioritizes privacy and control by using locally-hosted models:

def call_ollama(prompt, model="llama3.1:8b"):
    """Simple Ollama API call"""
    url = "http://localhost:11434/api/generate"

    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
        "options": {
            "temperature": 0.1,
            "top_p": 0.9
        }
    }

    try:
        response = requests.post(url, json=payload, timeout=30)
        response.raise_for_status()
        return response.json()['response']
    except Exception as e:
        return f"Error calling Ollama: {str(e)}"

The low temperature (0.1) encourages factual responses based on the provided context rather than creative generation. The 30-second timeout balances response time with the ability to process complex conversational contexts.

8.5.3 Health Monitoring and Model Discovery

The health check function provides visibility into available models and helps with deployment troubleshooting:

def test_ollama():
    """Test if Ollama is running"""
    try:
        response = requests.get("http://localhost:11434/api/tags", timeout=5)
        models = response.json()
        print("Available Ollama models:", [m['name'] for m in models.get('models', [])])
        return True
    except:
        print("Ollama not running. Start with: ollama serve")
        return False

Listing available models helps users verify their preferred model is loaded and provides guidance for model management.

8.5.4 Context Retrieval and Assembly

The retrieval function bridges the search system and the generation pipeline:

def retrieve_contexts(db_config, question, max_contexts=5, threshold=0.7):
    """Retrieve relevant contexts for RAG"""
    results = search_messages(db_config, question, limit=max_contexts, threshold=threshold)

    contexts = []
    for result in results:
        contexts.append(Context(
            text=result['text'],
            conversation_name=result['conversation_name'],
            sender=result['sender'],
            similarity=result['similarity']
        ))

    return contexts

This function transforms search results into the structured Context objects that the prompt formatting function expects, providing a clean abstraction layer.

8.5.5 Conversational Prompt Engineering

Prompts for conversation-based RAG require special consideration for the personal nature of the data and the conversational context:

def format_rag_prompt(question, contexts):
    """Format contexts into a prompt for the LLM"""
    if not contexts:
        return f"Question: {question}\n\nI don't have any relevant information to answer this question."

    context_parts = []
    for i, ctx in enumerate(contexts, 1):
        context_parts.append(
            f"Context {i} (from '{ctx.conversation_name}' - {ctx.sender}, similarity: {ctx.similarity:.3f}):\n"
            f"{ctx.text}"
        )

    context_section = "\n\n".join(context_parts)

    prompt = f"""You are a helpful assistant answering questions based on the user's conversation history with Claude.
Use ONLY the information provided in the contexts below to answer the question.
If the answer cannot be found in the contexts, say so clearly.

Question: {question}

Relevant contexts from your conversations:
{context_section}

Answer based on the above contexts:"""

    return prompt

The prompt design includes several important elements:

- Source Attribution: Shows which conversation and sender each context comes from

- Similarity Scores: Helps the LLM weight information appropriately

- Constraint Enforcement: Explicitly instructs the model to use only provided context

- Conversation Framing: Acknowledges this is the user’s personal conversation history

8.5.6 Complete RAG Pipeline with Performance Monitoring

The main RAG function orchestrates the entire retrieval-to-generation pipeline:

def answer_question(db_config, question, max_contexts=5, threshold=0.7):
    """Complete RAG pipeline to answer a question"""
    print(f"\nQuestion: {question}")

    # Step 1: Retrieve relevant contexts
    start_time = time.time()
    contexts = retrieve_contexts(db_config, question, max_contexts, threshold)
    retrieval_time = (time.time() - start_time) * 1000

    print(f"Retrieved {len(contexts)} contexts in {retrieval_time:.1f}ms")

    if not contexts:
        return "I couldn't find any relevant information in your conversation history to answer this question."

Performance monitoring helps identify bottlenecks and provides user feedback during processing. Retrieval should typically complete in under 100ms with proper indexing.

    # Step 2: Format prompt
    prompt = format_rag_prompt(question, contexts)

    # Step 3: Generate answer
    start_time = time.time()
    answer = call_ollama(prompt)
    generation_time = (time.time() - start_time) * 1000

    print(f"Generated answer in {generation_time:.1f}ms")
    print(f"Total time: {retrieval_time + generation_time:.1f}ms")

    return {
        'answer': answer,
        'contexts': contexts,
        'stats': {
            'retrieval_time_ms': retrieval_time,
            'generation_time_ms': generation_time,
            'num_contexts': len(contexts)
        }
    }

Returning detailed statistics and the source contexts enables downstream applications to provide transparency about how answers were generated and from which conversations.

8.6 Complete Web API with FastAPI

Building a web API transforms our conversation search system from a command-line tool into a reusable service that can power web applications, mobile apps, or integrate with other systems.

8.6.1 FastAPI Application Structure

FastAPI provides automatic API documentation, request validation, and high-performance async capabilities:

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List

app = FastAPI(title="Claude Conversation Search", version="1.0.0")

# Add CORS for web access
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global database config
DB_CONFIG = None

The CORS middleware enables web browser access from any origin, essential for development and client-side applications. In production, the allow_origins should be restricted to known domains.

8.6.2 Request Models with Validation

Pydantic models provide automatic request validation and API documentation:

class SearchRequest(BaseModel):
    query: str
    limit: int = 10
    threshold: float = 0.7

class RAGRequest(BaseModel):
    question: str
    max_contexts: int = 5
    threshold: float = 0.7

These models define default values for optional parameters and ensure type safety. FastAPI automatically generates OpenAPI documentation from these model definitions.

8.6.3 Search Endpoint Implementation

The search endpoint provides direct access to semantic search capabilities:

@app.post("/search")
async def search_endpoint(request: SearchRequest):
    """Search messages by semantic similarity"""
    try:
        results = search_messages(DB_CONFIG, request.query, request.limit, request.threshold)
        return {"results": results}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

Using POST instead of GET allows for complex query parameters and avoids URL length limitations. The try-catch block ensures proper HTTP error responses for debugging.

8.6.4 RAG Question-Answering Endpoint

The RAG endpoint provides the complete question-answering pipeline:

@app.post("/ask")
async def ask_endpoint(request: RAGRequest):
    """Ask a question using RAG"""
    try:
        result = answer_question(DB_CONFIG, request.question, request.max_contexts, request.threshold)
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

This endpoint returns the full result structure including answer, contexts, and performance statistics, giving clients complete visibility into the response generation process.

8.6.5 System Statistics and Monitoring

The stats endpoint provides operational visibility into the system’s data and health:

@app.get("/stats")
async def stats_endpoint():
    """Get database statistics"""
    try:
        conn = psycopg2.connect(**DB_CONFIG)
        cursor = conn.cursor()

        cursor.execute("SELECT COUNT(*) FROM conversations")
        total_conversations = cursor.fetchone()[0]

        cursor.execute("SELECT COUNT(*) FROM messages")
        total_messages = cursor.fetchone()[0]

        cursor.execute("SELECT COUNT(*) FROM message_embeddings")
        total_embeddings = cursor.fetchone()[0]

        cursor.close()
        conn.close()

        return {
            "total_conversations": total_conversations,
            "total_messages": total_messages,
            "total_embeddings": total_embeddings,
            "embedding_coverage": f"{(total_embeddings/total_messages)*100:.1f}%" if total_messages > 0 else "0%"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

The embedding coverage percentage helps operators understand if embedding generation is complete and identify when new messages need processing.

8.6.6 Server Startup and Configuration

The server startup function configures the global database connection and starts the ASGI server:

def start_api(db_config):
    """Start the FastAPI server"""
    global DB_CONFIG
    DB_CONFIG = db_config

    import uvicorn
    print("\nStarting API server at http://localhost:8000")
    print("- Search endpoint: POST /search")
    print("- RAG endpoint: POST /ask")
    print("- Stats endpoint: GET /stats")
    print("- API docs: http://localhost:8000/docs")

    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")

Binding to 0.0.0.0 makes the service accessible from other machines on the network, essential for container deployment or remote access. The automatic documentation at /docs provides an interactive interface for testing endpoints.

8.7 Demonstration and Sample Data

A complete application needs robust demonstration capabilities and sample data for testing and development.

8.7.1 Realistic Sample Data Generation

The sample data mimics real conversation patterns while covering different topic areas:

def create_sample_data(db_config):
    """Create sample conversation data for testing"""
    sample_conversations = [
        {
            "uuid": "conv_001",
            "name": "Python Learning Discussion",
            "created_at": "2024-01-15T10:00:00Z",
            "messages": [
                {
                    "uuid": "msg_001",
                    "sender": "Human",
                    "text": "I'm learning Python and struggling with list comprehensions. Can you help?",
                    "created_at": "2024-01-15T10:00:00Z"
                },
                {
                    "uuid": "msg_002",
                    "sender": "Assistant",
                    "text": "List comprehensions are a concise way to create lists. The basic syntax is [expression for item in iterable]. For example: squares = [x**2 for x in range(10)] creates a list of squares from 0 to 81.",
                    "created_at": "2024-01-15T10:01:00Z"
                },
                {
                    "uuid": "msg_003",
                    "sender": "Human",
                    "text": "That's helpful! What about filtering with conditions?",
                    "created_at": "2024-01-15T10:02:00Z"
                },
                {
                    "uuid": "msg_004",
                    "sender": "Assistant",
                    "text": "You can add conditions with 'if'. For example: even_squares = [x**2 for x in range(10) if x % 2 == 0] creates squares only for even numbers.",
                    "created_at": "2024-01-15T10:03:00Z"
                }
            ]
        },

This conversation demonstrates a typical learning interaction with follow-up questions and detailed explanations. The realistic timestamps and UUIDs match Claude’s export format.

8.7.2 Multi-Topic Sample Coverage

The second sample conversation covers a different domain to demonstrate cross-topic search capabilities:

        {
            "uuid": "conv_002",
            "name": "Machine Learning Concepts",
            "created_at": "2024-01-16T14:00:00Z",
            "messages": [
                {
                    "uuid": "msg_005",
                    "sender": "Human",
                    "text": "Explain the difference between supervised and unsupervised learning",
                    "created_at": "2024-01-16T14:00:00Z"
                },
                {
                    "uuid": "msg_006",
                    "sender": "Assistant",
                    "text": "Supervised learning uses labeled training data to learn patterns, like predicting house prices from features. Unsupervised learning finds patterns in data without labels, like clustering customers into groups.",
                    "created_at": "2024-01-16T14:01:00Z"
                },
                {
                    "uuid": "msg_007",
                    "sender": "Human",
                    "text": "What are some common algorithms for each type?",
                    "created_at": "2024-01-16T14:02:00Z"
                },
                {
                    "uuid": "msg_008",
                    "sender": "Assistant",
                    "text": "Supervised: Linear regression, decision trees, random forest, SVM, neural networks. Unsupervised: K-means clustering, hierarchical clustering, PCA, DBSCAN.",
                    "created_at": "2024-01-16T14:03:00Z"
                }
            ]
        }
    ]

Having conversations on different topics (Python programming and machine learning) allows testing cross-domain search and synthesis capabilities.

8.7.3 Sample Data Processing Pipeline

The sample data is converted to the expected JSON format and processed through the normal import pipeline:

    # Convert to the format expected by import_conversations
    import tempfile
    import json

    sample_data = {"conversations": sample_conversations}

    with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
        json.dump(sample_data, f)
        temp_file = f.name

    import_conversations(db_config, temp_file)

    # Clean up
    Path(temp_file).unlink()
    print("Sample data created")

Using the standard import pipeline ensures sample data processing exercises the same code paths as real data, improving test coverage and reliability.

8.7.4 Comprehensive System Demonstration

The main function orchestrates a complete system demonstration:

def main():
    """Main demonstration of the system"""
    print("=== Minimal Claude Conversation Search & RAG System ===\n")

    # Step 1: Setup database
    print("1. Setting up database...")
    db_config = setup_database()

    # Step 2: Check if we have data, if not create sample data
    print("\n2. Checking for existing data...")
    conn = psycopg2.connect(**db_config)
    cursor = conn.cursor()
    cursor.execute("SELECT COUNT(*) FROM messages")
    message_count = cursor.fetchone()[0]
    cursor.close()
    conn.close()

    if message_count == 0:
        print("No existing data found. Creating sample data...")
        create_sample_data(db_config)
    else:
        print(f"Found {message_count} existing messages")

The demonstration intelligently handles both fresh installations and systems with existing data, making it suitable for both development and production environments.

8.7.5 Progressive Feature Demonstration

Each step of the demonstration builds on the previous ones, showing the complete system capability:

    # Step 3: Generate embeddings
    print("\n3. Generating embeddings...")
    generate_message_embeddings(db_config)

    # Step 4: Check Ollama
    print("\n4. Checking Ollama...")
    ollama_available = test_ollama()

    # Step 5: Demo search
    print("\n5. Demonstrating semantic search...")

    demo_queries = [
        "Python programming help",
        "machine learning algorithms",
        "how to learn data science"
    ]

    for query in demo_queries[:1]:  # Just show one example
        print(f"\nSearching for: '{query}'")
        results = search_messages(db_config, query, limit=3)

        for i, result in enumerate(results, 1):
            print(f"\n{i}. Similarity: {result['similarity']:.3f}")
            print(f"   From: {result['conversation_name']}")
            print(f"   {result['sender']}: {result['text'][:100]}...")

The search demonstration shows similarity scores and conversation context, helping users understand how the semantic search performs on different query types.

8.7.6 RAG Demonstration with Conditional Execution

The RAG demonstration only runs if Ollama is available, gracefully degrading the demonstration experience:

    # Step 6: Demo RAG if Ollama is available
    if ollama_available:
        print("\n6. Demonstrating RAG...")
        question = "What did I learn about Python list comprehensions?"
        result = answer_question(db_config, question)

        if isinstance(result, dict):
            print(f"\nAnswer: {result['answer'][:200]}...")
            print(f"Used {result['stats']['num_contexts']} contexts")
        else:
            print(f"Answer: {result}")
    else:
        print("\n6. Skipping RAG demo (Ollama not available)")

    # Step 7: Offer to start API
    print("\n7. Would you like to start the API server? (y/n)")
    if input().lower().startswith('y'):
        start_api(db_config)
    else:
        print("Demo completed!")

The interactive API server option allows users to immediately explore the web interface, transitioning seamlessly from command-line demonstration to web-based interaction.

8.7.7 Production Import Functionality

For real-world usage, the system provides a separate import mode for processing actual Claude conversation exports:

def import_from_file():
    """Import conversations from a Claude export file"""
    print("=== Import Claude Conversations ===\n")

    file_path = input("Enter path to conversations.json file: ").strip()
    if not file_path:
        print("No file specified")
        return

    # Setup database
    db_config = setup_database()

    # Import conversations
    total_conversations, total_messages = import_conversations(db_config, file_path)

    if total_messages > 0:
        print(f"\n3. Generating embeddings...")
        generate_message_embeddings(db_config)

    print("\nImport completed!")

This function provides a streamlined workflow for users with real conversation data, handling the complete pipeline from file import to embedding generation.

8.7.8 Application Entry Points

The main function supports multiple execution modes based on command-line arguments:

if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "import":
        import_from_file()
    else:
        main()

This design pattern allows the same script to serve multiple purposes: demonstration with sample data, or production import of real conversation files.

Conclusion: A Complete Personal Knowledge System

This conversation search and RAG system demonstrates how vector databases can transform personal AI interactions from ephemeral exchanges into a persistent, queryable knowledge base. By combining PostgreSQL’s relational capabilities with pgvector’s semantic search, we’ve built a system that understands both the structure and meaning of conversational data.

Key Architectural Achievements

Scalable Data Design: The three-table schema balances normalization with query performance, supporting everything from personal conversation archives to enterprise chat systems.

Contextual Understanding: Unlike simple document search, our system preserves conversational context, recognizing that meaning in dialogue depends on surrounding exchanges.

Complete Infrastructure: The FastAPI web service provides a full interface for integration with other applications, making conversation history a programmable resource.

Privacy-First Approach: Local storage and processing ensure sensitive conversation data never leaves your control, while still providing powerful search and synthesis capabilities.

Performance and Scale Characteristics

With proper indexing, the system handles:

- Sub-second search across tens of thousands of messages

- Batch embedding generation processing thousands of messages efficiently

- Real-time web API responses suitable for interactive applications

- Incremental updates allowing continuous growth of conversation archives

Practical Applications

This foundation enables numerous applications:

- Personal knowledge management systems that learn from your interactions

- Research assistance tools that synthesize insights across conversation history

- Learning reinforcement platforms that help rediscover previously discussed concepts

- Decision support systems that surface relevant past discussions and outcomes

Technical Extensibility

The modular architecture supports natural extensions:

- Multi-user systems with conversation sharing and privacy controls

- Advanced RAG techniques like conversation summarization and trend analysis

- Integration pipelines for other chat platforms (Slack, Discord, Teams)

- Analytics dashboards for understanding conversation patterns and knowledge evolution

The Future of Conversational Memory

As AI assistants become more capable and conversations more valuable, systems like this bridge the gap between transient interactions and permanent knowledge. Rather than starting each conversation from scratch, we can build on the rich history of previous exchanges, creating a true conversational partnership that grows smarter over time.

The complete system we’ve built here - from database design through web API - provides the foundation for this future, where our conversations with AI become a living, searchable extension of our own memory and reasoning capabilities.

Chapter 9. Vector Query Language (VQL) — Draft Specification

A note for Early Release readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 9th chapter of the final book. The GitHub repo will be made active later on.

If you’d like to be actively involved in reviewing and commenting on this draft, please reach out to the editor at shunter@oreilly.com.

Vector Query Language (VQL) is a SQL-inspired query language designed for vector databases. It combines the familiar syntax of SQL with powerful vector operations, making it intuitive for developers while providing advanced capabilities for machine learning and AI applications.

9.1.1 Motivation

VQL is motivated by the needs of three communities:

Application Developers — who are used to and expect abstractions for databases and stored collections. The proliferation of vector databases with different implementation paradigms ties an application tightly to the vector database implementation and the elements of the database’s metaphors now leak into the application and application architecture. In contrast, relational databases have mature abstractions and application developers are insulated from underlying data structures and can program to the domain objects and relationships. This abstraction and a common access layer across vendors are both missing from the vector database landscape circa 2025. VQL provides these abstractions and a common access layer in one.

Database Creators and Maintainers — DevOps, SREs, and AI/ML data engineers need to be able to load, maintain, and optimize vector databases, often from multiple vendors. Scripts to run these activities need to be customized to each vendor’s offering and different DDL/DML syntax needs to be incorporated for each vendor. Even with AI, manual inspection of these scripts is necessary as they have a high, potentially destructive, impact.

AI and ML Researchers — in GenAI, signals and systems, topological data analysis and vector embeddings amongst others — they all need to write programs just to manipulate data. VQL serves as a macro language for composing transforms easily and running experimental algorithms on the data which run in the database. This short circuits the programming and program maintenance phase and changes it to writing VQL scripts which are easier to read, debug and maintain. This automation shortens the cycle time of experiments and increases velocity.

9.2 Core Concepts

9.2.1 Data Model

Collection: Top-level container for vector data (analogous to a database schema)

Table: Structured set of vectors with associated metadata (analogous to a database table)

Vector: Numerical array with fixed dimensions

Embedding: Specialized vector representing an encoded entity

Metadata: Additional structured information stored as native columns

9.2.2 Basic Syntax Structure

SELECT fields
FROM collection_name.table_name
VECTOR_OPERATION operation_parameters
WHERE metadata_filters
ORDER BY distance_metric
LIMIT top_k;

This is the general structure for VQL queries, combining SQL’s SELECT, FROM, WHERE, ORDER BY, and LIMIT with a new VECTOR_OPERATION clause to specify vector-specific tasks.

9.3 Vector Operations

9.3.1 Similarity Search

Basic Similarity Search

-- Basic similarity search
SELECT *
FROM ecommerce.product_vectors
SIMILARITY SEARCH [1.2, 0.8, -0.2, 0.5]
USING METRIC cosine
TOP K 10;

This query performs a basic similarity search on the product_vectors table within the ecommerce collection. It finds the top 10 most similar vectors to the provided query vector [1.2, 0.8, -0.2, 0.5] using the cosine similarity metric.

Similarity Search with Metadata Filtering

-- With metadata filtering
SELECT product_name, category, price
FROM ecommerce.product_vectors
SIMILARITY SEARCH [1.2, 0.8, -0.2, 0.5]
USING METRIC euclidean
WHERE category = 'electronics' AND price < 1000
TOP K 5
THRESHOLD 0.8;

This query refines the similarity search by adding metadata filters. It retrieves the product_name, category, and priceof the top 5 products that are similar to the query vector (using Euclidean distance) and belong to the ‘electronics’ category and have a price less than 1000. Additionally, it applies a THRESHOLD of 0.8, meaning only results with a similarity score (or distance, depending on metric interpretation) beyond this threshold will be returned.

9.3.2 Hybrid Search

-- Combining vector similarity with text search
HYBRID SEARCH (
    VECTOR [1.2, 0.8, -0.2] WEIGHT 0.7,
    TEXT "machine learning" WEIGHT 0.3
)
FROM research.paper_vectors
WHERE publication_year > 2020;

This query performs a hybrid search, combining vector similarity with traditional text search. It searches the paper_vectors table in the research collection. The search is weighted, giving 70% importance to vector similarity with [1.2, 0.8, -0.2] and 30% importance to text matching “machine learning”. Results are further filtered to include only papers published after 2020.

9.3.3 Range Search

-- Find all vectors within a distance threshold
SELECT *
FROM user_data.behavior_vectors
RANGE SEARCH [user_vector]
USING METRIC cosine
THRESHOLD 0.5;

This query finds all vectors within the behavior_vectors table of the user_data collection that are within a specified distance (THRESHOLD of 0.5) from a given user_vector, using the cosine distance metric. This is useful for finding all relevant items, not just the top K.

9.3.4 Batch Operations

-- Batch similarity search
BATCH SIMILARITY SEARCH (
    SELECT query_vectors FROM user_queries.batch_requests
)
FROM ecommerce.product_vectors
TOP K 5;

This query performs a batch similarity search. Instead of a single query vector, it takes a set of query vectors retrieved from user_queries.batch_requests and finds the top 5 similar products for each query vector in the ecommerce.product_vectors table. This is efficient for processing multiple queries simultaneously.

9.4 Linear Transformations

9.4.1 Matrix Operations

Define Transformation Matrices

-- Define transformation matrices
CREATE TRANSFORM rotation_matrix AS MATRIX [
    [cos(theta), -sin(theta)],
    [sin(theta), cos(theta)]
];

This statement defines a new named transformation rotation_matrix as a matrix. This matrix can then be reused in subsequent queries for applying rotations.

Apply Transformations

-- Apply transformations
SELECT TRANSFORM(vector_data, rotation_matrix) AS rotated_vector
FROM machine_learning.feature_vectors;

This query applies the previously defined rotation_matrix transformation to the vector_data column in the feature_vectors table of the machine_learning collection. The result is a new column named rotated_vectorcontaining the transformed vectors.

Chain Multiple Transformations

-- Chain multiple transformations
SELECT TRANSFORM_CHAIN(
    vector_data,
    [scaling_matrix, rotation_matrix, projection_matrix]
) AS transformed_vector
FROM machine_learning.feature_vectors;

This query applies a sequence of transformations to vector_data. The TRANSFORM_CHAIN function applies scaling_matrix, then rotation_matrix, and finally projection_matrix in that order, outputting the final result as transformed_vector.

9.4.2 Built-in Linear Transformations

Principal Component Analysis

-- Principal Component Analysis
SELECT PCA(vector_data, dimensions=50) AS reduced_vector
FROM analytics.high_dimensional_data;

This query applies Principal Component Analysis (PCA) to vector_data in high_dimensional_data, reducing its dimensionality to 50. The resulting lower-dimensional vectors are aliased as reduced_vector.

Random Projection

-- Random Projection
SELECT RANDOM_PROJECTION(vector_data, output_dim=100, seed=42) AS projected
FROM analytics.sparse_vectors;

This query applies a Random Projection to vector_data in sparse_vectors, reducing its dimensionality to 100 with a specified random seed for reproducibility. The result is aliased as projected.

Whitening Transform

-- Whitening Transform
SELECT WHITEN(vector_data) AS whitened_vector
FROM analytics.correlated_data;

This query applies a Whitening Transform to vector_data in correlated_data, decorrelating the features and scaling them to have unit variance. The output is whitened_vector.

Normalization

-- Normalization
SELECT NORMALIZE(vector_data) AS unit_vector
FROM analytics.unnormalized_data;

This query normalizes vector_data in unnormalized_data to a unit vector (e.g., L2 normalization), meaning its length will be 1. The result is unit_vector.

9.4.3 Linear Algebra Operations

Matrix Multiplication

-- Matrix multiplication
SELECT MATMUL(matrix1, matrix2) AS product;

This query performs matrix multiplication of matrix1 and matrix2, returning the resultant matrix as product.

Eigendecomposition

-- Eigendecomposition
SELECT EIGENDECOMPOSE(correlation_matrix) AS (
    eigenvalues,
    eigenvectors
);

This query performs eigendecomposition on correlation_matrix, returning two result sets: eigenvalues (a vector of eigenvalues) and eigenvectors (a matrix where columns are eigenvectors).

Singular Value Decomposition

-- Singular Value Decomposition
SELECT SVD(data_matrix) AS (U, S, V);

This query performs Singular Value Decomposition (SVD) on data_matrix, returning three matrices: U (left singular vectors), S (singular values, typically diagonal), and V (right singular vectors).

9.5 Nonlinear Transformations

9.5.1 Hyperbolic Space Operations

Poincaré Ball Projection

-- Project vectors into Poincaré ball model
SELECT POINCARE_PROJECTION(vector_data, radius=1.0) AS hyperbolic_vector
FROM hierarchy.tree_structures;

This query projects vectors from vector_data into the Poincaré ball model of hyperbolic space with a specified radius of 1.0. This is useful for representing hierarchical data. The output is hyperbolic_vector.

Hyperbolic Distance Calculation

-- Hyperbolic distance calculation
SELECT HYPERBOLIC_DISTANCE(
    vector1,
    vector2,
    model='poincare'
) AS h_distance;

This query calculates the hyperbolic distance between vector1 and vector2 within the specified poincare model of hyperbolic space, returning the distance as h_distance.

Möbius Addition

-- Möbius addition in hyperbolic space
SELECT MOBIUS_ADD(
    hyperbolic_vec1,
    hyperbolic_vec2,
    curvature=-1.0
) AS combined_vector;

This query performs Möbius addition on two hyperbolic vectors, hyperbolic_vec1 and hyperbolic_vec2, within a hyperbolic space defined by a curvature of -1.0. The result is combined_vector.

Exponential and Logarithmic Maps

-- Exponential map
SELECT EXPMAP(tangent_vector, base_point, model='poincare') AS hyperbolic_point;

This query applies the exponential map in the Poincaré model, mapping a tangent_vector from the tangent space at a base_point to a point in hyperbolic space, returned as hyperbolic_point.

-- Logarithmic map
SELECT LOGMAP(hyperbolic_point, base_point, model='poincare') AS tangent_vector;

This query applies the logarithmic map, which is the inverse of the exponential map, mapping a hyperbolic_point back to a tangent_vector in the tangent space at a base_point within the Poincaré model.

9.5.2 Manifold Operations

Riemannian Optimization

-- Riemannian optimization on manifold
OPTIMIZE ON MANIFOLD 'hyperbolic'
OBJECTIVE minimize_hyperbolic_distance(source_vectors, target_vectors)
USING riemannian_gradient_descent(learning_rate=0.01, max_iterations=1000);

This statement initiates an optimization process on a specified hyperbolic manifold. The objective is to minimize_hyperbolic_distance between source_vectors and target_vectors using the riemannian_gradient_descent algorithm with given learning_rate and max_iterations.

Parallel Transport

-- Parallel transport along geodesics
SELECT PARALLEL_TRANSPORT(
    vector,
    start_point,
    end_point,
    manifold='poincare'
) AS transported_vector;

This query performs parallel transport of a vector from a start_point to an end_point along a geodesic path on the specified poincare manifold. The resulting vector at the end_point is transported_vector.

Geodesic Interpolation

-- Geodesic interpolation
SELECT GEODESIC_INTERPOLATE(
    point1,
    point2,
    t=0.5,
    manifold='poincare'
) AS midpoint;

This query interpolates between point1 and point2 along a geodesic on the poincare manifold. The t=0.5 specifies finding the midpoint along this geodesic, returned as midpoint.

9.5.3 Advanced Nonlinear Transforms

Kernel Transformations

-- Kernel transformations
SELECT KERNEL_TRANSFORM(
    vector_data,
    kernel='rbf',
    gamma=0.1
) AS kernel_space;

This query applies a kernel transformation (specifically, a Radial Basis Function ‘rbf’ kernel with a gamma parameter of 0.1) to vector_data, mapping it into a higher-dimensional kernel space, returned as kernel_space.

Manifold Learning

ISOMAP

-- ISOMAP manifold learning
SELECT ISOMAP(vector_data, n_neighbors=5, n_components=2) AS manifold_vector;

This query applies the ISOMAP manifold learning algorithm to vector_data, reducing its dimensionality to 2 components while preserving geodesic distances, using 5 nearest neighbors. The result is manifold_vector.

UMAP

-- UMAP manifold learning
SELECT UMAP(vector_data, n_neighbors=15, min_dist=0.1) AS umap_vector;

This query applies the UMAP manifold learning algorithm to vector_data, creating a lower-dimensional embedding with 15 nearest neighbors and a minimum distance of 0.1, returned as umap_vector.

t-SNE

-- t-SNE manifold learning
SELECT TSNE(vector_data, perplexity=30, n_components=2) AS tsne_vector;

This query applies the t-SNE (t-Distributed Stochastic Neighbor Embedding) manifold learning algorithm to vector_data, reducing it to 2 components for visualization purposes, with a perplexity of 30. The output is tsne_vector.

9.5.4 Space Conversions

Convert Between Hyperbolic Models

-- Convert between hyperbolic models
SELECT CONVERT_SPACE(
    hyperbolic_vector,
    from='poincare',
    to='lorentz'
) AS lorentz_vector;

This query converts a hyperbolic_vector from the ‘poincare’ model to the ‘lorentz’ model of hyperbolic space, returning the transformed vector as lorentz_vector.

Project Between Different Curvatures

-- Project between different curvatures
SELECT CURVATURE_PROJECT(
    vector_data,
    source_curvature=-1.0,
    target_curvature=-0.5
) AS projected;

This query projects vector_data from a hyperbolic space with a source_curvature of -1.0 to one with a target_curvature of -0.5, returning the new vector as projected.

9.6 Vector Functions and Aggregations

9.6.1 Vector Functions

Basic Vector Operations

-- Get vector dimension
SELECT DIMENSION(embedding) AS vector_dim;

This query returns the dimensionality (number of components) of the embedding vector as vector_dim.

-- Calculate distance between vectors
SELECT DISTANCE(vector1, vector2, 'cosine') AS similarity;

This query calculates the cosine distance between vector1 and vector2, returning the result as similarity.

-- Concatenate vectors
SELECT CONCAT(vector1, vector2) AS combined_vector;

This query concatenates vector1 and vector2 into a single, longer vector, returned as combined_vector.

-- Dot product
SELECT DOT(vector1, vector2) AS dot_product;

This query calculates the dot product of vector1 and vector2, returning the scalar result as dot_product.

Vector Arithmetic

-- Vector addition
SELECT vector_data + influence_vector AS adjusted_vector;

This query performs vector addition, adding influence_vector to vector_data element-wise, returning the result as adjusted_vector.

-- Scalar multiplication
SELECT vector_data * scalar_value AS scaled_vector;

This query performs scalar multiplication, multiplying each element of vector_data by scalar_value, returning the result as scaled_vector.

9.6.2 Vector Aggregations

Compute Centroids

-- Compute centroids
SELECT AVG(user_vector) AS centroid
FROM analytics.user_profiles
GROUP BY user_category;

This query computes the average vector (centroid) for user_vector within each user_category in the user_profiles table.

Find Representative Vectors

-- Find representative vectors
SELECT MEDIAN_VECTOR(feature_vector) AS representative
FROM analytics.product_features
GROUP BY product_category;

This query computes the median vector (representative) for feature_vector within each product_category in the product_features table. The median vector is typically more robust to outliers than the mean.

9.7 Clustering and Analytics

9.7.1 Clustering Operations

K-means Clustering

-- K-means clustering
CLUSTER customer_data.behavior_vectors
USING KMEANS
CLUSTERS 5
OUTPUT AS customer_segments;

This statement performs K-means clustering on behavior_vectors in the customer_data collection, partitioning the data into 5 clusters. The cluster assignments or cluster centroids are output as customer_segments.

Hierarchical Clustering

-- Hierarchical clustering
CLUSTER content.document_vectors
USING HIERARCHICAL
LINKAGE 'ward'
OUTPUT AS doc_hierarchy;

This statement performs hierarchical clustering on document_vectors in the content collection, using the ‘ward’ linkage method. The resulting dendrogram or cluster assignments are output as doc_hierarchy.

9.7.2 Quality Control and Sampling

Diversity Sampling

-- Diversity sampling
SELECT *
FROM search.result_vectors
SIMILARITY SEARCH [query_vector]
WITH DIVERSITY FACTOR 0.3
TOP K 10;

This query combines a similarity search with diversity sampling. It finds the top 10 results similar to query_vector from result_vectors, but with a DIVERSITY FACTOR of 0.3, meaning it prioritizes returning diverse results in addition to highly similar ones.

Outlier Detection

-- Outlier detection
SELECT *, OUTLIER_SCORE(behavior_vector) AS anomaly_score
FROM analytics.user_behavior
WHERE OUTLIER_SCORE(behavior_vector) > 0.95;

This query identifies outliers in user_behavior by calculating an OUTLIER_SCORE for each behavior_vector. It then selects all data where this score exceeds 0.95, effectively flagging potential anomalies.

9.8 Index Management

9.8.1 Index Creation

Create Vector Index with Specific Parameters

-- Create vector index with specific parameters
CREATE INDEX product_vectors_idx
ON ecommerce.product_vectors
USING ALGORITHM 'hnsw'
WITH (
    M = 16,
    ef_construction = 200,
    distance_metric = 'cosine'
);

This statement creates an index named product_vectors_idx on the product_vectors table for efficient vector similarity searches. It specifies the hnsw algorithm with particular construction parameters (M, ef_construction) and the cosine distance_metric.

Create Composite Index for Filtered Searches

-- Create composite index for filtered searches
CREATE INDEX category_vector_idx
ON ecommerce.product_vectors (category, vector_data)
USING ALGORITHM 'ivfflat'
WITH (lists = 1000);

This statement creates a composite index category_vector_idx on both the category metadata column and the vector_data column. It uses the ivfflat algorithm, configured with 1000 lists, which is beneficial for queries that filter by category before performing vector searches.

9.8.2 Index Optimization

Rebuild Index

-- Rebuild index with new parameters
REBUILD INDEX product_vectors_idx
WITH (M = 32, ef_construction = 400);

This statement rebuilds the existing product_vectors_idx with new, potentially optimized, parameters for the HNSW algorithm (M and ef_construction), which can improve search performance or accuracy.

Analyze Index Performance

-- Analyze index performance
ANALYZE INDEX product_vectors_idx;

This statement analyzes the product_vectors_idx to gather statistics and provide insights into its performance, which can help in further optimization efforts.

9.9 Advanced Features

9.9.1 Custom Transform Functions

Define Custom Transformation

-- Define custom transformation
CREATE TRANSFORM normalize_and_scale AS (
    SELECT NORMALIZE(vector_data) * scale_factor AS result
    WHERE scale_factor = 1.5
);

This statement defines a reusable custom transformation named normalize_and_scale. It normalizes vector_data and then scales it by a scale_factor of 1.5, returning the result.

Composite Transformations

-- Composite transformations
CREATE TRANSFORM feature_extractor AS
    CHAIN(
        NORMALIZE(),
        PCA(dimensions=50),
        POINCARE_PROJECTION(radius=1.0)
    );

This statement creates a composite transformation feature_extractor by chaining three built-in transformations: NORMALIZE(), PCA() to 50 dimensions, and POINCARE_PROJECTION() with a radius of 1.0. This allows a complex sequence of operations to be applied with a single named transform.

9.9.2 Optimization and Learning

Learn Optimal Transformations

-- Learn optimal transformations
OPTIMIZE TRANSFORM learning_projection
OBJECTIVE minimize_distance(source_vectors, target_vectors)
USING gradient_descent(learning_rate=0.01, max_iterations=1000);

This statement initiates a process to learn an optimal transformation named learning_projection. The objective is to minimize_distance between source_vectors and target_vectors, using gradient_descent as the optimization algorithm with specified learning rate and iterations.

Validate Transformations

-- Validate transformations
VALIDATE TRANSFORM my_transform
CHECK (
    orthogonality = TRUE,
    determinant != 0,
    condition_number < 100
);

This statement validates the properties of a defined transform named my_transform. It checks if the transformation maintains orthogonality, has a non-zero determinant, and its condition_number is less than 100, ensuring its mathematical suitability.

9.10 Type System

9.10.1 Vector Types

VECTOR(dim)              -- Fixed-dimension vector
DYNAMIC_VECTOR           -- Variable-dimension vector
EMBEDDING                -- Specialized vector with metadata
HYPERBOLIC_VECTOR(model, dim)  -- Vector in hyperbolic space
MANIFOLD_POINT(type, coords)   -- Point on a specific manifold

These lines define various vector types supported by VQL, including fixed-dimension vectors (VECTOR(dim)), variable-dimension vectors (DYNAMIC_VECTOR), EMBEDDINGs (which include metadata), and specialized vectors for hyperbolic spaces and manifolds.

9.10.2 Transform Types

MATRIX(rows, cols)       -- Fixed-dimension matrix
TRANSFORM                -- Function type for transformations
TRANSFORM_CHAIN          -- Composite transformation
DISTANCE                 -- Numeric type for similarity scores

These lines define the types related to transformations, including fixed-dimension MATRIX, generic TRANSFORM functions, composite TRANSFORM_CHAINs, and DISTANCE as a numeric type for similarity scores.

9.10.3 Space Definitions

Define Custom Spaces

-- Define custom spaces
CREATE SPACE hyperbolic_space (
    model = 'poincare',
    curvature = -1.0,
    dimension = 50
);

This statement defines a custom named SPACE called hyperbolic_space with specific properties, including its model('poincare'), curvature (-1.0), and dimension (50). This allows for consistent handling of vectors within this defined space.

Validate Constraints

-- Validate constraints
VALIDATE IN hyperbolic_space
CHECK (
    point_norm < 1.0,
    satisfies_hyperbolic_constraints = TRUE
);

This statement validates vectors or points within the hyperbolic_space against defined constraints, specifically checking if the point_norm is less than 1.0 (a property of the Poincaré ball model) and if other hyperbolic_constraints are satisfied.

9.11 Error Handling and Validation

9.11.1 Dimension Handling

-- Handle dimension mismatches
ON DIMENSION MISMATCH
    THEN PAD_ZEROS | TRUNCATE | REJECT;

This statement specifies how VQL should handle situations where vector dimensions do not match requirements. Options include PAD_ZEROS (add zeros to match dimension), TRUNCATE (remove elements to match dimension), or REJECT (throw an error).

-- Invalid vector handling
ON INVALID VECTOR
    THEN NULL | REJECT | DEFAULT [default_vector];

This statement defines error handling for INVALID VECTOR inputs (e.g., containing NaN or infinite values). Options include returning NULL, REJECTing the operation, or substituting a DEFAULT predefined default_vector.

9.11.2 Quality Assurance

Validate Vector Quality

-- Validate vector quality
SELECT *
FROM analytics.user_data
WHERE VECTOR_QUALITY(user_vector) > 0.9;

This query selects data from user_data where the VECTOR_QUALITY score of user_vector exceeds 0.9, indicating high-quality or well-formed vectors.

Check for NaN or Infinite Values

-- Check for NaN or infinite values
SELECT *
FROM analytics.processed_data
WHERE IS_VALID_VECTOR(vector_data) = TRUE;

This query selects data from processed_data where vector_data is considered valid (i.e., does not contain NaN or infinite values), using the IS_VALID_VECTOR function.

9.12 Performance Considerations

9.12.1 Query Optimization

Lazy evaluation of vector operations

Automatic batching for large-scale operations

Index-aware query planning

Parallel processing capabilities

9.12.2 Caching Strategies

Frequently accessed vector caching

Transform result caching

Query plan caching

9.12.3 Memory Management

Streaming processing for large datasets

Efficient vector storage formats

Garbage collection for temporary vectors

9.13 Implementation Guidelines

9.13.1 Backend Integration

Support for multiple vector database backends

Pluggable distance metric implementations

Extensible transformation framework

9.13.2 API Design

RESTful API for remote query execution

Streaming results for large result sets

Async query execution capabilities

9.13.3 Security and Access Control

Role-based access to collections and tables

Query-level security policies

Vector data encryption support

9.14 Example Use Cases

9.14.1 Recommendation Systems

SELECT product_id, product_name, similarity_score
FROM ecommerce.product_vectors
SIMILARITY SEARCH (
    SELECT user_vector FROM analytics.user_profiles WHERE user_id = 12345
)
WHERE category IN ('electronics', 'books')
TOP K 20;

This query implements a recommendation system. It finds the top 20 product_id, product_name, and their similarity_score from ecommerce.product_vectors. The search is based on the user_vector of user_id = 12345, and results are filtered to only include products from the ‘electronics’ or ‘books’ categories.

9.14.2 Semantic Search

SELECT document_id, title, relevance
FROM content.document_vectors
HYBRID SEARCH (
    VECTOR [text_vector] WEIGHT 0.8,
    TEXT "artificial intelligence" WEIGHT 0.2
)
WHERE publication_date > '2020-01-01'
TOP K 50;

This query performs a semantic search for documents. It searches content.document_vectors using a hybrid approach: 80% weight on vector similarity with a given text_vector and 20% weight on text matching “artificial intelligence”. It retrieves the top 50 document_id, title, and relevance for documents published after ’2020-01-01’.

9.14.3 Anomaly Detection

SELECT user_id, behavior_vector, anomaly_score
FROM analytics.user_behavior
WHERE OUTLIER_SCORE(behavior_vector) > 0.95
ORDER BY anomaly_score DESC;

This query performs anomaly detection on user behavior. It selects user_id, behavior_vector, and an anomaly_score from analytics.user_behavior. It filters for users whose behavior_vector has an OUTLIER_SCORE greater than 0.95 and orders the results by the highest anomaly_score.

9.14.4 Hierarchical Clustering

SELECT product_id, category,
       HYPERBOLIC_HIERARCHY_EMBED(feature_vector, depth=3) AS hierarchical_pos
FROM ecommerce.product_features
ORDER BY HYPERBOLIC_TREE_DISTANCE(hierarchical_pos, root_category);

This query demonstrates hierarchical clustering using hyperbolic embeddings. It selects product_id, category, and a hierarchical_pos which is a hyperbolic embedding of the feature_vector at a depth of 3. Results are ordered by the HYPERBOLIC_TREE_DISTANCE from a root_category, allowing for exploration of hierarchical relationships.

9.15 Conclusion

This draft specification aims at providing a foundation for implementing vector query capabilities while maintaining SQL-like familiarity and extending into advanced mathematical operations. VQL bridges the gap between traditional database querying and modern vector operations, making it accessible to application developers while providing powerful capabilities for data engineering and AI/ML experiments.

Chapter 10. Advanced Embedding Techniques

A note for Early Release readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 10th chapter of the final book. The GitHub repo will be made active later on.

If you’d like to be actively involved in reviewing and commenting on this draft, please reach out to the editor at @oreilly.com.

Modern search systems need to handle multiple content types: photos, text snippets, and complex documents. Traditional approaches required separate systems for each modality, leading to fragmented search experiences and duplicated infrastructure. This chapter shows how recent advances in embedding technology enable unified search across diverse content types.

We’ll implement three specialized embedding techniques, each optimized for different content:

CLIP for searching images with natural language queries

Matryoshka embeddings for flexible-dimension text search

ColPali for searching within PDF documents without text extraction

By the end, you’ll have a single search system that can find relevant content across all three modalities, understanding both the technical implementation and the conceptual foundations that make each approach effective.

Section 1 Common Utilities

Let’s start by defining utilities we’ll use throughout the chapter.

1.1 Device Detection for Hardware Acceleration

All three models benefit from GPU acceleration. This utility detects the best available device:

import torch

def get_device():
    """
    Detect best available device for PyTorch models.
    Priority: Apple Silicon MPS > NVIDIA CUDA > CPU
    """
    if torch.backends.mps.is_available():
        return "mps"
    elif torch.cuda.is_available():
        return "cuda"
    else:
        return "cpu"

# Use throughout the chapter
device = get_device()

print(f"Using device: {device}”)

1.2 MaxSim Scoring Function

ColPali uses MaxSim (Maximum Similarity) scoring to compare multi-vector embeddings. We’ll define it once here:

import numpy as np

def maxsim_score(query_embeddings, document_embeddings):
    """
    Compute MaxSim score between query and document multi-vector embeddings.

    For each query token, find the maximum similarity to any document patch,
    then sum across all query tokens.

    Args:
        query_embeddings: (num_query_tokens, dim) array
        document_embeddings: (num_patches, dim) array

    Returns:
        float: MaxSim score
    """
    # Normalize embeddings for cosine similarity
    query_norm = query_embeddings / (np.linalg.norm(query_embeddings, axis=1, keepdims=True) + 1e-8)
    doc_norm = document_embeddings / (np.linalg.norm(document_embeddings, axis=1, keepdims=True) + 1e-8)

    # Compute similarity matrix: (num_query_tokens, num_patches)
    similarity_matrix = np.matmul(query_norm, doc_norm.T)

    # For each query token, find max similarity to any document patch
    max_similarities = np.max(similarity_matrix, axis=1)

    # Sum across query tokens
    return float(np.sum(max_similarities))

Section 2 Image Search with CLIP

2.1 Understanding CLIP

CLIP (Contrastive Language-Image Pre-training) represents a fundamental shift in how we think about image understanding. Unlike traditional computer vision models trained to classify images into fixed categories, CLIP learns to understand the relationship between images and natural language descriptions. This enables searching images with arbitrary text queries without retraining.

The key insight behind CLIP is that images and their descriptions should map to nearby points in a shared embedding space. By training on hundreds of millions of image-text pairs from the internet, CLIP learns rich visual and semantic representations that generalize to new concepts it has never explicitly seen.

2.2 When to Use CLIP

CLIP excels in scenarios where you need to search visual content with natural language queries, particularly when query patterns are diverse and unpredictable. Its zero-shot capability makes it invaluable when you can’t train on your specific data, and its cross-modal understanding enables matching images to text or vice versa. This flexibility makes CLIP the go-to choice for applications that require understanding the semantic relationship between visual and textual content without predefined categories.

Common applications include photo library search where users might query with natural descriptions like “sunset over mountains”, e-commerce product discovery systems that need to understand queries like “red leather jacket with zippers”, content moderation platforms that must find images matching specific policy-violating descriptions, and visual question answering systems that bridge the gap between what users ask and what images contain. In each of these cases, CLIP’s ability to understand arbitrary text-image relationships without task-specific training provides a significant advantage over traditional computer vision approaches.

2.3 Limitations and Tradeoffs

CLIP’s main limitations stem from its fixed 512-dimensional embeddings, which may be excessive overhead for simple use cases where smaller representations would suffice. The model struggles with fine-grained distinctions like differentiating between specific dog breeds or understanding technical diagrams, as it was trained on general web data rather than specialized domains. Additionally, CLIP is computationally expensive compared to traditional keyword search, requiring GPU acceleration for reasonable performance, and inherits biases from its web-scale training data that can lead to problematic results in certain contexts.

The key tradeoffs when using CLIP involve balancing accuracy against speed, as larger variants like ViT-L/14 provide better results but at significant computational cost. There’s also a fundamental tension between generalization and specialization - while CLIP works broadly across many domains out of the box, it may require fine-tuning to match the performance of specialized models on domain-specific tasks. Storage is another consideration, with each image requiring 2KB for embeddings (512 floats × 4 bytes), which adds up quickly for large image collections.

2.4 Loading CLIP

from transformers import CLIPModel, CLIPProcessor
from PIL import Image

# Load model and processor
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
clip_model.eval()

2.5 Basic Image-Text Matching

# Load an image
image = Image.open("photo.jpg")

# Define possible descriptions
texts = ["a cat sleeping", "a dog playing", "a bird flying"]

# Process inputs
inputs = clip_processor(text=texts, images=image, return_tensors="pt", padding=True)
inputs = {k: v.to(device) for k, v in inputs.items()}

# Get similarities
with torch.no_grad():
    outputs = clip_model(**inputs)
    logits_per_image = outputs.logits_per_image
    probs = logits_per_image.softmax(dim=1)

# Display results
for text, prob in zip(texts, probs[0]):
    print(f"{text}: {prob.item():.2%}")

2.6 Batch Processing Images

For efficiency, process multiple images at once:

from pathlib import Path

def encode_images_batch(image_paths, batch_size=32):
    """Encode multiple images efficiently with CLIP"""
    all_embeddings = []

    for i in range(0, len(image_paths), batch_size):
        batch_paths = image_paths[i:i+batch_size]
        images = [Image.open(p) for p in batch_paths]

        inputs = clip_processor(images=images, return_tensors="pt", padding=True)
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            image_features = clip_model.get_image_features(**inputs)
            # Normalize for cosine similarity
            image_features = image_features / image_features.norm(dim=-1, keepdim=True)

        all_embeddings.append(image_features.cpu().numpy())

    return np.vstack(all_embeddings)

2.7 Storing Image Embeddings in PostgreSQL

import psycopg2

# Connect and create table
conn = psycopg2.connect("dbname=image_search user=postgres password=password")
cur = conn.cursor()

cur.execute("""
    CREATE TABLE IF NOT EXISTS photos (
        id SERIAL PRIMARY KEY,
        file_path TEXT,
        embedding vector(512)
    )
""")

# Create HNSW index for fast similarity search
cur.execute("""
    CREATE INDEX IF NOT EXISTS photos_embedding_idx
    ON photos USING hnsw (embedding vector_cosine_ops)
""")
conn.commit()

2.8 Indexing and Searching Photos

# Index photos
image_paths = list(Path("photos").glob("*.jpg"))
embeddings = encode_images_batch(image_paths)

for path, embedding in zip(image_paths, embeddings):
    cur.execute(
        "INSERT INTO photos (file_path, embedding) VALUES (%s, %s)",
        (str(path), embedding.tolist())
    )
conn.commit()

# Search with text query
def search_photos(query_text, top_k=5):
    """Search photos using natural language"""
    # Encode text query
    inputs = clip_processor(text=[query_text], return_tensors="pt", padding=True)
    inputs = {k: v.to(device) for k, v in inputs.items()}

    with torch.no_grad():
        text_features = clip_model.get_text_features(**inputs)
        text_features = text_features / text_features.norm(dim=-1, keepdim=True)

    query_embedding = text_features[0].cpu().numpy()

    # Search database using cosine distance
    cur.execute("""
        SELECT file_path, 1 - (embedding <=> %s::vector) AS similarity
        FROM photos
        ORDER BY embedding <=> %s::vector
        LIMIT %s
    """, (query_embedding.tolist(), query_embedding.tolist(), top_k))

    return cur.fetchall()

# Example searches
results = search_photos("sunset over ocean")
for path, similarity in results:
    print(f"{path}: {similarity:.3f}")

Key Insight: CLIP’s 512-dimensional embeddings capture both visual and semantic information, enabling searches like “people laughing at a party” or “red sports car in the city” without any custom training.

Section 3: Matryoshka Embeddings for Text

Understanding Matryoshka Embeddings

Matryoshka Representation Learning addresses a fundamental challenge in vector search: the rigid tradeoff between embedding size and quality. Traditional models produce fixed-size embeddings, forcing you to choose between accuracy (large embeddings) and efficiency (small embeddings) at deployment time.

The Matryoshka approach, named after Russian nesting dolls, creates embeddings where information is hierarchically organized. The first dimensions capture the most important semantic information, with each subsequent dimension adding progressively finer details. This means a 768-dimensional Matryoshka embedding actually contains valid 512D, 256D, 128D, and even 64D representations within it.

This hierarchical structure emerges naturally during training through a modified loss function that optimizes accuracy at multiple truncation points simultaneously. The result is remarkable: you can achieve 95% of the full accuracy using just 128 dimensions, or 90% with only 64 dimensions.

When to Use Matryoshka Embeddings

Matryoshka embeddings shine in scenarios where you need to search millions or billions of documents while maintaining flexibility in your accuracy-speed tradeoffs. They’re particularly valuable when different queries have different latency requirements - some users need instant results while others can wait for higher accuracy. When storage costs are a concern or you want to implement sophisticated multi-stage retrieval pipelines, Matryoshka’s ability to use the same embedding at different dimensions becomes a game-changer. The approach also excels when resource constraints vary across deployment environments, allowing you to use smaller dimensions on mobile devices while leveraging full embeddings on powerful servers.

Common applications include large-scale document retrieval systems where the sheer volume of data makes traditional approaches prohibitive, mobile or edge deployments where limited resources demand efficient representations, real-time search systems with strict latency requirements that can’t afford full embedding comparisons for every query, and hierarchical filtering in recommendation systems where initial broad matching can use small embeddings before detailed ranking. In each case, Matryoshka’s flexibility allows you to optimize for your specific constraints without maintaining multiple separate models.

Limitations and Tradeoffs

The primary limitation of Matryoshka embeddings is that not all pre-trained models support truncation - you need models specifically trained with the Matryoshka loss function. While the accuracy loss is minimal at moderate truncation levels, aggressive truncation to very small dimensions can impact performance noticeably. Successfully deploying Matryoshka also requires a deep understanding of your accuracy/speed requirements to choose appropriate truncation points, and you may need to retrain existing models if they weren’t originally trained with the Matryoshka objective.

The core tradeoff with Matryoshka involves balancing dimension against accuracy - while 128D typically preserves about 95% of full performance and 64D maintains around 90%, these percentages can vary significantly based on your specific data and task. There’s also added complexity compared to fixed embeddings, as you need to manage truncation logic and potentially multiple index structures. Finally, you face a storage versus computation decision: you can store full embeddings and truncate at query time for maximum flexibility, or pre-compute and store multiple truncations for faster retrieval at the cost of increased storage.

3.1 Loading Matryoshka Models

from sentence_transformers import SentenceTransformer

# Load Matryoshka-capable model
text_model = SentenceTransformer("nomic-ai/nomic-embed-text-v1.5")

texts = [
    "Machine learning is a subset of artificial intelligence",
    "Deep learning uses neural networks with multiple layers"
]

# Encode at different dimensions
embedding_128d = text_model.encode(texts, truncate_dim=128)
embedding_256d = text_model.encode(texts, truncate_dim=256)
embedding_768d = text_model.encode(texts)  # Full dimension

print(f"128D shape: {embedding_128d.shape}")  # (2, 128)
print(f"768D shape: {embedding_768d.shape}")  # (2, 768)

# Important: First 128 dims of 768D embedding equal the 128D embedding
assert np.allclose(embedding_768d[:, :128], embedding_128d)

3.2 Multi-Dimension Database Schema

Store embeddings at multiple dimensions for tiered search:

# Create table with multiple embedding columns
cur.execute("""
    CREATE TABLE IF NOT EXISTS text_items (
        id SERIAL PRIMARY KEY,
        content TEXT,
        embedding_128 vector(128),
        embedding_256 vector(256),
        embedding_768 vector(768)
    )
""")

# Index each dimension
for dim in [128, 256, 768]:
    cur.execute(f"""
        CREATE INDEX IF NOT EXISTS text_embedding_{dim}_idx
        ON text_items USING hnsw (embedding_{dim} vector_cosine_ops)
    """)
conn.commit()

3.3 Indexing Text with Matryoshka

documents = [
    "Introduction to vector databases and similarity search",
    "Advanced indexing techniques for high-dimensional vectors",
    "Building scalable retrieval systems with embeddings"
]

for doc in documents:
    # Encode once at full dimension
    full_embedding = text_model.encode(doc)

    # Truncate to each target dimension
    emb_128 = full_embedding[:128].tolist()
    emb_256 = full_embedding[:256].tolist()
    emb_768 = full_embedding.tolist()

    cur.execute("""
        INSERT INTO text_items (content, embedding_128, embedding_256, embedding_768)
        VALUES (%s, %s, %s, %s)
    """, (doc, emb_128, emb_256, emb_768))

conn.commit()

3.4 Tiered Search Strategy

The power of Matryoshka embeddings becomes clear in multi-stage retrieval. This approach uses progressively higher dimensions to refine results, dramatically reducing computational cost while maintaining accuracy:

def search_text_tiered(query, top_k=10):
    """
    Three-tier search: 128D → 256D → 768D
    Each stage narrows candidates with increasing accuracy
    """
    # Encode query once
    query_embedding = text_model.encode(query)

    # Stage 1: Fast filtering with 128D (1000 candidates)
    query_128 = query_embedding[:128].tolist()
    cur.execute("""
        SELECT id, content, 1 - (embedding_128 <=> %s::vector) AS similarity
        FROM text_items
        ORDER BY embedding_128 <=> %s::vector
        LIMIT 1000
    """, (query_128, query_128))

    stage1_ids = [row[0] for row in cur.fetchall()]

    # Stage 2: Refinement with 256D (100 candidates)
    query_256 = query_embedding[:256].tolist()
    cur.execute("""
        SELECT id, content, 1 - (embedding_256 <=> %s::vector) AS similarity
        FROM text_items
        WHERE id = ANY(%s)
        ORDER BY embedding_256 <=> %s::vector
        LIMIT 100
    """, (query_256, stage1_ids, query_256))

    stage2_ids = [row[0] for row in cur.fetchall()]

    # Stage 3: Final ranking with 768D (top_k results)
    query_768 = query_embedding.tolist()
    cur.execute("""
        SELECT id, content, 1 - (embedding_768 <=> %s::vector) AS similarity
        FROM text_items
        WHERE id = ANY(%s)
        ORDER BY embedding_768 <=> %s::vector
        LIMIT %s
    """, (query_768, stage2_ids, query_768, top_k))

    return cur.fetchall()

# Example usage
results = search_text_tiered("How do vector databases work?")
for doc_id, content, similarity in results:
    print(f"{similarity:.3f}: {content[:60]}...")

Performance Benefit: Tiered search computes expensive 768D similarity only for final candidates, not the entire dataset. On a million-document corpus, this can reduce computation by 100x while maintaining 98%+ recall.

Section 4: Document Search with ColPali

4.1 Understanding ColPali

ColPali represents a paradigm shift in document retrieval. Traditional document search systems follow a fragile pipeline: extract text via OCR, chunk it into passages, embed each chunk, and hope nothing important was lost in the process. This approach fails catastrophically on documents where visual layout matters - tables become jumbled text, figures are ignored, and mathematical equations turn into gibberish.

ColPali takes a radically different approach: treat each document page as an image and use a Vision Language Model to understand it holistically. The model generates multiple embeddings per page (one per image patch), preserving spatial relationships and visual cues. When searching, it uses late interaction - comparing each query token to each document patch and aggregating the scores.

This approach works because modern Vision Language Models like PaliGemma have learned to read and understand visual content during pre-training. ColPali fine-tunes these models specifically for retrieval, teaching them to generate embeddings optimized for search rather than generation.

4.2 When to Use ColPali

ColPali excels when documents contain complex layouts with tables, figures, and diagrams that traditional text extraction would mangle or ignore entirely. It’s the ideal choice when visual formatting carries semantic meaning - such as highlighted text, color coding, or spatial relationships between elements that convey important information. The approach becomes essential when OCR quality is poor or inconsistent, as it bypasses text extraction entirely, and when documents mix text with visual elements like charts, equations, or handwritten annotations. For any use case where preserving exact layout is important for understanding the document’s meaning, ColPali’s visual approach offers capabilities that text-only methods simply cannot match.

Common applications include scientific paper retrieval where equations, figures, and experimental results tables are as important as the text; financial report search where quarterly results are presented in complex formatted tables and charts; technical documentation filled with architecture diagrams and code snippets; form and invoice processing where the spatial layout defines field relationships; and slide deck search where the visual design and layout are integral to the content’s meaning. In each of these scenarios, ColPali’s ability to understand documents as visual entities rather than pure text streams provides dramatically better search results.

4.3 Limitations and Tradeoffs

ColPali’s most significant limitation is its storage requirement - needing 100-500x more space than single-vector approaches due to its multi-vector representation, with each page requiring approximately 512KB compared to 2KB for traditional embeddings. The approach is also considerably slower than traditional text search, with query latencies of 100-500ms compared to 10-50ms for standard vector search, making it unsuitable for applications requiring instant responses. Additionally, ColPali is limited to visual understanding and cannot leverage document metadata like author, date, or tags that might be available in your system, and the model requires significant GPU memory for inference, making it expensive to deploy at scale.

The fundamental tradeoff with ColPali is storage versus accuracy - while the multi-vector storage enables incredibly fine-grained matching that can identify specific diagrams or tables within documents, this comes at the cost of storing hundreds of vectors per page. There’s also an inherent speed versus quality tradeoff where the late interaction mechanism that makes ColPali so effective at finding relevant content also makes it slower than approaches that compare single vectors. Finally, while ColPali works on any visual document out of the box, achieving optimal performance on specialized document types like medical records or legal contracts may require fine-tuning on domain-specific data, adding complexity to the deployment process.

4.4 Installing ColPali

pip install colpali-engine torch pillow pdf2image

4.4 Loading ColPali

from colpali_engine.models import ColPali, ColPaliProcessor
from pdf2image import convert_from_path

# Load model
colpali_model = ColPali.from_pretrained(
    "vidore/colpali-v1.2",
    torch_dtype=torch.float16 if device != "cpu" else torch.float32
).to(device)

colpali_processor = ColPaliProcessor.from_pretrained("vidore/colpali-v1.2")
colpali_model.eval()

4.5 Processing PDF Pages

ColPali converts each page to an image and generates multi-vector embeddings:

def process_pdf(pdf_path):
    """Convert PDF pages to multi-vector embeddings"""
    # Convert PDF to images (one per page)
    images = convert_from_path(pdf_path, dpi=150)

    page_embeddings = []
    for page_num, image in enumerate(images):
        inputs = colpali_processor.process_images([image]).to(device)

        with torch.no_grad():
            # Returns multi-vector embeddings: (num_patches, 128)
            embeddings = colpali_model(**inputs)

        page_embeddings.append({
            'page_num': page_num,
            'embeddings': embeddings[0].cpu().numpy(),
            'num_patches': embeddings.shape[1]
        })

    return page_embeddings

# Process a research paper
paper_embeddings = process_pdf("papers/attention_is_all_you_need.pdf")
print(f"Processed {len(paper_embeddings)} pages")
print(f"Page 1 has {paper_embeddings[0]['num_patches']} patch embeddings")

4.6 Encoding Text Queries

def encode_query(query_text):
    """Encode text query into multi-vector representation"""
    inputs = colpali_processor.process_queries([query_text]).to(device)

    with torch.no_grad():
        query_embeddings = colpali_model(**inputs)

    return query_embeddings[0].cpu().numpy()

# Example
query = encode_query("transformer architecture with attention mechanism")
print(f"Query has {query.shape[0]} token embeddings")

4.7 Searching Documents with MaxSim

Use the MaxSim function we defined earlier:

# Score query against document pages
query = encode_query("transformer attention mechanism")

for page in paper_embeddings:
    score = maxsim_score(query, page['embeddings'])
    print(f"Page {page['page_num']}: {score:.3f}")

4.8 Database Storage for Documents

Store multi-vector embeddings as serialized numpy arrays:

import io

# Create documents table
cur.execute("""
    CREATE TABLE IF NOT EXISTS documents (
        id SERIAL PRIMARY KEY,
        file_path TEXT,
        page_num INTEGER,
        num_patches INTEGER,
        embedding_data BYTEA
    )
""")

cur.execute("""
    CREATE INDEX IF NOT EXISTS doc_path_page_idx
    ON documents(file_path, page_num)
""")
conn.commit()

# Store document page
def store_document_page(file_path, page_num, embeddings):
    """Store multi-vector embeddings for a document page"""
    buffer = io.BytesIO()
    np.save(buffer, embeddings)

    cur.execute("""
        INSERT INTO documents (file_path, page_num, num_patches, embedding_data)
        VALUES (%s, %s, %s, %s)
    """, (file_path, page_num, embeddings.shape[0], buffer.getvalue()))

# Index documents
papers = Path("papers").glob("*.pdf")
for paper in papers:
    page_embeddings = process_pdf(paper)
    for page in page_embeddings:
        store_document_page(str(paper), page['page_num'], page['embeddings'])
    conn.commit()

4.9 Searching Stored Documents

def search_documents(query_text, top_k=5):
    """Search all stored documents"""
    query_embeddings = encode_query(query_text)

    # Retrieve all document pages
    cur.execute("SELECT id, file_path, page_num, embedding_data FROM documents")

    results = []
    for row in cur.fetchall():
        doc_id, file_path, page_num, embedding_bytes = row

        # Deserialize embeddings
        buffer = io.BytesIO(embedding_bytes)
        doc_embeddings = np.load(buffer)

        # Compute score using our utility function
        score = maxsim_score(query_embeddings, doc_embeddings)

        results.append({
            'document': file_path,
            'page': page_num,
            'score': score
        })

    # Sort and return top k
    results.sort(key=lambda x: x['score'], reverse=True)
    return results[:top_k]

# Example search
results = search_documents("attention mechanism architecture")
for r in results:
    print(f"{r['document']} (page {r['page']}): {r['score']:.3f}")

Key Advantage: ColPali finds relevant content without text extraction, preserving diagrams, equations, and layout information that would be lost in traditional pipelines.

Section 5 Unified Search Engine

Now we’ll combine all three approaches into a single search system. Instead of duplicating code, we’ll reference the implementations above and show how to integrate them.

5.1 Architecture Overview

The unified system brings together three specialized models, each handling different content types optimally. The key insight is that we don’t need a single model that handles everything - instead, we route queries to the appropriate specialist and combine results when needed.

class UnifiedSearchEngine:
    """
    Combines three search approaches:
    - CLIP for photos (Part 1)
    - Matryoshka for text (Part 2)
    - ColPali for documents (Part 3)
    """

    def __init__(self, db_connection_string):
        self.device = get_device()  # Use our utility function
        self.conn = psycopg2.connect(db_connection_string)
        self.cur = self.conn.cursor()

        # Load all models (details in Parts 1-3)
        self._load_models()
        self._initialize_database()

    def _load_models(self):
        """Load CLIP, Matryoshka, and ColPali models"""
        print(f"Initializing models on {self.device}...")

        # CLIP for images (see Part 1)
        self.clip_model = CLIPModel.from_pretrained(
            "openai/clip-vit-base-patch32"
        ).to(self.device)
        self.clip_processor = CLIPProcessor.from_pretrained(
            "openai/clip-vit-base-patch32"
        )

        # ColPali for documents (see Part 3)
        self.colpali_model = ColPali.from_pretrained(
            "vidore/colpali-v1.2",
            torch_dtype=torch.float16 if self.device != "cpu" else torch.float32
        ).to(self.device)
        self.colpali_processor = ColPaliProcessor.from_pretrained(
            "vidore/colpali-v1.2"
        )

        # Matryoshka for text (see Part 2)
        self.text_model = SentenceTransformer("nomic-ai/nomic-embed-text-v1.5")

    def _initialize_database(self):
        """Create tables for all content types (see Parts 1-3 for details)"""

        # Photos table (Part 1)
        self.cur.execute("""
            CREATE TABLE IF NOT EXISTS photos (
                id SERIAL PRIMARY KEY,
                file_path TEXT,
                embedding vector(512)
            )
        """)
        self.cur.execute("""
            CREATE INDEX IF NOT EXISTS photos_embedding_idx
            ON photos USING hnsw (embedding vector_cosine_ops)
        """)

        # Documents table (Part 3)
        self.cur.execute("""
            CREATE TABLE IF NOT EXISTS documents (
                id SERIAL PRIMARY KEY,
                file_path TEXT,
                page_num INTEGER,
                embedding_data BYTEA
            )
        """)

        # Text table (Part 2)
        self.cur.execute("""
            CREATE TABLE IF NOT EXISTS text_items (
                id SERIAL PRIMARY KEY,
                content TEXT,
                embedding_128 vector(128),
                embedding_256 vector(256),
                embedding_768 vector(768)
            )
        """)

        for dim in [128, 256, 768]:
            self.cur.execute(f"""
                CREATE INDEX IF NOT EXISTS text_embedding_{dim}_idx
                ON text_items USING hnsw (embedding_{dim} vector_cosine_ops)
            """)

        self.conn.commit()

5.2 Indexing Methods

Each method follows the pattern from earlier sections:

    def index_photo(self, image_path):
        """Index photo using CLIP (see Part 1 for standalone version)"""
        image = Image.open(image_path)

        inputs = self.clip_processor(images=image, return_tensors="pt")
        inputs = {k: v.to(self.device) for k, v in inputs.items()}

        with torch.no_grad():
            embedding = self.clip_model.get_image_features(**inputs)
            embedding = embedding / embedding.norm(dim=-1, keepdim=True)

        self.cur.execute(
            "INSERT INTO photos (file_path, embedding) VALUES (%s, %s)",
            (str(image_path), embedding[0].cpu().numpy().tolist())
        )
        self.conn.commit()

    def index_document(self, pdf_path):
        """Index PDF using ColPali (see Part 3 for standalone version)"""
        images = convert_from_path(pdf_path, dpi=150)

        for page_num, image in enumerate(images):
            inputs = self.colpali_processor.process_images([image]).to(self.device)

            with torch.no_grad():
                embeddings = self.colpali_model(**inputs)

            buffer = io.BytesIO()
            np.save(buffer, embeddings[0].cpu().numpy())

            self.cur.execute("""
                INSERT INTO documents (file_path, page_num, embedding_data)
                VALUES (%s, %s, %s)
            """, (str(pdf_path), page_num, buffer.getvalue()))

        self.conn.commit()

    def index_text(self, text_content):
        """Index text using Matryoshka (see Part 2 for standalone version)"""
        full_embedding = self.text_model.encode(text_content)

        self.cur.execute("""
            INSERT INTO text_items (content, embedding_128, embedding_256, embedding_768)
            VALUES (%s, %s, %s, %s)
        """, (
            text_content,
            full_embedding[:128].tolist(),
            full_embedding[:256].tolist(),
            full_embedding.tolist()
        ))
        self.conn.commit()

5.3 Search Methods

Each search method reuses logic from earlier sections:

    def search_photos(self, query_text, top_k=5):
        """Search photos - implements logic from Part 1"""
        inputs = self.clip_processor(text=[query_text], return_tensors="pt")
        inputs = {k: v.to(self.device) for k, v in inputs.items()}

        with torch.no_grad():
            query_embedding = self.clip_model.get_text_features(**inputs)
            query_embedding = query_embedding / query_embedding.norm(dim=-1, keepdim=True)

        query_np = query_embedding[0].cpu().numpy()

        self.cur.execute("""
            SELECT file_path, 1 - (embedding <=> %s::vector) AS similarity
            FROM photos
            ORDER BY embedding <=> %s::vector
            LIMIT %s
        """, (query_np.tolist(), query_np.tolist(), top_k))

        return [{"path": row[0], "score": row[1], "type": "photo"}
                for row in self.cur.fetchall()]

    def search_documents(self, query_text, top_k=5):
        """Search documents - implements logic from Part 3"""
        inputs = self.colpali_processor.process_queries([query_text]).to(self.device)

        with torch.no_grad():
            query_embeddings = self.colpali_model(**inputs)[0].cpu().numpy()

        self.cur.execute("SELECT id, file_path, page_num, embedding_data FROM documents")

        results = []
        for row in self.cur.fetchall():
            doc_id, file_path, page_num, embedding_data = row

            buffer = io.BytesIO(embedding_data)
            doc_embeddings = np.load(buffer)

            # Use our utility function
            score = maxsim_score(query_embeddings, doc_embeddings)

            results.append({
                "path": file_path,
                "page": page_num,
                "score": score,
                "type": "document"
            })

        results.sort(key=lambda x: x["score"], reverse=True)
        return results[:top_k]

    def search_text(self, query_text, top_k=10):
        """Search text with tiered retrieval - implements logic from Part 2"""
        query_embedding = self.text_model.encode(query_text)

        # Stage 1: 128D filtering
        query_128 = query_embedding[:128].tolist()
        self.cur.execute("""
            SELECT id FROM text_items
            ORDER BY embedding_128 <=> %s::vector
            LIMIT 100
        """, (query_128,))
        stage1_ids = [row[0] for row in self.cur.fetchall()]

        # Stage 2: 256D refinement
        query_256 = query_embedding[:256].tolist()
        self.cur.execute("""
            SELECT id FROM text_items
            WHERE id = ANY(%s)
            ORDER BY embedding_256 <=> %s::vector
            LIMIT 20
        """, (stage1_ids, query_256))
        stage2_ids = [row[0] for row in self.cur.fetchall()]

        # Stage 3: 768D final ranking
        query_768 = query_embedding.tolist()
        self.cur.execute("""
            SELECT content, 1 - (embedding_768 <=> %s::vector) AS similarity
            FROM text_items
            WHERE id = ANY(%s)
            ORDER BY embedding_768 <=> %s::vector
            LIMIT %s
        """, (query_768, stage2_ids, query_768, top_k))

        return [{"content": row[0], "score": row[1], "type": "text"}
                for row in self.cur.fetchall()]

    def search_all(self, query_text, top_k=5):
        """Search across all content types"""
        photo_results = self.search_photos(query_text, top_k)
        doc_results = self.search_documents(query_text, top_k)
        text_results = self.search_text(query_text, top_k)

        # Merge and sort by score
        all_results = photo_results + doc_results + text_results
        all_results.sort(key=lambda x: x["score"], reverse=True)

        return all_results[:top_k]

5.4 Using the Unified Engine

# Initialize
engine = UnifiedSearchEngine("dbname=unified_search user=postgres password=password")

# Index different content types
for photo in Path("photos").glob("*.jpg"):
    engine.index_photo(photo)

for doc in Path("documents").glob("*.pdf"):
    engine.index_document(doc)

for text in ["ML fundamentals", "Neural network basics"]:
    engine.index_text(text)

# Search across everything
results = engine.search_all("machine learning", top_k=5)

for r in results:
    if r['type'] == 'photo':
        print(f"[PHOTO] {r['path']}: {r['score']:.3f}")
    elif r['type'] == 'document':
        print(f"[DOC] {r['path']} p.{r['page']}: {r['score']:.3f}")
    else:
        print(f"[TEXT] {r['content'][:50]}...: {r['score']:.3f}")

Section 6 Production Considerations

6.1 Batch Processing for Efficiency

Process multiple items at once:

def index_photos_batch(engine, image_paths, batch_size=32):
    """Batch process photos for faster indexing"""
    for i in range(0, len(image_paths), batch_size):
        batch = image_paths[i:i+batch_size]
        images = [Image.open(p) for p in batch]

        inputs = engine.clip_processor(images=images, return_tensors="pt", padding=True)
        inputs = {k: v.to(engine.device) for k, v in inputs.items()}

        with torch.no_grad():
            embeddings = engine.clip_model.get_image_features(**inputs)
            embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)

        for path, embedding in zip(batch, embeddings):
            engine.cur.execute(
                "INSERT INTO photos (file_path, embedding) VALUES (%s, %s)",
                (str(path), embedding.cpu().numpy().tolist())
            )

        engine.conn.commit()

6.2 Query Caching

Cache frequently used queries:

from functools import lru_cache

class CachedSearchEngine(UnifiedSearchEngine):
    @lru_cache(maxsize=100)
    def _cached_text_embedding(self, query_text):
        """Cache text embeddings for repeated queries"""
        return self.text_model.encode(query_text)

6.3 Model Selection Tradeoffs

CLIP variants:

clip-vit-base-patch32: Balanced (512D embeddings)

clip-vit-large-patch14: Higher accuracy, slower

ColPali variants:

vidore/colpali-v1.2: Best accuracy

vidore/colsmol-v0.1: Faster inference, smaller

PDF Processing:

72 DPI: Fast, lower quality

150 DPI: Balanced (recommended)

300 DPI: High quality, slower

6.4 Intelligent Query Routing

Route queries to appropriate search based on content:

def intelligent_search(engine, query_text, top_k=5):
    """Route query based on keywords and characteristics"""
    query_lower = query_text.lower()

    if any(word in query_lower for word in ['photo', 'image', 'picture']):
        return engine.search_photos(query_text, top_k)
    elif any(word in query_lower for word in ['pdf', 'document', 'paper']):
        return engine.search_documents(query_text, top_k)
    elif len(query_text.split()) > 10:
        return engine.search_text(query_text, top_k)
    else:
        return engine.search_all(query_text, top_k)

Conclusion

This chapter demonstrated building a universal search engine handling three content types, each with its own strengths and optimal use cases. The unified system combines specialized models for each modality while maintaining a consistent search interface..

CLIP provides zero-shot image search with natural language, ideal when visual understanding matters

Matryoshka embeddings enable flexible precision-speed tradeoffs for text, perfect for large-scale systems

ColPali searches documents without text extraction, preserving visual information critical for technical content

Key Takeaways:

Use specialized models for different content types - one size doesn’t fit all

Normalize embeddings for cosine similarity to ensure consistent scoring

Leverage hardware acceleration (MPS/CUDA) for faster processing

Apply tiered search strategies for performance at scale

Batch process when possible to maximize throughput

Consider the tradeoffs between accuracy, speed, and storage for your use case

Each technique makes different tradeoffs, and understanding these helps you choose the right tool for your specific needs. CLIP trades storage for flexibility, Matryoshka trades a small accuracy loss for massive efficiency gains, and ColPali trades storage and speed for unprecedented document understanding. Together, they form a powerful toolkit for building modern search systems.