use promptforge_tool_picker::{Catalog as Descriptors, Config, ToolDescriptor, ToolId, ToolPicker};
use serde_json::json;
use crate::catalog::Catalog;
use crate::retrieval::{Candidate, Shortlist};
const SERVER: &str = "promptforge";
const SIMILARITY_FLOOR: f32 = 0.0;
#[derive(Debug)]
pub(crate) struct Index {
picker: ToolPicker,
}
impl Index {
pub(super) fn build(catalog: &Catalog) -> Option<Index> {
let config = compiled_config()?;
match ToolPicker::build(descriptors(catalog), config) {
Ok(picker) => Some(Index { picker }),
Err(error) => {
tracing::error!("need_prompt cannot answer: retrieval index build failed: {error}");
None
}
}
}
#[cfg(test)]
pub(super) fn build_with(
model: &promptforge_tool_picker::Model,
catalog: &Catalog,
) -> Option<Index> {
let config = compiled_config()?;
match ToolPicker::build_with_model(model, descriptors(catalog), config) {
Ok(picker) => Some(Index { picker }),
Err(error) => {
tracing::error!("the shared retrieval model could not index the catalog: {error}");
None
}
}
}
pub(super) fn rebuild(&self, catalog: &Catalog) -> Option<Index> {
match self.picker.rebuild(descriptors(catalog)) {
Ok(picker) => Some(Index { picker }),
Err(error) => {
tracing::warn!("need_prompt keeps its previous index: {error}");
None
}
}
}
pub(super) fn len(&self) -> usize {
self.picker.len()
}
pub(super) fn shortlist(&self, capability: &str, k: usize) -> Shortlist {
match self.picker.shortlist(capability, k) {
Ok(tools) => Shortlist::Candidates(tools.iter().map(candidate).collect()),
Err(error) => Shortlist::Failed(error.to_string()),
}
}
}
fn candidate(tool: &ToolDescriptor) -> Candidate {
Candidate {
name: tool.name().to_owned(),
description: tool.description().to_owned(),
}
}
pub(crate) fn descriptors(catalog: &Catalog) -> Descriptors {
catalog
.entries()
.iter()
.filter(|entry| entry.problem().is_none())
.map(|entry| {
ToolDescriptor::new(
ToolId::new(SERVER, entry.name()),
entry.description(),
args_schema(),
)
})
.collect()
}
fn args_schema() -> serde_json::Value {
json!({"type": "object", "properties": {"args": {"type": "string"}}})
}
fn compiled_config() -> Option<Config> {
match Config::default().with_similarity_floor(SIMILARITY_FLOOR) {
Ok(config) => Some(config),
Err(error) => {
tracing::error!(
"need_prompt cannot answer: the compiled-in retrieval policy is invalid: {error}"
);
None
}
}
}