use std::collections::BTreeMap;
use crate::{
AgentOutput, BoxError, ContentPart, Document, Documents, FunctionDefinition, Json, Message,
Resource,
};
pub trait CompletionFeatures: Sized {
fn completion(
&self,
req: CompletionRequest,
resources: Vec<Resource>,
) -> impl Future<Output = Result<AgentOutput, BoxError>> + Send;
fn model_name(&self) -> String;
}
#[derive(Debug, Clone, Default)]
pub struct CompletionRequest {
pub instructions: String,
pub role: Option<String>,
pub chat_history: Vec<Message>,
pub raw_history: Vec<Json>,
pub documents: Documents,
pub prompt: String,
pub content: Vec<ContentPart>,
pub tools: Vec<FunctionDefinition>,
pub tool_choice_required: bool,
pub temperature: Option<f64>,
pub max_output_tokens: Option<usize>,
pub output_schema: Option<Json>,
pub stop: Option<Vec<String>>,
pub model: Option<String>,
}
impl CompletionRequest {
pub fn context(mut self, id: String, text: String) -> Self {
self.documents.docs.push(Document {
content: text.into(),
metadata: BTreeMap::from([("id".to_string(), id.into())]),
});
self
}
pub fn append_documents(mut self, docs: Documents) -> Self {
self.documents.docs.extend(docs.docs);
self
}
pub fn append_tools(mut self, tools: Vec<FunctionDefinition>) -> Self {
self.tools.extend(tools);
self
}
}