Merge branch 'main' of github.com:plasma-umass/ChatDBG

.github/workflows/tests.yml

@@ -45,7 +45,7 @@ jobs:
 
     - name: install test dependencies
       run: |
-        python3 -m pip install pytest pytest-asyncio pytest-forked hypothesis
+        python3 -m pip install pytest pytest-asyncio pytest-forked hypothesis pytest-mock
         python3 -m pip install -r test/requirements.txt
         python3 -m pip install .
 

src/chatdbg/pdb_util/locals.py

@@ -82,7 +82,7 @@ def _repr_if_defined(obj: Any) -> bool:
     return result
 
 
-def _format_limited(value: Any, limit: int = 10, depth: int = 3) -> str:
+def _format_limited(value: Union[int, np.ndarray], limit: int = 10, depth: int = 3) -> str:
     def format_tuple(t, depth):
         return tuple([helper(x, depth) for x in t])
 

src/chatdbg/util/config.py

@@ -17,8 +17,8 @@
 
 
 def _chatdbg_get_env(
-    option_name: str, default_value: Union[int, bool, str]
-) -> Union[int, bool, str]:
+    option_name: str, default_value: Union[bool, int, str]
+) -> Union[bool, int, str]:
     env_name = "CHATDBG_" + option_name.upper()
     v = os.getenv(env_name, str(default_value))
     if type(default_value) == int:

src/chatdbg/util/prompts.py

@@ -2,12 +2,11 @@
 import os
 from chatdbg.util.config import chatdbg_config
 from .text import truncate_proportionally
-from types import NoneType
 from typing import Any, Callable, List, Union, Optional
 
 
 def _wrap_it(
-    before: str, text: Optional[str], after: str = "", maxlen: int = 2048
+    before: str, text: str, after: str = "", maxlen: int = 2048
 ) -> str:
     if text:
         text = truncate_proportionally(text, maxlen, 0.5)

src/chatdbg/util/text.py

@@ -18,7 +18,7 @@ def strip_ansi(s: str) -> str:
 
 
 def truncate_proportionally(
-    text: str, maxlen: int = 32000, top_proportion: float = 0.5
+    text: str, maxlen: int = 32000, top_proportion: Union[float, int] = 0.5
 ) -> str:
     """Omit part of a string if needed to make it fit in a maximum length."""
     if len(text) > maxlen:

src/chatdbg/util/trim.py

@@ -62,10 +62,10 @@ def _chunkify(messages, model):
 
 
 def trim_messages(
-    messages: List[Dict],  # list of JSON objects encoded as dicts
+    messages: List[Dict[str, str]],  # list of JSON objects encoded as dicts
     model: str,
     trim_ratio: float = 0.75,
-) -> List[Dict]:
+) -> list:
     """
     Strategy:
     - chunk messages: