Preserve additional context ordering between Rails and AIGW for Code Generations

Context

The Language Server aims to order the additional context sent to Rails or AIGW according to relevance. For example, the Open Tabs context are ordered according the most recently opened files. Later on, as we introduce more additional contexts (e.g. Imports), the Language Server will be responsible for ordering the combined list of contexts. For example, taking Open Tabs and Imports, the files that are in both lists would have priority in terms of ordering.

This ordering is introduced because request payloads have limits. For instance, Code Suggestions payloads sent from the Language Server have a 50KB limit.

In Code Generation, the additional contexts are divided into related_files and related_snippets in Rails, and then sent to AIGW as separate payload fields (see code). This currently does not affect the overall workflow since the AIGW does not limit Code Generation payloads.

However, if in the future, AIGW does introduce limits to Code Generation payloads, separating the additional context related_files and related_snippets in Rails would mean that Rails is no longer respecting the relevance ordering set by the Language Server.

References and Resources

Proposal

Update Rails and AIGW so that the relevance ordering set by the Language Server is preserved up to the request sent to AIGW.

In AIGW

  1. Update the Code Generation API so that it accepts a generic prompt_enhancer.related_content payload instead of separate prompt_enhancer.related_files and prompt_enhancer.related_snippets.

  2. Once the payload has "passed" the AIGW API checks, ie it has reached the API code generation function, we can now update the prompt_enhancer so that related_content is split into related_files and related_snippets as needed by the prompt template.

    Note: the prompt sent to Anthropic Claude is accompanied by instructions, it should be okay to split the additional context into related_files and related_snippet at this point in the workflow. What's important is that, between LS -> Rails -> AIGW, the additional context ordering does not change because this could be subjected to payload limits. We need to make sure that the most relevant additional contexts are prioritized.

Rails

Update the Code Generation payload so that Rails sends a prompt_enhancer.related_content parameter instead of a prompt_enhancer.related_files and related_snippets.

Additional recommended change

Currently, Rails is responsible for wrapping the related_files and related_snippets in tags (see code). We can change this so that the tags are only introduced in the AIGW prompt template, i.e.:

AIGW prompt template change diff

This is assuming that AIGW has already transformed the prompt_enhancer.related_content payload to related_files and related_snippets.

diff --git a/ai_gateway/prompts/definitions/code_suggestions/generations/jinja_template/dotcom/system.jinja b/ai_gateway/prompts/definitions/code_suggestions/generations/jinja_template/dotcom/system.jinja
index 6e343d3a..e8735481 100644
--- a/ai_gateway/prompts/definitions/code_suggestions/generations/jinja_template/dotcom/system.jinja
+++ b/ai_gateway/prompts/definitions/code_suggestions/generations/jinja_template/dotcom/system.jinja
@@ -34,13 +34,21 @@ Please use existing functions from these files and code snippets if possible whe
 
 {%- if related_files %}
 <related_files>
-{{("\n").join(related_files)}}
+{%- for related_file in related_files %}
+<file_content file_name="{% related_file.name %}">
+{% related_file.content %}
+</related_file>
+{% endfor %}
 </related_files>
 {%- endif %}
 
 {%- if related_snippets %}
 <related_snippets>
-{{("\n").join(related_snippets)}}
+{%- for related_snippet in related_snippets %}
+<file_content file_name="{% related_snippet.name %}">
+{% related_snippet.content %}
+</related_snippet>
+{% endfor %}
 </related_snippets>
 {%- endif %}
Rails payload change diff 
diff --git a/ee/lib/code_suggestions/prompts/code_generation/ai_gateway_messages.rb b/ee/lib/code_suggestions/prompts/code_generation/ai_gateway_messages.rb
index dcd4db77509f..933821be218d 100644
--- a/ee/lib/code_suggestions/prompts/code_generation/ai_gateway_messages.rb
+++ b/ee/lib/code_suggestions/prompts/code_generation/ai_gateway_messages.rb
@@ -67,29 +67,7 @@ def existing_code_block_params
         end
 
         def context_block_params
-          related_files = []
-          related_snippets = []
-
-          params[:context]&.each do |context|
-            if context[:type] == ::Ai::AdditionalContext::CODE_SUGGESTIONS_CONTEXT_TYPES[:file]
-              related_files << <<~FILE_CONTENT
-              <file_content file_name="#{context[:name]}">
-              #{context[:content]}
-              </file_content>
-              FILE_CONTENT
-            elsif context[:type] == ::Ai::AdditionalContext::CODE_SUGGESTIONS_CONTEXT_TYPES[:snippet]
-              related_snippets << <<~SNIPPET_CONTENT
-              <snippet_content name="#{context[:name]}">
-              #{context[:content]}
-              </snippet_content>
-              SNIPPET_CONTENT
-            end
-          end
-
-          {
-            related_files: related_files,
-            related_snippets: related_snippets
-          }
+          { related_content: params[:context] }
         end
 
         def libraries_block_params
Edited by 🤖 GitLab Bot 🤖