Skip to content
This repository was archived by the owner on May 13, 2026. It is now read-only.
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 77 additions & 0 deletions src/llm_registry/prompts/classification.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,58 @@ Return format: {{"sensitivity_level": <0-4>, "data_domain": "<domain>"}}"#
)
}

/// Valid interest categories for field-to-category mapping.
/// This is the single source of truth — used by the LLM prompt, the schema service,
/// and downstream discovery features.
pub const INTEREST_CATEGORIES: &[&str] = &[
"Photography",
"Cooking",
"Running",
"Software Engineering",
"Music",
"Travel",
"Fitness",
"Reading",
"Gaming",
"Finance",
"Gardening",
"Art & Design",
"Parenting",
"Health & Wellness",
"Sports",
"Movies & TV",
"Science",
"Writing",
"Fashion",
"Home Improvement",
"Pets",
"Automotive",
"Productivity",
"Social Media",
"Education",
];

/// Build the interest category classification prompt for a single field.
///
/// The LLM should return a JSON object with `interest_category` (one of the valid
/// categories, or null if the field doesn't map to a user interest).
pub fn build_interest_category_prompt(field_name: &str, description: &str) -> String {
let categories = INTEREST_CATEGORIES.join(", ");
format!(
r#"Classify this database field into a user interest category. Return ONLY a JSON object with one field, no explanation.

Field name: "{field_name}"
Description: "{description}"

Valid interest categories: {categories}

If this field clearly relates to one of the above interests, return that category.
If this field is a structural/metadata field (like id, hash, timestamp, source, content_hash) or doesn't map to any interest, return null.

Return format: {{"interest_category": "<category>" | null}}"#
)
}

#[cfg(test)]
mod tests {
use super::*;
Expand Down Expand Up @@ -60,4 +112,29 @@ mod tests {
assert!(prompt.contains(domain));
}
}

#[test]
fn interest_prompt_contains_field_name_and_description() {
let prompt =
build_interest_category_prompt("photo_album", "the album containing the photo");
assert!(prompt.contains("photo_album"));
assert!(prompt.contains("the album containing the photo"));
assert!(prompt.contains("interest_category"));
}

#[test]
fn interest_prompt_lists_all_categories() {
let prompt = build_interest_category_prompt("x", "y");
for category in INTEREST_CATEGORIES {
assert!(prompt.contains(category), "Missing category: {}", category);
}
}

