Spaces:
Sleeping
Sleeping
File size: 3,375 Bytes
d477d5c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
from __future__ import annotations
import schema
import yaml
from attrs import define
from yaml.resolver import Resolver
from griptape_statemachine.parsers.base_parser import BaseParser
STRUCTURE_SCHEMA = schema.Schema(
{
schema.Optional("model"): str,
schema.Optional("ruleset_ids"): [str],
schema.Optional("vector_stores"): [str],
schema.Optional("prompt_id"): str,
}
)
CONFIG_SCHEMA = schema.Schema(
{
"rulesets": schema.Schema(
{
str: schema.Schema(
{
"name": str,
"rules": [str],
}
)
}
),
# Added for vector stores
schema.Optional("vector_stores"): schema.Schema(
{
str: schema.Schema(
{
"file_path": str,
"file_type": str,
schema.Optional("max_tokens"): int,
}
)
}
),
"structures": schema.Schema({str: STRUCTURE_SCHEMA}),
"events": schema.Schema(
{
str: schema.Schema(
{
"transitions": [
schema.Schema(
{
"from": str,
"to": str,
schema.Optional("internal"): bool,
schema.Optional("on"): str,
schema.Optional("relevance"): str,
}
)
],
}
)
}
),
"states": schema.Schema(
{
str: schema.Schema(
{
schema.Optional(
schema.Or("initial", "final")
): bool, # pyright: ignore[reportArgumentType]
schema.Optional("structures"): schema.Schema(
{str: STRUCTURE_SCHEMA}
),
}
)
}
),
schema.Optional("prompts"): {
str: {schema.Optional("author_intent"): str, "prompt": str}
},
}
)
@define()
class UWConfigParser(BaseParser):
def __attrs_post_init__(self) -> None:
# remove resolver entries for On/Off/Yes/No
for ch in "OoYyNn":
if ch in Resolver.yaml_implicit_resolvers:
if len(Resolver.yaml_implicit_resolvers[ch]) == 1:
del Resolver.yaml_implicit_resolvers[ch]
else:
Resolver.yaml_implicit_resolvers[ch] = [
x
for x in Resolver.yaml_implicit_resolvers[ch]
if x[0] != "tag:yaml.org,2002:bool"
]
def parse(self) -> dict:
data = yaml.safe_load(self.file_path.read_text())
CONFIG_SCHEMA.validate(data)
return data
def update_and_save(self, config: dict) -> None:
with self.file_path.open("w") as file:
yaml.dump(config, file, default_flow_style=False, line_break="\n")
|