Switch to side-by-side view

--- a
+++ b/aitrika/base/aitrika_base.py
@@ -0,0 +1,131 @@
+from typing import Optional, Dict, List, Union
+import pandas as pd
+from aitrika.extractors.pubmed_extractor import PubMedExtractor
+from aitrika.extractors.llm_extractor import LLMExtractor
+from aitrika.utils.loader import loader
+from aitrika.llm.base_llm import BaseLLM
+
+
+class AItrikaBase:
+    def __init__(self, pubmed_id: Optional[str] = None):
+        self.pubmed_id = pubmed_id
+        self.pubmed_extractor = PubMedExtractor(pubmed_id) if pubmed_id else None
+        self.llm_extractor = LLMExtractor()
+
+    @loader(text="Extracting PubMed ID")
+    def extract_pubmed_id(self) -> str:
+        if not self.pubmed_extractor:
+            raise ValueError("PubMed ID not set")
+        return self.pubmed_extractor.extract_pubmed_id()
+
+    @loader(text="Extracting title")
+    def extract_title(self) -> str:
+        if not self.pubmed_extractor:
+            raise ValueError("PubMed ID not set")
+        return self.pubmed_extractor.extract_title()
+
+    @loader(text="Extracting abstract")
+    def extract_abstract(self) -> str:
+        if not self.pubmed_extractor:
+            raise ValueError("PubMed ID not set")
+        return self.pubmed_extractor.extract_abstract()
+
+    @loader(text="Extracting other abstract")
+    def extract_other_abstract(self) -> str:
+        if not self.pubmed_extractor:
+            raise ValueError("PubMed ID not set")
+        return self.pubmed_extractor.extract_other_abstract()
+
+    @loader(text="Extracting genes")
+    def extract_genes(self, dataframe: bool = False) -> Union[pd.DataFrame, str]:
+        if not self.pubmed_extractor:
+            raise ValueError("PubMed ID not set")
+        return self.pubmed_extractor.extract_genes(dataframe)
+
+    @loader(text="Extracting diseases")
+    def extract_diseases(self, dataframe: bool = False) -> Union[pd.DataFrame, str]:
+        if not self.pubmed_extractor:
+            raise ValueError("PubMed ID not set")
+        return self.pubmed_extractor.extract_diseases(dataframe)
+
+    @loader(text="Extracting species")
+    def extract_species(self, dataframe: bool = False) -> Union[pd.DataFrame, str]:
+        if not self.pubmed_extractor:
+            raise ValueError("PubMed ID not set")
+        return self.pubmed_extractor.extract_species(dataframe)
+
+    @loader(text="Extracting chemicals")
+    def extract_chemicals(self, dataframe: bool = False) -> Union[pd.DataFrame, str]:
+        if not self.pubmed_extractor:
+            raise ValueError("PubMed ID not set")
+        return self.pubmed_extractor.extract_chemicals(dataframe)
+
+    @loader(text="Extracting mutations")
+    def extract_mutations(self, dataframe: bool = False) -> Union[pd.DataFrame, str]:
+        if not self.pubmed_extractor:
+            raise ValueError("PubMed ID not set")
+        return self.pubmed_extractor.extract_mutations(dataframe)
+
+    @loader(text="Extracting associations between genes and diseases")
+    def extract_associations(
+        self, dataframe: bool = False
+    ) -> Union[pd.DataFrame, List[Dict]]:
+        if not self.pubmed_extractor:
+            raise ValueError("PubMed ID not set")
+        return self.pubmed_extractor.extract_associations(dataframe)
+
+    @loader(text="Extracting authors")
+    def extract_authors(self) -> str:
+        if not self.pubmed_extractor:
+            raise ValueError("PubMed ID not set")
+        return self.pubmed_extractor.extract_authors()
+
+    @loader(text="Extracting journal")
+    def extract_journal(self) -> str:
+        if not self.pubmed_extractor:
+            raise ValueError("PubMed ID not set")
+        return self.pubmed_extractor.extract_journal()
+
+    @loader(text="Extracting results")
+    def extract_results(self, llm: BaseLLM) -> str:
+        return self.llm_extractor.extract_results(llm)
+
+    @loader(text="Extracting bibliography")
+    def extract_bibliography(self, llm: BaseLLM) -> str:
+        return self.llm_extractor.extract_bibliography(llm)
+
+    @loader(text="Extracting methods")
+    def extract_methods(self, llm: BaseLLM) -> str:
+        return self.llm_extractor.extract_methods(llm)
+
+    @loader(text="Extracting introduction")
+    def extract_introduction(self, llm: BaseLLM) -> str:
+        return self.llm_extractor.extract_introduction(llm)
+
+    @loader(text="Extracting acknowledgements")
+    def extract_acknowledgements(self, llm: BaseLLM) -> str:
+        return self.llm_extractor.extract_acknowledgements(llm)
+
+    @loader(text="Extracting paper results")
+    def extract_paper_results(self, llm: BaseLLM) -> str:
+        return self.llm_extractor.extract_paper_results(llm)
+
+    @loader(text="Extracting effect sizes")
+    def extract_effect_sizes(self, llm: BaseLLM) -> str:
+        return self.llm_extractor.extract_effect_sizes(llm)
+
+    @loader(text="Extracting number of participants")
+    def extract_number_of_participants(self, llm: BaseLLM) -> str:
+        return self.llm_extractor.extract_number_of_participants(llm)
+
+    @loader(text="Extracting characteristics of participants")
+    def extract_characteristics_of_participants(self, llm: BaseLLM) -> str:
+        return self.llm_extractor.extract_characteristics_of_participants(llm)
+
+    @loader(text="Extracting interventions")
+    def extract_interventions(self, llm: BaseLLM) -> str:
+        return self.llm_extractor.extract_interventions(llm)
+
+    @loader(text="Extracting outcomes")
+    def extract_outcomes(self, llm: BaseLLM) -> str:
+        return self.llm_extractor.extract_outcomes(llm)