Avoid panicking when truncating code with special chars
Showing the changes to crates (2 of 2 files). Show all
Unified · Split
| @@ -67,11 +67,13 @@ impl EmbeddingProvider for DummyEmbeddings { · 66 unmodified lines | ||
| 67 | 67 | } |
| 68 | 68 | } |
| 69 | 69 | |
| 70 | +const INPUT_LIMIT: usize = 8190; | |
| 71 | + | |
| 70 | 72 | impl OpenAIEmbeddings { |
| 71 | − async fn truncate(span: String) -> String { | |
| 73 | + fn truncate(span: String) -> String { | |
| 72 | 74 | let mut tokens = OPENAI_BPE_TOKENIZER.encode_with_special_tokens(span.as_ref()); |
| 73 | − if tokens.len() > 8190 { | |
| 74 | − tokens.truncate(8190); | |
| 75 | + if tokens.len() > INPUT_LIMIT { | |
| 76 | + tokens.truncate(INPUT_LIMIT); | |
| 75 | 77 | let result = OPENAI_BPE_TOKENIZER.decode(tokens.clone()); |
| 76 | 78 | if result.is_ok() { |
| 77 | 79 | let transformed = result.unwrap(); |
| @@ -80,7 +82,7 @@ impl OpenAIEmbeddings { · 2 unmodified lines | ||
| 80 | 82 | } |
| 81 | 83 | } |
| 82 | 84 | |
| 83 | − return span.to_string(); | |
| 85 | + span | |
| 84 | 86 | } |
| 85 | 87 | |
| 86 | 88 | async fn send_request(&self, api_key: &str, spans: Vec<&str>) -> Result<Response<AsyncBody>> { |
| @@ -137,7 +139,7 @@ impl EmbeddingProvider for OpenAIEmbeddings { · 50 unmodified lines | ||
| 137 | 139 | // Don't worry about delaying bad request, as we can assume |
| 138 | 140 | // we haven't been rate limited yet. |
| 139 | 141 | for span in spans.iter_mut() { |
| 140 | − *span = Self::truncate(span.to_string()).await; | |
| 142 | + *span = Self::truncate(span.to_string()); | |
| 141 | 143 | } |
| 142 | 144 | } |
| 143 | 145 | StatusCode::OK => { |
| @@ -63,7 +63,7 @@ impl CodeContextRetriever { · 62 unmodified lines | ||
| 63 | 63 | ) { |
| 64 | 64 | // log::info!("-----MATCH-----"); |
| 65 | 65 | |
| 66 | − let mut name: Vec<&str> = vec![]; | |
| 66 | + let mut name = Vec::new(); | |
| 67 | 67 | let mut item: Option<&str> = None; |
| 68 | 68 | let mut offset: Option<usize> = None; |
| 69 | 69 | for capture in mat.captures { |
| @@ -91,11 +91,8 @@ impl CodeContextRetriever { · 21 unmodified lines | ||
| 91 | 91 | .replace("<language>", &pending_file.language.name().to_lowercase()) |
| 92 | 92 | .replace("<item>", item.unwrap()); |
| 93 | 93 | |
| 94 | − let mut truncated_span = context_span.clone(); | |
| 95 | − truncated_span.truncate(100); | |
| 96 | − | |
| 97 | 94 | // log::info!("Name: {:?}", name); |
| 98 | − // log::info!("Span: {:?}", truncated_span); | |
| 95 | + // log::info!("Span: {:?}", util::truncate(&context_span, 100)); | |
| 99 | 96 | |
| 100 | 97 | context_spans.push(context_span); |
| 101 | 98 | documents.push(Document { |