use std::future::Future;
use std::pin::Pin;
use tea_protocol::{ProfileId, ProtocolMetadata, RunId, SessionId};
use tea_tools::ToolSpec;
use crate::{ContextError, ContextProviderId, PromptModule};
pub const MAX_CONTEXT_TOOLS: usize = 256;
#[derive(Debug, Clone, PartialEq)]
pub struct ContextRequest {
profile_id: ProfileId,
session_id: SessionId,
run_id: Option<RunId>,
active_tools: Vec<ToolSpec>,
metadata: ProtocolMetadata,
}
impl ContextRequest {
pub fn new(
profile_id: ProfileId,
session_id: SessionId,
run_id: Option<RunId>,
mut active_tools: Vec<ToolSpec>,
metadata: ProtocolMetadata,
) -> Result<Self, ContextError> {
if active_tools.len() > MAX_CONTEXT_TOOLS {
return Err(ContextError::new(
crate::ContextErrorCode::BoundsExceeded,
"context request contains too many active tools",
));
}
active_tools.sort_by(|left, right| left.name().cmp(right.name()));
if active_tools
.windows(2)
.any(|tools| tools[0].name() == tools[1].name())
{
return Err(ContextError::new(
crate::ContextErrorCode::DuplicateIdentity,
"context request contains duplicate active tool names",
));
}
Ok(Self {
profile_id,
session_id,
run_id,
active_tools,
metadata,
})
}
#[must_use]
pub const fn profile_id(&self) -> &ProfileId {
&self.profile_id
}
#[must_use]
pub const fn session_id(&self) -> SessionId {
self.session_id
}
#[must_use]
pub const fn run_id(&self) -> Option<RunId> {
self.run_id
}
#[must_use]
pub fn active_tools(&self) -> &[ToolSpec] {
&self.active_tools
}
#[must_use]
pub const fn metadata(&self) -> &ProtocolMetadata {
&self.metadata
}
}
pub type ContextProviderFuture<'a> =
Pin<Box<dyn Future<Output = Result<Vec<PromptModule>, ContextError>> + Send + 'a>>;
pub trait ContextProvider: std::fmt::Debug + Send + Sync {
fn id(&self) -> &ContextProviderId;
fn provide(&self, request: ContextRequest) -> ContextProviderFuture<'_>;
}
#[derive(Debug, Clone)]
pub struct StaticContextProvider {
id: ContextProviderId,
modules: Vec<PromptModule>,
}
impl StaticContextProvider {
#[must_use]
pub const fn new(id: ContextProviderId, modules: Vec<PromptModule>) -> Self {
Self { id, modules }
}
}
impl ContextProvider for StaticContextProvider {
fn id(&self) -> &ContextProviderId {
&self.id
}
fn provide(&self, _request: ContextRequest) -> ContextProviderFuture<'_> {
let modules = self.modules.clone();
Box::pin(async move { Ok(modules) })
}
}