diff --git a/pk-cli/src/main.rs b/pk-cli/src/main.rs index 64458fd..f1cfea6 100644 --- a/pk-cli/src/main.rs +++ b/pk-cli/src/main.rs @@ -100,6 +100,12 @@ enum Cmd { /// Skip confirmation prompt when using --scope=shared #[arg(long, default_value_t = false)] yes: bool, + /// OKF entry type (e.g. Lesson, Gotcha, Decision). Defaults to Reference. + #[arg(long = "type")] + entry_type: Option, + /// Tag to add to the entry (repeatable), e.g. role:api-dev, vis:team. + #[arg(long = "tag")] + tags: Vec, }, /// Run a lint pass over the full knowledge base Lint { @@ -154,6 +160,9 @@ enum Cmd { max_bytes: usize, #[arg(long, value_enum, default_value = "hook")] format: ContextFormat, + /// Keep only entries carrying this tag (repeatable; all must match). + #[arg(long = "tag")] + tags: Vec, }, /// Publish immutable prompt-snapshot generations from local knowledge stores. Snapshot { @@ -287,6 +296,7 @@ async fn main() -> Result<()> { max_candidates, max_bytes, format, + tags, } = &cli.command { return run_context( @@ -296,6 +306,7 @@ async fn main() -> Result<()> { *max_candidates, *max_bytes, *format, + tags, cli.kb_dir.as_deref(), ) .await; @@ -348,6 +359,8 @@ async fn main() -> Result<()> { source, scope, yes, + entry_type, + tags, } => { // For shared scope, require confirmation unless --yes passed if matches!(scope, KbScope::Shared) && !yes { @@ -374,7 +387,9 @@ async fn main() -> Result<()> { } }; let doc = RawDoc::from_path(source_label, content); - let entry = librarian.compile(doc).await?; + let entry = librarian + .compile_with(doc, pk_librarian::CompileOverrides { entry_type, tags }) + .await?; commit_prompt_snapshot(&kb_dir, scope.snapshot_label(), store.snapshot().await?)?; println!("✓ compiled → {} [{}]", entry.title, entry.id); } @@ -691,6 +706,7 @@ async fn run_context( max_candidates: usize, max_bytes: usize, format: ContextFormat, + required_tags: &[String], explicit_project_kb: Option<&str>, ) -> Result<()> { let scopes = if requested_scopes.is_empty() { @@ -738,8 +754,18 @@ async fn run_context( // Score every entry before any budget applies: truncating first meant // only the first entries in snapshot order were ever considered, and a // failed scope's share of the budget was lost. - scored_count += snapshot.entries.len(); - candidates.extend(snapshot.entries.into_iter().filter_map(|entry| { + // `--tag` filters before scoring: an entry must carry every requested tag. + let entries: Vec<_> = snapshot + .entries + .into_iter() + .filter(|entry| { + required_tags + .iter() + .all(|tag| entry.tags.iter().any(|candidate| candidate == tag)) + }) + .collect(); + scored_count += entries.len(); + candidates.extend(entries.into_iter().filter_map(|entry| { let score = snapshot_score(query, &entry); (score > 0.0 || query.trim().is_empty()).then_some(ContextCandidate { scope, diff --git a/pk-cli/tests/tags_and_type.rs b/pk-cli/tests/tags_and_type.rs new file mode 100644 index 0000000..218c1bf --- /dev/null +++ b/pk-cli/tests/tags_and_type.rs @@ -0,0 +1,218 @@ +//! `pk ingest --type/--tag` classifies entries and `pk context --tag` filters on +//! those tags, so recall can target a role, team or visibility level. +//! +//! The compile step calls a model; the test points pk at a local +//! OpenAI-compatible endpoint that returns a fixed compile result per request, +//! so the real CLI, librarian, store and snapshot path run end to end. + +use serde_json::{json, Value}; +use std::{ + fs, + io::{Read, Write}, + net::TcpListener, + path::Path, + process::{Command, Output}, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, + thread, +}; + +/// Serves `/v1/chat/completions`, answering request N with compile JSON whose +/// title and content are `titles[N]`. +fn fake_model(titles: Vec<&'static str>) -> (String, thread::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let served = Arc::new(AtomicUsize::new(0)); + let handle = thread::spawn(move || { + for stream in listener.incoming() { + let mut stream = stream.unwrap(); + let mut request = Vec::new(); + let mut buffer = [0_u8; 16384]; + loop { + let count = stream.read(&mut buffer).unwrap(); + if count == 0 { + break; + } + request.extend_from_slice(&buffer[..count]); + let text = String::from_utf8_lossy(&request); + if let Some(split) = text.find("\r\n\r\n") { + let length = text[..split] + .lines() + .find_map(|line| { + line.to_ascii_lowercase() + .strip_prefix("content-length:") + .map(|value| value.trim().parse::().unwrap_or(0)) + }) + .unwrap_or(0); + if request.len() >= split + 4 + length { + break; + } + } + } + let index = served.fetch_add(1, Ordering::SeqCst); + let title = titles[index.min(titles.len() - 1)]; + let compiled = json!({ + "title": title, + "content": format!("{title} body with tagfiltertoken"), + "tags": ["model-tag"] + }); + let body = + json!({"choices": [{"message": {"content": compiled.to_string()}}]}).to_string(); + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len(), + body + ); + stream.write_all(response.as_bytes()).unwrap(); + if index + 1 >= titles.len() { + break; + } + } + }); + (format!("http://{address}/v1"), handle) +} + +fn pk(project: &Path, home: &Path, model_url: &str, args: &[&str], stdin: Option<&str>) -> Output { + let mut child = Command::new(env!("CARGO_BIN_EXE_pk")) + .current_dir(project) + .env("HOME", home) + .env("RUST_LOG", "error") + .env("CLOUD_LLM_URL", model_url) + .env("CLOUD_LLM_API_KEY", "test-key") + .args(args) + .stdin(std::process::Stdio::piped()) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .spawn() + .unwrap(); + if let Some(input) = stdin { + child + .stdin + .take() + .unwrap() + .write_all(input.as_bytes()) + .unwrap(); + } else { + drop(child.stdin.take()); + } + let output = child.wait_with_output().unwrap(); + assert!( + output.status.success(), + "pk {args:?} failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + output +} + +fn entry_frontmatter(project: &Path, title_slug_prefix: &str) -> String { + let wiki = project.join(".prometheus/knowledge/wiki"); + let file = fs::read_dir(&wiki) + .unwrap() + .flatten() + .map(|entry| entry.path()) + .find(|path| { + path.file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.starts_with(title_slug_prefix)) + }) + .unwrap_or_else(|| panic!("no wiki entry starting {title_slug_prefix} in {wiki:?}")); + let text = fs::read_to_string(file).unwrap(); + text.split("---").nth(1).unwrap_or_default().to_owned() +} + +#[test] +fn ingest_sets_type_and_tags_and_context_filters_by_tag() { + let fixture = tempfile::tempdir().unwrap(); + let home = fixture.path().join("home"); + let project = fixture.path().join("project"); + fs::create_dir_all(project.join(".git")).unwrap(); + fs::create_dir_all(&home).unwrap(); + let (url, server) = fake_model(vec!["Api role lesson", "Ui role lesson", "Untyped note"]); + + pk( + &project, + &home, + &url, + &["ingest", "--type", "Lesson", "--tag", "role:api-dev"], + Some("api lesson source\n"), + ); + pk( + &project, + &home, + &url, + &[ + "ingest", + "--type", + "Gotcha", + "--tag", + "role:ui-dev", + "--tag", + "vis:team", + ], + Some("ui lesson source\n"), + ); + pk( + &project, + &home, + &url, + &["ingest"], + Some("untyped note source\n"), + ); + server.join().unwrap(); + + let api = entry_frontmatter(&project, "api-role-lesson"); + assert!(api.contains("type: Lesson"), "{api}"); + assert!( + api.contains("role:api-dev") && api.contains("model-tag"), + "{api}" + ); + let ui = entry_frontmatter(&project, "ui-role-lesson"); + assert!(ui.contains("type: Gotcha"), "{ui}"); + assert!( + ui.contains("role:ui-dev") && ui.contains("vis:team"), + "{ui}" + ); + let untyped = entry_frontmatter(&project, "untyped-note"); + assert!( + untyped.contains("type: Reference"), + "default type must stay Reference: {untyped}" + ); + + let context = |tags: &[&str]| -> Value { + let mut args = vec![ + "context", + "tagfiltertoken", + "--scope", + "project", + "--format", + "json", + ]; + for tag in tags { + args.push("--tag"); + args.push(tag); + } + serde_json::from_slice(&pk(&project, &home, &url, &args, None).stdout).unwrap() + }; + let ids = |report: &Value| -> Vec { + report["results"] + .as_array() + .unwrap() + .iter() + .map(|item| item["id"].as_str().unwrap().to_owned()) + .collect() + }; + + let api_only = context(&["role:api-dev"]); + assert_eq!(ids(&api_only), ["api-role-lesson"], "{api_only:#}"); + + let both_tags = context(&["role:ui-dev", "vis:team"]); + assert_eq!(ids(&both_tags), ["ui-role-lesson"], "{both_tags:#}"); + + let no_match = context(&["role:api-dev", "vis:team"]); + assert!(ids(&no_match).is_empty(), "tags are all-of: {no_match:#}"); + + let unfiltered = context(&[]); + assert_eq!(ids(&unfiltered).len(), 3, "{unfiltered:#}"); +} diff --git a/pk-librarian/src/lib.rs b/pk-librarian/src/lib.rs index c87f622..66afed3 100644 --- a/pk-librarian/src/lib.rs +++ b/pk-librarian/src/lib.rs @@ -6,5 +6,5 @@ pub mod prompts; pub mod router; pub use keyword_extract::{extract_query, extract_query_multi_turn}; -pub use librarian::Librarian; +pub use librarian::{CompileOverrides, Librarian}; pub use router::{ModelRoute, ModelRouter}; diff --git a/pk-librarian/src/librarian.rs b/pk-librarian/src/librarian.rs index 0505ea0..1eb304a 100644 --- a/pk-librarian/src/librarian.rs +++ b/pk-librarian/src/librarian.rs @@ -37,6 +37,18 @@ impl Librarian { } pub async fn compile(&self, raw: RawDoc) -> PkResult { + self.compile_with(raw, CompileOverrides::default()).await + } + + /// Compile with caller-supplied classification: `entry_type` replaces the + /// generic OKF `Reference` default, and `tags` are added to the model's + /// tags (de-duplicated, caller tags first). Used by `pk ingest --type/--tag` + /// so lessons can be filtered by role/team/visibility at recall time. + pub async fn compile_with( + &self, + raw: RawDoc, + overrides: CompileOverrides, + ) -> PkResult { info!(source = %raw.source_path, "compiling raw doc"); let related = self.store.related_entries(&raw, 5).await?; @@ -51,6 +63,7 @@ impl Librarian { .map_err(|e| PkError::llm(e.to_string()))?; let mut entry = parse_compile_response(&response)?; + overrides.apply(&mut entry); // Ingest-time duplicate detection: the LLM synthesizes a fresh title // (and therefore a fresh ArticleId) on every call, even over content @@ -374,6 +387,36 @@ fn with_unique_ids(sources: Vec) -> Vec { .collect() } +/// Caller-supplied classification applied to a compiled entry. +#[derive(Debug, Clone, Default)] +pub struct CompileOverrides { + pub entry_type: Option, + pub tags: Vec, +} + +impl CompileOverrides { + fn apply(&self, entry: &mut WikiEntry) { + if let Some(entry_type) = self + .entry_type + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + entry.entry_type = Some(entry_type.to_owned()); + } + if !self.tags.is_empty() { + let mut tags: Vec = Vec::new(); + for tag in self.tags.iter().chain(entry.tags.iter()) { + let tag = tag.trim(); + if !tag.is_empty() && !tags.iter().any(|existing| existing == tag) { + tags.push(tag.to_owned()); + } + } + entry.tags = tags; + } + } +} + fn parse_compile_response(raw: &str) -> PkResult { #[derive(serde::Deserialize)] struct CompileOutput {