Tool unit tests

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - LLM evaluation logic
 - Tool implementations without langchain or langgraph dependencies
+- Unit tests for the migrated tools
 
 ## [0.3.3] - 30.10.2024
 

swarm_copy_tests/conftest.py

@@ -0,0 +1,189 @@
+"""Test configuration."""
+
+import json
+from pathlib import Path
+
+import pytest
+import pytest_asyncio
+from fastapi.testclient import TestClient
+from httpx import AsyncClient
+from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
+from langchain_core.messages import AIMessage
+from sqlalchemy import MetaData
+from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
+
+from neuroagent.app.config import Settings
+from neuroagent.app.dependencies import get_kg_token, get_settings
+from neuroagent.app.main import app
+from neuroagent.tools import GetMorphoTool
+
+
+@pytest.fixture(name="app_client")
+def client_fixture():
+    """Get client and clear app dependency_overrides."""
+    app_client = TestClient(app)
+    test_settings = Settings(
+        tools={
+            "literature": {
+                "url": "fake_literature_url",
+            },
+        },
+        knowledge_graph={
+            "base_url": "https://fake_url/api/nexus/v1",
+        },
+        openai={
+            "token": "fake_token",
+        },
+        keycloak={
+            "username": "fake_username",
+            "password": "fake_password",
+        },
+    )
+    app.dependency_overrides[get_settings] = lambda: test_settings
+    # mock keycloak authentication
+    app.dependency_overrides[get_kg_token] = lambda: "fake_token"
+    yield app_client
+    app.dependency_overrides.clear()
+
+
+@pytest.fixture(autouse=True, scope="session")
+def dont_look_at_env_file():
+    """Never look inside of the .env when running unit tests."""
+    Settings.model_config["env_file"] = None
+
+
+@pytest.fixture()
+def patch_required_env(monkeypatch):
+    monkeypatch.setenv("NEUROAGENT_TOOLS__LITERATURE__URL", "https://fake_url")
+    monkeypatch.setenv(
+        "NEUROAGENT_KNOWLEDGE_GRAPH__BASE_URL", "https://fake_url/api/nexus/v1"
+    )
+    monkeypatch.setenv("NEUROAGENT_OPENAI__TOKEN", "dummy")
+    monkeypatch.setenv("NEUROAGENT_KEYCLOAK__VALIDATE_TOKEN", "False")
+    monkeypatch.setenv("NEUROAGENT_KEYCLOAK__PASSWORD", "password")
+
+
+@pytest_asyncio.fixture(params=["sqlite", "postgresql"], name="db_connection")
+async def setup_sql_db(request, tmp_path):
+    db_type = request.param
+
+    # To start the postgresql database:
+    # docker run -it --rm -p 5432:5432 -e POSTGRES_USER=test -e POSTGRES_PASSWORD=password postgres:latest
+    path = (
+        f"sqlite+aiosqlite:///{tmp_path / 'test_db.db'}"
+        if db_type == "sqlite"
+        else "postgresql+asyncpg://test:password@localhost:5432"
+    )
+    if db_type == "postgresql":
+        try:
+            async with create_async_engine(path).connect() as conn:
+                pass
+        except Exception:
+            pytest.skip("Postgres database not connected")
+    yield path
+    if db_type == "postgresql":
+        metadata = MetaData()
+        engine = create_async_engine(path)
+        session = AsyncSession(bind=engine)
+        async with engine.begin() as conn:
+            await conn.run_sync(metadata.reflect)
+            await conn.run_sync(metadata.drop_all)
+
+        await session.commit()
+        await engine.dispose()
+        await session.aclose()
+
+
+@pytest.fixture
+def get_resolve_query_output():
+    with open("tests/data/resolve_query.json") as f:
+        outputs = json.loads(f.read())
+    return outputs
+
+
+@pytest.fixture
+def brain_region_json_path():
+    br_path = Path(__file__).parent / "data" / "brainregion_hierarchy.json"
+    return br_path
+
+
+@pytest.fixture
+async def fake_llm_with_tools(brain_region_json_path):
+    class FakeFuntionChatModel(GenericFakeChatModel):
+        def bind_tools(self, functions: list):
+            return self
+
+        def bind_functions(self, **kwargs):
+            return self
+
+    # If you need another fake response to use different tools,
+    # you can do in your test
+    # ```python
+    # llm, _ = await anext(fake_llm_with_tools)
+    # llm.responses = my_fake_responses
+    # ```
+    # and simply bind the corresponding tools
+    fake_responses = [
+        AIMessage(
+            content="",
+            additional_kwargs={
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_zHhwfNLSvGGHXMoILdIYtDVI",
+                        "function": {
+                            "arguments": '{"brain_region_id":"http://api.brain-map.org/api/v2/data/Structure/549"}',
+                            "name": "get-morpho-tool",
+                        },
+                        "type": "function",
+                    }
+                ]
+            },
+            response_metadata={"finish_reason": "tool_calls"},
+            id="run-3828644d-197b-401b-8634-e6ecf01c2e7c-0",
+            tool_calls=[
+                {
+                    "name": "get-morpho-tool",
+                    "args": {
+                        "brain_region_id": (
+                            "http://api.brain-map.org/api/v2/data/Structure/549"
+                        )
+                    },
+                    "id": "call_zHhwfNLSvGGHXMoILdIYtDVI",
+                }
+            ],
+        ),
+        AIMessage(
+            content="Great answer",
+            response_metadata={"finish_reason": "stop"},
+            id="run-42768b30-044a-4263-8c5c-da61429aa9da-0",
+        ),
+    ]
+
+    # If you use this tool in your test, DO NOT FORGET to mock the url response with the following snippet:
+    #
+    # ```python
+    # json_path = Path(__file__).resolve().parent.parent / "data" / "knowledge_graph.json"
+    # with open(json_path) as f:
+    #     knowledge_graph_response = json.load(f)
+
+    # httpx_mock.add_response(
+    #     url="http://fake_url",
+    #     json=knowledge_graph_response,
+    # )
+    # ```
+    # The http call is not mocked here because one might want to change the responses
+    # and the tools used.
+    async_client = AsyncClient()
+    tool = GetMorphoTool(
+        metadata={
+            "url": "http://fake_url",
+            "search_size": 2,
+            "httpx_client": async_client,
+            "token": "fake_token",
+            "brainregion_path": brain_region_json_path,
+        }
+    )
+
+    yield FakeFuntionChatModel(messages=iter(fake_responses)), [tool], fake_responses
+    await async_client.aclose()

