create a prompt store for registration

genai/generate.py

@@ -2,30 +2,7 @@
 
 import openai
 
-NOTEBOOK_CREATE_NEXT_CELL_PROCLAMATION = """
-As a coding assistant, your task is to help users write code in Python within Jupyter Notebooks. Provide comments and code for the user to read and edit, ensuring it can be run successfully. The user will be able to run the code in the cell and see the output.
-
-When the user is interacting with you their message will start with `%%assist`. Otherwise, they are running commands and getting output from the system.
-
-You can use markdown to format your response. For example, to create a code block, use
-
-```python
-# code
-```
-
-""".strip() # noqa: E501
-
-
-NOTEBOOK_ERROR_DIAGNOSER_PROCLAMATION = """
-As a coding assistant, you'll diagnose errors in Python code written in a Jupyter Notebook. Format your response using markdown. Making sure to include the language around code blocks, like
-
-```python
-# code
-```
-
-Provide concise code examples in your response which will be rendered in Markdown in the notebook. The user will not be able to respond to your response.
-""".strip() # noqa: E501
-
+from genai.prompts import PromptStore
 
 Completion = TypedDict(
  "Completion",
@@ -69,7 +46,7 @@ def deltas(completion: Iterator[StreamCompletion]) -> Iterator[str]:
  yield delta["content"]
 
 
-def generate_next_cell(
+def generate_next_from_history(
  context: List[Dict[str, str]],
  text: str,
  stream: bool = False,
@@ -80,7 +57,7 @@ def generate_next_cell(
  # Establish the context of the conversation
  {
  "role": "system",
- "content": NOTEBOOK_CREATE_NEXT_CELL_PROCLAMATION,
+ "content": PromptStore.assist_prompt,
  },
  # Presumably In, Out
  *context,
@@ -118,7 +95,7 @@ def generate_exception_suggestion(
  messages.append(
  {
  "role": "system",
- "content": NOTEBOOK_ERROR_DIAGNOSER_PROCLAMATION,
+ "content": PromptStore.exception_prompt,
  },
  )
 

genai/magics.py

@@ -6,7 +6,7 @@
 
 from genai.context import PastAssists, build_context
 from genai.display import GenaiMarkdown, Stage, can_handle_display_updates
-from genai.generate import generate_next_cell
+from genai.generate import generate_next_from_history
 from genai.tokens import trim_messages_to_fit_token_limit
 
 
@@ -104,6 +104,6 @@ def assist(line, cell):
  print("submission:", cell)
  print("messages:", messages)
 
- gm.consume(generate_next_cell(messages, cell_text, stream=stream))
+ gm.consume(generate_next_from_history(messages, cell_text, stream=stream))
 
  gm.stage = Stage.FINISHED

genai/prompts.py

@@ -0,0 +1,70 @@
+'''
+There are currently two main prompts, designed for use in Notebooks:
+
+* Coding assistance
+* Exception diagnosis
+
+The defaults used will come from the GlobalPromptStore. The user can override
+these defaults by passing in a custom prompt to set_exception_prompt or set_assist_prompt.
+
+>>> from genai.prompts import set_exception_prompt
+>>> set_exception_prompt("You will get an exception. Do something with it.")
+
+'''
+
+NOTEBOOK_ASSISTANCE_PROMPT = """
+As a coding assistant, your task is to help users write code in Python within Jupyter Notebooks. Provide comments and code for the user to read and edit, ensuring it can be run successfully. The user will be able to run the code in the cell and see the output.
+
+When the user is interacting with you their message will start with `%%assist`. Otherwise, they are running commands and getting output from the system.
+
+You can use markdown to format your response. For example, to create a code block, use
+
+```python
+# code
+```
+
+""".strip() # noqa: E501
+
+
+NOTEBOOK_ERROR_DIAGNOSER_PROCLAMATION = """
+As a coding assistant, you'll diagnose errors in Python code written in a Jupyter Notebook. Format your response using markdown. Making sure to include the language around code blocks, like
+
+```python
+# code
+```
+
+Provide concise code examples in your response which will be rendered in Markdown in the notebook. The user will not be able to respond to your response.
+""".strip() # noqa: E501
+
+
+class DefaultPromptStore:
+ def __init__(self):
+ self._assist_prompt = NOTEBOOK_ASSISTANCE_PROMPT
+ self._exception_prompt = NOTEBOOK_ERROR_DIAGNOSER_PROCLAMATION
+
+ @property
+ def assist_prompt(self) -> str:
+ return self._assist_prompt
+
+ @assist_prompt.setter
+ def assist_prompt(self, prompt: str) -> None:
+ self._assist_prompt = prompt
+
+ @property
+ def exception_prompt(self) -> str:
+ return self._exception_prompt
+
+ @exception_prompt.setter
+ def exception_prompt(self, prompt: str) -> None:
+ self._exception_prompt = prompt
+
+
+PromptStore = DefaultPromptStore()
+
+
+def set_assist_prompt(prompt: str) -> None:
+ PromptStore.assist_prompt = prompt
+
+
+def set_exception_prompt(prompt: str) -> None:
+ PromptStore.exception_prompt = prompt

tests/test_magics.py

@@ -1,6 +1,6 @@
 from unittest import mock
 
-from genai import generate
+from genai.prompts import PromptStore
 from genai.context import PastAssists
 
 
@@ -27,7 +27,7 @@ def test_assist_magic(create, ip):
  messages=[
  {
  "role": "system",
- "content": generate.NOTEBOOK_CREATE_NEXT_CELL_PROCLAMATION,
+ "content": PromptStore.assist_prompt,
  },
  {
  "role": "user",
@@ -66,7 +66,7 @@ def test_assist_magic_with_args(create, ip):
  messages=[
  {
  "role": "system",
- "content": generate.NOTEBOOK_CREATE_NEXT_CELL_PROCLAMATION,
+ "content": PromptStore.assist_prompt,
  },
  {
  "role": "user",
@@ -108,7 +108,7 @@ def test_assist_magic_with_fresh_arg(create, ip):
  messages=[
  {
  "role": "system",
- "content": generate.NOTEBOOK_CREATE_NEXT_CELL_PROCLAMATION,
+ "content": PromptStore.assist_prompt,
  },
  # Note that there is zero other context, due to running with --fresh
  {

tests/test_suggestions.py

@@ -1,9 +1,10 @@
 import sys
 from unittest import mock
 
-from genai import generate, suggestions
+from genai import suggestions
 from genai.context import PastErrors
 from genai.suggestions import can_handle_display_updates
+from genai.prompts import PromptStore
 
 
 def test_register():
@@ -58,7 +59,7 @@ def test_custom_exc(create, display, ip):
  assert len(kwargs["messages"]) == 3
  assert kwargs["messages"][0] == {
  "role": "system",
- "content": generate.NOTEBOOK_ERROR_DIAGNOSER_PROCLAMATION,
+ "content": PromptStore.exception_prompt,
  }
  assert kwargs["messages"][1] == {
  "role": "user",
@@ -131,7 +132,7 @@ def test_custom_exc_fallback_on_In(create, display, ip):
  assert len(kwargs["messages"]) == 3
  assert kwargs["messages"][0] == {
  "role": "system",
- "content": generate.NOTEBOOK_ERROR_DIAGNOSER_PROCLAMATION,
+ "content": PromptStore.exception_prompt,
  }
  assert kwargs["messages"][1] == {
  "role": "user",
@@ -201,7 +202,7 @@ def test_custom_exc_long_traceback(create, display, ip):
  assert len(kwargs["messages"]) == 3
  assert kwargs["messages"][0] == {
  "role": "system",
- "content": generate.NOTEBOOK_ERROR_DIAGNOSER_PROCLAMATION,
+ "content": PromptStore.exception_prompt,
  }
  assert kwargs["messages"][1] == {
  "role": "user",