| @@ -55,6 +55,7 @@ struct OpenAIEmbeddingUsage { · 54 unmodified lines |
| 55 | 55 | pub trait EmbeddingProvider: Sync + Send { |
| 56 | 56 | async fn embed_batch(&self, spans: Vec<&str>) -> Result<Vec<Vec<f32>>>; |
| 57 | 57 | fn count_tokens(&self, span: &str) -> usize; |
| 58 | + fn should_truncate(&self, span: &str) -> bool; |
| 58 | 59 | // fn truncate(&self, span: &str) -> Result<&str>; |
| 59 | 60 | } |
| 60 | 61 | |
| @@ -74,6 +75,20 @@ impl EmbeddingProvider for DummyEmbeddings { · 13 unmodified lines |
| 74 | 75 | let tokens = OPENAI_BPE_TOKENIZER.encode_with_special_tokens(span); |
| 75 | 76 | tokens.len() |
| 76 | 77 | } |
| 78 | + |
| 79 | + fn should_truncate(&self, span: &str) -> bool { |
| 80 | + self.count_tokens(span) > OPENAI_INPUT_LIMIT |
| 81 | + |
| 82 | + // let tokens = OPENAI_BPE_TOKENIZER.encode_with_special_tokens(span); |
| 83 | + // let Ok(output) = { |
| 84 | + // if tokens.len() > OPENAI_INPUT_LIMIT { |
| 85 | + // tokens.truncate(OPENAI_INPUT_LIMIT); |
| 86 | + // OPENAI_BPE_TOKENIZER.decode(tokens) |
| 87 | + // } else { |
| 88 | + // Ok(span) |
| 89 | + // } |
| 90 | + // }; |
| 91 | + } |
| 77 | 92 | } |
| 78 | 93 | |
| 79 | 94 | const OPENAI_INPUT_LIMIT: usize = 8190; |
| @@ -125,6 +140,10 @@ impl EmbeddingProvider for OpenAIEmbeddings { · 45 unmodified lines |
| 125 | 140 | tokens.len() |
| 126 | 141 | } |
| 127 | 142 | |
| 143 | + fn should_truncate(&self, span: &str) -> bool { |
| 144 | + self.count_tokens(span) > OPENAI_INPUT_LIMIT |
| 145 | + } |
| 146 | + |
| 128 | 147 | async fn embed_batch(&self, spans: Vec<&str>) -> Result<Vec<Vec<f32>>> { |
| 129 | 148 | const BACKOFF_SECONDS: [usize; 4] = [3, 5, 15, 45]; |
| 130 | 149 | const MAX_RETRIES: usize = 4; |