An attempt to fix tests (use init_graph_store where graph store is used)

docs/source/conf.py

@@ -51,4 +51,7 @@
 nbsphinx_allow_errors = True
 nbsphinx_execute = "never"
 
-github_url = "https://github.com/ShoggothAI/motleycrew"
+html_theme_options = {
+    "display_github": True,
+    "github_url": "https://github.com/ShoggothAI/motleycrew",
+}

tests/test_agents/test_crew.py

@@ -43,7 +43,9 @@ def tasks(self, request, crew, agent):
 
     def test_create_simple_task(self, crew, agent):
         assert len(crew.tasks) == 0
-        simple_task = crew.create_simple_task(description="task description", agent=agent)
+        simple_task = crew.create_simple_task(
+            description="task description", agent=agent
+        )
         assert isinstance(simple_task, SimpleTask)
         assert len(crew.tasks) == 1
         node = simple_task.node

tests/test_storage/test_kuzu_graph_store.py

@@ -35,7 +35,9 @@ def test_insert_new_node(self, database):
     def test_insert_node_and_retrieve(self, database):
         graph_store = MotleyKuzuGraphStore(database)
 
-        entity = Entity(int_param=1, optional_str_param="test", optional_list_str_param=["a", "b"])
+        entity = Entity(
+            int_param=1, optional_str_param="test", optional_list_str_param=["a", "b"]
+        )
         inserted_entity = graph_store.insert_node(entity)
         assert inserted_entity.id is not None
 
@@ -79,8 +81,12 @@ def test_create_relation(self, database):
         graph_store.create_relation(from_node=entity1, to_node=entity2, label="p")
 
         assert graph_store.check_relation_exists(from_node=entity1, to_node=entity2)
-        assert graph_store.check_relation_exists(from_node=entity1, to_node=entity2, label="p")
-        assert not graph_store.check_relation_exists(from_node=entity1, to_node=entity2, label="q")
+        assert graph_store.check_relation_exists(
+            from_node=entity1, to_node=entity2, label="p"
+        )
+        assert not graph_store.check_relation_exists(
+            from_node=entity1, to_node=entity2, label="q"
+        )
         assert not graph_store.check_relation_exists(from_node=entity2, to_node=entity1)
 
     def test_upsert_triplet(self, database):
@@ -93,8 +99,12 @@ def test_upsert_triplet(self, database):
         assert graph_store.check_node_exists(entity2)
 
         assert graph_store.check_relation_exists(from_node=entity1, to_node=entity2)
-        assert graph_store.check_relation_exists(from_node=entity1, to_node=entity2, label="p")
-        assert not graph_store.check_relation_exists(from_node=entity1, to_node=entity2, label="q")
+        assert graph_store.check_relation_exists(
+            from_node=entity1, to_node=entity2, label="p"
+        )
+        assert not graph_store.check_relation_exists(
+            from_node=entity1, to_node=entity2, label="q"
+        )
         assert not graph_store.check_relation_exists(from_node=entity2, to_node=entity1)
 
     def test_nodes_do_not_exist(self, database):
@@ -106,7 +116,9 @@ def test_nodes_do_not_exist(self, database):
         assert not graph_store.check_node_exists(entity2)
 
         assert not graph_store.check_relation_exists(from_node=entity1, to_node=entity2)
-        assert not graph_store.check_relation_exists(from_node=entity2, to_node=entity1, label="p")
+        assert not graph_store.check_relation_exists(
+            from_node=entity2, to_node=entity1, label="p"
+        )
 
     def test_relation_does_not_exist(self, database):
         graph_store = MotleyKuzuGraphStore(database)
@@ -141,25 +153,39 @@ def test_delete_entity_with_relations(self, database):
         graph_store.insert_node(entity1)
         graph_store.insert_node(entity2)
         graph_store.create_relation(from_node=entity1, to_node=entity2, label="p")
