Custom Error Support for LoomError

src/lib.rs

@@ -48,6 +48,7 @@
 #![feature(once_cell_try)]
 
 use std::{
+	error::Error,
 	fmt::{Debug, Display},
 	marker::PhantomData,
 	str::FromStr,
@@ -73,11 +74,11 @@ mod mock;
 mod tests;
 
 pub use storage::TapestryChestHandler;
-use types::{LoomError, SummaryModelTokens, WeaveError};
+use types::{LoomError, SummaryModelTokens};
 
 use crate::types::{PromptModelTokens, WrapperRole};
 
-pub type Result<T> = std::result::Result<T, LoomError>;
+pub type Result<T, U> = std::result::Result<T, LoomError<U>>;
 
 /// Represents a unique identifier for any arbitrary entity.
 ///
@@ -162,6 +163,7 @@ pub trait Llm<T: Config>:
 	type Response: Clone + Into<Option<String>> + Send;
 	/// Type representing the parameters for a prompt.
 	type Parameters: Debug + Clone + Send + Sync;
+	type PromptError: Error;
 
 	/// The maximum number of tokens that can be processed at once by an LLM model.
 	fn max_context_length(&self) -> Self::Tokens;
@@ -178,7 +180,7 @@ pub trait Llm<T: Config>:
 	/// Calculates the number of tokens in a string.
 	///
 	/// This may vary depending on the type of tokens used by the LLM. In the case of ChatGPT, can be calculated using the [tiktoken-rs](https://github.com/zurawiki/tiktoken-rs#counting-token-length) crate.
-	fn count_tokens(content: &str) -> Result<Self::Tokens>;
+	fn count_tokens(content: &str) -> Result<Self::Tokens, T>;
 	/// Prompt LLM with the supplied messages and parameters.
 	async fn prompt(
 		&self,
@@ -187,7 +189,7 @@ pub trait Llm<T: Config>:
 		msgs: Vec<Self::Request>,
 		params: &Self::Parameters,
 		max_tokens: Self::Tokens,
-	) -> Result<Self::Response>;
+	) -> Result<Self::Response, T>;
 	/// Compute cost of a message based on model.
 	fn compute_cost(&self, prompt_tokens: Self::Tokens, response_tokens: Self::Tokens) -> f64;
 	/// Calculate the upperbound of tokens allowed for the current [`Config::PromptModel`] before a
