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

Code before:
        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 if rerank_topk is not None else max(1, self.search_topk // 2)
        self.server: None | QASummaryRestServer = None

    def _init_schemas(self, default_llm_name: str | None = None) -> None:
        """Initialize API schemas with optional and non-optional arguments."""

        class PWAIQuerySchema(pw.Schema):
            prompt: str
            filters: str | None = pw.column_definition(default_value=None)
            model: str | None = pw.column_definition(default_value=default_llm_name)
            return_context_docs: bool = pw.column_definition(default_value=False)

        class SummarizeQuerySchema(pw.Schema):
            text_list: list[str]
            model: str | None = pw.column_definition(default_value=default_llm_name)

        self.AnswerQuerySchema = PWAIQuerySchema
        self.SummarizeQuerySchema = SummarizeQuerySchema
        self.RetrieveQuerySchema = self.indexer.RetrieveQuerySchema
        self.StatisticsQuerySchema = self.indexer.StatisticsQuerySchema
        self.InputsQuerySchema = self.indexer.InputsQuerySchema

    def _apply_reranking(self, pw_ai_results: pw.Table) -> pw.Table:
        """Apply reranking to retrieved documents."""

        @pw.udf

Review comment (suggestion):
I think I'd prefer if both `rerank_topk` and `reranker` need to be explicitly given. (i.e. if one of `reranker` and `rerank_topk` is `None`, set the other to `None` and print warning.) At the minimum, if you want to automatically set `rerank_topk`, there should be a warning, saying what `rerank_topk` is set to. 

Code after:
        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
        self.server: None | QASummaryRestServer = None

        # Check reranker settings mismatch
        if (self.reranker is None) != (rerank_topk is None):
            logging.warning(
                "Incomplete reranker configuration: "
                "Both 'reranker' and 'rerank_topk' must be specified. "
                "Reranking disabled."
            )
            self.reranker = None
            self.rerank_topk = None

    def _init_schemas(self, default_llm_name: str | None = None) -> None:
        """Initialize API schemas with optional and non-optional arguments."""

        class PWAIQuerySchema(pw.Schema):
            prompt: str
            filters: str | None = pw.column_definition(default_value=None)
            model: str | None = pw.column_definition(default_value=default_llm_name)
            return_context_docs: bool = pw.column_definition(default_value=False)

        class SummarizeQuerySchema(pw.Schema):
            text_list: list[str]
            model: str | None = pw.column_definition(default_value=default_llm_name)

        self.AnswerQuerySchema = PWAIQuerySchema
        self.SummarizeQuerySchema = SummarizeQuerySchema
        self.RetrieveQuerySchema = self.indexer.RetrieveQuerySchema
        self.StatisticsQuerySchema = self.indexer.StatisticsQuerySchema
        self.InputsQuerySchema = self.indexer.InputsQuerySchema