schema.py

from pydantic import BaseModel, constr
from enum import Enum
import outlines
import json
import copy
from outlines.samplers import Sampler, multinomial
import re

class CompletionStatus(str, Enum):
    todo = 'todo'
    done = 'done'
    inprogress = "inprogress"
class CompletionStatusOnlyNew(str, Enum):
    todo = 'todo'
    inprogress = "inprogress"
class InProgressCompletionStatus(str, Enum):
    inprogress = "inprogress"
    done = 'done'

class OnlyTODO(str, Enum):
    todo = 'todo'

class OnlyInProgress(str, Enum):
    todo = 'inprogress'

class BaseTask(BaseModel):
    status: CompletionStatus
    task_description: constr(max_length=150)
class BaseTaskOnlyTODO(BaseModel):
    status: OnlyTODO
    task_description: constr(max_length=150)
class BaseTaskOnlyNew(BaseModel):
    status: CompletionStatusOnlyNew
    task_description: constr(max_length=150)

class PTT(BaseModel):
    recon: list[BaseTask]
    initial_access: list[BaseTask]
    execution: list[BaseTask]
    post_exploitation: list[BaseTask]
    
def load_outlines(model_path:str, model_config: dict) -> tuple[outlines.models.LlamaCpp, Sampler]:
    llm = outlines.models.llamacpp(
        model_path,
        model_kwargs={
            "n_gpu_layers": model_config["n_gpu_layers"],
            "n_batch": model_config["n_batch"],
            "n_ctx": model_config["generate_len"]
        },
        device="cpu"
    )
    sampler =  multinomial(top_k=model_config["top_k"], top_p=model_config["top_p"], temperature=model_config["temperature"])
    return llm, sampler

def find_inprogress(ptt_dict: dict[str, list]) -> int:
    in_progress_num = 0
    for key in ptt_dict:
        for task in ptt_dict[key]:
            if task["status"] == "inprogress":
                in_progress_num += 1
    return in_progress_num

def find_inprogress_or_todo(ptt_dict: dict[str, list]) -> int:
    num = 0
    for key in ptt_dict:
        for task in ptt_dict[key]:
            if task["status"] == "inprogress" or task["status"] == "todo":
                num += 1
    return num

def find_inprogress_task(ptt_dict: dict[str, list]):
    for key in ptt_dict:
        for task in ptt_dict[key]:
            if task["status"] == "inprogress":
                return task["task_description"]
    assert False, f"There must be exactly one task that is inprogress in {ptt_dict}"


def str2dict(output: str):
    try:
        output = output.strip().replace("'", '"')
        return json.loads(output)
    except Exception as e:
        print("repr output: "+repr(output))
        raise Exception(e)

def get_current_status(ptt: dict):
    progress_ptt = copy.deepcopy(ptt)
    delete_indices = {}
    for key in progress_ptt:
        delete_indices[key] = []
        for i, task in enumerate(progress_ptt[key]):
            if task["status"] not in ["done", "inprogress"]:
                delete_indices[key].append(i)
    for key in delete_indices:
        for i in delete_indices[key][::-1]:
            progress_ptt[key].pop(i)
    return progress_ptt

def default_qa(template: str, question: str, llm, sampler, max_tokens: int = -1):
    generative_prompt = template.format(prompt=question)
    generator = outlines.generate.text(llm, sampler=sampler)
    if max_tokens == -1:
        return generator(generative_prompt)
    else:
        return generator(generative_prompt, max_tokens=max_tokens)

def input_parser(template: str, command_output: str, llm, sampler, max_tokens: int = 200):
    input_parser_prompt = template.format(prompt=command_output)
    generator = outlines.generate.text(llm, sampler=sampler)
    if max_tokens == -1:
        return generator(input_parser_prompt)
    else:
        return generator(input_parser_prompt, max_tokens=max_tokens)

def generative_module(template: str, ptt: dict, llm, sampler, current_history:str = "", max_tokens: int = -1, only_provide_currrent_status: bool = False, task: str | None = None):
    if only_provide_currrent_status:
        ptt = get_current_status(ptt)
    if task is None:
        task = find_inprogress_task(ptt)
    ptt = json.dumps(ptt)
    print(f"task is {task} given {ptt}")
    generative_prompt = template.format(ptt=ptt, prompt=task, history=current_history)
    generator = outlines.generate.text(llm, sampler=sampler)
    if max_tokens == -1:
        return generator(generative_prompt)
    else:
        return generator(generative_prompt, max_tokens=max_tokens)

def get_new_task_descriptions(original: dict[str, dict[str, str]], updated: dict[str, dict[str, str]]) -> list[str]:
    output: list[str] = []
    for key in updated:
        for task in updated[key]:
            task_description = task["task_description"]
            task_exists = False
            for original_task in original[key]:
                if task_description in original_task["task_description"]:
                    task_exists = True
            if not task_exists:
                output.append(task["task_description"])
    return output
