| @@ -55,6 +55,7 @@ struct OpenAIEmbeddingUsage { · 54 unmodified lines |
| 55 | pub trait EmbeddingProvider: Sync + Send { | 55 | pub trait EmbeddingProvider: Sync + Send { |
| 56 | async fn embed_batch(&self, spans: Vec<&str>) -> Result<Vec<Vec<f32>>>; | 56 | async fn embed_batch(&self, spans: Vec<&str>) -> Result<Vec<Vec<f32>>>; |
| 57 | fn count_tokens(&self, span: &str) -> usize; | 57 | fn count_tokens(&self, span: &str) -> usize; |
| | 58 | fn should_truncate(&self, span: &str) -> bool; |
| 58 | // fn truncate(&self, span: &str) -> Result<&str>; | 59 | // fn truncate(&self, span: &str) -> Result<&str>; |
| 59 | } | 60 | } |
| 60 | | 61 | |
| @@ -74,6 +75,20 @@ impl EmbeddingProvider for DummyEmbeddings { · 13 unmodified lines |
| 74 | let tokens = OPENAI_BPE_TOKENIZER.encode_with_special_tokens(span); | 75 | let tokens = OPENAI_BPE_TOKENIZER.encode_with_special_tokens(span); |
| 75 | tokens.len() | 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 | const OPENAI_INPUT_LIMIT: usize = 8190; | 94 | const OPENAI_INPUT_LIMIT: usize = 8190; |
| @@ -125,6 +140,10 @@ impl EmbeddingProvider for OpenAIEmbeddings { · 45 unmodified lines |
| 125 | tokens.len() | 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 | async fn embed_batch(&self, spans: Vec<&str>) -> Result<Vec<Vec<f32>>> { | 147 | async fn embed_batch(&self, spans: Vec<&str>) -> Result<Vec<Vec<f32>>> { |
| 129 | const BACKOFF_SECONDS: [usize; 4] = [3, 5, 15, 45]; | 148 | const BACKOFF_SECONDS: [usize; 4] = [3, 5, 15, 45]; |
| 130 | const MAX_RETRIES: usize = 4; | 149 | const MAX_RETRIES: usize = 4; |