Skip to content

Commit

Permalink
Infer options argument in infer_gen_constraints
Browse files Browse the repository at this point in the history
  • Loading branch information
parkervg committed Jun 19, 2024
1 parent 374e296 commit 1a98559
Showing 1 changed file with 9 additions and 0 deletions.
9 changes: 9 additions & 0 deletions blendsql/_sqlglot.py
Original file line number Diff line number Diff line change
Expand Up @@ -724,6 +724,9 @@ def infer_gen_constraints(self, start: int, end: int) -> dict:
- pattern: regular expression pattern lambda to use in constrained decoding with Model
- See `create_pattern` for more info on these pattern lambdas
- options: Optional str default to pass to `options` argument in a QAIngredient
- Will have the form '{table}::{column}'
"""

def create_pattern(
Expand Down Expand Up @@ -763,6 +766,12 @@ def create_pattern(
predicate_literals: List[str] = []
if start_node is not None:
predicate_literals = get_predicate_literals(start_node)
if isinstance(start_node, exp.EQ):
if isinstance(start_node.args["this"], exp.Column):
# This is valid for a default `options` set
added_kwargs[
"options"
] = f"{start_node.args['this'].args['table'].name}::{start_node.args['this'].args['this'].name}"
if len(predicate_literals) > 0:
if all(isinstance(x, bool) for x in predicate_literals):
output_type = "boolean"
Expand Down

0 comments on commit 1a98559

Please sign in to comment.