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

Code before:

        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
        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),

Review comment (suggestion):
```suggestion
        pw_ai_results_exploded = pw_ai_results.flatten(
            pw_ai_results.docs, origin_id='query_id'
        )
```

Code after:
            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

    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),
        )