refactoring

app.py

@@ -1,89 +1,75 @@
 import json
 import os
 import random
+from typing import List, Tuple
 
 import google.generativeai as genai
 import gradio as gr
 from jinja2 import Environment, FileSystemLoader
 
 GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY')
-
 genai.configure(api_key=GOOGLE_API_KEY)
 
-model = genai.GenerativeModel("gemini-1.5-flash-latest",
+MODEL = genai.GenerativeModel("gemini-1.5-flash-latest",
                               generation_config={"response_mime_type": "application/json"})
 
-
-def merge_games(clues, num_merges=10):
-    """Generates around 10 merges of words from the given clues.
-
-    Args:
-      clues: A list of clues, where each clue is a list containing the words, the answer, and the explanation.
-      num_merges: The approximate number of merges to generate (default: 10).
-
-    Returns:
-      A list of tuples, where each tuple contains the merged words and the indices of the selected rows.
-    """
-
-    merged = []
-    while len(merged) < num_merges:
-        num_rows = random.choice([3, 4])
-        selected_rows = random.sample(range(len(clues)), num_rows)
-        words_list = [word for row in [clues[i][0] for i in selected_rows] for word in row]
-        if len(words_list) in [8, 9]:
-            merged.append((words_list, selected_rows))
-
-    return merged
-
-
-example_clues = [
+# Example clues and groupings
+EXAMPLE_CLUES = [
     (['ARROW', 'TIE', 'HONOR'], 'BOW', 'such as a bow and arrow, a bow tie, or a bow as a sign of honor'),
     (['DOG', 'TREE'], 'BARK', 'such as the sound a dog makes, or a tree is made of bark'),
-    (['MONEY', 'RIVER', 'ROB', 'BLOOD'], 'CRIME', 'such as money being stolen, a river being a potential crime scene, '
-                                                  'robbery, or blood being a result of a violent crime'),
-    (['BEEF', 'TURKEY', 'FIELD', 'GRASS'], 'GROUND',
-     'such as ground beef, a turkey being a ground-dwelling bird, a field or grass being a type of ground'),
-    (['BANK', 'GUITAR', 'LIBRARY'], 'NOTE',
-     'such as a bank note, a musical note on a guitar, or a note being a written comment in a library book'),
+    (['MONEY', 'RIVER', 'ROB', 'BLOOD'], 'CRIME', 'such as money being stolen, a river being a potential crime scene, robbery, or blood being a result of a violent crime'),
+    (['BEEF', 'TURKEY', 'FIELD', 'GRASS'], 'GROUND', 'such as ground beef, a turkey being a ground-dwelling bird, a field or grass being a type of ground'),
+    (['BANK', 'GUITAR', 'LIBRARY'], 'NOTE', 'such as a bank note, a musical note on a guitar, or a note being a written comment in a library book'),
     (['ROOM', 'PIANO', 'TYPEWRITER'], 'KEYS', 'such as a room key, piano keys, or typewriter keys'),
     (['TRAFFIC', 'RADAR', 'PHONE'], 'SIGNAL', 'such as traffic signals, radar signals, or phone signals'),
-    (['FENCE', 'PICTURE', 'COOKIE'], 'FRAME',
-     'such as a frame around a yard, a picture frame, or a cookie cutter being a type of frame'),
+    (['FENCE', 'PICTURE', 'COOKIE'], 'FRAME', 'such as a frame around a yard, a picture frame, or a cookie cutter being a type of frame'),
     (['YARN', 'VIOLIN', 'DRESS'], 'STRING', 'strings like material, instrument, clothing fastener'),
-    (['JUMP', 'FLOWER', 'CLOCK'], 'SPRING',
-     'such as jumping, flowers blooming in the spring, or a clock having a sprint component'),
-    (['SPY', 'KNIFE'], 'WAR',
-     'Both relate to aspects of war, such as spies being involved in war or knives being used as weapons'),
+    (['JUMP', 'FLOWER', 'CLOCK'], 'SPRING', 'such as jumping, flowers blooming in the spring, or a clock having a sprint component'),
+    (['SPY', 'KNIFE'], 'WAR', 'Both relate to aspects of war, such as spies being involved in war or knives being used as weapons'),
     (['STADIUM', 'SHOE', 'FIELD'], 'SPORT', 'Sports like venues, equipment, playing surfaces'),
-    (['TEACHER', 'CLUB'], 'SCHOOL',
-     'such as a teacher being a school staff member or a club being a type of school organization'),
+    (['TEACHER', 'CLUB'], 'SCHOOL', 'such as a teacher being a school staff member or a club being a type of school organization'),
     (['CYCLE', 'ARMY', 'COURT', 'FEES'], 'CHARGE', 'charges like electricity, battle, legal, payments'),
-    (['FRUIT', 'MUSIC', 'TRAFFIC', 'STUCK'], 'JAM',
-     'Jams such as fruit jam, a music jam session, traffic jam, or being stuck in a jam'),
-    (['POLICE', 'DOG', 'THIEF'], 'CRIME',
-     'such as police investigating crimes, dogs being used to detect crimes, or a thief committing a crime'),
-    (['ARCTIC', 'SHUT', 'STAMP'], 'SEAL',
-     'such as the Arctic being home to seals, or shutting a seal on an envelope, or a stamp being a type of seal'),
+    (['FRUIT', 'MUSIC', 'TRAFFIC', 'STUCK'], 'JAM', 'Jams such as fruit jam, a music jam session, traffic jam, or being stuck in a jam'),
+    (['POLICE', 'DOG', 'THIEF'], 'CRIME', 'such as police investigating crimes, dogs being used to detect crimes, or a thief committing a crime'),
+    (['ARCTIC', 'SHUT', 'STAMP'], 'SEAL', 'such as the Arctic being home to seals, or shutting a seal on an envelope, or a stamp being a type of seal'),
 ]
 
-example_groupings = []
-merges = merge_games(example_clues, 5)
-for merged_words, indices in merges:
-    groups = [{
-        "words": example_clues[i][0],
-        "clue": example_clues[i][1],
-        "explanation": example_clues[i][2]
-    } for i in indices]
-    example_groupings.append((merged_words, json.dumps(groups, separators=(',', ':'))))
-
-
-def render_jinja2_template(template, system_prompt, history, query):
+def create_example_groupings(clues: List[Tuple[List[str], str, str]], num_groups: int = 5) -> List[Tuple[List[str], str]]:
+    """Creates example groupings from the given clues."""
+    merged = create_random_word_groups(clues, num_groups)
+    return [
+        (
+            merged_words,
+            json.dumps([{
+                "words": clues[i][0],
+                "clue": clues[i][1],
+                "explanation": clues[i][2]
+            } for i in indices], separators=(',', ':'))
+        )
+        for merged_words, indices in merged
+    ]
+
+EXAMPLE_GROUPINGS = create_example_groupings(EXAMPLE_CLUES)
+
+def create_random_word_groups(clues: List[Tuple[List[str], str, str]], num_groups: int = 10) -> List[Tuple[List[str], List[int]]]:
+    """Creates random groups of words from the given clues."""
+    word_groups = []
+    while len(word_groups) < num_groups:
+        group_size = random.choice([3, 4])
+        selected_indices = random.sample(range(len(clues)), group_size)
+        words = [word for row in [clues[i][0] for i in selected_indices] for word in row]
+        if len(words) in [8, 9]:
+            word_groups.append((words, selected_indices))
+    return word_groups
+
+def render_template(template: str, system_prompt: str, history: List[Tuple], query: str) -> str:
+    """Renders a Jinja2 template with the given parameters."""
     env = Environment(loader=FileSystemLoader('.'))
     template = env.from_string(template)
     return template.render(system_prompt=system_prompt, history=history, query=query)
 
-
-def group_words(words):
+def group_words(words: List[str]) -> List[dict]:
+    """Groups the given words using the AI model."""
     template = '''
     {% for example in history %}
     INPUT:
@@ -101,24 +87,15 @@ def group_words(words):
     Return: Groups
     '''
 
-    grouping_system_prompt = ("You are an assistant for the game Codenames. Your task is to help players by grouping a "
-                              "given group of secrets into 3 to 4 groups of 2 to 4 words. Each group should consist of secrets that "
-                              "share a common theme or other word connections such as homonym, hypernyms or synonyms. Avoid clues that are not too generic or not unique enough to be guessed easily")
-
-    prompt = render_jinja2_template(template, grouping_system_prompt, example_groupings, words)
-    # print(prompt)
-    raw_response = model.generate_content(
-        prompt,
-        generation_config={'top_k': 3, 'temperature': 1.1})
-    response = json.loads(raw_response.text)
-
-    print("Grouping words:", words)
+    grouping_prompt = ("You are an assistant for the game Codenames. Group the given words into 3 to 4 sets of 2 to 4 words each. "
+                       "Each group should share a common theme or word connection. Avoid generic or easily guessable clues.")
 
-    print("Got groupings: ", json.dumps(response, indent=4))
-    return response
+    prompt = render_template(template, grouping_prompt, EXAMPLE_GROUPINGS, words)
+    response = MODEL.generate_content(prompt, generation_config={'top_k': 3, 'temperature': 1.1})
+    return json.loads(response.text)
 
-
-def generate_clues(group):
+def generate_clue(group: List[str]) -> dict:
+    """Generates a clue for the given group of words using the AI model."""
     template = '''
     {% for example in history %}
     INPUT:
@@ -130,99 +107,76 @@ def generate_clues(group):
     {{ query }}
     OUTPUT:
 
-
     {{ system }}
 
     Clue = {'clue': str, 'explanation': str}
     Return: Clue
     '''
 
-    clue_system_prompt = ("You are a codenames game companion. Your task is to give a single word clue related to "
-                          "a given group of words. You will only respond with a single word clue. The clue can be a common theme or other word connections such as homonym, hypernyms or synonyms. Avoid clues that are not too generic or not unique enough to be guessed easily")
-
-    prompt = render_jinja2_template(template, clue_system_prompt, example_clues, group)
+    clue_prompt = ("As a Codenames game companion, provide a single-word clue for the given group of words. "
+                   "The clue should relate to a common theme or word connection. DO NOT reuse any of the given "
+                   "words as a clue. Avoid generic or easily guessable clues.")
 
-    raw_response = model.generate_content(
-        prompt,
-        generation_config={'top_k': 3, 'temperature': 1.1})
-    response = json.loads(raw_response.text)
+    prompt = render_template(template, clue_prompt, EXAMPLE_CLUES, group)
+    response = MODEL.generate_content(prompt, generation_config={'top_k': 3, 'temperature': 1.1})
+    return json.loads(response.text)
 
-    print("Generating clues for: ", group)
 
-    print("Got clue: ", json.dumps(response, indent=4))
-    return response
+def process_image(img) -> gr.update:
+    """Processes the uploaded image and extracts words for the game."""
+    prompt = ('Identify the words in this Codenames game image. Provide only a list of words in capital letters. '
+              'Group these words into 6 or 8 sets that can be guessed together using a single-word clue. '
+              'Respond with JSON in the format: {"Game": <list of words in the game>}')
+    response = MODEL.generate_content([prompt, img], stream=True)
+    response.resolve()
+    words = json.loads(response.text)['Game']
+    return gr.update(choices=words, value=words)
 
+def pad_or_truncate(lst: List, n: int = 4) -> List:
+    """Ensures the list has exactly n elements, padding with None if necessary."""
+    truncated_lst = lst[:n]
+    return truncated_lst + (n - len(truncated_lst)) * [{}]
 
-def process_image(img):
-    raw_response = model.generate_content(['Identify the words in this game of Codenames. Provide only a list of '
-                                           'words. Provide the words in capital letters only. Group these words into '
-                                           '6 or 8 groups that can be guessed together using a single word clue in '
-                                           'the game of codenames. Give a response as json of the form: {"Game": '
-                                           '<list of words in the game>}', img], stream=True)
-    raw_response.resolve()
-    response = json.loads(raw_response.text)
-    words = response['Game']
-    return gr.update(choices=words, value=words)
+def group_words_callback(words: List[str]) -> List[gr.update]:
+    """Callback function for grouping words."""
+    groups = group_words(words)
+    groups = pad_or_truncate(groups, 4)
+    return [gr.update(value=groups[i].get("words", ""), choices=words, info=groups[i].get("explanation","")) for i in range(4)]
 
+def generate_clue_callback(group: List[str]) -> gr.update:
+    """Callback function for generating clues."""
+    clue = generate_clue(group)
+    return gr.update(value=clue['clue'], info=clue['explanation'])
 
-with gr.Blocks() as demo:
-    gr.Markdown("# *Codenames* clue generator")
-    gr.Markdown("Provide a list of words to generate a clue")
+# UI Setup
+with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
+    gr.Markdown("# *Codenames* Clue Generator")
+    gr.Markdown("Provide a list of words to generate clues")
 
     with gr.Row():
         game_image = gr.Image(type="pil")
-        word_list_input = gr.Dropdown(label="Detected words",
-                                      choices=[],
-                                      multiselect=True,
-                                      interactive=True)
+        word_list_input = gr.Dropdown(label="Detected words", choices=[], multiselect=True, interactive=True)
 
     with gr.Row():
         detect_words_button = gr.Button("Detect Words")
         group_words_button = gr.Button("Group Words")
 
-    dropdowns, buttons, outputs = [], [], []
+    group_inputs, clue_buttons, clue_outputs = [], [], []
 
     for i in range(4):
         with gr.Row():
-            group_input = gr.Dropdown(label=f"Group {i + 1}",
-                                      choices=[],
-                                      allow_custom_value=True,
-                                      multiselect=True,
-                                      interactive=True)
+            group_input = gr.Dropdown(label=f"Group {i + 1}", choices=[], allow_custom_value=True, multiselect=True, interactive=True)
             clue_button = gr.Button("Generate Clue", size='sm')
             clue_output = gr.Textbox(label=f"Clue {i + 1}")
-            dropdowns.append(group_input)
-            buttons.append(clue_button)
-            outputs.append(clue_output)
-
-
-    def pad_or_truncate(lst, n=4):
-        # Ensure the length of the list is at most n
-        truncated_lst = lst[:n]
-        return truncated_lst + (n - len(truncated_lst)) * [None]
-
-
-    def group_words_callback(words):
-        groups = group_words(words)
-        groups = pad_or_truncate(groups, 4)
-        print("Got groups: ", groups, type(groups))
-        return [gr.update(value=groups[i]["words"], choices=words, info=groups[i]["explanation"]) for i in range(4)]
-
-
-    def generate_clues_callback(group):
-        print("Generating clues: ", group)
-        g = generate_clues(group)
-        return gr.update(value=g['clue'], info=g['explanation'])
-
+            group_inputs.append(group_input)
+            clue_buttons.append(clue_button)
+            clue_outputs.append(clue_output)
 
-    detect_words_button.click(fn=process_image,
-                              inputs=game_image,
-                              outputs=[word_list_input])
-    group_words_button.click(fn=group_words_callback,
-                             inputs=word_list_input,
-                             outputs=dropdowns)
+    # Event handlers
+    detect_words_button.click(fn=process_image, inputs=game_image, outputs=[word_list_input])
+    group_words_button.click(fn=group_words_callback, inputs=word_list_input, outputs=group_inputs)
 
     for i in range(4):
-        buttons[i].click(generate_clues_callback, inputs=dropdowns[i], outputs=outputs[i])
+        clue_buttons[i].click(generate_clue_callback, inputs=group_inputs[i], outputs=clue_outputs[i])
 
 demo.launch(share=True, debug=True)