killerstorm · March 26, 2023 17:08 · Mar 26, 2023
diff --git a/codecode_yaml_scale.py b/codecode_yaml_scale.py
@@ -0,0 +1,244 @@
+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)