use std::collections::BTreeSet;
use std::future::Future;
use std::time::{Duration, Instant};
use af_agent::ContextRequest;
use af_agent_session::{text, Event};
use af_llm::ChatMessage;
use crate::{AgentRuntime, CancellationToken, EventWriter, RuntimeError, TurnRequest};
pub(super) async fn await_extension<T>(
label: &str,
cancellation: CancellationToken,
deadline: Instant,
future: impl Future<Output = T>,
) -> Result<T, RuntimeError> {
tokio::pin!(future);
let timeout = tokio::time::sleep(deadline.saturating_duration_since(Instant::now()));
tokio::pin!(timeout);
let error = tokio::select! {
result = &mut future => return Ok(result),
_ = &mut timeout => RuntimeError::ExtensionDeadline(label.into()),
_ = cancellation.cancelled() => RuntimeError::Cancelled,
};
cancellation.cancel();
let _ = tokio::time::timeout(Duration::from_secs(1), &mut future).await;
Err(error)
}
impl AgentRuntime {
pub(super) async fn context_messages(
&self,
request: &TurnRequest,
step: u32,
query: &str,
writer: &dyn EventWriter,
cancellation: CancellationToken,
) -> Result<Vec<ChatMessage>, RuntimeError> {
let extension_cancellation = cancellation.child();
let deadline = Instant::now() + self.limits.provider_deadline;
let context_request = ContextRequest {
request: request.context.clone(),
session_id: request.session_id.clone(),
run_id: request.run_id.clone(),
step,
query: query.into(),
cancellation: extension_cancellation.clone(),
deadline,
};
let mut seen = BTreeSet::new();
let mut contributions = Vec::new();
for contributor in &self.contexts {
for contribution in await_extension(
"context contributor",
extension_cancellation.clone(),
deadline,
contributor.contribute(&context_request),
)
.await?
.map_err(RuntimeError::Context)?
{
contribution.validate().map_err(RuntimeError::Context)?;
if !seen.insert(contribution.id.clone()) {
return Err(RuntimeError::Context(format!(
"duplicate context contribution {}",
contribution.id
)));
}
contributions.push(contribution);
}
}
contributions.sort_by(|left, right| left.id.cmp(&right.id));
if !contributions.is_empty() {
writer
.append(
contributions
.iter()
.map(|contribution| Event::ContextInjected {
run_id: request.run_id.clone(),
step,
contribution_id: contribution.id.clone(),
source: contribution.source.clone(),
version: contribution.version.clone(),
authority: contribution.authority.as_str().into(),
form: contribution.form.clone(),
content: text(contribution.content.clone()),
})
.collect(),
)
.await?;
}
Ok(contributions
.into_iter()
.map(|contribution| contribution.model_message())
.collect())
}
}