huginn-daemon: source index (chunker, git mirrors, forge catalog, refresh with embedder failover, SourceIndex on Qdrant), huginn-sources bin, daemon/worker source search not gated by the mind's index
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
Meta Vens
parent
5074181ebf
commit
1c227f5907
15 changed files with 2006 additions and 26 deletions
|
|
@ -17,6 +17,12 @@ path = "src/lib.rs"
|
|||
name = "huginn-daemon"
|
||||
path = "src/main.rs"
|
||||
|
||||
# The source index's one writer: a oneshot job that brings `huginn_sources` to
|
||||
# the forge's current state and exits. No daemon calls any of its writes.
|
||||
[[bin]]
|
||||
name = "huginn-sources"
|
||||
path = "src/bin/huginn-sources.rs"
|
||||
|
||||
[dependencies]
|
||||
anyhow = "1"
|
||||
base64 = "0.22"
|
||||
|
|
|
|||
93
crates/huginn-daemon/src/bin/huginn-sources.rs
Normal file
93
crates/huginn-daemon/src/bin/huginn-sources.rs
Normal file
|
|
@ -0,0 +1,93 @@
|
|||
//! The source index's refresh job: one run brings the `huginn_sources`
|
||||
//! collection to the forge's current state, journals one line per repo, and
|
||||
//! exits nonzero if any repo failed or the collection could not be used.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::PathBuf;
|
||||
use std::process::ExitCode;
|
||||
|
||||
use anyhow::{Context, Result, ensure};
|
||||
use huginn_daemon::index::ollama::OllamaEmbedder;
|
||||
use huginn_daemon::index::qdrant::QdrantIndex;
|
||||
use huginn_daemon::sources::failover::FailoverEmbedder;
|
||||
use huginn_daemon::sources::forge::HttpForge;
|
||||
use huginn_daemon::sources::refresh;
|
||||
|
||||
/// Where the forge answers when the flag is absent: the loopback address of the
|
||||
/// forge on the host that runs this job.
|
||||
const DEFAULT_FORGE_URL: &str = "http://127.0.0.1:3310";
|
||||
|
||||
struct Options {
|
||||
state_root: PathBuf,
|
||||
qdrant_url: String,
|
||||
/// In order of preference; the first that serves the run's model embeds.
|
||||
ollama_urls: Vec<String>,
|
||||
embedding_model: String,
|
||||
forge_url: String,
|
||||
}
|
||||
|
||||
fn parse(args: impl Iterator<Item = String>) -> Result<Options> {
|
||||
let mut args = args;
|
||||
let mut values: BTreeMap<String, String> = BTreeMap::new();
|
||||
let mut ollama_urls = Vec::new();
|
||||
while let Some(name) = args.next() {
|
||||
let name = name.strip_prefix("--").with_context(|| format!("expected --option, got {name:?}"))?;
|
||||
ensure!(
|
||||
matches!(name, "state-root" | "qdrant-url" | "ollama-url" | "embedding-model" | "forge-url"),
|
||||
"unsupported huginn-sources option --{name}"
|
||||
);
|
||||
let value = args.next().with_context(|| format!("missing value for --{name}"))?;
|
||||
if name == "ollama-url" {
|
||||
ensure!(value.starts_with("http://"), "--ollama-url must be a plain http:// address, got {value:?}");
|
||||
ollama_urls.push(value);
|
||||
} else {
|
||||
ensure!(values.insert(name.to_owned(), value).is_none(), "duplicate option --{name}");
|
||||
}
|
||||
}
|
||||
let take = |name: &str| values.get(name).cloned().with_context(|| format!("--{name} is required"));
|
||||
let url = |name: &str, value: String| {
|
||||
ensure!(value.starts_with("http://"), "--{name} must be a plain http:// address, got {value:?}");
|
||||
Ok(value)
|
||||
};
|
||||
let state_root = PathBuf::from(take("state-root")?);
|
||||
ensure!(state_root.is_absolute(), "--state-root must be an absolute path");
|
||||
let embedding_model = take("embedding-model")?;
|
||||
ensure!(!embedding_model.is_empty(), "--embedding-model must name a model");
|
||||
Ok(Options {
|
||||
state_root,
|
||||
qdrant_url: url("qdrant-url", take("qdrant-url")?)?,
|
||||
ollama_urls: {
|
||||
ensure!(!ollama_urls.is_empty(), "at least one --ollama-url is required; repeat it in order of preference");
|
||||
ollama_urls
|
||||
},
|
||||
embedding_model,
|
||||
forge_url: url("forge-url", values.get("forge-url").cloned().unwrap_or_else(|| DEFAULT_FORGE_URL.to_owned()))?,
|
||||
})
|
||||
}
|
||||
|
||||
fn main() -> ExitCode {
|
||||
let options = match parse(std::env::args().skip(1)) {
|
||||
Ok(options) => options,
|
||||
Err(error) => {
|
||||
eprintln!("huginn-sources: {error:#}");
|
||||
return ExitCode::from(2);
|
||||
}
|
||||
};
|
||||
let mut embedder = FailoverEmbedder::new(
|
||||
options.ollama_urls.iter().map(|url| (url.clone(), OllamaEmbedder::new(url, &options.embedding_model))).collect(),
|
||||
);
|
||||
let mut index = QdrantIndex::new(&options.qdrant_url);
|
||||
let mut forge = HttpForge::new(&options.forge_url);
|
||||
let mut journal = |report: &huginn_daemon::sources::Report| println!("{report}");
|
||||
match refresh(&mut embedder, &mut index, &mut forge, &options.state_root, &mut journal) {
|
||||
Ok(refreshed) if refreshed.failed.is_empty() => ExitCode::SUCCESS,
|
||||
Ok(refreshed) => {
|
||||
eprintln!("huginn-sources: {} repos failed: {}", refreshed.failed.len(), refreshed.failed.join(", "));
|
||||
ExitCode::from(1)
|
||||
}
|
||||
Err(error) => {
|
||||
eprintln!("huginn-sources: {error:#}");
|
||||
ExitCode::from(1)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -6,7 +6,7 @@ use chrono::{DateTime, Utc};
|
|||
use cultnet_rs::Selection;
|
||||
use huginn_mind::eureka_pipeline::{PipelineRef, Slug};
|
||||
use huginn_mind::wire::{HuginnMindRequest, HuginnMindResponse, IndexStatus};
|
||||
use huginn_mind::{Members, Mind, MindRefusal, MindStore, PipelineAdmissionOutcome, SemanticQuery};
|
||||
use huginn_mind::{Members, Mind, MindRefusal, MindStore, PipelineAdmissionOutcome, SemanticQuery, SourceHits, SourceQuery};
|
||||
|
||||
/// Where the index attaches. Called after a batch commits, with the refs the
|
||||
/// batch landed, derived writes included; the sink may read each through
|
||||
|
|
@ -29,6 +29,11 @@ pub trait IndexSink<S: MindStore> {
|
|||
fn search(&mut self, text: &str, top_k: u32, among: Vec<String>) -> Result<SearchTicket, String>;
|
||||
/// How long after it was asked a search is worth answering: the index's
|
||||
/// own bound on it, which the loop waiting for the answer must share.
|
||||
/// Starts a source search. The source index is not the mind's: it needs no
|
||||
/// member ids, and nothing about the mind's own index state refuses it.
|
||||
fn search_sources(&mut self, query: &SourceQuery) -> Result<SearchTicket, String>;
|
||||
/// Every source search that finished since the last call, by ticket.
|
||||
fn sources_searched(&mut self) -> Vec<(SearchTicket, Result<SourceHits, String>)>;
|
||||
fn search_deadline(&self) -> std::time::Duration;
|
||||
/// Every search that finished since the last call, by ticket: its
|
||||
/// candidates and scores, or why it failed.
|
||||
|
|
@ -53,9 +58,22 @@ pub type Hits = Vec<(PipelineRef, f32)>;
|
|||
/// `Daemon::finish` with the index's answer.
|
||||
pub struct Search {
|
||||
pub ticket: SearchTicket,
|
||||
selection: Selection,
|
||||
semantic: SemanticQuery,
|
||||
members: Members,
|
||||
pending: Pending,
|
||||
}
|
||||
|
||||
/// What the asker is owed: a ranked page of the mind's documents, or the
|
||||
/// source index's chunks as the index holds them.
|
||||
enum Pending {
|
||||
Mind { selection: Selection, semantic: SemanticQuery, members: Members },
|
||||
Sources,
|
||||
}
|
||||
|
||||
/// What a finished search came to: the mind's candidates, the source index's
|
||||
/// hits, or why it failed.
|
||||
pub enum Found {
|
||||
Mind(Hits),
|
||||
Sources(SourceHits),
|
||||
Failed(String),
|
||||
}
|
||||
|
||||
/// What a request came to: an answer, or a search the index has yet to
|
||||
|
|
@ -170,6 +188,15 @@ impl<S: MindStore, I: IndexSink<S>> Daemon<S, I> {
|
|||
HuginnMindRequest::Query { selection, semantic: Some(semantic), .. } => {
|
||||
return self.begin_search(selection, semantic);
|
||||
}
|
||||
HuginnMindRequest::SearchSources { query, .. } => {
|
||||
if let Err(refusal) = query.check() {
|
||||
return Handled::Answered(HuginnMindResponse::Refused(refusal));
|
||||
}
|
||||
return match self.index.search_sources(&query) {
|
||||
Ok(ticket) => Handled::Searching(Search { ticket, pending: Pending::Sources }),
|
||||
Err(detail) => Handled::Answered(HuginnMindResponse::Refused(MindRefusal::Unavailable { detail })),
|
||||
};
|
||||
}
|
||||
};
|
||||
Handled::Answered(answer)
|
||||
}
|
||||
|
|
@ -193,7 +220,7 @@ impl<S: MindStore, I: IndexSink<S>> Daemon<S, I> {
|
|||
}
|
||||
let among = members.ids.iter().map(|id| id.id.0.clone()).collect();
|
||||
match self.index.search(&semantic.text.0, semantic.top_k, among) {
|
||||
Ok(ticket) => Handled::Searching(Search { ticket, selection, semantic, members }),
|
||||
Ok(ticket) => Handled::Searching(Search { ticket, pending: Pending::Mind { selection, semantic, members } }),
|
||||
Err(detail) => refused(MindRefusal::Unavailable { detail }),
|
||||
}
|
||||
}
|
||||
|
|
@ -204,8 +231,10 @@ impl<S: MindStore, I: IndexSink<S>> Daemon<S, I> {
|
|||
}
|
||||
|
||||
/// Searches the index has finished since this was last called.
|
||||
pub fn searched(&mut self) -> Vec<(SearchTicket, Result<Hits, String>)> {
|
||||
self.index.searched()
|
||||
pub fn searched(&mut self) -> Vec<(SearchTicket, Found)> {
|
||||
let mind = self.index.searched().into_iter().map(|(ticket, found)| (ticket, found.map_or_else(Found::Failed, Found::Mind)));
|
||||
let sources = self.index.sources_searched().into_iter().map(|(ticket, found)| (ticket, found.map_or_else(Found::Failed, Found::Sources)));
|
||||
mind.chain(sources).collect()
|
||||
}
|
||||
|
||||
/// The client has stopped waiting for `ticket`.
|
||||
|
|
@ -217,12 +246,15 @@ impl<S: MindStore, I: IndexSink<S>> Daemon<S, I> {
|
|||
/// through the mind (`Mind::rank` answers as of the snapshot the members
|
||||
/// were taken at and drops what is not a member), or, when the index could not answer, a typed
|
||||
/// `Unavailable`. Never an empty page in place of a failure.
|
||||
pub fn finish(&mut self, search: Search, found: Result<Hits, String>) -> HuginnMindResponse {
|
||||
let hits = match found {
|
||||
Ok(hits) => hits,
|
||||
Err(detail) => return HuginnMindResponse::Refused(MindRefusal::Unavailable { detail }),
|
||||
pub fn finish(&mut self, search: Search, found: Found) -> HuginnMindResponse {
|
||||
let unavailable = |detail: String| HuginnMindResponse::Refused(MindRefusal::Unavailable { detail });
|
||||
let (selection, semantic, members, hits) = match (search.pending, found) {
|
||||
(_, Found::Failed(detail)) => return unavailable(detail),
|
||||
(Pending::Sources, Found::Sources(hits)) => return HuginnMindResponse::Sources(hits),
|
||||
(Pending::Mind { selection, semantic, members }, Found::Mind(hits)) => (selection, semantic, members, hits),
|
||||
_ => return unavailable("the index answered another kind of search than the one asked".into()),
|
||||
};
|
||||
match self.mind.rank(&search.selection, &search.semantic, &search.members, &hits) {
|
||||
match self.mind.rank(&selection, &semantic, &members, &hits) {
|
||||
Ok(ranked) => {
|
||||
if let Some(note) = dropped_note(&ranked.dropped) {
|
||||
eprintln!("{note}");
|
||||
|
|
@ -497,6 +529,14 @@ pub(crate) mod tests {
|
|||
Err("this index cannot search".into())
|
||||
}
|
||||
|
||||
fn search_sources(&mut self, _query: &SourceQuery) -> Result<SearchTicket, String> {
|
||||
Err("this index cannot search".into())
|
||||
}
|
||||
|
||||
fn sources_searched(&mut self) -> Vec<(SearchTicket, Result<SourceHits, String>)> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn search_deadline(&self) -> std::time::Duration {
|
||||
std::time::Duration::ZERO
|
||||
}
|
||||
|
|
@ -539,6 +579,14 @@ pub(crate) mod tests {
|
|||
Err("this index cannot search".into())
|
||||
}
|
||||
|
||||
fn search_sources(&mut self, _query: &SourceQuery) -> Result<SearchTicket, String> {
|
||||
Err("this index cannot search".into())
|
||||
}
|
||||
|
||||
fn sources_searched(&mut self) -> Vec<(SearchTicket, Result<SourceHits, String>)> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn search_deadline(&self) -> std::time::Duration {
|
||||
std::time::Duration::ZERO
|
||||
}
|
||||
|
|
@ -565,6 +613,14 @@ pub(crate) mod tests {
|
|||
Err("this index cannot search".into())
|
||||
}
|
||||
|
||||
fn search_sources(&mut self, _query: &SourceQuery) -> Result<SearchTicket, String> {
|
||||
Err("this index cannot search".into())
|
||||
}
|
||||
|
||||
fn sources_searched(&mut self) -> Vec<(SearchTicket, Result<SourceHits, String>)> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn search_deadline(&self) -> std::time::Duration {
|
||||
std::time::Duration::ZERO
|
||||
}
|
||||
|
|
|
|||
|
|
@ -8,6 +8,8 @@ use std::time::Duration;
|
|||
|
||||
use anyhow::{Result, bail, ensure};
|
||||
|
||||
use crate::sources::{DescribedSources, SourceIndex, SourcePoint, SourceRow, SourcesMeta};
|
||||
|
||||
use super::{CollectionMeta, Described, Embedder, Hit, ModelIdentity, Point, VectorIndex, point_id};
|
||||
|
||||
pub(crate) fn identity(digest: &str) -> ModelIdentity {
|
||||
|
|
@ -162,6 +164,9 @@ pub(crate) struct IndexState {
|
|||
#[derive(Clone, Default)]
|
||||
pub(crate) struct FakeIndex {
|
||||
pub(crate) state: Arc<Mutex<IndexState>>,
|
||||
/// The source collection, held beside the mind's so one fake serves a
|
||||
/// worker that owns both.
|
||||
pub(crate) sources: FakeSources,
|
||||
}
|
||||
|
||||
impl FakeIndex {
|
||||
|
|
@ -241,3 +246,138 @@ impl VectorIndex for FakeIndex {
|
|||
pub(crate) fn pair() -> (FakeEmbedder, FakeIndex) {
|
||||
(FakeEmbedder::new("d1"), FakeIndex::default())
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct SourcesState {
|
||||
pub(crate) described: Option<DescribedSources>,
|
||||
pub(crate) points: BTreeMap<String, SourcePoint>,
|
||||
/// How many times the collection was recreated.
|
||||
pub(crate) recreated: u32,
|
||||
pub(crate) upsert_calls: u32,
|
||||
/// Every upsert call refuses once this many have succeeded.
|
||||
pub(crate) fail_upserts_after: Option<u32>,
|
||||
pub(crate) cursor_writes: Vec<std::collections::BTreeMap<String, String>>,
|
||||
pub(crate) deletes: Vec<(String, Option<Vec<String>>)>,
|
||||
/// The repo filter and limit of every search, in order.
|
||||
pub(crate) searches: Vec<(Option<String>, u32)>,
|
||||
}
|
||||
|
||||
/// An in-memory source collection: a brute-force cosine search over what it
|
||||
/// holds, so a test sees what a filter and a limit do.
|
||||
#[derive(Clone, Default)]
|
||||
pub(crate) struct FakeSources {
|
||||
pub(crate) state: Arc<Mutex<SourcesState>>,
|
||||
}
|
||||
|
||||
impl FakeSources {
|
||||
pub(crate) fn plant(&self, described: DescribedSources) {
|
||||
self.state.lock().unwrap().described = Some(described);
|
||||
}
|
||||
|
||||
pub(crate) fn cursors(&self) -> BTreeMap<String, String> {
|
||||
match &self.state.lock().unwrap().described {
|
||||
Some(DescribedSources::Labelled { cursors, .. }) => cursors.clone(),
|
||||
_ => BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// (repo, path, chunk_index, commit) of every point held, sorted.
|
||||
pub(crate) fn held(&self) -> Vec<(String, String, u32, String)> {
|
||||
let state = self.state.lock().unwrap();
|
||||
let mut held: Vec<_> = state
|
||||
.points
|
||||
.values()
|
||||
.map(|point| (point.payload.repo.clone(), point.payload.path.clone(), point.payload.chunk_index, point.payload.commit.clone()))
|
||||
.collect();
|
||||
held.sort();
|
||||
held
|
||||
}
|
||||
|
||||
pub(crate) fn ids(&self) -> BTreeSet<String> {
|
||||
self.state.lock().unwrap().points.keys().cloned().collect()
|
||||
}
|
||||
}
|
||||
|
||||
fn cosine(a: &[f32], b: &[f32]) -> f32 {
|
||||
let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
|
||||
let norm = |v: &[f32]| v.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
dot / (norm(a) * norm(b)).max(f32::MIN_POSITIVE)
|
||||
}
|
||||
|
||||
impl SourceIndex for FakeSources {
|
||||
fn describe_sources(&mut self, _within: Duration) -> Result<DescribedSources> {
|
||||
Ok(self.state.lock().unwrap().described.clone().unwrap_or(DescribedSources::Absent))
|
||||
}
|
||||
|
||||
fn recreate_sources(&mut self, meta: &SourcesMeta) -> Result<()> {
|
||||
let mut state = self.state.lock().unwrap();
|
||||
state.recreated += 1;
|
||||
state.points.clear();
|
||||
state.described = Some(DescribedSources::Labelled { meta: meta.clone(), cursors: BTreeMap::new() });
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn set_cursors(&mut self, cursors: &BTreeMap<String, String>) -> Result<()> {
|
||||
let mut state = self.state.lock().unwrap();
|
||||
state.cursor_writes.push(cursors.clone());
|
||||
let Some(DescribedSources::Labelled { cursors: held, .. }) = &mut state.described else { bail!("no source collection") };
|
||||
*held = cursors.clone();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn delete_by(&mut self, repo: &str, paths: Option<&[String]>) -> Result<()> {
|
||||
let mut state = self.state.lock().unwrap();
|
||||
state.deletes.push((repo.to_owned(), paths.map(<[String]>::to_vec)));
|
||||
state.points.retain(|_, point| {
|
||||
!(point.payload.repo == repo && paths.is_none_or(|paths| paths.contains(&point.payload.path)))
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn upsert_sources(&mut self, points: &[SourcePoint]) -> Result<()> {
|
||||
let mut state = self.state.lock().unwrap();
|
||||
if state.fail_upserts_after.is_some_and(|after| state.upsert_calls >= after) {
|
||||
bail!("the vector store refused the upsert");
|
||||
}
|
||||
state.upsert_calls += 1;
|
||||
for point in points {
|
||||
state.points.insert(point.id.clone(), point.clone());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn search_sources(&mut self, vector: &[f32], repo: Option<&str>, limit: u32, _within: Duration) -> Result<Vec<SourceRow>> {
|
||||
let mut state = self.state.lock().unwrap();
|
||||
state.searches.push((repo.map(str::to_owned), limit));
|
||||
let mut rows: Vec<SourceRow> = state
|
||||
.points
|
||||
.values()
|
||||
.filter(|point| repo.is_none_or(|repo| point.payload.repo == repo))
|
||||
.map(|point| SourceRow { score: cosine(vector, &point.vector), payload: point.payload.clone() })
|
||||
.collect();
|
||||
rows.sort_by(|a, b| b.score.total_cmp(&a.score));
|
||||
rows.truncate(limit as usize);
|
||||
Ok(rows)
|
||||
}
|
||||
}
|
||||
|
||||
impl SourceIndex for FakeIndex {
|
||||
fn describe_sources(&mut self, within: Duration) -> Result<DescribedSources> {
|
||||
self.sources.describe_sources(within)
|
||||
}
|
||||
fn recreate_sources(&mut self, meta: &SourcesMeta) -> Result<()> {
|
||||
self.sources.recreate_sources(meta)
|
||||
}
|
||||
fn set_cursors(&mut self, cursors: &BTreeMap<String, String>) -> Result<()> {
|
||||
self.sources.set_cursors(cursors)
|
||||
}
|
||||
fn delete_by(&mut self, repo: &str, paths: Option<&[String]>) -> Result<()> {
|
||||
self.sources.delete_by(repo, paths)
|
||||
}
|
||||
fn upsert_sources(&mut self, points: &[SourcePoint]) -> Result<()> {
|
||||
self.sources.upsert_sources(points)
|
||||
}
|
||||
fn search_sources(&mut self, vector: &[f32], repo: Option<&str>, limit: u32, within: Duration) -> Result<Vec<SourceRow>> {
|
||||
self.sources.search_sources(vector, repo, limit, within)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -105,6 +105,8 @@ use huginn_mind::{INDEX_TEXT_VERSION, IndexEntry, Mind, MindStore};
|
|||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::daemon::{Hits, IndexSink, SearchTicket};
|
||||
use crate::sources::{self, SourceIndex};
|
||||
use huginn_mind::{SourceHits, SourceQuery};
|
||||
use ollama::QUERY_INSTRUCTION;
|
||||
|
||||
/// The `managed_by` a collection carries when this organ made it.
|
||||
|
|
@ -262,17 +264,17 @@ pub enum Advance {
|
|||
|
||||
/// The time one operation has left. Every call the operation makes is bounded
|
||||
/// by what remains, not by a bound of its own, so their sum cannot outlast it.
|
||||
struct Budget {
|
||||
pub(crate) struct Budget {
|
||||
end: Instant,
|
||||
total: Duration,
|
||||
}
|
||||
|
||||
impl Budget {
|
||||
fn new(total: Duration) -> Self {
|
||||
pub(crate) fn new(total: Duration) -> Self {
|
||||
Self { end: Instant::now() + total, total }
|
||||
}
|
||||
|
||||
fn left(&self) -> Result<Duration> {
|
||||
pub(crate) fn left(&self) -> Result<Duration> {
|
||||
let left = self.end.saturating_duration_since(Instant::now());
|
||||
ensure!(!left.is_zero(), "the time allowed ({:?}) has run out", self.total);
|
||||
Ok(left)
|
||||
|
|
@ -669,6 +671,14 @@ struct Queued {
|
|||
due: Instant,
|
||||
}
|
||||
|
||||
/// A source search waiting for the worker. It is not the mind's: nothing about
|
||||
/// the mind's reconciliation gates it.
|
||||
struct QueuedSource {
|
||||
ticket: SearchTicket,
|
||||
query: SourceQuery,
|
||||
due: Instant,
|
||||
}
|
||||
|
||||
struct Shared {
|
||||
inbox: Vec<IndexEntry>,
|
||||
/// The mind's identity, once a commit has revealed it to a mind that had none.
|
||||
|
|
@ -681,6 +691,8 @@ struct Shared {
|
|||
searches: VecDeque<Queued>,
|
||||
/// Searches the worker finished and the serve loop has not collected.
|
||||
found: Vec<(SearchTicket, Result<Hits, String>)>,
|
||||
source_searches: VecDeque<QueuedSource>,
|
||||
sources_found: Vec<(SearchTicket, Result<SourceHits, String>)>,
|
||||
next_ticket: u64,
|
||||
}
|
||||
|
||||
|
|
@ -711,7 +723,7 @@ fn wait_for_work(shared: &Lock, wait: Wait) {
|
|||
if guard.closed {
|
||||
return;
|
||||
}
|
||||
if wakes_on_work && (!guard.inbox.is_empty() || guard.mind.is_some() || !guard.searches.is_empty()) {
|
||||
if !guard.source_searches.is_empty() || (wakes_on_work && (!guard.inbox.is_empty() || guard.mind.is_some() || !guard.searches.is_empty())) {
|
||||
return;
|
||||
}
|
||||
let left = deadline.saturating_duration_since(Instant::now());
|
||||
|
|
@ -792,6 +804,21 @@ fn serve_search<E: Embedder, V: VectorIndex>(projector: &mut Projector<E, V>, sh
|
|||
/// owed (documents to write, a reconciliation to make) unless the index
|
||||
/// cannot be asked at all, in which case the search is taken to be refused:
|
||||
/// waiting would not change the answer.
|
||||
/// Answers the oldest queued source search, whatever state the mind's
|
||||
/// projection is in: it asks the embedder and the source collection and nothing
|
||||
/// of the projector's reconciliation.
|
||||
fn serve_source_search<E: Embedder, V: VectorIndex + SourceIndex>(projector: &mut Projector<E, V>, shared: &Lock) {
|
||||
let mut guard = lock(shared);
|
||||
if guard.closed {
|
||||
return;
|
||||
}
|
||||
let Some(QueuedSource { ticket, query, due }) = guard.source_searches.pop_front() else { return };
|
||||
drop(guard);
|
||||
let found = sources::search(&mut projector.embedder, &mut projector.index, &query, due.saturating_duration_since(Instant::now()))
|
||||
.map_err(|error| format!("{error:#}"));
|
||||
lock(shared).sources_found.push((ticket, found));
|
||||
}
|
||||
|
||||
fn take_search<E: Embedder, V: VectorIndex>(projector: &Projector<E, V>, guard: &mut Shared) -> Option<Queued> {
|
||||
if !guard.inbox.is_empty() {
|
||||
return None;
|
||||
|
|
@ -852,7 +879,7 @@ fn absorb_into<E: Embedder, V: VectorIndex>(
|
|||
true
|
||||
}
|
||||
|
||||
fn work<E: Embedder, V: VectorIndex>(mut projector: Projector<E, V>, shared: Lock, backoff: Backoff) {
|
||||
fn work<E: Embedder, V: VectorIndex + SourceIndex>(mut projector: Projector<E, V>, shared: Lock, backoff: Backoff) {
|
||||
let mut failing: Option<(StepError, u32)> = None;
|
||||
let mut delay = backoff.initial;
|
||||
// When the worker next verifies the model and the collection whether or
|
||||
|
|
@ -878,6 +905,7 @@ fn work<E: Embedder, V: VectorIndex>(mut projector: Projector<E, V>, shared: Loc
|
|||
}
|
||||
publish(&shared, projector.progress_status());
|
||||
serve_search(&mut projector, &shared);
|
||||
serve_source_search(&mut projector, &shared);
|
||||
// A model or collection the search found changed has left the
|
||||
// reconciliation to redo: do it now, not at the next wake.
|
||||
if advance == Advance::Idle && !projector.stale() {
|
||||
|
|
@ -894,7 +922,17 @@ fn work<E: Embedder, V: VectorIndex>(mut projector: Projector<E, V>, shared: Loc
|
|||
publish(&shared, projector.failure_status(&error, attempts));
|
||||
refuse_searches(&shared, &unsearchable(&projector.failure_status(&error, attempts)).unwrap_or_default());
|
||||
failing = Some((error, attempts));
|
||||
wait_for_work(&shared, Wait::Backoff(delay));
|
||||
// The mind's index is down; a source search does not depend on
|
||||
// it, so it is answered while the backoff runs out.
|
||||
let until = Instant::now() + delay;
|
||||
loop {
|
||||
serve_source_search(&mut projector, &shared);
|
||||
let left = until.saturating_duration_since(Instant::now());
|
||||
if left.is_zero() || lock(&shared).closed {
|
||||
break;
|
||||
}
|
||||
wait_for_work(&shared, Wait::Backoff(left));
|
||||
}
|
||||
delay = backoff.after(delay);
|
||||
}
|
||||
}
|
||||
|
|
@ -928,7 +966,7 @@ impl WorkerSink {
|
|||
) -> Self
|
||||
where
|
||||
E: Embedder + Send + 'static,
|
||||
V: VectorIndex + Send + 'static,
|
||||
V: VectorIndex + SourceIndex + Send + 'static,
|
||||
{
|
||||
let identified = mind.is_some();
|
||||
let mut projector = Projector::new(embedder, index, instance, mind);
|
||||
|
|
@ -941,6 +979,8 @@ impl WorkerSink {
|
|||
closed: false,
|
||||
searches: VecDeque::new(),
|
||||
found: Vec::new(),
|
||||
source_searches: VecDeque::new(),
|
||||
sources_found: Vec::new(),
|
||||
next_ticket: 0,
|
||||
}),
|
||||
Condvar::new(),
|
||||
|
|
@ -1017,10 +1057,10 @@ impl<S: MindStore> IndexSink<S> for WorkerSink {
|
|||
if let Some(why) = unsearchable(&status) {
|
||||
return Err(why);
|
||||
}
|
||||
if guard.searches.len() >= SEARCHES_QUEUED_MAX {
|
||||
if guard.searches.len() + guard.source_searches.len() >= SEARCHES_QUEUED_MAX {
|
||||
return Err(format!(
|
||||
"the semantic index has {} searches waiting and takes no more until it has answered some",
|
||||
guard.searches.len()
|
||||
guard.searches.len() + guard.source_searches.len()
|
||||
));
|
||||
}
|
||||
let ticket = SearchTicket(guard.next_ticket);
|
||||
|
|
@ -1031,6 +1071,26 @@ impl<S: MindStore> IndexSink<S> for WorkerSink {
|
|||
Ok(ticket)
|
||||
}
|
||||
|
||||
fn search_sources(&mut self, query: &SourceQuery) -> Result<SearchTicket, String> {
|
||||
let mut guard = lock(&self.shared);
|
||||
if guard.searches.len() + guard.source_searches.len() >= SEARCHES_QUEUED_MAX {
|
||||
return Err(format!(
|
||||
"the index has {} searches waiting and takes no more until it has answered some",
|
||||
guard.searches.len() + guard.source_searches.len()
|
||||
));
|
||||
}
|
||||
let ticket = SearchTicket(guard.next_ticket);
|
||||
guard.next_ticket += 1;
|
||||
guard.source_searches.push_back(QueuedSource { ticket, query: query.clone(), due: Instant::now() + self.deadline });
|
||||
drop(guard);
|
||||
self.shared.1.notify_all();
|
||||
Ok(ticket)
|
||||
}
|
||||
|
||||
fn sources_searched(&mut self) -> Vec<(SearchTicket, Result<SourceHits, String>)> {
|
||||
std::mem::take(&mut lock(&self.shared).sources_found)
|
||||
}
|
||||
|
||||
fn search_deadline(&self) -> Duration {
|
||||
self.deadline
|
||||
}
|
||||
|
|
@ -1043,7 +1103,9 @@ impl<S: MindStore> IndexSink<S> for WorkerSink {
|
|||
/// the worker already holds cannot be recalled; its answer is collected
|
||||
/// and dropped by the loop, which no longer waits for it.
|
||||
fn abandon(&mut self, ticket: SearchTicket) {
|
||||
lock(&self.shared).searches.retain(|queued| queued.ticket != ticket);
|
||||
let mut guard = lock(&self.shared);
|
||||
guard.searches.retain(|queued| queued.ticket != ticket);
|
||||
guard.source_searches.retain(|queued| queued.ticket != ticket);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1061,6 +1123,7 @@ pub(crate) enum Method {
|
|||
Put,
|
||||
Post,
|
||||
Delete,
|
||||
Patch,
|
||||
}
|
||||
|
||||
impl Http {
|
||||
|
|
@ -1084,6 +1147,7 @@ impl Http {
|
|||
Method::Delete => self.agent.delete(&url).call(),
|
||||
Method::Post => self.agent.post(&url).header("content-type", json).send(body.unwrap_or_default()),
|
||||
Method::Put => self.agent.put(&url).header("content-type", json).send(body.unwrap_or_default()),
|
||||
Method::Patch => self.agent.patch(&url).header("content-type", json).send(body.unwrap_or_default()),
|
||||
};
|
||||
let mut response = sent.with_context(|| format!("no answer from {url}"))?;
|
||||
let status = response.status().as_u16();
|
||||
|
|
|
|||
|
|
@ -1,12 +1,14 @@
|
|||
//! The vector-store port over Qdrant's HTTP API. With `ollama.rs`, the xenos
|
||||
//! boundary: JSON exists here and nowhere else in the organ.
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{Context, Result, ensure};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::sources::{COLLECTION, DescribedSources, SourceIndex, SourcePayload, SourcePoint, SourceRow, SourcesMeta};
|
||||
|
||||
use super::{CollectionMeta, INDEX_CALL_TIMEOUT, Described, Hit, Http, Method, Point, VectorIndex};
|
||||
|
||||
/// One Qdrant. Every path names a collection this organ was told to own; it
|
||||
|
|
@ -163,3 +165,138 @@ impl VectorIndex for QdrantIndex {
|
|||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
fn payload_json(payload: &SourcePayload) -> Value {
|
||||
json!({
|
||||
"repo": payload.repo,
|
||||
"path": payload.path,
|
||||
"commit": payload.commit,
|
||||
"blob": payload.blob,
|
||||
"line_start": payload.line_start,
|
||||
"line_end": payload.line_end,
|
||||
"chunk_index": payload.chunk_index,
|
||||
"ext": payload.ext,
|
||||
"text": payload.text,
|
||||
})
|
||||
}
|
||||
|
||||
fn payload_of(payload: &Value) -> Result<SourcePayload> {
|
||||
let text = |key: &str| payload[key].as_str().map(str::to_owned).with_context(|| format!("a chunk carries no {key}"));
|
||||
let number = |key: &str| {
|
||||
payload[key].as_u64().and_then(|value| u32::try_from(value).ok()).with_context(|| format!("a chunk carries no {key}"))
|
||||
};
|
||||
Ok(SourcePayload {
|
||||
repo: text("repo")?,
|
||||
path: text("path")?,
|
||||
commit: text("commit")?,
|
||||
blob: text("blob")?,
|
||||
line_start: number("line_start")?,
|
||||
line_end: number("line_end")?,
|
||||
chunk_index: number("chunk_index")?,
|
||||
ext: text("ext")?,
|
||||
text: text("text")?,
|
||||
})
|
||||
}
|
||||
|
||||
/// The source collection's label and cursors, as they sit in its metadata.
|
||||
fn sources_described(metadata: &Value) -> DescribedSources {
|
||||
let text = |key: &str| metadata[key].as_str().map(str::to_owned);
|
||||
let Some(managed_by) = text("managed_by") else { return DescribedSources::Unlabelled };
|
||||
let number = |key: &str| metadata[key].as_u64().and_then(|value| u32::try_from(value).ok()).unwrap_or(0);
|
||||
let cursors: BTreeMap<String, String> = metadata["cursors"]
|
||||
.as_object()
|
||||
.map(|held| held.iter().filter_map(|(repo, commit)| Some((repo.clone(), commit.as_str()?.to_owned()))).collect())
|
||||
.unwrap_or_default();
|
||||
DescribedSources::Labelled {
|
||||
meta: SourcesMeta {
|
||||
managed_by,
|
||||
corpus: text("corpus").unwrap_or_default(),
|
||||
model: text("model").unwrap_or_default(),
|
||||
model_digest: text("model_digest").unwrap_or_default(),
|
||||
dimensions: number("dimensions"),
|
||||
chunker_version: number("chunker_version"),
|
||||
},
|
||||
cursors,
|
||||
}
|
||||
}
|
||||
|
||||
impl SourceIndex for QdrantIndex {
|
||||
fn describe_sources(&mut self, within: Duration) -> Result<DescribedSources> {
|
||||
let http = self.http.bounded(within);
|
||||
let (status, answer) = self.ask_with(&http, Method::Get, &format!("/collections/{COLLECTION}"), None, &[200, 404])?;
|
||||
if status == 404 {
|
||||
return Ok(DescribedSources::Absent);
|
||||
}
|
||||
Ok(sources_described(&answer["result"]["config"]["metadata"]))
|
||||
}
|
||||
|
||||
fn recreate_sources(&mut self, meta: &SourcesMeta) -> Result<()> {
|
||||
let path = format!("/collections/{COLLECTION}");
|
||||
self.ask(Method::Delete, &format!("{path}?wait=true"), None, &[200, 404])?;
|
||||
let body = json!({
|
||||
"vectors": { "size": meta.dimensions, "distance": "Cosine", "on_disk": true },
|
||||
"metadata": {
|
||||
"managed_by": meta.managed_by,
|
||||
"corpus": meta.corpus,
|
||||
"model": meta.model,
|
||||
"model_digest": meta.model_digest,
|
||||
"dimensions": meta.dimensions,
|
||||
"chunker_version": meta.chunker_version,
|
||||
"cursors": {},
|
||||
},
|
||||
});
|
||||
self.ask(Method::Put, &format!("{path}?wait=true"), Some(&body), &[200])?;
|
||||
for field in ["repo", "path"] {
|
||||
let index = json!({ "field_name": field, "field_schema": "keyword" });
|
||||
self.ask(Method::Put, &format!("{path}/index?wait=true"), Some(&index), &[200])?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn set_cursors(&mut self, cursors: &BTreeMap<String, String>) -> Result<()> {
|
||||
// Qdrant merges metadata by top-level key and replaces a nested object
|
||||
// whole, so this is the whole cursors object and nothing beside it.
|
||||
let body = json!({ "metadata": { "cursors": cursors } });
|
||||
self.ask(Method::Patch, &format!("/collections/{COLLECTION}?wait=true"), Some(&body), &[200])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn delete_by(&mut self, repo: &str, paths: Option<&[String]>) -> Result<()> {
|
||||
let mut must = vec![json!({ "key": "repo", "match": { "value": repo } })];
|
||||
if let Some(paths) = paths {
|
||||
must.push(json!({ "key": "path", "match": { "any": paths } }));
|
||||
}
|
||||
let body = json!({ "filter": { "must": must } });
|
||||
self.ask(Method::Post, &format!("/collections/{COLLECTION}/points/delete?wait=true"), Some(&body), &[200])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn upsert_sources(&mut self, points: &[SourcePoint]) -> Result<()> {
|
||||
let points: Vec<Value> = points
|
||||
.iter()
|
||||
.map(|point| json!({ "id": point.id, "vector": point.vector, "payload": payload_json(&point.payload) }))
|
||||
.collect();
|
||||
self.ask(Method::Put, &format!("/collections/{COLLECTION}/points?wait=true"), Some(&json!({ "points": points })), &[200])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn search_sources(&mut self, vector: &[f32], repo: Option<&str>, limit: u32, within: Duration) -> Result<Vec<SourceRow>> {
|
||||
let http = self.http.bounded(within);
|
||||
let mut body = json!({ "vector": vector, "limit": limit, "with_payload": true, "with_vector": false });
|
||||
if let Some(repo) = repo {
|
||||
body["filter"] = json!({ "must": [{ "key": "repo", "match": { "value": repo } }] });
|
||||
}
|
||||
let (_, answer) =
|
||||
self.ask_with(&http, Method::Post, &format!("/collections/{COLLECTION}/points/search"), Some(&body), &[200])?;
|
||||
let found = answer["result"].as_array().context("Qdrant's search has no result")?;
|
||||
found
|
||||
.iter()
|
||||
.map(|point| {
|
||||
Ok(SourceRow {
|
||||
score: point["score"].as_f64().context("a hit carries no score")? as f32,
|
||||
payload: payload_of(&point["payload"])?,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1968,3 +1968,90 @@ fn the_sink_owns_the_searchs_deadline_and_tells_the_loop() {
|
|||
assert_eq!(<WorkerSink as IndexSink<OwnedRedbMessagePackBackingStore>>::search_deadline(&sink), deadline);
|
||||
assert_eq!(Daemon::new(mind, sink).search_deadline(), deadline);
|
||||
}
|
||||
|
||||
fn ask_sources(sink: &mut WorkerSink, text: &str) -> Result<SearchTicket, String> {
|
||||
let query = SourceQuery { text: huginn_mind::eureka_pipeline::Line(text.into()), top_k: 5, repo: None };
|
||||
<WorkerSink as IndexSink<OwnedRedbMessagePackBackingStore>>::search_sources(sink, &query)
|
||||
}
|
||||
|
||||
fn sources_answers(sink: &mut WorkerSink, count: usize) -> Vec<(SearchTicket, Result<SourceHits, String>)> {
|
||||
let mut all = Vec::new();
|
||||
eventually("the source searches being answered", || {
|
||||
all.extend(<WorkerSink as IndexSink<OwnedRedbMessagePackBackingStore>>::sources_searched(sink));
|
||||
all.len() >= count
|
||||
});
|
||||
all
|
||||
}
|
||||
|
||||
/// The mind's reconciliation does not gate a source search: where the sink
|
||||
/// refuses a mind search because the mind's index is being rebuilt, failing or
|
||||
/// refused, it takes a source search.
|
||||
#[test]
|
||||
fn source_search_is_not_gated_by_the_minds_index_state() {
|
||||
use IndexStatus as S;
|
||||
for status in [
|
||||
S::Reconciling { pending: 4 },
|
||||
S::Failing { pending: 1, attempts: 1, error: "boom".into() },
|
||||
S::Refused { pending: 1, attempts: 1, reason: "not ours".into() },
|
||||
] {
|
||||
let shared = shared_with(Vec::new(), None, Vec::new(), false);
|
||||
lock(&shared).status = status.clone();
|
||||
let mut sink = WorkerSink { shared, identified: true, deadline: SEARCH_DEADLINE };
|
||||
assert!(ask_sink(&mut sink, "anything", 3).is_err(), "{status:?} refuses a mind search");
|
||||
assert!(ask_sources(&mut sink, "anything").is_ok(), "{status:?} must not refuse a source search");
|
||||
}
|
||||
}
|
||||
|
||||
/// A worker whose mind collection is another mind's (it cannot advance and
|
||||
/// backs off) still answers a source search while it waits, and the answer is
|
||||
/// the source collection's, not an error from the mind's.
|
||||
#[test]
|
||||
fn a_worker_stuck_on_the_minds_collection_still_answers_a_source_search() {
|
||||
use crate::sources::{DescribedSources, SourcePayload, SourcePoint, SourcesMeta};
|
||||
let (embedder, index) = pair();
|
||||
let mind_label = Described::Labelled(CollectionMeta {
|
||||
managed_by: MANAGED_BY.into(),
|
||||
instance: INSTANCE.into(),
|
||||
mind: "another-mind".into(),
|
||||
model: "model".into(),
|
||||
model_digest: "d1".into(),
|
||||
dimensions: 4,
|
||||
index_text_version: INDEX_TEXT_VERSION,
|
||||
});
|
||||
index.plant(&collection_name(&slug(INSTANCE)), mind_label);
|
||||
let meta = SourcesMeta {
|
||||
managed_by: MANAGED_BY.into(),
|
||||
corpus: "forge:GameCult".into(),
|
||||
model: "model".into(),
|
||||
model_digest: "d1".into(),
|
||||
dimensions: 4,
|
||||
chunker_version: 1,
|
||||
};
|
||||
index.sources.plant(DescribedSources::Labelled { meta, cursors: Default::default() });
|
||||
index.sources.state.lock().unwrap().points.insert(
|
||||
"p1".into(),
|
||||
SourcePoint {
|
||||
id: "p1".into(),
|
||||
vector: vec![1.0; 4],
|
||||
payload: SourcePayload {
|
||||
repo: "Huginn".into(),
|
||||
path: "a.rs".into(),
|
||||
commit: "abc".into(),
|
||||
blob: "b".into(),
|
||||
line_start: 1,
|
||||
line_end: 2,
|
||||
chunk_index: 0,
|
||||
ext: "rs".into(),
|
||||
text: "fn a() {}".into(),
|
||||
},
|
||||
},
|
||||
);
|
||||
let mut sink = WorkerSink::spawn(embedder, index.clone(), &slug(INSTANCE), Some(MIND.into()), entries_of(1), quick());
|
||||
let refused = sink.wait_status(|status| matches!(status, IndexStatus::Refused { .. }));
|
||||
assert!(matches!(refused, IndexStatus::Refused { .. }), "{refused:?}");
|
||||
ask_sources(&mut sink, "a function").unwrap();
|
||||
let answered = sources_answers(&mut sink, 1).remove(0).1.expect("the source search answers");
|
||||
assert_eq!(answered.hits.len(), 1);
|
||||
assert_eq!((answered.hits[0].repo.as_str(), answered.hits[0].path.as_str()), ("Huginn", "a.rs"));
|
||||
assert_eq!(answered.indexed_under.digest, "d1");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -17,9 +17,10 @@ pub mod bodies;
|
|||
pub mod daemon;
|
||||
pub mod index;
|
||||
pub mod serve;
|
||||
pub mod sources;
|
||||
|
||||
#[cfg(test)]
|
||||
mod semantic_tests;
|
||||
|
||||
pub use daemon::{Daemon, Handled, Hits, IndexSink, Search, SearchTicket, runtime_id};
|
||||
pub use daemon::{Daemon, Found, Handled, Hits, IndexSink, Search, SearchTicket, runtime_id};
|
||||
pub use serve::{Options, ServeOptions, parse_options, run, startup};
|
||||
|
|
|
|||
|
|
@ -522,14 +522,14 @@ pub fn run<S: MindStore, I: IndexSink<S>>(
|
|||
fn collect_searches<S: MindStore, I: IndexSink<S>>(
|
||||
daemon: &mut Daemon<S, I>,
|
||||
waiting: &BTreeMap<SearchTicket, Waiting>,
|
||||
) -> Vec<(SearchTicket, Result<crate::daemon::Hits, String>)> {
|
||||
) -> Vec<(SearchTicket, crate::daemon::Found)> {
|
||||
let deadline = daemon.search_deadline();
|
||||
let mut collected = daemon.searched();
|
||||
let now = Instant::now();
|
||||
for (ticket, asked) in waiting {
|
||||
if asked.until <= now {
|
||||
daemon.abandon(*ticket);
|
||||
collected.push((*ticket, Err(format!("the semantic search did not finish within {deadline:?}"))));
|
||||
collected.push((*ticket, crate::daemon::Found::Failed(format!("the search did not finish within {deadline:?}"))));
|
||||
}
|
||||
}
|
||||
collected
|
||||
|
|
|
|||
134
crates/huginn-daemon/src/sources/chunk.rs
Normal file
134
crates/huginn-daemon/src/sources/chunk.rs
Normal file
|
|
@ -0,0 +1,134 @@
|
|||
//! The source chunker: line windows, no overlap. A chunk is the unit the index
|
||||
//! holds and a search answers, so its line range is a citation.
|
||||
|
||||
/// Bumped whenever the windows or the embedded text would change for the same
|
||||
/// blob: it is part of every chunk's identity and of the collection's label, so
|
||||
/// a change rebuilds the collection instead of mixing two chunkings.
|
||||
pub const CHUNKER_VERSION: u32 = 1;
|
||||
|
||||
/// The most lines a chunk holds.
|
||||
pub const CHUNK_LINES_MAX: usize = 48;
|
||||
/// The most characters a chunk's text holds.
|
||||
pub const CHUNK_CHARS_MAX: usize = 1_400;
|
||||
|
||||
/// One window of a file: its 1-based inclusive line range and its text.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Chunk {
|
||||
pub line_start: u32,
|
||||
pub line_end: u32,
|
||||
pub text: String,
|
||||
}
|
||||
|
||||
/// The file cut into windows of at most `CHUNK_LINES_MAX` lines or
|
||||
/// `CHUNK_CHARS_MAX` characters. CRLF is read as LF and a trailing newline ends
|
||||
/// the last line without starting another. A line longer than the character
|
||||
/// bound is a chunk of its own, cut at the bound on a character boundary.
|
||||
/// A window of blank lines holds nothing to find and is dropped.
|
||||
pub fn chunk(text: &str) -> Vec<Chunk> {
|
||||
let normalised = text.replace("\r\n", "\n");
|
||||
let body = normalised.strip_suffix('\n').unwrap_or(&normalised);
|
||||
if body.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
let mut chunks = Vec::new();
|
||||
let mut window: Vec<&str> = Vec::new();
|
||||
let mut window_chars = 0_usize;
|
||||
let mut window_start = 1_u32;
|
||||
let flush = |window: &mut Vec<&str>, window_chars: &mut usize, window_start: u32, chunks: &mut Vec<Chunk>| {
|
||||
if !window.is_empty() {
|
||||
let joined = window.join("\n");
|
||||
if !joined.trim().is_empty() {
|
||||
let end = window_start + u32::try_from(window.len()).unwrap_or(u32::MAX) - 1;
|
||||
chunks.push(Chunk { line_start: window_start, line_end: end, text: joined });
|
||||
}
|
||||
}
|
||||
window.clear();
|
||||
*window_chars = 0;
|
||||
};
|
||||
for (index, line) in body.split('\n').enumerate() {
|
||||
let number = u32::try_from(index + 1).unwrap_or(u32::MAX);
|
||||
let length = line.chars().count();
|
||||
if length > CHUNK_CHARS_MAX {
|
||||
flush(&mut window, &mut window_chars, window_start, &mut chunks);
|
||||
let cut: String = line.chars().take(CHUNK_CHARS_MAX).collect();
|
||||
chunks.push(Chunk { line_start: number, line_end: number, text: cut });
|
||||
continue;
|
||||
}
|
||||
let joined = window_chars + length + usize::from(!window.is_empty());
|
||||
if !window.is_empty() && (window.len() >= CHUNK_LINES_MAX || joined > CHUNK_CHARS_MAX) {
|
||||
flush(&mut window, &mut window_chars, window_start, &mut chunks);
|
||||
}
|
||||
if window.is_empty() {
|
||||
window_start = number;
|
||||
window_chars = length;
|
||||
} else {
|
||||
window_chars += length + 1;
|
||||
}
|
||||
window.push(line);
|
||||
}
|
||||
flush(&mut window, &mut window_chars, window_start, &mut chunks);
|
||||
chunks
|
||||
}
|
||||
|
||||
/// What is embedded for a chunk: where it lives, then what it says. The path
|
||||
/// is part of the meaning of a line of code.
|
||||
pub fn embedded_text(repo: &str, path: &str, chunk: &Chunk) -> String {
|
||||
format!("{repo}/{path}\n{}", chunk.text)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn lines(count: usize) -> String {
|
||||
(1..=count).map(|n| format!("line {n}\n")).collect()
|
||||
}
|
||||
|
||||
/// Every chunk's range is the lines it holds, counted from one, whatever
|
||||
/// the file's line endings or last newline.
|
||||
#[test]
|
||||
fn chunk_lines_are_exact() {
|
||||
for text in [lines(100), lines(100).replace('\n', "\r\n"), lines(100).trim_end().to_owned()] {
|
||||
let source: Vec<String> =
|
||||
text.replace("\r\n", "\n").trim_end_matches('\n').split('\n').map(str::to_owned).collect();
|
||||
let chunks = chunk(&text);
|
||||
assert_eq!(chunks.len(), 3, "100 lines at 48 per window");
|
||||
let mut next = 1;
|
||||
for held in &chunks {
|
||||
assert_eq!(held.line_start, next, "no gap and no overlap");
|
||||
let expected = source[held.line_start as usize - 1..held.line_end as usize].join("\n");
|
||||
assert_eq!(held.text, expected, "lines {}-{}", held.line_start, held.line_end);
|
||||
assert!(!held.text.contains('\r'));
|
||||
next = held.line_end + 1;
|
||||
}
|
||||
assert_eq!(next, 101, "every line is in a chunk");
|
||||
}
|
||||
assert_eq!(chunk("a\n"), vec![Chunk { line_start: 1, line_end: 1, text: "a".into() }]);
|
||||
assert!(chunk("").is_empty() && chunk("\n").is_empty() && chunk("\n\n \n").is_empty());
|
||||
}
|
||||
|
||||
/// A window closes at the character bound before it closes at the line
|
||||
/// bound, and a line past the bound stands alone, cut on a boundary.
|
||||
#[test]
|
||||
fn a_window_closes_at_the_character_bound_and_a_long_line_stands_alone() {
|
||||
let wide = "x".repeat(700);
|
||||
let chunks = chunk(&format!("{wide}\n{wide}\n{wide}\n"));
|
||||
let ranges: Vec<(u32, u32)> = chunks.iter().map(|held| (held.line_start, held.line_end)).collect();
|
||||
assert_eq!(ranges, [(1, 1), (2, 2), (3, 3)], "two 700-char lines would join to 1,401");
|
||||
let fits = chunk(&format!("{}\n{}\n", "y".repeat(699), "y".repeat(700)));
|
||||
assert_eq!(fits.len(), 1, "699 + newline + 700 is exactly 1,400");
|
||||
|
||||
let long = "\u{e9}".repeat(CHUNK_CHARS_MAX + 50);
|
||||
let chunks = chunk(&format!("before\n{long}\nafter\n"));
|
||||
let shown: Vec<(u32, u32, usize)> =
|
||||
chunks.iter().map(|held| (held.line_start, held.line_end, held.text.chars().count())).collect();
|
||||
assert_eq!(shown, [(1, 1, 6), (2, 2, CHUNK_CHARS_MAX), (3, 3, 5)]);
|
||||
assert!(chunks[1].text.chars().all(|c| c == '\u{e9}'));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_embedded_text_names_the_file_and_the_payload_text_does_not() {
|
||||
let held = Chunk { line_start: 3, line_end: 4, text: "fn a() {}".into() };
|
||||
assert_eq!(embedded_text("Huginn", "src/a.rs", &held), "Huginn/src/a.rs\nfn a() {}");
|
||||
}
|
||||
}
|
||||
150
crates/huginn-daemon/src/sources/failover.rs
Normal file
150
crates/huginn-daemon/src/sources/failover.rs
Normal file
|
|
@ -0,0 +1,150 @@
|
|||
//! The writer's embedder: an ordered list of endpoints serving one model. The
|
||||
//! first that answers with the run's model identity embeds; one that is
|
||||
//! unreachable, failing, or serving another model (name, digest or length) is
|
||||
//! skipped and named in the log, and is never written with. Daemons do not use
|
||||
//! this: a query embeds on its daemon's own embedder.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{Result, bail};
|
||||
|
||||
use crate::index::{Embedder, ModelIdentity};
|
||||
|
||||
/// Endpoints in the order they are preferred, each with the name the log uses.
|
||||
/// The first identity read of a run pins the model; every endpoint must then
|
||||
/// report that identity before it is written with.
|
||||
pub struct FailoverEmbedder<E: Embedder> {
|
||||
endpoints: Vec<(String, E)>,
|
||||
pinned: Option<ModelIdentity>,
|
||||
}
|
||||
|
||||
impl<E: Embedder> FailoverEmbedder<E> {
|
||||
pub fn new(endpoints: Vec<(String, E)>) -> Self {
|
||||
Self { endpoints, pinned: None }
|
||||
}
|
||||
|
||||
/// The first endpoint at or after `from` that answers with the pinned
|
||||
/// identity. The others are named in the log and in the error if none does.
|
||||
fn serving(&mut self, from: usize, within: Duration) -> Result<usize> {
|
||||
let mut skipped = Vec::new();
|
||||
for at in from..self.endpoints.len() {
|
||||
let name = self.endpoints[at].0.clone();
|
||||
match self.endpoints[at].1.model_identity(within) {
|
||||
Ok(identity) => {
|
||||
let pinned = self.pinned.get_or_insert_with(|| identity.clone());
|
||||
if *pinned == identity {
|
||||
return Ok(at);
|
||||
}
|
||||
eprintln!("huginn-sources: embedder {name} serves {identity:?}, not {pinned:?}; skipped");
|
||||
skipped.push(format!("{name}: another model"));
|
||||
}
|
||||
Err(error) => {
|
||||
eprintln!("huginn-sources: embedder {name} is unavailable: {error:#}; trying the next");
|
||||
skipped.push(format!("{name}: {error:#}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
bail!("no embedder endpoint can serve this run: {}", skipped.join("; "))
|
||||
}
|
||||
}
|
||||
|
||||
impl<E: Embedder> Embedder for FailoverEmbedder<E> {
|
||||
fn model_identity(&mut self, within: Duration) -> Result<ModelIdentity> {
|
||||
self.serving(0, within)?;
|
||||
Ok(self.pinned.clone().expect("serving pins the model"))
|
||||
}
|
||||
|
||||
fn embed(&mut self, texts: &[String], within: Duration) -> Result<Vec<Vec<f32>>> {
|
||||
let mut failures = Vec::new();
|
||||
let mut from = 0;
|
||||
while from < self.endpoints.len() {
|
||||
let at = match self.serving(from, within) {
|
||||
Ok(at) => at,
|
||||
Err(error) => bail!("no embedder endpoint embedded the batch: {}{error:#}", failures.join("; ")),
|
||||
};
|
||||
match self.endpoints[at].1.embed(texts, within) {
|
||||
Ok(vectors) => return Ok(vectors),
|
||||
Err(error) => {
|
||||
let name = &self.endpoints[at].0;
|
||||
eprintln!("huginn-sources: embedder {name} failed an embed: {error:#}; trying the next");
|
||||
failures.push(format!("{name}: {error:#}; "));
|
||||
from = at + 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
bail!("no embedder endpoint embedded the batch: {}", failures.concat())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::index::fakes::FakeEmbedder;
|
||||
|
||||
const WITHIN: Duration = Duration::from_secs(1);
|
||||
|
||||
fn failover(spec: &[(&str, &str)]) -> (FailoverEmbedder<FakeEmbedder>, Vec<FakeEmbedder>) {
|
||||
let fakes: Vec<FakeEmbedder> = spec.iter().map(|(_, digest)| FakeEmbedder::new(digest)).collect();
|
||||
let endpoints = spec.iter().zip(&fakes).map(|((name, _), fake)| ((*name).to_owned(), fake.clone())).collect();
|
||||
(FailoverEmbedder::new(endpoints), fakes)
|
||||
}
|
||||
|
||||
fn used(fake: &FakeEmbedder) -> usize {
|
||||
fake.texts().len()
|
||||
}
|
||||
|
||||
/// The first healthy endpoint in order embeds; the others are not touched.
|
||||
#[test]
|
||||
fn the_first_healthy_endpoint_embeds() {
|
||||
let (mut embedder, fakes) = failover(&[("raven", "d1"), ("starfire", "d1"), ("yggdrasil", "d1")]);
|
||||
embedder.embed(&["a".into()], WITHIN).unwrap();
|
||||
assert_eq!((used(&fakes[0]), used(&fakes[1]), used(&fakes[2])), (1, 0, 0));
|
||||
}
|
||||
|
||||
/// A down endpoint hands the batch to the next, the CPU box embeds when
|
||||
/// both desktops are down, and the preferred one is used again once it
|
||||
/// answers.
|
||||
#[test]
|
||||
fn a_down_endpoint_hands_over_to_the_next_and_comes_back() {
|
||||
let (mut embedder, fakes) = failover(&[("raven", "d1"), ("starfire", "d1"), ("yggdrasil", "d1")]);
|
||||
fakes[0].state.lock().unwrap().down = true;
|
||||
embedder.embed(&["a".into()], WITHIN).unwrap();
|
||||
assert_eq!((used(&fakes[0]), used(&fakes[1]), used(&fakes[2])), (0, 1, 0));
|
||||
fakes[1].state.lock().unwrap().down = true;
|
||||
embedder.embed(&["b".into()], WITHIN).unwrap();
|
||||
assert_eq!(used(&fakes[2]), 1, "both desktops down, the CPU embeds");
|
||||
fakes[0].state.lock().unwrap().down = false;
|
||||
embedder.embed(&["c".into()], WITHIN).unwrap();
|
||||
assert_eq!(used(&fakes[0]), 1, "the preferred endpoint is used again when it answers");
|
||||
}
|
||||
|
||||
/// An endpoint serving another digest or another length is skipped and
|
||||
/// never written with, even though it answers; with nothing left the run
|
||||
/// fails naming each endpoint.
|
||||
#[test]
|
||||
fn a_mismatched_endpoint_is_skipped_and_never_embeds() {
|
||||
let (mut embedder, fakes) = failover(&[("raven", "d1"), ("starfire", "d2"), ("yggdrasil", "d1")]);
|
||||
embedder.model_identity(WITHIN).unwrap();
|
||||
fakes[0].state.lock().unwrap().down = true;
|
||||
embedder.embed(&["a".into()], WITHIN).unwrap();
|
||||
assert_eq!((used(&fakes[1]), used(&fakes[2])), (0, 1), "starfire serves d2 and is skipped");
|
||||
fakes[2].state.lock().unwrap().identity.dimensions = 8;
|
||||
let error = format!("{:#}", embedder.embed(&["b".into()], WITHIN).unwrap_err());
|
||||
assert!(error.contains("raven") && error.contains("starfire: another model") && error.contains("yggdrasil: another model"), "{error}");
|
||||
assert_eq!(used(&fakes[1]), 0);
|
||||
}
|
||||
|
||||
/// With every endpoint down the identity read fails naming them all, and
|
||||
/// the first one back up serves.
|
||||
#[test]
|
||||
fn with_every_endpoint_down_the_run_fails_naming_them_all() {
|
||||
let (mut embedder, fakes) = failover(&[("raven", "d1"), ("yggdrasil", "d1")]);
|
||||
for fake in &fakes {
|
||||
fake.state.lock().unwrap().down = true;
|
||||
}
|
||||
let error = format!("{:#}", embedder.model_identity(WITHIN).unwrap_err());
|
||||
assert!(error.contains("raven") && error.contains("yggdrasil"), "{error}");
|
||||
fakes[1].state.lock().unwrap().down = false;
|
||||
assert_eq!(embedder.model_identity(WITHIN).unwrap().digest, "d1");
|
||||
}
|
||||
}
|
||||
93
crates/huginn-daemon/src/sources/forge.rs
Normal file
93
crates/huginn-daemon/src/sources/forge.rs
Normal file
|
|
@ -0,0 +1,93 @@
|
|||
//! The forge's catalog over its own loopback API. JSON for the forge is spoken
|
||||
//! here and in `index/qdrant.rs`, and nowhere else.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{Context, Result, ensure};
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{Forge, SourceRepo};
|
||||
use crate::index::{Http, Method};
|
||||
|
||||
/// The organisation whose repositories are the corpus.
|
||||
pub const ORG: &str = "GameCult";
|
||||
|
||||
const PAGE: usize = 50;
|
||||
const CALL_TIMEOUT: Duration = Duration::from_secs(60);
|
||||
|
||||
/// The forge as one HTTP address. Repos are listed through the organisation's
|
||||
/// API, the unauthenticated view of it: what an anonymous reader may see is
|
||||
/// what is indexed.
|
||||
pub struct HttpForge {
|
||||
http: Http,
|
||||
base: String,
|
||||
}
|
||||
|
||||
impl HttpForge {
|
||||
pub fn new(base_url: &str) -> Self {
|
||||
Self { http: Http::new(base_url, CALL_TIMEOUT), base: base_url.trim_end_matches('/').to_owned() }
|
||||
}
|
||||
}
|
||||
|
||||
/// The repos worth indexing in one page of the forge's answer: public,
|
||||
/// unarchived, not a mirror, not empty. A flag the forge does not state counts
|
||||
/// against the repo.
|
||||
pub(crate) fn indexable(page: &Value) -> Result<Vec<SourceRepo>> {
|
||||
let repos = page.as_array().context("the forge's repo page is not a list")?;
|
||||
let flag = |repo: &Value, key: &str| repo[key].as_bool().unwrap_or(true);
|
||||
Ok(repos
|
||||
.iter()
|
||||
.filter(|repo| !flag(repo, "private") && !flag(repo, "archived") && !flag(repo, "mirror") && !flag(repo, "empty"))
|
||||
.filter_map(|repo| {
|
||||
Some(SourceRepo { name: repo["name"].as_str()?.to_owned(), default_branch: repo["default_branch"].as_str()?.to_owned() })
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
impl Forge for HttpForge {
|
||||
fn catalog(&mut self) -> Result<Vec<SourceRepo>> {
|
||||
let mut found = Vec::new();
|
||||
for page in 1.. {
|
||||
let path = format!("/api/v1/orgs/{ORG}/repos?page={page}&limit={PAGE}");
|
||||
let (status, answer) = self.http.request(Method::Get, &path, None)?;
|
||||
ensure!(status == 200, "the forge answered {status} to {path}");
|
||||
let parsed: Value = serde_json::from_str(&answer).context("the forge's repo page is not JSON")?;
|
||||
let listed = parsed.as_array().map_or(0, Vec::len);
|
||||
found.extend(indexable(&parsed)?);
|
||||
if listed < PAGE {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(found)
|
||||
}
|
||||
|
||||
fn clone_url(&self, name: &str) -> String {
|
||||
format!("{}/{ORG}/{name}.git", self.base)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
/// Only public, live, original, non-empty repos are the corpus; a repo the
|
||||
/// forge does not say is safe is not indexed.
|
||||
#[test]
|
||||
fn only_public_unarchived_original_nonempty_repos_are_indexed() {
|
||||
let repo = |name: &str, private: bool, archived: bool, mirror: bool, empty: bool| {
|
||||
json!({ "name": name, "default_branch": "main", "private": private, "archived": archived, "mirror": mirror, "empty": empty })
|
||||
};
|
||||
let page = json!([
|
||||
repo("Open", false, false, false, false),
|
||||
repo("Secret", true, false, false, false),
|
||||
repo("Old", false, true, false, false),
|
||||
repo("Copy", false, false, true, false),
|
||||
repo("Void", false, false, false, true),
|
||||
{ "name": "Silent", "default_branch": "main" },
|
||||
]);
|
||||
let names: Vec<String> = indexable(&page).unwrap().into_iter().map(|repo| repo.name).collect();
|
||||
assert_eq!(names, ["Open"]);
|
||||
}
|
||||
}
|
||||
202
crates/huginn-daemon/src/sources/git.rs
Normal file
202
crates/huginn-daemon/src/sources/git.rs
Normal file
|
|
@ -0,0 +1,202 @@
|
|||
//! Git, by shelling out: bare mirrors of the forge's repos, what changed
|
||||
//! between two commits, and the blobs at one. Nothing here knows an index or an
|
||||
//! embedder; the mirrors are cache and can be deleted whenever.
|
||||
|
||||
use std::io::{BufRead, BufReader, Read, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};
|
||||
|
||||
use anyhow::{Context, Result, bail, ensure};
|
||||
|
||||
/// A blob over this many bytes is not indexed.
|
||||
pub const BLOB_BYTES_MAX: u64 = 256 * 1024;
|
||||
|
||||
/// Extensions of files that are not text worth finding, VoidBot's list, plus
|
||||
/// Unity's serialised YAML (huge, machine written) and lock files.
|
||||
const SKIPPED_EXTENSIONS: &[&str] = &[
|
||||
"7z", "avi", "bmp", "class", "dll", "dylib", "eot", "exe", "gif", "gz", "ico", "jar", "jpeg", "jpg", "lock", "map", "mp3",
|
||||
"mp4", "mov", "ogg", "otf", "pdf", "png", "rar", "so", "svg", "tar", "ttf", "wav", "webm", "webp", "woff", "woff2", "zip",
|
||||
"meta", "prefab", "unity", "asset", "mat", "anim", "controller", "physicmaterial",
|
||||
];
|
||||
|
||||
/// Lock files whose extension is not `.lock`.
|
||||
const SKIPPED_NAMES: &[&str] = &["package-lock.json", "pnpm-lock.yaml", "yarn.lock", "bun.lockb", "composer.lock"];
|
||||
|
||||
const LFS_POINTER: &[u8] = b"version https://git-lfs.github.com/spec/v1";
|
||||
|
||||
/// Whether a path is not worth reading, by name alone.
|
||||
pub fn skipped_path(path: &str) -> bool {
|
||||
let name = path.rsplit('/').next().unwrap_or(path).to_ascii_lowercase();
|
||||
if SKIPPED_NAMES.contains(&name.as_str()) {
|
||||
return true;
|
||||
}
|
||||
name.rsplit_once('.').is_some_and(|(_, extension)| SKIPPED_EXTENSIONS.contains(&extension))
|
||||
}
|
||||
|
||||
/// The text of a blob worth indexing, or `None`: a NUL byte, bytes that are
|
||||
/// not UTF-8, or a Git LFS pointer make it not text.
|
||||
pub fn indexable_text(bytes: &[u8]) -> Option<String> {
|
||||
if bytes.contains(&0) || bytes.starts_with(LFS_POINTER) {
|
||||
return None;
|
||||
}
|
||||
String::from_utf8(bytes.to_vec()).ok()
|
||||
}
|
||||
|
||||
/// One file of a tree: its path, the blob it holds and that blob's size.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct TreeFile {
|
||||
pub path: String,
|
||||
pub blob: String,
|
||||
pub size: u64,
|
||||
}
|
||||
|
||||
fn git(dir: &Path, args: &[&str]) -> Command {
|
||||
let mut command = Command::new("git");
|
||||
command.arg("-C").arg(dir).args(args).env("GIT_TERMINAL_PROMPT", "0");
|
||||
command
|
||||
}
|
||||
|
||||
fn run(mut command: Command, what: &str) -> Result<Vec<u8>> {
|
||||
let output = command.output().with_context(|| format!("starting git to {what}"))?;
|
||||
ensure!(
|
||||
output.status.success(),
|
||||
"git could not {what}: {}",
|
||||
String::from_utf8_lossy(&output.stderr).trim().chars().take(400).collect::<String>()
|
||||
);
|
||||
Ok(output.stdout)
|
||||
}
|
||||
|
||||
/// A bare mirror of one forge repo's default branch.
|
||||
pub struct Mirror {
|
||||
dir: PathBuf,
|
||||
branch: String,
|
||||
}
|
||||
|
||||
impl Mirror {
|
||||
pub fn path_for(root: &Path, name: &str) -> PathBuf {
|
||||
root.join(format!("{name}.git"))
|
||||
}
|
||||
|
||||
pub fn new(root: &Path, name: &str, branch: &str) -> Self {
|
||||
Self { dir: Self::path_for(root, name), branch: branch.to_owned() }
|
||||
}
|
||||
|
||||
/// Fetches the default branch, creating the mirror if it is absent, and
|
||||
/// returns the commit it now points at. The branch is forced: a rewritten
|
||||
/// history is followed, not refused.
|
||||
pub fn fetch(&self, remote: &str) -> Result<String> {
|
||||
if !self.dir.join("HEAD").exists() {
|
||||
std::fs::create_dir_all(&self.dir).with_context(|| format!("creating {}", self.dir.display()))?;
|
||||
run(Command::new("git").arg("init").arg("--bare").arg("--quiet").arg(&self.dir), "initialise a mirror")?;
|
||||
}
|
||||
let refspec = format!("+refs/heads/{0}:refs/heads/{0}", self.branch);
|
||||
run(git(&self.dir, &["fetch", "--quiet", "--no-tags", remote, &refspec]), "fetch the default branch")?;
|
||||
self.head()
|
||||
}
|
||||
|
||||
pub fn head(&self) -> Result<String> {
|
||||
let out = run(git(&self.dir, &["rev-parse", "--verify", &format!("refs/heads/{}", self.branch)]), "read the branch head")?;
|
||||
Ok(String::from_utf8_lossy(&out).trim().to_owned())
|
||||
}
|
||||
|
||||
/// Whether `old` is an ancestor of `new`. A commit the mirror no longer
|
||||
/// holds is not.
|
||||
pub fn is_ancestor(&self, old: &str, new: &str) -> bool {
|
||||
git(&self.dir, &["merge-base", "--is-ancestor", old, new]).status().is_ok_and(|status| status.success())
|
||||
}
|
||||
|
||||
/// The paths that differ between two commits, a rename counted as a delete
|
||||
/// and an add.
|
||||
pub fn changed_paths(&self, old: &str, new: &str) -> Result<Vec<String>> {
|
||||
let out = run(git(&self.dir, &["diff", "--no-renames", "--name-only", "-z", old, new]), "diff two commits")?;
|
||||
Ok(out.split(|byte| *byte == 0).filter(|path| !path.is_empty()).map(|path| String::from_utf8_lossy(path).into_owned()).collect())
|
||||
}
|
||||
|
||||
/// Every regular file at `commit`; symlinks and submodules are not files.
|
||||
pub fn files(&self, commit: &str) -> Result<Vec<TreeFile>> {
|
||||
let out = run(git(&self.dir, &["ls-tree", "-r", "-z", "-l", commit]), "list a tree")?;
|
||||
let mut files = Vec::new();
|
||||
for entry in out.split(|byte| *byte == 0).filter(|entry| !entry.is_empty()) {
|
||||
let text = String::from_utf8_lossy(entry);
|
||||
let (meta, path) = text.split_once('\t').with_context(|| format!("an ls-tree entry has no tab: {text:?}"))?;
|
||||
let mut fields = meta.split_whitespace();
|
||||
let (mode, kind, blob, size) = (fields.next(), fields.next(), fields.next(), fields.next());
|
||||
let (Some(mode), Some("blob"), Some(blob), Some(size)) = (mode, kind, blob, size) else { continue };
|
||||
if mode == "120000" {
|
||||
continue;
|
||||
}
|
||||
let Ok(size) = size.parse::<u64>() else { continue };
|
||||
files.push(TreeFile { path: path.to_owned(), blob: blob.to_owned(), size });
|
||||
}
|
||||
Ok(files)
|
||||
}
|
||||
|
||||
/// A reader of blobs by id, one `cat-file --batch` process.
|
||||
pub fn blobs(&self) -> Result<Blobs> {
|
||||
let mut child = git(&self.dir, &["cat-file", "--batch"])
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::null())
|
||||
.spawn()
|
||||
.context("starting git cat-file")?;
|
||||
let stdin = child.stdin.take().context("git cat-file has no stdin")?;
|
||||
let stdout = BufReader::new(child.stdout.take().context("git cat-file has no stdout")?);
|
||||
Ok(Blobs { child, stdin, stdout })
|
||||
}
|
||||
}
|
||||
|
||||
/// `git cat-file --batch`, asked one blob at a time.
|
||||
pub struct Blobs {
|
||||
child: Child,
|
||||
stdin: ChildStdin,
|
||||
stdout: BufReader<ChildStdout>,
|
||||
}
|
||||
|
||||
impl Blobs {
|
||||
pub fn read(&mut self, blob: &str) -> Result<Vec<u8>> {
|
||||
writeln!(self.stdin, "{blob}").and_then(|()| self.stdin.flush()).context("asking git cat-file for a blob")?;
|
||||
let mut header = String::new();
|
||||
self.stdout.read_line(&mut header).context("reading git cat-file's header")?;
|
||||
let mut fields = header.split_whitespace();
|
||||
let (Some(_), Some(kind), Some(size)) = (fields.next(), fields.next(), fields.next()) else {
|
||||
bail!("git cat-file does not hold blob {blob}: {}", header.trim());
|
||||
};
|
||||
ensure!(kind == "blob", "git cat-file says {blob} is a {kind}");
|
||||
let size: usize = size.parse().with_context(|| format!("a blob size of {size:?}"))?;
|
||||
let mut bytes = vec![0; size];
|
||||
self.stdout.read_exact(&mut bytes).context("reading a blob")?;
|
||||
let mut newline = [0_u8; 1];
|
||||
self.stdout.read_exact(&mut newline).context("reading a blob's end")?;
|
||||
Ok(bytes)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Blobs {
|
||||
fn drop(&mut self) {
|
||||
let _ = self.child.kill();
|
||||
let _ = self.child.wait();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn binary_unity_and_lock_paths_are_skipped_by_name_and_source_is_not() {
|
||||
for path in ["a.png", "dir/B.PNG", "x.lock", "Assets/a.prefab", "Assets/a.unity", "a.cs.meta", "p/package-lock.json", "Cargo.lock"] {
|
||||
assert!(skipped_path(path), "{path}");
|
||||
}
|
||||
for path in ["src/main.rs", "README.md", "docs/lock.md", "Assets/Scripts/Player.cs", "Makefile", ".gitignore", "locksmith"] {
|
||||
assert!(!skipped_path(path), "{path}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_nul_byte_an_lfs_pointer_and_invalid_utf8_are_not_text() {
|
||||
assert_eq!(indexable_text(b"fn main() {}\n").as_deref(), Some("fn main() {}\n"));
|
||||
assert_eq!(indexable_text(b"a\0b"), None);
|
||||
assert_eq!(indexable_text(b"version https://git-lfs.github.com/spec/v1\noid sha256:ab\nsize 12\n"), None);
|
||||
assert_eq!(indexable_text(&[0xff, 0xfe, b'a']), None);
|
||||
}
|
||||
}
|
||||
409
crates/huginn-daemon/src/sources/mod.rs
Normal file
409
crates/huginn-daemon/src/sources/mod.rs
Normal file
|
|
@ -0,0 +1,409 @@
|
|||
//! The source index: the public GameCult forge repos' code and docs as chunks
|
||||
//! in one Qdrant collection, `huginn_sources`, for semantic search.
|
||||
//!
|
||||
//! Two truths, one writer. The forge is the truth of which repos exist and what
|
||||
//! they hold; the collection is derived from it and can be deleted whenever.
|
||||
//! `refresh` is the only writer of the collection, run by the `huginn-sources`
|
||||
//! oneshot job: first build, incremental refresh, rebuild after a model or
|
||||
//! chunker change, and retirement are all the one routine (delete by filter,
|
||||
//! ingest at a commit, rewrite the cursors). The daemon only reads: `search`
|
||||
//! is its whole part, and nothing in `serve` or the worker calls a write.
|
||||
//!
|
||||
//! A cursor is the commit whose chunks the collection holds for a repo. It
|
||||
//! lives in the collection's own metadata, so it is lost with the collection
|
||||
//! and the two cannot disagree about what is held. It is advanced only after
|
||||
//! the repo's deletes and upserts are acknowledged: it may lag and never leads.
|
||||
//! The cursors object is written whole by this job alone, once per repo.
|
||||
|
||||
pub mod chunk;
|
||||
pub mod failover;
|
||||
pub mod forge;
|
||||
pub mod git;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::path::Path;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use anyhow::{Context, Result, bail, ensure};
|
||||
use huginn_mind::{SourceHit, SourceHits, SourceModel, SourceQuery};
|
||||
|
||||
use crate::index::{BATCH, BULK_EMBED_TIMEOUT, Budget, Embedder, INDEX_CALL_TIMEOUT, MANAGED_BY, point_id};
|
||||
use chunk::{CHUNKER_VERSION, embedded_text};
|
||||
use git::{BLOB_BYTES_MAX, Mirror, TreeFile};
|
||||
|
||||
/// The collection this organ's source index lives in.
|
||||
pub const COLLECTION: &str = "huginn_sources";
|
||||
|
||||
/// What the collection says its corpus is.
|
||||
pub const CORPUS: &str = "forge:GameCult";
|
||||
|
||||
/// The instruction Qwen3 embedding models take in front of a query about code.
|
||||
pub const SOURCE_QUERY_INSTRUCTION: &str =
|
||||
"Given a question about a codebase, retrieve the source code and documentation passages that answer it";
|
||||
|
||||
/// Paths named in one delete call.
|
||||
const DELETE_PATHS_PER_CALL: usize = 256;
|
||||
|
||||
/// What the collection says about how it was made. A model, a digest, a length
|
||||
/// or a chunker version that differs from the embedder's and this build's makes
|
||||
/// every vector in it incomparable, so it is rebuilt.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct SourcesMeta {
|
||||
pub managed_by: String,
|
||||
pub corpus: String,
|
||||
pub model: String,
|
||||
pub model_digest: String,
|
||||
pub dimensions: u32,
|
||||
pub chunker_version: u32,
|
||||
}
|
||||
|
||||
/// What the vector store holds under `COLLECTION`.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum DescribedSources {
|
||||
Absent,
|
||||
/// A collection that carries no metadata this organ wrote.
|
||||
Unlabelled,
|
||||
Labelled { meta: SourcesMeta, cursors: BTreeMap<String, String> },
|
||||
}
|
||||
|
||||
/// What travels beside a chunk's vector.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct SourcePayload {
|
||||
pub repo: String,
|
||||
pub path: String,
|
||||
pub commit: String,
|
||||
pub blob: String,
|
||||
pub line_start: u32,
|
||||
pub line_end: u32,
|
||||
pub chunk_index: u32,
|
||||
pub ext: String,
|
||||
pub text: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct SourcePoint {
|
||||
pub id: String,
|
||||
pub vector: Vec<f32>,
|
||||
pub payload: SourcePayload,
|
||||
}
|
||||
|
||||
/// One chunk a vector search returned and how near it scored.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct SourceRow {
|
||||
pub score: f32,
|
||||
pub payload: SourcePayload,
|
||||
}
|
||||
|
||||
/// The vector store as the source index uses it. The writes are
|
||||
/// `huginn-sources`'s alone; `describe_sources` and `search_sources` are what a
|
||||
/// daemon calls.
|
||||
pub trait SourceIndex {
|
||||
fn describe_sources(&mut self, within: Duration) -> Result<DescribedSources>;
|
||||
/// Drops the collection if it exists and creates it empty with this
|
||||
/// metadata, no cursors, and keyword indexes on `repo` and `path`.
|
||||
fn recreate_sources(&mut self, meta: &SourcesMeta) -> Result<()>;
|
||||
/// Replaces the cursors object whole.
|
||||
fn set_cursors(&mut self, cursors: &BTreeMap<String, String>) -> Result<()>;
|
||||
/// Deletes a repo's chunks: those at `paths`, or all of them when `None`.
|
||||
/// Acknowledged before it returns.
|
||||
fn delete_by(&mut self, repo: &str, paths: Option<&[String]>) -> Result<()>;
|
||||
/// Writes chunks, acknowledged before it returns.
|
||||
fn upsert_sources(&mut self, points: &[SourcePoint]) -> Result<()>;
|
||||
/// The `limit` nearest chunks, best first, in `repo` alone when given.
|
||||
fn search_sources(&mut self, vector: &[f32], repo: Option<&str>, limit: u32, within: Duration) -> Result<Vec<SourceRow>>;
|
||||
}
|
||||
|
||||
/// A repo the forge lists as indexable.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct SourceRepo {
|
||||
pub name: String,
|
||||
pub default_branch: String,
|
||||
}
|
||||
|
||||
/// The forge: which repos exist, and where to fetch them from.
|
||||
pub trait Forge {
|
||||
fn catalog(&mut self) -> Result<Vec<SourceRepo>>;
|
||||
fn clone_url(&self, name: &str) -> String;
|
||||
}
|
||||
|
||||
/// What one repo's refresh did, for the journal.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Report {
|
||||
pub name: String,
|
||||
pub old: Option<String>,
|
||||
pub new: String,
|
||||
pub files: u32,
|
||||
pub chunks: u32,
|
||||
pub seconds: f64,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for Report {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(
|
||||
formatter,
|
||||
"{} {}..{} files={} chunks={} seconds={:.1}",
|
||||
self.name,
|
||||
self.old.as_deref().map_or("-", |old| &old[..old.len().min(7)]),
|
||||
&self.new[..self.new.len().min(7)],
|
||||
self.files,
|
||||
self.chunks,
|
||||
self.seconds
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Where the mirrors live under the state root.
|
||||
pub fn mirrors_dir(state_root: &Path) -> std::path::PathBuf {
|
||||
state_root.join("sources").join("repos")
|
||||
}
|
||||
|
||||
/// The id of a chunk: content-addressed by where it is, which blob it was cut
|
||||
/// from, the chunker and its position, so the same blob chunked the same way
|
||||
/// is the same point.
|
||||
pub fn chunk_id(repo: &str, path: &str, blob: &str, index: u32) -> String {
|
||||
point_id(&format!("{repo}\0{path}\0{blob}\0{CHUNKER_VERSION}\0{index}"))
|
||||
}
|
||||
|
||||
fn extension(path: &str) -> String {
|
||||
let name = path.rsplit('/').next().unwrap_or(path);
|
||||
name.rsplit_once('.').map(|(_, extension)| extension.to_ascii_lowercase()).unwrap_or_default()
|
||||
}
|
||||
|
||||
/// The names the refresh reports as failed; empty when every repo converged.
|
||||
#[derive(Debug, Default)]
|
||||
pub struct Refreshed {
|
||||
pub failed: Vec<String>,
|
||||
}
|
||||
|
||||
/// Brings the collection to the forge's current state. A collection that
|
||||
/// exists without Huginn's label is refused untouched. One repo's failure is
|
||||
/// logged and the others go on; the caller exits nonzero if any failed.
|
||||
pub fn refresh<E: Embedder, S: SourceIndex, F: Forge>(
|
||||
embedder: &mut E,
|
||||
index: &mut S,
|
||||
forge: &mut F,
|
||||
state_root: &Path,
|
||||
on_repo: &mut dyn FnMut(&Report),
|
||||
) -> Result<Refreshed> {
|
||||
let identity = embedder.model_identity(BULK_EMBED_TIMEOUT).context("reading the embedding model's identity")?;
|
||||
let wanted = SourcesMeta {
|
||||
managed_by: MANAGED_BY.into(),
|
||||
corpus: CORPUS.into(),
|
||||
model: identity.name.clone(),
|
||||
model_digest: identity.digest.clone(),
|
||||
dimensions: identity.dimensions,
|
||||
chunker_version: CHUNKER_VERSION,
|
||||
};
|
||||
let mut cursors = match index.describe_sources(INDEX_CALL_TIMEOUT).context("describing the source collection")? {
|
||||
DescribedSources::Absent => {
|
||||
index.recreate_sources(&wanted).context("creating the source collection")?;
|
||||
BTreeMap::new()
|
||||
}
|
||||
DescribedSources::Unlabelled => {
|
||||
bail!("collection {COLLECTION} exists and carries no Huginn metadata; it is not this job's to rewrite")
|
||||
}
|
||||
DescribedSources::Labelled { meta, cursors } if meta == wanted => cursors,
|
||||
DescribedSources::Labelled { .. } => {
|
||||
index.recreate_sources(&wanted).context("rebuilding the source collection")?;
|
||||
BTreeMap::new()
|
||||
}
|
||||
};
|
||||
|
||||
let catalog = forge.catalog().context("listing the forge's repos")?;
|
||||
let mirrors = mirrors_dir(state_root);
|
||||
retire(index, &catalog, &mirrors, &mut cursors)?;
|
||||
|
||||
let mut refreshed = Refreshed::default();
|
||||
for repo in &catalog {
|
||||
let started = Instant::now();
|
||||
match refresh_repo(embedder, index, forge, &mirrors, repo, &identity.dimensions, &mut cursors) {
|
||||
Ok(mut report) => {
|
||||
report.seconds = started.elapsed().as_secs_f64();
|
||||
on_repo(&report);
|
||||
}
|
||||
Err(error) => {
|
||||
eprintln!("huginn-sources: {} failed, the others go on: {error:#}", repo.name);
|
||||
refreshed.failed.push(repo.name.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(refreshed)
|
||||
}
|
||||
|
||||
/// Drops every repo the catalog no longer lists: its chunks, its cursor and
|
||||
/// its mirror, in that order, so a failure leaves the cursor to find it again.
|
||||
fn retire<S: SourceIndex>(
|
||||
index: &mut S,
|
||||
catalog: &[SourceRepo],
|
||||
mirrors: &Path,
|
||||
cursors: &mut BTreeMap<String, String>,
|
||||
) -> Result<()> {
|
||||
let listed = |name: &str| catalog.iter().any(|repo| repo.name.eq_ignore_ascii_case(name));
|
||||
let mut gone: BTreeSet<String> = cursors.keys().filter(|name| !listed(name)).cloned().collect();
|
||||
if let Ok(entries) = std::fs::read_dir(mirrors) {
|
||||
for entry in entries.flatten() {
|
||||
if let Some(name) = entry.file_name().to_string_lossy().strip_suffix(".git")
|
||||
&& !listed(name)
|
||||
{
|
||||
gone.insert(name.to_owned());
|
||||
}
|
||||
}
|
||||
}
|
||||
for name in gone {
|
||||
index.delete_by(&name, None).with_context(|| format!("deleting the chunks of retired {name}"))?;
|
||||
if cursors.remove(&name).is_some() {
|
||||
index.set_cursors(cursors).context("dropping a retired repo's cursor")?;
|
||||
}
|
||||
let mirror = Mirror::path_for(mirrors, &name);
|
||||
if mirror.exists() {
|
||||
std::fs::remove_dir_all(&mirror).with_context(|| format!("removing {}", mirror.display()))?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn refresh_repo<E: Embedder, S: SourceIndex, F: Forge>(
|
||||
embedder: &mut E,
|
||||
index: &mut S,
|
||||
forge: &F,
|
||||
mirrors: &Path,
|
||||
repo: &SourceRepo,
|
||||
dimensions: &u32,
|
||||
cursors: &mut BTreeMap<String, String>,
|
||||
) -> Result<Report> {
|
||||
let mirror = Mirror::new(mirrors, &repo.name, &repo.default_branch);
|
||||
let head = mirror.fetch(&forge.clone_url(&repo.name))?;
|
||||
let old = cursors.get(&repo.name).cloned();
|
||||
let mut report = Report { name: repo.name.clone(), old: old.clone(), new: head.clone(), files: 0, chunks: 0, seconds: 0.0 };
|
||||
if old.as_deref() == Some(head.as_str()) {
|
||||
return Ok(report);
|
||||
}
|
||||
// Only an ancestor's diff says what changed; anything else (a first
|
||||
// build, a rewritten history) starts the repo over.
|
||||
let changed: Option<BTreeSet<String>> = match &old {
|
||||
Some(old) if mirror.is_ancestor(old, &head) => Some(mirror.changed_paths(old, &head)?.into_iter().collect()),
|
||||
_ => None,
|
||||
};
|
||||
match &changed {
|
||||
Some(paths) => {
|
||||
let paths: Vec<String> = paths.iter().cloned().collect();
|
||||
for slice in paths.chunks(DELETE_PATHS_PER_CALL) {
|
||||
index.delete_by(&repo.name, Some(slice)).context("deleting the changed paths' chunks")?;
|
||||
}
|
||||
}
|
||||
None => index.delete_by(&repo.name, None).context("deleting the repo's chunks")?,
|
||||
}
|
||||
let files: Vec<TreeFile> = mirror
|
||||
.files(&head)?
|
||||
.into_iter()
|
||||
.filter(|file| changed.as_ref().is_none_or(|paths| paths.contains(&file.path)))
|
||||
.filter(|file| file.size <= BLOB_BYTES_MAX && !git::skipped_path(&file.path))
|
||||
.collect();
|
||||
let mut blobs = mirror.blobs()?;
|
||||
let mut batch = Batch::default();
|
||||
for file in &files {
|
||||
let Some(text) = git::indexable_text(&blobs.read(&file.blob)?) else { continue };
|
||||
let chunks = chunk::chunk(&text);
|
||||
if chunks.is_empty() {
|
||||
continue;
|
||||
}
|
||||
report.files += 1;
|
||||
for (position, held) in chunks.iter().enumerate() {
|
||||
let chunk_index = u32::try_from(position).unwrap_or(u32::MAX);
|
||||
batch.texts.push(embedded_text(&repo.name, &file.path, held));
|
||||
batch.payloads.push((
|
||||
chunk_id(&repo.name, &file.path, &file.blob, chunk_index),
|
||||
SourcePayload {
|
||||
repo: repo.name.clone(),
|
||||
path: file.path.clone(),
|
||||
commit: head.clone(),
|
||||
blob: file.blob.clone(),
|
||||
line_start: held.line_start,
|
||||
line_end: held.line_end,
|
||||
chunk_index,
|
||||
ext: extension(&file.path),
|
||||
text: held.text.clone(),
|
||||
},
|
||||
));
|
||||
report.chunks += 1;
|
||||
if batch.texts.len() >= BATCH {
|
||||
batch.write(embedder, index, *dimensions)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
batch.write(embedder, index, *dimensions)?;
|
||||
cursors.insert(repo.name.clone(), head);
|
||||
index.set_cursors(cursors).context("advancing the cursor")?;
|
||||
Ok(report)
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct Batch {
|
||||
texts: Vec<String>,
|
||||
payloads: Vec<(String, SourcePayload)>,
|
||||
}
|
||||
|
||||
impl Batch {
|
||||
/// Embeds and writes what has gathered; the batch is empty afterwards.
|
||||
fn write<E: Embedder, S: SourceIndex>(&mut self, embedder: &mut E, index: &mut S, dimensions: u32) -> Result<()> {
|
||||
if self.texts.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let vectors = embedder.embed(&self.texts, BULK_EMBED_TIMEOUT).context("embedding chunks")?;
|
||||
ensure!(vectors.len() == self.texts.len(), "the embedder returned {} vectors for {} chunks", vectors.len(), self.texts.len());
|
||||
let mut points = Vec::with_capacity(vectors.len());
|
||||
for ((id, payload), vector) in self.payloads.drain(..).zip(vectors) {
|
||||
ensure!(vector.len() == dimensions as usize, "the embedder returned a vector of {} where the collection holds {dimensions}", vector.len());
|
||||
points.push(SourcePoint { id, vector, payload });
|
||||
}
|
||||
self.texts.clear();
|
||||
index.upsert_sources(&points).context("writing chunks")
|
||||
}
|
||||
}
|
||||
|
||||
/// A source search, which a daemon answers from the collection as it finds it:
|
||||
/// it embeds the query with its own embedder, and refuses if the collection
|
||||
/// was built under another model, because vectors of two models do not compare.
|
||||
pub fn search<E: Embedder, S: SourceIndex>(embedder: &mut E, index: &mut S, query: &SourceQuery, within: Duration) -> Result<SourceHits> {
|
||||
let budget = Budget::new(within);
|
||||
let identity = embedder.model_identity(budget.left()?).context("reading the embedding model's identity")?;
|
||||
let built = match index.describe_sources(budget.left()?).context("describing the source collection")? {
|
||||
DescribedSources::Absent => bail!("source index not built"),
|
||||
DescribedSources::Unlabelled => bail!("source index not built: {COLLECTION} carries no Huginn metadata"),
|
||||
DescribedSources::Labelled { meta, .. } => meta,
|
||||
};
|
||||
ensure!(
|
||||
built.model == identity.name && built.model_digest == identity.digest && built.dimensions == identity.dimensions,
|
||||
"the source index was built under {} {} ({} dimensions) and this daemon embeds under {} {} ({} dimensions)",
|
||||
built.model,
|
||||
built.model_digest,
|
||||
built.dimensions,
|
||||
identity.name,
|
||||
identity.digest,
|
||||
identity.dimensions
|
||||
);
|
||||
let text = format!("Instruct: {SOURCE_QUERY_INSTRUCTION}\nQuery: {}", query.text.0);
|
||||
let mut vectors = embedder.embed(&[text], budget.left()?).context("embedding the query")?;
|
||||
ensure!(vectors.len() == 1, "the embedder returned {} vectors for one query", vectors.len());
|
||||
let vector = vectors.remove(0);
|
||||
ensure!(vector.len() == identity.dimensions as usize, "the embedder returned a query vector of {} where the collection holds {}", vector.len(), identity.dimensions);
|
||||
let rows = index.search_sources(&vector, query.repo.as_ref().map(|repo| repo.0.as_str()), query.top_k, budget.left()?)?;
|
||||
Ok(SourceHits {
|
||||
indexed_under: SourceModel { name: built.model, digest: built.model_digest },
|
||||
hits: rows
|
||||
.into_iter()
|
||||
.map(|row| SourceHit {
|
||||
repo: row.payload.repo,
|
||||
path: row.payload.path,
|
||||
commit: row.payload.commit,
|
||||
line_start: row.payload.line_start,
|
||||
line_end: row.payload.line_end,
|
||||
score: row.score,
|
||||
text: row.payload.text,
|
||||
})
|
||||
.collect(),
|
||||
})
|
||||
}
|
||||
408
crates/huginn-daemon/src/sources/tests.rs
Normal file
408
crates/huginn-daemon/src/sources/tests.rs
Normal file
|
|
@ -0,0 +1,408 @@
|
|||
//! The refresh against a fake vector store, a fake embedder and real git: a
|
||||
//! working repo on disk stands in for the forge, and each test commits to it as
|
||||
//! history would.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::Command;
|
||||
|
||||
use anyhow::Result;
|
||||
use huginn_mind::eureka_pipeline::{Line, Short};
|
||||
use huginn_mind::SourceQuery;
|
||||
|
||||
use super::*;
|
||||
use crate::index::fakes::{FakeEmbedder, FakeSources};
|
||||
|
||||
fn git_in(dir: &Path, args: &[&str]) -> String {
|
||||
let output = Command::new("git")
|
||||
.arg("-C")
|
||||
.arg(dir)
|
||||
.args(["-c", "user.name=t", "-c", "user.email=t@t", "-c", "commit.gpgsign=false", "-c", "core.autocrlf=false"])
|
||||
.args(args)
|
||||
.output()
|
||||
.unwrap();
|
||||
assert!(output.status.success(), "git {args:?}: {}", String::from_utf8_lossy(&output.stderr));
|
||||
String::from_utf8_lossy(&output.stdout).trim().to_owned()
|
||||
}
|
||||
|
||||
/// A working repo on disk, committed to by the test.
|
||||
struct Remote {
|
||||
dir: tempfile::TempDir,
|
||||
}
|
||||
|
||||
impl Remote {
|
||||
fn new() -> Self {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
git_in(dir.path(), &["init", "--quiet", "-b", "main"]);
|
||||
Self { dir }
|
||||
}
|
||||
|
||||
fn path(&self) -> PathBuf {
|
||||
self.dir.path().to_path_buf()
|
||||
}
|
||||
|
||||
/// Writes files, removes paths, commits everything and returns the commit.
|
||||
fn commit(&self, write: &[(&str, &[u8])], remove: &[&str]) -> String {
|
||||
for (path, bytes) in write {
|
||||
let full = self.dir.path().join(path);
|
||||
std::fs::create_dir_all(full.parent().unwrap()).unwrap();
|
||||
std::fs::write(full, bytes).unwrap();
|
||||
}
|
||||
for path in remove {
|
||||
std::fs::remove_file(self.dir.path().join(path)).unwrap();
|
||||
}
|
||||
git_in(self.dir.path(), &["add", "-A"]);
|
||||
git_in(self.dir.path(), &["commit", "--quiet", "--allow-empty", "-m", "c"]);
|
||||
git_in(self.dir.path(), &["rev-parse", "HEAD"])
|
||||
}
|
||||
}
|
||||
|
||||
struct FakeForge {
|
||||
repos: Vec<(SourceRepo, PathBuf)>,
|
||||
}
|
||||
|
||||
impl FakeForge {
|
||||
fn of(repos: &[(&str, &Remote)]) -> Self {
|
||||
Self {
|
||||
repos: repos
|
||||
.iter()
|
||||
.map(|(name, remote)| (SourceRepo { name: (*name).into(), default_branch: "main".into() }, remote.path()))
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Forge for FakeForge {
|
||||
fn catalog(&mut self) -> Result<Vec<SourceRepo>> {
|
||||
Ok(self.repos.iter().map(|(repo, _)| repo.clone()).collect())
|
||||
}
|
||||
|
||||
fn clone_url(&self, name: &str) -> String {
|
||||
self.repos.iter().find(|(repo, _)| repo.name == name).map_or_else(|| "/nowhere".into(), |(_, path)| path.display().to_string())
|
||||
}
|
||||
}
|
||||
|
||||
struct Rig {
|
||||
root: tempfile::TempDir,
|
||||
embedder: FakeEmbedder,
|
||||
index: FakeSources,
|
||||
}
|
||||
|
||||
impl Rig {
|
||||
fn new() -> Self {
|
||||
Self { root: tempfile::tempdir().unwrap(), embedder: FakeEmbedder::new("d1"), index: FakeSources::default() }
|
||||
}
|
||||
|
||||
fn run(&mut self, forge: &mut FakeForge) -> (Result<Refreshed>, Vec<Report>) {
|
||||
let mut reports = Vec::new();
|
||||
let refreshed = refresh(&mut self.embedder, &mut self.index, forge, self.root.path(), &mut |report| reports.push(report.clone()));
|
||||
(refreshed, reports)
|
||||
}
|
||||
|
||||
fn converged(&mut self, forge: &mut FakeForge) -> Vec<Report> {
|
||||
let (refreshed, reports) = self.run(forge);
|
||||
assert!(refreshed.unwrap().failed.is_empty());
|
||||
reports
|
||||
}
|
||||
|
||||
fn mirror(&self, name: &str) -> PathBuf {
|
||||
Mirror::path_for(&mirrors_dir(self.root.path()), name)
|
||||
}
|
||||
}
|
||||
|
||||
fn held(repo: &str, path: &str, chunk: u32, commit: &str) -> (String, String, u32, String) {
|
||||
(repo.into(), path.into(), chunk, commit.into())
|
||||
}
|
||||
|
||||
/// An empty collection plus a repo at a commit ends with that commit's chunks
|
||||
/// and the cursor on it.
|
||||
#[test]
|
||||
fn first_build_indexes_head() {
|
||||
let remote = Remote::new();
|
||||
let a = remote.commit(&[("src/a.rs", b"fn a() {}\nfn b() {}\n"), ("README.md", b"# Demo\n")], &[]);
|
||||
let mut rig = Rig::new();
|
||||
let reports = rig.converged(&mut FakeForge::of(&[("Demo", &remote)]));
|
||||
assert_eq!(
|
||||
rig.index.held(),
|
||||
[held("Demo", "README.md", 0, &a), held("Demo", "src/a.rs", 0, &a)],
|
||||
"one chunk per small file, at the head commit"
|
||||
);
|
||||
assert_eq!(rig.index.cursors().get("Demo"), Some(&a));
|
||||
assert_eq!((reports.len(), reports[0].files, reports[0].chunks, reports[0].old.clone()), (1, 2, 2, None));
|
||||
let point = rig.index.state.lock().unwrap().points.values().find(|point| point.payload.path == "src/a.rs").unwrap().clone();
|
||||
assert_eq!((point.payload.line_start, point.payload.line_end, point.payload.ext.as_str()), (1, 2, "rs"));
|
||||
assert_eq!(point.payload.text, "fn a() {}\nfn b() {}", "the payload text is the chunk alone");
|
||||
assert!(rig.embedder.texts().contains(&"Demo/src/a.rs\nfn a() {}\nfn b() {}".to_owned()), "the embedded text leads with repo and path");
|
||||
assert!(reports[0].to_string().starts_with("Demo -..") && reports[0].to_string().contains("files=2 chunks=2"));
|
||||
}
|
||||
|
||||
/// From A to B with one file edited and one deleted, only those two paths'
|
||||
/// points change: untouched files keep their point ids, and the cursor moves.
|
||||
#[test]
|
||||
fn refresh_replaces_only_changed_paths() {
|
||||
let remote = Remote::new();
|
||||
let a = remote.commit(&[("a.rs", b"one\n"), ("b.rs", b"two\n"), ("c.rs", b"three\n")], &[]);
|
||||
let mut rig = Rig::new();
|
||||
let mut forge = FakeForge::of(&[("Demo", &remote)]);
|
||||
rig.converged(&mut forge);
|
||||
let ids_before = rig.index.ids();
|
||||
let untouched: Vec<String> = {
|
||||
let state = rig.index.state.lock().unwrap();
|
||||
state.points.iter().filter(|(_, point)| point.payload.path == "c.rs").map(|(id, _)| id.clone()).collect()
|
||||
};
|
||||
let embedded_before = rig.embedder.texts().len();
|
||||
|
||||
let b = remote.commit(&[("a.rs", b"one edited\n")], &["b.rs"]);
|
||||
let reports = rig.converged(&mut forge);
|
||||
|
||||
assert_eq!(rig.index.held(), [held("Demo", "a.rs", 0, &b), held("Demo", "c.rs", 0, &a)]);
|
||||
assert!(untouched.iter().all(|id| rig.index.ids().contains(id)), "c.rs's point id survived");
|
||||
assert!(ids_before.len() == 3 && rig.index.ids().len() == 2);
|
||||
assert_eq!(rig.index.cursors().get("Demo"), Some(&b));
|
||||
assert_eq!(rig.embedder.texts().len() - embedded_before, 1, "only the edited file is embedded again");
|
||||
assert_eq!((reports[0].old.clone(), reports[0].files), (Some(a), 1));
|
||||
}
|
||||
|
||||
/// A head equal to the cursor is the whole refresh: nothing is embedded,
|
||||
/// written, deleted or advanced.
|
||||
#[test]
|
||||
fn unchanged_repo_makes_no_embed_call() {
|
||||
let remote = Remote::new();
|
||||
remote.commit(&[("a.rs", b"one\n")], &[]);
|
||||
let mut rig = Rig::new();
|
||||
let mut forge = FakeForge::of(&[("Demo", &remote)]);
|
||||
rig.converged(&mut forge);
|
||||
let (embeds, upserts, deletes, cursor_writes) = {
|
||||
let state = rig.index.state.lock().unwrap();
|
||||
(rig.embedder.state.lock().unwrap().embedded.len(), state.upsert_calls, state.deletes.len(), state.cursor_writes.len())
|
||||
};
|
||||
let reports = rig.converged(&mut forge);
|
||||
let state = rig.index.state.lock().unwrap();
|
||||
assert_eq!(rig.embedder.state.lock().unwrap().embedded.len(), embeds, "no embed call");
|
||||
assert_eq!((state.upsert_calls, state.deletes.len(), state.cursor_writes.len()), (upserts, deletes, cursor_writes));
|
||||
assert_eq!((reports[0].files, reports[0].chunks), (0, 0));
|
||||
}
|
||||
|
||||
/// An upsert that fails mid-repo leaves the cursor on the old commit, and the
|
||||
/// next run converges to the points a clean run of the same history makes.
|
||||
#[test]
|
||||
fn cursor_never_leads() {
|
||||
let remote = Remote::new();
|
||||
let many: Vec<(String, Vec<u8>)> = (0..40).map(|n| (format!("f{n}.rs"), format!("file {n}\n").into_bytes())).collect();
|
||||
let files: Vec<(&str, &[u8])> = many.iter().map(|(path, bytes)| (path.as_str(), bytes.as_slice())).collect();
|
||||
let a = remote.commit(&files, &[]);
|
||||
let mut rig = Rig::new();
|
||||
let mut forge = FakeForge::of(&[("Demo", &remote)]);
|
||||
rig.converged(&mut forge);
|
||||
|
||||
let edited: Vec<(String, Vec<u8>)> = (0..40).map(|n| (format!("f{n}.rs"), format!("file {n} edited\n").into_bytes())).collect();
|
||||
let files: Vec<(&str, &[u8])> = edited.iter().map(|(path, bytes)| (path.as_str(), bytes.as_slice())).collect();
|
||||
let b = remote.commit(&files, &[]);
|
||||
{
|
||||
let mut state = rig.index.state.lock().unwrap();
|
||||
state.fail_upserts_after = Some(state.upsert_calls + 1);
|
||||
}
|
||||
let (refreshed, _) = rig.run(&mut forge);
|
||||
assert_eq!(refreshed.unwrap().failed, ["Demo"], "one repo failed, the run finished");
|
||||
assert_eq!(rig.index.cursors().get("Demo"), Some(&a), "the cursor never leads the points");
|
||||
|
||||
rig.index.state.lock().unwrap().fail_upserts_after = None;
|
||||
rig.converged(&mut forge);
|
||||
assert_eq!(rig.index.cursors().get("Demo"), Some(&b));
|
||||
let mut clean = Rig::new();
|
||||
clean.converged(&mut forge);
|
||||
assert_eq!(rig.index.held(), clean.index.held());
|
||||
assert_eq!(rig.index.ids(), clean.index.ids());
|
||||
}
|
||||
|
||||
/// One repo failing is logged and the others go on.
|
||||
#[test]
|
||||
fn one_repos_failure_does_not_stop_the_others() {
|
||||
let good = Remote::new();
|
||||
let g = good.commit(&[("a.rs", b"one\n")], &[]);
|
||||
let nowhere = Remote::new();
|
||||
nowhere.commit(&[("a.rs", b"x\n")], &[]);
|
||||
let mut forge = FakeForge::of(&[("Broken", &nowhere), ("Good", &good)]);
|
||||
forge.repos[0].1 = PathBuf::from("/nowhere/at/all");
|
||||
let mut rig = Rig::new();
|
||||
let (refreshed, reports) = rig.run(&mut forge);
|
||||
assert_eq!(refreshed.unwrap().failed, ["Broken"]);
|
||||
assert_eq!(rig.index.held(), [held("Good", "a.rs", 0, &g)]);
|
||||
assert_eq!(reports.len(), 1);
|
||||
}
|
||||
|
||||
/// A history that no longer contains the cursor starts the repo over: the
|
||||
/// old commit's files that the new history lacks do not linger.
|
||||
#[test]
|
||||
fn a_rewritten_history_replaces_the_whole_repo() {
|
||||
let remote = Remote::new();
|
||||
remote.commit(&[("keep.rs", b"keep\n")], &[]);
|
||||
let old = remote.commit(&[("stale.rs", b"stale\n")], &[]);
|
||||
let mut rig = Rig::new();
|
||||
let mut forge = FakeForge::of(&[("Demo", &remote)]);
|
||||
rig.converged(&mut forge);
|
||||
git_in(&remote.path(), &["reset", "--quiet", "--hard", "HEAD~1"]);
|
||||
let new = remote.commit(&[("other.rs", b"other\n")], &[]);
|
||||
rig.converged(&mut forge);
|
||||
assert_ne!(old, new);
|
||||
assert_eq!(rig.index.cursors().get("Demo"), Some(&new));
|
||||
let paths: Vec<String> = rig.index.held().into_iter().map(|(_, path, _, _)| path).collect();
|
||||
assert_eq!(paths, ["keep.rs", "other.rs"], "stale.rs was in the old history only");
|
||||
}
|
||||
|
||||
/// A repo the forge stops listing leaves no points, no cursor key and no
|
||||
/// mirror; the one still listed is untouched.
|
||||
#[test]
|
||||
fn retired_repo_leaves_nothing() {
|
||||
let (keep, gone) = (Remote::new(), Remote::new());
|
||||
let k = keep.commit(&[("a.rs", b"a\n")], &[]);
|
||||
gone.commit(&[("b.rs", b"b\n")], &[]);
|
||||
let mut rig = Rig::new();
|
||||
rig.converged(&mut FakeForge::of(&[("Keep", &keep), ("Gone", &gone)]));
|
||||
assert!(rig.mirror("Gone").exists() && rig.index.cursors().contains_key("Gone"));
|
||||
rig.converged(&mut FakeForge::of(&[("Keep", &keep)]));
|
||||
assert_eq!(rig.index.held(), [held("Keep", "a.rs", 0, &k)]);
|
||||
assert_eq!(rig.index.cursors().keys().collect::<Vec<_>>(), ["Keep"]);
|
||||
assert!(!rig.mirror("Gone").exists() && rig.mirror("Keep").exists());
|
||||
}
|
||||
|
||||
/// A mirror no cursor knows (a first ingest that never finished) is retired
|
||||
/// with its points too.
|
||||
#[test]
|
||||
fn a_retired_repo_with_no_cursor_is_still_cleaned() {
|
||||
let keep = Remote::new();
|
||||
keep.commit(&[("a.rs", b"a\n")], &[]);
|
||||
let mut rig = Rig::new();
|
||||
let mut forge = FakeForge::of(&[("Keep", &keep)]);
|
||||
rig.converged(&mut forge);
|
||||
std::fs::create_dir_all(rig.mirror("Orphan")).unwrap();
|
||||
rig.converged(&mut forge);
|
||||
assert!(!rig.mirror("Orphan").exists());
|
||||
assert!(rig.index.state.lock().unwrap().deletes.iter().any(|(repo, paths)| repo == "Orphan" && paths.is_none()));
|
||||
}
|
||||
|
||||
/// Another digest or chunker version recreates the collection and refills it;
|
||||
/// an unlabelled collection is refused untouched.
|
||||
#[test]
|
||||
fn rebuild_on_model_or_chunker_change() {
|
||||
let remote = Remote::new();
|
||||
let a = remote.commit(&[("a.rs", b"one\n")], &[]);
|
||||
let mut rig = Rig::new();
|
||||
let mut forge = FakeForge::of(&[("Demo", &remote)]);
|
||||
rig.converged(&mut forge);
|
||||
assert_eq!(rig.index.state.lock().unwrap().recreated, 1, "created once, empty");
|
||||
rig.converged(&mut forge);
|
||||
assert_eq!(rig.index.state.lock().unwrap().recreated, 1, "same label, no rebuild");
|
||||
|
||||
rig.embedder = FakeEmbedder::new("d2");
|
||||
rig.converged(&mut forge);
|
||||
assert_eq!(rig.index.state.lock().unwrap().recreated, 2, "another digest rebuilds");
|
||||
assert_eq!(rig.index.held(), [held("Demo", "a.rs", 0, &a)], "and refills from the cursorless collection");
|
||||
|
||||
let mut old_chunker = match rig.index.state.lock().unwrap().described.clone() {
|
||||
Some(DescribedSources::Labelled { meta, cursors }) => (meta, cursors),
|
||||
other => panic!("{other:?}"),
|
||||
};
|
||||
old_chunker.0.chunker_version = CHUNKER_VERSION - 1;
|
||||
rig.index.plant(DescribedSources::Labelled { meta: old_chunker.0, cursors: old_chunker.1 });
|
||||
rig.converged(&mut forge);
|
||||
assert_eq!(rig.index.state.lock().unwrap().recreated, 3, "another chunker version rebuilds");
|
||||
|
||||
rig.index.plant(DescribedSources::Unlabelled);
|
||||
let before = rig.index.held();
|
||||
let (refreshed, _) = rig.run(&mut forge);
|
||||
assert!(format!("{:#}", refreshed.unwrap_err()).contains("no Huginn metadata"));
|
||||
assert_eq!((rig.index.state.lock().unwrap().recreated, rig.index.held()), (3, before), "untouched");
|
||||
}
|
||||
|
||||
/// A NUL blob, an LFS pointer, an oversize file, Unity YAML and a lock file
|
||||
/// yield no chunks; the plain file beside them does.
|
||||
#[test]
|
||||
fn skips_binary_lfs_oversize() {
|
||||
let remote = Remote::new();
|
||||
let huge = vec![b'x'; 300 * 1024];
|
||||
let big_but_ok = vec![b'y'; 100 * 1024];
|
||||
let c = remote.commit(
|
||||
&[
|
||||
("nul.txt", b"abc\0def\n"),
|
||||
("lfs.txt", b"version https://git-lfs.github.com/spec/v1\noid sha256:aa\nsize 9\n"),
|
||||
("huge.txt", &huge),
|
||||
("scene.prefab", b"%YAML 1.1\n"),
|
||||
("Cargo.lock", b"[[package]]\n"),
|
||||
("ok.txt", b"hello\n"),
|
||||
("big.txt", &big_but_ok),
|
||||
],
|
||||
&[],
|
||||
);
|
||||
let mut rig = Rig::new();
|
||||
rig.converged(&mut FakeForge::of(&[("Demo", &remote)]));
|
||||
let paths: Vec<String> = rig.index.held().into_iter().map(|(_, path, _, _)| path).collect();
|
||||
assert_eq!(paths, ["big.txt", "ok.txt"], "100 KiB is under the 256 KiB bound; one long line is one chunk");
|
||||
assert_eq!(rig.index.cursors().get("Demo"), Some(&c));
|
||||
}
|
||||
|
||||
fn query(text: &str, top_k: u32, repo: Option<&str>) -> SourceQuery {
|
||||
SourceQuery { text: Line(text.into()), top_k, repo: repo.map(|repo| Short(repo.into())) }
|
||||
}
|
||||
|
||||
/// A daemon whose embedder is not the one the collection was built under
|
||||
/// refuses and names both identities; a missing collection says it is not built.
|
||||
#[test]
|
||||
fn source_search_refuses_foreign_model() {
|
||||
let remote = Remote::new();
|
||||
remote.commit(&[("a.rs", b"one\n")], &[]);
|
||||
let mut rig = Rig::new();
|
||||
let within = Duration::from_secs(5);
|
||||
let error = search(&mut FakeEmbedder::new("d1"), &mut rig.index, &query("one", 3, None), within).unwrap_err();
|
||||
assert_eq!(format!("{error:#}"), "source index not built");
|
||||
|
||||
rig.converged(&mut FakeForge::of(&[("Demo", &remote)]));
|
||||
let mut foreign = FakeEmbedder::new("d2");
|
||||
let error = format!("{:#}", search(&mut foreign, &mut rig.index, &query("one", 3, None), within).unwrap_err());
|
||||
assert!(error.contains("built under model d1") && error.contains("embeds under model d2"), "{error}");
|
||||
assert!(foreign.state.lock().unwrap().embedded.is_empty(), "nothing was embedded for a foreign model");
|
||||
|
||||
foreign.state.lock().unwrap().identity.dimensions = 8;
|
||||
let error = format!("{:#}", search(&mut FakeEmbedder::new("d1"), &mut rig.index, &query("one", 3, None), within).map(|_| ()).err().map_or(String::new(), |e| format!("{e:#}")));
|
||||
assert_eq!(error, "", "the same model still answers");
|
||||
}
|
||||
|
||||
/// A search answers the chunks as the collection holds them, nearest first,
|
||||
/// within the repo asked for and no more than `top_k`; the query carries the
|
||||
/// instruction and not the path prefix.
|
||||
#[test]
|
||||
fn a_source_search_answers_held_chunks_within_a_repo_and_a_limit() {
|
||||
let (one, two) = (Remote::new(), Remote::new());
|
||||
let c1 = one.commit(&[("a.rs", b"alpha\n"), ("b.rs", b"beta\nbeta\n")], &[]);
|
||||
two.commit(&[("z.rs", b"zeta\n")], &[]);
|
||||
let mut rig = Rig::new();
|
||||
rig.converged(&mut FakeForge::of(&[("One", &one), ("Two", &two)]));
|
||||
let within = Duration::from_secs(5);
|
||||
|
||||
let all = search(&mut rig.embedder, &mut rig.index, &query("alpha", 10, None), within).unwrap();
|
||||
assert_eq!(all.hits.len(), 3);
|
||||
assert_eq!((all.indexed_under.name.as_str(), all.indexed_under.digest.as_str()), ("model", "d1"));
|
||||
assert!(all.hits.windows(2).all(|pair| pair[0].score >= pair[1].score), "nearest first");
|
||||
|
||||
let in_one = search(&mut rig.embedder, &mut rig.index, &query("alpha", 10, Some("One")), within).unwrap();
|
||||
assert_eq!(in_one.hits.len(), 2);
|
||||
assert!(in_one.hits.iter().all(|hit| hit.repo == "One" && hit.commit == c1));
|
||||
let b = in_one.hits.iter().find(|hit| hit.path == "b.rs").unwrap();
|
||||
assert_eq!((b.line_start, b.line_end, b.text.as_str()), (1, 2, "beta\nbeta"));
|
||||
|
||||
let limited = search(&mut rig.embedder, &mut rig.index, &query("alpha", 1, None), within).unwrap();
|
||||
assert_eq!(limited.hits.len(), 1);
|
||||
let asked = rig.embedder.texts().last().cloned().unwrap();
|
||||
assert_eq!(asked, format!("Instruct: {SOURCE_QUERY_INSTRUCTION}\nQuery: alpha"));
|
||||
}
|
||||
|
||||
/// Chunk ids are the content address: the same blob chunked the same way is
|
||||
/// the same point, and any part of the address changes it.
|
||||
#[test]
|
||||
fn a_chunk_id_is_its_content_address() {
|
||||
let base = chunk_id("R", "p", "blob", 0);
|
||||
assert_eq!(base, chunk_id("R", "p", "blob", 0));
|
||||
for other in [chunk_id("S", "p", "blob", 0), chunk_id("R", "q", "blob", 0), chunk_id("R", "p", "blob2", 0), chunk_id("R", "p", "blob", 1)] {
|
||||
assert_ne!(base, other);
|
||||
}
|
||||
assert_ne!(chunk_id("R", "pb", "lob", 0), chunk_id("R", "p", "blob", 0), "fields are separated, not concatenated");
|
||||
}
|
||||
Loading…
Reference in a new issue