Skip to content

Instantly share code, notes, and snippets.

@killerstorm
Created March 26, 2023 17:08
Show Gist options
  • Save killerstorm/c8d64d29adfad55d9201caecfbdfb7a8 to your computer and use it in GitHub Desktop.
Save killerstorm/c8d64d29adfad55d9201caecfbdfb7a8 to your computer and use it in GitHub Desktop.
import requests
import yaml
from typing import List, Dict, Optional
class PlanItem:
def __init__(self, id: str, description: str, status: str, subtasks: List['PlanItem'] = None):
self.id = id
self.status = status
self.description = description
self.subtasks = subtasks or []
def to_dict(self) -> dict:
return {
'id': self.id,
'status': self.status,
'description': self.description,
'subtasks': [subtask.to_dict() for subtask in self.subtasks]
}
class CodeBlock:
def __init__(self, path: str, tag: str, code: str):
self.path = path
self.tag = tag
self.code = code
def to_dict(self) -> dict:
return {
'path': self.path,
'tag': self.tag,
'code': self.code
}
class CodeConstructionContext:
def __init__(self, specification: str, plan: List[PlanItem], code_blocks: List[Optional[CodeBlock]],
code_index: Dict[str, List[str]],
compiler_error_messages: List[str], test_results: List[str]):
self.specification = specification
self.plan = plan
self.code_index = code_index
self.code_blocks = code_blocks
self.compiler_error_messages = compiler_error_messages
self.test_results = test_results
class CodeBlockUpdate:
def __init__(self, path: str, tag: str, code: str):
self.path, self.tag, self.code = path, tag, code
def parse_plan_update(data: dict) -> PlanItem:
id = data['id']
status = data['status']
description = data.get('description')
subtasks = [parse_plan_update(subtask) for subtask in data.get('subtasks', [])]
return PlanItem(id, description, status, subtasks)
class CodeBlockRequest:
def __init__(self, path: str, tag: str):
self.path, self.tag = path, tag
class CodeConstructionOutput:
def __init__(self, code_block_updates: List[CodeBlockUpdate], plan_updates: List[PlanItem],
code_block_requests: List[CodeBlockRequest]):
self.code_block_updates, self.plan_updates, self.code_block_requests = code_block_updates, plan_updates, code_block_requests
class GenerationEngineState:
code_blocks: Dict[str, List[CodeBlock]]
plan: List[PlanItem]
def __init__(self, specification: str):
self.specification = specification
self.plan = [
PlanItem(id='start', status='todo', description='Implement a first sketch of the solution'),
PlanItem(id='plan', status='todo', description='Formulate a more detailed plan'),
PlanItem(id='test', status='todo', description='Create tests'),
]
self.code_blocks = {}
def get_initial_context(self) -> CodeConstructionContext:
return CodeConstructionContext(self.specification, self.plan, [], {}, [], [])
def _find_code_block(self, path: str, tag: str) -> Optional[CodeBlock]:
if path not in self.code_blocks:
return None
for block in self.code_blocks[path]:
if block.tag == tag:
return block
return None
def _apply_code_blocks(self, code_block_updates: List[CodeBlockUpdate]):
for update in code_block_updates:
if update.path not in self.code_blocks:
self.code_blocks[update.path] = []
path_blocks = self.code_blocks[update.path]
for i, block in enumerate(path_blocks):
if block.tag == update.tag:
path_blocks[i] = CodeBlock(update.path, update.tag, update.code)
break
else:
path_blocks.append(CodeBlock(update.path, update.tag, update.code))
def _apply_plan_updates(self, plan_updates: List[PlanItem]):
def update_plan_item(item: PlanItem, update: PlanItem) -> PlanItem:
if item.id == update.id:
item.status = update.status
if update.description:
item.description = update.description
for update_subtask in update.subtasks:
for item_subtask in item.subtasks:
if item_subtask.id == update_subtask.id:
update_plan_item(item_subtask, update_subtask)
break
else:
item.subtasks.append(update_subtask)
return item
for update in plan_updates:
item_found = False
for i, item in enumerate(self.plan):
if item.id == update.id:
self.plan[i] = update_plan_item(item, update)
item_found = True
break
if not item_found:
self.plan.append(PlanItem(update.id, update.description, update.status, update.subtasks))
def _resolve_code_block_requests(self, code_block_requests: List[CodeBlockRequest]) -> List[CodeBlock]:
blocks = [self._find_code_block(request.path, request.tag) for request in code_block_requests]
# remove None values
return [block for block in blocks if block]
def _generate_code_index(self) -> Dict[str, List[str]]:
return {path: [block.tag for block in blocks] for path, blocks in self.code_blocks.items()}
def apply_step(self, output: CodeConstructionOutput) -> CodeConstructionContext:
self._apply_code_blocks(output.code_block_updates)
self._apply_plan_updates(output.plan_updates)
requested_code_blocks = self._resolve_code_block_requests(output.code_block_requests)
return CodeConstructionContext(self.specification, self.plan, requested_code_blocks,
self._generate_code_index(), [], [])
def parse_output(data: dict) -> CodeConstructionOutput:
code_block_updates = [CodeBlockUpdate(d['path'], d['tag'], d['code']) for d in data.get('code_blocks', [])]
plan_updates = [parse_plan_update(d) for d in data.get('plan_updates', [])]
code_block_requests = [CodeBlockRequest(d['path'], d['tag']) for d in data.get('code_block_requests', [])]
return CodeConstructionOutput(code_block_updates, plan_updates, code_block_requests)
class CodeConstructionDriver:
def __init__(self, api_key: str):
self.api_key = api_key
#openai.api_key = api_key
self.state = None
def initialize(self, specification: str):
self.state = GenerationEngineState(specification)
def generate_code(self, context: CodeConstructionContext) -> CodeConstructionOutput:
if not self.state:
raise Exception("Call initialize() first")
content = yaml.dump({
'specification': self.state.specification,
'plan': [item.to_dict() for item in context.plan],
'code_blocks': [block.to_dict() for block in context.code_blocks],
'code_index': context.code_index
})
messages = [
{"role": "system", "content":
"""You're a code construction AI which creates code iteratively.
In each step a context is supplied. It includes
specification, plan, code blocks requested in the previous step.
Generated response should be in YAML format as in example below
```
code_blocks: # list of code blocks to be inserted or updated
- path: "src/main_yaml.py"
tag: "def_main"
code: "new code here"
plan_updates:
- id: "start"
status: "done"
- id: "new task"
status: "in_progress
description: "description of new task"
subtasks:
- id: "subtask1.1"
status: "todo"
description: "subtask1.1 description"
code_block_requests: # list of code blocks requested in the next step
- path: "path_to_previously_created_file.py"
tag: "tag2"
```
code_block_requests are used to get previously written code into the context for the next step.
code_block_updates are the newly generated code, possibly updating existing code blocks.
Do not include "```" in the response, include only raw YAML data.
"""},
{"role": "user", "content": content}
]
print(content)
headers = {"Authorization":"Basic"}
response_1 = requests.post(
"URL to scale AI deployment",
json={"input": {"content": content}},
headers=headers
)
print(response_1)
response = response_1.json()["output"]
print("Response:")
print(response)
print("Response ENDS")
output_data = yaml.load(response, Loader=yaml.SafeLoader)
return parse_output(output_data)
def run(self, specification: str, num_steps: int):
self.initialize(specification)
context = self.state.get_initial_context()
for _ in range(num_steps):
output = self.generate_code(context)
context = self.state.apply_step(output)
if __name__ == "__main__":
driver = CodeConstructionDriver("openAI key actually not used")
driver.run("Write utility similar to grep in C++", 10)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment