Skip to content

GeneSetAnnotator

The main annotation engine for gene set analysis.

Overview

GeneSetAnnotator is the core class that orchestrates:

  1. Gene filtering by p-value and log2FC
  2. Pathway enrichment analysis
  3. Prompt construction
  4. LLM invocation

Basic Usage

from gs2txt import GeneSetAnnotator
from gs2txt.llm import OpenAIProvider

provider = OpenAIProvider(api_key="...", model_id="gpt-4")
annotator = GeneSetAnnotator(llm_provider=provider)

result = annotator.annotate(deg_df)

Class Reference

gs2txt.core.GeneSetAnnotator

Main class for gene set annotation using LLMs.

Source code in gs2txt/core.py 
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
class GeneSetAnnotator:
    """Main class for gene set annotation using LLMs."""

    def __init__(
        self,
        llm_provider: BaseLLMProvider,
        enrichment_method: Optional[Union[str, BaseEnrichment]] = "pathway",
        prompt_builder: Optional[PromptBuilder] = None,
        **enrichment_kwargs
    ):
        """
        Initialize annotator.

        Parameters
        ----------
        llm_provider : BaseLLMProvider
            LLM provider instance
        enrichment_method : str or BaseEnrichment
            Enrichment method or instance
        prompt_builder : PromptBuilder, optional
            Custom prompt builder
        **enrichment_kwargs
            Additional parameters for enrichment

        Examples
        --------
        >>> from gs2txt import GeneSetAnnotator
        >>> from gs2txt.llm.base import OpenAIProvider
        >>>
        >>> provider = OpenAIProvider(api_key="sk-...", model_id="gpt-4")
        >>> annotator = GeneSetAnnotator(llm_provider=provider)
        >>> result = annotator.annotate(deg_df)
        """
        self.llm_provider = llm_provider

        # Setup enrichment
        if isinstance(enrichment_method, str):
            self.enrichment = create_enrichment(enrichment_method, **enrichment_kwargs)
        elif isinstance(enrichment_method, BaseEnrichment):
            self.enrichment = enrichment_method
        else:
            self.enrichment = None

        # Setup prompt builder
        self.prompt_builder = prompt_builder or PromptBuilder()

    def _filter_genes_by_statistics(
        self,
        deg_df: pd.DataFrame,
        pvalue_threshold: float = 0.05,
        log2fc_threshold: float = 1.0,
        pvalue_column: str = "pvalue",
        log2fc_column: str = "logFC",
        max_gene_num: int = 60,
    ) -> list[str]:
        """
        Filter genes by statistical criteria (pvalue, log2FC, gene count).

        Parameters
        ----------
        deg_df : pd.DataFrame
            Input dataframe with gene column and optional statistical columns
        pvalue_threshold : float, default 0.05
            P-value threshold for filtering (genes with pvalue <= threshold are kept)
        log2fc_threshold : float, default 1.0
            Log2 fold-change threshold (genes with |log2FC| >= threshold are kept)
        pvalue_column : str, default "pvalue"
            Column name for p-values
        log2fc_column : str, default "logFC"
            Column name for log2 fold-change values
        max_gene_num : int, default 60
            Maximum number of genes to return

        Returns
        -------
        List[str]
            Filtered list of gene names, sorted by p-value (if available)

        Notes
        -----
        - If statistical columns are missing, only max_gene_num limit is applied
        - Genes are sorted by p-value (ascending) before limiting to max_gene_num
        - All genes as a complete list for enrichment and LLM input
        """
        df = deg_df.copy()

        # Apply p-value filtering if column exists
        if pvalue_column in df.columns:
            df = df[df[pvalue_column] <= pvalue_threshold]

        # Apply log2FC filtering if column exists
        # Support both "log2FoldChange" and "logFC" column names
        fc_col = None
        if log2fc_column in df.columns:
            fc_col = log2fc_column
        elif "log2FoldChange" in df.columns:
            fc_col = "log2FoldChange"

        if fc_col is not None:
            df = df[abs(df[fc_col]) >= log2fc_threshold]

        # Sort by p-value for prioritization (if available)
        if pvalue_column in df.columns:
            df = df.sort_values(pvalue_column)

        # Extract gene names and limit to max_gene_num
        genes = df["gene"].dropna().astype(str).tolist()[:max_gene_num]
        return genes

    def annotate(
        self,
        deg_df: pd.DataFrame,
        max_gene_num: int = 60,
        max_pathway_num: int = 10,
        pathways: Optional[list[str]] = None,
        compute_enrichment: bool = True,
        additional_context: Optional[str] = None,
        # Differential gene filtering parameters
        pvalue_threshold: float = 0.05,
        log2fc_threshold: float = 1.0,
        pvalue_column: str = "pvalue",
        log2fc_column: str = "logFC",
    ) -> str:
        """
        Annotate a gene set with biological process description.

        Parameters
        ----------
        deg_df : pd.DataFrame
            DataFrame with at least a 'gene' column
        max_gene_num : int
            Maximum genes to use
        max_pathway_num : int
            Maximum pathways to use
        pathways : List[str], optional
            Pre-computed pathway terms (skip enrichment if provided)
        compute_enrichment : bool
            Whether to compute enrichment if pathways not provided
        additional_context : str, optional
            Additional context to include in prompt (e.g., PPI info)
        pvalue_threshold : float, default 0.05
            P-value threshold for gene filtering (genes with pvalue <= threshold)
        log2fc_threshold : float, default 1.0
            Log2 fold-change threshold (genes with |log2FC| >= threshold)
        pvalue_column : str, default "pvalue"
            Column name for p-values in deg_df
        log2fc_column : str, default "logFC"
            Column name for log2 fold-change in deg_df

        Returns
        -------
        str
            LLM-generated biological process annotation

        Raises
        ------
        ValueError
            If deg_df doesn't contain 'gene' column

        Examples
        --------
        >>> # Basic usage
        >>> result = annotator.annotate(deg_df)

        >>> # With pre-computed pathways
        >>> result = annotator.annotate(
        ...     deg_df,
        ...     pathways=["Inflammatory response", "Apoptosis"]
        ... )

        >>> # With additional context
        >>> ppi_info = "Hub genes: TP53, MYC"
        >>> result = annotator.annotate(deg_df, additional_context=ppi_info)
        """
        # Validate input - return empty string for empty input
        if deg_df is None or len(deg_df) == 0:
            return ""

        if "gene" not in deg_df.columns:
            raise ValueError("deg_df must contain a 'gene' column")

        # Extract genes with statistical filtering
        genes = self._filter_genes_by_statistics(
            deg_df,
            pvalue_threshold=pvalue_threshold,
            log2fc_threshold=log2fc_threshold,
            pvalue_column=pvalue_column,
            log2fc_column=log2fc_column,
            max_gene_num=max_gene_num,
        )

        # Return empty string if no genes after filtering
        if not genes:
            return ""

        # Get pathways
        pathway_terms = None
        if pathways is not None:
            pathway_terms = pathways[:max_pathway_num]
        elif compute_enrichment and self.enrichment is not None:
            try:
                enr_results = self.enrichment.enrich(genes)
                if enr_results is not None and len(enr_results) > 0:
                    enr_results = enr_results.sort_values("Adjusted P-value")
                    pathway_terms = (
                        enr_results["Term"].astype(str).tolist()[:max_pathway_num]
                    )
            except Exception as e:
                print(f"Warning: Enrichment failed: {e}")

        # Build prompt
        messages = self.prompt_builder.build(
            genes=genes, pathways=pathway_terms, additional_context=additional_context
        )

        # Generate annotation
        try:
            return self.llm_provider.generate(messages)
        except Exception as e:
            return f"Process: Failed\nError: {str(e)}"

    def annotate_detailed(
        self,
        deg_df: pd.DataFrame,
        max_gene_num: int = 60,
        max_pathway_num: int = 10,
        pathways: Optional[list[str]] = None,
        compute_enrichment: bool = True,
        additional_context: Optional[str] = None,
        # Differential gene filtering parameters
        pvalue_threshold: float = 0.05,
        log2fc_threshold: float = 1.0,
        pvalue_column: str = "pvalue",
        log2fc_column: str = "logFC",
    ) -> dict[str, Any]:
        """
        Annotate a gene set with biological process description and return detailed info.

        Parameters
        ----------
        deg_df : pd.DataFrame
            DataFrame with at least a 'gene' column
        max_gene_num : int
            Maximum genes to use
        max_pathway_num : int
            Maximum pathways to use
        pathways : List[str], optional
            Pre-computed pathway terms (skip enrichment if provided)
        compute_enrichment : bool
            Whether to compute enrichment if pathways not provided
        additional_context : str, optional
            Additional context to include in prompt (e.g., PPI info)
        pvalue_threshold : float, default 0.05
            P-value threshold for gene filtering
        log2fc_threshold : float, default 1.0
            Log2 fold-change threshold
        pvalue_column : str, default "pvalue"
            Column name for p-values in deg_df
        log2fc_column : str, default "logFC"
            Column name for log2 fold-change in deg_df

        Returns
        -------
        Dict[str, Any]
            Dictionary containing:
            - annotation: LLM-generated biological process annotation
            - pathways: comma-separated pathway terms (or empty string)
            - ppis: additional context / PPI info (or empty string)
            - final_prompt: the final prompt sent to LLM (or empty string)
        """
        result = {
            "annotation": "",
            "pathways": "",
            "ppis": "",
            "final_prompt": ""
        }

        # Validate input - return empty result for empty input
        if deg_df is None or len(deg_df) == 0:
            return result

        if "gene" not in deg_df.columns:
            raise ValueError("deg_df must contain a 'gene' column")

        # Extract genes with statistical filtering
        genes = self._filter_genes_by_statistics(
            deg_df,
            pvalue_threshold=pvalue_threshold,
            log2fc_threshold=log2fc_threshold,
            pvalue_column=pvalue_column,
            log2fc_column=log2fc_column,
            max_gene_num=max_gene_num,
        )

        # Return empty result if no genes after filtering
        if not genes:
            return result

        # Get pathways
        pathway_terms = None
        if pathways is not None:
            pathway_terms = pathways[:max_pathway_num]
        elif compute_enrichment and self.enrichment is not None:
            try:
                enr_results = self.enrichment.enrich(genes)
                if enr_results is not None and len(enr_results) > 0:
                    enr_results = enr_results.sort_values("Adjusted P-value")
                    pathway_terms = (
                        enr_results["Term"].astype(str).tolist()[:max_pathway_num]
                    )
            except Exception as e:
                print(f"Warning: Enrichment failed: {e}")

        # Store pathways in result
        if pathway_terms:
            result["pathways"] = ", ".join(pathway_terms)

        # Store PPIs/additional context in result
        if additional_context:
            result["ppis"] = additional_context

        # Build prompt
        messages = self.prompt_builder.build(
            genes=genes, pathways=pathway_terms, additional_context=additional_context
        )

        # Store final prompt (user message content)
        user_msg = next((m["content"] for m in messages if m["role"] == "user"), "")
        result["final_prompt"] = user_msg

        # Generate annotation
        try:
            result["annotation"] = self.llm_provider.generate(messages)
        except Exception as e:
            result["annotation"] = f"Process: Failed\nError: {str(e)}"

        return result

