Remove unnecessary poetry.lock file and fix tests in OpenAI SDK cookb…

…ook code (#46)

* Fix tests

Signed-off-by: Sid Murching <[email protected]>

* Update reqs

Signed-off-by: Sid Murching <[email protected]>

* remove dbconnect dep

Signed-off-by: Sid Murching <[email protected]>

* Skip some tests in GHA

Signed-off-by: Sid Murching <[email protected]>

---------

Signed-off-by: Sid Murching <[email protected]>

.github/workflows/unit-test.yml

@@ -24,6 +24,11 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install -r dev/dev_requirements.txt
-      - name: Test with pytest
+      - name: Test agent app sample code
         run: |
+          cd agent_app_sample_code 
           pytest
+      - name: Test OpenAI SDK sample code
+        run: |
+          cd openai_sdk_agent_app_sample_code 
+          pytest          

.gitignore

@@ -14,4 +14,5 @@ mlruns/
 
 _scratch_pad/
 openai_sdk_agent_app_sample_code/_scratch_pad/
-.vscode/
+.vscode/
+*poetry.lock

agent_app_sample_code/tests/conftest.py

@@ -1,6 +1,6 @@
 import sys
 import os
 
-# Add the root directory to sys.path, so that we can treat directories like
-# agent_app_sample_code as modules
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+# Add the parent directory to sys.path, so that we can treat directories like
+# `utils` as modules
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

agent_app_sample_code/tests/test_file_loading.py

@@ -5,14 +5,13 @@
 import pandas as pd
 from typing import TypedDict
 
-from agent_app_sample_code.utils.file_loading import load_files_to_df, apply_parsing_udf
-from agent_app_sample_code.utils.typed_dicts_to_spark_schema import typed_dicts_to_spark_schema
+from utils.file_loading import load_files_to_df, apply_parsing_udf
+from utils.typed_dicts_to_spark_schema import typed_dicts_to_spark_schema
 
 @pytest.fixture(scope="module")
 def spark():
     return (
-        pyspark.sql.SparkSession.builder
-        .master("local[1]")
+        pyspark.sql.SparkSession.builder.master("local[1]")
         .config("spark.driver.bindAddress", "127.0.0.1")
         .config("spark.task.maxFailures", "1")  # avoid retry failed spark tasks
         .getOrCreate()

dev/dev_requirements.txt

@@ -2,7 +2,14 @@ jupyter-book
 livereload
 black
 pytest
-mlflow-skinny[databricks]
+mlflow-skinny[databricks]>=2.18.0
+pydantic>=2.9.2
+databricks-agents
+databricks-sdk[openai]
 databricks-vectorsearch
-pyspark
+pyyaml
+git+https://github.com/unitycatalog/unitycatalog.git#subdirectory=ai/core
+git+https://github.com/unitycatalog/unitycatalog.git#subdirectory=ai/integrations/openai
+tabulate
 pandas
+pyspark

openai_sdk_agent_app_sample_code/requirements.txt

@@ -8,4 +8,5 @@ git+https://github.com/unitycatalog/unitycatalog.git#subdirectory=ai/core
 git+https://github.com/unitycatalog/unitycatalog.git#subdirectory=ai/integrations/openai
 tabulate
 pandas
-databricks-connect==15.1.0
+pyspark
+databricks-connect==15.1.0

openai_sdk_agent_app_sample_code/tests/conftest.py

@@ -1,6 +1,6 @@
 import sys
 import os
 
-# Add the root directory to sys.path, so that we can treat directories like
-# openai_sdk_agent_app_sample_code as modules
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+# Add the parent directory to sys.path, so that we can treat directories like
+# `cookbook` as modules
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

openai_sdk_agent_app_sample_code/tests/test_file_loading.py → openai_sdk_agent_app_sample_code/tests/test_data_pipeline_utils.py

@@ -5,8 +5,8 @@
 import pandas as pd
 from typing import TypedDict
 
-from utils.file_loading import load_files_to_df, apply_parsing_udf
-from utils.typed_dicts_to_spark_schema import typed_dicts_to_spark_schema
+from cookbook.data_pipeline.parse_docs import load_files_to_df, apply_parsing_fn
+from cookbook.data_pipeline.utils.typed_dicts_to_spark_schema import typed_dicts_to_spark_schema
 
 
 @pytest.fixture(scope="module")
@@ -73,7 +73,7 @@ class ParserReturnValue(TypedDict):
     doc_uri: str  # do not change this name
 
 
-def test_apply_parsing_udf(spark, example_files_dir):
+def test_apply_parsing_fn(spark, example_files_dir):
     def _mock_file_parser(
         raw_doc_contents_bytes: bytes,
         doc_path: str,
@@ -88,7 +88,7 @@ def _mock_file_parser(
 
     temp_dir, file_1, file_2 = example_files_dir
     raw_files_df = load_files_to_df(spark, str(temp_dir)).orderBy("path")
-    parsed_df = apply_parsing_udf(
+    parsed_df = apply_parsing_fn(
         raw_files_df,
         _mock_file_parser,
         parsed_df_schema=typed_dicts_to_spark_schema(ParserReturnValue),

openai_sdk_agent_app_sample_code/tools/test_code_exec_as_uc_tool.py

@@ -1,10 +1,15 @@
 import pytest
 from cookbook.tools.uc_tool import UCTool
+import os
 
 CATALOG = "ep"
 SCHEMA = "cookbook_local_test"
 
 
+if os.getenv("GITHUB_ACTIONS") == "true":
+    pytest.skip("Skipping all tests in this module in CI, as "
+                "Databricks auth is not available there", allow_module_level=True)
+
 @pytest.fixture
 def python_exec():
     """Fixture to provide the python_exec function from UCTool."""

openai_sdk_agent_app_sample_code/tools/test_sample_tool_uc.py

@@ -1,5 +1,10 @@
 import pytest
 from cookbook.tools.uc_tool import UCTool
+import os
+
+if os.getenv("GITHUB_ACTIONS") == "true":
+    pytest.skip("Skipping all tests in this module in CI, as "
+                "Databricks auth is not available there", allow_module_level=True)
 
 # Load the function from the UCTool versus locally
 @pytest.fixture