Repository: geekan/MetaGPT
File: metagpt/provider/anthropic_api.py (Python)

Code before:

    def get_choice_text(self, resp: Message) -> str:
        return resp.content[0].text

    async def _achat_completion(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> Message:
        resp: Message = await self.aclient.messages.create(**self._const_kwargs(messages))
        self._update_costs(resp.usage, self.model)
        return resp

    async def acompletion(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> Message:
        return await self._achat_completion(messages, timeout=self.get_timeout(timeout))

    async def _achat_completion_stream(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> str:
        stream = await self.aclient.messages.create(**self._const_kwargs(messages, stream=True))
        collected_content = []
        usage = Usage(input_tokens=0, output_tokens=0)
        async for event in stream:
            event_type = event.type
            if event_type == "message_start":
                usage.input_tokens = event.message.usage.input_tokens
                usage.output_tokens = event.message.usage.output_tokens
            elif event_type == "content_block_delta":
                content = event.delta.text
                log_llm_stream(content)
                if self.stream_pipe:
                    self.stream_pipe.set_message(content)

                collected_content.append(content)
            elif event_type == "message_delta":
                usage.output_tokens = event.usage.output_tokens  # update final output_tokens

        log_llm_stream("\n")
        self._update_costs(usage)
        full_content = "".join(collected_content)
        return full_content


Review comment (refactor):
The functionality provided by this code `self.stream_pipe.set_message(content)` duplicates that of `log_llm_stream(content)`. You can use the `set_llm_stream_logfunc` method to make `log_llm_stream` call your `stream_pipe.set_message` method.

Code after:

    def get_choice_text(self, resp: Message) -> str:
        return resp.content[0].text

    async def _achat_completion(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> Message:
        resp: Message = await self.aclient.messages.create(**self._const_kwargs(messages))
        self._update_costs(resp.usage, self.model)
        return resp

    async def acompletion(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> Message:
        return await self._achat_completion(messages, timeout=self.get_timeout(timeout))

    async def _achat_completion_stream(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> str:
        stream = await self.aclient.messages.create(**self._const_kwargs(messages, stream=True))
        collected_content = []
        usage = Usage(input_tokens=0, output_tokens=0)
        async for event in stream:
            event_type = event.type
            if event_type == "message_start":
                usage.input_tokens = event.message.usage.input_tokens
                usage.output_tokens = event.message.usage.output_tokens
            elif event_type == "content_block_delta":
                content = event.delta.text
                log_llm_stream(content)
                collected_content.append(content)
            elif event_type == "message_delta":
                usage.output_tokens = event.usage.output_tokens  # update final output_tokens

        log_llm_stream("\n")
        self._update_costs(usage)
        full_content = "".join(collected_content)
        return full_content