Repository: jerryjliu/gpt_index
File: llama-index-core/llama_index/core/llms/mock.py (Python)

Code before:
            elif isinstance(block, CitationBlock):
                if isinstance(block.cited_content, TextBlock):
                    content_parts.append(block.cited_content.text)
                else:
                    content_parts.append(
                        self._data_from_binary(
                            data=block.cited_content.image or b"",
                            block_type=block.cited_content.block_type,
                        )
                    )
            elif isinstance(block, ToolCallBlock):
                self.tool_calls.append(block)
            else:
                pass
        return "".join(content_parts)

    def stream_chat(
        self, messages: Sequence[ChatMessage], **kwargs: Any
    ) -> ChatResponseGen:
        content = self._content_from_content_blocks(messages[-1].blocks)
        if not content:
            content = "<empty>"
        response_msg = ChatMessage(role="assistant", content=content)

        def _gen():
            yield ChatResponse(
                message=response_msg,
                delta=content,
                raw={"content": content},
            )

        return _gen()

    async def astream_chat(
        self, messages: Sequence[ChatMessage], **kwargs: Any
    ) -> ChatResponseAsyncGen:
        content = self._content_from_content_blocks(messages[-1].blocks)
        if not content:
            content = "<empty>"
        response_msg = ChatMessage(role="assistant", content=content)

        async def _gen():
            yield ChatResponse(
                message=response_msg,
                delta=content,
                raw={"content": content},
            )

        return _gen()

    def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:

Review comment (suggestion):
I wonder if this would be slightly more usable if there was an option to pass in a map of input message(s) -> output messages? (this way people could trigger tool calls in testing). Not sure on the UX for that though

Code after:
                        data=block.cited_content.image
                        if isinstance(block.cited_content.image, bytes)
                        else b"",
                        block_type=block.cited_content.block_type,
                    )
                )
        elif isinstance(block, ToolCallBlock):
            if tool_calls is not None:
                tool_calls.append(block)
        else:
            pass
    return "".join(content_parts)


BlockToContentCallback: TypeAlias = Callable[
    [list[ContentBlock], Optional[list[ToolCallBlock]]], str
]


class MockFunctionCallingLLM(MockLLM):
    tool_calls: List[ToolCallBlock] = Field(default_factory=list)
    blocks_to_content_callback: BlockToContentCallback = Field(
        default=_default_blocks_to_content_callback
    )

    def __init__(
        self,
        max_tokens: Optional[int] = None,
        callback_manager: Optional[CallbackManager] = None,
        system_prompt: Optional[str] = None,
        messages_to_prompt: Optional[MessagesToPromptType] = None,
        completion_to_prompt: Optional[CompletionToPromptType] = None,
        pydantic_program_mode: PydanticProgramMode = PydanticProgramMode.DEFAULT,
        blocks_to_content_callback: Optional[BlockToContentCallback] = None,
    ) -> None:
        super().__init__(
            max_tokens=max_tokens,
            callback_manager=callback_manager or CallbackManager([]),
            system_prompt=system_prompt,
            messages_to_prompt=messages_to_prompt,
            completion_to_prompt=completion_to_prompt,
            pydantic_program_mode=pydantic_program_mode,
        )
        self.blocks_to_content_callback = (
            blocks_to_content_callback or self.blocks_to_content_callback
        )

    @property
    def metadata(self) -> LLMMetadata:
        return LLMMetadata(is_function_calling_model=True)