Repository: pathwaycom/pathway
File: python/pathway/xpacks/llm/question_answering.py (Python)

Code before:
    ... )
    >>> prompt_template = "Answer the question. Context: {context}. Question: {query}"  # doctest: +SKIP
    >>> rag = BaseRAGQuestionAnswerer(  # doctest: +SKIP
    ...     llm=chat,
    ...     indexer=vector_server,
    ...     reranker=reranker,
    ...     prompt_template=prompt_template,
    ... )
    >>> app = QASummaryRestServer(app_host, app_port, rag)  # doctest: +SKIP
    >>> app.run()  # doctest: +SKIP
    """  # noqa: E501

    def __init__(
        self,
        llm: BaseChat,
        indexer: VectorStoreServer | DocumentStore,
        *,
        default_llm_name: str | None = None,
        prompt_template: str | Callable[[str, str], str] | pw.UDF = prompts.prompt_qa,
        context_processor: (
            BaseContextProcessor | Callable[[list[dict] | list[Doc]], str] | pw.UDF
        ) = SimpleContextProcessor(),
        summarize_template: pw.UDF = prompts.prompt_summarize,
        search_topk: int = 6,
        reranker: LLMReranker | CrossEncoderReranker | EncoderReranker | FlashRankReranker = None,
        rerank_topk: int = None,
    ) -> None:

        self.llm = llm
        self.indexer = indexer
        self.reranker = reranker

        if default_llm_name is None:
            default_llm_name = llm.model

        self._init_schemas(default_llm_name)

        self.prompt_udf = _get_RAG_prompt_udf(prompt_template)

        if isinstance(context_processor, BaseContextProcessor):
            self.docs_to_context_transformer = context_processor.as_udf()
        elif isinstance(context_processor, pw.UDF):
            self.docs_to_context_transformer = context_processor
        elif callable(context_processor):
            self.docs_to_context_transformer = pw.udf(context_processor)
        else:
            raise ValueError(
                "Context processor must be type of one of the following: \
                             ~BaseContextProcessor | Callable[[list[dict] | list[Doc]], str] | ~pw.UDF"
            )


Review comment (suggestion):
```suggestion
        reranker: pw.UDF | None = None,
        rerank_topk: int | None = None,
```
Let's relax type of reranker. The best would be to have a Reranker class, that all others inherit from, but otherwise this will allow other rerankers to be implemented

Code after:
    ...     reranker=reranker,
    ...     prompt_template=prompt_template,
    ... )
    >>> app = QASummaryRestServer(app_host, app_port, rag)  # doctest: +SKIP
    >>> app.run()  # doctest: +SKIP
    """  # noqa: E501

    def __init__(
        self,
        llm: BaseChat,
        indexer: VectorStoreServer | DocumentStore,
        *,
        default_llm_name: str | None = None,
        prompt_template: str | Callable[[str, str], str] | pw.UDF = prompts.prompt_qa,
        context_processor: (
            BaseContextProcessor | Callable[[list[dict] | list[Doc]], str] | pw.UDF
        ) = SimpleContextProcessor(),
        summarize_template: pw.UDF = prompts.prompt_summarize,
        search_topk: int = 6,
        reranker: pw.UDF | None = None,
        rerank_topk: int | None = None,
    ) -> None:

        self.llm = llm
        self.indexer = indexer
        self.reranker = reranker

        if default_llm_name is None:
            default_llm_name = llm.model

        self._init_schemas(default_llm_name)

        self.prompt_udf = _get_RAG_prompt_udf(prompt_template)

        if isinstance(context_processor, BaseContextProcessor):
            self.docs_to_context_transformer = context_processor.as_udf()
        elif isinstance(context_processor, pw.UDF):
            self.docs_to_context_transformer = context_processor
        elif callable(context_processor):
            self.docs_to_context_transformer = pw.udf(context_processor)
        else:
            raise ValueError(
                "Context processor must be type of one of the following: \
                             ~BaseContextProcessor | Callable[[list[dict] | list[Doc]], str] | ~pw.UDF"
            )

        self._pending_endpoints: list[tuple] = []

        self.summarize_template = summarize_template
        self.search_topk = search_topk
        self.rerank_topk = rerank_topk