Merge remote-tracking branch 'origin/main' into eureka-body/heartbeat-joint-lever
This commit is contained in:
Meta Vens
commit
bc59cb9025
30 changed files with 3725 additions and 50 deletions
|
|
@ -26,7 +26,9 @@ Upstream: `https://github.com/GameCult/Huginn.git`
|
|||
restart compares the mind with the collection and writes what is missing. An
|
||||
unreachable Qdrant or Ollama shows in `whoami` as the index's status and is
|
||||
retried; admission never waits on it. The daemon requires `--qdrant-url`,
|
||||
`--ollama-url` and `--embedding-model`. A `query` carrying `semantic` is
|
||||
`--ollama-url` and `--embedding-model`; `--ollama-url` repeats for each GPU
|
||||
host and `--reserve-ollama-url` for each CPU host, and the GPU hosts embed
|
||||
together while the CPU hosts serve only when none can. A `query` carrying `semantic` is
|
||||
answered through the same selection as any other: the index supplies
|
||||
candidate ids and scores, and the mind keeps only the documents it holds and
|
||||
the selection admits (ask for `in_force` to exclude resolved ones), ordered by
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ use cultnet_rs::{PROJECTION_DOCUMENT, Selection};
|
|||
use huginn_mind::eureka_pipeline::{Line, PipelineDocument, PipelineRef, Short, Slug, pipeline_key};
|
||||
use huginn_mind::{
|
||||
BATCH_MAX, Faculty, HuginnMindRequest, HuginnMindResponse, MindRefusal, PipelineAdmissionBatch, PipelineAdmissionOutcome,
|
||||
PipelinePageItems, PipelineProvenance, PipelineSelectionPage, SemanticQuery,
|
||||
PipelinePageItems, PipelineProvenance, PipelineSelectionPage, SemanticQuery, SourceQuery,
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
|
|
@ -83,6 +83,20 @@ enum Command {
|
|||
#[arg(long)]
|
||||
documents: bool,
|
||||
},
|
||||
/// Search the public forge repos' code and docs by meaning.
|
||||
Sources {
|
||||
/// What to find, in words.
|
||||
text: String,
|
||||
/// Look in this repo alone (its forge name).
|
||||
#[arg(long)]
|
||||
repo: Option<String>,
|
||||
/// How many chunks (default 10, at most 50).
|
||||
#[arg(long)]
|
||||
top_k: Option<u32>,
|
||||
/// Print the answer as JSON.
|
||||
#[arg(long)]
|
||||
json: bool,
|
||||
},
|
||||
#[command(flatten)]
|
||||
Recipe(Recipe),
|
||||
/// Admit a batch of documents, whole or not at all, and read it back.
|
||||
|
|
@ -168,6 +182,7 @@ fn run(settings: &Settings, command: Command) -> Result<u8, Trouble> {
|
|||
let semantic = near.map(|text| SemanticQuery { text: Line(text), top_k: top_k.unwrap_or(DEFAULT_TOP_K) });
|
||||
query(settings, selection, semantic, None, documents)
|
||||
}
|
||||
Command::Sources { text, repo, top_k, json } => sources(settings, text, repo, top_k, json),
|
||||
Command::Recipe(recipe) => {
|
||||
let ask = recipe.ask()?;
|
||||
query(settings, door::selection_value(&ask.selection)?, ask.semantic, ask.page.cursor, ask.page.documents)
|
||||
|
|
@ -249,6 +264,22 @@ fn query(
|
|||
}
|
||||
}
|
||||
|
||||
fn sources(settings: &Settings, text: String, repo: Option<String>, top_k: Option<u32>, json: bool) -> Result<u8, Trouble> {
|
||||
let query = SourceQuery { text: Line(text), top_k: top_k.unwrap_or(DEFAULT_TOP_K), repo: repo.map(Short) };
|
||||
match settings.ask(|instance| HuginnMindRequest::SearchSources { instance: instance.clone(), query })? {
|
||||
HuginnMindResponse::Sources(found) => {
|
||||
if json {
|
||||
say(serde_json::to_string_pretty(&found).map_err(|error| Trouble::Internal(error.to_string()))?);
|
||||
} else {
|
||||
say(render::source_hits(&found));
|
||||
}
|
||||
Ok(0)
|
||||
}
|
||||
HuginnMindResponse::Refused(refusal) => Ok(refused(&refusal)),
|
||||
other => Err(unexpected(&other)),
|
||||
}
|
||||
}
|
||||
|
||||
fn show(page: &PipelineSelectionPage) {
|
||||
match &page.items {
|
||||
PipelinePageItems::Headers(headers) => {
|
||||
|
|
@ -392,6 +423,7 @@ fn unexpected(answer: &HuginnMindResponse) -> Trouble {
|
|||
HuginnMindResponse::Query(_) => "query",
|
||||
HuginnMindResponse::Persona(_) => "a persona",
|
||||
HuginnMindResponse::PersonaPut(_) => "a persona put",
|
||||
HuginnMindResponse::Sources(_) => "source hits",
|
||||
HuginnMindResponse::Refused(_) => "a refusal",
|
||||
HuginnMindResponse::Deferred(_) => "a deferral",
|
||||
};
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
use std::fmt::Write;
|
||||
|
||||
use huginn_mind::{IndexStatus, PersonaStatus, PipelineDocumentSummary, PipelineDocumentView, PipelineStatus, PipelineStatusSummary};
|
||||
use huginn_mind::{SourceHits, IndexStatus, PersonaStatus, PipelineDocumentSummary, PipelineDocumentView, PipelineStatus, PipelineStatusSummary};
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
|
||||
|
|
@ -96,6 +96,32 @@ pub fn header(item: &PipelineDocumentSummary) -> String {
|
|||
line.chars().take(HEADER_MAX_CHARS).collect()
|
||||
}
|
||||
|
||||
/// Source hits as the CLI prints them: per hit a citation line,
|
||||
/// `repo:path:start-end @commit7 score`, then the chunk's text, and a blank line
|
||||
/// between hits. Nothing for no hits.
|
||||
pub fn source_hits(found: &SourceHits) -> String {
|
||||
found
|
||||
.hits
|
||||
.iter()
|
||||
.map(|hit| {
|
||||
format!(
|
||||
"{}:{}:{}-{} @{} {:.3}
|
||||
{}
|
||||
",
|
||||
hit.repo,
|
||||
hit.path,
|
||||
hit.line_start,
|
||||
hit.line_end,
|
||||
&hit.commit[..hit.commit.len().min(7)],
|
||||
hit.score,
|
||||
hit.text
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("
|
||||
")
|
||||
}
|
||||
|
||||
/// A fact as one token: a string, a number, a flag, a reference's id, or a
|
||||
/// list of those joined by commas. Anything deeper is not a header fact.
|
||||
fn scalar(value: &Value) -> Option<String> {
|
||||
|
|
@ -121,6 +147,35 @@ mod tests {
|
|||
|
||||
use super::*;
|
||||
|
||||
/// The citation line is what an agent copies: repo, path, the inclusive
|
||||
/// line range and the first seven of the commit.
|
||||
#[test]
|
||||
fn a_source_hit_prints_as_a_citation_then_its_text() {
|
||||
let hit = |path: &str, text: &str| huginn_mind::SourceHit {
|
||||
repo: "Huginn".into(),
|
||||
path: path.into(),
|
||||
commit: "0123456789abcdef".into(),
|
||||
line_start: 3,
|
||||
line_end: 9,
|
||||
score: 0.87654,
|
||||
text: text.into(),
|
||||
};
|
||||
let found = SourceHits {
|
||||
indexed_under: huginn_mind::SourceModel { name: "m".into(), digest: "d".into() },
|
||||
hits: vec![hit("a.rs", "one"), hit("b/c.rs", "two")],
|
||||
};
|
||||
assert_eq!(source_hits(&found), "Huginn:a.rs:3-9 @0123456 0.877
|
||||
one
|
||||
|
||||
Huginn:b/c.rs:3-9 @0123456 0.877
|
||||
two
|
||||
");
|
||||
let empty = SourceHits { indexed_under: found.indexed_under.clone(), hits: vec![] };
|
||||
assert_eq!(source_hits(&empty), "");
|
||||
let json = serde_json::to_string(&found).unwrap();
|
||||
assert_eq!(serde_json::from_str::<SourceHits>(&json).unwrap(), found, "--json round-trips");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_fact_is_one_token_or_nothing() {
|
||||
assert_eq!(scalar(&json!("")), None);
|
||||
|
|
|
|||
|
|
@ -19,7 +19,10 @@ use huginn_mind::{
|
|||
AdmissionFacts, HuginnMindRequest, HuginnMindResponse, IndexStatus, MindRefusal, MindStatus, PersonaStatus, PipelineAdmissionOutcome, PipelineDocumentView,
|
||||
PipelinePageItems, PipelineProvenance, PipelineSelectionPage, PipelineStatus,
|
||||
};
|
||||
use huginn_daemon::{Daemon, Hits, IndexSink, SearchTicket};
|
||||
use huginn_mind::MindStore;
|
||||
use serde_json::{Value, json};
|
||||
use std::time::Duration;
|
||||
use tempfile::TempDir;
|
||||
|
||||
mod common;
|
||||
|
|
@ -1047,3 +1050,160 @@ fn a_conflicting_batch_is_the_minds_no() {
|
|||
assert_eq!(ran.code, 1);
|
||||
assert_eq!(ran.out, format!("Conflict\nconflicts {CAMPAIGN}:campaign:self\n"));
|
||||
}
|
||||
|
||||
/// The chunks a source index holds, served from memory: each text embeds to one
|
||||
/// vector and each held chunk scores the same, so the search is the daemon's own
|
||||
/// path (`sources::search`, its identity check and its citation fields) with
|
||||
/// nothing about ranking in it.
|
||||
struct Desk {
|
||||
built_under: &'static str,
|
||||
embeds_under: &'static str,
|
||||
held: Vec<huginn_daemon::sources::SourcePayload>,
|
||||
tickets: u64,
|
||||
done: Vec<(SearchTicket, Result<huginn_mind::SourceHits, String>)>,
|
||||
}
|
||||
|
||||
/// An embedder that sends every text to one vector, under a digest.
|
||||
struct Flat(&'static str);
|
||||
|
||||
impl huginn_daemon::index::Embedder for Flat {
|
||||
fn model_identity(&mut self, _within: Duration) -> anyhow::Result<huginn_daemon::index::ModelIdentity> {
|
||||
Ok(huginn_daemon::index::ModelIdentity { name: "flat".into(), digest: self.0.into(), dimensions: 2 })
|
||||
}
|
||||
|
||||
fn embed(&mut self, texts: &[String], _within: Duration) -> anyhow::Result<Vec<Vec<f32>>> {
|
||||
Ok(texts.iter().map(|_| vec![1.0, 0.0]).collect())
|
||||
}
|
||||
}
|
||||
|
||||
impl huginn_daemon::sources::SourceIndex for Desk {
|
||||
fn describe_sources(&mut self, _within: Duration) -> anyhow::Result<huginn_daemon::sources::DescribedSources> {
|
||||
Ok(huginn_daemon::sources::DescribedSources::Labelled {
|
||||
meta: huginn_daemon::sources::SourcesMeta {
|
||||
managed_by: "huginn".into(),
|
||||
corpus: huginn_daemon::sources::CORPUS.into(),
|
||||
model: "flat".into(),
|
||||
model_digest: self.built_under.into(),
|
||||
dimensions: 2,
|
||||
chunker_version: huginn_daemon::sources::chunk::CHUNKER_VERSION,
|
||||
},
|
||||
cursors: Default::default(),
|
||||
})
|
||||
}
|
||||
|
||||
fn recreate_sources(&mut self, _meta: &huginn_daemon::sources::SourcesMeta) -> anyhow::Result<()> {
|
||||
anyhow::bail!("a daemon never writes the source index")
|
||||
}
|
||||
|
||||
fn set_cursors(&mut self, _cursors: &std::collections::BTreeMap<String, String>) -> anyhow::Result<()> {
|
||||
anyhow::bail!("a daemon never writes the source index")
|
||||
}
|
||||
|
||||
fn delete_by(&mut self, _repo: &str, _paths: Option<&[String]>) -> anyhow::Result<()> {
|
||||
anyhow::bail!("a daemon never writes the source index")
|
||||
}
|
||||
|
||||
fn upsert_sources(&mut self, _points: &[huginn_daemon::sources::SourcePoint]) -> anyhow::Result<()> {
|
||||
anyhow::bail!("a daemon never writes the source index")
|
||||
}
|
||||
|
||||
fn search_sources(
|
||||
&mut self,
|
||||
_vector: &[f32],
|
||||
repo: Option<&str>,
|
||||
limit: u32,
|
||||
_within: Duration,
|
||||
) -> anyhow::Result<Vec<huginn_daemon::sources::SourceRow>> {
|
||||
Ok(self
|
||||
.held
|
||||
.iter()
|
||||
.filter(|payload| repo.is_none_or(|repo| payload.repo == repo))
|
||||
.take(limit as usize)
|
||||
.map(|payload| huginn_daemon::sources::SourceRow { score: 0.9, payload: payload.clone() })
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: MindStore> IndexSink<S> for Desk {
|
||||
fn committed(&mut self, _mind: &huginn_mind::Mind<S>, _writes: &[PipelineRef]) -> anyhow::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
fn status(&self) -> IndexStatus {
|
||||
IndexStatus::Current
|
||||
}
|
||||
fn search(&mut self, _text: &str, _top_k: u32, _among: Vec<String>) -> Result<SearchTicket, String> {
|
||||
Err("this index holds sources only".into())
|
||||
}
|
||||
fn search_sources(&mut self, query: &huginn_mind::SourceQuery) -> Result<SearchTicket, String> {
|
||||
self.tickets += 1;
|
||||
let ticket = SearchTicket(self.tickets);
|
||||
let mut embedder = Flat(self.embeds_under);
|
||||
let found = huginn_daemon::sources::search(&mut embedder, self, query, Duration::from_secs(5)).map_err(|error| format!("{error:#}"));
|
||||
self.done.push((ticket, found));
|
||||
Ok(ticket)
|
||||
}
|
||||
fn sources_searched(&mut self) -> Vec<(SearchTicket, Result<huginn_mind::SourceHits, String>)> {
|
||||
std::mem::take(&mut self.done)
|
||||
}
|
||||
fn search_deadline(&self) -> Duration {
|
||||
Duration::from_secs(5)
|
||||
}
|
||||
fn searched(&mut self) -> Vec<(SearchTicket, Result<Hits, String>)> {
|
||||
Vec::new()
|
||||
}
|
||||
fn abandon(&mut self, _ticket: SearchTicket) {}
|
||||
}
|
||||
|
||||
/// A daemon over a temporary mind whose source index was built under
|
||||
/// `built_under` and whose embedder reports `embeds_under`.
|
||||
fn source_daemon(built_under: &'static str, embeds_under: &'static str) -> Server {
|
||||
let chunk = |repo: &str, path: &str, line_start: u32, line_end: u32, text: &str| huginn_daemon::sources::SourcePayload {
|
||||
repo: repo.into(),
|
||||
path: path.into(),
|
||||
commit: "abcdef1234567890".into(),
|
||||
blob: "b".into(),
|
||||
line_start,
|
||||
line_end,
|
||||
chunk_index: 0,
|
||||
ext: "rs".into(),
|
||||
text: text.into(),
|
||||
};
|
||||
let desk = Desk {
|
||||
built_under,
|
||||
embeds_under,
|
||||
held: vec![
|
||||
chunk("Demo", "src/lib.rs", 3, 9, "fn cursor() {}"),
|
||||
chunk("Demo", "README.md", 1, 2, "# Demo"),
|
||||
chunk("Other", "x.rs", 1, 1, "fn other() {}"),
|
||||
],
|
||||
tickets: 0,
|
||||
done: Vec::new(),
|
||||
};
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let mut daemon = Daemon::new(huginn_mind::Mind::open(root.path(), &slug(INSTANCE)).unwrap(), desk);
|
||||
let admitted = answered(daemon.handle(HuginnMindRequest::Admit(batch(INSTANCE, vec![identity(INSTANCE)])), chrono::Utc::now()));
|
||||
assert!(matches!(admitted, HuginnMindResponse::Admit(PipelineAdmissionOutcome::Committed { .. })), "{admitted:?}");
|
||||
serve(root, daemon)
|
||||
}
|
||||
|
||||
/// `huginn sources` end to end: the binary asks a daemon over CultNet, the
|
||||
/// daemon answers from its source index through the worker path, and the
|
||||
/// citation is printed for the chunk the filter and limit chose.
|
||||
#[test]
|
||||
fn sources_prints_a_citation_for_a_search_answered_by_a_daemon() {
|
||||
let server = source_daemon("d1", "d1");
|
||||
let found = at(server.addr, &["sources", "where is the cursor kept", "--repo", "Demo", "--top-k", "1"], None);
|
||||
assert_eq!(found.code, 0, "{}", found.err);
|
||||
assert!(found.out.starts_with("Demo:src/lib.rs:3-9 @abcdef1 0.900\nfn cursor() {}"), "{}", found.out);
|
||||
assert!(!found.out.contains("README") && !found.out.contains("Other"), "{}", found.out);
|
||||
}
|
||||
|
||||
/// A daemon whose embedder is not the model the index was built under refuses
|
||||
/// the search naming both, and the CLI exits 1 with that word.
|
||||
#[test]
|
||||
fn sources_refuses_naming_both_models_when_the_daemon_embeds_under_another() {
|
||||
let server = source_daemon("d1", "d2");
|
||||
let refused = at(server.addr, &["sources", "anything"], None);
|
||||
assert_eq!(refused.code, 1, "{}", refused.err);
|
||||
assert!(refused.out.contains("d1") && refused.out.contains("d2"), "{}", refused.out);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -92,6 +92,12 @@ impl<S: MindStore> IndexSink<S> for Inert {
|
|||
fn search(&mut self, _text: &str, _top_k: u32, _among: Vec<String>) -> Result<SearchTicket, String> {
|
||||
Err("this index cannot search".into())
|
||||
}
|
||||
fn search_sources(&mut self, _query: &huginn_mind::SourceQuery) -> Result<SearchTicket, String> {
|
||||
Err("this index cannot search".into())
|
||||
}
|
||||
fn sources_searched(&mut self) -> Vec<(SearchTicket, Result<huginn_mind::SourceHits, String>)> {
|
||||
Vec::new()
|
||||
}
|
||||
fn search_deadline(&self) -> Duration {
|
||||
Duration::ZERO
|
||||
}
|
||||
|
|
@ -142,7 +148,10 @@ impl Drop for Server {
|
|||
}
|
||||
}
|
||||
|
||||
pub fn serve(root: TempDir, mut daemon: TestDaemon) -> Server {
|
||||
pub fn serve<I: IndexSink<OwnedRedbMessagePackBackingStore> + Send + 'static>(
|
||||
root: TempDir,
|
||||
mut daemon: Daemon<OwnedRedbMessagePackBackingStore, I>,
|
||||
) -> Server {
|
||||
let mut hub = bind("127.0.0.1:0".parse().unwrap(), &daemon.runtime_id()).unwrap();
|
||||
let addr = hub.local_addr().unwrap();
|
||||
let stopping = Arc::new(AtomicBool::new(false));
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
99
crates/huginn-daemon/src/bin/huginn-sources.rs
Normal file
99
crates/huginn-daemon/src/bin/huginn-sources.rs
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
//! 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::pool;
|
||||
use huginn_daemon::index::qdrant::QdrantIndex;
|
||||
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,
|
||||
/// The GPU hosts: all take batches at once.
|
||||
ollama_urls: Vec<String>,
|
||||
/// The CPU host: takes batches only while no GPU host can.
|
||||
reserve_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();
|
||||
let mut reserve_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" | "reserve-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" || name == "reserve-ollama-url" {
|
||||
ensure!(value.starts_with("http://"), "--{name} must be a plain http:// address, got {value:?}");
|
||||
if name == "ollama-url" { ollama_urls.push(value) } else { reserve_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 for each GPU host");
|
||||
ollama_urls
|
||||
},
|
||||
reserve_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 = pool(&options.ollama_urls, &options.reserve_urls, &options.embedding_model);
|
||||
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.is_clean() => ExitCode::SUCCESS,
|
||||
Ok(refreshed) => {
|
||||
if !refreshed.failed.is_empty() {
|
||||
eprintln!("huginn-sources: {} repos failed: {}", refreshed.failed.len(), refreshed.failed.join(", "));
|
||||
}
|
||||
if let Some(why) = &refreshed.retirement_withheld {
|
||||
eprintln!("huginn-sources: no repo was retired: {why}");
|
||||
}
|
||||
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,11 +58,27 @@ pub type Hits = Vec<(PipelineRef, f32)>;
|
|||
/// `Daemon::finish` with the index's answer.
|
||||
pub struct Search {
|
||||
pub ticket: SearchTicket,
|
||||
/// What the mind owes the asker of a mind search; `None` for a source
|
||||
/// search, whose answer is the source index's hits as it holds them.
|
||||
mind: Option<MindSearch>,
|
||||
}
|
||||
|
||||
/// A mind search in flight: the asked selection and query, and the members the
|
||||
/// mind chose before the index was asked.
|
||||
struct MindSearch {
|
||||
selection: Selection,
|
||||
semantic: SemanticQuery,
|
||||
members: Members,
|
||||
}
|
||||
|
||||
/// 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
|
||||
/// finish.
|
||||
pub enum Handled {
|
||||
|
|
@ -170,6 +191,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, mind: None }),
|
||||
Err(detail) => Handled::Answered(HuginnMindResponse::Refused(MindRefusal::Unavailable { detail })),
|
||||
};
|
||||
}
|
||||
};
|
||||
Handled::Answered(answer)
|
||||
}
|
||||
|
|
@ -193,7 +223,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, mind: Some(MindSearch { selection, semantic, members }) }),
|
||||
Err(detail) => refused(MindRefusal::Unavailable { detail }),
|
||||
}
|
||||
}
|
||||
|
|
@ -204,8 +234,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 +249,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.mind, found) {
|
||||
(_, Found::Failed(detail)) => return unavailable(detail),
|
||||
(None, Found::Sources(hits)) => return HuginnMindResponse::Sources(hits),
|
||||
(Some(asked), Found::Mind(hits)) => (asked.selection, asked.semantic, asked.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 +532,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
|
||||
}
|
||||
|
|
@ -508,6 +551,83 @@ pub(crate) mod tests {
|
|||
fn abandon(&mut self, _ticket: SearchTicket) {}
|
||||
}
|
||||
|
||||
/// An index that answers source searches from a script and counts how often
|
||||
/// it was asked.
|
||||
struct SourcesIndex {
|
||||
asked: Arc<Mutex<u32>>,
|
||||
hits: SourceHits,
|
||||
}
|
||||
|
||||
impl<S: MindStore> IndexSink<S> for SourcesIndex {
|
||||
fn committed(&mut self, _mind: &Mind<S>, _writes: &[PipelineRef]) -> anyhow::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn status(&self) -> IndexStatus {
|
||||
IndexStatus::Current
|
||||
}
|
||||
|
||||
fn search(&mut self, _text: &str, _top_k: u32, _among: Vec<String>) -> Result<SearchTicket, String> {
|
||||
Err("this index cannot search the mind".into())
|
||||
}
|
||||
|
||||
fn search_sources(&mut self, _query: &SourceQuery) -> Result<SearchTicket, String> {
|
||||
*self.asked.lock().unwrap() += 1;
|
||||
Ok(SearchTicket(7))
|
||||
}
|
||||
|
||||
fn sources_searched(&mut self) -> Vec<(SearchTicket, Result<SourceHits, String>)> {
|
||||
vec![(SearchTicket(7), Ok(self.hits.clone()))]
|
||||
}
|
||||
|
||||
fn search_deadline(&self) -> std::time::Duration {
|
||||
std::time::Duration::ZERO
|
||||
}
|
||||
|
||||
fn searched(&mut self) -> Vec<(SearchTicket, Result<Hits, String>)> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn abandon(&mut self, _ticket: SearchTicket) {}
|
||||
}
|
||||
|
||||
/// A source search is checked at the door before the index is asked, is
|
||||
/// answered with the hits the index found and nothing joined from the mind,
|
||||
/// turns a failure into `Unavailable`, and is refused for another mind.
|
||||
#[test]
|
||||
fn a_source_search_reaches_the_index_only_when_well_formed_and_answers_its_hits() {
|
||||
use huginn_mind::{SourceHit, SourceModel};
|
||||
let hits = SourceHits {
|
||||
indexed_under: SourceModel { name: "m".into(), digest: "d".into() },
|
||||
hits: vec![SourceHit { repo: "R".into(), path: "p.rs".into(), commit: "c".into(), line_start: 1, line_end: 2, score: 0.5, text: "t".into() }],
|
||||
};
|
||||
let asked = Arc::new(Mutex::new(0));
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let mind = Mind::open(root.path(), &slug(INSTANCE)).unwrap();
|
||||
let mut daemon = Daemon::new(mind, SourcesIndex { asked: Arc::clone(&asked), hits: hits.clone() });
|
||||
let request = |instance: &str, top_k: u32| HuginnMindRequest::SearchSources {
|
||||
instance: slug(instance),
|
||||
query: SourceQuery { text: huginn_mind::eureka_pipeline::Line("find".into()), top_k, repo: None },
|
||||
};
|
||||
|
||||
let Handled::Searching(search) = daemon.handle(request(INSTANCE, 5), now()) else { panic!("a well-formed source search is a search") };
|
||||
assert_eq!((search.ticket, *asked.lock().unwrap()), (SearchTicket(7), 1));
|
||||
let mut finished = daemon.searched();
|
||||
let (ticket, found) = finished.remove(0);
|
||||
assert_eq!(ticket, SearchTicket(7));
|
||||
assert_eq!(daemon.finish(search, found), HuginnMindResponse::Sources(hits));
|
||||
|
||||
let Handled::Searching(search) = daemon.handle(request(INSTANCE, 5), now()) else { panic!() };
|
||||
let failed = daemon.finish(search, Found::Failed("the index fell over".into()));
|
||||
assert!(matches!(failed, HuginnMindResponse::Refused(MindRefusal::Unavailable { ref detail }) if detail.contains("fell over")), "{failed:?}");
|
||||
|
||||
for refused in [request(INSTANCE, 0), request(INSTANCE, 51), request(OTHER, 5)] {
|
||||
let answer = daemon.handle(refused, now()).answered();
|
||||
assert!(matches!(answer, HuginnMindResponse::Refused(_)), "{answer:?}");
|
||||
}
|
||||
assert_eq!(*asked.lock().unwrap(), 2, "a refused query never reached the index");
|
||||
}
|
||||
|
||||
/// A daemon over the instance's mind under `state_root`, with no index.
|
||||
pub(crate) fn open_unindexed(
|
||||
state_root: &Path,
|
||||
|
|
@ -539,6 +659,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 +693,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,143 @@ 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>>)>,
|
||||
/// A delete naming this repo refuses.
|
||||
pub(crate) fail_deletes_of: Option<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();
|
||||
if state.fail_deletes_of.as_deref() == Some(repo) {
|
||||
bail!("the vector store refused the delete");
|
||||
}
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -89,6 +89,7 @@
|
|||
//! model, the collection's label) forgets it.
|
||||
|
||||
pub mod ollama;
|
||||
pub mod pool;
|
||||
pub mod qdrant;
|
||||
|
||||
#[cfg(test)]
|
||||
|
|
@ -105,6 +106,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.
|
||||
|
|
@ -152,6 +155,13 @@ pub struct CollectionMeta {
|
|||
pub index_text_version: u32,
|
||||
}
|
||||
|
||||
impl CollectionMeta {
|
||||
/// The model the collection was built with.
|
||||
pub fn model_identity(&self) -> ModelIdentity {
|
||||
ModelIdentity { name: self.model.clone(), digest: self.model_digest.clone(), dimensions: self.dimensions }
|
||||
}
|
||||
}
|
||||
|
||||
/// What the vector store holds under a collection's name.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum Described {
|
||||
|
|
@ -189,6 +199,10 @@ pub struct Hit {
|
|||
}
|
||||
|
||||
pub trait Embedder {
|
||||
/// The model the collection about to be served was built with, or `None`
|
||||
/// when there is no collection yet. Called before `model_identity`. One
|
||||
/// endpoint has nothing to do with it; a pool holds its endpoints to it.
|
||||
fn built_under(&mut self, _built: Option<&ModelIdentity>) {}
|
||||
/// The model's identity, or an error once `within` has passed.
|
||||
fn model_identity(&mut self, within: Duration) -> Result<ModelIdentity>;
|
||||
/// One vector per text, in order, or an error once `within` has passed.
|
||||
|
|
@ -262,17 +276,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)
|
||||
|
|
@ -414,10 +428,12 @@ impl<E: Embedder, V: VectorIndex> Projector<E, V> {
|
|||
let Some(mind) = self.mind.clone() else {
|
||||
return Err(StepError::Unavailable(anyhow::anyhow!("the mind has no identity yet")));
|
||||
};
|
||||
let described = self.index.describe(&self.collection, INDEX_CALL_TIMEOUT).context("describing the collection")?;
|
||||
self.embedder.built_under(if let Described::Labelled(found) = &described { Some(found.model_identity()) } else { None }.as_ref());
|
||||
let identity =
|
||||
self.embedder.model_identity(BULK_EMBED_TIMEOUT).context("reading the embedding model's identity")?;
|
||||
let wanted = self.collection_meta(&mind, &identity);
|
||||
match self.index.describe(&self.collection, INDEX_CALL_TIMEOUT).context("describing the collection")? {
|
||||
match described {
|
||||
Described::Absent => {
|
||||
// A collection lost while the worker ran was an index; one that
|
||||
// was never made is a first build.
|
||||
|
|
@ -669,6 +685,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 +705,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 +737,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());
|
||||
|
|
@ -787,6 +813,21 @@ fn serve_search<E: Embedder, V: VectorIndex>(projector: &mut Projector<E, V>, sh
|
|||
lock(shared).found.push((ticket, found));
|
||||
}
|
||||
|
||||
/// 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));
|
||||
}
|
||||
|
||||
/// The search at the front of the queue, if it may be served now. Never while
|
||||
/// a document sits in the inbox, and never while anything ahead of it is
|
||||
/// owed (documents to write, a reconciliation to make) unless the index
|
||||
|
|
@ -852,7 +893,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 +919,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 +936,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 +980,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 +993,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 +1071,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 +1085,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 +1117,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,12 +1137,23 @@ pub(crate) enum Method {
|
|||
Put,
|
||||
Post,
|
||||
Delete,
|
||||
Patch,
|
||||
}
|
||||
|
||||
impl Http {
|
||||
pub(crate) fn new(base_url: &str, timeout: Duration) -> Self {
|
||||
let agent: ureq::Agent =
|
||||
ureq::Agent::config_builder().timeout_global(Some(timeout)).http_status_as_error(false).build().into();
|
||||
Self::connecting_within(base_url, timeout, None)
|
||||
}
|
||||
|
||||
/// As `new`, with a bound on the connect alone, so a host that never
|
||||
/// answers the connection is not waited on for the whole `timeout`.
|
||||
pub(crate) fn connecting_within(base_url: &str, timeout: Duration, connect: Option<Duration>) -> Self {
|
||||
let agent: ureq::Agent = ureq::Agent::config_builder()
|
||||
.timeout_global(Some(timeout))
|
||||
.timeout_connect(connect)
|
||||
.http_status_as_error(false)
|
||||
.build()
|
||||
.into();
|
||||
Self { agent, base: base_url.trim_end_matches('/').to_owned() }
|
||||
}
|
||||
|
||||
|
|
@ -1084,6 +1171,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();
|
||||
|
|
|
|||
|
|
@ -14,15 +14,43 @@ use super::{Embedder, Http, Method, ModelIdentity};
|
|||
pub const QUERY_INSTRUCTION: &str =
|
||||
"Given a question about a project's pipeline state, retrieve the documents that answer it";
|
||||
|
||||
/// What an embed call may take, sized to the batch and never beyond the
|
||||
/// caller's budget: a host that froze after accepting the batch is given up on
|
||||
/// after `floor + per_text * texts`, not after the caller's whole `within`.
|
||||
#[derive(Clone, Copy)]
|
||||
pub struct EmbedBounds {
|
||||
/// The bound on connecting alone.
|
||||
pub connect: Duration,
|
||||
/// What any batch is allowed, a model load included.
|
||||
pub floor: Duration,
|
||||
/// What each text adds.
|
||||
pub per_text: Duration,
|
||||
}
|
||||
|
||||
impl EmbedBounds {
|
||||
pub const DEFAULT: Self =
|
||||
Self { connect: Duration::from_secs(3), floor: Duration::from_secs(20), per_text: Duration::from_millis(500) };
|
||||
|
||||
fn for_batch(&self, texts: usize, within: Duration) -> Duration {
|
||||
(self.floor + self.per_text.saturating_mul(u32::try_from(texts).unwrap_or(u32::MAX))).min(within)
|
||||
}
|
||||
}
|
||||
|
||||
/// One model on one Ollama. Every call carries its own bound.
|
||||
pub struct OllamaEmbedder {
|
||||
base_url: String,
|
||||
model: String,
|
||||
bounds: EmbedBounds,
|
||||
}
|
||||
|
||||
impl OllamaEmbedder {
|
||||
pub fn new(base_url: &str, model: &str) -> Self {
|
||||
Self { base_url: base_url.to_owned(), model: model.to_owned() }
|
||||
Self { base_url: base_url.to_owned(), model: model.to_owned(), bounds: EmbedBounds::DEFAULT }
|
||||
}
|
||||
|
||||
pub fn bounded_by(mut self, bounds: EmbedBounds) -> Self {
|
||||
self.bounds = bounds;
|
||||
self
|
||||
}
|
||||
|
||||
fn ask(&self, http: &Http, method: Method, path: &str, body: Option<&Value>) -> Result<Value> {
|
||||
|
|
@ -57,7 +85,8 @@ impl Embedder for OllamaEmbedder {
|
|||
}
|
||||
|
||||
fn embed(&mut self, texts: &[String], within: Duration) -> Result<Vec<Vec<f32>>> {
|
||||
let bounded = Http::new(&self.base_url, within);
|
||||
let allowed = self.bounds.for_batch(texts.len(), within);
|
||||
let bounded = Http::connecting_within(&self.base_url, allowed, Some(self.bounds.connect.min(allowed)));
|
||||
let answer = self.ask(&bounded, Method::Post, "/api/embed", Some(&json!({ "model": self.model, "input": texts })))?;
|
||||
let embeddings = answer["embeddings"].as_array().context("Ollama's answer has no embeddings")?;
|
||||
embeddings
|
||||
|
|
@ -73,3 +102,23 @@ impl Embedder for OllamaEmbedder {
|
|||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// One pool over Ollama endpoints serving `model`: the `primaries` (GPU hosts)
|
||||
/// share the work and the `reserve` (CPU hosts) serve only when none of them
|
||||
/// can. The one embedder client of every Huginn process that embeds.
|
||||
pub fn pool(primaries: &[String], reserve: &[String], model: &str) -> super::pool::EmbedderPool<OllamaEmbedder> {
|
||||
let endpoint = |url: &String| (url.clone(), OllamaEmbedder::new(url, model));
|
||||
super::pool::EmbedderPool::new(primaries.iter().map(endpoint).collect(), reserve.iter().map(endpoint).collect())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn a_batch_is_given_its_floor_and_what_each_text_adds_but_never_more_than_the_budget() {
|
||||
let bounds = EmbedBounds { connect: Duration::from_secs(1), floor: Duration::from_secs(10), per_text: Duration::from_secs(2) };
|
||||
assert_eq!(bounds.for_batch(3, Duration::from_secs(100)), Duration::from_secs(16));
|
||||
assert_eq!(bounds.for_batch(3, Duration::from_secs(12)), Duration::from_secs(12));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
714
crates/huginn-daemon/src/index/pool.rs
Normal file
714
crates/huginn-daemon/src/index/pool.rs
Normal file
|
|
@ -0,0 +1,714 @@
|
|||
//! A pool of embedder endpoints serving one model, as one `Embedder`.
|
||||
//!
|
||||
//! The endpoints are two tiers. The primaries (the GPU hosts) all take work at
|
||||
//! once: a batch is cut into contiguous pieces, one per healthy primary, and
|
||||
//! the pieces are embedded concurrently and joined back in order. The reserve
|
||||
//! (the CPU host) takes work only while no primary is healthy. An endpoint is
|
||||
//! admitted only after its model identity (name, digest, dimensions) equals the
|
||||
//! pool's; one that is unreachable, failing or serving another model drops out,
|
||||
//! is named in the log, is never written with, and is read again once
|
||||
//! `retry_after` has passed. A piece that fails on one endpoint is embedded
|
||||
//! again on another, so a batch fails only when no endpoint is left to take it.
|
||||
//!
|
||||
//! The model of the pool is the one the collection was built with: the caller
|
||||
//! hands it over with `built_under` before it reads the identity, and an
|
||||
//! endpoint serving anything else never takes work, even when it is the only
|
||||
//! one up, so a run cannot rebuild an index around a foreign model. Only when
|
||||
//! there is no collection yet does the pool adopt a model: the first one
|
||||
//! reported in list order, primaries first, once every endpoint has answered.
|
||||
//!
|
||||
//! An endpoint's identity is read off the batch's clock. Each read runs on its
|
||||
//! own thread under `IDENTITY_WITHIN`, and a batch waits for the reads only
|
||||
//! while no healthy primary exists to start on: a primary being read keeps the
|
||||
//! reserve waiting, so a returning GPU is used and not passed over for the CPU,
|
||||
//! but an endpoint that drops packets never delays work a healthy one can do.
|
||||
|
||||
use std::panic::AssertUnwindSafe;
|
||||
use std::sync::{Arc, Condvar, Mutex, MutexGuard, PoisonError};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use anyhow::{Result, bail};
|
||||
|
||||
use super::{Embedder, ModelIdentity};
|
||||
|
||||
/// How long an endpoint that dropped out is left alone by default.
|
||||
pub const RETRY_AFTER: Duration = Duration::from_secs(60);
|
||||
|
||||
/// The bound on one identity read, connect and answer together. Reading a
|
||||
/// model's identity loads nothing, so an endpoint that needs longer is dark.
|
||||
pub const IDENTITY_WITHIN: Duration = Duration::from_secs(5);
|
||||
|
||||
/// One piece's outcome: the member that took it, where it starts, how long it is,
|
||||
/// and what came back.
|
||||
type Answer = (usize, usize, usize, Result<Vec<Vec<f32>>>);
|
||||
|
||||
struct Health {
|
||||
name: String,
|
||||
reserve: bool,
|
||||
/// `None` while the endpoint is in the pool; otherwise when it may be read
|
||||
/// again, and why it left.
|
||||
out: Option<(Instant, String)>,
|
||||
probing: bool,
|
||||
/// What the last read found, until the roster has judged it.
|
||||
seen: Option<Result<ModelIdentity, String>>,
|
||||
}
|
||||
|
||||
struct Roster {
|
||||
pinned: Option<ModelIdentity>,
|
||||
members: Vec<Health>,
|
||||
}
|
||||
|
||||
struct Shared {
|
||||
roster: Mutex<Roster>,
|
||||
changed: Condvar,
|
||||
}
|
||||
|
||||
fn locked<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
|
||||
mutex.lock().unwrap_or_else(PoisonError::into_inner)
|
||||
}
|
||||
|
||||
impl Roster {
|
||||
fn healthy(&self, reserve: bool) -> Vec<usize> {
|
||||
(0..self.members.len()).filter(|at| self.members[*at].reserve == reserve && self.members[*at].out.is_none()).collect()
|
||||
}
|
||||
|
||||
/// The members to work with now: the healthy primaries, or if none, the
|
||||
/// healthy reserve.
|
||||
fn active(&self) -> Vec<usize> {
|
||||
let primaries = self.healthy(false);
|
||||
if primaries.is_empty() { self.healthy(true) } else { primaries }
|
||||
}
|
||||
|
||||
fn reading(&self, reserve: bool) -> bool {
|
||||
self.members.iter().any(|member| member.reserve == reserve && member.probing)
|
||||
}
|
||||
|
||||
/// Whether a batch may start: a healthy primary is there to start on; else
|
||||
/// no primary is still being read, and the reserve is healthy or nothing is
|
||||
/// left being read. Without a model to hold endpoints to, it waits for every
|
||||
/// read, so the model adopted does not depend on who answered first.
|
||||
fn ready(&self) -> bool {
|
||||
if self.pinned.is_none() && self.members.iter().any(|member| member.probing) {
|
||||
return false;
|
||||
}
|
||||
if !self.healthy(false).is_empty() {
|
||||
return true;
|
||||
}
|
||||
if self.reading(false) {
|
||||
return false;
|
||||
}
|
||||
!self.healthy(true).is_empty() || !self.reading(true)
|
||||
}
|
||||
|
||||
/// Adopts a model when there is none and every read is in, then lets in the
|
||||
/// endpoints that serve it and puts out the rest.
|
||||
fn judge(&mut self, retry_after: Duration) {
|
||||
if self.pinned.is_none() && !self.members.iter().any(|member| member.probing) {
|
||||
self.pinned =
|
||||
self.members.iter().find_map(|member| member.seen.as_ref().and_then(|seen| seen.as_ref().ok())).cloned();
|
||||
}
|
||||
let pinned = self.pinned.clone();
|
||||
for member in &mut self.members {
|
||||
let why = match (&pinned, &member.seen) {
|
||||
(_, None) | (None, Some(Ok(_))) => continue,
|
||||
(_, Some(Err(why))) => why.clone(),
|
||||
(Some(pinned), Some(Ok(identity))) if pinned == identity => {
|
||||
member.out = None;
|
||||
member.seen = None;
|
||||
continue;
|
||||
}
|
||||
(Some(pinned), Some(Ok(identity))) => format!("serves {identity:?}, not {pinned:?}"),
|
||||
};
|
||||
member.seen = None;
|
||||
eprintln!("huginn embedder pool: {} is out of the pool: {why}", member.name);
|
||||
member.out = Some((Instant::now() + retry_after, why));
|
||||
}
|
||||
}
|
||||
|
||||
fn unavailable(&self) -> anyhow::Error {
|
||||
let named: Vec<String> = self
|
||||
.members
|
||||
.iter()
|
||||
.map(|member| format!("{}: {}", member.name, member.out.as_ref().map_or("failed this batch", |(_, why)| why.as_str())))
|
||||
.collect();
|
||||
anyhow::anyhow!("no embedder endpoint can take the batch: {}", named.join("; "))
|
||||
}
|
||||
}
|
||||
|
||||
/// Reads one endpoint's identity and tells the roster. Runs on its own thread.
|
||||
fn probe<E: Embedder>(shared: &Shared, embedder: &Mutex<E>, at: usize, retry_after: Duration) {
|
||||
let read = std::panic::catch_unwind(AssertUnwindSafe(|| locked(embedder).model_identity(IDENTITY_WITHIN)))
|
||||
.unwrap_or_else(|_| Err(anyhow::anyhow!("the embedder panicked")));
|
||||
let mut roster = locked(&shared.roster);
|
||||
roster.members[at].probing = false;
|
||||
roster.members[at].seen = Some(read.map_err(|error| format!("{error:#}")));
|
||||
roster.judge(retry_after);
|
||||
drop(roster);
|
||||
shared.changed.notify_all();
|
||||
}
|
||||
|
||||
pub struct EmbedderPool<E: Embedder + Send + 'static> {
|
||||
embedders: Vec<Arc<Mutex<E>>>,
|
||||
shared: Arc<Shared>,
|
||||
retry_after: Duration,
|
||||
}
|
||||
|
||||
impl<E: Embedder + Send + 'static> EmbedderPool<E> {
|
||||
/// `primaries` share the work; `reserve` serves only when none of them can.
|
||||
/// Every endpoint starts out, to be admitted by its first identity read.
|
||||
pub fn new(primaries: Vec<(String, E)>, reserve: Vec<(String, E)>) -> Self {
|
||||
let mut embedders = Vec::new();
|
||||
let mut members = Vec::new();
|
||||
let tiers = primaries.into_iter().map(|primary| (false, primary)).chain(reserve.into_iter().map(|spare| (true, spare)));
|
||||
for (is_reserve, (name, embedder)) in tiers {
|
||||
embedders.push(Arc::new(Mutex::new(embedder)));
|
||||
members.push(Health { name, reserve: is_reserve, out: Some((Instant::now(), "not yet read".into())), probing: false, seen: None });
|
||||
}
|
||||
let shared = Arc::new(Shared { roster: Mutex::new(Roster { pinned: None, members }), changed: Condvar::new() });
|
||||
Self { embedders, shared, retry_after: RETRY_AFTER }
|
||||
}
|
||||
|
||||
pub fn retrying_after(mut self, retry_after: Duration) -> Self {
|
||||
self.retry_after = retry_after;
|
||||
self
|
||||
}
|
||||
|
||||
/// Starts a read of every endpoint that is out and due, each on its own
|
||||
/// thread, and waits for them only until a batch can start (`Roster::ready`)
|
||||
/// or `within` has passed.
|
||||
fn admit(&self, within: Duration) {
|
||||
let started = Instant::now();
|
||||
let mut roster = locked(&self.shared.roster);
|
||||
for at in 0..self.embedders.len() {
|
||||
let member = &mut roster.members[at];
|
||||
if member.probing || member.out.as_ref().is_none_or(|(due, _)| *due > started) {
|
||||
continue;
|
||||
}
|
||||
member.probing = true;
|
||||
let (shared, embedder, retry_after) = (Arc::clone(&self.shared), Arc::clone(&self.embedders[at]), self.retry_after);
|
||||
std::thread::spawn(move || probe(&shared, &embedder, at, retry_after));
|
||||
}
|
||||
while !roster.ready() {
|
||||
let left = (started + within).saturating_duration_since(Instant::now());
|
||||
if left.is_zero() {
|
||||
break;
|
||||
}
|
||||
roster = self.shared.changed.wait_timeout(roster, left).unwrap_or_else(PoisonError::into_inner).0;
|
||||
}
|
||||
}
|
||||
|
||||
fn drop_out(&self, at: usize, why: String) {
|
||||
let mut roster = locked(&self.shared.roster);
|
||||
eprintln!("huginn embedder pool: {} is out of the pool: {why}", roster.members[at].name);
|
||||
roster.members[at].out = Some((Instant::now() + self.retry_after, why));
|
||||
}
|
||||
|
||||
fn active(&self) -> Vec<usize> {
|
||||
locked(&self.shared.roster).active()
|
||||
}
|
||||
|
||||
/// Waits until no read is running: a test's way to say the pool has heard
|
||||
/// every endpoint.
|
||||
#[cfg(test)]
|
||||
fn settle(&self) {
|
||||
let mut roster = locked(&self.shared.roster);
|
||||
while roster.members.iter().any(|member| member.probing) {
|
||||
roster = self.shared.changed.wait_timeout(roster, Duration::from_millis(50)).unwrap_or_else(PoisonError::into_inner).0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<E: Embedder + Send + 'static> Embedder for EmbedderPool<E> {
|
||||
fn built_under(&mut self, built: Option<&ModelIdentity>) {
|
||||
let mut roster = locked(&self.shared.roster);
|
||||
if roster.pinned.as_ref() == built {
|
||||
return;
|
||||
}
|
||||
roster.pinned = built.cloned();
|
||||
for member in &mut roster.members {
|
||||
member.seen = None;
|
||||
member.out = Some((Instant::now(), "not yet read under the collection's model".into()));
|
||||
}
|
||||
}
|
||||
|
||||
fn model_identity(&mut self, within: Duration) -> Result<ModelIdentity> {
|
||||
self.admit(within);
|
||||
let roster = locked(&self.shared.roster);
|
||||
match &roster.pinned {
|
||||
Some(pinned) if !roster.active().is_empty() => Ok(pinned.clone()),
|
||||
_ => Err(roster.unavailable()),
|
||||
}
|
||||
}
|
||||
|
||||
fn embed(&mut self, texts: &[String], within: Duration) -> Result<Vec<Vec<f32>>> {
|
||||
// One budget for the batch: a piece moved off a failed endpoint is
|
||||
// embedded under what is left of it, not under a fresh `within`.
|
||||
let deadline = Instant::now() + within;
|
||||
self.admit(within);
|
||||
let mut vectors: Vec<Option<Vec<f32>>> = vec![None; texts.len()];
|
||||
// Pieces still owed, as (first index, length) over `texts`.
|
||||
let mut owed = vec![(0, texts.len())];
|
||||
while !owed.is_empty() {
|
||||
let mut active = self.active();
|
||||
if active.is_empty() {
|
||||
// Every member that took the last round failed; the reserve may
|
||||
// still be on its first read.
|
||||
self.admit(deadline.saturating_duration_since(Instant::now()));
|
||||
active = self.active();
|
||||
}
|
||||
if active.is_empty() {
|
||||
bail!("{}", locked(&self.shared.roster).unavailable());
|
||||
}
|
||||
let left = deadline.saturating_duration_since(Instant::now());
|
||||
if left.is_zero() {
|
||||
bail!("the batch's budget of {within:?} is spent; {}", locked(&self.shared.roster).unavailable());
|
||||
}
|
||||
// Each owed piece is cut evenly over the active members.
|
||||
let mut jobs: Vec<(usize, usize, usize)> = Vec::new(); // (member, first, length)
|
||||
for (first, length) in owed.drain(..) {
|
||||
let share = length.div_ceil(active.len()).max(1);
|
||||
for (n, start) in (first..first + length).step_by(share).enumerate() {
|
||||
jobs.push((active[n % active.len()], start, share.min(first + length - start)));
|
||||
}
|
||||
}
|
||||
let mut answers: Vec<Answer> = Vec::new();
|
||||
std::thread::scope(|scope| {
|
||||
let mut busy = vec![false; self.embedders.len()];
|
||||
let mut running = Vec::new();
|
||||
for (member, first, length) in jobs {
|
||||
// A member holds one job per round; later pieces for it wait for the next.
|
||||
if std::mem::replace(&mut busy[member], true) {
|
||||
owed.push((first, length));
|
||||
} else {
|
||||
let (piece, embedder) = (&texts[first..first + length], &self.embedders[member]);
|
||||
running.push((member, first, length, scope.spawn(move || locked(embedder).embed(piece, left))));
|
||||
}
|
||||
}
|
||||
for (member, first, length, handle) in running {
|
||||
let answer = handle.join().unwrap_or_else(|_| Err(anyhow::anyhow!("the embedder thread panicked")));
|
||||
answers.push((member, first, length, answer));
|
||||
}
|
||||
});
|
||||
for (member, first, length, answer) in answers {
|
||||
match answer {
|
||||
Ok(found) if found.len() == length => {
|
||||
for (offset, vector) in found.into_iter().enumerate() {
|
||||
vectors[first + offset] = Some(vector);
|
||||
}
|
||||
}
|
||||
Ok(found) => {
|
||||
self.drop_out(member, format!("returned {} vectors for {length} texts", found.len()));
|
||||
owed.push((first, length));
|
||||
}
|
||||
Err(error) => {
|
||||
self.drop_out(member, format!("failed an embed: {error:#}"));
|
||||
owed.push((first, length));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(vectors.into_iter().map(|vector| vector.expect("every piece was embedded")).collect())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::index::fakes::{FakeEmbedder, Gate, identity};
|
||||
|
||||
const WITHIN: Duration = Duration::from_secs(1);
|
||||
|
||||
fn pool(primaries: &[(&str, &str)], reserve: &[(&str, &str)]) -> (EmbedderPool<FakeEmbedder>, Vec<FakeEmbedder>) {
|
||||
let build = |spec: &[(&str, &str)]| -> Vec<(String, FakeEmbedder)> {
|
||||
spec.iter().map(|(name, digest)| ((*name).to_owned(), FakeEmbedder::new(digest))).collect()
|
||||
};
|
||||
let (primaries, reserve) = (build(primaries), build(reserve));
|
||||
let handles = primaries.iter().chain(&reserve).map(|(_, fake)| fake.clone()).collect();
|
||||
(EmbedderPool::new(primaries, reserve).retrying_after(Duration::ZERO), handles)
|
||||
}
|
||||
|
||||
/// Hears every endpoint once, so a test starts from a pool that knows who is up.
|
||||
fn warm(pool: &EmbedderPool<FakeEmbedder>) {
|
||||
pool.admit(WITHIN);
|
||||
pool.settle();
|
||||
}
|
||||
|
||||
fn texts(count: usize) -> Vec<String> {
|
||||
(0..count).map(|n| "x".repeat(n + 1)).collect()
|
||||
}
|
||||
|
||||
fn used(fake: &FakeEmbedder) -> usize {
|
||||
fake.texts().len()
|
||||
}
|
||||
|
||||
fn down(fake: &FakeEmbedder, is_down: bool) {
|
||||
fake.state.lock().unwrap().down = is_down;
|
||||
}
|
||||
|
||||
/// Both primaries take a batch at the same moment: each is inside its embed
|
||||
/// before either is let go, and the vectors come back in the texts' order.
|
||||
#[test]
|
||||
fn both_primaries_embed_one_batch_concurrently_and_the_order_is_kept() {
|
||||
let gate = Gate::shut();
|
||||
let a = FakeEmbedder::new("d1").behind(&gate);
|
||||
let b = FakeEmbedder::new("d1").behind(&gate);
|
||||
let reserve = FakeEmbedder::new("d1");
|
||||
let mut pool = EmbedderPool::new(vec![("raven".into(), a.clone()), ("starfire".into(), b.clone())], vec![("yggdrasil".into(), reserve.clone())]);
|
||||
warm(&pool);
|
||||
let batch = texts(8);
|
||||
let opener = {
|
||||
let gate = gate.clone();
|
||||
std::thread::spawn(move || {
|
||||
gate.wait_arrived(2);
|
||||
gate.open();
|
||||
})
|
||||
};
|
||||
let vectors = pool.embed(&batch, WITHIN).unwrap();
|
||||
opener.join().unwrap();
|
||||
assert_eq!((used(&a), used(&b), used(&reserve)), (4, 4, 0), "an even share each, the reserve idle");
|
||||
assert_eq!(vectors.len(), 8);
|
||||
for (text, vector) in batch.iter().zip(&vectors) {
|
||||
assert_eq!(vector[0], text.len() as f32, "the vector at each position is that text's");
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/// The reserve takes work only while no primary is in the pool, and gives
|
||||
/// it back when a primary returns.
|
||||
#[test]
|
||||
fn the_reserve_serves_only_while_both_primaries_are_out() {
|
||||
let (mut pool, fakes) = pool(&[("raven", "d1"), ("starfire", "d1")], &[("yggdrasil", "d1")]);
|
||||
down(&fakes[0], true);
|
||||
warm(&pool);
|
||||
pool.embed(&texts(4), WITHIN).unwrap();
|
||||
assert_eq!((used(&fakes[0]), used(&fakes[1]), used(&fakes[2])), (0, 4, 0), "one GPU left takes all of it");
|
||||
down(&fakes[1], true);
|
||||
pool.embed(&texts(3), WITHIN).unwrap();
|
||||
assert_eq!(used(&fakes[2]), 3, "both GPUs out, the CPU takes it");
|
||||
down(&fakes[0], false);
|
||||
pool.embed(&texts(2), WITHIN).unwrap();
|
||||
assert_eq!((used(&fakes[0]), used(&fakes[2])), (2, 3), "a GPU back in the pool, the CPU is left alone");
|
||||
}
|
||||
|
||||
/// A piece that fails on one endpoint is embedded on another within the
|
||||
/// same call, in order, and the failed endpoint is not asked again until
|
||||
/// its retry time.
|
||||
#[test]
|
||||
fn a_failing_endpoint_drops_out_its_piece_moves_on_and_it_is_retried_later() {
|
||||
let (mut pool, fakes) = pool(&[("raven", "d1"), ("starfire", "d1")], &[("yggdrasil", "d1")]);
|
||||
warm(&pool);
|
||||
down(&fakes[0], true);
|
||||
let batch = texts(6);
|
||||
let vectors = pool.embed(&batch, WITHIN).unwrap();
|
||||
assert_eq!(used(&fakes[1]), 6, "starfire took raven's piece as well as its own");
|
||||
for (text, vector) in batch.iter().zip(&vectors) {
|
||||
assert_eq!(vector[0], text.len() as f32);
|
||||
}
|
||||
down(&fakes[0], false);
|
||||
warm(&pool);
|
||||
pool.embed(&texts(2), WITHIN).unwrap();
|
||||
assert_eq!(used(&fakes[0]), 1, "back in the pool at the next admission");
|
||||
}
|
||||
|
||||
/// With a retry time that has not come, a dropped endpoint is not read and
|
||||
/// not used; once it has come, it is.
|
||||
#[test]
|
||||
fn a_dropped_endpoint_stays_out_until_its_retry_time() {
|
||||
let (pool, fakes) = pool(&[("raven", "d1"), ("starfire", "d1")], &[]);
|
||||
let mut pool = pool.retrying_after(Duration::from_millis(400));
|
||||
warm(&pool);
|
||||
down(&fakes[0], true);
|
||||
pool.embed(&texts(2), WITHIN).unwrap();
|
||||
down(&fakes[0], false);
|
||||
let reads = fakes[0].state.lock().unwrap().identity_bounds.len();
|
||||
pool.embed(&texts(2), WITHIN).unwrap();
|
||||
assert_eq!((used(&fakes[0]), fakes[0].state.lock().unwrap().identity_bounds.len()), (0, reads), "not read, not used");
|
||||
std::thread::sleep(Duration::from_millis(450));
|
||||
warm(&pool);
|
||||
pool.embed(&texts(2), WITHIN).unwrap();
|
||||
assert_eq!(used(&fakes[0]), 1, "due, read and used");
|
||||
}
|
||||
|
||||
/// An endpoint serving another digest or another length is never written
|
||||
/// with, even when it is the only one answering; the log names it.
|
||||
#[test]
|
||||
fn a_mismatched_endpoint_is_never_written_with() {
|
||||
let (mut pool, fakes) = pool(&[("raven", "d1"), ("starfire", "d2")], &[("yggdrasil", "d1")]);
|
||||
pool.embed(&texts(4), WITHIN).unwrap();
|
||||
assert_eq!((used(&fakes[0]), used(&fakes[1]), used(&fakes[2])), (4, 0, 0));
|
||||
down(&fakes[0], true);
|
||||
pool.embed(&texts(2), WITHIN).unwrap();
|
||||
assert_eq!((used(&fakes[1]), used(&fakes[2])), (0, 2), "starfire is skipped for the CPU, which serves the pool's model");
|
||||
fakes[2].state.lock().unwrap().identity.dimensions = 8;
|
||||
down(&fakes[2], true);
|
||||
assert!(pool.embed(&texts(1), WITHIN).is_err(), "nothing left that serves the pool's model");
|
||||
down(&fakes[2], false);
|
||||
let error = format!("{:#}", pool.embed(&texts(1), WITHIN).unwrap_err());
|
||||
assert!(error.contains("raven") && error.contains("starfire") && error.contains("yggdrasil"), "{error}");
|
||||
assert_eq!(used(&fakes[1]), 0);
|
||||
}
|
||||
|
||||
/// With every endpoint down, the identity read and the embed both fail
|
||||
/// naming each; the pool's model is what the first endpoint reported.
|
||||
#[test]
|
||||
fn with_every_endpoint_down_the_pool_fails_naming_them_all() {
|
||||
let (mut pool, fakes) = pool(&[("raven", "d1")], &[("yggdrasil", "d1")]);
|
||||
assert_eq!(pool.model_identity(WITHIN).unwrap().digest, "d1");
|
||||
for fake in &fakes {
|
||||
down(fake, true);
|
||||
}
|
||||
let error = format!("{:#}", pool.embed(&texts(2), WITHIN).unwrap_err());
|
||||
assert!(error.contains("raven") && error.contains("yggdrasil"), "{error}");
|
||||
let (mut cold, cold_fakes) = self::pool(&[("raven", "d1")], &[("yggdrasil", "d1")]);
|
||||
down(&cold_fakes[0], true);
|
||||
down(&cold_fakes[1], true);
|
||||
assert!(format!("{:#}", cold.model_identity(WITHIN).unwrap_err()).contains("raven"));
|
||||
}
|
||||
|
||||
/// An endpoint that never answers an identity read costs a batch nothing:
|
||||
/// across several retry windows every batch finishes at the healthy
|
||||
/// endpoint's pace, however long the caller's budget is, and the dark
|
||||
/// endpoint is read under the pool's own short bound, never the caller's.
|
||||
#[test]
|
||||
fn a_dark_endpoint_never_delays_a_batch_a_healthy_one_can_take() {
|
||||
let (pool, fakes) = pool(&[("raven", "d1"), ("starfire", "d1")], &[("yggdrasil", "d1")]);
|
||||
let mut pool = pool.retrying_after(Duration::from_millis(100));
|
||||
fakes[1].state.lock().unwrap().hang_identity = true;
|
||||
pool.built_under(Some(&identity("d1")));
|
||||
let budget = Duration::from_secs(300);
|
||||
for _ in 0..4 {
|
||||
let started = Instant::now();
|
||||
pool.embed(&texts(4), budget).unwrap();
|
||||
assert!(started.elapsed() < Duration::from_millis(1500), "a batch waited {:?} on a dark endpoint", started.elapsed());
|
||||
std::thread::sleep(Duration::from_millis(150));
|
||||
}
|
||||
assert_eq!((used(&fakes[0]), used(&fakes[1]), used(&fakes[2])), (16, 0, 0));
|
||||
let bounds = fakes[1].state.lock().unwrap().identity_bounds.clone();
|
||||
assert!(!bounds.is_empty() && bounds.iter().all(|bound| *bound == IDENTITY_WITHIN), "{bounds:?}");
|
||||
}
|
||||
|
||||
/// The model of the pool is the collection's: an endpoint serving another
|
||||
/// is not written with, however alone it is, and the identity is not given.
|
||||
#[test]
|
||||
fn the_pool_holds_endpoints_to_the_collections_model() {
|
||||
let (mut pool, fakes) = pool(&[("raven", "d2"), ("starfire", "d1")], &[("yggdrasil", "d2")]);
|
||||
pool.built_under(Some(&identity("d1")));
|
||||
pool.embed(&texts(4), WITHIN).unwrap();
|
||||
assert_eq!((used(&fakes[0]), used(&fakes[1]), used(&fakes[2])), (0, 4, 0), "only the endpoint serving the collection's model works");
|
||||
down(&fakes[1], true);
|
||||
assert!(pool.embed(&texts(1), WITHIN).is_err(), "a different model is not a substitute");
|
||||
assert!(pool.model_identity(WITHIN).is_err(), "and the pool does not claim the model of another");
|
||||
assert_eq!((used(&fakes[0]), used(&fakes[2])), (0, 0));
|
||||
down(&fakes[1], false);
|
||||
assert_eq!(pool.model_identity(WITHIN).unwrap().digest, "d1");
|
||||
}
|
||||
|
||||
/// With no collection to hold to, the model is the first in list order that
|
||||
/// answers, whichever answers first; handing the pool a collection's model
|
||||
/// changes it, and handing it the same one again reads nothing.
|
||||
#[test]
|
||||
fn the_pool_adopts_a_model_only_without_a_collection_and_a_new_pin_rereads() {
|
||||
let (slow, quick) = (FakeEmbedder::new("d1"), FakeEmbedder::new("d2"));
|
||||
slow.state.lock().unwrap().identity_delay = Duration::from_millis(200);
|
||||
let mut pool = EmbedderPool::new(vec![("raven".into(), slow.clone()), ("starfire".into(), quick.clone())], vec![]);
|
||||
assert_eq!(pool.model_identity(WITHIN).unwrap().digest, "d1", "list order, not the quickest");
|
||||
pool.built_under(Some(&identity("d1")));
|
||||
let reads = quick.state.lock().unwrap().identity_bounds.len();
|
||||
pool.model_identity(WITHIN).unwrap();
|
||||
assert_eq!(quick.state.lock().unwrap().identity_bounds.len(), reads, "the same model again reads nothing");
|
||||
pool.built_under(Some(&identity("d2")));
|
||||
assert_eq!(pool.model_identity(WITHIN).unwrap().digest, "d2");
|
||||
pool.embed(&texts(2), WITHIN).unwrap();
|
||||
assert_eq!((used(&slow), used(&quick)), (0, 2));
|
||||
pool.built_under(None);
|
||||
assert_eq!(pool.model_identity(WITHIN).unwrap().digest, "d1", "no collection, the pool adopts again");
|
||||
}
|
||||
|
||||
/// A batch waits only for what is coming: with every primary out and not due,
|
||||
/// the reserve takes it at once, and with nothing left it fails at once,
|
||||
/// however long the caller's budget.
|
||||
#[test]
|
||||
fn a_batch_does_not_wait_for_reads_that_are_not_running() {
|
||||
let (pool, fakes) = pool(&[("raven", "d1"), ("starfire", "d1")], &[("yggdrasil", "d1")]);
|
||||
let mut pool = pool.retrying_after(Duration::from_secs(3600));
|
||||
warm(&pool);
|
||||
let budget = Duration::from_secs(10);
|
||||
down(&fakes[0], true);
|
||||
down(&fakes[1], true);
|
||||
let started = Instant::now();
|
||||
pool.embed(&texts(2), budget).unwrap();
|
||||
assert_eq!(used(&fakes[2]), 2);
|
||||
assert!(started.elapsed() < Duration::from_secs(2), "the reserve waited {:?} for GPUs that are not being read", started.elapsed());
|
||||
down(&fakes[2], true);
|
||||
let started = Instant::now();
|
||||
assert!(pool.embed(&texts(1), budget).is_err());
|
||||
assert!(started.elapsed() < Duration::from_secs(2), "the failure waited {:?} for reads that are not running", started.elapsed());
|
||||
}
|
||||
|
||||
/// An endpoint whose read failed is not read again before its retry time.
|
||||
#[test]
|
||||
fn an_endpoint_that_failed_its_read_is_not_read_again_before_its_retry_time() {
|
||||
let (pool, fakes) = pool(&[("raven", "d1"), ("starfire", "d1")], &[]);
|
||||
let pool = pool.retrying_after(Duration::from_secs(3600));
|
||||
down(&fakes[0], true);
|
||||
warm(&pool);
|
||||
let reads = fakes[0].state.lock().unwrap().identity_bounds.len();
|
||||
warm(&pool);
|
||||
assert_eq!(fakes[0].state.lock().unwrap().identity_bounds.len(), reads);
|
||||
}
|
||||
|
||||
/// A primary that is still being read keeps the reserve waiting, so the
|
||||
/// first batch goes to the GPU and not to the CPU that answered first.
|
||||
#[test]
|
||||
fn a_primary_still_being_read_is_waited_for_before_the_reserve_is_used() {
|
||||
let (mut pool, fakes) = pool(&[("raven", "d1")], &[("yggdrasil", "d1")]);
|
||||
fakes[0].state.lock().unwrap().identity_delay = Duration::from_millis(200);
|
||||
pool.built_under(Some(&identity("d1")));
|
||||
pool.embed(&texts(3), WITHIN).unwrap();
|
||||
assert_eq!((used(&fakes[0]), used(&fakes[1])), (3, 0));
|
||||
}
|
||||
|
||||
// Real Ollama clients against sockets that go wrong after the pool has
|
||||
// admitted them.
|
||||
|
||||
use std::io::{Read, Write};
|
||||
use std::net::{TcpListener, TcpStream};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
use crate::index::ollama::{EmbedBounds, OllamaEmbedder};
|
||||
|
||||
const TAGS: &str = r#"{"models":[{"name":"m:1","model":"m:1","digest":"d1"}]}"#;
|
||||
const SHOW: &str = r#"{"model_info":{"qwen3.embedding_length":1}}"#;
|
||||
const QUICK: EmbedBounds =
|
||||
EmbedBounds { connect: Duration::from_millis(200), floor: Duration::from_millis(300), per_text: Duration::ZERO };
|
||||
|
||||
/// An Ollama on a local port. `stalls` makes it take an embed request and
|
||||
/// never answer; `gone` closes the listener, so the next connect is refused.
|
||||
/// An embedding is `[text length]`, so order is checkable.
|
||||
struct Endpoint {
|
||||
url: String,
|
||||
stalls: Arc<AtomicBool>,
|
||||
gone: Arc<AtomicBool>,
|
||||
listening: std::thread::JoinHandle<()>,
|
||||
}
|
||||
|
||||
fn endpoint() -> Endpoint {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
listener.set_nonblocking(true).unwrap();
|
||||
let url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let (stalls, gone) = (Arc::new(AtomicBool::new(false)), Arc::new(AtomicBool::new(false)));
|
||||
let (stalling, leaving) = (Arc::clone(&stalls), Arc::clone(&gone));
|
||||
let listening = std::thread::spawn(move || {
|
||||
while !leaving.load(Ordering::SeqCst) {
|
||||
match listener.accept() {
|
||||
Ok((stream, _)) => {
|
||||
let stalling = Arc::clone(&stalling);
|
||||
std::thread::spawn(move || answer(stream, &stalling));
|
||||
}
|
||||
Err(_) => std::thread::sleep(Duration::from_millis(5)),
|
||||
}
|
||||
}
|
||||
});
|
||||
Endpoint { url, stalls, gone, listening }
|
||||
}
|
||||
|
||||
fn answer(mut stream: TcpStream, stalls: &AtomicBool) {
|
||||
stream.set_nonblocking(false).unwrap();
|
||||
let mut buffer = Vec::new();
|
||||
let (head, length) = loop {
|
||||
let mut chunk = [0_u8; 4096];
|
||||
match stream.read(&mut chunk) {
|
||||
Ok(0) | Err(_) => return,
|
||||
Ok(read) => buffer.extend_from_slice(&chunk[..read]),
|
||||
}
|
||||
let text = String::from_utf8_lossy(&buffer).into_owned();
|
||||
if let Some(end) = text.find("\r\n\r\n") {
|
||||
let length = text.lines().find_map(|line| line.to_ascii_lowercase().strip_prefix("content-length:").map(|v| v.trim().parse().unwrap())).unwrap_or(0);
|
||||
break (end + 4, length);
|
||||
}
|
||||
};
|
||||
while buffer.len() < head + length {
|
||||
let mut chunk = [0_u8; 4096];
|
||||
match stream.read(&mut chunk) {
|
||||
Ok(0) | Err(_) => return,
|
||||
Ok(read) => buffer.extend_from_slice(&chunk[..read]),
|
||||
}
|
||||
}
|
||||
let request = String::from_utf8_lossy(&buffer[..head]).into_owned();
|
||||
let body = if request.starts_with("GET /api/tags") {
|
||||
TAGS.to_owned()
|
||||
} else if request.starts_with("POST /api/show") {
|
||||
SHOW.to_owned()
|
||||
} else {
|
||||
if stalls.load(Ordering::SeqCst) {
|
||||
std::thread::sleep(Duration::from_secs(30));
|
||||
return;
|
||||
}
|
||||
let asked: serde_json::Value = serde_json::from_slice(&buffer[head..head + length]).unwrap();
|
||||
let vectors: Vec<Vec<usize>> = asked["input"].as_array().unwrap().iter().map(|text| vec![text.as_str().unwrap().len()]).collect();
|
||||
serde_json::json!({ "embeddings": vectors }).to_string()
|
||||
};
|
||||
let _ = write!(stream, "HTTP/1.1 200 OK\r\ncontent-length: {}\r\n\r\n{body}", body.len());
|
||||
}
|
||||
|
||||
fn real_pool(endpoints: &[&Endpoint], bounds: EmbedBounds) -> EmbedderPool<OllamaEmbedder> {
|
||||
let members = endpoints.iter().map(|at| (at.url.clone(), OllamaEmbedder::new(&at.url, "m:1").bounded_by(bounds))).collect();
|
||||
let pool = EmbedderPool::new(members, Vec::new()).retrying_after(Duration::from_secs(3600));
|
||||
warm_real(&pool);
|
||||
pool
|
||||
}
|
||||
|
||||
fn warm_real(pool: &EmbedderPool<OllamaEmbedder>) {
|
||||
pool.admit(Duration::from_secs(5));
|
||||
pool.settle();
|
||||
}
|
||||
|
||||
fn lengths(vectors: &[Vec<f32>]) -> Vec<f32> {
|
||||
vectors.iter().map(|vector| vector[0]).collect()
|
||||
}
|
||||
|
||||
/// An endpoint that takes the connection and then says nothing costs the
|
||||
/// batch its own bound, not the caller's budget; its share goes to the
|
||||
/// healthy endpoint, in order, and it stays out until its retry time.
|
||||
#[test]
|
||||
fn an_endpoint_that_stalls_mid_batch_costs_the_batch_its_bound_not_the_budget() {
|
||||
let (healthy, stalling) = (endpoint(), endpoint());
|
||||
let mut pool = real_pool(&[&healthy, &stalling], QUICK);
|
||||
stalling.stalls.store(true, Ordering::SeqCst);
|
||||
let started = Instant::now();
|
||||
let vectors = pool.embed(&texts(8), Duration::from_secs(20)).unwrap();
|
||||
assert!(started.elapsed() < Duration::from_secs(3), "took {:?}", started.elapsed());
|
||||
assert_eq!(lengths(&vectors), (1..=8).map(|n| n as f32).collect::<Vec<_>>());
|
||||
assert!(locked(&pool.shared.roster).members[1].out.is_some());
|
||||
let again = Instant::now();
|
||||
pool.embed(&texts(4), Duration::from_secs(20)).unwrap();
|
||||
assert!(again.elapsed() < Duration::from_millis(250), "the dropped endpoint was tried again: {:?}", again.elapsed());
|
||||
}
|
||||
|
||||
/// An endpoint that stops listening after it was admitted: the same.
|
||||
#[test]
|
||||
fn an_endpoint_that_refuses_mid_batch_costs_the_batch_nothing_the_healthy_one_could_not_do() {
|
||||
let (healthy, leaving) = (endpoint(), endpoint());
|
||||
let mut pool = real_pool(&[&healthy, &leaving], QUICK);
|
||||
leaving.gone.store(true, Ordering::SeqCst);
|
||||
leaving.listening.join().unwrap();
|
||||
let started = Instant::now();
|
||||
let vectors = pool.embed(&texts(6), Duration::from_secs(20)).unwrap();
|
||||
assert!(started.elapsed() < Duration::from_secs(3), "took {:?}", started.elapsed());
|
||||
assert_eq!(lengths(&vectors), (1..=6).map(|n| n as f32).collect::<Vec<_>>());
|
||||
assert!(locked(&pool.shared.roster).members[1].out.is_some());
|
||||
}
|
||||
|
||||
/// The bound never exceeds what the caller gave: with one endpoint, stalled
|
||||
/// and a floor far above the budget, the batch fails at the budget.
|
||||
#[test]
|
||||
fn a_stall_is_never_waited_on_beyond_the_callers_budget() {
|
||||
let only = endpoint();
|
||||
let generous = EmbedBounds { floor: Duration::from_secs(60), ..QUICK };
|
||||
let mut pool = real_pool(&[&only], generous);
|
||||
only.stalls.store(true, Ordering::SeqCst);
|
||||
let started = Instant::now();
|
||||
assert!(pool.embed(&texts(2), Duration::from_millis(400)).is_err());
|
||||
assert!(started.elapsed() < Duration::from_secs(3), "took {:?}", started.elapsed());
|
||||
}
|
||||
}
|
||||
|
|
@ -1,18 +1,23 @@
|
|||
//! 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
|
||||
/// never lists or touches another.
|
||||
pub struct QdrantIndex {
|
||||
http: Http,
|
||||
/// The collection the source port reads and writes: `huginn_sources`,
|
||||
/// unless a live test points it at a scratch one.
|
||||
sources: String,
|
||||
}
|
||||
|
||||
/// The ids fetched per scroll page.
|
||||
|
|
@ -20,7 +25,13 @@ const SCROLL_PAGE: u64 = 1000;
|
|||
|
||||
impl QdrantIndex {
|
||||
pub fn new(base_url: &str) -> Self {
|
||||
Self { http: Http::new(base_url, INDEX_CALL_TIMEOUT) }
|
||||
Self { http: Http::new(base_url, INDEX_CALL_TIMEOUT), sources: COLLECTION.to_owned() }
|
||||
}
|
||||
|
||||
/// The source port on another collection, for a live test's scratch one.
|
||||
pub fn with_sources_collection(mut self, collection: &str) -> Self {
|
||||
self.sources = collection.to_owned();
|
||||
self
|
||||
}
|
||||
|
||||
fn ask(&self, method: Method, path: &str, body: Option<&Value>, ok: &[u16]) -> Result<(u16, Value)> {
|
||||
|
|
@ -163,3 +174,144 @@ 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 collection = self.sources.clone();
|
||||
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 collection = self.sources.clone();
|
||||
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<()> {
|
||||
let collection = self.sources.clone();
|
||||
// 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 collection = self.sources.clone();
|
||||
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 collection = self.sources.clone();
|
||||
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 collection = self.sources.clone();
|
||||
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()
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1143,6 +1143,8 @@ fn shared_with(
|
|||
closed,
|
||||
searches: searches.into(),
|
||||
found: Vec::new(),
|
||||
source_searches: VecDeque::new(),
|
||||
sources_found: Vec::new(),
|
||||
next_ticket: 0,
|
||||
}),
|
||||
Condvar::new(),
|
||||
|
|
@ -1968,3 +1970,117 @@ 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");
|
||||
}
|
||||
|
||||
/// Source searches take their own tickets, share the queue's bound with the
|
||||
/// mind's searches, and an abandon removes the one ticket named from either.
|
||||
#[test]
|
||||
fn source_searches_have_distinct_tickets_share_the_queue_bound_and_abandon_by_ticket() {
|
||||
let shared = shared_with(Vec::new(), None, Vec::new(), false);
|
||||
let mut sink = WorkerSink { shared, identified: true, deadline: SEARCH_DEADLINE };
|
||||
let first = ask_sources(&mut sink, "a").unwrap();
|
||||
let second = ask_sources(&mut sink, "b").unwrap();
|
||||
let mind = ask_sink(&mut sink, "m", 3).unwrap();
|
||||
assert!(first != second && second != mind && first != mind, "{first:?} {second:?} {mind:?}");
|
||||
|
||||
abandon(&mut sink, first);
|
||||
{
|
||||
let guard = lock(&sink.shared);
|
||||
let left: Vec<SearchTicket> = guard.source_searches.iter().map(|queued| queued.ticket).collect();
|
||||
assert_eq!((left, guard.searches.len()), (vec![second], 1), "only the named ticket went");
|
||||
}
|
||||
|
||||
for _ in 0..SEARCHES_QUEUED_MAX - 2 {
|
||||
ask_sources(&mut sink, "fill").unwrap();
|
||||
}
|
||||
assert!(ask_sources(&mut sink, "over").is_err(), "the bound counts both kinds");
|
||||
assert!(ask_sink(&mut sink, "over", 3).is_err(), "and refuses both kinds");
|
||||
abandon(&mut sink, second);
|
||||
assert!(ask_sink(&mut sink, "now", 3).is_ok(), "a freed slot is a slot");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -7,13 +7,13 @@ use std::sync::atomic::AtomicBool;
|
|||
|
||||
use anyhow::{Context, Result};
|
||||
use huginn_daemon::index::Backoff;
|
||||
use huginn_daemon::index::ollama::OllamaEmbedder;
|
||||
use huginn_daemon::index::ollama::pool;
|
||||
use huginn_daemon::index::qdrant::QdrantIndex;
|
||||
use huginn_daemon::serve::{ServeOptions, parse_options, run, startup};
|
||||
|
||||
fn main() -> Result<()> {
|
||||
let options = parse_options(std::env::args().skip(1))?;
|
||||
let embedder = OllamaEmbedder::new(&options.ollama_url, &options.embedding_model);
|
||||
let embedder = pool(&options.ollama_urls, &options.reserve_ollama_urls, &options.embedding_model);
|
||||
let index = QdrantIndex::new(&options.qdrant_url);
|
||||
let (mut daemon, mut hub, registry) = match startup(&options, embedder, index, Backoff::default()) {
|
||||
Ok(started) => started,
|
||||
|
|
|
|||
|
|
@ -101,7 +101,7 @@ impl Default for ServeOptions {
|
|||
}
|
||||
}
|
||||
|
||||
/// The six things the operator supplies. Nothing here reads the environment:
|
||||
/// The things the operator supplies. Nothing here reads the environment:
|
||||
/// how Idunn supplies `--bind` is Cut 14's. The index has no off switch: an
|
||||
/// unreachable Qdrant or Ollama is a degraded projection, reported by
|
||||
/// `whoami`, not a configuration.
|
||||
|
|
@ -111,23 +111,34 @@ pub struct Options {
|
|||
pub instance: Slug,
|
||||
pub bind: SocketAddr,
|
||||
pub qdrant_url: String,
|
||||
pub ollama_url: String,
|
||||
/// The GPU hosts: all take embedding at once. One is the plain case.
|
||||
pub ollama_urls: Vec<String>,
|
||||
/// The CPU host: takes embedding only while no GPU host can.
|
||||
pub reserve_ollama_urls: Vec<String>,
|
||||
pub embedding_model: String,
|
||||
}
|
||||
|
||||
/// Exactly `--state-root`, `--instance`, `--bind`, `--qdrant-url`,
|
||||
/// `--ollama-url` and `--embedding-model`, each required, each once. The two
|
||||
/// URLs are plain `http://`: this build carries no TLS.
|
||||
/// `--ollama-url` and `--embedding-model`, each required, each once, except
|
||||
/// `--ollama-url`, which repeats for each GPU host, and the optional
|
||||
/// `--reserve-ollama-url`, which repeats for each CPU host. The URLs are plain
|
||||
/// `http://`: this build carries no TLS.
|
||||
pub fn parse_options(args: impl Iterator<Item = String>) -> Result<Options> {
|
||||
let mut args = args;
|
||||
let mut values: BTreeMap<String, String> = BTreeMap::new();
|
||||
let (mut ollama_urls, mut reserve_ollama_urls) = (Vec::new(), 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" | "instance" | "bind" | "qdrant-url" | "ollama-url" | "embedding-model"),
|
||||
matches!(name, "state-root" | "instance" | "bind" | "qdrant-url" | "ollama-url" | "reserve-ollama-url" | "embedding-model"),
|
||||
"unsupported Huginn option --{name}"
|
||||
);
|
||||
let value = args.next().with_context(|| format!("missing value for --{name}"))?;
|
||||
if name == "ollama-url" || name == "reserve-ollama-url" {
|
||||
ensure!(value.starts_with("http://"), "--{name} must be a plain http:// address, got {value:?}");
|
||||
if name == "ollama-url" { ollama_urls.push(value) } else { reserve_ollama_urls.push(value) }
|
||||
continue;
|
||||
}
|
||||
ensure!(values.insert(name.to_owned(), value).is_none(), "duplicate Huginn option --{name}");
|
||||
}
|
||||
let take = |name: &str| -> Result<String> {
|
||||
|
|
@ -148,7 +159,11 @@ pub fn parse_options(args: impl Iterator<Item = String>) -> Result<Options> {
|
|||
instance: Slug(take("instance")?),
|
||||
bind: bind.parse().with_context(|| format!("--bind must be an ip:port, got {bind:?}"))?,
|
||||
qdrant_url: url("qdrant-url")?,
|
||||
ollama_url: url("ollama-url")?,
|
||||
ollama_urls: {
|
||||
ensure!(!ollama_urls.is_empty(), "--ollama-url is required; repeat it for each GPU host");
|
||||
ollama_urls
|
||||
},
|
||||
reserve_ollama_urls,
|
||||
embedding_model,
|
||||
})
|
||||
}
|
||||
|
|
@ -167,7 +182,7 @@ pub fn startup<E, V>(
|
|||
) -> Result<(Daemon<OwnedRedbMessagePackBackingStore, WorkerSink>, CultNetRudpServerHub, CultNetSchemaRegistry)>
|
||||
where
|
||||
E: Embedder + Send + 'static,
|
||||
V: VectorIndex + Send + 'static,
|
||||
V: VectorIndex + crate::sources::SourceIndex + Send + 'static,
|
||||
{
|
||||
let mind = Mind::open(&options.state_root, &options.instance)?;
|
||||
let hub = bind(options.bind, &runtime_id(&options.instance))?;
|
||||
|
|
@ -522,14 +537,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
|
||||
|
|
@ -1347,7 +1362,8 @@ mod tests {
|
|||
instance: slug(INSTANCE),
|
||||
bind: port,
|
||||
qdrant_url: "http://127.0.0.1:1".into(),
|
||||
ollama_url: "http://127.0.0.1:1".into(),
|
||||
ollama_urls: vec!["http://127.0.0.1:1".into()],
|
||||
reserve_ollama_urls: Vec::new(),
|
||||
embedding_model: "a-model".into(),
|
||||
};
|
||||
let started = |options: &Options| {
|
||||
|
|
@ -1369,7 +1385,7 @@ mod tests {
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn options_are_exactly_the_six_the_operator_supplies() {
|
||||
fn options_are_the_ones_the_operator_supplies() {
|
||||
let root = if cfg!(windows) { r"C:\state" } else { "/state" };
|
||||
let args = |values: &[&str]| values.iter().map(|value| value.to_string()).collect::<Vec<_>>().into_iter();
|
||||
let good = [
|
||||
|
|
@ -1383,7 +1399,8 @@ mod tests {
|
|||
instance: slug("yggdrasil"),
|
||||
bind: "127.0.0.1:17872".parse().unwrap(),
|
||||
qdrant_url: "http://127.0.0.1:6333".into(),
|
||||
ollama_url: "http://10.77.0.4:11434".into(),
|
||||
ollama_urls: vec!["http://10.77.0.4:11434".into()],
|
||||
reserve_ollama_urls: Vec::new(),
|
||||
embedding_model: "qwen3-embedding:0.6b".into(),
|
||||
}
|
||||
);
|
||||
|
|
@ -1411,5 +1428,11 @@ mod tests {
|
|||
] {
|
||||
assert!(parse_options(args(&bad)).is_err(), "{bad:?}");
|
||||
}
|
||||
// The GPU hosts and the CPU host repeat; a lone --ollama-url is the plain case.
|
||||
let pooled = [good.to_vec(), vec!["--ollama-url", "http://10.77.0.2:11434", "--reserve-ollama-url", "http://127.0.0.1:11434"]].concat();
|
||||
let parsed = parse_options(args(&pooled)).unwrap();
|
||||
assert_eq!(parsed.ollama_urls, ["http://10.77.0.4:11434", "http://10.77.0.2:11434"]);
|
||||
assert_eq!(parsed.reserve_ollama_urls, ["http://127.0.0.1:11434"]);
|
||||
assert!(parse_options(args(&[good.to_vec(), vec!["--reserve-ollama-url", "127.0.0.1:11434"]].concat())).is_err());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
153
crates/huginn-daemon/src/sources/chunk.rs
Normal file
153
crates/huginn-daemon/src/sources/chunk.rs
Normal file
|
|
@ -0,0 +1,153 @@
|
|||
//! 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() {}");
|
||||
}
|
||||
|
||||
/// The character bound is inclusive: a line of exactly the bound is kept
|
||||
/// whole, one more character is cut; and the running count of a window
|
||||
/// includes each joining newline.
|
||||
#[test]
|
||||
fn the_character_bound_is_inclusive_and_counts_the_joining_newlines() {
|
||||
let exact = "z".repeat(CHUNK_CHARS_MAX);
|
||||
let kept = chunk(&format!("{exact}\n"));
|
||||
assert_eq!((kept.len(), kept[0].text.chars().count()), (1, CHUNK_CHARS_MAX));
|
||||
let over = chunk(&format!("{exact}z\nnext\n"));
|
||||
let shown: Vec<(u32, u32, usize)> = over.iter().map(|held| (held.line_start, held.line_end, held.text.chars().count())).collect();
|
||||
assert_eq!(shown, [(1, 1, CHUNK_CHARS_MAX), (2, 2, 4)]);
|
||||
|
||||
// 13 lines of 100 characters are 1,312 with their newlines; the 14th would make 1,413.
|
||||
let hundred = "h".repeat(100);
|
||||
let text: String = (0..20).map(|_| format!("{hundred}\n")).collect();
|
||||
let ranges: Vec<(u32, u32)> = chunk(&text).iter().map(|held| (held.line_start, held.line_end)).collect();
|
||||
assert_eq!(ranges, [(1, 13), (14, 20)]);
|
||||
}
|
||||
}
|
||||
178
crates/huginn-daemon/src/sources/forge.rs
Normal file
178
crates/huginn-daemon/src/sources/forge.rs
Normal file
|
|
@ -0,0 +1,178 @@
|
|||
//! 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::{Catalog, Forge, SourceRepo};
|
||||
|
||||
/// 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 {
|
||||
agent: ureq::Agent,
|
||||
base: String,
|
||||
}
|
||||
|
||||
impl HttpForge {
|
||||
pub fn new(base_url: &str) -> Self {
|
||||
let agent: ureq::Agent =
|
||||
ureq::Agent::config_builder().timeout_global(Some(CALL_TIMEOUT)).http_status_as_error(false).build().into();
|
||||
Self { agent, base: base_url.trim_end_matches('/').to_owned() }
|
||||
}
|
||||
|
||||
/// One GET: the status, the count of repos the forge says it holds
|
||||
/// (`X-Total-Count`, which only a header carries, so this does not go
|
||||
/// through `index::Http`), and the body.
|
||||
fn page(&self, path: &str) -> Result<(u16, Option<usize>, String)> {
|
||||
let url = format!("{}{path}", self.base);
|
||||
let mut response = self.agent.get(&url).call().with_context(|| format!("no answer from {url}"))?;
|
||||
let status = response.status().as_u16();
|
||||
let total = response.headers().get("x-total-count").and_then(|value| value.to_str().ok()).and_then(|value| value.trim().parse().ok());
|
||||
let body = response.body_mut().read_to_string().with_context(|| format!("reading the answer from {url}"))?;
|
||||
Ok((status, total, body))
|
||||
}
|
||||
}
|
||||
|
||||
/// Why the pages listed are not the whole catalog, or `None` when the repos
|
||||
/// listed (before any are filtered out) are the count the forge states. A forge
|
||||
/// that states no count cannot vouch for its list.
|
||||
pub(crate) fn shortfall(listed: usize, stated: Option<usize>) -> Option<String> {
|
||||
match stated {
|
||||
None => Some("the forge stated no X-Total-Count".to_owned()),
|
||||
Some(total) if total == listed => None,
|
||||
Some(total) => Some(format!("the forge counts {total} repos and {listed} were listed")),
|
||||
}
|
||||
}
|
||||
|
||||
/// 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<Catalog> {
|
||||
let mut found = Vec::new();
|
||||
let (mut listed_in_all, mut stated) = (0, None);
|
||||
for page in 1.. {
|
||||
let path = format!("/api/v1/orgs/{ORG}/repos?page={page}&limit={PAGE}");
|
||||
let (status, total, answer) = self.page(&path)?;
|
||||
ensure!(status == 200, "the forge answered {status} to {path}");
|
||||
stated = stated.or(total);
|
||||
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);
|
||||
listed_in_all += listed;
|
||||
found.extend(indexable(&parsed)?);
|
||||
if listed < PAGE {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(Catalog { repos: found, shortfall: shortfall(listed_in_all, stated) })
|
||||
}
|
||||
|
||||
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"]);
|
||||
}
|
||||
|
||||
/// Serves each page in `pages` once, to the request that names it, with
|
||||
/// `total` as its `X-Total-Count` when given.
|
||||
fn serve(pages: Vec<String>, total: Option<usize>) -> String {
|
||||
use std::io::{Read, Write};
|
||||
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
std::thread::spawn(move || {
|
||||
for _ in 0..pages.len() {
|
||||
let (mut stream, _) = listener.accept().unwrap();
|
||||
let mut buffer = [0_u8; 2048];
|
||||
let read = stream.read(&mut buffer).unwrap();
|
||||
let request = String::from_utf8_lossy(&buffer[..read]).into_owned();
|
||||
let body = if request.contains("page=1&") { &pages[0] } else { &pages[1] };
|
||||
let count = total.map_or_else(String::new, |total| format!("x-total-count: {total}\r\n"));
|
||||
write!(stream, "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\n{count}content-length: {}\r\nconnection: close\r\n\r\n{body}", body.len()).unwrap();
|
||||
}
|
||||
});
|
||||
format!("http://{address}")
|
||||
}
|
||||
|
||||
fn page_of(range: std::ops::Range<usize>) -> String {
|
||||
let repos: Vec<Value> = range.map(|n| json!({ "name": format!("r{n}"), "default_branch": "main", "private": false, "archived": false, "mirror": false, "empty": false })).collect();
|
||||
Value::Array(repos).to_string()
|
||||
}
|
||||
|
||||
/// The catalog reads every page until one is short: a full first page
|
||||
/// does not end it, and a short second page does; clone URLs name the org.
|
||||
#[test]
|
||||
fn the_catalog_reads_pages_until_a_short_one_and_clones_from_the_org() {
|
||||
let base = serve(vec![page_of(0..PAGE), page_of(PAGE..PAGE + 3)], Some(PAGE + 3));
|
||||
let mut forge = HttpForge::new(&format!("{base}/"));
|
||||
let catalog = forge.catalog().unwrap();
|
||||
let names: Vec<String> = catalog.repos.into_iter().map(|repo| repo.name).collect();
|
||||
assert_eq!(names.len(), PAGE + 3);
|
||||
assert_eq!((names[0].as_str(), names[PAGE + 2].as_str()), ("r0", "r52"));
|
||||
assert_eq!(catalog.shortfall, None, "the pages sum to the forge's count");
|
||||
assert_eq!(forge.clone_url("Demo"), format!("{base}/GameCult/Demo.git"));
|
||||
}
|
||||
|
||||
/// A catalog whose pages do not sum to the count the forge states is
|
||||
/// marked short: a forge that caps a page below ours stops the paging at
|
||||
/// the first short page, and an empty answer or a missing count vouches for
|
||||
/// nothing.
|
||||
#[test]
|
||||
fn a_catalog_that_does_not_sum_to_the_forges_count_is_short() {
|
||||
let truncated = serve(vec![page_of(0..30)], Some(62));
|
||||
let catalog = HttpForge::new(&truncated).catalog().unwrap();
|
||||
assert_eq!(catalog.repos.len(), 30);
|
||||
assert!(catalog.shortfall.as_deref().unwrap().contains("counts 62 repos and 30 were listed"));
|
||||
|
||||
let empty = serve(vec!["[]".into()], Some(62));
|
||||
let catalog = HttpForge::new(&empty).catalog().unwrap();
|
||||
assert!(catalog.repos.is_empty() && catalog.shortfall.is_some());
|
||||
|
||||
let unstated = serve(vec![page_of(0..3)], None);
|
||||
assert!(HttpForge::new(&unstated).catalog().unwrap().shortfall.unwrap().contains("X-Total-Count"));
|
||||
}
|
||||
}
|
||||
204
crates/huginn-daemon/src/sources/git.rs
Normal file
204
crates/huginn-daemon/src/sources/git.rs
Normal file
|
|
@ -0,0 +1,204 @@
|
|||
//! 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()))?;
|
||||
let mut init = Command::new("git");
|
||||
init.arg("init").arg("--bare").arg("--quiet").arg(&self.dir);
|
||||
run(init, "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);
|
||||
}
|
||||
}
|
||||
472
crates/huginn-daemon/src/sources/mod.rs
Normal file
472
crates/huginn-daemon/src/sources/mod.rs
Normal file
|
|
@ -0,0 +1,472 @@
|
|||
//! 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.
|
||||
//! It is dropped before the repo's first delete, so a run that fails between
|
||||
//! the delete and the upserts leaves no cursor, and the repo is ingested whole.
|
||||
//! The cursors object is written whole by this job alone, once per repo.
|
||||
|
||||
pub mod chunk;
|
||||
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, ModelIdentity, 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,
|
||||
}
|
||||
|
||||
impl SourcesMeta {
|
||||
/// The model the collection was built with.
|
||||
pub fn model_identity(&self) -> ModelIdentity {
|
||||
ModelIdentity { name: self.model.clone(), digest: self.model_digest.clone(), dimensions: self.dimensions }
|
||||
}
|
||||
}
|
||||
|
||||
/// 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 repos the forge listed as indexable, and whether the list is the whole
|
||||
/// of what the forge holds. A list that is not whole cannot say which repos
|
||||
/// are gone.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Catalog {
|
||||
pub repos: Vec<SourceRepo>,
|
||||
/// Why the list may be missing repos; `None` when it is the forge's count.
|
||||
pub shortfall: Option<String>,
|
||||
}
|
||||
|
||||
/// The forge: which repos exist, and where to fetch them from.
|
||||
pub trait Forge {
|
||||
fn catalog(&mut self) -> Result<Catalog>;
|
||||
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()
|
||||
}
|
||||
|
||||
/// What a refresh reports: the names that failed, and why nothing was retired
|
||||
/// if the catalog could not be trusted to say which repos are gone.
|
||||
#[derive(Debug, Default)]
|
||||
pub struct Refreshed {
|
||||
pub failed: Vec<String>,
|
||||
/// Set when retirement was skipped because the catalog was empty or short.
|
||||
pub retirement_withheld: Option<String>,
|
||||
}
|
||||
|
||||
impl Refreshed {
|
||||
/// Every repo converged and the catalog could be trusted to retire from.
|
||||
pub fn is_clean(&self) -> bool {
|
||||
self.failed.is_empty() && self.retirement_withheld.is_none()
|
||||
}
|
||||
}
|
||||
|
||||
/// 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 described = index.describe_sources(INDEX_CALL_TIMEOUT).context("describing the source collection")?;
|
||||
embedder.built_under(if let DescribedSources::Labelled { meta, .. } = &described { Some(meta.model_identity()) } else { None }.as_ref());
|
||||
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 described {
|
||||
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);
|
||||
let mut refreshed = Refreshed::default();
|
||||
// Retirement reads the catalog as the complete list of what exists. An
|
||||
// empty or short list would retire everything it fails to name, so then
|
||||
// nothing is retired and the run says so.
|
||||
let withheld = match (&catalog.shortfall, catalog.repos.is_empty()) {
|
||||
(Some(why), _) => Some(format!("the forge's catalog is incomplete ({why})")),
|
||||
(None, true) => Some("the forge's catalog lists no indexable repos".to_owned()),
|
||||
(None, false) => None,
|
||||
};
|
||||
match withheld {
|
||||
Some(why) => {
|
||||
eprintln!("huginn-sources: retiring nothing: {why}");
|
||||
refreshed.retirement_withheld = Some(why);
|
||||
}
|
||||
None => retire(index, &catalog.repos, &mirrors, &mut cursors, &mut refreshed.failed),
|
||||
}
|
||||
|
||||
for repo in &catalog.repos {
|
||||
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.
|
||||
/// One repo's failure is logged and named in `failed`; the others go on. Names
|
||||
/// are compared as written: a repo renamed only by case is a new repo to the
|
||||
/// index (its points carry the old spelling), so the old one is retired.
|
||||
fn retire<S: SourceIndex>(
|
||||
index: &mut S,
|
||||
catalog: &[SourceRepo],
|
||||
mirrors: &Path,
|
||||
cursors: &mut BTreeMap<String, String>,
|
||||
failed: &mut Vec<String>,
|
||||
) {
|
||||
let listed = |name: &str| catalog.iter().any(|repo| repo.name == 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 {
|
||||
if let Err(error) = retire_one(index, &name, mirrors, cursors) {
|
||||
eprintln!("huginn-sources: retiring {name} failed, the others go on: {error:#}");
|
||||
failed.push(name);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn retire_one<S: SourceIndex>(index: &mut S, name: &str, mirrors: &Path, cursors: &mut BTreeMap<String, String>) -> Result<()> {
|
||||
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,
|
||||
};
|
||||
// From the first delete on, the collection holds no single commit's chunks
|
||||
// for this repo until the run finishes. Say so before changing anything:
|
||||
// the cursor is dropped, so a run that dies half way leaves a repo the next
|
||||
// run re-ingests whole, whatever head it finds then. A cursor left at the
|
||||
// old commit would let a later head that reverts the files hide the hole.
|
||||
if cursors.remove(&repo.name).is_some() {
|
||||
index.set_cursors(cursors).context("dropping the cursor before the repo's chunks change")?;
|
||||
}
|
||||
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 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,
|
||||
};
|
||||
let identity = embedder.model_identity(budget.left()?).context("reading the embedding model's identity")?;
|
||||
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(),
|
||||
})
|
||||
}
|
||||
546
crates/huginn-daemon/src/sources/tests.rs
Normal file
546
crates/huginn-daemon/src/sources/tests.rs
Normal file
|
|
@ -0,0 +1,546 @@
|
|||
//! 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)>,
|
||||
shortfall: Option<String>,
|
||||
}
|
||||
|
||||
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(),
|
||||
shortfall: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Forge for FakeForge {
|
||||
fn catalog(&mut self) -> Result<Catalog> {
|
||||
Ok(Catalog { repos: self.repos.iter().map(|(repo, _)| repo.clone()).collect(), shortfall: self.shortfall.clone() })
|
||||
}
|
||||
|
||||
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 no cursor (no commit is held whole), 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"), None, "no commit's chunks are held whole, so no cursor says one is");
|
||||
assert_ne!(a, b);
|
||||
|
||||
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());
|
||||
}
|
||||
|
||||
/// A run that fails between a repo's deletes and its upserts leaves chunks of
|
||||
/// no single commit. A head that then reverts the files to the old content
|
||||
/// makes an empty diff from the old commit, which would keep the deleted
|
||||
/// chunks deleted; the repo is ingested whole instead.
|
||||
#[test]
|
||||
fn a_failed_run_then_a_revert_converges() {
|
||||
let remote = Remote::new();
|
||||
let at_a: Vec<(String, Vec<u8>)> = (0..40).map(|n| (format!("f{n}.rs"), format!("file {n}
|
||||
").into_bytes())).collect();
|
||||
let original: Vec<(&str, &[u8])> = at_a.iter().map(|(path, bytes)| (path.as_str(), bytes.as_slice())).collect();
|
||||
remote.commit(&original, &[]);
|
||||
let mut rig = Rig::new();
|
||||
let mut forge = FakeForge::of(&[("Demo", &remote)]);
|
||||
rig.converged(&mut forge);
|
||||
|
||||
let at_b: Vec<(String, Vec<u8>)> = (0..40).map(|n| (format!("f{n}.rs"), format!("file {n} edited
|
||||
").into_bytes())).collect();
|
||||
let edited: Vec<(&str, &[u8])> = at_b.iter().map(|(path, bytes)| (path.as_str(), bytes.as_slice())).collect();
|
||||
remote.commit(&edited, &[]);
|
||||
{
|
||||
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"]);
|
||||
assert!(rig.index.held().len() < 40, "the failed run left a partial repo");
|
||||
|
||||
rig.index.state.lock().unwrap().fail_upserts_after = None;
|
||||
let c = remote.commit(&original, &[]);
|
||||
rig.converged(&mut forge);
|
||||
assert_eq!(rig.index.cursors().get("Demo"), Some(&c));
|
||||
let mut clean = Rig::new();
|
||||
clean.converged(&mut forge);
|
||||
assert_eq!(rig.index.held().len(), 40, "every file of the reverted head is held");
|
||||
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());
|
||||
}
|
||||
|
||||
/// An empty catalog, and a catalog the forge says is short, retire nothing and
|
||||
/// the run says why; the same forge answering whole retires what it dropped.
|
||||
#[test]
|
||||
fn an_empty_or_short_catalog_retires_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)]));
|
||||
let both = rig.index.held();
|
||||
|
||||
let mut empty = FakeForge::of(&[]);
|
||||
let (refreshed, reports) = rig.run(&mut empty);
|
||||
let refreshed = refreshed.unwrap();
|
||||
assert!(reports.is_empty() && refreshed.retirement_withheld.as_deref().unwrap().contains("no indexable repos"));
|
||||
assert_eq!(rig.index.held(), both, "an empty answer deleted nothing");
|
||||
assert!(rig.mirror("Gone").exists() && rig.mirror("Keep").exists());
|
||||
assert_eq!(rig.index.cursors().len(), 2);
|
||||
|
||||
let mut short = FakeForge::of(&[("Keep", &keep)]);
|
||||
short.shortfall = Some("the forge counts 2 repos and 1 were listed".into());
|
||||
let (refreshed, _) = rig.run(&mut short);
|
||||
assert!(refreshed.unwrap().retirement_withheld.unwrap().contains("incomplete"));
|
||||
assert_eq!(rig.index.held(), both, "a short answer deleted nothing");
|
||||
assert!(rig.mirror("Gone").exists());
|
||||
|
||||
let (refreshed, _) = rig.run(&mut FakeForge::of(&[("Keep", &keep)]));
|
||||
assert_eq!(refreshed.unwrap().retirement_withheld, None);
|
||||
assert_eq!(rig.index.held(), [held("Keep", "a.rs", 0, &k)], "a whole answer retires what the forge dropped");
|
||||
}
|
||||
|
||||
/// A retirement that fails is that repo's failure alone: it is named, the
|
||||
/// other retirements and the live repos go on, and the failed one keeps its
|
||||
/// cursor and mirror so the next run finds it again.
|
||||
#[test]
|
||||
fn a_retire_error_skips_that_repo_and_the_run_goes_on() {
|
||||
let (keep, first, second) = (Remote::new(), Remote::new(), Remote::new());
|
||||
let k = keep.commit(&[("a.rs", b"a\n")], &[]);
|
||||
first.commit(&[("b.rs", b"b\n")], &[]);
|
||||
second.commit(&[("c.rs", b"c\n")], &[]);
|
||||
let mut rig = Rig::new();
|
||||
rig.converged(&mut FakeForge::of(&[("Keep", &keep), ("First", &first), ("Second", &second)]));
|
||||
let edited = keep.commit(&[("a.rs", b"a edited\n")], &[]);
|
||||
rig.index.state.lock().unwrap().fail_deletes_of = Some("First".into());
|
||||
|
||||
let (refreshed, reports) = rig.run(&mut FakeForge::of(&[("Keep", &keep)]));
|
||||
assert_eq!(refreshed.unwrap().failed, ["First"]);
|
||||
assert!(!rig.mirror("Second").exists() && rig.mirror("First").exists());
|
||||
assert_eq!(rig.index.cursors().into_keys().collect::<Vec<_>>(), ["First", "Keep"]);
|
||||
assert_eq!(reports.len(), 1, "the live repo was refreshed after the failed retirement");
|
||||
assert_ne!(k, edited);
|
||||
|
||||
rig.index.state.lock().unwrap().fail_deletes_of = None;
|
||||
rig.converged(&mut FakeForge::of(&[("Keep", &keep)]));
|
||||
assert!(!rig.mirror("First").exists());
|
||||
assert_eq!(rig.index.cursors().into_keys().collect::<Vec<_>>(), ["Keep"]);
|
||||
assert_eq!(rig.index.held(), [held("Keep", "a.rs", 0, &edited)]);
|
||||
}
|
||||
|
||||
/// A repo the forge now spells with another case is another repo to the
|
||||
/// index: the points under the old spelling do not stay behind.
|
||||
#[test]
|
||||
fn a_repo_renamed_only_by_case_keeps_no_stale_points() {
|
||||
let remote = Remote::new();
|
||||
remote.commit(&[("a.rs", b"a\n")], &[]);
|
||||
let mut rig = Rig::new();
|
||||
rig.converged(&mut FakeForge::of(&[("Demo", &remote)]));
|
||||
let head = remote.commit(&[("b.rs", b"b\n")], &[]);
|
||||
rig.converged(&mut FakeForge::of(&[("demo", &remote)]));
|
||||
let repos: BTreeSet<String> = rig.index.held().into_iter().map(|(repo, _, _, _)| repo).collect();
|
||||
assert_eq!(repos, BTreeSet::from(["demo".to_owned()]));
|
||||
assert_eq!(rig.index.cursors().into_iter().collect::<Vec<_>>(), [("demo".to_owned(), head)]);
|
||||
}
|
||||
|
||||
/// A run is clean only when no repo failed and retirement was not withheld.
|
||||
#[test]
|
||||
fn a_run_is_clean_only_when_nothing_failed_and_retirement_was_not_withheld() {
|
||||
assert!(Refreshed::default().is_clean());
|
||||
assert!(!Refreshed { failed: vec!["Demo".into()], retirement_withheld: None }.is_clean());
|
||||
assert!(!Refreshed { failed: Vec::new(), retirement_withheld: Some("short".into()) }.is_clean());
|
||||
}
|
||||
|
||||
/// 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"));
|
||||
let recreated = rig.index.state.lock().unwrap().recreated;
|
||||
assert_eq!((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");
|
||||
}
|
||||
|
||||
/// Chunks are embedded in batches of at most `BATCH`, however many a repo has.
|
||||
#[test]
|
||||
fn chunks_are_embedded_in_bounded_batches() {
|
||||
let remote = Remote::new();
|
||||
let many: Vec<(String, Vec<u8>)> = (0..70).map(|n| (format!("f{n}.txt"), format!("file {n}\n").into_bytes())).collect();
|
||||
let files: Vec<(&str, &[u8])> = many.iter().map(|(path, bytes)| (path.as_str(), bytes.as_slice())).collect();
|
||||
remote.commit(&files, &[]);
|
||||
let mut rig = Rig::new();
|
||||
rig.converged(&mut FakeForge::of(&[("Demo", &remote)]));
|
||||
let sizes: Vec<usize> = rig.embedder.state.lock().unwrap().embedded.iter().map(Vec::len).collect();
|
||||
assert_eq!(sizes, [BATCH, BATCH, 6]);
|
||||
assert_eq!(rig.index.held().len(), 70);
|
||||
}
|
||||
|
|
@ -132,3 +132,62 @@ fn serde_json_free_query(vector: &[f32]) -> String {
|
|||
let numbers: Vec<String> = vector.iter().map(|value| value.to_string()).collect();
|
||||
format!("{{\"vector\":[{}],\"limit\":1}}", numbers.join(","))
|
||||
}
|
||||
|
||||
/// One forge repo, however many the forge lists: the live check refreshes a
|
||||
/// single small repo, not the corpus.
|
||||
struct OneRepo {
|
||||
inner: huginn_daemon::sources::forge::HttpForge,
|
||||
name: String,
|
||||
}
|
||||
|
||||
impl huginn_daemon::sources::Forge for OneRepo {
|
||||
fn catalog(&mut self) -> anyhow::Result<huginn_daemon::sources::Catalog> {
|
||||
// One repo of the forge on purpose: this list is a subset, not the forge's count.
|
||||
let repos = self.inner.catalog()?.repos.into_iter().filter(|repo| repo.name == self.name).collect();
|
||||
Ok(huginn_daemon::sources::Catalog { repos, shortfall: None })
|
||||
}
|
||||
|
||||
fn clone_url(&self, name: &str) -> String {
|
||||
self.inner.clone_url(name)
|
||||
}
|
||||
}
|
||||
|
||||
/// Refreshes one small public repo into a scratch source collection, searches
|
||||
/// it, and deletes the collection. Needs `HUGINN_SMOKE_FORGE` (the forge's
|
||||
/// loopback address, so run it on Yggdrasil) and `HUGINN_SMOKE_REPO` beside
|
||||
/// the other three variables. The collection is the scratch one alone;
|
||||
/// `huginn_sources` is never named.
|
||||
#[test]
|
||||
#[ignore = "needs a live Qdrant, Ollama and forge, and writes a scratch collection"]
|
||||
fn a_small_repo_is_refreshed_into_a_scratch_source_collection_and_found() {
|
||||
use huginn_daemon::sources::{Forge, forge::HttpForge, refresh};
|
||||
let (qdrant, ollama, model) = (env("HUGINN_SMOKE_QDRANT"), env("HUGINN_SMOKE_OLLAMA"), env("HUGINN_SMOKE_MODEL"));
|
||||
let (forge, repo) = (env("HUGINN_SMOKE_FORGE"), env("HUGINN_SMOKE_REPO"));
|
||||
let stamp = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_nanos();
|
||||
let scratch = format!("huginn-smoke-sources-{:x}-{stamp:x}", std::process::id());
|
||||
let root = std::env::temp_dir().join(&scratch);
|
||||
|
||||
let mut embedder = OllamaEmbedder::new(&ollama, &model);
|
||||
let mut index = QdrantIndex::new(&qdrant).with_sources_collection(&scratch);
|
||||
let mut forge = OneRepo { inner: HttpForge::new(&forge), name: repo.clone() };
|
||||
assert_eq!(forge.catalog().unwrap().repos.len(), 1, "{repo} is a public, live, non-empty repo on this forge");
|
||||
let mut lines = Vec::new();
|
||||
let refreshed = refresh(&mut embedder, &mut index, &mut forge, &root, &mut |report| lines.push(report.to_string()));
|
||||
let found = refreshed
|
||||
.as_ref()
|
||||
.ok()
|
||||
.and_then(|_| {
|
||||
let query = huginn_mind::SourceQuery {
|
||||
text: huginn_mind::eureka_pipeline::Line("what does this repository do".into()),
|
||||
top_k: 3,
|
||||
repo: None,
|
||||
};
|
||||
huginn_daemon::sources::search(&mut embedder, &mut index, &query, Duration::from_secs(60)).ok()
|
||||
});
|
||||
let _ = raw("DELETE", &format!("{qdrant}/collections/{scratch}?wait=true"), None);
|
||||
let _ = std::fs::remove_dir_all(&root);
|
||||
println!("journal: {lines:?}");
|
||||
assert!(refreshed.unwrap().failed.is_empty());
|
||||
let found = found.expect("a search answers from the scratch collection");
|
||||
assert!(!found.hits.is_empty() && found.hits.iter().all(|hit| hit.repo == repo));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ pub use store::{MindStore, OwnedRedbMessagePackBackingStore};
|
|||
pub use wire::{
|
||||
DeferredAnswer, HuginnMindRequest, HuginnMindResponse, IndexStatus, MAX_DEFERRED_BODY_BYTES, MIND_REQUEST_SCHEMA, MIND_REQUEST_SCHEMA_JSON,
|
||||
MIND_RESPONSE_SCHEMA, MIND_RESPONSE_SCHEMA_JSON, MIND_SERVICE_ID, MindStatus, PersonaStatus,
|
||||
SOURCE_TOP_K_MAX, SourceHit, SourceHits, SourceModel, SourceQuery,
|
||||
};
|
||||
|
||||
/// The leaf, whole, through the one crate that pins its rev: the daemon and
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@
|
|||
|
||||
use cultcache_rs::DatabaseEntry;
|
||||
use cultnet_rs::{CultMeshCdnArtifactManifest, Selection};
|
||||
use eureka_pipeline::{PIPELINE_SCHEMA_EPOCH, PipelineKind, PipelineRef, Slug};
|
||||
use eureka_pipeline::{Line, PIPELINE_SCHEMA_EPOCH, PipelineKind, PipelineRef, Short, Slug};
|
||||
use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
|
@ -88,6 +88,73 @@ pub enum HuginnMindRequest {
|
|||
state: Value,
|
||||
expected_updated_at: Option<String>,
|
||||
},
|
||||
/// Semantic search over the source index: the public forge repos' code and
|
||||
/// docs, which no mind holds. The daemon embeds the text and asks the
|
||||
/// index it does not write.
|
||||
SearchSources { instance: Slug, query: SourceQuery },
|
||||
}
|
||||
|
||||
/// The most hits a source search answers.
|
||||
pub const SOURCE_TOP_K_MAX: u32 = 50;
|
||||
|
||||
/// A source search: the text to find chunks near, how many, and optionally the
|
||||
/// one forge repo (by name) to look in.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
|
||||
pub struct SourceQuery {
|
||||
pub text: Line,
|
||||
pub top_k: u32,
|
||||
pub repo: Option<Short>,
|
||||
}
|
||||
|
||||
impl SourceQuery {
|
||||
/// The door: text that says something, and between one and
|
||||
/// `SOURCE_TOP_K_MAX` hits. A refusal names the field.
|
||||
pub fn check(&self) -> Result<(), MindRefusal> {
|
||||
let refused = |field: &str, value: Option<String>, message: &str| MindRefusal::SelectionInvalid {
|
||||
field: field.into(),
|
||||
value,
|
||||
message: message.into(),
|
||||
};
|
||||
if self.text.0.trim().is_empty() {
|
||||
return Err(refused("text", None, "a source search needs text to find"));
|
||||
}
|
||||
if !(1..=SOURCE_TOP_K_MAX).contains(&self.top_k) {
|
||||
return Err(refused("top_k", Some(self.top_k.to_string()), "a source search asks for 1 to 50 hits"));
|
||||
}
|
||||
if self.repo.as_ref().is_some_and(|repo| repo.0.trim().is_empty()) {
|
||||
return Err(refused("repo", None, "a repo filter names a repo or is absent"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// The model a source collection was built under, and so the one a query must
|
||||
/// be embedded by. The address of the embedder is no part of it.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
|
||||
pub struct SourceModel {
|
||||
pub name: String,
|
||||
pub digest: String,
|
||||
}
|
||||
|
||||
/// One chunk the index holds, as the index holds it: the repo it belongs to,
|
||||
/// the path in that repo, the commit whose blob it was cut from, its 1-based
|
||||
/// inclusive line range, and its text.
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
|
||||
pub struct SourceHit {
|
||||
pub repo: String,
|
||||
pub path: String,
|
||||
pub commit: String,
|
||||
pub line_start: u32,
|
||||
pub line_end: u32,
|
||||
pub score: f32,
|
||||
pub text: String,
|
||||
}
|
||||
|
||||
/// The answer to a source search, nearest first.
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
|
||||
pub struct SourceHits {
|
||||
pub indexed_under: SourceModel,
|
||||
pub hits: Vec<SourceHit>,
|
||||
}
|
||||
|
||||
impl HuginnMindRequest {
|
||||
|
|
@ -100,6 +167,7 @@ impl HuginnMindRequest {
|
|||
Self::Query { .. } => "query",
|
||||
Self::PersonaGet { .. } => "persona_get",
|
||||
Self::PersonaPut { .. } => "persona_put",
|
||||
Self::SearchSources { .. } => "search_sources",
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -111,7 +179,8 @@ impl HuginnMindRequest {
|
|||
Self::View { instance, .. }
|
||||
| Self::Query { instance, .. }
|
||||
| Self::PersonaGet { instance }
|
||||
| Self::PersonaPut { instance, .. } => Some(instance),
|
||||
| Self::PersonaPut { instance, .. }
|
||||
| Self::SearchSources { instance, .. } => Some(instance),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -128,6 +197,7 @@ pub enum HuginnMindResponse {
|
|||
Refused(MindRefusal),
|
||||
Persona(Option<PersonaStateView>),
|
||||
PersonaPut(PersonaPutOutcome),
|
||||
Sources(SourceHits),
|
||||
/// The answer did not fit one send and is being held for the client to
|
||||
/// fetch. Nothing in a mind produces this: the daemon's serve module does,
|
||||
/// in place of an answer, and the answer it stands for is the named
|
||||
|
|
@ -259,6 +329,10 @@ mod tests {
|
|||
state: serde_json::json!({ "updatedAt": "2026-10-10T10:00:00Z", "n": [1, 2.5, null, "x"] }),
|
||||
expected_updated_at: Some("2026-10-10T09:00:00Z".into()),
|
||||
},
|
||||
HuginnMindRequest::SearchSources {
|
||||
instance: slug(INSTANCE),
|
||||
query: SourceQuery { text: Line("reconcile".into()), top_k: 5, repo: Some(Short("Huginn".into())) },
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
|
|
@ -296,6 +370,18 @@ mod tests {
|
|||
})),
|
||||
HuginnMindResponse::PersonaPut(PersonaPutOutcome::Committed { receipt_id: "mind-commit-2".into() }),
|
||||
HuginnMindResponse::PersonaPut(PersonaPutOutcome::Refused(MindRefusal::PersonaNotCanonical)),
|
||||
HuginnMindResponse::Sources(SourceHits {
|
||||
indexed_under: SourceModel { name: "m".into(), digest: "d".into() },
|
||||
hits: vec![SourceHit {
|
||||
repo: "Huginn".into(),
|
||||
path: "a.rs".into(),
|
||||
commit: "abc".into(),
|
||||
line_start: 1,
|
||||
line_end: 2,
|
||||
score: 0.5,
|
||||
text: "x".into(),
|
||||
}],
|
||||
}),
|
||||
HuginnMindResponse::Deferred(DeferredAnswer {
|
||||
manifest: cultnet_rs::pack_content("huginn.mind_response", "package", "", "", "", b"body", 4)
|
||||
.unwrap()
|
||||
|
|
@ -310,7 +396,7 @@ mod tests {
|
|||
#[test]
|
||||
fn the_wire_vocabulary_is_the_minds_methods_and_status_is_derived_from_the_response() {
|
||||
let names: Vec<&str> = requests().iter().map(|request| request.operation()).collect();
|
||||
assert_eq!(names, ["whoami", "admit", "view", "query", "persona_get", "persona_put"]);
|
||||
assert_eq!(names, ["whoami", "admit", "view", "query", "persona_get", "persona_put", "search_sources"]);
|
||||
for request in requests() {
|
||||
let declared = request.instance().is_none();
|
||||
assert_eq!(declared, matches!(request, HuginnMindRequest::Whoami), "{request:?}");
|
||||
|
|
@ -336,6 +422,27 @@ mod tests {
|
|||
}
|
||||
}
|
||||
|
||||
/// A source search is refused before any index is asked when it says
|
||||
/// nothing, asks for no hits or more than the cap, or names a blank repo.
|
||||
#[test]
|
||||
fn a_source_query_without_text_or_within_one_to_fifty_hits_is_refused_by_field() {
|
||||
let query = |text: &str, top_k: u32, repo: Option<&str>| SourceQuery {
|
||||
text: Line(text.into()),
|
||||
top_k,
|
||||
repo: repo.map(|repo| Short(repo.into())),
|
||||
};
|
||||
let field = |query: SourceQuery| match query.check() {
|
||||
Err(MindRefusal::SelectionInvalid { field, .. }) => field,
|
||||
other => panic!("expected a refusal, got {other:?}"),
|
||||
};
|
||||
assert_eq!(query("x", 1, None).check(), Ok(()));
|
||||
assert_eq!(query("x", SOURCE_TOP_K_MAX, Some("Huginn")).check(), Ok(()));
|
||||
assert_eq!(field(query(" ", 5, None)), "text");
|
||||
assert_eq!(field(query("x", 0, None)), "top_k");
|
||||
assert_eq!(field(query("x", SOURCE_TOP_K_MAX + 1, None)), "top_k");
|
||||
assert_eq!(field(query("x", 5, Some(" "))), "repo");
|
||||
}
|
||||
|
||||
/// A new client reads an old daemon's page: the five fields it had, no
|
||||
/// `scores`. This is what keeps rollback by flip safe.
|
||||
#[test]
|
||||
|
|
|
|||
|
|
@ -894,7 +894,7 @@ struct Addressees {
|
|||
}
|
||||
|
||||
fn addressees(config: &Config, mentions: &[String], names: &BTreeMap<&str, &str>) -> Addressees {
|
||||
let to_you = mentions.iter().any(|id| *id == config.own_sender);
|
||||
let to_you = mentions.contains(&config.own_sender);
|
||||
let others = mentions.iter().filter(|id| **id != config.own_sender).map(|id| names.get(id.as_str()).copied().unwrap_or(id.as_str()).to_string()).collect();
|
||||
Addressees { to_you, others }
|
||||
}
|
||||
|
|
|
|||
|
|
@ -92,6 +92,9 @@ pub const SEED_VARIABLE: &str = "MIND_SELF_CONFIG_SEED";
|
|||
const TOKEN_FD_VARIABLE: &str = "CLAUDE_CODE_OAUTH_TOKEN_FILE_DESCRIPTOR";
|
||||
/// Tells `claude` to scrub credentials from the environment of what it starts.
|
||||
const ENV_SCRUB_VARIABLE: &str = "CLAUDE_CODE_SUBPROCESS_ENV_SCRUB";
|
||||
/// Optional, set by the unit: `off` turns that scrub off for a run whose user holds root, where it
|
||||
/// protects nothing and its sandbox's no_new_privs would only stop `sudo`. Anything else keeps it on.
|
||||
const SCRUB_VARIABLE: &str = "MIND_SELF_SCRUB";
|
||||
/// The name the start message tells her to admit as.
|
||||
pub const AGENT: &str = "mind-self-run";
|
||||
/// The namespace of the session id: the run id is the name in it.
|
||||
|
|
@ -674,6 +677,11 @@ fn child_environment(unit: &[(OsString, OsString)]) -> Vec<(&OsString, &OsString
|
|||
unit.iter().filter(|(name, _)| child_may_see(name)).map(|(name, value)| (name, value)).collect()
|
||||
}
|
||||
|
||||
/// Whether the unit turned the child's subprocess scrub off: only `MIND_SELF_SCRUB=off` does.
|
||||
fn scrub_is_off(unit: &[(OsString, OsString)]) -> bool {
|
||||
unit.iter().any(|(name, value)| name == SCRUB_VARIABLE && value == "off")
|
||||
}
|
||||
|
||||
/// Starts `program` with a duplicate of the token descriptor open in it, and
|
||||
/// closes that duplicate here the moment the child exists. The lease's own
|
||||
/// descriptor is close-on-exec, so this duplicate is the child's only way to the
|
||||
|
|
@ -686,7 +694,7 @@ fn spawn_program(program: &Path, args: &[String], token: &TokenFd, unit: &[(OsSt
|
|||
command.args(args).stdin(Stdio::piped()).stdout(Stdio::piped()).stderr(Stdio::inherit());
|
||||
command.env_clear().envs(child_environment(unit));
|
||||
command.env("HOME", &config.0).env(CONFIG_DIR_VARIABLE, &config.0);
|
||||
command.env(TOKEN_FD_VARIABLE, inheritable.number().to_string()).env(ENV_SCRUB_VARIABLE, "1");
|
||||
command.env(TOKEN_FD_VARIABLE, inheritable.number().to_string()).env(ENV_SCRUB_VARIABLE, if scrub_is_off(unit) { "0" } else { "1" });
|
||||
// Its own process group, so that a signal reaches what the child started too.
|
||||
#[cfg(unix)]
|
||||
std::os::unix::process::CommandExt::process_group(&mut command, 0);
|
||||
|
|
|
|||
|
|
@ -156,3 +156,22 @@ fn the_parent_closes_the_duplicate_it_lent_the_child() {
|
|||
child.stop(Stop::Kill);
|
||||
drain(&mut *child);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_child_scrubs_its_subprocesses_unless_the_unit_says_off() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let program = script(&dir, "echo \"scrub=$CLAUDE_CODE_SUBPROCESS_ENV_SCRUB\"");
|
||||
let scrub_with = |value: Option<&str>| {
|
||||
let mut unit = unit_environment();
|
||||
unit.extend(value.map(|value| (SCRUB_VARIABLE.into(), value.into())));
|
||||
let config = FreshConfig::create(dir.path(), None).unwrap();
|
||||
let mut child = spawn_program(&program, &[], &token_holding("tok-canary-7f3a"), &unit, config).unwrap();
|
||||
child.close_input();
|
||||
let (lines, _) = drain(&mut *child);
|
||||
field(&lines, "scrub")
|
||||
};
|
||||
assert_eq!(scrub_with(None), "1", "unset keeps the scrub on");
|
||||
assert_eq!(scrub_with(Some("on")), "1", "any other value keeps it on");
|
||||
assert_eq!(scrub_with(Some("OFF")), "1", "only the exact word turns it off");
|
||||
assert_eq!(scrub_with(Some("off")), "0", "a unit for a user that holds root turns it off");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -132,6 +132,31 @@
|
|||
"required": [
|
||||
"PersonaPut"
|
||||
]
|
||||
},
|
||||
{
|
||||
"description": "Semantic search over the source index: the public forge repos' code and\ndocs, which no mind holds. The daemon embeds the text and asks the\nindex it does not write.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"SearchSources": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"instance": {
|
||||
"$ref": "#/$defs/Slug"
|
||||
},
|
||||
"query": {
|
||||
"$ref": "#/$defs/SourceQuery"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"instance",
|
||||
"query"
|
||||
]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"SearchSources"
|
||||
]
|
||||
}
|
||||
],
|
||||
"$defs": {
|
||||
|
|
@ -1951,6 +1976,34 @@
|
|||
"type": "string",
|
||||
"maxLength": 64
|
||||
},
|
||||
"SourceQuery": {
|
||||
"description": "A source search: the text to find chunks near, how many, and optionally the\none forge repo (by name) to look in.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"repo": {
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/$defs/Short"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
},
|
||||
"text": {
|
||||
"$ref": "#/$defs/Line"
|
||||
},
|
||||
"top_k": {
|
||||
"type": "integer",
|
||||
"format": "uint32",
|
||||
"minimum": 0
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"text",
|
||||
"top_k"
|
||||
]
|
||||
},
|
||||
"StructuralDelta": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
|
|
|
|||
|
|
@ -101,6 +101,18 @@
|
|||
"PersonaPut"
|
||||
]
|
||||
},
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"Sources": {
|
||||
"$ref": "#/$defs/SourceHits"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"required": [
|
||||
"Sources"
|
||||
]
|
||||
},
|
||||
{
|
||||
"description": "The answer did not fit one send and is being held for the client to\nfetch. Nothing in a mind produces this: the daemon's serve module does,\nin place of an answer, and the answer it stands for is the named\nMessagePack of the `HuginnMindResponse` the mind gave.",
|
||||
"type": "object",
|
||||
|
|
@ -4252,6 +4264,82 @@
|
|||
"type": "string",
|
||||
"maxLength": 64
|
||||
},
|
||||
"SourceHit": {
|
||||
"description": "One chunk the index holds, as the index holds it: the repo it belongs to,\nthe path in that repo, the commit whose blob it was cut from, its 1-based\ninclusive line range, and its text.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"commit": {
|
||||
"type": "string"
|
||||
},
|
||||
"line_end": {
|
||||
"type": "integer",
|
||||
"format": "uint32",
|
||||
"minimum": 0
|
||||
},
|
||||
"line_start": {
|
||||
"type": "integer",
|
||||
"format": "uint32",
|
||||
"minimum": 0
|
||||
},
|
||||
"path": {
|
||||
"type": "string"
|
||||
},
|
||||
"repo": {
|
||||
"type": "string"
|
||||
},
|
||||
"score": {
|
||||
"type": "number",
|
||||
"format": "float"
|
||||
},
|
||||
"text": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"repo",
|
||||
"path",
|
||||
"commit",
|
||||
"line_start",
|
||||
"line_end",
|
||||
"score",
|
||||
"text"
|
||||
]
|
||||
},
|
||||
"SourceHits": {
|
||||
"description": "The answer to a source search, nearest first.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"hits": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"$ref": "#/$defs/SourceHit"
|
||||
}
|
||||
},
|
||||
"indexed_under": {
|
||||
"$ref": "#/$defs/SourceModel"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"indexed_under",
|
||||
"hits"
|
||||
]
|
||||
},
|
||||
"SourceModel": {
|
||||
"description": "The model a source collection was built under, and so the one a query must\nbe embedded by. The address of the embedder is no part of it.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"digest": {
|
||||
"type": "string"
|
||||
},
|
||||
"name": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"name",
|
||||
"digest"
|
||||
]
|
||||
},
|
||||
"StructuralDelta": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
|
|
|
|||
Loading…
Reference in a new issue