Merge pull request #159 from codestoryai/features/fix-how-prompt-gets…

…-constructed [sidecar] use proper calcluations for constructing the prompt
codestoryai · Oct 19, 2023 · c487481 · c487481
2 parents 46bc1d8 + 649b433
commit c487481
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 31 deletions.
diff --git a/src/.DS_Store b/src/.DS_Store
diff --git a/src/agent/search.rs b/src/agent/search.rs
@@ -804,38 +804,50 @@ impl Agent {
         match self.get_last_conversation_message_agent_state() {
             &AgentState::Explain => {
                 if let Some(user_selected_context_slice) = user_selected_context {
-                    prompt += "##### SELECTED CODE CONTEXT #####\n";
-                    for user_selected_context in user_selected_context_slice.iter() {
-                        let snippet = user_selected_context
-                            .data
-                            .lines()
-                            .enumerate()
-                            .map(|(i, line)| {
-                                format!(
-                                    "{} {line}\n",
-                                    i + user_selected_context.start_line as usize + 1
+                    let selected_code_context = "##### SELECTED CODE CONTEXT #####\n";
+                    let selected_code_header_tokens =
+                        bpe.encode_ordinary(&selected_code_context).len();
+                    if selected_code_header_tokens
+                        >= remaining_prompt_tokens - self.model.prompt_tokens_limit
+                    {
+                        info!("we can't set selected selection because of prompt limit");
+                    } else {
+                        prompt += "##### SELECTED CODE CONTEXT #####\n";
+                        remaining_prompt_tokens -= selected_code_header_tokens;
+
+                        for user_selected_context in user_selected_context_slice.iter() {
+                            let snippet = user_selected_context
+                                .data
+                                .lines()
+                                .enumerate()
+                                .map(|(i, line)| {
+                                    format!(
+                                        "{} {line}\n",
+                                        i + user_selected_context.start_line as usize + 1
+                                    )
+                                })
+                                .collect::<String>();
+
+                            let formatted_string = format!(
+                                "### {} ###\n{snippet}\n\n",
+                                self.get_absolute_path(
+                                    self.reporef(),
+                                    &user_selected_context.file_path
                                 )
-                            })
-                            .collect::<String>();
-
-                        let formatted_string = format!(
-                            "### {} ###\n{snippet}\n\n",
-                            self.get_absolute_path(
-                                self.reporef(),
-                                &user_selected_context.file_path
-                            )
-                        );
-
-                        let snippet_tokens = bpe.encode_ordinary(&formatted_string).len();
-                        if snippet_tokens
-                            >= remaining_prompt_tokens - self.model.prompt_tokens_limit
-                        {
-                            info!("breaking at {} tokens", remaining_prompt_tokens);
-                            break;
-                        }
+                            );
+
+                            let snippet_tokens = bpe.encode_ordinary(&formatted_string).len();
+                            if snippet_tokens
+                                >= remaining_prompt_tokens - self.model.prompt_tokens_limit
+                            {
+                                info!("breaking at {} tokens", remaining_prompt_tokens);
+                                break;
+                            }
+                            prompt += &formatted_string;
 
-                        // Make sure we are always in the context limit
-                        remaining_prompt_tokens -= snippet_tokens;
+                            // Make sure we are always in the context limit
+                            remaining_prompt_tokens -= snippet_tokens;
+                        }
                     }
                 }
             }

diff --git a/src/indexes/schema.rs b/src/indexes/schema.rs
@@ -192,7 +192,7 @@ impl CodeSnippet {
 
         let unique_hash = builder.add_text_field("unique_hash", STRING | STORED);
 
-        let repo_disk_path = dbg!(builder.add_text_field("repo_disk_path", STRING));
+        let repo_disk_path = builder.add_text_field("repo_disk_path", STRING);
         let repo_ref = builder.add_text_field("repo_ref", STRING | STORED);
         let repo_name = builder.add_text_field("repo_name", code_snippet_tokenizer.clone());
         let relative_path = builder.add_text_field("relative_path", code_snippet_tokenizer.clone());