Upload 6 files

entities/pipeline_request.py

@@ -0,0 +1,50 @@
+from pydantic import BaseModel
+
+
+class PipelineRequest(BaseModel):
+    customer_name: str
+    task_type: str
+    
+    
+class PubMedPipelineRequest(PipelineRequest):
+    query: str
+    direction: str
+    start_year: int
+    end_year: int
+    size: int
+    translate_to_cn: bool = False
+    task_type: str = "pubmed"
+    model_names: list[str] = ["gpt-4o-mini"]
+    do_compare: bool = False
+
+
+class PubMedPlusPipelineRequest(PipelineRequest):
+    query: str
+    direction: str
+    start_year: int
+    end_year: int
+    size: int
+    translate_to_cn: bool = False
+    task_type: str = "pubmed_plus"
+    model_names: list[str] = ["gpt-4o-mini"]
+    do_compare: bool = False
+    do_refine: bool = False
+    
+    
+class ArxivPipelineRequest(PipelineRequest):
+    query: str
+    direction: str
+    start_date: str
+    end_date: str
+    size: int
+    translate_to_cn: bool = False
+    task_type: str = "arxiv"
+    model_names: list[str] = ["gpt-4o-mini"]
+    do_compare: bool = False
+    
+
+pipeline_request_factory = {
+    "pubmed": PubMedPipelineRequest,
+    "pubmed_plus": PubMedPlusPipelineRequest,
+    "arxiv": ArxivPipelineRequest
+}

entities/task.py

@@ -0,0 +1,149 @@
+import json
+
+from typing import Dict
+from loguru import logger
+
+
+# 定义任务
+class Task:
+    
+    @classmethod
+    def load_from_json(cls, json_data):
+        task = cls(**json_data)
+        return task
+    
+    def __init__(self):
+        pass
+    
+    def save_to_json(self):
+        return json.dumps(vars(self))
+
+
+class PubMedTask(Task):
+    
+    def __init__(
+        self,
+        query: str = "",
+        direction: str = "",
+        start_year: int = 2018,
+        end_year: int = 2023,
+        size: int = 10,
+        uuid: int = 0,
+        customer_name: str = "test",
+        status: dict[str:int] = {},
+        status_string: dict[str:str] = {"overall": "in queue"},
+        translate_to_cn: bool = False,
+        model_names: list[str] = ["gpt-4o-mini"],
+        task_type: str = "pubmed",
+        do_compare: bool = False
+    ):
+        # basic query information
+        self.query = query
+        self.direction = direction
+        self.start_year = start_year
+        self.end_year = end_year
+        self.size = size
+        self.translate_to_cn = translate_to_cn
+        self.model_names = model_names
+        self.task_type = task_type
+        self.do_compare = do_compare
+        
+        # task information
+        self.uuid = uuid
+        self.customer_name = customer_name
+        
+        # task status
+        self.status = status
+        self.status_string = status_string
+        
+
+class PubMedPlusTask(Task):
+    
+    def __init__(
+        self,
+        query: str = "",
+        direction: str = "",
+        start_year: int = 2018,
+        end_year: int = 2023,
+        size: int = 10,
+        uuid: int = 0,
+        customer_name: str = "test",
+        status: dict[str:int] = {},
+        status_string: dict[str:str] = {"overall": "in queue"},
+        translate_to_cn: bool = False,
+        model_names: list[str] = ["gpt-4o-mini"],
+        task_type: str = "pubmed",
+        do_compare: bool = False,
+        do_refine: bool = False
+    ):
+        # basic query information
+        self.query = query
+        self.direction = direction
+        self.start_year = start_year
+        self.end_year = end_year
+        self.size = size
+        self.translate_to_cn = translate_to_cn
+        self.model_names = model_names
+        self.task_type = task_type
+        self.do_compare = do_compare
+        self.do_refine = do_refine
+        
+        # task information
+        self.uuid = uuid
+        self.customer_name = customer_name
+        
+        # task status
+        self.status = status
+        self.status_string = status_string   
+                
+          
+class ArxivTask(Task):
+    
+    def __init__(
+        self,
+        query: str = "",
+        direction: str = "",
+        start_date: str = "2018-01-01",
+        end_date: str = "2023-12-31",
+        size: int = 10,
+        uuid: int = 0,
+        customer_name: str = "test",
+        status: dict[str:int] = {},
+        status_string: dict[str:str] = {"overall": "in queue"},
+        translate_to_cn: bool = False,
+        model_names: list[str] = ["gpt-4o-mini"],
+        task_type: str = "arxiv",
+        do_compare: bool = False
+    ):
+        # basic query information
+        self.query = query
+        if direction:
+            self.direction = direction
+        else:
+            self.direction = query
+        self.start_date = start_date
+        self.end_date = end_date
+        self.size = size
+        self.translate_to_cn = translate_to_cn
+        self.model_names = model_names
+        self.task_type = task_type
+        self.do_compare = do_compare
+
+        # task information
+        self.uuid = uuid
+        self.customer_name = customer_name
+        
+        # task information
+        self.uuid = uuid
+        self.customer_name = customer_name
+        
+        # task status
+        self.status = status
+        self.status_string = status_string
+        
+        
+task_factory = {
+    "pubmed": PubMedTask,
+    "pubmed_plus": PubMedPlusTask,
+    "arxiv": ArxivTask
+}