support database class testing

examples/minimal_demo.py

@@ -1,4 +1,4 @@
-from typing import Dict, List
+from beartype.typing import Dict, List
 
 from research_town.dbs import (
     AgentProfile,

research_town/agents/agent_base.py

@@ -1,7 +1,7 @@
 from datetime import datetime
-from typing import Any, Dict, List, Tuple
 
 from beartype import beartype
+from beartype.typing import Any, Dict, List, Tuple
 
 from ..dbs import (
     AgentAgentDiscussionLog,

research_town/dbs/agent_db.py

@@ -1,7 +1,7 @@
 import json
 import uuid
-from typing import Any, Dict, List, Optional
 
+from beartype.typing import Any, Dict, List, Optional
 from pydantic import BaseModel, Field
 
 

research_town/dbs/env_db.py

@@ -1,7 +1,7 @@
 import json
 import uuid
-from typing import Any, Dict, List, Optional, Type, TypeVar
 
+from beartype.typing import Any, Dict, List, Optional, Type, TypeVar
 from pydantic import BaseModel, Field
 
 T = TypeVar('T', bound=BaseModel)

research_town/dbs/paper_db.py

@@ -1,7 +1,7 @@
 import json
 import uuid
-from typing import Any, Dict, List, Optional
 
+from beartype.typing import Any, Dict, List, Optional
 from pydantic import BaseModel, Field
 
 from ..utils.paper_collector import get_daily_papers
@@ -33,7 +33,7 @@ def __init__(self) -> None:
     def add_paper(self, paper: PaperProfile) -> None:
         self.data[paper.pk] = paper
 
-    def update_paper(self, paper_pk: str, updates: Dict[str, Optional[str]]) -> bool:
+    def update_paper(self, paper_pk: str, updates: Dict[str, Any]) -> bool:
         if paper_pk in self.data:
             for key, value in updates.items():
                 if value is not None:

research_town/dbs/progress_db.py

@@ -1,7 +1,7 @@
 import json
 import uuid
-from typing import Any, Dict, List, Optional, Type, TypeVar
 
+from beartype.typing import Any, Dict, List, Optional, Type, TypeVar
 from pydantic import BaseModel, Field
 
 T = TypeVar('T', bound=BaseModel)

research_town/envs/env_base.py

@@ -1,4 +1,4 @@
-from typing import List
+from beartype.typing import List
 
 from ..agents.agent_base import BaseResearchAgent
 from ..dbs import AgentProfile, EnvLogDB

research_town/envs/env_paper_rebuttal.py

@@ -1,6 +1,5 @@
-from typing import Dict, List, Tuple
-
 from beartype import beartype
+from beartype.typing import Dict, List, Tuple
 
 from ..dbs import (
     AgentPaperMetaReviewLog,

research_town/envs/env_paper_submission.py

@@ -1,6 +1,5 @@
-from typing import Dict, List
-
 from beartype import beartype
+from beartype.typing import Dict, List
 
 from ..agents.agent_base import BaseResearchAgent
 from ..dbs import (

research_town/evaluators/output_format.py

@@ -1,6 +1,5 @@
 
-from typing import Type, TypeVar
-
+from beartype.typing import Type, TypeVar
 from pydantic import BaseModel, Extra, Field, validator
 
 T = TypeVar('T', bound=BaseModel)

research_town/evaluators/quality_evaluator.py

@@ -1,6 +1,7 @@
 
 import re
-from typing import Any
+
+from beartype.typing import Any
 
 from ..utils.decorator import parsing_error_exponential_backoff
 from ..utils.eval_prompter import (

research_town/utils/agent_collector.py

@@ -1,6 +1,5 @@
-from typing import Any, Dict, List, Tuple
-
 from arxiv import Client, Search
+from beartype.typing import Any, Dict, List, Tuple
 from tqdm import tqdm
 
 

research_town/utils/agent_prompter.py

@@ -1,6 +1,5 @@
-from typing import Dict, List, Union
-
 from beartype import beartype
+from beartype.typing import Dict, List, Union
 
 from .model_prompting import model_prompting
 from .paper_collector import get_related_papers

research_town/utils/decorator.py

@@ -1,8 +1,15 @@
 import math
 import time
 from functools import wraps
-from typing import Any, Callable, List, Optional, TypeVar, cast
 
+from beartype.typing import (
+    Any,
+    Callable,
+    List,
+    Optional,
+    TypeVar,
+    cast,
+)
 from pydantic import BaseModel
 
 INF = float(math.inf)

research_town/utils/eval_prompter.py

@@ -1,6 +1,5 @@
-from typing import Dict
-
 from beartype import beartype
+from beartype.typing import Dict
 
 from .model_prompting import model_prompting
 

research_town/utils/logging.py

@@ -1,5 +1,6 @@
 import logging
-from typing import Dict, List, Union
+
+from beartype.typing import Dict, List, Union
 
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 def logging_callback(messages: Union[List[Dict[str, str]], None] = None) -> None:

research_town/utils/model_prompting.py

@@ -1,7 +1,6 @@
-from typing import List, Optional
-
 import litellm
 from beartype import beartype
+from beartype.typing import List, Optional
 
 from .decorator import api_calling_error_exponential_backoff
 

research_town/utils/paper_collector.py

@@ -1,11 +1,11 @@
 import datetime
-from typing import Any, Dict, List, Tuple
 from xml.etree import ElementTree
 
 import arxiv
 import faiss
 import requests
 import torch
+from beartype.typing import Any, Dict, List, Tuple
 from transformers import BertModel, BertTokenizer
 
 ATOM_NAMESPACE = "{http://www.w3.org/2005/Atom}"

research_town/utils/serializer.py

@@ -1,6 +1,6 @@
 import importlib
-from typing import Any, Dict, List, Set, Tuple, Union
 
+from beartype.typing import Any, Dict, List, Set, Tuple, Union
 from pydantic import BaseModel
 
 

research_town/utils/string_mapper.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Union
+from beartype.typing import Dict, List, Union
 
 
 def map_idea_list_to_str(ideas: List[Dict[str, str]]) -> str:

research_town/utils/tools.py

@@ -1,7 +1,8 @@
 import datetime
 import json
 import os
-from typing import Any, Dict
+
+from beartype.typing import Any, Dict
 
 
 def show_time() -> str: