Possible error with nests of nested queries in Python

[mlinegar]

I'm running into an issue when trying to call LMQL functions inside of other LMQL functions. I've produced a minimal example that shows what I'm running into, and attached it below.

In the example I'm trying to yes/no responses following some reasoning from the LLM. This works as expected when I try to do so in separate functions (e.g. in two_step_test below), but not when I do so in a single wrapper function (e.g. in one_step_test below).

Am I doing something wrong? Issue 189 seems related, but it's not clear to me what's going on.

import lmql
import nest_asyncio
nest_asyncio.apply()

# basic version of query
@lmql.query(chunksize=16)
def yes_no():
    '''lmql
    "Respond only with one of: 'yes' or 'no'."
    "A:[FLAG]" where FLAG in ["yes", "no"]
    return "yes" in FLAG.lower()
    '''
# simple reasoning
@lmql.query
def chain_of_thought():
    '''lmql
    "Think through the question step by step. [THOUGHTS]"
    return THOUGHTS
    '''
# combine reasoning and yes_no into one function
@lmql.query
def yes_no_reasoning():
    '''lmql
    "You must choose between 'yes' and 'no'."
    "[COT: chain_of_thought]"
    print("Inner COT: " + COT)
    "[ANS: yes_no]"
    return ANS
    '''
# test either yes_no or yes_no_reasoning
@lmql.query
def test_q(question, two_step=False):
    '''lmql
    "Q: {question}"
    if two_step:
        "[COT: chain_of_thought]"
        print("Outer COT: " + COT)
        "[ANS: yes_no]"    
    else: 
        "[ANS: yes_no_reasoning]"
    print(ANS)
    return ANS
    '''
# works as expected!
two_step_test = test_q("Is the sky blue?", two_step=True)
# Outer COT: 
# A: Yes, the sky appears blue to the human eye during the day. This is because of a phenomenon called Rayleigh scattering, where the shorter blue wavelengths of light are scattered more in the Earth's atmosphere compared to longer wavelengths like red. This results in the blue color we see when we look up at the sky. However, the sky can also appear different colors at different times of the day, such as orange during sunrise and sunset, and black at night. Additionally, the color of the sky can also be affected by factors such as pollution, clouds, and weather conditions.
# True
# seems to keep looping
one_step_test = test_q("Is the sky blue?", two_step=False)
# Inner COT: Think through the question step by step.
# Yes, the sky appears blue to the human eye during the day due to the scattering of sunlight by the Earth's atmosphere. However, at night or in certain weather conditions, the sky may appear to be a different color. So, the answer is yes, the sky is typically perceived as blue.
# Inner COT:  sky appears blue to the human eye during the day due to the scattering of sunlight by the Earth's atmosphere. However, at night or in certain weather conditions, the sky may appear to be a different color. So, the answer is yes, the sky is typically perceived as blue.Respond only with one of: 'yes' or 'no'.A: Yes
# Inner COT:  to the scattering of sunlight by the Earth's atmosphere. However, at night or in certain weather conditions, the sky may appear to be a different color. So, the answer is yes, the sky is typically perceived as blue.Respond only with one of: 'yes' or 'no'.A: YesRespond only with one of: 'yes' or 'no'.A: Yes
# Inner COT:  However, at night or in certain weather conditions, the sky may appear to be a different color. So, the answer is yes, the sky is typically perceived as blue.Respond only with one of: 'yes' or 'no'.A: YesRespond only with one of: 'yes' or 'no'.A: YesRespond only with one of: 'yes' or 'no'.A: Yes
# Inner COT:  sky may appear to be a different color. So, the answer is yes, the sky is typically perceived as blue.Respond only with one of: 'yes' or 'no'.A: YesRespond only with one of: 'yes' or 'no'.A: YesRespond only with one of: 'yes' or 'no'.A: YesRespond only with one of: 'yes' or 'no'.A: Yes
# Inner COT:  the answer is yes, the sky is typically perceived as blue.Respond only with one of: 'yes' or 'no'.A: YesRespond only with one of: 'yes' or 'no'.A: YesRespond only with one of: 'yes' or 'no'.A: YesRespond only with one of: 'yes' or 'no'.A: YesRespond only with one of: 'yes' or 'no'.A: Yes
# ...

[lbeurerkellner]

Thanks for reporting this. It seems to be a bug, I was able to reproduce it.

The extra level of indirection via yes_no_reasoning seems to cause an issue in the nested query function calling. A workaround for now is to remove this level of indirection. We will have to investigate further, before we can fix this.

[mlinegar]

Thanks! Being able to have the setup of yes_no_reasoning would be helpful, as I often try to have many layers of indirection (so as to reuse code as much as possible).

I'm extremely appreciative of the library and all the work you do to maintain it!