Initial commit

app.py

@@ -0,0 +1,253 @@
+"""
+Cognitive Debriefing App - Respondent Interface
+
+Author: Dr Musashi Hinck
+
+
+Respondent-facing app. Reads arguments from request (in form of shareable link)
+
+TODO:
+- uuid from request for wandb
+
+"""
+from __future__ import annotations
+
+import logging
+import json
+import wandb
+import gradio as gr
+import openai
+
+from base64 import urlsafe_b64decode
+
+logger = logging.getLogger(__name__)
+
+from utils import PromptTemplate, convert_gradio_to_openai
+
+
+# %% Initialization
+client = openai.OpenAI()
+
+
+# %% (functions)
+def decode_config(config_dta: str) -> dict[str, str | float]:
+    "Read base64_url encoded json and loads into configuration"
+    config_str: str = urlsafe_b64decode(config_dta)
+    config: dict = json.loads(config_str)
+    return config
+
+
+def load_config(request: gr.Request):
+    "Read parameters from request header"
+    config = decode_config(request.query_params['dta'])
+    survey_question = config['question']
+    survey_template = config['template']
+    initial_message = config['initial_message']
+    model_args = {'model': config['model'], 'temperature': config['temperature']}
+    return survey_question, survey_template, initial_message, model_args
+
+# Post-loading
+def update_system_message(question: str, template: PromptTemplate | str) -> str:
+    """
+    On questionBox|templateBox update, read questionBox|templateBox and update SystemMessageBox
+    """
+    return template.format(question=question)
+
+
+def reset_interview() -> tuple[list[list[str | None]], gr.Button, gr.Button]:
+    wandb.finish()
+    gr.Info("Interview reset.")
+    return (
+        [],
+        gr.Button("Start Interview", visible=True),
+        gr.Button("Reply", visible=False),
+        gr.Button("Save Survey", visible=False, variant="secondary"),
+        gr.Button("Reset Survey", visible=False, variant="stop"),
+    )
+
+
+def initialize_interview(
+    system_message: str, first_question: str, model_args: dict[str, str | float]
+) -> tuple[list[list[str | None]], gr.Textbox, gr.Button, gr.Button]:
+    "Read system prompt and start interview"
+    if len(first_question) == 0:
+        first_question = call_openai(
+            [], system_message, client, model_args, stream=False
+        )
+    # Use fixed prompt
+    chat_history = [[None, first_question]]
+    return (
+        chat_history,
+        gr.Textbox(
+            placeholder="Type response here.", interactive=True, show_label=False
+        ),
+        gr.Button("Start Interview", visible=False),
+        gr.Button("Reset Survey", visible=True, variant="stop"),
+    )
+
+
+def initialize_tracker(
+    model_args: dict[str, str | float], question: str, template: PromptTemplate
+):
+    "Initializes wandb run for interview"
+    run_config = model_args | {"question": question, "template": str(template)}
+    wandb.init(project="cognitive-debrief", config=run_config, tags=["dev"])
+
+
+def save_interview(
+    chat_history: list[list[str | None]],
+) -> None:
+    chat_data = []
+    for pair in chat_history:
+        for i, role in enumerate(["user", "system"]):
+            if pair[i] is not None:
+                chat_data += [[role, pair[i]]]
+    chat_table = wandb.Table(data=chat_data, columns=["role", "message"])
+    gr.Info("Uploading interview transcript to WandB...")
+    wandb.log({"chat_history": chat_table})
+
+
+def call_openai(
+    messages: list[dict[str, str]],
+    system_message: str | None,
+    client: openai.Client,
+    model_args: dict,
+    stream: bool = False,
+):
+    "Utility function for calling OpenAI chat. Expects formatted messages."
+    if not messages:
+        messages = []
+    if system_message:
+        messages = [{"role": "system", "content": system_message}] + messages
+    try:
+        response = client.chat.completions.create(
+            messages=messages, **model_args, stream=stream
+        )
+        if stream:
+            for chunk in response:
+                yield chunk.choices[0].message.content
+        else:
+            content = response.choices[0].message.content
+            return content
+    except openai.APIConnectionError | openai.APIStatusError as e:
+        error_msg = (
+            "API unreachable.\n" f"STATUS_CODE: {e.status_code}" f"ERROR: {e.response}"
+        )
+        gr.Error(error_msg)
+        logger.error(error_msg)
+    except openai.RateLimitError:
+        warning_msg = "Hit rate limit. Wait a moment and retry."
+        gr.Warning(warning_msg)
+        logger.warning(warning_msg)
+
+
+def user_message(
+    message: str, chat_history: list[list[str | None]]
+) -> tuple[str, list[list[str | None]]]:
+    "Displays user message immediately."
+    return "", chat_history + [[message, None]]
+
+
+def bot_message(
+    chat_history: list[list[str | None]],
+    system_message: str,
+    model_args: dict[str, str | float],
+) -> list[list[str | None]]:
+    # Prep messages
+    user_msg = chat_history[-1][0]
+    messages = convert_gradio_to_openai(chat_history[:-1])
+    messages = (
+        [{"role": "system", "content": system_message}]
+        + messages
+        + [{"role": "user", "content": user_msg}]
+    )
+    response = client.chat.completions.create(
+        messages=messages, stream=True, **model_args
+    )
+    # Streaming
+    chat_history[-1][1] = ""
+    for chunk in response:
+        delta = chunk.choices[0].delta.content
+        if delta:
+            chat_history[-1][1] += delta
+            yield chat_history
+
+
+# LAYOUT
+with gr.Blocks() as demo:
+    gr.Markdown("# Cognitive Debriefing Prototype")
+
+    # Hidden values
+    surveyQuestion = gr.Textbox(visible=False)
+    surveyTemplate = gr.Textbox(visible=False)
+    initialMessage = gr.Textbox(visible=False)
+    systemMessage = gr.Textbox(visible=False)
+    modelArgs = gr.State(value={"model": "", "temperature": ""})
+
+    # Debugging
+    with gr.Accordion("Debugging Panel", open=False):
+        debugPane = gr.Textbox(show_label=False, lines=8)
+        debugRequest = gr.Button('Read Request')
+        debugRequest.click(load_config, outputs=[debugPane])
+
+
+    ## RESPONDENT
+    chatDisplay = gr.Chatbot(
+        show_label=False,
+    )
+    chatInput = gr.Textbox(
+        placeholder="Click 'Start Interview' to begin.",
+        interactive=False,
+        show_label=False,
+    )
+    startInterview = gr.Button("Start Interview", variant="primary")
+    resetButton = gr.Button("Reset Survey", visible=False, variant="stop")
+
+    ## INTERACTIONS
+    startInterview.click(
+        load_config,
+        inputs=None,
+        outputs=[
+            surveyQuestion,
+            surveyTemplate,
+            initialMessage,
+            modelArgs,
+        ]
+    ).then(
+        update_system_message,
+        inputs=[surveyQuestion, surveyTemplate],
+        outputs=[systemMessage],
+    ).then(
+        initialize_interview,
+        inputs=[systemMessage, initialMessage, modelArgs],
+        outputs=[
+            chatDisplay,
+            chatInput,
+            startInterview,
+            resetButton,
+        ],
+    ).then(initialize_tracker, inputs=[modelArgs, surveyQuestion, surveyTemplate])
+
+    chatInput.submit(
+        user_message,
+        inputs=[chatInput, chatDisplay],
+        outputs=[chatInput, chatDisplay],
+        queue=False
+    ).then(
+        bot_message,
+        inputs=[chatDisplay, systemMessage, modelArgs],
+        outputs=[chatDisplay])
+
+    resetButton.click(
+        save_interview,
+        [chatDisplay]
+    ).then(
+        reset_interview,
+        outputs=[chatDisplay, startInterview, resetButton],
+        show_progress=False,
+    )
+
+
+if __name__ == "__main__":
+    # Testing
+    demo.launch()

