Stop streaming

Author: simondanielsson

src/common/tools/review_pull_request.py

@@ -26,7 +26,7 @@ async def arun(
         cls: type[ReviewPullRequest],
         *args: Any,  # noqa: ARG003
         **kwargs: Any,
-    ) -> AsyncIterator[str]:
+    ) -> str:
         """Use the tool asynchronously.
 
         :return: Tool output.
@@ -47,7 +47,7 @@ async def arun(
             patch=request.patch,
             file_content=request.file_content,
         )
-        return ask_llm(
+        return await ask_llm(
             system_prompt=system_prompt,
             user_prompt=user_prompt,
         )

src/reviews/services/pull_requests_service.py

@@ -66,11 +66,7 @@ async def _review_file_diff(
     :param semaphore: The semaphore.
     :return: A mapping from the file name to the review content.
     """
-    answer = ""
     async with semaphore:
-        review_content_iterator = await ReviewPullRequest.arun(request=file_diff)
-
-        async for review_content in review_content_iterator:
-            answer += review_content
+        answer = await ReviewPullRequest.arun(request=file_diff)
 
     return {file_diff.filename: answer}

src/utils/llm.py

@@ -39,7 +39,7 @@ async def ask_llm(
     user_prompt: str,
     memory: list[BaseMessage] | None = None,
     llm_model: str = config.provider_to_llm[config.llm_provider],
-) -> AsyncIterator[str]:
+) -> str:
     """Ask the LLM for a response.
 
     :param system_prompt: The system prompt.
@@ -54,6 +54,4 @@ async def ask_llm(
     ]
     chat_model = get_chat_model(llm_model=llm_model)
 
-    answer_iterator = chat_model.astream(input=messages)
-    async for answer in answer_iterator:
-        yield answer.content
+    return (await chat_model.ainvoke(input=messages)).content