use std::io::{BufRead, BufReader, Read, Write};
use std::path::Path;
use serde::{Deserialize, Serialize};
use serde_json::{json, Map, Value};
use crate::error::{MnemoError, Result};
use crate::memory::{Memory, MemoryType, Metric, Scope};
use crate::result_cache::{CachePutOpts, SemanticCachePutOpts, DEFAULT_SEMANTIC_THRESHOLD};
use crate::store::{Mnemo, RecallRequest};
const PROTOCOL_VERSION: &str = "2025-06-18";
const PASSPHRASE_ENV: &str = "MNEMO_PASSPHRASE";
#[doc(hidden)]
pub fn __fuzz_parse_request_line(line: &str) -> std::result::Result<(), String> {
serde_json::from_str::<Request>(line)
.map(|_| ())
.map_err(|e| e.to_string())
}
#[derive(Debug, Deserialize)]
struct Request {
#[serde(default)]
jsonrpc: String,
#[serde(default)]
id: Option<Value>,
method: String,
#[serde(default)]
params: Value,
}
#[derive(Debug, Serialize)]
struct Response {
jsonrpc: &'static str,
id: Value,
#[serde(skip_serializing_if = "Option::is_none")]
result: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<RpcError>,
}
#[derive(Debug, Serialize)]
struct RpcError {
code: i32,
message: String,
#[serde(skip_serializing_if = "Option::is_none")]
data: Option<Value>,
}
impl RpcError {
fn new(code: i32, message: impl Into<String>) -> Self {
Self {
code,
message: message.into(),
data: None,
}
}
}
const PARSE_ERROR: i32 = -32700;
const INVALID_REQUEST: i32 = -32600;
const METHOD_NOT_FOUND: i32 = -32601;
const INVALID_PARAMS: i32 = -32602;
const INTERNAL_ERROR: i32 = -32603;
pub fn serve_stdio(path: &Path) -> Result<()> {
serve_stdio_with_embedder(path, None)
}
pub fn serve_stdio_with_embedder(
path: &Path,
embed_spec: Option<(String, String, bool, std::time::Duration)>,
) -> Result<()> {
let stdin = std::io::stdin();
let stdout = std::io::stdout();
serve_with_streams(path, stdin.lock(), stdout.lock(), embed_spec)
}
pub fn serve_with_streams<R: Read, W: Write>(
path: &Path,
reader: R,
mut writer: W,
embed_spec: Option<(String, String, bool, std::time::Duration)>,
) -> Result<()> {
let passphrase = std::env::var(PASSPHRASE_ENV).map_err(|_| {
MnemoError::Invalid(format!(
"{PASSPHRASE_ENV} is required to serve — set it in the shell before starting"
))
})?;
eprintln!(
"mnemo mcp: opening {} (protocol {PROTOCOL_VERSION})",
path.display()
);
let mut db = Mnemo::open(path, &passphrase)?;
if let Some((cmd, name, normalize, timeout)) = embed_spec {
let identity = crate::EmbedderIdentity {
name,
dimensions: db.dimensions(),
normalize,
};
let embedder = crate::EmbedCmd::new(&cmd, identity)?.with_timeout(timeout);
db.set_embedder(embedder)?;
eprintln!(
"mnemo mcp: embedder ready (text params accepted, timeout {}s)",
timeout.as_secs()
);
}
eprintln!("mnemo mcp: ready ({} live memories)", db.len());
let reader = BufReader::new(reader);
for line in reader.lines() {
let line = line.map_err(MnemoError::Io)?;
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
let (id, out) = match serde_json::from_str::<Request>(trimmed) {
Ok(req) => {
let id = req.id.clone();
let is_notification = id.is_none();
let result = handle(&mut db, &req);
if is_notification {
continue;
}
(id.unwrap_or(Value::Null), result)
}
Err(e) => (
Value::Null,
Err(RpcError::new(PARSE_ERROR, format!("parse error: {e}"))),
),
};
let resp = match out {
Ok(value) => Response {
jsonrpc: "2.0",
id,
result: Some(value),
error: None,
},
Err(err) => Response {
jsonrpc: "2.0",
id,
result: None,
error: Some(err),
},
};
let line =
serde_json::to_string(&resp).map_err(|e| MnemoError::Serialize(e.to_string()))?;
writeln!(writer, "{line}").map_err(MnemoError::Io)?;
writer.flush().map_err(MnemoError::Io)?;
}
eprintln!("mnemo mcp: stdin closed, shutting down");
Ok(())
}
fn handle(db: &mut Mnemo, req: &Request) -> std::result::Result<Value, RpcError> {
if req.jsonrpc != "2.0" {
return Err(RpcError::new(
INVALID_REQUEST,
format!("expected jsonrpc \"2.0\", got {:?}", req.jsonrpc),
));
}
match req.method.as_str() {
"initialize" => Ok(initialize_result(&req.params)),
"notifications/initialized" | "notifications/cancelled" => Ok(Value::Null),
"ping" => Ok(json!({})),
"tools/list" => Ok(tools_list()),
"tools/call" => tools_call(db, &req.params),
other => Err(RpcError::new(
METHOD_NOT_FOUND,
format!("method '{other}' not implemented"),
)),
}
}
fn initialize_result(params: &Value) -> Value {
let client_version = params
.get("protocolVersion")
.and_then(|v| v.as_str())
.unwrap_or(PROTOCOL_VERSION)
.to_string();
json!({
"protocolVersion": client_version,
"capabilities": {
"tools": { "listChanged": false },
},
"serverInfo": {
"name": "mnemo",
"version": env!("CARGO_PKG_VERSION"),
},
})
}
fn tools_list() -> Value {
json!({
"tools": [
{
"name": "about",
"description": "Read-only briefing for this .mnemo file: manifest scaffold plus any memories tagged as onboarding. Run this first to learn how the file expects to be used.",
"inputSchema": {
"type": "object",
"properties": {},
"additionalProperties": false
}
},
{
"name": "remember",
"description": "Insert or overwrite a memory. Pass EITHER `vector` (explicit embedding) OR `text` (embedded by the server if it was started with --embed-cmd / --embed-name). Persists to disk before returning. Returns `{id}`.",
"inputSchema": {
"type": "object",
"required": ["content", "memory_type"],
"properties": {
"content": { "type": "string" },
"memory_type": { "type": "string", "enum": ["episodic", "semantic", "procedural", "working"] },
"vector": { "type": "array", "items": { "type": "number" }, "description": "Explicit embedding; mutually exclusive with `text`." },
"text": { "type": "string", "description": "Content to embed via the server's registered embedder; mutually exclusive with `vector`. Requires the server to have been started with --embed-cmd." },
"importance": { "type": "number", "minimum": 0, "maximum": 1 },
"agent_id": { "type": "string" },
"session_id": { "type": "string" },
"ttl_secs": { "type": "integer" },
"shared": { "type": "boolean" },
"metadata": { "type": "object" }
},
"additionalProperties": false
}
},
{
"name": "recall",
"description": "Score every live memory (or the ANN-narrowed candidate set) against a query and return the top-k. Pass EITHER `query` (explicit vector) OR `text` (embedded by the server; requires --embed-cmd at serve start). Does not track access on this handle.",
"inputSchema": {
"type": "object",
"properties": {
"query": { "type": "array", "items": { "type": "number" }, "description": "Explicit query vector; mutually exclusive with `text`." },
"text": { "type": "string", "description": "Query text; embedded by the server's registered embedder. Mutually exclusive with `query`." },
"top_k": { "type": "integer", "default": 10 },
"types": { "type": "array", "items": { "type": "string" } },
"agent": { "type": "string" },
"metric": { "type": "string", "enum": ["cosine", "l2", "dot"], "default": "cosine" }
},
"additionalProperties": false
}
},
{
"name": "forget",
"description": "Soft-delete a memory by ULID. Space is reclaimed on the next `mnemo compact`. Persists to disk before returning.",
"inputSchema": {
"type": "object",
"required": ["id"],
"properties": {
"id": { "type": "string", "description": "ULID of the memory to delete" }
},
"additionalProperties": false
}
},
{
"name": "list",
"description": "Return every live memory as JSON. Sorted by creation time, oldest first. Vectors are omitted by default to keep payloads small.",
"inputSchema": {
"type": "object",
"properties": {
"include_vectors": { "type": "boolean", "default": false }
},
"additionalProperties": false
}
},
{
"name": "snapshot_list",
"description": "List every restorable snapshot (one per committed flush). Each entry gives the transaction ID, creation timestamp, and memory count.",
"inputSchema": {
"type": "object",
"properties": {},
"additionalProperties": false
}
},
{
"name": "stats",
"description": "Summary statistics: memory count, deleted count, agents present, snapshot count, page cache occupancy.",
"inputSchema": {
"type": "object",
"properties": {},
"additionalProperties": false
}
},
{
"name": "cache_put",
"description": "Insert or overwrite an exact-key cache entry. Value is a UTF-8 string; use `content_type` to label the payload. Persists to disk before returning. Namespace and key are caller-chosen; the key is SHA-256 hashed by the engine.",
"inputSchema": {
"type": "object",
"required": ["namespace", "key", "value"],
"properties": {
"namespace": { "type": "string" },
"key": { "type": "string" },
"value": { "type": "string" },
"content_type": { "type": "string", "default": "text" },
"ttl_secs": { "type": "integer", "minimum": 0 }
},
"additionalProperties": false
}
},
{
"name": "cache_get",
"description": "Look up an exact-key cache entry. Returns `{value, content_type, hit}` on a hit, `{hit: false}` on miss (missing or TTL-expired). Bumps access stats catalog-only.",
"inputSchema": {
"type": "object",
"required": ["namespace", "key"],
"properties": {
"namespace": { "type": "string" },
"key": { "type": "string" }
},
"additionalProperties": false
}
},
{
"name": "cache_stats",
"description": "Cache statistics for a namespace (omit for the whole cache). Returns entries, bytes, hits, misses, hit_rate, evictions.",
"inputSchema": {
"type": "object",
"properties": {
"namespace": { "type": "string" }
},
"additionalProperties": false
}
},
{
"name": "cache_put_semantic",
"description": "Insert an entry with an embedding vector for semantic recall (Phase 10.2). Pass EITHER `vector` + `model` (explicit) OR `text` (embedded by the server; `model` defaults to the server embedder's identity name). Later cache_get_semantic calls under a different `model` won't see this entry (a different-model hit would be a semantics bug).",
"inputSchema": {
"type": "object",
"required": ["namespace", "key", "value"],
"properties": {
"namespace": { "type": "string" },
"key": { "type": "string" },
"vector": { "type": "array", "items": { "type": "number" }, "description": "Explicit embedding; mutually exclusive with `text`. Requires `model`." },
"text": { "type": "string", "description": "Content to embed via the server's registered embedder; mutually exclusive with `vector`." },
"value": { "type": "string" },
"model": { "type": "string", "description": "Required with `vector`; optional with `text` (defaults to the embedder's identity name)." },
"content_type": { "type": "string", "default": "text" },
"ttl_secs": { "type": "integer", "minimum": 0 }
},
"additionalProperties": false
}
},
{
"name": "cache_get_semantic",
"description": "Semantic recall over the cache: top-1 cosine over live vectored entries in the namespace whose `model` matches. Pass EITHER `query` + `model` OR `text` (which embeds via the server and defaults the model to the embedder's identity). Returns `{hit: true, value, similarity, ...}` iff similarity >= threshold, else `{hit: false}`. Threshold defaults to 0.97.",
"inputSchema": {
"type": "object",
"required": ["namespace"],
"properties": {
"namespace": { "type": "string" },
"query": { "type": "array", "items": { "type": "number" }, "description": "Explicit query vector; mutually exclusive with `text`. Requires `model`." },
"text": { "type": "string", "description": "Query text; embedded by the server. Mutually exclusive with `query`." },
"model": { "type": "string", "description": "Required with `query`; optional with `text`." },
"threshold": { "type": "number", "minimum": -1.0, "maximum": 1.0 }
},
"additionalProperties": false
}
}
]
})
}
fn tools_call(db: &mut Mnemo, params: &Value) -> std::result::Result<Value, RpcError> {
let name = params
.get("name")
.and_then(|v| v.as_str())
.ok_or_else(|| RpcError::new(INVALID_PARAMS, "missing tool name"))?;
let empty = json!({});
let args = params.get("arguments").unwrap_or(&empty);
let text = match name {
"about" => tool_about(db)?,
"remember" => tool_remember(db, args)?,
"recall" => tool_recall(db, args)?,
"forget" => tool_forget(db, args)?,
"list" => tool_list(db, args)?,
"snapshot_list" => tool_snapshot_list(db)?,
"stats" => tool_stats(db)?,
"cache_put" => tool_cache_put(db, args)?,
"cache_get" => tool_cache_get(db, args)?,
"cache_stats" => tool_cache_stats(db, args)?,
"cache_put_semantic" => tool_cache_put_semantic(db, args)?,
"cache_get_semantic" => tool_cache_get_semantic(db, args)?,
other => {
return Err(RpcError::new(
METHOD_NOT_FOUND,
format!("unknown tool '{other}'"),
))
}
};
Ok(json!({
"content": [ { "type": "text", "text": text } ],
"isError": false,
}))
}
fn embed_text_via_server(
db: &mut Mnemo,
tool: &str,
text: &str,
) -> std::result::Result<Vec<f32>, RpcError> {
let embedder = db.embedder().ok_or_else(|| {
RpcError::new(
INVALID_PARAMS,
format!(
"{tool}: 'text' requires an embedder on the server. \
Start with: mnemo serve --mcp <file> --embed-cmd '<cmd>' --embed-name <name>"
),
)
})?;
let mut vs = embedder.embed(&[text]).map_err(mnemo_err)?;
vs.pop().ok_or_else(|| {
RpcError::new(
INTERNAL_ERROR,
format!("{tool}: embedder returned no vector for the input text"),
)
})
}
fn mnemo_err(e: MnemoError) -> RpcError {
RpcError::new(INTERNAL_ERROR, e.to_string())
}
fn tool_about(db: &mut Mnemo) -> std::result::Result<String, RpcError> {
let mems = db.about().map_err(mnemo_err)?;
let arr: Vec<Value> = mems.iter().map(memory_to_json).collect();
Ok(serde_json::to_string_pretty(&arr).unwrap_or_else(|_| "[]".into()))
}
fn tool_remember(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
let content = args
.get("content")
.and_then(|v| v.as_str())
.ok_or_else(|| RpcError::new(INVALID_PARAMS, "remember: missing 'content'"))?;
let mtype_str = args
.get("memory_type")
.and_then(|v| v.as_str())
.ok_or_else(|| RpcError::new(INVALID_PARAMS, "remember: missing 'memory_type'"))?;
let mtype = MemoryType::parse(mtype_str).ok_or_else(|| {
RpcError::new(
INVALID_PARAMS,
format!("remember: unknown memory_type '{mtype_str}'"),
)
})?;
let vector = match (args.get("vector"), args.get("text")) {
(Some(_), Some(_)) => {
return Err(RpcError::new(
INVALID_PARAMS,
"remember: pass EITHER 'vector' or 'text', not both",
));
}
(Some(v), None) => v
.as_array()
.ok_or_else(|| {
RpcError::new(INVALID_PARAMS, "remember: 'vector' must be a JSON array")
})?
.iter()
.map(|v| v.as_f64().map(|f| f as f32))
.collect::<Option<Vec<f32>>>()
.ok_or_else(|| {
RpcError::new(
INVALID_PARAMS,
"remember: 'vector' must be an array of numbers",
)
})?,
(None, Some(t)) => {
let text = t.as_str().ok_or_else(|| {
RpcError::new(INVALID_PARAMS, "remember: 'text' must be a string")
})?;
embed_text_via_server(db, "remember", text)?
}
(None, None) => {
return Err(RpcError::new(
INVALID_PARAMS,
"remember: missing 'vector' or 'text'",
))
}
};
let mut mem = Memory::new(content, mtype, vector);
if let Some(agent) = args.get("agent_id").and_then(|v| v.as_str()) {
mem = mem.with_agent(agent);
}
if let Some(sess) = args.get("session_id").and_then(|v| v.as_str()) {
mem = mem.with_session(sess);
}
if let Some(imp) = args.get("importance").and_then(|v| v.as_f64()) {
mem = mem.with_importance(imp as f32);
}
if let Some(ttl) = args.get("ttl_secs").and_then(|v| v.as_i64()) {
mem = mem.with_ttl(ttl);
}
if let Some(shared) = args.get("shared").and_then(|v| v.as_bool()) {
if shared {
mem = mem.with_scope(Scope::Shared);
}
}
if let Some(meta) = args.get("metadata").and_then(|v| v.as_object()) {
for (k, v) in meta {
mem.metadata.insert(k.clone(), v.clone());
}
}
let id = db.remember(mem).map_err(mnemo_err)?;
db.flush().map_err(mnemo_err)?;
Ok(json!({ "id": id.to_string() }).to_string())
}
fn tool_recall(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
let query = match (args.get("query"), args.get("text")) {
(Some(_), Some(_)) => {
return Err(RpcError::new(
INVALID_PARAMS,
"recall: pass EITHER 'query' or 'text', not both",
))
}
(Some(v), None) => v
.as_array()
.ok_or_else(|| RpcError::new(INVALID_PARAMS, "recall: 'query' must be a JSON array"))?
.iter()
.map(|v| v.as_f64().map(|f| f as f32))
.collect::<Option<Vec<f32>>>()
.ok_or_else(|| {
RpcError::new(
INVALID_PARAMS,
"recall: 'query' must be an array of numbers",
)
})?,
(None, Some(t)) => {
let text = t
.as_str()
.ok_or_else(|| RpcError::new(INVALID_PARAMS, "recall: 'text' must be a string"))?;
embed_text_via_server(db, "recall", text)?
}
(None, None) => {
return Err(RpcError::new(
INVALID_PARAMS,
"recall: missing 'query' or 'text'",
))
}
};
let top_k = args.get("top_k").and_then(|v| v.as_u64()).unwrap_or(10) as usize;
let metric = match args
.get("metric")
.and_then(|v| v.as_str())
.unwrap_or("cosine")
{
"cosine" => Metric::Cosine,
"l2" => Metric::L2,
"dot" => Metric::Dot,
other => {
return Err(RpcError::new(
INVALID_PARAMS,
format!("recall: unknown metric '{other}'"),
))
}
};
let mut req = RecallRequest::new(query)
.top_k(top_k)
.metric(metric)
.track_access(false);
if let Some(types) = args.get("types").and_then(|v| v.as_array()) {
let parsed: Vec<MemoryType> = types
.iter()
.filter_map(|t| t.as_str().and_then(MemoryType::parse))
.collect();
if !parsed.is_empty() {
req = req.types(parsed);
}
}
if let Some(agent) = args.get("agent").and_then(|v| v.as_str()) {
req = req.agent(agent.to_string());
}
let hits = db.recall(&req).map_err(mnemo_err)?;
let out: Vec<Value> = hits
.into_iter()
.map(|h| {
let mut obj = memory_to_json(&h.memory);
if let Some(map) = obj.as_object_mut() {
map.insert("score".into(), json!(h.score));
map.insert("similarity".into(), json!(h.similarity));
}
obj
})
.collect();
Ok(serde_json::to_string_pretty(&out).unwrap_or_else(|_| "[]".into()))
}
fn tool_forget(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
let id_str = args
.get("id")
.and_then(|v| v.as_str())
.ok_or_else(|| RpcError::new(INVALID_PARAMS, "forget: missing 'id'"))?;
let ulid = crate::Ulid::from_string(id_str)
.map_err(|_| RpcError::new(INVALID_PARAMS, format!("forget: invalid ULID '{id_str}'")))?;
db.delete(&ulid).map_err(mnemo_err)?;
db.flush().map_err(mnemo_err)?;
Ok(json!({ "deleted": id_str }).to_string())
}
fn tool_list(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
let include_vectors = args
.get("include_vectors")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let mems = db.memories().map_err(mnemo_err)?;
let out: Vec<Value> = mems
.iter()
.map(|m| {
let mut v = memory_to_json(m);
if !include_vectors {
if let Some(o) = v.as_object_mut() {
o.remove("vector");
}
}
v
})
.collect();
Ok(serde_json::to_string_pretty(&out).unwrap_or_else(|_| "[]".into()))
}
fn tool_snapshot_list(db: &mut Mnemo) -> std::result::Result<String, RpcError> {
let snaps = db.snapshots();
let out: Vec<Value> = snaps
.into_iter()
.map(|s| {
json!({
"txn_id": s.txn_id,
"created_at": s.created_at,
"memory_count": s.memory_count,
})
})
.collect();
Ok(serde_json::to_string_pretty(&out).unwrap_or_else(|_| "[]".into()))
}
fn tool_stats(db: &mut Mnemo) -> std::result::Result<String, RpcError> {
let s = db.stats().map_err(mnemo_err)?;
let (cache_pages, cache_capacity) = db.page_cache_stats();
let snapshot_count = db.snapshots().len();
Ok(json!({
"memories": s.memories,
"deleted": s.deleted,
"dimensions": s.dimensions,
"file_bytes": s.file_bytes,
"agents": s.agents,
"encrypted": s.encrypted,
"created_at": s.created_at,
"wal_pages": s.wal_pages,
"index": s.index,
"snapshot_count": snapshot_count,
"cache_pages": cache_pages,
"cache_capacity": cache_capacity,
})
.to_string())
}
fn tool_cache_put(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
let ns = args
.get("namespace")
.and_then(|v| v.as_str())
.ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put: missing 'namespace'"))?;
let key = args
.get("key")
.and_then(|v| v.as_str())
.ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put: missing 'key'"))?;
let value = args
.get("value")
.and_then(|v| v.as_str())
.ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put: missing 'value'"))?;
let content_type = args
.get("content_type")
.and_then(|v| v.as_str())
.unwrap_or("text")
.to_string();
let ttl_secs = args.get("ttl_secs").and_then(|v| v.as_u64());
let opts = CachePutOpts {
content_type,
ttl_secs,
cost_hint_ms: None,
};
db.cache_put(ns, key, value.as_bytes(), opts)
.map_err(mnemo_err)?;
db.flush().map_err(mnemo_err)?;
Ok(json!({ "cached": true, "namespace": ns, "key": key, "bytes": value.len() }).to_string())
}
fn tool_cache_get(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
let ns = args
.get("namespace")
.and_then(|v| v.as_str())
.ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_get: missing 'namespace'"))?;
let key = args
.get("key")
.and_then(|v| v.as_str())
.ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_get: missing 'key'"))?;
match db.cache_get(ns, key).map_err(mnemo_err)? {
Some(hit) => Ok(json!({
"hit": true,
"value": String::from_utf8_lossy(&hit.value),
"content_type": hit.content_type,
"created_at": hit.created_at,
"accessed_at": hit.accessed_at,
"access_count": hit.access_count,
})
.to_string()),
None => Ok(json!({ "hit": false }).to_string()),
}
}
fn tool_cache_stats(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
let ns = args.get("namespace").and_then(|v| v.as_str());
let s = db.cache_stats(ns);
Ok(json!({
"namespace": ns,
"entries": s.entries,
"bytes": s.bytes,
"hits": s.hits,
"misses": s.misses,
"hit_rate": s.hit_rate,
"evictions": s.evictions,
})
.to_string())
}
fn tool_cache_put_semantic(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
let ns = args
.get("namespace")
.and_then(|v| v.as_str())
.ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put_semantic: missing 'namespace'"))?;
let key = args
.get("key")
.and_then(|v| v.as_str())
.ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put_semantic: missing 'key'"))?;
let value = args
.get("value")
.and_then(|v| v.as_str())
.ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put_semantic: missing 'value'"))?;
let (vector, model) = match (args.get("vector"), args.get("text")) {
(Some(_), Some(_)) => {
return Err(RpcError::new(
INVALID_PARAMS,
"cache_put_semantic: pass EITHER 'vector' or 'text', not both",
));
}
(Some(v), None) => {
let vec = v
.as_array()
.ok_or_else(|| {
RpcError::new(
INVALID_PARAMS,
"cache_put_semantic: 'vector' must be a JSON array",
)
})?
.iter()
.map(|v| v.as_f64().map(|f| f as f32))
.collect::<Option<Vec<f32>>>()
.ok_or_else(|| {
RpcError::new(
INVALID_PARAMS,
"cache_put_semantic: 'vector' must be an array of numbers",
)
})?;
let model = args
.get("model")
.and_then(|v| v.as_str())
.ok_or_else(|| {
RpcError::new(INVALID_PARAMS, "cache_put_semantic: missing 'model'")
})?
.to_string();
(vec, model)
}
(None, Some(t)) => {
let text = t.as_str().ok_or_else(|| {
RpcError::new(
INVALID_PARAMS,
"cache_put_semantic: 'text' must be a string",
)
})?;
let vec = embed_text_via_server(db, "cache_put_semantic", text)?;
let model = match args.get("model").and_then(|v| v.as_str()) {
Some(m) => m.to_string(),
None => {
db.embedder()
.expect("embed_text_via_server would have errored otherwise")
.identity()
.name
}
};
(vec, model)
}
(None, None) => {
return Err(RpcError::new(
INVALID_PARAMS,
"cache_put_semantic: missing 'vector' or 'text'",
))
}
};
let content_type = args
.get("content_type")
.and_then(|v| v.as_str())
.unwrap_or("text")
.to_string();
let ttl_secs = args.get("ttl_secs").and_then(|v| v.as_u64());
let opts = SemanticCachePutOpts {
content_type,
ttl_secs,
cost_hint_ms: None,
model,
};
db.cache_put_semantic(ns, key, vector, value.as_bytes(), opts)
.map_err(mnemo_err)?;
db.flush().map_err(mnemo_err)?;
Ok(json!({ "cached": true, "namespace": ns, "key": key, "bytes": value.len() }).to_string())
}
fn tool_cache_get_semantic(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
let ns = args
.get("namespace")
.and_then(|v| v.as_str())
.ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_get_semantic: missing 'namespace'"))?;
let (query, model_owned) = match (args.get("query"), args.get("text")) {
(Some(_), Some(_)) => {
return Err(RpcError::new(
INVALID_PARAMS,
"cache_get_semantic: pass EITHER 'query' or 'text', not both",
));
}
(Some(q), None) => {
let vec = q
.as_array()
.ok_or_else(|| {
RpcError::new(
INVALID_PARAMS,
"cache_get_semantic: 'query' must be a JSON array",
)
})?
.iter()
.map(|v| v.as_f64().map(|f| f as f32))
.collect::<Option<Vec<f32>>>()
.ok_or_else(|| {
RpcError::new(
INVALID_PARAMS,
"cache_get_semantic: 'query' must be an array of numbers",
)
})?;
let model = args
.get("model")
.and_then(|v| v.as_str())
.ok_or_else(|| {
RpcError::new(INVALID_PARAMS, "cache_get_semantic: missing 'model'")
})?
.to_string();
(vec, model)
}
(None, Some(t)) => {
let text = t.as_str().ok_or_else(|| {
RpcError::new(
INVALID_PARAMS,
"cache_get_semantic: 'text' must be a string",
)
})?;
let vec = embed_text_via_server(db, "cache_get_semantic", text)?;
let model = match args.get("model").and_then(|v| v.as_str()) {
Some(m) => m.to_string(),
None => {
db.embedder()
.expect("embed_text_via_server would have errored otherwise")
.identity()
.name
}
};
(vec, model)
}
(None, None) => {
return Err(RpcError::new(
INVALID_PARAMS,
"cache_get_semantic: missing 'query' or 'text'",
))
}
};
let model = model_owned.as_str();
let threshold = args
.get("threshold")
.and_then(|v| v.as_f64())
.map(|f| f as f32)
.unwrap_or(DEFAULT_SEMANTIC_THRESHOLD);
match db
.cache_get_semantic(ns, &query, threshold, model)
.map_err(mnemo_err)?
{
Some((hit, similarity)) => Ok(json!({
"hit": true,
"similarity": similarity,
"value": String::from_utf8_lossy(&hit.value),
"content_type": hit.content_type,
"created_at": hit.created_at,
"accessed_at": hit.accessed_at,
"access_count": hit.access_count,
})
.to_string()),
None => Ok(json!({ "hit": false }).to_string()),
}
}
fn memory_to_json(m: &Memory) -> Value {
serde_json::to_value(m).unwrap_or_else(|_| {
let mut map = Map::new();
map.insert("id".into(), json!(m.id.to_string()));
Value::Object(map)
})
}