Failed to call the "abatch" method of OpenAI when the config "max_concurrency" is not set to 1

[pursuemoon]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

import asyncio

from typing import List

from langchain_openai import OpenAI, ChatOpenAI
from langchain_core.runnables import Runnable, RunnableSequence, RunnableParallel
from langchain_core.language_models.base import BaseLanguageModel
from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate
from langchain_core.output_parsers import StrOutputParser


def init_translation_chain():
    # Declare an OpenAI object.
    model: BaseLanguageModel = OpenAI(
        openai_api_key="EMPTY",
        openai_api_base=LLM_URL,
        model_name=LLM_MODEL,
        n=1,
        temperature=0,
        max_tokens=4096,
        frequency_penalty=0.2,
        presence_penalty=2,
        streaming=True
    )

    prompt = PromptTemplate(
        input_variables=['text'],
        template=(
            "USER: Please translate the following text into Chinese:\n\n{text}"
            "ASSISTANT: "
        ),
    )

    output_parser = StrOutputParser()
    chain: RunnableSequence = prompt | model | output_parser

    return chain

async def ask_llm_for_batch_translation(text_list: List[str], concurrency: int):
    texts = [{"text": text} for text in text_list]

    chain = init_translation_chain()

    return await chain.abatch(texts, config={"max_concurrency": concurrency})

async def _test():
    # This works.
    for it in await ask_llm_for_batch_translation(['How are you?', 'How do you do?'], 1):
        print(it)

    # This does not work.
    # for it in await ask_llm_for_batch_translation(['How are you?', 'How do you do?'], 2):
    #     print(it)

if __name__ == '__main__':
    asyncio.run(_test())

Description

I am writing a asynchronous function to process batch translation with a chain constructed by LCEL and with OpenAI. But it raise a ValueError when the value of the configure max_concurrency is set to more than 1. At the mean time, this function works when the value of max_concurrency is not 1. If I change to use LLMChain, this error will not happen, which is expected.

I am confused because calling with a batch size more than 1 is normal and reasonable. It is expected that no matter how high max_concurrency is set, batch or abatch could work.

I am really interested to use advanced LCEL feature when developing. Thank everyone to help me answer this question!

System Info

System Information

OS: Linux
OS Version: #1 SMP Tue Jun 20 11:48:01 UTC 2023
Python Version: 3.11.8 (main, Feb 26 2024, 21:39:34) [GCC 11.2.0]

Package Information

langchain_core: 0.1.40
langchain: 0.1.14
langchain_community: 0.0.31
langsmith: 0.1.40
langchain_openai: 0.1.1
langchain_text_splitters: 0.0.1

Packages not installed (Not Necessarily a Problem)

The following packages were not found:

langgraph
langserve

[dosubot[bot]]

Hey @pursuemoon! Great to see you diving deep into LangChain again. How's everything been since our last chat?

Based on the code you've shared and the error you're encountering, it seems like the issue arises from the way LangChain's batch and abatch methods handle concurrency. These methods attempt to process inputs in batches based on the max_concurrency value but do not implement an asynchronous execution model that supports true concurrency beyond a single thread or coroutine. Instead, they recursively call themselves with smaller batches, resetting max_concurrency to None to avoid further subdivision, which does not leverage asynchronous I/O operations or thread pools to execute in parallel.

To resolve this issue and utilize higher concurrency levels, you need to modify the batch and abatch methods to support concurrent execution of model predictions. This can be achieved by using asynchronous I/O operations with asyncio for the abatch method or by employing a thread pool executor for the batch method to handle synchronous calls in a concurrent manner. Here's an example modification for the abatch method using asyncio.gather to achieve concurrency:

async def abatch(
    self,
    inputs: List[LanguageModelInput],
    config: Optional[Union[RunnableConfig, List[RunnableConfig]]] = None,
    *,
    return_exceptions: bool = False,
    **kwargs: Any,
) -> List[str]:
    if not inputs:
        return []
    config = get_config_list(config, len(inputs))
    max_concurrency = config[0].get("max_concurrency", 1)  # Default to 1 if not set

    # Split inputs into chunks based on max_concurrency
    input_chunks = [inputs[i:i + max_concurrency] for i in range(0, len(inputs), max_concurrency)]

    results = []
    for chunk in input_chunks:
        # Use asyncio.gather to run predictions concurrently within each chunk
        chunk_results = await asyncio.gather(*[
            self.ainvoke(input, config=config[i], **kwargs) for i, input in enumerate(chunk)
        ], return_exceptions=return_exceptions)
        results.extend(chunk_results)

    return results

This modification ensures that within each chunk of inputs, predictions are run concurrently up to the limit specified by max_concurrency. Note that this is a simplified example and might need adjustments based on the full context of your application, especially error handling and configuration management for each input.

Regarding your question about any known limitations or restrictions in the OpenAI model or the langchain_openai package that could cause this issue, there are several points to consider:

1. The OpenAI model requires an API key for authentication. If the OPENAI_API_KEY is not provided or is incorrect, it will result in an authentication error.
2. The BaseOpenAI class has restrictions on streaming results. Specifically, it cannot stream results when the number of completions (n) is greater than 1 or when best_of is greater than 1.
3. The _update_token_usage method updates the token usage based on the response from the OpenAI API. If the response does not contain the expected usage keys, this could lead to incorrect token usage tracking.
4. The _generate and _agenerate methods check for errors in the response from the OpenAI API. If an error is present, it raises a ValueError. However, if the error handling is not robust enough or if the OpenAI API changes its error response format, it could lead to issues in error detection.
5. The get_token_ids method mentions that tiktoken is not supported for Python versions less than 3.8. This could limit the usability of the package for users on older Python versions.
6. The modelname_to_contextsize method has a hardcoded mapping of model names to their maximum context sizes. If a model name is not in this mapping or if OpenAI introduces new models, this method will raise a ValueError, indicating an unknown model.
7. The max_tokens_for_prompt method calculates the maximum number of tokens that can be generated for a given prompt. This calculation depends on the model's maximum context size and the number of tokens in the prompt. If the prompt is too long, it might not leave enough tokens available for generating a response, leading to limitations in the output.
8. The BaseOpenAI class allows for the customization of the HTTP client used for making API requests. If the custom client is not properly configured (e.g., missing timeout settings or incorrect proxy settings), it could lead to issues in making API requests.

These limitations and restrictions highlight the importance of proper configuration, error handling, and understanding the operational constraints of the OpenAI model and the langchain_openai package.

Sources

• libs/core/langchain_core/language_models/llms.py
• libs/partners/openai/langchain_openai/llms/base.py

About Dosu

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.

[pursuemoon]

If I use LLMChain, the method abatch or batch could work correctly with the value of max_concurrency more than 1. So I don't think the problem lies in the implementation of OpenAI. Code could be write like this:

chain = LLMChain(prompt=prompt, llm=model, output_parser=output_parser)
texts = [{"text": text} for text in text_list]
result_list = await chain.abatch(texts, config={"max_concurrency": 8})

[Schlomoh]

Did you manage to get this working concurrently? I'm currently also struggling with concurrent invocations, only working like half of the time. The code you proposed @pursuemoon is similar to mine, which seems to be overseeing something...