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", generation_config={"response_mime_type": "application/json"}) # 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'), (['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'), ] 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: List[str]) -> List[dict]: """Groups the given words using the AI model.""" 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_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.") 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_clue(group: List[str]) -> dict: """Generates a clue for the given group of words using the AI model.""" 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_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.") 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) 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": }') 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 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']) # 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) with gr.Row(): detect_words_button = gr.Button("Detect Words") group_words_button = gr.Button("Group Words") 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) clue_button = gr.Button("Generate Clue", size='sm') clue_output = gr.Textbox(label=f"Clue {i + 1}") group_inputs.append(group_input) clue_buttons.append(clue_button) clue_outputs.append(clue_output) # 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): clue_buttons[i].click(generate_clue_callback, inputs=group_inputs[i], outputs=clue_outputs[i]) demo.launch(share=True, debug=True)