use crate::projector::{slug, CorpusKind, Projector, Situation};
use serde::{Deserialize, Serialize};
use std::collections::BTreeSet;
use std::path::PathBuf;
#[derive(Debug, Clone, Serialize)]
pub struct ColumnProfile {
pub name: String, pub raw_name: String,
pub kind: String, pub distinct: usize,
pub distinct_ratio: f64,
pub nulls: usize,
pub numeric_min: Option<f64>,
pub numeric_max: Option<f64>,
pub candidate_key: bool,
pub samples: Vec<String>,
}
#[derive(Debug, Clone, Serialize)]
pub struct Profile {
pub rows: usize,
pub columns: Vec<ColumnProfile>,
}
fn parse_num(s: &str) -> Option<f64> {
let t = s.trim().replace([',', '$', '%'], "");
if t.is_empty() {
return None;
}
t.parse::<f64>().ok()
}
fn is_bool(s: &str) -> bool {
matches!(s.trim().to_lowercase().as_str(), "true" | "false" | "yes" | "no" | "y" | "n")
}
pub fn profile_csv(path: &std::path::Path) -> std::io::Result<Profile> {
let mut rdr = csv::ReaderBuilder::new().flexible(true).from_path(path)?;
let headers: Vec<String> = rdr.headers()?.iter().map(|h| h.to_string()).collect();
let n = headers.len();
let mut distinct: Vec<BTreeSet<String>> = vec![BTreeSet::new(); n];
let mut nulls = vec![0usize; n];
let mut num_ok = vec![0usize; n];
let mut bool_ok = vec![0usize; n];
let mut nonnull = vec![0usize; n];
let mut mn = vec![f64::INFINITY; n];
let mut mx = vec![f64::NEG_INFINITY; n];
let mut rows = 0usize;
for rec in rdr.records() {
let rec = match rec {
Ok(r) => r,
Err(_) => continue,
};
rows += 1;
for i in 0..n {
let cell = rec.get(i).unwrap_or("").trim();
if cell.is_empty() {
nulls[i] += 1;
continue;
}
nonnull[i] += 1;
if distinct[i].len() < 10_000 {
distinct[i].insert(cell.to_string());
}
if let Some(v) = parse_num(cell) {
num_ok[i] += 1;
mn[i] = mn[i].min(v);
mx[i] = mx[i].max(v);
}
if is_bool(cell) {
bool_ok[i] += 1;
}
}
}
let columns = (0..n)
.map(|i| {
let nn = nonnull[i].max(1);
let distinct_ct = distinct[i].len();
let ratio = distinct_ct as f64 / rows.max(1) as f64;
let numeric = num_ok[i] as f64 / nn as f64 > 0.9 && distinct_ct > 1;
let boolean = bool_ok[i] as f64 / nn as f64 > 0.9;
let kind = if boolean {
"boolean"
} else if numeric {
"numeric"
} else if ratio < 0.5 || distinct_ct <= 50 {
"categorical"
} else {
"text"
};
ColumnProfile {
name: slug(&headers[i]),
raw_name: headers[i].clone(),
kind: kind.to_string(),
distinct: distinct_ct,
distinct_ratio: (ratio * 1000.0).round() / 1000.0,
nulls: nulls[i],
numeric_min: numeric.then_some(mn[i]),
numeric_max: numeric.then_some(mx[i]),
candidate_key: nulls[i] == 0 && distinct_ct == rows && rows > 0,
samples: distinct[i].iter().take(8).cloned().collect(),
}
})
.collect();
Ok(Profile { rows, columns })
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FacetCol {
pub column: String,
#[serde(default)]
pub facet: Option<String>,
}
fn default_bins() -> usize {
5
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MeasureCol {
pub column: String,
#[serde(default)]
pub facet: Option<String>,
#[serde(default = "default_bins")]
pub bins: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RelationSpec {
pub name: String,
pub head: String,
pub tail: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProjectionSpec {
#[serde(default = "row_str")]
pub situation: String,
#[serde(default)]
pub facets: Vec<FacetCol>,
#[serde(default)]
pub measures: Vec<MeasureCol>,
#[serde(default)]
pub relations: Vec<RelationSpec>,
#[serde(default)]
pub notes: String,
}
fn row_str() -> String {
"row".to_string()
}
impl ProjectionSpec {
pub fn default_for(profile: &Profile) -> ProjectionSpec {
let mut facets = Vec::new();
let mut measures = Vec::new();
for c in &profile.columns {
if c.candidate_key {
continue;
}
match c.kind.as_str() {
"categorical" | "boolean" => facets.push(FacetCol { column: c.name.clone(), facet: None }),
"numeric" => measures.push(MeasureCol { column: c.name.clone(), facet: None, bins: 5 }),
_ => {}
}
}
ProjectionSpec { situation: "row".into(), facets, measures, relations: Vec::new(), notes: "deterministic default (no LLM)".into() }
}
pub fn validated(mut self, profile: &Profile) -> ProjectionSpec {
let known: BTreeSet<&str> = profile.columns.iter().map(|c| c.name.as_str()).collect();
let known_slug = |s: &str| known.contains(slug(s).as_str());
self.facets.retain(|f| known_slug(&f.column));
self.measures.retain(|m| known_slug(&m.column));
self.relations.retain(|r| known_slug(&r.head) && known_slug(&r.tail));
if self.facets.is_empty() && self.measures.is_empty() {
return ProjectionSpec::default_for(profile);
}
self
}
}
pub struct SpecProjector {
path: PathBuf,
spec: ProjectionSpec,
header_slugs: Vec<String>,
raw_headers: Vec<String>,
edges: std::collections::HashMap<String, Vec<f64>>,
}
impl SpecProjector {
pub fn open(path: impl Into<PathBuf>, spec: ProjectionSpec) -> std::io::Result<SpecProjector> {
let path = path.into();
let mut rdr = csv::ReaderBuilder::new().flexible(true).from_path(&path)?;
let raw_headers: Vec<String> = rdr.headers()?.iter().map(|h| h.to_string()).collect();
let header_slugs: Vec<String> = raw_headers.iter().map(|h| slug(h)).collect();
let idx_of = |col: &str| header_slugs.iter().position(|h| h == &slug(col));
let mut vals: std::collections::HashMap<String, Vec<f64>> = std::collections::HashMap::new();
for m in &spec.measures {
if idx_of(&m.column).is_some() {
vals.insert(slug(&m.column), Vec::new());
}
}
if !vals.is_empty() {
let mut rdr2 = csv::ReaderBuilder::new().flexible(true).from_path(&path)?;
for rec in rdr2.records().flatten() {
for m in &spec.measures {
if let Some(i) = idx_of(&m.column) {
if let Some(v) = rec.get(i).and_then(parse_num) {
vals.get_mut(&slug(&m.column)).unwrap().push(v);
}
}
}
}
}
let edges = spec
.measures
.iter()
.filter_map(|m| {
let key = slug(&m.column);
vals.get(&key).map(|v| (key, quantile_edges(v, m.bins.clamp(2, 12))))
})
.collect();
Ok(SpecProjector { path, spec, header_slugs, raw_headers, edges })
}
}
fn quantile_edges(vals: &[f64], bins: usize) -> Vec<f64> {
let mut v: Vec<f64> = vals.iter().copied().filter(|x| x.is_finite()).collect();
v.sort_by(|a, b| a.partial_cmp(b).unwrap());
v.dedup();
if v.len() < bins {
return Vec::new();
}
let mut edges = vec![v[0]];
edges.extend((1..bins).map(|k| v[(k * v.len()) / bins]));
edges.push(v[v.len() - 1]);
edges
}
fn fmt_num(x: f64) -> String {
if x.fract().abs() < 1e-9 {
format!("{}", x as i64)
} else {
format!("{x:.2}")
}
}
fn bucket_label(value: f64, edges: &[f64]) -> String {
if edges.len() < 2 {
return fmt_num(value);
}
for w in edges.windows(2) {
if value < w[1] || (w[1] - edges[edges.len() - 1]).abs() < f64::EPSILON {
return format!("{} to {}", fmt_num(w[0]), fmt_num(w[1]));
}
}
format!("{} to {}", fmt_num(edges[edges.len() - 2]), fmt_num(edges[edges.len() - 1]))
}
impl Projector for SpecProjector {
fn columns(&self) -> Vec<String> {
self.raw_headers.clone()
}
fn kind(&self) -> CorpusKind {
CorpusKind::Csv
}
fn source(&self) -> String {
self.path.display().to_string()
}
fn project(self: Box<Self>, sink: &mut dyn FnMut(Situation)) -> std::io::Result<()> {
let idx_of = |col: &str| self.header_slugs.iter().position(|h| h == &slug(col));
let facet_name = |col: &str, given: &Option<String>| given.clone().map(|f| slug(&f)).unwrap_or_else(|| slug(col));
let mut rdr = csv::ReaderBuilder::new().flexible(true).from_path(&self.path)?;
for rec in rdr.records().flatten() {
let cells: Vec<String> = rec.iter().map(|c| c.to_string()).collect();
let mut tokens: Vec<String> = Vec::new();
for f in &self.spec.facets {
if let Some(i) = idx_of(&f.column) {
let v = cells.get(i).map(|s| s.trim()).unwrap_or("");
if !v.is_empty() {
tokens.push(format!("{}/{}", facet_name(&f.column, &f.facet), slug(v)));
}
}
}
for m in &self.spec.measures {
if let Some(i) = idx_of(&m.column) {
if let Some(v) = cells.get(i).and_then(|s| parse_num(s)) {
let edges = self.edges.get(&slug(&m.column)).map(|e| e.as_slice()).unwrap_or(&[]);
let label = bucket_label(v, edges);
tokens.push(format!("{}/{}", facet_name(&m.column, &m.facet), slug(&label)));
}
}
}
for r in &self.spec.relations {
if let (Some(hi), Some(ti)) = (idx_of(&r.head), idx_of(&r.tail)) {
let (hv, tv) = (cells.get(hi).map(|s| s.trim()).unwrap_or(""), cells.get(ti).map(|s| s.trim()).unwrap_or(""));
if !hv.is_empty() && !tv.is_empty() {
tokens.push(format!("rel/{}/+/{}/{}", slug(&r.name), slug(&r.head), slug(hv)));
tokens.push(format!("rel/{}/-/{}/{}", slug(&r.name), slug(&r.tail), slug(tv)));
}
}
}
sink(Situation::new(tokens, cells));
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::projector::Projector;
fn tmp_csv(name: &str, body: &str) -> PathBuf {
let p = std::env::temp_dir().join(format!("steeldb_discover_{}_{name}.csv", std::process::id()));
std::fs::write(&p, body).unwrap();
p
}
#[test]
fn profiles_and_projects() {
let p = tmp_csv(
"proj",
"id,make,year,price\n1,toyota,2020,25000\n2,honda,2021,30000\n3,toyota,2020,41000\n4,kia,2019,18000\n5,honda,2021,22000\n6,toyota,2019,35000\n",
);
let profile = profile_csv(&p).unwrap();
assert_eq!(profile.rows, 6);
let by = |n: &str| profile.columns.iter().find(|c| c.name == n).unwrap();
assert!(by("id").candidate_key);
assert_eq!(by("make").kind, "categorical");
assert_eq!(by("year").kind, "numeric");
let spec = ProjectionSpec::default_for(&profile);
assert!(spec.facets.iter().any(|f| f.column == "make"));
assert!(spec.measures.iter().any(|m| m.column == "year"));
assert!(!spec.facets.iter().any(|f| f.column == "id"));
let mut situations = Vec::new();
Box::new(SpecProjector::open(&p, spec).unwrap()).project(&mut |s| situations.push(s)).unwrap();
assert_eq!(situations.len(), 6);
let all: Vec<String> = situations.iter().flat_map(|s| s.tokens.clone()).collect();
assert!(all.iter().any(|t| t == "make/toyota"));
assert!(all.iter().any(|t| t.starts_with("year/")));
}
#[test]
fn spec_validation_drops_unknown_columns() {
let p = tmp_csv("valid", "a,b\nx,1\ny,2\n");
let profile = profile_csv(&p).unwrap();
let spec = ProjectionSpec {
situation: "row".into(),
facets: vec![FacetCol { column: "nonexistent".into(), facet: None }],
measures: vec![],
relations: vec![],
notes: String::new(),
}
.validated(&profile);
assert!(spec.facets.iter().all(|f| f.column != "nonexistent"));
}
}
#[cfg(feature = "agent")]
pub use propose::propose_spec;
#[cfg(feature = "agent")]
mod propose {
use super::{Profile, ProjectionSpec};
use crate::agent::provider::LlmProvider;
use crate::agent::types::Msg;
const SPEC_SYS: &str = "\
You are a data engineer designing how to project a table into a queryable hypergraph. You are given a \
column profile. Decide, per column, whether it is a FACET (a categorical dimension to filter/group by), \
a MEASURE (a numeric quantity to bucketize into ranges), or ignore it (free-text or an id/key). Also \
propose useful RELATIONS between two columns when one clearly acts on another. Do FEATURE ENGINEERING: \
prefer low-cardinality categoricals as facets; treat continuous numerics as measures; drop high- \
cardinality identifiers. Respond with ONLY a JSON object, no prose:
{\"situation\":\"row\",\"facets\":[{\"column\":\"<name>\"}],\"measures\":[{\"column\":\"<name>\",\"bins\":5}],\"relations\":[{\"name\":\"<verb>\",\"head\":\"<col>\",\"tail\":\"<col>\"}],\"notes\":\"one line\"}";
pub async fn propose_spec(provider: &dyn LlmProvider, profile: &Profile) -> ProjectionSpec {
let profile_json = serde_json::to_string_pretty(profile).unwrap_or_default();
let user = format!("Column profile ({} rows):\n{profile_json}\n\nDesign the projection spec. /no_think", profile.rows);
let turn = match provider.chat(SPEC_SYS, &[Msg::user_text(user)], &[]).await {
Ok(t) => t,
Err(_) => return ProjectionSpec::default_for(profile),
};
match extract_json(&turn.text).and_then(|j| serde_json::from_value::<ProjectionSpec>(j).ok()) {
Some(spec) => spec.validated(profile),
None => ProjectionSpec::default_for(profile),
}
}
fn extract_json(text: &str) -> Option<serde_json::Value> {
let dethunk = strip_think(text);
let cleaned = dethunk.replace("```json", "```");
let body = match cleaned.split_once("```") {
Some((_, rest)) => rest.split_once("```").map(|(b, _)| b).unwrap_or(rest).to_string(),
None => cleaned,
};
let (a, b) = (body.find('{')?, body.rfind('}')?);
if b <= a {
return None;
}
serde_json::from_str(&body[a..=b]).ok()
}
fn strip_think(text: &str) -> String {
let mut s = text.to_string();
for (open, close) in [("<think>", "</think>"), ("<thinking>", "</thinking>")] {
while let (Some(a), Some(b)) = (s.find(open), s.find(close)) {
if b > a {
s.replace_range(a..b + close.len(), "");
} else {
break;
}
}
if let Some(b) = s.find(close) {
s = s[b + close.len()..].to_string();
}
}
s
}
}