| 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 | |
| 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) |
huncholane