#[test]
fn interest_categories_are_non_empty() {
assert!(!INTEREST_CATEGORIES.is_empty());
for cat in INTEREST_CATEGORIES {
assert!(!cat.is_empty());
}
}
}
12 changes: 11 additions & 1 deletion src/schema/types/declarative_schemas.rs
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,8 @@ impl<'de> serde::Deserialize<'de> for DeclarativeSchemaDefinition {
#[serde(default)]
field_data_classifications: HashMap<String, DataClassification>,
#[serde(default)]
field_interest_categories: HashMap<String, String>,
#[serde(default)]
ref_fields: HashMap<String, String>,
#[serde(default)]
field_types: HashMap<String, crate::schema::types::field_value_type::FieldValueType>,
Expand Down Expand Up @@ -207,9 +209,10 @@ impl<'de> serde::Deserialize<'de> for DeclarativeSchemaDefinition {
.insert(field_name, classifications);
}

// Preserve field_descriptions, field_data_classifications, ref_fields, field_types and identity_hash
// Preserve field_descriptions, field_data_classifications, field_interest_categories, ref_fields, field_types and identity_hash
schema.field_descriptions = helper.field_descriptions;
schema.field_data_classifications = helper.field_data_classifications;
schema.field_interest_categories = helper.field_interest_categories;
schema.ref_fields = helper.ref_fields;
schema.field_types = helper.field_types;
schema.identity_hash = helper.identity_hash;
Expand Down Expand Up @@ -267,6 +270,11 @@ pub struct DeclarativeSchemaDefinition {
/// Maps field_name -> DataClassification. Required for new fields at schema creation.
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub field_data_classifications: HashMap<String, DataClassification>,
/// Interest categories for each field (e.g. "Photography", "Cooking", "Running").
/// Assigned by the schema service from the canonical field registry.
/// Maps field_name -> interest category string.
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub field_interest_categories: HashMap<String, String>,
/// Reference fields that point to child schemas
/// Maps field_name -> child_schema_name
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
Expand Down Expand Up @@ -324,6 +332,7 @@ impl PartialEq for DeclarativeSchemaDefinition {
&& self.field_molecule_uuids == other.field_molecule_uuids
&& self.field_classifications == other.field_classifications
&& self.field_data_classifications == other.field_data_classifications
&& self.field_interest_categories == other.field_interest_categories
&& self.ref_fields == other.ref_fields
&& self.field_types == other.field_types
&& self.identity_hash == other.identity_hash
Expand Down Expand Up @@ -463,6 +472,7 @@ impl DeclarativeSchemaDefinition {
field_classifications: HashMap::new(),
field_descriptions: HashMap::new(),
field_data_classifications: HashMap::new(),
field_interest_categories: HashMap::new(),
ref_fields: HashMap::new(),
field_types: HashMap::new(),
identity_hash: None,
Expand Down
173 changes: 168 additions & 5 deletions src/schema_service/classify.rs
Original file line number Diff line number Diff line change
@@ -1,16 +1,18 @@
//! Field sensitivity classification inference.
//! Field sensitivity classification and interest category inference.
//!
//! Determines the (sensitivity_level, data_domain) for new canonical fields
//! based on the field's description. The schema service is the sole authority
//! on data classification.
//! Determines the (sensitivity_level, data_domain) and interest_category for
//! new canonical fields based on the field's description. The schema service
//! is the sole authority on both data classification and interest categories.
//!
//! Strategy for new fields without an existing canonical match:
//! 1. Caller-provided classification → use it
//! 2. LLM call using field description (requires ANTHROPIC_API_KEY)
//! 3. No fallback — returns error. Incorrect classification is worse than no schema.

use crate::llm_registry::models;
use crate::llm_registry::prompts::classification::build_classification_prompt;
use crate::llm_registry::prompts::classification::{
build_classification_prompt, build_interest_category_prompt, INTEREST_CATEGORIES,
};
use crate::schema::types::data_classification::DataClassification;

/// Classify a field using LLM analysis of its description.
Expand Down Expand Up @@ -121,6 +123,151 @@ pub async fn classify_with_llm(
Ok(classification)
}

/// Classify a field's interest category using LLM analysis of its description.
/// Returns `Ok(None)` if the field doesn't map to any interest category (structural fields).
/// Returns `Err` only on LLM communication failures.
pub async fn classify_interest_category_with_llm(
field_name: &str,
description: &str,
) -> Result<Option<String>, String> {
let api_key = std::env::var("ANTHROPIC_API_KEY").map_err(|_| {
"Schema service cannot classify interest categories: ANTHROPIC_API_KEY not set.".to_string()
})?;
if api_key.trim().is_empty() {
return Err(
"Schema service cannot classify interest categories: ANTHROPIC_API_KEY is empty"
.to_string(),
);
}

let prompt = build_interest_category_prompt(field_name, description);

let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(
models::TIMEOUT_CLASSIFICATION,
))
.no_proxy()
.build()
.map_err(|e| {
format!(
"Failed to create HTTP client for interest classification: {}",
e
)
})?;

let request_body = serde_json::json!({
"model": models::ANTHROPIC_HAIKU,
"messages": [{"role": "user", "content": prompt}],
"max_tokens": models::MAX_TOKENS_CLASSIFICATION,
"temperature": models::TEMPERATURE_DETERMINISTIC
});

let response = client
.post(format!("{}/v1/messages", models::ANTHROPIC_API_URL))
.header("x-api-key", &api_key)
.header("anthropic-version", models::ANTHROPIC_API_VERSION)
.header("Content-Type", "application/json")
.json(&request_body)
.send()
.await
.map_err(|e| {
format!(
"Interest category LLM call failed for field '{}': {}",
field_name, e
)
})?;

if !response.status().is_success() {
return Err(format!(
"Interest category LLM call returned status {} for field '{}'",
response.status(),
field_name
));
}

let resp: serde_json::Value = response.json().await.map_err(|e| {
format!(
"Failed to parse LLM response for interest category of field '{}': {}",
field_name, e
)
})?;

let text = resp
.get("content")
.and_then(|c| c.as_array())
.and_then(|a| a.first())
.and_then(|c| c.get("text"))
.and_then(|t| t.as_str())
.ok_or_else(|| {
format!(
"LLM response missing content text for interest category of field '{}'",
field_name
)
})?;

// Parse the JSON response — try raw text first, then extract from markdown fence
let parsed: serde_json::Value = serde_json::from_str(text)
.or_else(|_| {
let trimmed = text.trim();
let json_str = trimmed
.strip_prefix("```json")
.or_else(|| trimmed.strip_prefix("```"))
.and_then(|s| s.strip_suffix("```"))
.unwrap_or(trimmed)
.trim();
serde_json::from_str(json_str)
})
.map_err(|e| {
format!(
"Failed to parse LLM interest category for field '{}': {} (raw: {})",
field_name, e, text
)
})?;

let category = parsed
.get("interest_category")
.and_then(|v| v.as_str())
.map(|s| s.to_string());

// Validate against known categories
let validated = category.filter(|cat| {
INTEREST_CATEGORIES
.iter()
.any(|valid| valid.eq_ignore_ascii_case(cat))
});

if let Some(ref cat) = validated {
crate::log_feature!(
crate::logging::features::LogFeature::Schema,
info,
"LLM classified field '{}' interest category as '{}'",
field_name,
cat
);
}

Ok(validated)
}

/// Infer interest category for a new canonical field.
/// Returns `Ok(None)` for structural fields or when the API key is missing.
/// Interest category is best-effort — missing it doesn't block schema creation.
pub async fn infer_interest_category(field_name: &str, description: &str) -> Option<String> {
match classify_interest_category_with_llm(field_name, description).await {
Ok(category) => category,
Err(e) => {
crate::log_feature!(
crate::logging::features::LogFeature::Schema,
warn,
"Interest category classification failed for field '{}': {} (non-blocking)",
field_name,
e
);
None
}
}
}

/// Infer classification for a new canonical field.
/// Returns an error if classification cannot be determined — no silent fallbacks.
///
Expand Down Expand Up @@ -171,4 +318,20 @@ mod tests {
}
}
}

#[tokio::test]
async fn infer_interest_category_returns_none_without_api_key() {
// Without ANTHROPIC_API_KEY, should return None (non-blocking)
let result = infer_interest_category("photo_album", "the album containing the photo").await;
// Either returns a valid category (if API key is set) or None
if let Some(ref cat) = result {
assert!(
INTEREST_CATEGORIES
.iter()
.any(|valid| valid.eq_ignore_ascii_case(cat)),
"Invalid category: {}",
cat
);
}
}
}
Loading
Loading