mirror of
https://github.com/hwchase17/langchain.git
synced 2025-08-11 22:04:37 +00:00
core[patch]: fix nested sections for mustache templating (#23747)
The prompt template variable detection only worked for singly-nested sections because we just kept track of whether we were in a section and then set that to false as soon as we encountered an end block. i.e. the following: ``` {{#outerSection}} {{variableThatShouldntShowUp}} {{#nestedSection}} {{nestedVal}} {{/nestedSection}} {{anotherVariableThatShouldntShowUp}} {{/outerSection}} ``` Would yield `['outerSection', 'anotherVariableThatShouldntShowUp']` as input_variables (whereas it should just yield `['outerSection']`). This fixes that by keeping track of the current depth and using a stack.
This commit is contained in:
parent
acc8fb3ead
commit
acc457f645
@ -97,19 +97,18 @@ def mustache_template_vars(
|
||||
) -> Set[str]:
|
||||
"""Get the variables from a mustache template."""
|
||||
vars: Set[str] = set()
|
||||
in_section = False
|
||||
section_depth = 0
|
||||
for type, key in mustache.tokenize(template):
|
||||
if type == "end":
|
||||
in_section = False
|
||||
elif in_section:
|
||||
continue
|
||||
section_depth -= 1
|
||||
elif (
|
||||
type in ("variable", "section", "inverted section", "no escape")
|
||||
and key != "."
|
||||
and section_depth == 0
|
||||
):
|
||||
vars.add(key.split(".")[0])
|
||||
if type in ("section", "inverted section"):
|
||||
in_section = True
|
||||
if type in ("section", "inverted section"):
|
||||
section_depth += 1
|
||||
return vars
|
||||
|
||||
|
||||
@ -122,12 +121,15 @@ def mustache_schema(
|
||||
"""Get the variables from a mustache template."""
|
||||
fields = {}
|
||||
prefix: Tuple[str, ...] = ()
|
||||
section_stack: List[Tuple[str, ...]] = []
|
||||
for type, key in mustache.tokenize(template):
|
||||
if key == ".":
|
||||
continue
|
||||
if type == "end":
|
||||
prefix = prefix[: -key.count(".")]
|
||||
if section_stack:
|
||||
prefix = section_stack.pop()
|
||||
elif type in ("section", "inverted section"):
|
||||
section_stack.append(prefix)
|
||||
prefix = prefix + tuple(key.split("."))
|
||||
fields[prefix] = False
|
||||
elif type in ("variable", "no escape"):
|
||||
|
@ -141,6 +141,115 @@ def test_mustache_prompt_from_template() -> None:
|
||||
},
|
||||
}
|
||||
|
||||
# more complex nested section/context variables
|
||||
template = """This{{#foo}}
|
||||
{{bar}}
|
||||
{{#baz}}
|
||||
{{qux}}
|
||||
{{/baz}}
|
||||
{{quux}}
|
||||
{{/foo}}is a test."""
|
||||
prompt = PromptTemplate.from_template(template, template_format="mustache")
|
||||
assert prompt.format(
|
||||
foo={"bar": "yo", "baz": [{"qux": "wassup"}], "quux": "hello"}
|
||||
) == (
|
||||
"""This
|
||||
yo
|
||||
wassup
|
||||
hello
|
||||
is a test."""
|
||||
)
|
||||
assert prompt.input_variables == ["foo"]
|
||||
assert prompt.input_schema.schema() == {
|
||||
"title": "PromptInput",
|
||||
"type": "object",
|
||||
"properties": {"foo": {"$ref": "#/definitions/foo"}},
|
||||
"definitions": {
|
||||
"foo": {
|
||||
"title": "foo",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"bar": {"title": "Bar", "type": "string"},
|
||||
"baz": {"$ref": "#/definitions/baz"},
|
||||
"quux": {"title": "Quux", "type": "string"},
|
||||
},
|
||||
},
|
||||
"baz": {
|
||||
"title": "baz",
|
||||
"type": "object",
|
||||
"properties": {"qux": {"title": "Qux", "type": "string"}},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
# triply nested section/context variables
|
||||
template = """This{{#foo}}
|
||||
{{bar}}
|
||||
{{#baz.qux}}
|
||||
{{#barfoo}}
|
||||
{{foobar}}
|
||||
{{/barfoo}}
|
||||
{{foobar}}
|
||||
{{/baz.qux}}
|
||||
{{quux}}
|
||||
{{/foo}}is a test."""
|
||||
prompt = PromptTemplate.from_template(template, template_format="mustache")
|
||||
assert prompt.format(
|
||||
foo={
|
||||
"bar": "yo",
|
||||
"baz": {
|
||||
"qux": [
|
||||
{"foobar": "wassup"},
|
||||
{"foobar": "yoyo", "barfoo": {"foobar": "hello there"}},
|
||||
]
|
||||
},
|
||||
"quux": "hello",
|
||||
}
|
||||
) == (
|
||||
"""This
|
||||
yo
|
||||
wassup
|
||||
hello there
|
||||
yoyo
|
||||
hello
|
||||
is a test."""
|
||||
)
|
||||
assert prompt.input_variables == ["foo"]
|
||||
assert prompt.input_schema.schema() == {
|
||||
"title": "PromptInput",
|
||||
"type": "object",
|
||||
"properties": {"foo": {"$ref": "#/definitions/foo"}},
|
||||
"definitions": {
|
||||
"foo": {
|
||||
"title": "foo",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"bar": {"title": "Bar", "type": "string"},
|
||||
"baz": {"$ref": "#/definitions/baz"},
|
||||
"quux": {"title": "Quux", "type": "string"},
|
||||
},
|
||||
},
|
||||
"baz": {
|
||||
"title": "baz",
|
||||
"type": "object",
|
||||
"properties": {"qux": {"$ref": "#/definitions/qux"}},
|
||||
},
|
||||
"qux": {
|
||||
"title": "qux",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"foobar": {"title": "Foobar", "type": "string"},
|
||||
"barfoo": {"$ref": "#/definitions/barfoo"},
|
||||
},
|
||||
},
|
||||
"barfoo": {
|
||||
"title": "barfoo",
|
||||
"type": "object",
|
||||
"properties": {"foobar": {"title": "Foobar", "type": "string"}},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
# section/context variables with repeats
|
||||
template = """This{{#foo}}
|
||||
{{bar}}
|
||||
|
Loading…
Reference in New Issue
Block a user