Added cell types and other tests

BlueBrain · Nov 20, 2024 · 2a4b0ad · 2a4b0ad
1 parent d5a2ccf
commit 2a4b0ad
Show file tree

Hide file tree

Showing 7 changed files with 821 additions and 0 deletions.
diff --git a/swarm_copy_tests/conftest.py b/swarm_copy_tests/conftest.py
@@ -0,0 +1,210 @@
+"""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="settings")
+def settings():
+    return 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",
+        },
+    )
+
+
+@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()
diff --git a/swarm_copy_tests/data/KG_brain_regions_hierarchy_test.json b/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"
+}
diff --git a/swarm_copy_tests/data/brainregion_hierarchy.json b/swarm_copy_tests/data/brainregion_hierarchy.json
diff --git a/swarm_copy_tests/data/kg_cell_types_hierarchy_test.json b/swarm_copy_tests/data/kg_cell_types_hierarchy_test.json
@@ -0,0 +1,13 @@
+{
+  "@context": "https://neuroshapes.org",
+  "@id": "http://bbp.epfl.ch/neurosciencegraph/ontologies/core/celltypes",
+  "@type": "Ontology",
+  "preferredNamespacePrefix": "celltypes",
+  "versionInfo": "R1660",
+  "defines": [
+    {"@id": "http://bbp.epfl.ch/neurosciencegraph/ontologies/etypes/cACint", "@type": "Class", "label": "cACint", "subClassOf": ["https://bbp.epfl.ch/ontologies/core/bmo/NeuronElectricalType", "https://neuroshapes.org/EType", "https://bbp.epfl.ch/ontologies/core/bmo/BrainCellType"], "definition": "Continuous accommodating interneuron electrical type", "color": "#108b8b", "prefLabel": "Continuous accommodating interneuron electrical type", "notation": "cACint", "atlasRelease": {"@id": "https://bbp.epfl.ch/neurosciencegraph/data/brainatlasrelease/c96c71a8-4c0d-4bc1-8a1a-141d9ed6693d", "@type": "BrainAtlasRelease", "_rev": 45}},
+    {"@id": "http://bbp.epfl.ch/neurosciencegraph/ontologies/mtypes/GCL_GC", "@type": "Class", "label": "GCL_GC", "subClassOf": ["https://bbp.epfl.ch/ontologies/core/mtypes/HippocampusMType", "https://bbp.epfl.ch/ontologies/core/bmo/NeuronMorphologicalType", "https://neuroshapes.org/MType", "https://bbp.epfl.ch/ontologies/core/bmo/BrainCellType"], "atlasRelease": {"@id": "https://bbp.epfl.ch/neurosciencegraph/data/brainatlasrelease/c96c71a8-4c0d-4bc1-8a1a-141d9ed6693d", "@type": "BrainAtlasRelease", "_rev": 45}},
+    {"@id": "http://bbp.epfl.ch/neurosciencegraph/ontologies/mtypes/L23_PTPC", "@type": "Class", "label": "L23_PTPC", "subClassOf": ["https://bbp.epfl.ch/ontologies/core/bmo/HumanNeocortexMType", "https://neuroshapes.org/PyramidalNeuron", "https://bbp.epfl.ch/ontologies/core/bmo/NeuronMorphologicalType", "https://neuroshapes.org/MType", "https://bbp.epfl.ch/ontologies/core/bmo/BrainCellType"], "notation": "L2_MC", "atlasRelease": {"@id": "https://bbp.epfl.ch/neurosciencegraph/data/brainatlasrelease/c96c71a8-4c0d-4bc1-8a1a-141d9ed6693d", "@type": "BrainAtlasRelease", "_rev": 45}}
+  ],
+  "label": "Cell Types Ontology"
+}