Improve content generation prompt to reduce over-generation (#16333)

I focused on cases where we're inserting doc comments or annotations
above symbols.

I added 5 new examples to the content generation prompt, covering
various scenarios:

1. Inserting documentation for a Rust struct
2. Writing docstrings for a Python class
3. Adding comments to a TypeScript method
4. Adding a derive attribute to a Rust struct
5. Adding a decorator to a Python class

These examples demonstrate how to handle different languages and common
tasks like adding documentation, attributes, and decorators.

To improve context integration, I've made the following changes:

1. Added a `transform_context_range` that includes 3 lines before and
after the transform range
2. Introduced `rewrite_section_prefix` and `rewrite_section_suffix` to
provide more context around the section being rewritten
3. Updated the prompt template to include this additional context in a
separate code snippet

Release Notes:

- Reduced instances of over-generation when inserting docs or
annotations above a symbol.

crates/assistant/src/inline_assistant.rs

@@ -45,6 +45,7 @@ use std::{
     task::{self, Poll},
     time::{Duration, Instant},
 };
+use text::OffsetRangeExt as _;
 use theme::ThemeSettings;
 use ui::{prelude::*, CheckboxWithLabel, IconButtonShape, Popover, Tooltip};
 use util::{RangeExt, ResultExt};
@@ -2354,6 +2355,15 @@ impl Codegen {
             return Err(anyhow::anyhow!("invalid transformation range"));
         };
 
+        let mut transform_context_range = transform_range.to_point(&transform_buffer);
+        transform_context_range.start.row = transform_context_range.start.row.saturating_sub(3);
+        transform_context_range.start.column = 0;
+        transform_context_range.end =
+            (transform_context_range.end + Point::new(3, 0)).min(transform_buffer.max_point());
+        transform_context_range.end.column =
+            transform_buffer.line_len(transform_context_range.end.row);
+        let transform_context_range = transform_context_range.to_offset(&transform_buffer);
+
         let selected_ranges = self
             .selected_ranges
             .iter()
@@ -2376,6 +2386,7 @@ impl Codegen {
                 transform_buffer,
                 transform_range,
                 selected_ranges,
+                transform_context_range,
             )
             .map_err(|e| anyhow::anyhow!("Failed to generate content prompt: {}", e))?;
 

crates/assistant/src/prompts.rs

@@ -16,7 +16,9 @@ pub struct ContentPromptContext {
     pub document_content: String,
     pub user_prompt: String,
     pub rewrite_section: String,
-    pub rewrite_section_with_selections: String,
+    pub rewrite_section_prefix: String,
+    pub rewrite_section_suffix: String,
+    pub rewrite_section_with_edits: String,
     pub has_insertion: bool,
     pub has_replacement: bool,
 }
@@ -173,6 +175,7 @@ impl PromptBuilder {
         buffer: BufferSnapshot,
         transform_range: Range<usize>,
         selected_ranges: Vec<Range<usize>>,
+        transform_context_range: Range<usize>,
     ) -> Result<String, RenderError> {
         let content_type = match language_name {
             None | Some("Markdown" | "Plain Text") => "text",
@@ -202,6 +205,7 @@ impl PromptBuilder {
         for chunk in buffer.text_for_range(truncated_before) {
             document_content.push_str(chunk);
         }
+
         document_content.push_str("<rewrite_this>\n");
         for chunk in buffer.text_for_range(transform_range.clone()) {
             document_content.push_str(chunk);
@@ -217,7 +221,17 @@ impl PromptBuilder {
             rewrite_section.push_str(chunk);
         }
 
-        let rewrite_section_with_selections = {
+        let mut rewrite_section_prefix = String::new();
+        for chunk in buffer.text_for_range(transform_context_range.start..transform_range.start) {
+            rewrite_section_prefix.push_str(chunk);
+        }
+
+        let mut rewrite_section_suffix = String::new();
+        for chunk in buffer.text_for_range(transform_range.end..transform_context_range.end) {
+            rewrite_section_suffix.push_str(chunk);
+        }
+
+        let rewrite_section_with_edits = {
             let mut section_with_selections = String::new();
             let mut last_end = 0;
             for selected_range in &selected_ranges {
@@ -254,7 +268,9 @@ impl PromptBuilder {
             document_content,
             user_prompt,
             rewrite_section,
-            rewrite_section_with_selections,
+            rewrite_section_prefix,
+            rewrite_section_suffix,
+            rewrite_section_with_edits,
             has_insertion,
             has_replacement,
         };