-        assert graph_store.check_relation_exists(from_node=entity1, to_node=entity2) is True
+        assert (
+            graph_store.check_relation_exists(from_node=entity1, to_node=entity2)
+            is True
+        )
 
         graph_store.delete_node(entity1)
         assert graph_store.check_node_exists(entity1) is False
         assert graph_store.check_node_exists(entity2) is True
-        assert graph_store.check_relation_exists(from_node=entity1, to_node=entity2) is False
+        assert (
+            graph_store.check_relation_exists(from_node=entity1, to_node=entity2)
+            is False
+        )
 
     def test_set_property(self, database):
         graph_store = MotleyKuzuGraphStore(database)
         entity = Entity(int_param=1)
         graph_store.insert_node(entity)
         assert entity.optional_str_param is None
-        assert graph_store.get_node_by_class_and_id(Entity, entity.id).optional_str_param is None
+        assert (
+            graph_store.get_node_by_class_and_id(Entity, entity.id).optional_str_param
+            is None
+        )
 
         entity.optional_str_param = "test"
-        assert graph_store.get_node_by_class_and_id(Entity, entity.id).optional_str_param == "test"
+        assert (
+            graph_store.get_node_by_class_and_id(Entity, entity.id).optional_str_param
+            == "test"
+        )
 
         entity.optional_list_str_param = ["a", "b"]
-        assert graph_store.get_node_by_class_and_id(Entity, entity.id).optional_list_str_param == [
+        assert graph_store.get_node_by_class_and_id(
+            Entity, entity.id
+        ).optional_list_str_param == [
             "a",
             "b",
         ]

tests/test_tasks/test_simple_task.py

@@ -1,23 +1,20 @@
 import pytest
 
 from langchain_community.tools import DuckDuckGoSearchRun
-import kuzu
 
 from motleycrew.crew import MotleyCrew
 from motleycrew.agents.langchain.openai_tools_react import ReactOpenAIToolsAgent
-from motleycrew.storage import MotleyKuzuGraphStore
+from motleycrew.storage.graph_store_utils import init_graph_store
 from motleycrew.tasks.simple import (
     SimpleTask,
     SimpleTaskUnit,
     compose_simple_task_prompt_with_dependencies,
 )
 
 
-@pytest.fixture
-def graph_store(tmpdir):
-    db_path = tmpdir / "test_db"
-    db = kuzu.Database(str(db_path))
-    graph_store = MotleyKuzuGraphStore(db)
+@pytest.fixture(scope="session")
+def graph_store():
+    graph_store = init_graph_store()
     return graph_store
 
 
@@ -64,7 +61,9 @@ def test_get_next_unit(self, tasks, crew):
         task1, task2 = tasks
         crew.add_dependency(task1, task2)
         assert task2.get_next_unit() is None
-        prompt = compose_simple_task_prompt_with_dependencies(task1.description, task1.get_units())
+        prompt = compose_simple_task_prompt_with_dependencies(
+            task1.description, task1.get_units()
+        )
         expected_unit = SimpleTaskUnit(
             name=task1.name,
             prompt=prompt,

tests/test_tasks/test_task.py

@@ -2,13 +2,12 @@
 
 import pytest
 
-import kuzu
 from langchain_core.runnables import Runnable
 
 from motleycrew import MotleyCrew
 from motleycrew.tools import MotleyTool
 from motleycrew.tasks import Task, TaskUnitType, TaskUnit
-from motleycrew.storage import MotleyKuzuGraphStore
+from motleycrew.storage.graph_store_utils import init_graph_store
 from motleycrew.common.exceptions import TaskDependencyCycleError
 
 
@@ -28,11 +27,9 @@ def create_dummy_task(crew: MotleyCrew, name: str):
     )
 
 
-@pytest.fixture
-def graph_store(tmpdir):
-    db_path = tmpdir / "test_db"
-    db = kuzu.Database(str(db_path))
-    graph_store = MotleyKuzuGraphStore(db)
+@pytest.fixture(scope="session")
+def graph_store():
+    graph_store = init_graph_store()
     return graph_store