docstrings in prompts!

app.py

@@ -49,7 +49,7 @@ canvas {
   <canvas id="canvas"></canvas>
   <div class="playpause">▶</div>
 </div>
-blank canvas here indicates that some of the shadertoy specific functions are not yet supported with this implementation (like #define I believe). you can always copy and paste the code into a shadertoy.com window to try.
+\nblank canvas here indicates that some of the shadertoy specific functions are not yet supported with this implementation (like #define I believe). you can always copy and paste the code into a shadertoy.com window to try.
 </body>
 <!--
 for most samples webgl-utils only provides shader compiling/linking and
@@ -227,7 +227,7 @@ main();
 
 def make_iframe(shader_code): #keep a single function?
      script = make_script(shader_code)
-     return f"""<iframe width="640" height="512" srcdoc=\'{script}\' allowfullscreen></iframe>"""
+     return f"""<iframe width="640" height="420" srcdoc=\'{script}\' allowfullscreen></iframe>"""
     
 
 intro_text = """
@@ -266,10 +266,11 @@ outro_text ="""
  - [x] generate whole functions (seems to work quite well)
  - [] dropdown for model selection (from curated list or all supported models?)
  - [] generation history stating which function and orig/generated returns. (use State ??). do it as comments in the code?
- - [] display errros/issues to the user (raise gr.Error could be one idea, but highlighting in the code would be awesome)
+ - [~] display errros/issues to the user (raise gr.Error could be one idea, but highlighting in the code would be awesome) currently adds a comment to the code.
  - [] generate whole shaders (via prompts guidance, recursive from errors)
  - [] accordion with generation parameters (as pipeline_kwargs?) look up starcoder playround and take "inspiration" from there
  - [] support FIM task for better model context
+ - [~] include some context for prompt (title, comments before a functions) - now works with the first comment inside a function body (has to be first)
  - [] gradio examples
 
 ### Notes:
@@ -310,10 +311,11 @@ def grab_sample(sample_idx):
 def _parse_functions(in_code):
     """
     returns all functions in the code as their actual nodes.
-
+    includes any comment made directly after the function definition or diretly after #copilot trigger
     """
     tree = parser.parse(bytes(in_code, "utf8"))
     funcs = [n for n in tree.root_node.children if n.type == "function_definition"]
+
     return funcs
 
 PIPE = None
@@ -411,12 +413,18 @@ def alter_body(old_code, func_id: str, funcs_list: list, pipeline=PIPE):
         pipeline = _make_pipeline("Vipitis/santacoder-finetuned-Shadertoys-fine")
 
     func_start_idx = _line_chr2char(old_code, func_node.start_point[0], func_node.start_point[1])
-    identifier_str = func_node.child_by_field_name("type").text.decode() + " " + func_node.child_by_field_name("declarator").text.decode()
+    identifier_str = func_node.child_by_field_name("type").text.decode() + " " + func_node.child_by_field_name("declarator").text.decode() #func_start_idx:body_start_idx?
     body_node = func_node.child_by_field_name("body")
     body_start_idx = _line_chr2char(old_code, body_node.start_point[0], body_node.start_point[1])
     body_end_idx = _line_chr2char(old_code, body_node.end_point[0], body_node.end_point[1])
     print(f"{old_code[body_start_idx:body_end_idx]=}")
-    model_context = identifier_str # just this
+    model_context = identifier_str # base case
+    # add any comments at the beginning of the function to the model_context
+    second_child = func_node.child_by_field_name("body").children[1] #might error out?
+    if second_child.type == "comment":
+        # print(second_child.text.decode())
+        model_context += second_child.text.decode()
+        print(f"{model_context=}")
     num_new_tokens = max(160,(body_end_idx - body_start_idx) + 10) #TODO: approximation, we do have early stopping? maybe also use a number instead? HARD MAX for performance limits.
     
     print(f"generating up to {num_new_tokens} after {model_context!r}")
@@ -429,13 +437,13 @@ def alter_body(old_code, func_id: str, funcs_list: list, pipeline=PIPE):
         first_gened_func = _parse_functions(id_with_generation)[0] # truncate generation to a single function?
     except IndexError:
         print("generation wasn't a full function.")
-        altered_code = old_code[:body_start_idx] + generation + "//the generation didn't complete the function!\n" + old_code[body_end_idx:] #needs a newline to break out of the comment.
+        altered_code = old_code[:func_start_idx] + model_context + generation + "//the generation didn't complete the function!\n" + old_code[body_end_idx:] #needs a newline to break out of the comment.
         return altered_code, pipeline
         # raise gr.Error(f"didn't generate a full function: {generation!r}]")
     print(f"{first_gened_func=}")
     generated_body = first_gened_func.child_by_field_name("body").text.decode()
     print(f"{generated_body=}")
-    altered_code = old_code[:body_start_idx] + generated_body + old_code[body_end_idx:]
+    altered_code = old_code[:func_start_idx] + model_context + generated_body + old_code[body_end_idx:]
     return altered_code, pipeline
 
 def add_history(func_id, orig_rtn, gened_rtn, history):