add code generation

generate.py

@@ -0,0 +1,70 @@
+import os
+import re
+from llms import build_prompt, get_completion, PROMPT_TEMPLATE, SYSTEM_PROMPT
+
+
+def generate_code(problem_description: str, prompt_template: str, system_prompt: str, model: str, temperature: float):
+    prompt = build_prompt(problem_description, prompt_template)
+    completion = get_completion(prompt, system_prompt, model, temperature)
+    return completion
+
+
+def load_txt_file(file_path: str):
+    with open(file_path, "r") as file:
+        return file.read()
+
+def load_problem_description(day: int):
+    day_str = f"{day:02d}"
+    file_path = f"day{day_str}/problem.txt"
+    return load_txt_file(file_path)
+
+
+def extract_code(completion: str) -> str:
+    """Extracts the code from the completion, should be contained in ```python ...``` code block."""
+    code_block = re.search(r"```python\s*([\s\S]*?)\s*```", completion)
+    if code_block:
+        return code_block.group(1)
+
+def get_solution_file_path(day: int, model: str):
+    """Returns the path to the solution file for the given day and model."""
+    day_str = f"{day:02d}"
+    return f"day{day_str}/solution_{model}.py"
+
+
+def save_code(day: int, code: str, model: str):
+    """Saves the code to the solution file for the given day and model."""
+    file_path = get_solution_file_path(day, model)
+    with open(file_path, "w") as file:
+        file.write(code)
+    print(f"Saved code to {file_path}")
+
+
+all_models = {
+    "openai": ["gpt-4o"],
+    "gemini": ["gemini-1.5-pro"],
+    "anthropic": ["claude-3-5-sonnet-20241022"],
+}
+
+
+if __name__ == "__main__":
+
+    for day in range(1, 26):
+        problem_description = load_problem_description(day)
+        print(f"***Generating code for day {day}***")
+
+        for provider in all_models:
+            for model in all_models[provider]:
+                print("-" * 80)
+                print(f"Generating code for {provider} {model}")
+
+                if os.path.exists(get_solution_file_path(day, model)):
+                    print(f"Skipping {provider} {model} for day {day} because it already exists")
+                    continue
+
+                prompt = build_prompt(problem_description, PROMPT_TEMPLATE)
+                completion = get_completion(provider=provider, user_prompt=prompt, system_prompt=SYSTEM_PROMPT, model=model, temperature=0)
+                code = extract_code(completion)
+                save_code(day, code, model)
+                print("-" * 80)
+
+        print("*" * 80)

llms.py

@@ -0,0 +1,126 @@
+import os
+import logging
+
+import google.generativeai as genai
+from openai import OpenAI
+import anthropic
+
+
+DEFAULT_PROVIDER = "openai"
+DEFAULT_MODEL = "gpt-4o-mini"
+
+# Load API keys
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+AI_STUDIO_API_KEY = os.getenv("AI_STUDIO_API_KEY")
+ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
+
+assert (
+    OPENAI_API_KEY
+), "OPENAI_API_KEY is not set, please set it in your environment variables."
+assert (
+    AI_STUDIO_API_KEY
+), "AI_STUDIO_API_KEY is not set, please set it in your environment variables."
+assert (
+    ANTHROPIC_API_KEY
+), "ANTHROPIC_API_KEY is not set, please set it in your environment variables."
+
+# Initialize clients
+openai_client = OpenAI(api_key=OPENAI_API_KEY)
+genai.configure(api_key=AI_STUDIO_API_KEY)
+anthropic_client = anthropic.Anthropic(api_key=ANTHROPIC_API_KEY)
+
+
+SYSTEM_PROMPT = "You are a programming assistant. You are solving the 2024 advent of code challenge."
+PROMPT_TEMPLATE = """You are solving the 2024 advent of code challenge.
+You will be provided the description of each challenge. You are to provide the solution to each given challenge.
+1) You can reason and explain your logic before writing the code.
+2) You must write the code such that it can be parsed into an actual python file.
+3) It will be parsed by the evaluator, so it must be valid python code.
+4) All of the code must be in a single code block, delimited by ```python and ```.
+5) To count as a proper submission, the code must print the result to each question asked.
+6) Each question will have a single string as an answer. Make sure to print it that string, and nothing else.
+7) The actual input to the question will be provided in a file relative to the python file, e.g. "./input.txt". You must read and parse from the file accordingly. You can safely assume the file will always be relative to the python file.
+
+Here is an example of a proper submission:
+
+You reasoning goes here ...
+
+```python
+
+
+file = "input.txt"
+
+def your_function(...)
+    ...
+
+...
+print(result1)
+
+
+def your_other_function(...)
+    ...
+
+...
+print(result2)
+
+```
+
+Here is today's challenge description:
+{problem_description}
+"""
+
+
+def build_prompt(
+    problem_description: str, prompt_template: str = PROMPT_TEMPLATE
+) -> str:
+    return prompt_template.format(problem_description=problem_description)
+
+
+def get_completion(
+    provider: str,
+    user_prompt: str,
+    system_prompt: str,
+    model: str,
+    temperature: float,
+) -> str:
+    """
+    Unified function to get completions from various LLM providers.
+    """
+    if provider == "openai":
+        completion = openai_client.chat.completions.create(
+            model=model,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": user_prompt},
+            ],
+            temperature=temperature,
+        )
+        # logger.info("Completion: %s", completion)
+        return completion.choices[0].message.content
+
+    elif provider == "gemini":
+        model = genai.GenerativeModel(
+            model_name=model,
+            system_instruction=system_prompt,
+        )
+        response = model.generate_content(
+            user_prompt,
+            generation_config=genai.types.GenerationConfig(temperature=temperature),
+        )
+
+        # logger.info("reponse: %s", response)
+        return response.text
+
+    elif provider == "anthropic":
+        response = anthropic_client.messages.create(
+            model=model,
+            max_tokens=2048,
+            temperature=temperature,
+            system=system_prompt,
+            messages=[{"role": "user", "content": user_prompt}],
+        )
+        # logger.info("Response: %s", response)
+        return response.content[0].text
+
+    else:
+        raise ValueError(f"Unknown provider: {provider}")