Spaces:
Running
Running
| # TODO: create a csv parser | |
| from __future__ import annotations | |
| from ast import Lambda | |
| import contextlib | |
| import csv | |
| from pathlib import Path | |
| from typing import TYPE_CHECKING, Callable | |
| import yaml | |
| if TYPE_CHECKING: | |
| from io import TextIOWrapper | |
| class CsvParser: | |
| def __init__(self, directory: str) -> None: | |
| self.yaml_path = Path.joinpath(Path.cwd(), Path(f"{directory}/config.yaml")) | |
| self.csv_directory = Path.joinpath(Path.cwd(), Path(f"{directory}/csv_files")) | |
| csv_files = Path(self.csv_directory).glob("*") | |
| self.csv_file_paths = [file for file in csv_files if file.is_file()] | |
| def csv_parser(self) -> None: | |
| """This is going to take in a big csv, split it, and put it in config.yaml""" | |
| # This is going to parse multiple different csv files this time. | |
| split_csv = {} | |
| for csv_file in self.csv_file_paths: | |
| with Path.open(csv_file, "r", newline="") as csvfile: | |
| self.split_csv(csvfile, split_csv) | |
| # split_csv should have all the information | |
| yaml_data = yaml.safe_load(self.yaml_path.read_text()) | |
| # Rulesets CHANGE | |
| try: | |
| yaml_data["rulesets"] = self.csv_rulesets( | |
| split_csv["Ruleset ID"] | |
| ) # Rulesets | |
| except KeyError: | |
| print("No rulesets") | |
| # Agents DONE | |
| try: | |
| yaml_data["structures"] = self.csv_agents( | |
| split_csv["Agent ID"] | |
| ) # Agent Definitions | |
| except KeyError: | |
| print("No structures") | |
| # States | |
| # Tailoring (affects the states section only) CHANGE | |
| if "State ID to Tailor" in split_csv: | |
| try: | |
| yaml_data["states"] = self.csv_states( | |
| split_csv["State ID"], # State Definitions | |
| split_csv["State ID to Tailor"], # Agent Tailoring State ID | |
| ) | |
| except KeyError: | |
| print(" no states") | |
| else: | |
| try: | |
| yaml_data["states"] = self.csv_states( | |
| split_csv["State ID"], # State Definitions | |
| [], # Agent Tailoring State ID | |
| ) | |
| except KeyError: | |
| print(" no states") | |
| try: | |
| yaml_data["prompts"] = self.csv_prompts(split_csv["Prompt ID"]) | |
| except KeyError: | |
| print("no prompts") | |
| # # Transitioning (affects event section) DONE | |
| try: | |
| yaml_data["events"] = self.csv_transition_id( | |
| split_csv["Transition ID"] | |
| ) # State Transitions | |
| except KeyError: | |
| print("No transitions") | |
| # That's all folks! | |
| self.update_and_save(yaml_data) | |
| def split_csv(self, csv_file: TextIOWrapper, all_information: dict) -> None: | |
| """Takes in a csv_file, and splits it into a dictionary that is headed by each of the sections. | |
| Hooray! | |
| """ | |
| reader = csv.reader(csv_file) | |
| # Get the header of the section | |
| header = next(reader) | |
| header = header[0] # Go to the meat of it (get rid of descriptive header) | |
| current_information = [] | |
| for row in reader: | |
| key = row[0] | |
| # If the row is empty and/or has no value in the first column. | |
| if key == ",,": | |
| continue | |
| current_information.append({key: row[1:]}) | |
| all_information[header] = current_information | |
| def csv_kbs(self, kb_info: list) -> dict: | |
| dictionary = {} | |
| for row in kb_info: | |
| key, value = row.popitem() | |
| if key and value[0] and value[1]: | |
| dictionary[key] = {"file_path": value[0], "file_type": value[1]} | |
| return dictionary | |
| def csv_rulesets(self, ruleset_info: list) -> dict: | |
| dictionary = {} | |
| for row in ruleset_info: | |
| key, value = row.popitem() | |
| if key and value[0] and value[1]: | |
| rules = [ | |
| rule.strip().strip('"').lstrip("- ") | |
| for rule in value[1].split("\n") | |
| if rule.strip() | |
| ] | |
| dictionary[key] = { | |
| "name": value[0], | |
| "rules": rules, | |
| } # Will have to check this. | |
| return dictionary | |
| def csv_prompts(self, prompt_info: list) -> dict: | |
| dictionary = {} | |
| for row in prompt_info: | |
| key, value = row.popitem() | |
| if key and value[0]: | |
| dictionary[key] = {"prompt": value[0]} | |
| if value[1]: | |
| dictionary[key]["author_intent"] = value[1] | |
| return dictionary | |
| def csv_agents(self, agent_info: list) -> dict: | |
| dictionary = {} | |
| for row in agent_info: | |
| key, value = row.popitem() | |
| if key: | |
| ruleset_ids = [] | |
| if value[0]: | |
| ruleset_ids = [rule_id.strip() for rule_id in value[0].split(",")] | |
| config = { | |
| "model": "gpt-4o", | |
| "ruleset_ids": ruleset_ids, | |
| } | |
| # If there is a global KB used | |
| if value[1]: | |
| config["vector_stores"] = [value[1]] | |
| # If there is a global prompt used (can be overrided by state specfic) | |
| if value[2]: | |
| config["prompt_id"] = value[2] | |
| # If there is a model override | |
| if value[4]: | |
| config["model"] = value[4] | |
| dictionary[key] = config | |
| return dictionary | |
| def csv_states(self, state_info: list, tailor_info: list) -> dict: | |
| states = {} | |
| for row in state_info: | |
| key, value = row.popitem() | |
| if not key: | |
| continue | |
| if key == "start": | |
| states[key] = {"initial": True} | |
| elif key == "end": | |
| states[key] = {"final": True} | |
| else: | |
| states[key] = {} | |
| if value[0] and value[0] != "none": | |
| agent_list = {name.strip(): {} for name in value[0].split(",")} | |
| states[key]["structures"] = agent_list | |
| for row in tailor_info: | |
| tailor, value = row.popitem() | |
| if not tailor: | |
| continue | |
| structures = ( | |
| states[tailor]["structures"] | |
| if tailor in states and "structures" in states[tailor] | |
| else {} | |
| ) | |
| structure = value | |
| structure_name = structure[0] | |
| # if ruleset | |
| try: | |
| structure_ruleset = structure[1] | |
| structure_ruleset_list = [] | |
| for item in structure_ruleset.split(","): | |
| if item.strip() != "": | |
| structure_ruleset_list.append(item.strip()) | |
| if len(structure_ruleset_list): | |
| structures[structure_name] = { | |
| "ruleset_ids": structure_ruleset_list, | |
| } | |
| except KeyError: | |
| structures[structure_name] = {} | |
| try: | |
| if structure[2]: | |
| structures[structure_name]["prompt_id"] = structure[2] | |
| except KeyError: | |
| pass | |
| states[tailor] = {"structures": structures} | |
| return states | |
| def csv_transition_id(self, transition_info: list) -> dict: | |
| events = {} | |
| for row in transition_info: | |
| key, value = row.popitem() | |
| if key and value[0] and value[1]: | |
| if key in events: | |
| # Add the transition if there already are transitions | |
| events[key]["transitions"].append( | |
| {"from": value[0], "to": value[1]} | |
| ) | |
| else: | |
| # create the first transition | |
| events[key] = { | |
| "transitions": [ | |
| {"from": value[0], "to": value[1]}, | |
| ] | |
| } | |
| return events | |
| def update_and_save(self, config: dict) -> None: | |
| with self.yaml_path.open("w") as file: | |
| yaml.dump(config, file, default_flow_style=False, line_break="\n") | |
| if __name__ == "__main__": | |
| CsvParser("uw_programmatic").csv_parser() | |