use saya_agent::OverrideFindingDto;
use serde::{Deserialize, Serialize};
use super::turn_table::TurnObjectTable;
use crate::agent::tools::DrainedObservations;
use crate::connection::ConnectionRegistry;
use crate::contracts::{RecallReceipt, SuppliedContract};
#[allow(dead_code)]
pub const MAX_PROMPT_BYTES: usize = 4096;
#[allow(dead_code)]
pub const MAX_ANSWER_BYTES: usize = 8192;
#[allow(dead_code)]
pub const MAX_TURN_RECORD_BYTES: usize = 16_384;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[allow(dead_code)]
pub struct SuppliedClaimDto {
pub claim_id: String,
pub kind: String,
pub value: String,
pub column: Option<String>,
pub status: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[allow(dead_code)]
pub struct SuppliedContractDto {
pub profile: String,
pub object: String,
pub schema_state: String,
pub claims: Vec<SuppliedClaimDto>,
}
impl From<&SuppliedContract> for SuppliedContractDto {
fn from(c: &SuppliedContract) -> Self {
Self {
profile: c.profile.clone(),
object: c.object.clone(),
schema_state: c.schema_state.to_string(),
claims: c
.claims
.iter()
.map(|cl| SuppliedClaimDto {
claim_id: cl.claim_id.as_str().to_string(),
kind: cl.kind.to_string(),
value: cl.value.clone(),
column: cl.column.clone(),
status: cl.status.as_str().to_string(),
})
.collect(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[allow(dead_code)]
pub struct TurnRecord {
pub prompt: String,
pub assistant_answer: String,
pub object_table: TurnObjectTable,
pub user_corrections: Vec<String>,
pub override_findings: Vec<OverrideFindingDto>,
pub supplied_claims: Vec<SuppliedContractDto>,
}
impl TurnRecord {
#[allow(dead_code)]
pub fn assemble(
prompt: &str,
assistant_answer: &str,
registry: &ConnectionRegistry,
observations: &DrainedObservations,
receipt: Option<&RecallReceipt>,
overrides: &[OverrideFindingDto],
) -> Self {
let bounded_prompt = truncate_utf8(prompt, MAX_PROMPT_BYTES);
let bounded_answer = truncate_utf8(assistant_answer, MAX_ANSWER_BYTES);
let primary_profile = registry.primary_name();
let mut object_table = TurnObjectTable::new();
for obs in &observations.observations {
let prof = obs
.profile
.as_ref()
.and_then(|p| registry.name_for_identity(p.as_str()))
.unwrap_or(primary_profile);
for obj_parts in &obs.objects {
let qualified = obj_parts.join(".");
object_table.register(prof, &qualified, &obs.columns);
}
}
let mut supplied_dtos = Vec::new();
if let Some(r) = receipt {
for contract in &r.supplied {
supplied_dtos.push(SuppliedContractDto::from(contract));
let cols: Vec<String> = contract
.claims
.iter()
.filter_map(|c| c.column.clone())
.collect();
object_table.register(&contract.profile, &contract.object, &cols);
}
}
let user_corrections = extract_user_corrections(&bounded_prompt);
Self {
prompt: bounded_prompt,
assistant_answer: bounded_answer,
object_table,
user_corrections,
override_findings: overrides.to_vec(),
supplied_claims: supplied_dtos,
}
}
}
#[allow(dead_code)]
fn truncate_utf8(s: &str, max_bytes: usize) -> String {
if s.len() <= max_bytes {
return s.to_string();
}
let mut boundary = max_bytes;
while boundary > 0 && !s.is_char_boundary(boundary) {
boundary -= 1;
}
s[..boundary].to_string()
}
#[allow(dead_code)]
fn extract_user_corrections(prompt: &str) -> Vec<String> {
prompt
.lines()
.flat_map(|line| line.split(['.', '!', '?']))
.map(str::trim)
.filter(|s| {
let lower = s.to_lowercase();
lower.starts_with("no")
|| lower.starts_with("actually")
|| lower.starts_with("instead")
|| lower.starts_with("remember")
|| lower.contains("should be")
})
.map(ToString::to_string)
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
use saya_types::ClaimId;
#[test]
fn test_turn_record_bounds_input_strings() {
let huge_prompt = "p".repeat(50_000);
let huge_answer = "a".repeat(50_000);
let obs = DrainedObservations {
observations: Vec::new(),
truncated: false,
};
let registry = ConnectionRegistry::new("primary");
let record = TurnRecord::assemble(&huge_prompt, &huge_answer, ®istry, &obs, None, &[]);
assert_eq!(record.prompt.len(), MAX_PROMPT_BYTES);
assert_eq!(record.assistant_answer.len(), MAX_ANSWER_BYTES);
}
#[test]
fn test_extract_user_corrections() {
let prompt = "Actually, orders are sorted by date. No, use return_date.";
let corrections = extract_user_corrections(prompt);
assert_eq!(corrections.len(), 2);
}
#[test]
fn test_dto_conversion_from_supplied_contract() {
let contract = SuppliedContract {
profile: "p1".into(),
object: "catalog.public.orders".into(),
schema_state: "current",
claims: vec![crate::contracts::SuppliedClaim {
claim_id: ClaimId::parse("c-1").unwrap(),
kind: "default_time_column",
value: "created_at".into(),
column: Some("created_at".into()),
status: saya_types::ClaimStatus::Confirmed,
}],
};
let dto = SuppliedContractDto::from(&contract);
assert_eq!(dto.profile, "p1");
assert_eq!(dto.claims.len(), 1);
assert_eq!(dto.claims[0].claim_id, "c-1");
assert_eq!(dto.claims[0].value, "created_at");
}
}