@@ -301,12 +303,10 @@ impl<T: Config> TapestryFragment<T> {
 	/// Add a [`ContextMessage`] to the `context_messages` list.
 	///
 	/// Also increments the `context_tokens` by the number of tokens in the message.
-	fn push_message(&mut self, msg: ContextMessage<T>) -> Result<()> {
+	fn push_message(&mut self, msg: ContextMessage<T>) -> Result<(), T> {
 		let tokens = T::PromptModel::count_tokens(&msg.content)?;
 		let new_token_count = self.context_tokens.checked_add(&tokens).ok_or_else(|| {
-			LoomError::from(WeaveError::BadConfig(
-				"Number of tokens exceeds max tokens for model".to_string(),
-			))
+			LoomError::BadConfig("Number of tokens exceeds max tokens for model".to_string())
 		})?;
 
 		trace!("Pushing message: {:?}, new token count: {}", msg, new_token_count);
@@ -319,7 +319,7 @@ impl<T: Config> TapestryFragment<T> {
 	/// Add a [`ContextMessage`] to the `context_messages` list.
 	///
 	/// Also increments the `context_tokens` by the number of tokens in the message.
-	fn extend_messages(&mut self, msgs: Vec<ContextMessage<T>>) -> Result<()> {
+	fn extend_messages(&mut self, msgs: Vec<ContextMessage<T>>) -> Result<(), T> {
 		let total_new_tokens = msgs
 			.iter()
 			.map(|m| T::PromptModel::count_tokens(&m.content).unwrap())
@@ -332,9 +332,7 @@ impl<T: Config> TapestryFragment<T> {
 		trace!("Extending messages with token sum: {}", sum);
 
 		let new_token_count = self.context_tokens.checked_add(&sum).ok_or_else(|| {
-			LoomError::from(WeaveError::BadConfig(
-				"Number of tokens exceeds max tokens for model".to_string(),
-			))
+			LoomError::BadConfig("Number of tokens exceeds max tokens for model".to_string())
 		})?;
 
 		// Update the token count and messages only if all checks pass

src/loom.rs

@@ -6,7 +6,7 @@ use tracing::{debug, error, instrument, trace};
 use crate::{
 	types::{
 		LoomError, PromptModelRequest, PromptModelTokens, SummaryModelTokens, VecPromptMsgsDeque,
-		WeaveError, WrapperRole, ASSISTANT_ROLE, SYSTEM_ROLE,
+		WrapperRole, ASSISTANT_ROLE, SYSTEM_ROLE,
 	},
 	Config, ContextMessage, Llm, LlmConfig, TapestryChestHandler, TapestryFragment, TapestryId,
 };
@@ -52,7 +52,7 @@ impl<T: Config> Loom<T> {
 		tapestry_id: TID,
 		instructions: String,
 		mut msgs: Vec<ContextMessage<T>>,
-	) -> Result<(<<T as Config>::PromptModel as Llm<T>>::Response, u64, bool), LoomError> {
+	) -> Result<(<<T as Config>::PromptModel as Llm<T>>::Response, u64, bool), LoomError<T>> {
 		let instructions_ctx_msg =
 			Self::build_context_message(SYSTEM_ROLE.into(), instructions, None);
 		let instructions_req_msg: PromptModelRequest<T> = instructions_ctx_msg.clone().into();
@@ -150,7 +150,7 @@ impl<T: Config> Loom<T> {
 		trace!("Max completion tokens available: {:?}", max_completion_tokens);
 
 		if max_completion_tokens.is_zero() {
-			return Err(LoomError::from(WeaveError::MaxCompletionTokensIsZero).into());
+			return Err(LoomError::MaxCompletionTokensIsZero.into());
 		}
 
 		trace!("Prompting LLM with request messages");
@@ -208,7 +208,7 @@ impl<T: Config> Loom<T> {
 		summary_model_config: &LlmConfig<T, T::SummaryModel>,
 		tapestry_fragment: &TapestryFragment<T>,
 		summary_max_tokens: SummaryModelTokens<T>,
-	) -> Result<String, LoomError> {
+	) -> Result<String, LoomError<T>> {
 		trace!(
 			"Generating summary with max tokens: {:?}, for tapestry fragment: {:?}",
 			summary_max_tokens,

src/mock.rs

@@ -23,49 +23,52 @@ impl TapestryChestHandler<MockConfig> for MockChest {
 		_tapestry_id: &TID,
 		_tapestry_fragment: TapestryFragment<MockConfig>,
 		_increment: bool,
-	) -> crate::Result<u64> {
+	) -> crate::Result<u64, MockConfig> {
 		Ok(0)
 	}
 
 	async fn save_tapestry_metadata<TID: TapestryId, M: Debug + Clone + Send + Sync>(
 		&self,
 		_tapestry_id: TID,
 		_metadata: M,
-	) -> crate::Result<()> {
+	) -> crate::Result<(), MockConfig> {
 		Ok(())
 	}
 
 	async fn get_instance_index<TID: TapestryId>(
 		&self,
 		_tapestry_id: TID,
-	) -> crate::Result<Option<u16>> {
+	) -> crate::Result<Option<u16>, MockConfig> {
 		Ok(Some(0))
 	}
 
 	async fn get_tapestry_fragment<TID: TapestryId>(
 		&self,
 		_tapestry_id: TID,
 		_instance: Option<u64>,
-	) -> crate::Result<Option<TapestryFragment<MockConfig>>> {
+	) -> crate::Result<Option<TapestryFragment<MockConfig>>, MockConfig> {
 		Ok(Some(TapestryFragment { context_tokens: 0, context_messages: vec![] }))
 	}
 
 	async fn get_tapestry_metadata<TID: TapestryId, M: DeserializeOwned>(
 		&self,
 		_tapestry_id: TID,
-	) -> crate::Result<Option<M>> {
+	) -> crate::Result<Option<M>, MockConfig> {
 		Ok(Some(serde_json::from_str("{}").unwrap()))
 	}
 
-	async fn delete_tapestry<TID: TapestryId>(&self, _tapestry_id: TID) -> crate::Result<()> {
+	async fn delete_tapestry<TID: TapestryId>(
+		&self,
+		_tapestry_id: TID,
+	) -> crate::Result<(), MockConfig> {
 		Ok(())
 	}
 
 	async fn delete_tapestry_fragment<TID: TapestryId>(
 		&self,
 		_tapestry_id: TID,
 		_instance: Option<u64>,
-	) -> crate::Result<()> {
+	) -> crate::Result<(), MockConfig> {
 		Ok(())
 	}
 }
@@ -104,17 +107,14 @@ impl Llm<MockConfig> for MockLlm {
 	type Parameters = ();
 	type Request = MockLlmRequest;
 	type Response = MockLlmResponse;
+	type PromptError = MockPromptError;
 
-	fn count_tokens(content: &str) -> Result<Self::Tokens> {
+	fn count_tokens(content: &str) -> Result<Self::Tokens, MockConfig> {
 		let bpe = p50k_base().unwrap();
 		let tokens = bpe.encode_with_special_tokens(&content.to_string());
 
 		tokens.len().try_into().map_err(|_| {
-			LoomError::from(WeaveError::BadConfig(format!(
-				"Number of tokens exceeds max tokens for model: {}",
-				content
-			)))
-			.into()
+			LoomError::Llm(MockPromptError::BadConfig("Token count exceeds u16".to_string()))
 		})
 	}
 
@@ -133,7 +133,7 @@ impl Llm<MockConfig> for MockLlm {
 		_msgs: Vec<Self::Request>,
 		_params: &Self::Parameters,
 		_max_tokens: Self::Tokens,
-	) -> Result<Self::Response> {
+	) -> Result<Self::Response, MockConfig> {
 		Ok(MockLlmResponse {})
 	}
 
@@ -203,3 +203,9 @@ impl From<MockLlmResponse> for Option<String> {
 		Some("TestLlmResponse".to_string())
 	}
 }
+
+#[derive(Debug, Clone, thiserror::Error)]
+pub enum MockPromptError {
+	#[error("Bad configuration: {0}")]
+	BadConfig(String),
+}

src/storage/mod.rs

@@ -41,22 +41,22 @@ pub trait TapestryChestHandler<T: Config> {
 		tapestry_id: &TID,
 		tapestry_fragment: TapestryFragment<T>,
 		increment: bool,
-	) -> crate::Result<u64>;
+	) -> crate::Result<u64, T>;
 	/// Save tapestry metadata.
 	///
 	/// Based on application use cases, you can add aditional data for a given [`TapestryId`]
 	async fn save_tapestry_metadata<TID: TapestryId, M: Serialize + Debug + Clone + Send + Sync>(
 		&self,
 		tapestry_id: TID,
 		metadata: M,
-	) -> crate::Result<()>;
+	) -> crate::Result<(), T>;
 	/// Retrieves the index of a tapestry.
 	///
 	/// Returns None if the tapestry does not exist.
 	async fn get_instance_index<TID: TapestryId>(
 		&self,
 		tapestry_id: TID,
-	) -> crate::Result<Option<u16>>;
+	) -> crate::Result<Option<u16>, T>;
 	/// Retrieves the last tapestry fragment, or a fragment at a specified instance.
 	///
 	/// # Parameters
@@ -73,18 +73,18 @@ pub trait TapestryChestHandler<T: Config> {
 		&self,
 		tapestry_id: TID,
 		instance: Option<u64>,
-	) -> crate::Result<Option<TapestryFragment<T>>>;
+	) -> crate::Result<Option<TapestryFragment<T>>, T>;
 	/// Retrieves the last tapestry metadata, or a metadata at a specified instance.
 	async fn get_tapestry_metadata<TID: TapestryId, M: DeserializeOwned + Send + Sync>(
 		&self,
 		tapestry_id: TID,
-	) -> crate::Result<Option<M>>;
+	) -> crate::Result<Option<M>, T>;
 	/// Deletes a tapestry and all its instances.
-	async fn delete_tapestry<TID: TapestryId>(&self, tapestry_id: TID) -> crate::Result<()>;
+	async fn delete_tapestry<TID: TapestryId>(&self, tapestry_id: TID) -> crate::Result<(), T>;
 	/// Deletes a tapestry fragment.
 	async fn delete_tapestry_fragment<TID: TapestryId>(
 		&self,
 		tapestry_id: TID,
 		instance: Option<u64>,
-	) -> crate::Result<()>;
+	) -> crate::Result<(), T>;
 }