Cohere LLM astream_complete

[anirbanbasu]

Hi,

Is there any reason why calling generator = await llm.astream_complete(prompt) on an instance of Cohere LLM (from llama_index.llms.cohere import Cohere) will throw the following error while there is no problem with Groq, Open AI, Anthropic and Ollama?

TypeError: object async_generator can't be used in 'await' expression

I am baffled. Is this intended?

(Note that API key has been set correctly using the api_key parameter during initialising an object of the Cohere class.)

[dosubot[bot]]

The error TypeError: object async_generator can't be used in 'await' expression occurs because generator = await llm.astream_complete(prompt) is trying to await an asynchronous generator, which is not allowed in Python. Instead, you should iterate over the asynchronous generator using async for.

Here is the correct way to handle the asynchronous generator:

async for response in llm.astream_complete(prompt):
    # Process each response chunk
    print(response)

This approach works for Cohere LLM and should be used instead of await for asynchronous generators ^[1].

To continue talking to Dosu, mention @dosu.

[dosubot[bot]]

The error TypeError: 'async for' requires an object with __aiter__ method, got coroutine occurs because the object being iterated over with async for is a coroutine, not an asynchronous iterator.

In the provided code, the mock_acompletion_with_retry and mock_achat_with_retry functions are defined as asynchronous functions that return a dictionary. These functions are used to mock the behavior of the Cohere LLM's asynchronous methods. However, they return a dictionary directly, which is not an asynchronous iterator.

To fix this, you need to ensure that the object returned by these functions is an asynchronous iterator. You can achieve this by using an asynchronous generator. Here is an example of how you can modify the mock_acompletion_with_retry function to return an asynchronous generator:

async def mock_acompletion_with_retry(*args: Any, **kwargs: Any):
    response = cohere.responses.Generations.from_dict(
        {
            "id": "21caa4c4-6b88-45f7-b144-14ef4985384c",
            "generations": [
                {
                    "id": "b5e2bb70-bc9c-4f86-a22e-5b5fd13a3482",
                    "text": "\n\nThis is indeed a test",
                    "finish_reason": "COMPLETE",
                }
            ],
            "prompt": "test prompt",
            "meta": {"api_version": {"version": "1"}},
        },
        return_likelihoods=False,
    )
    yield response

Similarly, modify the mock_achat_with_retry function:

async def mock_achat_with_retry(*args: Any, **kwargs: Any):
    response = cohere.responses.Chat.from_dict(
        {
            "chatlog": None,
            "citations": None,
            "conversation_id": None,
            "documents": None,
            "generation_id": "357d15b3-9bd4-4170-9439-2e4cef2242c8",
            "id": "25c3632f-2d2a-4e15-acbd-804b976d0568",
            "is_search_required": None,
            "message": "test prompt",
            "meta": {"api_version": {"version": "1"}},
            "preamble": None,
            "prompt": None,
            "response_id": "25c3632f-2d2a-4e15-acbd-804b976d0568",
            "search_queries": None,
            "search_results": None,
            "text": "\n\nThis is indeed a test",
            "token_count": {
                "billed_tokens": 66,
                "prompt_tokens": 64,
                "response_tokens": 9,
                "total_tokens": 73,
            },
        },
        client=None,
        message="test_prompt",
    )
    yield response

Additionally, you can use an AsyncIterator class to wrap the iterable and make it compatible with async for. Here is an example of how you can define and use such a class:

class AsyncIterator:
    def __init__(self, iterable: Iterable) -> None:
        self._iterable = iter(iterable)

    def __aiter__(self) -> AsyncIterator:
        return self

    async def __anext__(self) -> Any:
        try:
            return next(self._iterable)
        except StopIteration:
            raise StopAsyncIteration

# Example usage
async def mock_acompletion_with_retry(*args: Any, **kwargs: Any):
    response = cohere.responses.Generations.from_dict(
        {
            "id": "21caa4c4-6b88-45f7-b144-14ef4985384c",
            "generations": [
                {
                    "id": "b5e2bb70-bc9c-4f86-a22e-5b5fd13a3482",
                    "text": "\n\nThis is indeed a test",
                    "finish_reason": "COMPLETE",
                }
            ],
            "prompt": "test prompt",
            "meta": {"api_version": {"version": "1"}},
        },
        return_likelihoods=False,
    )
    return AsyncIterator([response])

With these changes, the async for loop will work correctly because the functions now return asynchronous generators or an AsyncIterator, which are compatible with async for ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[anirbanbasu]

The error TypeError: 'async for' requires an object with __aiter__ method, got coroutine occurs because the object being iterated over with async for is a coroutine, not an asynchronous iterator.

This is absolutely right. That is why I am trying to call the the astream_complete on the LLM that should generate an asynchronous iterator.

Like I said, it works perfectly fine with other LLMs but not with Cohere. I find that very odd.

[anirbanbasu]

@dosu How can I fix this problem? Or, is it a problem in the Cohere LLM implementation of LlamaIndex?

[dosubot[bot]]

The issue with astream_complete not generating an asynchronous iterator in the Cohere LLM while it works fine with other LLMs in LlamaIndex could be due to the implementation of the acompletion_with_retry function or the way the asynchronous response is handled in the astream_complete method.

Here is the relevant code snippet for astream_complete:

@llm_completion_callback()
async def astream_complete(
    self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
    all_kwargs = self._get_all_kwargs(**kwargs)
    all_kwargs["stream"] = True

    response = await acompletion_with_retry(
        aclient=self._aclient,
        max_retries=self.max_retries,
        chat=False,
        prompt=prompt,
        **all_kwargs,
    )

    async def gen() -> CompletionResponseAsyncGen:
        content = ""
        async for r in response:
            content_delta = r.text
            content += content_delta
            yield CompletionResponse(
                text=content, delta=content_delta, raw=r._asdict()
            )

    return gen()

To fix the issue, ensure that acompletion_with_retry correctly returns an asynchronous iterable response. If acompletion_with_retry does not yield results asynchronously, astream_complete will not function as expected.

Here is a possible fix:

1. Verify that acompletion_with_retry is implemented to handle asynchronous streaming correctly.
2. Ensure that the response object returned by acompletion_with_retry supports asynchronous iteration.

Additionally, based on the PaLM class implementation in LlamaIndex, ensure that the stream_complete method in your custom LLM class is correctly implemented to support asynchronous iteration. The astream_complete method in PaLM raises a NotImplementedError, indicating that streaming completion is not supported:

@llm_completion_callback()
def stream_complete(
    self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseGen:
    """Stream the answer to a query.

    NOTE: this is a beta feature. Will try to build or use
    better abstractions about response handling.

    Args:
        prompt (str): Prompt to use for prediction.

    Returns:
        str: The predicted answer.
    """
    raise NotImplementedError(
        "PaLM does not support streaming completion in LlamaIndex currently."
    )

Ensure that your acompletion_with_retry function is correctly implemented to support asynchronous iteration. If acompletion_with_retry is correctly implemented and the astream_complete method supports asynchronous iteration, the astream_complete method should work as expected. If the LLM you are using does not support streaming completion, you may need to implement or use a different method that supports this feature ^{[1][2][3][4][5]}.

To continue talking to Dosu, mention @dosu.

[anirbanbasu]

Thanks. Seems like this should be raised as an issue so that the Cohere LLM code is fixed.