blob: ce87e494796fcd9b7782f2a94ce33535c862f803 [file]
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
//! Scalar canonical evidence for an assistant row's model and output producer.
use crate::StoreError;
use maka_agent_facts::history::RowPosition;
use maka_agent_ports::session::history::HistoryModel;
use sqlx::sqlite::SqliteWorkerConnection;
use sqlx::{Row, SqliteConnection};
#[derive(Clone, Debug, PartialEq, Eq)]
pub(super) enum Producer {
Model { invocation: String, step: String },
Executor { invocation: String },
}
pub(super) struct Metadata {
pub model: Option<HistoryModel>,
pub producer: Option<Producer>,
}
// Drive each lookup from the canonical sequence/operation/step index. Membership
// validates each candidate; the history UNION must not drive a whole-session scan.
const PRODUCER: &str = "SELECT e.invocation_id, e.kind, e.operation_id
FROM runtime_events e WHERE e.sequence = ?2 AND e.sequence <= ?3
AND (e.event_session = ?1 OR EXISTS(
SELECT 1 FROM session_history_members h WHERE h.session_id = ?1 AND h.sequence = e.sequence))";
const REQUESTED: &str = "SELECT json_extract(e.event_json, '$.fact.model_id')
FROM runtime_events e WHERE e.operation_id = ?3 AND e.kind = 'model_requested'
AND e.invocation_id = ?2 AND e.sequence <= ?4
AND json_extract(e.event_json, '$.fact.step_id') = ?3
AND (e.event_session = ?1 OR EXISTS(
SELECT 1 FROM session_history_members h WHERE h.session_id = ?1 AND h.sequence = e.sequence))";
const OBSERVED: &str = "SELECT e.sequence, json_extract(e.event_json, '$.fact.event.data.model')
FROM runtime_events e WHERE e.invocation_id = ?2
AND CAST(json_extract(e.event_json, '$.fact.step_id') AS TEXT) = ?3
AND json_extract(e.event_json, '$.fact.step_id') = ?3
AND e.sequence <= ?4 AND e.kind = 'model_observed' AND e.event_json IS NOT NULL
AND json_extract(e.event_json, '$.fact.event.kind') = 'response_metadata'
AND json_type(e.event_json, '$.fact.event.data.model') = 'text'
AND (e.event_session = ?1 OR EXISTS(
SELECT 1 FROM session_history_members h WHERE h.session_id = ?1 AND h.sequence = e.sequence))
ORDER BY e.sequence DESC LIMIT 1";
const COMPLETED: &str = "SELECT e.sequence, json_extract(e.event_json, '$.fact.output.model')
FROM runtime_events e WHERE e.operation_id = ?3 AND e.kind = 'model_completed'
AND e.invocation_id = ?2 AND e.sequence <= ?4
AND json_extract(e.event_json, '$.fact.step_id') = ?3
AND json_type(e.event_json, '$.fact.output.model') = 'text'
AND (e.event_session = ?1 OR EXISTS(
SELECT 1 FROM session_history_members h WHERE h.session_id = ?1 AND h.sequence = e.sequence))";
const EXECUTOR: &str = "SELECT EXISTS(SELECT 1 FROM runtime_events e
WHERE e.invocation_id = ?2 AND e.kind = 'executor_started' AND e.sequence <= ?3
AND (e.event_session = ?1 OR EXISTS(
SELECT 1 FROM session_history_members h WHERE h.session_id = ?1 AND h.sequence = e.sequence)))";
pub(super) async fn resolve(
tx: &mut SqliteConnection,
session: &str,
position: RowPosition,
through: i64,
) -> Result<Metadata, StoreError> {
let session = session.to_owned();
crate::connection::worker_query(tx, move |worker| {
Box::pin(async move { resolve_on_worker(worker, &session, position, through).await })
})
.await
}
pub(super) async fn resolve_on_worker(
tx: &mut SqliteWorkerConnection<'_>,
session: &str,
position: RowPosition,
through: i64,
) -> Result<Metadata, StoreError> {
let producer_sequence = super::read::sql_position(position)?.0;
let producer = sqlx::query(PRODUCER)
.bind(session)
.bind(producer_sequence)
.bind(through)
.fetch_optional(&mut *tx)
.await?
.ok_or_else(invalid)?;
let invocation: String = producer.try_get(0)?;
let kind: String = producer.try_get(1)?;
if kind == "message_imported" {
let (content, model): (String, Option<String>) = sqlx::query_as(
"SELECT json_extract(event_json, '$.fact.record.content.kind'),
json_extract(event_json, '$.fact.record.content.model')
FROM runtime_events WHERE sequence = ?",
)
.bind(producer_sequence)
.fetch_one(&mut *tx)
.await?;
if content != "assistant" {
return Err(invalid());
}
return Ok(Metadata {
model: model.map(|model_id| HistoryModel::Imported { model_id }),
producer: None,
});
}
if matches!(kind.as_str(), "executor_observed" | "executor_completed") {
return Ok(Metadata {
model: None,
producer: (kind == "executor_completed").then_some(Producer::Executor { invocation }),
});
}
let step: Option<String> = if kind == "invocation_ended" {
let executor: bool = sqlx::query_scalar(EXECUTOR)
.bind(session)
.bind(&invocation)
.bind(producer_sequence)
.fetch_one(&mut *tx)
.await?;
if executor {
return Ok(Metadata {
model: None,
producer: None,
});
}
// Resolve the unfinished step at the actual sealing fact, never from a
// later public fence that may contain unrelated model steps.
super::evidence::pending_step_on_worker(tx, &invocation, producer_sequence).await?
} else if kind == "model_observed" {
sqlx::query_scalar("SELECT json_extract(event_json, '$.fact.step_id') FROM runtime_events WHERE sequence = ?")
.bind(producer_sequence).fetch_one(&mut *tx).await?
} else if matches!(kind.as_str(), "model_completed" | "model_interrupted") {
producer.try_get(2)?
} else {
return Err(invalid());
};
let step = step.ok_or_else(invalid)?;
let requested: Option<String> = sqlx::query_scalar(REQUESTED)
.bind(session)
.bind(&invocation)
.bind(&step)
.bind(through)
.fetch_optional(&mut *tx)
.await?
.flatten();
let requested_model_id = requested.ok_or_else(invalid)?;
let observed: Option<(i64, String)> = sqlx::query_as(OBSERVED)
.bind(session)
.bind(&invocation)
.bind(&step)
.bind(through)
.fetch_optional(&mut *tx)
.await?;
let completed: Option<(i64, String)> = sqlx::query_as(COMPLETED)
.bind(session)
.bind(&invocation)
.bind(&step)
.bind(through)
.fetch_optional(&mut *tx)
.await?;
let reported_model_id = observed
.into_iter()
.chain(completed)
.max_by_key(|(sequence, _)| *sequence)
.map(|(_, model)| model);
Ok(Metadata {
model: Some(HistoryModel::Native {
requested_model_id,
reported_model_id,
}),
producer: matches!(kind.as_str(), "model_observed" | "model_completed")
.then_some(Producer::Model { invocation, step }),
})
}
fn invalid() -> StoreError {
StoreError::InvalidTransition("invalid canonical history model evidence".into())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::EventLog;
use serde_json::json;
#[tokio::test]
async fn identity_queries_seek_scalar_candidates_instead_of_scanning_step_observations() {
let temp = tempfile::tempdir().unwrap();
let log = EventLog::open(&temp.path().join("model-identity.sqlite"))
.await
.unwrap();
log.connection.run(|db| Box::pin(async move {
// This SQL-plan fixture deliberately writes only columns used by the
// scalar queries; canonical admission behavior is covered by history.
let mut sequence = 0i64;
for step in ["with-report", "without-report"] {
sequence += 1;
let requested = json!({"invocation":{"session_id":"source"},"fact":{"kind":"model_requested","step_id":step,"model_id":format!("alias-{step}")}});
sqlx::query("INSERT INTO event_log(sequence,event_id,invocation_id,kind,operation_id,event_json) VALUES(?,?, 'invocation','model_requested',?,?)")
.bind(sequence).bind(format!("event-{sequence}")).bind(step).bind(requested.to_string()).execute(&mut *db).await?;
for delta in 0..128 {
sequence += 1;
let event = if step == "with-report" && delta == 0 {
json!({"kind":"response_metadata","data":{"model":"reported-early"}})
} else { json!({"kind":"part_delta","data":{"id":"part","text":"text"}}) };
let observed = json!({"invocation":{"session_id":"source"},"fact":{"kind":"model_observed","step_id":step,"event":event}});
sqlx::query("INSERT INTO event_log(sequence,event_id,invocation_id,kind,event_json) VALUES(?,?,'invocation','model_observed',?)")
.bind(sequence).bind(format!("event-{sequence}")).bind(observed.to_string()).execute(&mut *db).await?;
}
sequence += 1;
let completed = json!({"invocation":{"session_id":"source"},"fact":{"kind":"model_completed","step_id":step,"output":{"model":if step=="with-report" {Some("reported-final")} else {None},"parts":[{"text":"large output".repeat(16384)}]}}});
sqlx::query("INSERT INTO event_log(sequence,event_id,invocation_id,kind,operation_id,event_json) VALUES(?,?,'invocation','model_completed',?,?)")
.bind(sequence).bind(format!("event-{sequence}")).bind(step).bind(completed.to_string()).execute(&mut *db).await?;
}
for (query, expected_index) in [(REQUESTED, "operation_fact"), (OBSERVED, "model_response_identity"), (COMPLETED, "operation_fact")] {
let explain = format!("EXPLAIN QUERY PLAN {query}");
let rows = sqlx::query(sqlx::AssertSqlSafe(explain.as_str())).bind("source").bind("invocation").bind("with-report").bind(sequence)
.fetch_all(&mut *db).await?;
let plan = rows.iter().map(|row| row.get::<String,_>(3)).collect::<Vec<_>>().join("\n");
assert!(plan.contains(&format!("USING INDEX {expected_index}")), "{plan}");
assert!(!plan.contains("SCAN event_log") && !plan.contains("SCAN h") && !plan.contains("USE TEMP B-TREE"), "{plan}");
}
let requested: String = sqlx::query_scalar(REQUESTED).bind("source").bind("invocation").bind("with-report").bind(sequence).fetch_one(&mut *db).await?;
assert_eq!(requested, "alias-with-report");
let observed: (i64, String) = sqlx::query_as(OBSERVED).bind("source").bind("invocation").bind("with-report").bind(sequence).fetch_one(&mut *db).await?;
assert_eq!(observed, (2, "reported-early".into()));
let completed: (i64, String) = sqlx::query_as(COMPLETED).bind("source").bind("invocation").bind("with-report").bind(sequence).fetch_one(&mut *db).await?;
assert_eq!(completed, (130, "reported-final".into()));
let absent: Option<(i64, String)> = sqlx::query_as(OBSERVED).bind("source").bind("invocation").bind("without-report").bind(sequence).fetch_optional(&mut *db).await?;
assert_eq!(absent, None);
Ok(())
})).await.unwrap();
log.close().await.unwrap();
}
}