requirements.txt

@@ -0,0 +1,3 @@
+gradio
+openai
+wandb

utils.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+from string import Formatter
+
+
+# General Class
+class PromptTemplate(str):
+    """More robust String Formatter. Takes a string and parses out the keywords."""
+
+    def __init__(self, template) -> None:
+        self.template: str = template
+        self.variables: list[str] = self.parse_template()
+
+    def parse_template(self) -> list[str]:
+        "Returns template variables"
+        return [
+            fn for _, fn, _, _ in Formatter().parse(self.template) if fn is not None
+        ]
+
+    def format(self, *args, **kwargs) -> str:
+        """
+        Formats the template string with the given arguments.
+        Provides slightly more informative error handling.
+
+        :param args: Positional arguments for unnamed placeholders.
+        :param kwargs: Keyword arguments for named placeholders.
+        :return: Formatted string.
+        :raises: ValueError if arguments do not match template variables.
+        """
+        # If keyword arguments are provided, check if they match the template variables
+        if kwargs and set(kwargs) != set(self.variables):
+            raise ValueError("Keyword arguments do not match template variables.")
+
+        # If positional arguments are provided, check if their count matches the number of template variables
+        if args and len(args) != len(self.variables):
+            raise ValueError(
+                "Number of arguments does not match the number of template variables."
+            )
+
+        # Check if a dictionary is passed as a single positional argument
+        if len(args) == 1 and isinstance(args[0], dict):
+            arg_dict = args[0]
+            if set(arg_dict) != set(self.variables):
+                raise ValueError("Dictionary keys do not match template variables.")
+            return self.template.format(**arg_dict)
+
+        # Check for the special case where both args and kwargs are empty, which means self.variables must also be empty
+        if not args and not kwargs and self.variables:
+            raise ValueError("No arguments provided, but template expects variables.")
+
+        # Use the arguments to format the template
+        try:
+            return self.template.format(*args, **kwargs)
+        except KeyError as e:
+            raise ValueError(f"Missing a keyword argument: {e}")
+
+    @classmethod
+    def from_file(cls, file_path: str) -> PromptTemplate:
+        with open(file_path, encoding="utf-8") as file:
+            template_content = file.read()
+        return cls(template_content)
+
+    def dump_prompt(self, file_path: str) -> None:
+        with open(file_path, "w", encoding="utf-8") as file:
+            file.write(self.template)
+            file.close()
+
+
+def convert_gradio_to_openai(
+    chat_history: list[list[str | None]],
+) -> list[dict[str, str]]:
+    "Converts gradio chat format -> openai chat request format"
+    messages = []
+    for pair in chat_history:  # [(user), (assistant)]
+        for i, role in enumerate(["user", "assistant"]):
+            if not ((pair[i] is None) or (pair[i] == "")):
+                messages += [{"role": role, "content": pair[i]}]
+    return messages
+
+
+def convert_openai_to_gradio(
+    messages: list[dict[str, str]]
+) -> list[list[str, str | None]]:
+    "Converts openai chat request format -> gradio chat format"
+    chat_history = []
+    if messages[0]["role"] != "user":
+        messages.insert(0, {"role": "user", "content": None})
+    for i in range(0, len(messages), 2):
+        chat_history.append([messages[i]["content"], messages[i + 1]["content"]])
+    return chat_history