From 978d4c54c0d5ac71fa0407d26ff06f16268e4ec3 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Fri, 10 Mar 2023 07:25:33 +0800
Subject: chore: opitmize code, move MAX_TOKENS related to config/message.rs

---
 src/config/message.rs | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

(limited to 'src/config/message.rs')

diff --git a/src/config/message.rs b/src/config/message.rs
index 514a2fa..3fd5935 100644
--- a/src/config/message.rs
+++ b/src/config/message.rs
@@ -1,6 +1,9 @@
+use crate::utils::count_tokens;
+
+use anyhow::{bail, Result};
 use serde::{Deserialize, Serialize};
 
-use crate::utils::count_tokens;
+pub const MAX_TOKENS: usize = 4096;
 
 #[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct Message {
@@ -25,6 +28,14 @@ pub enum MessageRole {
     User,
 }
 
+pub fn within_max_tokens_limit(messages: &[Message]) -> Result<()> {
+    let tokens = num_tokens_from_messages(messages);
+    if tokens >= MAX_TOKENS {
+        bail!("Exceed max tokens limit")
+    }
+    Ok(())
+}
+
 pub fn num_tokens_from_messages(messages: &[Message]) -> usize {
     let mut num_tokens = 0;
     for message in messages.iter() {
-- 
cgit v1.2.3