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

Code before:

        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
        def add_score_to_doc(doc: pw.Json, score: float) -> dict:
            return {**doc.as_dict(), "reranker_score": score}

        # Flatten docs into rows
        pw_ai_results_exploded = pw_ai_results.flatten(pw_ai_results.docs, origin_id='query_id')

        # Apply reranker
        pw_ai_results_scored = pw_ai_results_exploded.select(
            pw.this.prompt,
            pw.this.model,
            pw.this.filters,
            pw.this.return_context_docs,
            pw.this.query_id,
            doc=pw.this.docs,
            reranker_score=self.reranker(pw.this.docs["text"], pw.this.prompt)
        )

        pw_ai_results_scored = pw_ai_results_scored.await_futures()

        # Add score and sort
        pw_ai_results_scored = pw_ai_results_scored.with_columns(
            sort_key=-pw.this.reranker_score,
            doc_with_score=add_score_to_doc(pw.this.doc, pw.this.reranker_score)
        )

        # Reassemble documents
        pw_ai_results = pw_ai_results_scored.groupby(pw.this.query_id, sort_by=pw.this.sort_key).reduce(
            query_id=pw.this.query_id,
            prompt=pw.reducers.any(pw.this.prompt),
            model=pw.reducers.any(pw.this.model),
            filters=pw.reducers.any(pw.this.filters),
            return_context_docs=pw.reducers.any(pw.this.return_context_docs),
            docs=pw.reducers.tuple(pw.this.doc_with_score),
        ).with_id(pw.this.query_id)

        # Keep only top k
        return pw_ai_results.with_columns(
            docs=_limit_documents(pw.this.docs, k=self.rerank_topk)
        )


Review comment (bug):
```suggestion
            reranker_score=self.reranker(pw.this.docs["text"], pw.this.prompt),  # type: ignore
```
mypy fails to realize that reranker cannot be None here

Code after:

        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
        def add_score_to_doc(doc: pw.Json, score: float) -> dict:
            return {**doc.as_dict(), "reranker_score": score}

        # Flatten docs into rows
        pw_ai_results_exploded = pw_ai_results.flatten(
            pw_ai_results.docs, origin_id='query_id'
        )

        # Apply reranker
        pw_ai_results_scored = pw_ai_results_exploded.select(
            pw.this.prompt,
            pw.this.model,
            pw.this.filters,
            pw.this.return_context_docs,
            pw.this.query_id,
            doc=pw.this.docs,
            reranker_score=self.reranker(pw.this.docs["text"], pw.this.prompt),  # type: ignore
        )

        pw_ai_results_scored = pw_ai_results_scored.await_futures()

        # Add score and sort
        pw_ai_results_scored = pw_ai_results_scored.with_columns(
            sort_key=-pw.this.reranker_score,
            doc_with_score=add_score_to_doc(pw.this.doc, pw.this.reranker_score),
        )

        # Reassemble documents
        pw_ai_results = (
            pw_ai_results_scored.groupby(pw.this.query_id, sort_by=pw.this.sort_key)
            .reduce(
                query_id=pw.this.query_id,
                prompt=pw.reducers.any(pw.this.prompt),
                model=pw.reducers.any(pw.this.model),
                filters=pw.reducers.any(pw.this.filters),
                return_context_docs=pw.reducers.any(pw.this.return_context_docs),