🔧 Update import paths for langchain tools and improve test coverage f…

…or chain functionalities
shroominic · Apr 4, 2024 · bc22070 · bc22070
1 parent 50421d8
commit bc22070
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 23 deletions.
diff --git a/examples/use_additional_tools.py b/examples/use_additional_tools.py
@@ -8,9 +8,8 @@
 import io
 from typing import Any
 
-from langchain.tools import BaseTool
-
 from codeinterpreterapi import CodeInterpreterSession
+from langchain_core.tools import BaseTool
 
 
 class ExampleKnowledgeBaseTool(BaseTool):

diff --git a/tests/chain_test.py b/tests/chain_test.py
@@ -1,37 +1,61 @@
-from codeinterpreterapi.chains import remove_download_link, get_file_modifications
+import asyncio
+
+from codeinterpreterapi.chains import (
+    aget_file_modifications,
+    aremove_download_link,
+    get_file_modifications,
+    remove_download_link,
+)
+from langchain_openai import ChatOpenAI
+
+llm = ChatOpenAI(model="gpt-3.5-turbo")
+
+remove_download_link_example = (
+    "I have created the plot to your dataset.\n\n"
+    "Link to the file [here](sandbox:/plot.png)."
+)
+
+base_code = """
+    import matplotlib.pyplot as plt
+
+    x = list(range(1, 11))
+    y = [29, 39, 23, 32, 4, 43, 43, 23, 43, 77]
+
+    plt.plot(x, y, marker='o')
+    plt.xlabel('Index')
+    plt.ylabel('Value')
+    plt.title('Data Plot')
+    """
+code_with_mod = base_code + "\nplt.savefig('plot.png')"
+
+code_no_mod = base_code + "\nplt.show()"
 
 
 def test_remove_download_link() -> None:
-    example = (
-        "I have created the plot to your dataset.\n\n"
-        "Link to the file [here](sandbox:/plot.png)."
-    )
     assert (
-        remove_download_link(example).formatted_response.strip()
+        remove_download_link(remove_download_link_example, llm=llm).strip()
         == "I have created the plot to your dataset."
     )
 
 
-def test_get_file_modifications() -> None:
-    base_code = """
-        import matplotlib.pyplot as plt
+async def test_remove_download_link_async() -> None:
+    assert (
+        await aremove_download_link(remove_download_link_example, llm=llm)
+    ).strip() == "I have created the plot to your dataset."
 
-        x = list(range(1, 11))
-        y = [29, 39, 23, 32, 4, 43, 43, 23, 43, 77]
 
-        plt.plot(x, y, marker='o')
-        plt.xlabel('Index')
-        plt.ylabel('Value')
-        plt.title('Data Plot')
-        """
-    code_with_mod = base_code + "\nplt.savefig('plot.png')"
+def test_get_file_modifications() -> None:
+    assert get_file_modifications(code_with_mod, llm=llm) == ["plot.png"]
+    assert get_file_modifications(code_no_mod, llm=llm) == []
 
-    code_no_mod = base_code + "\nplt.show()"
 
-    assert get_file_modifications(code_with_mod).modifications == ["plot.png"]
-    assert get_file_modifications(code_no_mod).modifications == []
+async def test_get_file_modifications_async() -> None:
+    assert await aget_file_modifications(code_with_mod, llm=llm) == ["plot.png"]
+    assert await aget_file_modifications(code_no_mod, llm=llm) == []
 
 
 if __name__ == "__main__":
-    # test_remove_download_link()
+    test_remove_download_link()
+    asyncio.run(test_remove_download_link_async())
     test_get_file_modifications()
+    asyncio.run(test_get_file_modifications_async())