def reasoning_module(template: str, prompt: str, past_history: str, ptt: dict, llm, sampler, force_add_task={"recon": 1}, update_status=True, todo_task_descriptions:list[str]=["Obtain a secret file with a hash in it"], max_spaces:int=0):
    force_add_task = copy.deepcopy(force_add_task)
    num_inprogress = find_inprogress(ptt)
    assert num_inprogress == 1, f"The number of current inprogress tasks must be 1 in {ptt}"
    llm_prompt = template.format(ptt=json.dumps(ptt), history=past_history, prompt=prompt).strip()
    if update_status:
        original, inprogress_set = update_completion_status_outlines(llm, llm_prompt, ptt, sampler, todo_task_descriptions=todo_task_descriptions, inprogress_always_set=False)
    else:
        original = ptt
        inprogress_set = True
    new_task_descriptions: list[str] = []
    while True:
        updated = add_new_items_outlines(llm, llm_prompt, original, sampler, inprogress_set, max_spaces=max_spaces, force_add_task=force_add_task)
        # If new items are added they do not get deleted
        added_tasks = get_new_task_descriptions(original, updated)
        for added_task in added_tasks:
            new_task_descriptions.append(added_task)
        num_inprogress = find_inprogress(updated)
        print("generated")
        print(updated)
        print(f"Number of inprogress tasks: ", num_inprogress)
        if num_inprogress == 0:
            ptt = updated
            llm_prompt = template.format(ptt=json.dumps(ptt), history=past_history, prompt=prompt).strip()

            if update_status:
                print("Updating status")
                original, inprogress_set = update_completion_status_outlines(llm, llm_prompt, ptt, sampler, todo_task_descriptions=todo_task_descriptions, inprogress_always_set=True, new_task_descriptions=new_task_descriptions)
                num_inprogress = find_inprogress(original)
                print("generated")
                print(original)
                assert num_inprogress == 1
                if num_inprogress == 1:
                    return original
            else:
                # shouldn't be possible to come here as if update_staus is set to False, then we only add todo tasks to the existing ptt 
                # and do nothing else
                raise Exception("If we are not updating status, there must already be one inprogress task")
        if num_inprogress == 1:
            return updated
        

def description_first_status_later(ptt: dict) -> dict:
    output = {}
    for key in ptt:
        output[key] = []
        for elem in ptt[key]:
            task_description = elem["task_description"]
            status = elem["status"]
            output[key].append({"task_description": task_description, "status": status})
    return output
def status_first_description_later(ptt: dict) -> dict:
    output = {}
    for key in ptt:
        output[key] = []
        for elem in ptt[key]:
            task_description = elem["task_description"]
            status = elem["status"]
            output[key].append({"status": status, "task_description": task_description})
    return output

def update_completion_status_outlines(llm, prompt: str, ptt: dict, sampler, todo_task_descriptions:list[str]=["Obtain a secret file with a hash in it"], inprogress_always_set: bool = True, new_task_descriptions: list[str] = []):
    """
    For updating status, we want to first see the task description then decide on the status.
    This is because we want our model to first read the task description before choosing a status
    """
    min_number_of_todo_inprogress = len(todo_task_descriptions)+1
    current_number_of_todo_inprogress = find_inprogress_or_todo(ptt)
    assert current_number_of_todo_inprogress >= min_number_of_todo_inprogress
    status_schema = CompletionStatus if current_number_of_todo_inprogress > min_number_of_todo_inprogress else CompletionStatusOnlyNew
    # below is incomplete as we need the current schema's str for this to work
    choice_sampler =  multinomial(top_k=50, top_p=1, temperature=0.3)
    ptt = description_first_status_later(ptt)
    ptt = json.dumps(ptt)
    original_prompt_len = len(prompt)
    choices: list[str] = []
    for status in status_schema:
        choices += [status.value]
    generator =outlines.generate.choice(llm, choices, sampler=choice_sampler)
    ptt_list = ptt.split("todo")
    output = []
    for elem in ptt_list:
        if "inprogress" not in elem:
            output.append((elem, "todo"))
        else:
            # there is only one in progress in todo list
            elems = elem.split("inprogress")
            output.append((elems[0], "inprogress"))
            output.append((elems[1], "todo"))

    ptt_list = output
    inprogress_set = False
    for i, elem in enumerate(ptt_list):
        elem_task = elem[0]
        elem_curr_status = elem[1]
        prompt+=elem_task
        task_set = False
        # If task is a task we want to set as todo
        for todo_task_description in todo_task_descriptions:
            if todo_task_description in elem_task:
                output = "todo"
                task_set = True
        if task_set:
            prompt+=output
            continue
        new_task = False
        for new_task_description in new_task_descriptions:
            if new_task_description in elem_task:
                new_task = True
        if new_task:
            status_schema = CompletionStatusOnlyNew
            choices: list[str] = []
            for status in status_schema:
                choices += [status.value]
            generator =outlines.generate.choice(llm, choices, sampler=choice_sampler)
        if i == len(ptt_list)-1:
            output = ""
        elif inprogress_set:
            output = "todo"
        elif current_number_of_todo_inprogress == min_number_of_todo_inprogress and not inprogress_set and inprogress_always_set:
            print("Forcing output to be inprogress")
            output = "inprogress"
        elif elem_curr_status == "inprogress":
            assert not new_task
            inprogress_choices: list[str] = []
            for status in InProgressCompletionStatus:
                inprogress_choices += [status.value]
            generator =outlines.generate.choice(llm, inprogress_choices, sampler=choice_sampler)
            output = generator(prompt)
        else:
            output = generator(prompt)
        if output == "done":
            current_number_of_todo_inprogress -= 1
        status_schema = CompletionStatus
        choices: list[str] = []
        for status in status_schema:
            choices += [status.value]
        generator =outlines.generate.choice(llm, choices, sampler=choice_sampler)
        if "inprogress" in output:
            inprogress_set = True
        prompt+=output
    return str2dict(prompt[original_prompt_len:]), inprogress_set

