Repository: imartinez/privateGPT
File: private_gpt/ui/ui.py (Python)

Code before:

class PrivateGptUi:
    def __init__(self) -> None:
        self._ingest_service = root_injector.get(IngestService)
        self._chat_service = root_injector.get(ChatService)
        self._chunks_service = root_injector.get(ChunksService)

        # Cache the UI blocks
        self._ui_block = None

    @staticmethod
    def _curate_sources(sources: list[Chunk]) -> set[Source]:
        return {
            Source(
                file=chunk.document.doc_metadata["file_name"]
                if chunk.document.doc_metadata
                and "file_name" in chunk.document.doc_metadata
                else "-",
                page=chunk.document.doc_metadata["page_label"]
                if chunk.document.doc_metadata
                and "page_label" in chunk.document.doc_metadata
                else "-",
                text=chunk.text,
            )
            for chunk in sources
        }

    def _chat(self, message: str, history: list[list[str]], mode: str, *_: Any) -> Any:
        def yield_deltas(completion_gen: CompletionGen) -> Iterable[str]:
            full_response: str = ""
            stream = completion_gen.response
            for delta in stream:
                if isinstance(delta, str):
                    full_response += str(delta)
                elif isinstance(delta, ChatResponse):
                    full_response += delta.delta or ""
                yield full_response

            if completion_gen.sources:
                full_response += SOURCES_SEPARATOR
                cur_sources = self._curate_sources(completion_gen.sources)
                sources_text = "\n\n\n".join(
                    f"{index}. {source.file} (page {source.page})"
                    for index, source in enumerate(cur_sources, start=1)
                )
                full_response += sources_text
            yield full_response

        def build_history() -> list[ChatMessage]:
            history_messages: list[ChatMessage] = list(
                itertools.chain(

Review comment (suggestion):
Could have been in `Source` class

```python
# Usage would look like
Source.curated_from_chunks(sources)
```

Code after:
    @staticmethod
    def curate_sources(sources: list[Chunk]) -> set["Source"]:
        curated_sources = set()

        for chunk in sources:
            doc_metadata = chunk.document.doc_metadata

            file_name = doc_metadata.get("file_name", "-") if doc_metadata else "-"
            page_label = doc_metadata.get("page_label", "-") if doc_metadata else "-"

            source = Source(file=file_name, page=page_label, text=chunk.text)
            curated_sources.add(source)

        return curated_sources


class PrivateGptUi:
    def __init__(self) -> None:
        self._ingest_service = root_injector.get(IngestService)
        self._chat_service = root_injector.get(ChatService)
        self._chunks_service = root_injector.get(ChunksService)

        # Cache the UI blocks
        self._ui_block = None

    def _chat(self, message: str, history: list[list[str]], mode: str, *_: Any) -> Any:
        def yield_deltas(completion_gen: CompletionGen) -> Iterable[str]:
            full_response: str = ""
            stream = completion_gen.response
            for delta in stream:
                if isinstance(delta, str):
                    full_response += str(delta)
                elif isinstance(delta, ChatResponse):
                    full_response += delta.delta or ""
                yield full_response

            if completion_gen.sources:
                full_response += SOURCES_SEPARATOR
                cur_sources = Source.curate_sources(completion_gen.sources)
                sources_text = "\n\n\n".join(
                    f"{index}. {source.file} (page {source.page})"
                    for index, source in enumerate(cur_sources, start=1)
                )
                full_response += sources_text
            yield full_response

        def build_history() -> list[ChatMessage]:
            history_messages: list[ChatMessage] = list(
                itertools.chain(
                    *[
                        [