chromium/components/history_embeddings/vector_database.cc

// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "components/history_embeddings/vector_database.h"

#include <queue>

#include "base/ranges/algorithm.h"
#include "base/strings/string_util.h"
#include "base/timer/elapsed_timer.h"
#include "components/history_embeddings/history_embeddings_features.h"

namespace history_embeddings {

// Standard normalized magnitude for all embeddings.
constexpr float kUnitLength = …;

// Close enough to be considered near zero.
constexpr float kEpsilon = …;

////////////////////////////////////////////////////////////////////////////////

UrlPassages::UrlPassages(history::URLID url_id,
                         history::VisitID visit_id,
                         base::Time visit_time)
    : … { … }
UrlPassages::~UrlPassages() = default;
UrlPassages::UrlPassages(const UrlPassages&) = default;
UrlPassages& UrlPassages::operator=(const UrlPassages&) = default;
UrlPassages::UrlPassages(UrlPassages&&) = default;
UrlPassages& UrlPassages::operator=(UrlPassages&&) = default;
bool UrlPassages::operator==(const UrlPassages& other) const { … }

////////////////////////////////////////////////////////////////////////////////

Embedding::Embedding(std::vector<float> data) : … { … }
Embedding::Embedding() = default;
Embedding::Embedding(std::vector<float> data, size_t passage_word_count)
    : … { … }
Embedding::~Embedding() = default;
Embedding::Embedding(const Embedding&) = default;
Embedding& Embedding::operator=(const Embedding&) = default;
Embedding::Embedding(Embedding&&) = default;
Embedding& Embedding::operator=(Embedding&&) = default;
bool Embedding::operator==(const Embedding&) const = default;

size_t Embedding::Dimensions() const { … }

float Embedding::Magnitude() const { … }

void Embedding::Normalize() { … }

float Embedding::ScoreWith(SearchInfo& search_info,
                           const std::string& other_passage,
                           const Embedding& other_embedding) const { … }

////////////////////////////////////////////////////////////////////////////////

UrlEmbeddings::UrlEmbeddings() : … { … }
UrlEmbeddings::UrlEmbeddings(history::URLID url_id,
                             history::VisitID visit_id,
                             base::Time visit_time)
    : … { … }
UrlEmbeddings::UrlEmbeddings(const UrlPassages& url_passages)
    : … { … }
UrlEmbeddings::~UrlEmbeddings() = default;
UrlEmbeddings::UrlEmbeddings(UrlEmbeddings&&) = default;
UrlEmbeddings& UrlEmbeddings::operator=(UrlEmbeddings&&) = default;
UrlEmbeddings::UrlEmbeddings(const UrlEmbeddings&) = default;
UrlEmbeddings& UrlEmbeddings::operator=(const UrlEmbeddings&) = default;
bool UrlEmbeddings::operator==(const UrlEmbeddings&) const = default;

float UrlEmbeddings::BestScoreWith(SearchInfo& search_info,
                                   const Embedding& query,
                                   const proto::PassagesValue& passages,
                                   size_t search_minimum_word_count) const { … }

////////////////////////////////////////////////////////////////////////////////

ScoredUrl::ScoredUrl(history::URLID url_id,
                     history::VisitID visit_id,
                     base::Time visit_time,
                     float score)
    : … { … }
ScoredUrl::~ScoredUrl() = default;
ScoredUrl::ScoredUrl(ScoredUrl&&) = default;
ScoredUrl& ScoredUrl::operator=(ScoredUrl&&) = default;
ScoredUrl::ScoredUrl(const ScoredUrl&) = default;
ScoredUrl& ScoredUrl::operator=(const ScoredUrl&) = default;

////////////////////////////////////////////////////////////////////////////////

SearchInfo::SearchInfo() = default;
SearchInfo::SearchInfo(SearchInfo&&) = default;
SearchInfo::~SearchInfo() = default;

////////////////////////////////////////////////////////////////////////////////

UrlPassagesEmbeddings::UrlPassagesEmbeddings(history::URLID url_id,
                                             history::VisitID visit_id,
                                             base::Time visit_time)
    : … { … }
UrlPassagesEmbeddings::UrlPassagesEmbeddings(const UrlPassagesEmbeddings&) =
    default;
UrlPassagesEmbeddings& UrlPassagesEmbeddings::operator=(
    const UrlPassagesEmbeddings&) = default;
bool UrlPassagesEmbeddings::operator==(const UrlPassagesEmbeddings&) const =
    default;

////////////////////////////////////////////////////////////////////////////////

SearchInfo VectorDatabase::FindNearest(
    std::optional<base::Time> time_range_start,
    size_t count,
    const Embedding& query,
    base::RepeatingCallback<bool()> is_search_halted) { … }

////////////////////////////////////////////////////////////////////////////////

VectorDatabaseInMemory::VectorDatabaseInMemory() = default;
VectorDatabaseInMemory::~VectorDatabaseInMemory() = default;

void VectorDatabaseInMemory::SaveTo(VectorDatabase* database) { … }

size_t VectorDatabaseInMemory::GetEmbeddingDimensions() const { … }

bool VectorDatabaseInMemory::AddUrlData(UrlPassagesEmbeddings url_data) { … }

std::unique_ptr<VectorDatabase::UrlDataIterator>
VectorDatabaseInMemory::MakeUrlDataIterator(
    std::optional<base::Time> time_range_start) { … }

}  // namespace history_embeddings