better error message context window

dcSpark · Nov 16, 2024 · 3d7af09 · 3d7af09
1 parent 75a99cc
commit 3d7af09
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 0 deletions.
diff --git a/shinkai-bin/shinkai-node/src/llm_provider/error.rs b/shinkai-bin/shinkai-node/src/llm_provider/error.rs
@@ -85,6 +85,7 @@ pub enum LLMProviderError {
     ToolRetrievalError(String),
     ToolSearchError(String),
     AgentNotFound(String),
+    MessageTooLargeForLLM { max_tokens: usize, used_tokens: usize },
 }
 
 impl fmt::Display for LLMProviderError {
@@ -175,6 +176,9 @@ impl fmt::Display for LLMProviderError {
             LLMProviderError::ToolRetrievalError(s) => write!(f, "Tool retrieval error: {}", s),
             LLMProviderError::ToolSearchError(s) => write!(f, "Tool search error: {}", s),
             LLMProviderError::AgentNotFound(s) => write!(f, "Agent not found: {}", s),
+            LLMProviderError::MessageTooLargeForLLM { max_tokens, used_tokens } => {
+                write!(f, "Message too large for LLM: Used {} tokens, but the maximum allowed is {}.", used_tokens, max_tokens)
+            }
         }
     }
 }
@@ -255,6 +259,7 @@ impl LLMProviderError {
             LLMProviderError::ToolRetrievalError(_) => "ToolRetrievalError",
             LLMProviderError::ToolSearchError(_) => "ToolSearchError",
             LLMProviderError::AgentNotFound(_) => "AgentNotFound",
+            LLMProviderError::MessageTooLargeForLLM { .. } => "MessageTooLargeForLLM",
         };
 
         let error_message = format!("{}", self);

diff --git a/shinkai-bin/shinkai-node/src/llm_provider/providers/shared/ollama_api.rs b/shinkai-bin/shinkai-node/src/llm_provider/providers/shared/ollama_api.rs
@@ -56,6 +56,19 @@ pub fn ollama_prepare_messages(model: &LLMProviderInterface, prompt: Prompt) ->
         &ModelCapabilitiesManager::num_tokens_from_llama3,
     )?;
 
+    // Check if the prompt is not empty but messages are empty
+    if !prompt.sub_prompts.is_empty() && chat_completion_messages.is_empty() {
+        // Calculate the number of tokens used by the prompt
+        let (_, used_tokens) = prompt.generate_chat_completion_messages(
+            Some("tool".to_string()),
+            &ModelCapabilitiesManager::num_tokens_from_llama3,
+        );
+        return Err(LLMProviderError::MessageTooLargeForLLM {
+            max_tokens: max_input_tokens,
+            used_tokens,
+        });
+    }
+
     // Get a more accurate estimate of the number of used tokens
     let used_tokens = ModelCapabilitiesManager::num_tokens_from_llama3(&chat_completion_messages);
 

diff --git a/shinkai-bin/shinkai-node/src/managers/model_capabilities_manager.rs b/shinkai-bin/shinkai-node/src/managers/model_capabilities_manager.rs
@@ -438,6 +438,7 @@ impl ModelCapabilitiesManager {
             model_type if model_type.starts_with("qwen2.5:14b") => 128_000,
             model_type if model_type.starts_with("qwen2.5:32b") => 128_000,
             model_type if model_type.starts_with("qwen2.5:72b") => 128_000,
+            model_type if model_type.starts_with("qwen2.5-coder") => 128_000,
             model_type if model_type.starts_with("aya") => 32_000,
             model_type if model_type.starts_with("wizardlm2") => 8_000,
             model_type if model_type.starts_with("phi2") => 4_000,