def add_whitespace(prompt, llm, sampler, max_spaces=4):
    if max_spaces == 0:
        return ""
    whitespace_regex = r"[ \t\r\n]*"
    generator = outlines.generate.regex(
        llm,
        whitespace_regex,
        sampler=sampler
    )
    whitespace = generator(prompt, max_tokens=max_spaces)
    return whitespace
def add_new_items_outlines(llm, prompt: str, ptt: dict, sampler, inprogress_set: bool = False, max_spaces=0, force_add_task: dict[str, int] = {}):
    """
    For this, we want our model to first output a completion status and then a task description. This is because we want it to know the
    status before thinking of what kind of task it is
    """
    """
    schema format:
    {"recon": [
        {"task_description": "Perform a full port scan", "status": "done"},
        {"task_description": "Determine the purpose of each open port", "status": "todo"}
    ], 
    "initial_access": [], 
    "execution": [], 
    "post_exploitation": []}
    """
    ptt = status_first_description_later(ptt)
    ptt = json.dumps(ptt)
    original_prompt_len = len(prompt)
    choice_sampler =  multinomial(top_k=50, top_p=1, temperature=1.0)


    if inprogress_set:
        task_schema = BaseTaskOnlyTODO
    else:
        task_schema = BaseTaskOnlyNew


    continue_choices = [",",  "],"]
    tree_dict = json.loads(ptt)
    prompt += add_whitespace(prompt, llm, sampler, max_spaces)
    prompt += "{"
    prompt += add_whitespace(prompt, llm, sampler, max_spaces)
    for i, key in enumerate(tree_dict):
        prompt += f'"{key}": ['
        prompt += add_whitespace(prompt, llm, sampler, max_spaces)
        if len(tree_dict[key]) > 0:
            for task in tree_dict[key]:
                task_description = task["task_description"]
                status = task["status"]
                prompt += '{"status":' + f'"{status}", "task_description": "{task_description}"' +'},'
            prompt = prompt[:-1]
            # The model always chooses to end the todo list
            if force_add_task.get(key, 0) > 0:
                force_add_task[key] -= 1
                prompt += ","
                prompt += add_whitespace(prompt, llm, sampler, max_spaces)
            else:
                generator =outlines.generate.choice(llm, continue_choices, sampler=choice_sampler)
                output = generator(prompt)
                prompt+=output
                if "]," in output:
                    if len(tree_dict) -1 != i:
                        prompt += add_whitespace(prompt, llm, sampler, max_spaces)
                    continue
                prompt += add_whitespace(prompt, llm, sampler, max_spaces)
        else:
            if force_add_task.get(key, 0) > 0:
                force_add_task[key] -= 1
                prompt += add_whitespace(prompt, llm, sampler, max_spaces)
            else:
                generator =outlines.generate.choice(llm, ['{"',  "],"], sampler=choice_sampler)
                if "]," in output:
                    prompt+=output
                    continue
        while True:
            generator = outlines.generate.json(llm, task_schema, sampler=sampler, whitespace_pattern=" \t\n\r")
            output = generator(prompt)
            safe_task_description = re.escape(output.task_description.replace("'",'"')).replace('\\','')
            output = json.dumps({"status": output.status.value, "task_description": f"{safe_task_description}"})
            if "inprogress" in output:
                task_schema =  BaseTaskOnlyTODO
            prompt += output
            if force_add_task.get(key, 0) > 0:
                force_add_task[key] -= 1
                prompt += ","
                prompt += add_whitespace(prompt, llm, sampler, max_spaces)
            else:
                generator =outlines.generate.choice(llm, continue_choices, sampler=sampler)
                output = generator(prompt)
                prompt += output
                if output == "],":
                    if len(tree_dict) -1 != i:
                        prompt += add_whitespace(prompt, llm, sampler, max_spaces)
                    break
    prompt = prompt[:-1] +"}"

    return str2dict(prompt[original_prompt_len:])