Functions

__init__(llm_provider, enrichment_method='pathway', prompt_builder=None, **enrichment_kwargs)

Initialize annotator.

Parameters:

Name Type Description Default
llm_provider BaseLLMProvider

LLM provider instance

 required
enrichment_method str or BaseEnrichment

Enrichment method or instance

 'pathway'
prompt_builder PromptBuilder

Custom prompt builder

 None
**enrichment_kwargs

Additional parameters for enrichment

 {}

Examples:

>>> from gs2txt import GeneSetAnnotator
>>> from gs2txt.llm.base import OpenAIProvider
>>>
>>> provider = OpenAIProvider(api_key="sk-...", model_id="gpt-4")
>>> annotator = GeneSetAnnotator(llm_provider=provider)
>>> result = annotator.annotate(deg_df)
Source code in gs2txt/core.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def __init__(
    self,
    llm_provider: BaseLLMProvider,
    enrichment_method: Optional[Union[str, BaseEnrichment]] = "pathway",
    prompt_builder: Optional[PromptBuilder] = None,
    **enrichment_kwargs
):
    """
    Initialize annotator.

    Parameters
    ----------
    llm_provider : BaseLLMProvider
        LLM provider instance
    enrichment_method : str or BaseEnrichment
        Enrichment method or instance
    prompt_builder : PromptBuilder, optional
        Custom prompt builder
    **enrichment_kwargs
        Additional parameters for enrichment

    Examples
    --------
    >>> from gs2txt import GeneSetAnnotator
    >>> from gs2txt.llm.base import OpenAIProvider
    >>>
    >>> provider = OpenAIProvider(api_key="sk-...", model_id="gpt-4")
    >>> annotator = GeneSetAnnotator(llm_provider=provider)
    >>> result = annotator.annotate(deg_df)
    """
    self.llm_provider = llm_provider

    # Setup enrichment
    if isinstance(enrichment_method, str):
        self.enrichment = create_enrichment(enrichment_method, **enrichment_kwargs)
    elif isinstance(enrichment_method, BaseEnrichment):
        self.enrichment = enrichment_method
    else:
        self.enrichment = None

    # Setup prompt builder
    self.prompt_builder = prompt_builder or PromptBuilder()
annotate(deg_df, max_gene_num=60, max_pathway_num=10, pathways=None, compute_enrichment=True, additional_context=None, pvalue_threshold=0.05, log2fc_threshold=1.0, pvalue_column='pvalue', log2fc_column='logFC')

Annotate a gene set with biological process description.

Parameters:

Name Type Description Default
deg_df DataFrame

DataFrame with at least a 'gene' column

 required
max_gene_num int

Maximum genes to use

 60
max_pathway_num int

Maximum pathways to use

 10
pathways List[str]

Pre-computed pathway terms (skip enrichment if provided)

 None
compute_enrichment bool

Whether to compute enrichment if pathways not provided

 True
additional_context str

Additional context to include in prompt (e.g., PPI info)

 None
pvalue_threshold float

P-value threshold for gene filtering (genes with pvalue <= threshold)

 0.05
log2fc_threshold float

Log2 fold-change threshold (genes with |log2FC| >= threshold)

 1.0
pvalue_column str

Column name for p-values in deg_df

 "pvalue"
log2fc_column str

Column name for log2 fold-change in deg_df

 "logFC"

Returns:

Type Description
str

LLM-generated biological process annotation

Raises:

Type Description
ValueError

If deg_df doesn't contain 'gene' column

Examples:

>>> # Basic usage
>>> result = annotator.annotate(deg_df)

>>> # With pre-computed pathways
>>> result = annotator.annotate(
...     deg_df,
...     pathways=["Inflammatory response", "Apoptosis"]
... )

>>> # With additional context
>>> ppi_info = "Hub genes: TP53, MYC"
>>> result = annotator.annotate(deg_df, additional_context=ppi_info)
Source code in gs2txt/core.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
def annotate(
    self,
    deg_df: pd.DataFrame,
    max_gene_num: int = 60,
    max_pathway_num: int = 10,
    pathways: Optional[list[str]] = None,
    compute_enrichment: bool = True,
    additional_context: Optional[str] = None,
    # Differential gene filtering parameters
    pvalue_threshold: float = 0.05,
    log2fc_threshold: float = 1.0,
    pvalue_column: str = "pvalue",
    log2fc_column: str = "logFC",
) -> str:
    """
    Annotate a gene set with biological process description.

    Parameters
    ----------
    deg_df : pd.DataFrame
        DataFrame with at least a 'gene' column
    max_gene_num : int
        Maximum genes to use
    max_pathway_num : int
        Maximum pathways to use
    pathways : List[str], optional
        Pre-computed pathway terms (skip enrichment if provided)
    compute_enrichment : bool
        Whether to compute enrichment if pathways not provided
    additional_context : str, optional
        Additional context to include in prompt (e.g., PPI info)
    pvalue_threshold : float, default 0.05
        P-value threshold for gene filtering (genes with pvalue <= threshold)
    log2fc_threshold : float, default 1.0
        Log2 fold-change threshold (genes with |log2FC| >= threshold)
    pvalue_column : str, default "pvalue"
        Column name for p-values in deg_df
    log2fc_column : str, default "logFC"
        Column name for log2 fold-change in deg_df

    Returns
    -------
    str
        LLM-generated biological process annotation

    Raises
    ------
    ValueError
        If deg_df doesn't contain 'gene' column

    Examples
    --------
    >>> # Basic usage
    >>> result = annotator.annotate(deg_df)

    >>> # With pre-computed pathways
    >>> result = annotator.annotate(
    ...     deg_df,
    ...     pathways=["Inflammatory response", "Apoptosis"]
    ... )

    >>> # With additional context
    >>> ppi_info = "Hub genes: TP53, MYC"
    >>> result = annotator.annotate(deg_df, additional_context=ppi_info)
    """
    # Validate input - return empty string for empty input
    if deg_df is None or len(deg_df) == 0:
        return ""

    if "gene" not in deg_df.columns:
        raise ValueError("deg_df must contain a 'gene' column")

    # Extract genes with statistical filtering
    genes = self._filter_genes_by_statistics(
        deg_df,
        pvalue_threshold=pvalue_threshold,
        log2fc_threshold=log2fc_threshold,
        pvalue_column=pvalue_column,
        log2fc_column=log2fc_column,
        max_gene_num=max_gene_num,
    )

    # Return empty string if no genes after filtering
    if not genes:
        return ""

    # Get pathways
    pathway_terms = None
    if pathways is not None:
        pathway_terms = pathways[:max_pathway_num]
    elif compute_enrichment and self.enrichment is not None:
        try:
            enr_results = self.enrichment.enrich(genes)
            if enr_results is not None and len(enr_results) > 0:
                enr_results = enr_results.sort_values("Adjusted P-value")
                pathway_terms = (
                    enr_results["Term"].astype(str).tolist()[:max_pathway_num]
                )
        except Exception as e:
            print(f"Warning: Enrichment failed: {e}")

    # Build prompt
    messages = self.prompt_builder.build(
        genes=genes, pathways=pathway_terms, additional_context=additional_context
    )

    # Generate annotation
    try:
        return self.llm_provider.generate(messages)
    except Exception as e:
        return f"Process: Failed\nError: {str(e)}"
annotate_detailed(deg_df, max_gene_num=60, max_pathway_num=10, pathways=None, compute_enrichment=True, additional_context=None, pvalue_threshold=0.05, log2fc_threshold=1.0, pvalue_column='pvalue', log2fc_column='logFC')

Annotate a gene set with biological process description and return detailed info.

Parameters:

Name Type Description Default
deg_df DataFrame

DataFrame with at least a 'gene' column

 required
max_gene_num int

Maximum genes to use

 60
max_pathway_num int

Maximum pathways to use

 10
pathways List[str]

Pre-computed pathway terms (skip enrichment if provided)

 None
compute_enrichment bool

Whether to compute enrichment if pathways not provided

 True
additional_context str

Additional context to include in prompt (e.g., PPI info)

 None
pvalue_threshold float

P-value threshold for gene filtering

 0.05
log2fc_threshold float

Log2 fold-change threshold

 1.0
pvalue_column str

Column name for p-values in deg_df

 "pvalue"
log2fc_column str

Column name for log2 fold-change in deg_df

 "logFC"

Returns:

Type Description
Dict[str, Any]

Dictionary containing: - annotation: LLM-generated biological process annotation - pathways: comma-separated pathway terms (or empty string) - ppis: additional context / PPI info (or empty string) - final_prompt: the final prompt sent to LLM (or empty string)
Source code in gs2txt/core.py 
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
def annotate_detailed(
    self,
    deg_df: pd.DataFrame,
    max_gene_num: int = 60,
    max_pathway_num: int = 10,
    pathways: Optional[list[str]] = None,
    compute_enrichment: bool = True,
    additional_context: Optional[str] = None,
    # Differential gene filtering parameters
    pvalue_threshold: float = 0.05,
    log2fc_threshold: float = 1.0,
    pvalue_column: str = "pvalue",
    log2fc_column: str = "logFC",
) -> dict[str, Any]:
    """
    Annotate a gene set with biological process description and return detailed info.

    Parameters
    ----------
    deg_df : pd.DataFrame
        DataFrame with at least a 'gene' column
    max_gene_num : int
        Maximum genes to use
    max_pathway_num : int
        Maximum pathways to use
    pathways : List[str], optional
        Pre-computed pathway terms (skip enrichment if provided)
    compute_enrichment : bool
        Whether to compute enrichment if pathways not provided
    additional_context : str, optional
        Additional context to include in prompt (e.g., PPI info)
    pvalue_threshold : float, default 0.05
        P-value threshold for gene filtering
    log2fc_threshold : float, default 1.0
        Log2 fold-change threshold
    pvalue_column : str, default "pvalue"
        Column name for p-values in deg_df
    log2fc_column : str, default "logFC"
        Column name for log2 fold-change in deg_df

    Returns
    -------
    Dict[str, Any]
        Dictionary containing:
        - annotation: LLM-generated biological process annotation
        - pathways: comma-separated pathway terms (or empty string)
        - ppis: additional context / PPI info (or empty string)
        - final_prompt: the final prompt sent to LLM (or empty string)
    """
    result = {
        "annotation": "",
        "pathways": "",
        "ppis": "",
        "final_prompt": ""
    }

    # Validate input - return empty result for empty input
    if deg_df is None or len(deg_df) == 0:
        return result

    if "gene" not in deg_df.columns:
        raise ValueError("deg_df must contain a 'gene' column")

    # Extract genes with statistical filtering
    genes = self._filter_genes_by_statistics(
        deg_df,
        pvalue_threshold=pvalue_threshold,
        log2fc_threshold=log2fc_threshold,
        pvalue_column=pvalue_column,
        log2fc_column=log2fc_column,
        max_gene_num=max_gene_num,
    )

    # Return empty result if no genes after filtering
    if not genes:
        return result

    # Get pathways
    pathway_terms = None
    if pathways is not None:
        pathway_terms = pathways[:max_pathway_num]
    elif compute_enrichment and self.enrichment is not None:
        try:
            enr_results = self.enrichment.enrich(genes)
            if enr_results is not None and len(enr_results) > 0:
                enr_results = enr_results.sort_values("Adjusted P-value")
                pathway_terms = (
                    enr_results["Term"].astype(str).tolist()[:max_pathway_num]
                )
        except Exception as e:
            print(f"Warning: Enrichment failed: {e}")

    # Store pathways in result
    if pathway_terms:
        result["pathways"] = ", ".join(pathway_terms)

    # Store PPIs/additional context in result
    if additional_context:
        result["ppis"] = additional_context

    # Build prompt
    messages = self.prompt_builder.build(
        genes=genes, pathways=pathway_terms, additional_context=additional_context
    )

    # Store final prompt (user message content)
    user_msg = next((m["content"] for m in messages if m["role"] == "user"), "")
    result["final_prompt"] = user_msg

    # Generate annotation
    try:
        result["annotation"] = self.llm_provider.generate(messages)
    except Exception as e:
        result["annotation"] = f"Process: Failed\nError: {str(e)}"

    return result

Constructor

GeneSetAnnotator(
    llm_provider: BaseLLMProvider,
    enrichment_method: Optional[str] = "pathway",
    prompt_builder: Optional[PromptBuilder] = None
)

Parameters

Parameter Type Default Description
llm_provider BaseLLMProvider Required LLM provider instance
enrichment_method str or None "pathway" Enrichment method: "pathway", "custom", or None
prompt_builder PromptBuilder None Custom prompt builder

Example

from gs2txt import GeneSetAnnotator
from gs2txt.llm import OpenAIProvider
from gs2txt.prompts.builder import PromptBuilder

# Basic setup
annotator = GeneSetAnnotator(
    llm_provider=OpenAIProvider(api_key="...", model_id="gpt-4")
)

# With custom prompt
custom_builder = PromptBuilder(
    system_template="You are a cancer genomics expert...",
    user_template="Analyze these cancer genes: {genes}"
)
annotator = GeneSetAnnotator(
    llm_provider=provider,
    prompt_builder=custom_builder
)

# Without enrichment
annotator = GeneSetAnnotator(
    llm_provider=provider,
    enrichment_method=None
)

Methods

annotate()

Main method to generate annotation for a gene set.

def annotate(
    self,
    deg_df: pd.DataFrame,
    max_gene_num: int = 60,
    max_pathway_num: int = 10,
    pathways: Optional[List[str]] = None,
    compute_enrichment: bool = True,
    additional_context: Optional[str] = None,
    pvalue_threshold: float = 0.05,
    log2fc_threshold: float = 1.0
) -> str

Parameters

Parameter Type Default Description
deg_df pd.DataFrame Required DataFrame with 'gene' column
max_gene_num int 60 Maximum genes to include
max_pathway_num int 10 Maximum pathways to include
pathways List[str] None Pre-computed pathway list
compute_enrichment bool True Whether to run enrichment
additional_context str None Extra context (e.g., PPI)
pvalue_threshold float 0.05 P-value filter threshold
log2fc_threshold float 1.0 Log2FC filter threshold

Returns

str - LLM-generated annotation text

Example

result = annotator.annotate(
    deg_df,
    max_gene_num=100,
    pvalue_threshold=0.01,
    additional_context="PPI hub genes: TP53, MYC"
)

annotate_detailed()

Returns detailed results including intermediate data.

def annotate_detailed(
    self,
    deg_df: pd.DataFrame,
    **kwargs
) -> dict

Returns

{
    "annotation": str,      # LLM result
    "genes": List[str],     # Filtered genes
    "pathways": List[str],  # Enriched pathways
    "prompt": str           # Final prompt sent to LLM
}

Input DataFrame Format

The input DataFrame must have a gene column. Optional columns:

Column Description
gene Gene symbols (required)
pvalue P-values for filtering
logFC Log2 fold change for filtering

Example

import pandas as pd

deg_df = pd.DataFrame({
    "gene": ["TP53", "MYC", "BRCA1"],
    "pvalue": [0.001, 0.002, 0.003],
    "logFC": [2.3, 1.8, -1.5]
})

Error Handling

try:
    result = annotator.annotate(deg_df)
except ValueError as e:
    # Missing 'gene' column
    print(f"Invalid input: {e}")
except Exception as e:
    # LLM API error
    print(f"LLM error: {e}")

Thread Safety

GeneSetAnnotator instances are thread-safe for read operations. However, creating separate instances per thread is recommended for parallel processing.

from concurrent.futures import ThreadPoolExecutor

def process_geneset(name, df):
    annotator = GeneSetAnnotator(llm_provider=provider)
    return annotator.annotate(df)

with ThreadPoolExecutor(max_workers=4) as executor:
    results = executor.map(
        lambda x: process_geneset(x[0], x[1]),
        gene_sets.items()
    )