swarm_copy_tests/data/KG_brain_regions_hierarchy_test.json

@@ -0,0 +1,165 @@
+{
+    "@context": "https://neuroshapes.org",
+    "@id": "http://bbp.epfl.ch/neurosciencegraph/ontologies/core/brainregion",
+    "@type": "Ontology",
+    "versionInfo": "R103",
+    "hasHierarchyView": [
+        {
+            "@id": "https://bbp.epfl.ch/ontologies/core/bmo/BrainLayer",
+            "label": "Layer",
+            "description": "Layer based hierarchy",
+            "hasParentHierarchyProperty": "isLayerPartOf",
+            "hasChildrenHierarchyProperty": "hasLayerPart",
+            "hasLeafHierarchyProperty": "hasLayerLeafRegionPart"
+        },
+        {
+            "@id": "https://neuroshapes.org/BrainRegion",
+            "label": "BrainRegion",
+            "description": "Atlas default brain region hierarchy",
+            "hasParentHierarchyProperty": "isPartOf",
+            "hasChildrenHierarchyProperty": "hasPart",
+            "hasLeafHierarchyProperty": "hasLeafRegionPart"
+        }
+    ],
+    "atlasRelease": {
+        "@id": "https://bbp.epfl.ch/neurosciencegraph/data/4906ab85-694f-469d-962f-c0174e901885",
+        "@type": "BrainAtlasRelease",
+        "_rev": 8
+    },
+    "defines": [
+        {
+            "@id": "http://api.brain-map.org/api/v2/data/Structure/1",
+            "@type": "Class",
+            "atlas_id": 424,
+            "color_hex_triplet": "FF4C3E",
+            "graph_order": 775,
+            "hemisphere_id": 3,
+            "st_level": 9,
+            "identifier": "1",
+            "isPartOf": [
+                "http://api.brain-map.org/api/v2/data/Structure/2"
+            ],
+            "isDefinedBy": "http://bbp.epfl.ch/neurosciencegraph/ontologies/core/brainregion",
+            "subClassOf": [
+                "https://neuroshapes.org/BrainRegion"
+            ],
+            "delineates": [
+                "http://purl.obolibrary.org/obo/UBERON_0014594"
+            ],
+            "regionVolume": {
+                "unitCode": "cubic micrometer",
+                "value": 108296875.0
+            },
+            "regionVolumeRatioToWholeBrain": {
+                "unitCode": "cubic micrometer",
+                "value": 0.00021400307558019888
+            },
+            "representedInAnnotation": true,
+            "hasHierarchyView": [
+                "https://neuroshapes.org/BrainRegion"
+            ],
+            "prefLabel": "Tuberomammillary nucleus, ventral part",
+            "label": "Tuberomammillary nucleus, ventral part",
+            "notation": "TMv",
+            "altLabel": "TMv",
+            "atlasRelease": {
+                "@id": "https://bbp.epfl.ch/neurosciencegraph/data/4906ab85-694f-469d-962f-c0174e901885",
+                "@type": "BrainAtlasRelease",
+                "_rev": 8
+            }
+        },
+        {
+            "@id": "http://api.brain-map.org/api/v2/data/Structure/2",
+            "@type": "Class",
+            "atlas_id": 425,
+            "color_hex_triplet": "FF90FF",
+            "graph_order": 834,
+            "hemisphere_id": 3,
+            "st_level": 9,
+            "hasPart": [
+                "http://api.brain-map.org/api/v2/data/Structure/503",
+                "http://api.brain-map.org/api/v2/data/Structure/511",
+                "http://api.brain-map.org/api/v2/data/Structure/614454425",
+                "http://api.brain-map.org/api/v2/data/Structure/494"
+            ],
+            "identifier": "2",
+            "isDefinedBy": "http://bbp.epfl.ch/neurosciencegraph/ontologies/core/brainregion",
+            "subClassOf": [
+                "https://neuroshapes.org/BrainRegion"
+            ],
+            "regionVolume": {
+                "unitCode": "cubic micrometer",
+                "value": 2036828125.0
+            },
+            "regionVolumeRatioToWholeBrain": {
+                "unitCode": "cubic micrometer",
+                "value": 0.004024931311990765
+            },
+            "representedInAnnotation": true,
+            "hasLeafRegionPart": [
+                "http://api.brain-map.org/api/v2/data/Structure/494",
+                "http://api.brain-map.org/api/v2/data/Structure/614454425",
+                "http://api.brain-map.org/api/v2/data/Structure/511",
+                "http://api.brain-map.org/api/v2/data/Structure/503"
+            ],
+            "hasHierarchyView": [
+                "https://neuroshapes.org/BrainRegion"
+            ],
+            "prefLabel": "Superior colliculus, motor related, intermediate gray layer",
+            "label": "Superior colliculus, motor related, intermediate gray layer",
+            "notation": "SCig",
+            "altLabel": "SCig",
+            "atlasRelease": {
+                "@id": "https://bbp.epfl.ch/neurosciencegraph/data/4906ab85-694f-469d-962f-c0174e901885",
+                "@type": "BrainAtlasRelease",
+                "_rev": 8
+            }
+        },
+        {
+            "@id": "http://api.brain-map.org/api/v2/data/Structure/3",
+            "@type": "Class",
+            "atlas_id": 424,
+            "color_hex_triplet": "FF4C3E",
+            "graph_order": 775,
+            "hemisphere_id": 3,
+            "st_level": 9,
+            "identifier": "3",
+            "isPartOf": [
+                "http://api.brain-map.org/api/v2/data/Structure/2"
+            ],
+            "isDefinedBy": "http://bbp.epfl.ch/neurosciencegraph/ontologies/core/brainregion",
+            "subClassOf": [
+                "https://neuroshapes.org/BrainRegion"
+            ],
+            "delineates": [
+                "http://purl.obolibrary.org/obo/UBERON_0014594"
+            ],
+            "regionVolume": {
+                "unitCode": "cubic micrometer",
+                "value": 108296875.0
+            },
+            "regionVolumeRatioToWholeBrain": {
+                "unitCode": "cubic micrometer",
+                "value": 0.00021400307558019888
+            },
+            "representedInAnnotation": true,
+            "hasHierarchyView": [
+                "https://neuroshapes.org/BrainRegion"
+            ],
+            "prefLabel": "Tuberomammillary nucleus, ventral part",
+            "label": "Primary Motor Cortex",
+            "notation": "TMv",
+            "altLabel": "TMv",
+            "atlasRelease": {
+                "@id": "https://bbp.epfl.ch/neurosciencegraph/data/4906ab85-694f-469d-962f-c0174e901885",
+                "@type": "BrainAtlasRelease",
+                "_rev": 8
+            }
+        }
+    ],
+    "derivation": [
+        {},
+        {}
+    ],
+    "label": "Brain Region Ontology"
+}