From 15633a2c35e68a1b90e0e9fe4406e6f5ac8ba7cd Mon Sep 17 00:00:00 2001 From: Nico Arqueros Date: Tue, 22 Oct 2024 22:44:43 -0500 Subject: [PATCH] fix --- .../src/lance_db/ollama_embedding_fn.rs | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/shinkai-libs/shinkai-lancedb/src/lance_db/ollama_embedding_fn.rs b/shinkai-libs/shinkai-lancedb/src/lance_db/ollama_embedding_fn.rs index cb526539f..07a562cec 100644 --- a/shinkai-libs/shinkai-lancedb/src/lance_db/ollama_embedding_fn.rs +++ b/shinkai-libs/shinkai-lancedb/src/lance_db/ollama_embedding_fn.rs @@ -36,9 +36,19 @@ impl OllamaEmbeddingFunction { } }; + // Get the maximum token count for the model + let max_tokens = self.model_type.max_input_token_count(); + + // Truncate the prompt if it exceeds the maximum token count + let truncated_prompt = if prompt.len() > max_tokens { + &prompt[..max_tokens] + } else { + prompt + }; + let request_body = serde_json::json!({ "model": model_str, - "prompt": prompt + "prompt": truncated_prompt // Use the truncated prompt }); let full_url = if self.api_url.ends_with('/') {