huncholanea bin full of cool programs
'broken codegen '
huncholane4 files changed, +129 -0
+3-0.gitignore
| @@ -0,0 +1,3 @@ | ||
| 1 | +__pycache__/ | |
| 2 | +testdir/ | |
| 3 | +datadir/ |
+2-0bin/codegen
| @@ -0,0 +1,2 @@ | ||
| 1 | +#!/bin/bash | |
| 2 | +python ~/.hygobin/python_scripts/codegen/cli.py $@ |
+12-0python_scripts/codegen/cli.py
| @@ -0,0 +1,12 @@ | ||
| 1 | +import click | |
| 2 | +from .codegen import Codegen | |
| 3 | +from pathlib import Path | |
| 4 | + | |
| 5 | + | |
| 6 | +@click.command() | |
| 7 | +@click.argument("prompt") | |
| 8 | +def cli(prompt): | |
| 9 | + codegen = Codegen() | |
| 10 | + codegen.load_session_data(str(Path.cwd())) | |
| 11 | + if codegen.is_new_session_data: | |
| 12 | + codegen.init_session() |
+112-0python_scripts/codegen/codegen.py
| @@ -0,0 +1,112 @@ | ||
| 1 | +from pathlib import Path | |
| 2 | +from typing import Dict | |
| 3 | +from openai import Client | |
| 4 | +from openai.types.beta import Assistant, Thread | |
| 5 | +from openai.types.beta.threads import Run | |
| 6 | +import json | |
| 7 | + | |
| 8 | +from pydantic import BaseModel, Field | |
| 9 | + | |
| 10 | + | |
| 11 | +class CodegenResponse(BaseModel): | |
| 12 | + file_data: Dict[str, str] = Field( | |
| 13 | + ..., | |
| 14 | + description="The relative path to put the file and the content of the file.", | |
| 15 | + ) | |
| 16 | + run_script: str = Field(..., description="The script to run the code.") | |
| 17 | + production_script: str = Field( | |
| 18 | + ..., description="The script to run the code in production." | |
| 19 | + ) | |
| 20 | + test_script: str = Field(..., description="The script to test the code.") | |
| 21 | + | |
| 22 | + | |
| 23 | +class Codegen: | |
| 24 | + assistant_data: Dict[str, Dict[str, str]] | |
| 25 | + assistant_id = None | |
| 26 | + thread_id = None | |
| 27 | + session_name: str | |
| 28 | + assistant: Assistant | |
| 29 | + thread: Thread | |
| 30 | + | |
| 31 | + @property | |
| 32 | + def is_new_session_data(self): | |
| 33 | + return self.assistant_id is None or self.thread_id is None | |
| 34 | + | |
| 35 | + def __init__( | |
| 36 | + self, | |
| 37 | + datadir="~/.hygobin/datadir", | |
| 38 | + model="gpt-4o-mini", | |
| 39 | + extensions=[ | |
| 40 | + "*.py", | |
| 41 | + "*.html", | |
| 42 | + "*.rs", | |
| 43 | + "*.sh", | |
| 44 | + "*.ts", | |
| 45 | + "*.js", | |
| 46 | + "*.jsx", | |
| 47 | + "*.tsx", | |
| 48 | + "*.json", | |
| 49 | + ], | |
| 50 | + tools=[{"type": "code_interpreter"}], | |
| 51 | + instructions="You are a software developer working on a project.", | |
| 52 | + ): | |
| 53 | + self.datadir = Path(datadir).expanduser() | |
| 54 | + self.assistant_file = self.datadir / "sessions.json" | |
| 55 | + self.instructions = instructions | |
| 56 | + self.tools = tools | |
| 57 | + self.model = model | |
| 58 | + self.extensions = extensions | |
| 59 | + self.assistant_file.touch(exist_ok=True) | |
| 60 | + with open(self.assistant_file, "r") as f: | |
| 61 | + self.assistant_data = json.load(f) | |
| 62 | + self.client = Client() | |
| 63 | + | |
| 64 | + def init_session(self): | |
| 65 | + # Create a new assistant | |
| 66 | + self.assistant = self.client.beta.assistants.create( | |
| 67 | + name=self.session_name, | |
| 68 | + model=self.model, | |
| 69 | + tools=self.tools, | |
| 70 | + instructions=self.instructions, | |
| 71 | + ) | |
| 72 | + self.thread = self.client.beta.threads.create() | |
| 73 | + self.assistant_id = self.assistant.id | |
| 74 | + self.thread_id = self.thread.id | |
| 75 | + run = self.client.beta.threads.create_and_run_poll( | |
| 76 | + thread=self.thread, | |
| 77 | + assistant_id=self.assistant_id, | |
| 78 | + model=self.model, | |
| 79 | + response_format=CodegenResponse, | |
| 80 | + ) | |
| 81 | + | |
| 82 | + def get_current_file_data(self): | |
| 83 | + file_data = {} | |
| 84 | + for file in Path(".").rglob(self.extensions): | |
| 85 | + with open(file, "r") as f: | |
| 86 | + file_data[file.absolute()] = f.read() | |
| 87 | + return file_data | |
| 88 | + | |
| 89 | + def handle_run(self, run: Run): | |
| 90 | + if run.status == "completed": | |
| 91 | + message = self.client.beta.threads.messages.retrieve(message_id=run.id) | |
| 92 | + print(message) | |
| 93 | + | |
| 94 | + def load_session_data(self, session_name: str): | |
| 95 | + self.session_name = session_name | |
| 96 | + with open(self.assistant_file, "r") as f: | |
| 97 | + self.assistant_data = json.load(f) | |
| 98 | + try: | |
| 99 | + self.assistant_id = self.assistant_data[self.session_name]["assistant_id"] | |
| 100 | + self.thread_id = self.assistant_data[self.session_name]["thread_id"] | |
| 101 | + except KeyError: | |
| 102 | + self.assistant_data[self.session_name] = {} | |
| 103 | + self.assistant_id = None | |
| 104 | + self.thread_id = None | |
| 105 | + | |
| 106 | + def save_session(self): | |
| 107 | + self.assistant_data[self.session_name] = { | |
| 108 | + "assistant_id": self.assistant_id, | |
| 109 | + "thread_id": self.thread_id, | |
| 110 | + } | |
| 111 | + with open(self.assistant_file, "w") as f: | |
| 112 | + json.dump(self.assistant_data, f) |