app.py

import json
import os
import random

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",
                              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 = [
    (['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'),
    (['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'),
    (['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'),
    (['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'),
    (['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'),
]

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):
    env = Environment(loader=FileSystemLoader('.'))
    template = env.from_string(template)
    return template.render(system_prompt=system_prompt, history=history, query=query)


def group_words(words):
    template = '''
    {% for example in history %}
    INPUT:
    {{ example[0] }}
    OUTPUT:
    {{ example[1] }}
    {% endfor %}
    INPUT:
    {{ query }}
    OUTPUT:

    {{ system }}

    Groups = {'words': list[str], 'clue': str, 'explanation': str}
    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)

    print("Got groupings: ", json.dumps(response, indent=4))
    return response


def generate_clues(group):
    template = '''
    {% for example in history %}
    INPUT:
    {{ example[0] }}
    OUTPUT:
    { 'clue':{{ example[1] }}, 'explanation':{{ example[2] }} }
    {% endfor %}
    INPUT:
    {{ 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)

    raw_response = model.generate_content(
        prompt,
        generation_config={'top_k': 3, 'temperature': 1.1})
    response = json.loads(raw_response.text)

    print("Generating clues for: ", group)

    print("Got clue: ", json.dumps(response, indent=4))
    return response


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)


with gr.Blocks() as demo:
    gr.Markdown("# *Codenames* clue generator")
    gr.Markdown("Provide a list of words to generate a clue")

    with gr.Row():
        game_image = gr.Image(type="pil")
        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 = [], [], []

    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)
            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'])


    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)

    for i in range(4):
        buttons[i].click(generate_clues_callback, inputs=dropdowns[i], outputs=outputs[i])

demo.launch(share=True, debug=True)