Fix bug with system instructions not being passed

llm.py

@@ -38,9 +38,12 @@ Again, please generate outputs for these placeholders: {placeholders}
 """.strip()
 
 
-def _make_model(model_name: str, temperature: float) -> genai.GenerativeModel:
+def _make_model(
+  model_name: str, system_instruction: str = "", temperature: float = 1.0
+) -> genai.GenerativeModel:
   return genai.GenerativeModel(
     model_name,
+    system_instruction=system_instruction,
     generation_config={
       "temperature": temperature,
       "top_p": 0.95,
@@ -51,7 +54,7 @@ def _make_model(model_name: str, temperature: float) -> genai.GenerativeModel:
 
 
 def generate_prompt(task_description: str, model_name: str, temperature: float) -> str:
-  model = _make_model(model_name, temperature)
+  model = _make_model(model_name, temperature=temperature)
   prompt = _GENERATE_PROMPT.format(task=task_description)
   return model.generate_content(prompt).text
 
@@ -59,7 +62,7 @@ def generate_prompt(task_description: str, model_name: str, temperature: float)
 def generate_variables(
   prompt: str, variable_names: list[str], model_name: str, temperature: float
 ) -> dict[str, str]:
-  model = _make_model(model_name, temperature)
+  model = _make_model(model_name, temperature=temperature)
   output = (
     model.generate_content(
       _GENERATE_VARIABLES_PROMPT.format(placeholders=", ".join(variable_names))
@@ -70,6 +73,8 @@ def generate_variables(
   return json.loads(output)
 
 
-def run_prompt(prompt_with_variables: str, model_name: str, temperature: float) -> str:
-  model = _make_model(model_name, temperature)
-  return model.generate_content(prompt_with_variables).text
+def run_prompt(
+  prompt_with_variables: str, system_instruction: str, model_name: str, temperature: float
+) -> str:
+  model = _make_model(model_name, temperature=temperature, system_instruction=system_instruction)
+  return model.generate_content(prompt_with_variables, request_options={"timeout": 120}).text

main.py

@@ -163,7 +163,6 @@ def on_click_system_instructions_header(e: me.ClickEvent):
 
 def on_click_eval_run(e: me.ClickEvent):
   state = me.state(State)
-  print(e.key)
   _, prompt_version, response_index, selected_prompt_response_index = e.key.split("_")
   prompt = find_prompt(state.prompts, int(prompt_version))
   selected_prompt = find_prompt(state.prompts, state.version)
@@ -183,7 +182,9 @@ def on_click_eval_run(e: me.ClickEvent):
   prompt_text = prompt.prompt
   for name, value in response["variables"].items():
     prompt_text = prompt_text.replace("{{" + name + "}}", value)
-  response["output"] = llm.run_prompt(prompt_text, prompt.model, prompt.model_temperature)
+  response["output"] = llm.run_prompt(
+    prompt_text, prompt.system_instructions, prompt.model, prompt.model_temperature
+  )
 
 
 def on_click_run(e: me.ClickEvent):
@@ -228,7 +229,9 @@ def on_click_run(e: me.ClickEvent):
   prompt = state.prompt
   for name, value in prompt_variables.items():
     prompt = prompt.replace("{{" + name + "}}", value)
-  state.response = llm.run_prompt(prompt, state.model, state.model_temperature)
+  state.response = llm.run_prompt(
+    prompt, state.system_instructions, state.model, state.model_temperature
+  )
   state.prompts[-1].responses.append(dict(output=state.response, variables=prompt_variables))