Asyncio exception when initializing LLMRails inside function

[lauradang]

I am running my FastAPI application via uvicorn main:app --reload --host 0.0.0.0 --loop asyncio via this bug fix: https://github.com/NVIDIA/NeMo-Guardrails/issues/112. However, I am running into another error when initializing LLMRails inside a function.

Erroneous code:

async def run_chain_with_guardrails(
    query: str,
    model_id: str,
    chain: Chain,
    config_path: Optional[str] = None
) -> dict:
    ''' Runs chain with guardrails
    '''
    with open(config_path, 'r') as file:
        colang_config = file.read()

    rails_config = RailsConfig.from_content(COLANG_CONFIG, YAML_CONFIG % model_id)
    rails_app = LLMRails(rails_config)

    rails_app.register_action(chain, name='chain')

    user_input = [{'role': 'user', 'content': query}]

    start_ts = time.time()
    result = await rails_app.generate_async(messages=user_input)
    end_ts = time.time()

    return result | {'response_time_ms': (end_ts - start_ts) * 1000}

Error traceback:

ERROR:asyncio:Exception in callback Task.task_wakeup(<Future finished result=None>)
handle: <Handle Task.task_wakeup(<Future finished result=None>)>
Traceback (most recent call last):
  File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/asyncio/events.py", line 80, in _run
    self._context.run(self._callback, *self._args)
RuntimeError: Cannot enter into task <Task pending name='Task-1' coro=<Server.serve() running at /Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/uvicorn/server.py:81> wait_for=<Future finished result=None> cb=[_run_until_complete_cb() at /Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/asyncio/base_events.py:184, <1 more>, WorkerThread.stop()]> while another task <Task pending name='Task-12' coro=<RequestResponseCycle.run_asgi() running at /Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/uvicorn/protocols/http/httptools_impl.py:435> cb=[set.discard()]> is being executed.

This only occurs when pinging the API at startup. When pinging the endpoint again, the error goes away (but the first request will always error out with this method).

However, when moving the initialization of LLMRails() to the top-level, this error goes away. I want the user to be able to specify certain things in the configuration files, so having the initialization in a function is necessary.

Code that has no errors: - but we cannot allow the user to specify rails_config (must be loaded at startup if in top-level)

rails_config = RailsConfig.from_content(COLANG_CONFIG, YAML_CONFIG % 'model_id')
rails_app = LLMRails(rails_config)

async def run_chain_with_guardrails(
    query: str,
    model_id: str,
    chain: Chain,
    config_path: Optional[str] = None
) -> dict:
    ''' Runs chain with guardrails (Streaming currently not supported)
    '''
    rails_app.register_action(chain, name='chain')

    user_input = [{
        'role': 'user',
        'content': query
    }]

    start_ts = time.time()
    result = await rails_app.generate_async(messages=user_input)
    end_ts = time.time()

    return result | {'response_time_ms': (end_ts - start_ts) * 1000}

[drazvan]

@lauradang : this has to do with some unexpected side-effects from nest_asyncio. We don't have a workaround just yet. We'll followup on this.

[sidharthsingla-oa]

Hi, Is it fixed and do we have a followup on it? Thanks!

[drazvan]

@sidharthsingla-oa : we've not looked into this recently. Are you getting the same error? What's your setup / steps to reproduce?

[sidharthsingla-oa]

@drazvan I'm using this code:

nest_asyncio.apply()

LOGGER = logging.getLogger(__name__)

class NemoGuard:
    def __init__(self, rails_config, chat_model):
        self.rails = LLMRails(rails_config, llm=chat_model, verbose=False)

    async def validate(self, message):
        messages = [
            {"role": "user", "content": message}
        ]
        ##new_message = self.rails.generate(messages=messages)
        ##new_message = await asyncio.run(self.rails.generate_async(messages=messages))
        
        new_message = await self.rails.generate_async(messages=messages)

        return new_message
                
  Error:
  

> [nemoguardrails.rails.llm.llmrails] - INFO: request:jwra2vkp --- :: Total processing took 2.35 seconds.
[nemoguardrails.rails.llm.llmrails] - INFO: request:jwra2vkp --- :: Stats: 2 total calls, 0.8056309223175049 total time, 931 total tokens, 911 total prompt tokens, 20 total completion tokens
 [uvicorn.error] - ERROR: request:jwra2vkp Exception in ASGI application
Traceback (most recent call last):
  File "/lib/python3.11/site-packages/uvicorn/protocols/http/h11_impl.py", line 429, in run_asgi
    result = await app(  # type: ignore[func-returns-value]
    
  I have tried uncommenting and using the commented statements as well one by one but the issue is the same. Also, tried using commenting nest_asyncio.
  
  Thanks!

[drazvan]

And how exactly is the NemoGuard class used? I see the error stack involves uvicorn. Do you have a server that invokes NemoGuard.validate at some point? Also, can you paste the full error stack? Thanks!