refactor: update package layout

tablegpt · Nov 6, 2024 · 5bb1810 · 5bb1810
1 parent 0253209
commit 5bb1810
Show file tree

Hide file tree

Showing 9 changed files with 9 additions and 12 deletions.
diff --git a/src/tablegpt/agent/data_analyzer.py b/src/tablegpt/agent/data_analyzer.py
@@ -16,8 +16,8 @@
 from langgraph.prebuilt import ToolNode
 
 from tablegpt.agent.output_parser import MarkdownOutputParser
-from tablegpt.chains import create_hazard_classifier
 from tablegpt.retriever import format_columns
+from tablegpt.safety import create_hazard_classifier
 from tablegpt.tools import IPythonTool, markdown_console_template, process_content
 from tablegpt.utils import filter_contents
 

diff --git a/src/tablegpt/agent/file_reading.py → src/tablegpt/agent/file_reading/__init__.py b/src/tablegpt/agent/file_reading.py → src/tablegpt/agent/file_reading/__init__.py
@@ -11,14 +11,14 @@
 from langgraph.prebuilt import ToolNode
 from pybox.base import BasePyBoxManager  # noqa: TCH002
 
-from tablegpt.chains.data_normalizer import (
+from tablegpt.agent.file_reading.data_normalizer import (
     get_data_normalize_chain,
     get_table_reformat_chain,
     wrap_normalize_code,
 )
-from tablegpt.chains.translation import get_translation_chain
 from tablegpt.errors import NoAttachmentsError
 from tablegpt.tools import IPythonTool, markdown_console_template
+from tablegpt.translation import create_translator
 from tablegpt.utils import get_raw_table_info
 
 if TYPE_CHECKING:
@@ -83,7 +83,7 @@ def create_file_reading_workflow(
 
     translation_chain = None
     if locale is not None:
-        translation_chain = get_translation_chain(llm=llm)
+        translation_chain = create_translator(llm=llm)
 
     tools = [IPythonTool(pybox_manager=pybox_manager, cwd=workdir, session_id=session_id)]
     tool_executor = ToolNode(tools)

diff --git a/src/tablegpt/chains/data_normalizer.py → ...gpt/agent/file_reading/data_normalizer.py b/src/tablegpt/chains/data_normalizer.py → ...gpt/agent/file_reading/data_normalizer.py
diff --git a/src/tablegpt/chains/__init__.py b/src/tablegpt/chains/__init__.py
diff --git a/src/tablegpt/chains/safety.py → src/tablegpt/safety.py b/src/tablegpt/chains/safety.py → src/tablegpt/safety.py
diff --git a/src/tablegpt/chains/translation.py → src/tablegpt/translation.py b/src/tablegpt/chains/translation.py → src/tablegpt/translation.py
@@ -23,6 +23,6 @@
 output_parser = StrOutputParser()
 
 
-def get_translation_chain(llm: BaseLanguageModel) -> Runnable:
+def create_translator(llm: BaseLanguageModel) -> Runnable:
     """return the guard chain runnable."""
     return translation_prompt_template | llm | output_parser
diff --git a/tests/chains/__init__.py → tests/agent/file_reading/__init__.py b/tests/chains/__init__.py → tests/agent/file_reading/__init__.py
diff --git a/tests/chains/test_data_normalizer.py → ...gent/file_reading/test_data_normalizer.py b/tests/chains/test_data_normalizer.py → ...gent/file_reading/test_data_normalizer.py
@@ -1,7 +1,7 @@
 import unittest
 
 from langchain_core.exceptions import OutputParserException
-from tablegpt.chains.data_normalizer import (
+from tablegpt.agent.file_reading.data_normalizer import (
     CodeOutputParser,
     ListListOutputParser,
     ListTupleOutputParser,

diff --git a/tests/chains/test_safety.py → tests/test_safety.py b/tests/chains/test_safety.py → tests/test_safety.py
@@ -1,7 +1,7 @@
 import unittest
 from unittest import mock
 
-from tablegpt.chains.safety import HazardOutputParser
+from tablegpt.safety import HazardOutputParser
 
 
 class TestHazardOutputParser(unittest.TestCase):
@@ -21,7 +21,7 @@ def test_parse_unknown(self):
         assert result == ("unknown", None)
 
     @mock.patch(
-        "tablegpt.chains.safety.hazard_categories",
+        "tablegpt.safety.hazard_categories",
         {
             "hazard1": "category1",
         },
@@ -31,7 +31,7 @@ def test_parse_unsafe_with_known_category(self):
         assert result == ("unsafe", "category1")
 
     @mock.patch(
-        "tablegpt.chains.safety.hazard_categories",
+        "tablegpt.safety.hazard_categories",
         {
             "hazard1": "category1",
         },