use std::collections::BTreeMap;
use sim_kernel::{ContentId, Cx, Error, Expr, Result};
use crate::{
capabilities::ai_runner_cache_capability,
content_id::{content_id_for_expr, request_expr_content_id},
objects::content_id_expr,
};
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct PlanCacheKey {
request_content_id: ContentId,
plan_content_id: ContentId,
}
impl PlanCacheKey {
pub fn new(request_content_id: ContentId, plan_content_id: ContentId) -> Self {
Self {
request_content_id,
plan_content_id,
}
}
pub fn for_request_plan(request: &Expr, plan: &Expr) -> Result<Self> {
Ok(Self::new(
request_expr_content_id(request)?,
content_id_for_expr(plan)?,
))
}
pub fn request_content_id(&self) -> &ContentId {
&self.request_content_id
}
pub fn plan_content_id(&self) -> &ContentId {
&self.plan_content_id
}
pub fn to_expr(&self) -> Expr {
Expr::Map(vec![
field(
"request-content-id",
content_id_expr(&self.request_content_id),
),
field("plan-content-id", content_id_expr(&self.plan_content_id)),
])
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum PlanCacheMode {
Disabled,
#[default]
ReadThrough,
ReadOnly,
WriteOnly,
Refresh,
}
impl PlanCacheMode {
pub fn from_expr(expr: &Expr) -> Result<Self> {
match expr {
Expr::String(name) => Self::from_name(name),
Expr::Symbol(symbol) if symbol.namespace.is_none() => Self::from_name(&symbol.name),
Expr::List(items) => match items.as_slice() {
[Expr::Symbol(head), Expr::String(name)]
if head.namespace.is_none() && head.name.as_ref() == "plan/atom" =>
{
Self::from_name(name)
}
_ => Err(Error::Eval(
"plan/cache mode must be a symbol, string, or atom".to_owned(),
)),
},
_ => Err(Error::Eval(
"plan/cache mode must be a symbol, string, or atom".to_owned(),
)),
}
}
pub fn from_name(name: &str) -> Result<Self> {
match name {
"off" | "none" | "disabled" => Ok(Self::Disabled),
"read-through" | "read-write" | "default" => Ok(Self::ReadThrough),
"read" | "read-only" => Ok(Self::ReadOnly),
"write" | "write-only" => Ok(Self::WriteOnly),
"refresh" => Ok(Self::Refresh),
other => Err(Error::Eval(format!("unsupported plan/cache mode {other}"))),
}
}
pub fn reads(self) -> bool {
matches!(self, Self::ReadThrough | Self::ReadOnly)
}
pub fn writes(self) -> bool {
matches!(self, Self::ReadThrough | Self::WriteOnly | Self::Refresh)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PlanCacheWriteTarget {
Memory,
Persistent,
}
#[derive(Clone, Debug, Default)]
pub struct OpenAiPlanCache {
entries: BTreeMap<PlanCacheKey, Expr>,
}
impl OpenAiPlanCache {
pub fn new() -> Self {
Self::default()
}
pub fn get(&self, key: &PlanCacheKey) -> Option<&Expr> {
self.entries.get(key)
}
pub fn put(
&mut self,
cx: &mut Cx,
target: PlanCacheWriteTarget,
key: PlanCacheKey,
response: Expr,
) -> Result<()> {
if target == PlanCacheWriteTarget::Persistent {
cx.require(&ai_runner_cache_capability())?;
}
self.entries.insert(key, response);
Ok(())
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
}
use sim_value::build::entry as field;
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cache_modes_report_read_and_write_behavior() {
let cases = [
("disabled", false, false),
("read-through", true, true),
("read-only", true, false),
("write-only", false, true),
("refresh", false, true),
];
for (name, reads, writes) in cases {
let mode = PlanCacheMode::from_name(name).unwrap();
assert_eq!(mode.reads(), reads, "{name} read behavior");
assert_eq!(mode.writes(), writes, "{name} write behavior");
}
}
}