diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 369e6de..a20df73 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -34,7 +34,7 @@ jobs: - name: Set up environment run: | pip install --upgrade pip wheel setuptools - pip install bandit[toml]==1.7.4 ruff==0.5.5 + pip install bandit[toml]==1.7.4 ruff==0.6.7 - name: Linting check run: | bandit -qr -c pyproject.toml src/ diff --git a/.gitignore b/.gitignore index db7f261..cabbcce 100644 --- a/.gitignore +++ b/.gitignore @@ -152,3 +152,8 @@ cython_debug/ # static files generated from Django application using `collectstatic` media static + +# database stuff +*db +*.db-shm +*.db-wal \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index d71d1fc..2cffd59 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- Add get morphoelectric (me) model tool + ## [0.1.1] - 26.09.2024 ### Fixed diff --git a/src/neuroagent/app/config.py b/src/neuroagent/app/config.py index 9055def..1e4b01a 100644 --- a/src/neuroagent/app/config.py +++ b/src/neuroagent/app/config.py @@ -114,6 +114,14 @@ class SettingsGetMorpho(BaseModel): model_config = ConfigDict(frozen=True) +class SettingsGetMEModel(BaseModel): + """Get ME Model settings.""" + + search_size: int = 10 + + model_config = ConfigDict(frozen=True) + + class SettingsKnowledgeGraph(BaseModel): """Knowledge graph API settings.""" @@ -157,6 +165,7 @@ class SettingsTools(BaseModel): morpho: SettingsGetMorpho = SettingsGetMorpho() trace: SettingsTrace = SettingsTrace() kg_morpho_features: SettingsKGMorpho = SettingsKGMorpho() + me_model: SettingsGetMEModel = SettingsGetMEModel() model_config = ConfigDict(frozen=True) diff --git a/src/neuroagent/app/dependencies.py b/src/neuroagent/app/dependencies.py index 1294660..430491d 100644 --- a/src/neuroagent/app/dependencies.py +++ b/src/neuroagent/app/dependencies.py @@ -28,6 +28,7 @@ from neuroagent.multi_agents import BaseMultiAgent, SupervisorMultiAgent from neuroagent.tools import ( ElectrophysFeatureTool, + GetMEModelTool, GetMorphoTool, GetTracesTool, KGMorphoFeatureTool, @@ -304,6 +305,25 @@ def get_morphology_feature_tool( return tool +def get_me_model_tool( + settings: Annotated[Settings, Depends(get_settings)], + token: Annotated[str, Depends(get_kg_token)], + httpx_client: Annotated[AsyncClient, Depends(get_httpx_client)], +) -> GetMEModelTool: + """Load get ME model tool.""" + tool = GetMEModelTool( + metadata={ + "url": settings.knowledge_graph.url, + "token": token, + "httpx_client": httpx_client, + "search_size": settings.tools.me_model.search_size, + "brainregion_path": settings.knowledge_graph.br_saving_path, + "celltypes_path": settings.knowledge_graph.ct_saving_path, + } + ) + return tool + + def get_language_model( settings: Annotated[Settings, Depends(get_settings)], ) -> ChatOpenAI: @@ -369,6 +389,7 @@ def get_agent( ElectrophysFeatureTool, Depends(get_electrophys_feature_tool) ], traces_tool: Annotated[GetTracesTool, Depends(get_traces_tool)], + me_model_tool: Annotated[GetMEModelTool, Depends(get_me_model_tool)], settings: Annotated[Settings, Depends(get_settings)], ) -> BaseAgent | BaseMultiAgent: """Get the generative question answering service.""" @@ -397,6 +418,7 @@ def get_agent( kg_morpho_feature_tool, electrophys_feature_tool, traces_tool, + me_model_tool, ] logger.info("Load simple agent") return SimpleAgent(llm=llm, tools=tools) # type: ignore diff --git a/src/neuroagent/tools/__init__.py b/src/neuroagent/tools/__init__.py index 7ce3aee..af6761e 100644 --- a/src/neuroagent/tools/__init__.py +++ b/src/neuroagent/tools/__init__.py @@ -1,6 +1,7 @@ """Tools folder.""" from neuroagent.tools.electrophys_tool import ElectrophysFeatureTool, FeaturesOutput +from neuroagent.tools.get_me_model_tool import GetMEModelTool from neuroagent.tools.get_morpho_tool import GetMorphoTool, KnowledgeGraphOutput from neuroagent.tools.kg_morpho_features_tool import ( KGMorphoFeatureOutput, @@ -35,4 +36,5 @@ "ParagraphMetadata", "ResolveBrainRegionTool", "TracesOutput", + "GetMEModelTool", ] diff --git a/src/neuroagent/tools/get_me_model_tool.py b/src/neuroagent/tools/get_me_model_tool.py new file mode 100644 index 0000000..d82c185 --- /dev/null +++ b/src/neuroagent/tools/get_me_model_tool.py @@ -0,0 +1,266 @@ +"""Module defining the Get ME Model tool.""" + +import logging +from typing import Any, Literal, Optional, Type + +from langchain_core.tools import ToolException +from pydantic import BaseModel, Field + +from neuroagent.cell_types import get_celltypes_descendants +from neuroagent.tools.base_tool import BaseToolOutput, BasicTool +from neuroagent.utils import get_descendants_id + +logger = logging.getLogger(__name__) + + +class InputGetMEModel(BaseModel): + """Inputs of the knowledge graph API.""" + + brain_region_id: str = Field(description="ID of the brain region of interest.") + mtype_id: Optional[str] = Field( + default=None, description="ID of the M-type of interest." + ) + etype_id: Optional[ + Literal[ + "bAC", + "bIR", + "bNAC", + "bSTUT", + "cAC", + "cIR", + "cNAC", + "cSTUT", + "dNAC", + "dSTUT", + ] + ] = Field(default=None, description="ID of the E-type of interest.") + + +class MEModelOutput(BaseToolOutput): + """Output schema for the knowledge graph API.""" + + me_model_id: str + me_model_name: str | None + me_model_description: str | None + mtype: str | None + etype: str | None + + brain_region_id: str + brain_region_label: str | None + + subject_species_label: str | None + subject_age: str | None + + +class GetMEModelTool(BasicTool): + """Class defining the Get ME Model logic.""" + + name: str = "get-me-model-tool" + description: str = """Searches a neuroscience based knowledge graph to retrieve neuron morpho-electric model names, IDs and descriptions. + Requires a 'brain_region_id' which is the ID of the brain region of interest as registered in the knowledge graph. To get this ID, please use the `resolve-brain-region-tool` first. + Ideally, the user should also provide an 'mtype_id' and/or an 'etype_id' to filter the search results. But in case they are not provided, the search will return all models that match the brain region. + The output is a list of ME models, containing: + - The brain region ID. + - The brain region name. + - The subject species name. + - The subject age. + - The model ID. + - The model name. + - The model description. + The model ID is in the form of an HTTP(S) link such as 'https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/...'.""" + metadata: dict[str, Any] + args_schema: Type[BaseModel] = InputGetMEModel + + def _run(self) -> None: + pass + + async def _arun( + self, + brain_region_id: str, + mtype_id: str | None = None, + etype_id: str | None = None, + ) -> list[MEModelOutput]: + """From a brain region ID, extract ME models. + + Parameters + ---------- + brain_region_id + ID of the brain region of interest (of the form http://api.brain-map.org/api/v2/data/Structure/...) + mtype_id + ID of the mtype of the model + etype_id + ID of the etype of the model + + Returns + ------- + list of MEModelOutput to describe the model and its metadata, or an error dict. + """ + logger.info( + f"Entering Get ME Model tool. Inputs: {brain_region_id=}, {mtype_id=}, {etype_id=}" + ) + try: + # From the brain region ID, get the descendants. + hierarchy_ids = get_descendants_id( + brain_region_id, json_path=self.metadata["brainregion_path"] + ) + logger.info( + f"Found {len(list(hierarchy_ids))} children of the brain ontology." + ) + + if mtype_id: + mtype_ids = set( + get_celltypes_descendants(mtype_id, self.metadata["celltypes_path"]) + ) + logger.info( + f"Found {len(list(mtype_ids))} children of the cell types ontology for mtype." + ) + else: + mtype_ids = None + + if etype_id: + etype_ids = set( + get_celltypes_descendants(etype_id, self.metadata["celltypes_path"]) + ) + logger.info( + f"Found {len(list(etype_ids))} children of the cell types ontology for etype." + ) + else: + etype_ids = None + + # Create the ES query to query the KG. + entire_query = self.create_query( + brain_regions_ids=hierarchy_ids, + mtype_ids=mtype_ids, + etype_ids=etype_ids, + ) + + # Send the query to get ME models. + response = await self.metadata["httpx_client"].post( + url=self.metadata["url"], + headers={"Authorization": f"Bearer {self.metadata['token']}"}, + json=entire_query, + ) + + # Process the output and return. + return self._process_output(response.json()) + + except Exception as e: + raise ToolException(str(e), self.name) + + def create_query( + self, + brain_regions_ids: set[str], + mtype_ids: set[str] | None = None, + etype_ids: set[str] | None = None, + ) -> dict[str, Any]: + """Create ES query out of the BR, mtype, and etype IDs. + + Parameters + ---------- + brain_regions_ids + IDs of the brain region of interest (of the form http://api.brain-map.org/api/v2/data/Structure/...) + mtype_id + ID of the mtype of the model + etype_id + ID of the etype of the model + + Returns + ------- + dict containing the elasticsearch query to send to the KG. + """ + # At least one of the children brain region should match. + conditions = [ + { + "bool": { + "should": [ + {"term": {"brainRegion.@id.keyword": hierarchy_id}} + for hierarchy_id in brain_regions_ids + ] + } + }, + {"term": {"@type.keyword": "https://neuroshapes.org/MEModel"}}, + {"term": {"deprecated": False}}, + ] + + if mtype_ids: + # The correct mtype should match. For now + # It is a one term should condition, but eventually + # we will resolve the subclasses of the mtypes. + # They will all be appended here. + conditions.append( + { + "bool": { + "should": [ + {"match": {"mType.label": mtype_id}} + for mtype_id in mtype_ids + ] + } + } + ) + + if etype_ids: + # The correct etype should match. + conditions.append( + { + "bool": { + "should": [ + {"match": {"eType.label": etype_id}} + for etype_id in etype_ids + ] + } + } + ) + + # Assemble the query to return ME models. + entire_query = { + "size": self.metadata["search_size"], + "track_total_hits": True, + "query": {"bool": {"must": conditions}}, + "sort": {"createdAt": {"order": "desc", "unmapped_type": "keyword"}}, + } + return entire_query + + @staticmethod + def _process_output(output: Any) -> list[MEModelOutput]: + """Process output to fit the MEModelOutput pydantic class defined above. + + Parameters + ---------- + output + Raw output of the _arun method, which comes from the KG + + Returns + ------- + list of MEModelOutput to describe the model and its metadata. + """ + formatted_output = [ + MEModelOutput( + me_model_id=res["_source"]["@id"], + me_model_name=res["_source"].get("name"), + me_model_description=res["_source"].get("description"), + mtype=( + res["_source"]["mType"].get("label") + if "mType" in res["_source"] + else None + ), + etype=( + res["_source"]["eType"].get("label") + if "eType" in res["_source"] + else None + ), + brain_region_id=res["_source"]["brainRegion"]["@id"], + brain_region_label=res["_source"]["brainRegion"].get("label"), + subject_species_label=( + res["_source"]["subjectSpecies"].get("label") + if "subjectSpecies" in res["_source"] + else None + ), + subject_age=( + res["_source"]["subjectAge"].get("label") + if "subjectAge" in res["_source"] + else None + ), + ) + for res in output["hits"]["hits"] + ] + return formatted_output diff --git a/tests/agents/test_simple_agent.py b/tests/agents/test_simple_agent.py index a31cf91..0f0ca1d 100644 --- a/tests/agents/test_simple_agent.py +++ b/tests/agents/test_simple_agent.py @@ -4,6 +4,7 @@ from pathlib import Path import pytest + from neuroagent.agents import AgentOutput, AgentStep, SimpleAgent diff --git a/tests/agents/test_simple_chat_agent.py b/tests/agents/test_simple_chat_agent.py index c8420a3..72f2498 100644 --- a/tests/agents/test_simple_chat_agent.py +++ b/tests/agents/test_simple_chat_agent.py @@ -6,6 +6,7 @@ import pytest from langchain_core.messages import HumanMessage, ToolMessage from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver + from neuroagent.agents import AgentOutput, AgentStep, SimpleChatAgent diff --git a/tests/app/database/test_threads.py b/tests/app/database/test_threads.py index 7807826..dfc3313 100644 --- a/tests/app/database/test_threads.py +++ b/tests/app/database/test_threads.py @@ -1,13 +1,14 @@ """Test of the thread router.""" import pytest +from sqlalchemy import MetaData, create_engine +from sqlalchemy.orm import Session +from sqlalchemy.sql.expression import Select + from neuroagent.app.config import Settings from neuroagent.app.dependencies import get_language_model, get_settings from neuroagent.app.main import app from neuroagent.app.routers.database.schemas import GetThreadsOutput -from sqlalchemy import MetaData, create_engine -from sqlalchemy.orm import Session -from sqlalchemy.sql.expression import Select def test_create_thread(patch_required_env, app_client, db_connection): diff --git a/tests/app/database/test_tools.py b/tests/app/database/test_tools.py index a5f55e0..03f2cf0 100644 --- a/tests/app/database/test_tools.py +++ b/tests/app/database/test_tools.py @@ -1,6 +1,7 @@ """Test of the tool router.""" import pytest + from neuroagent.app.config import Settings from neuroagent.app.dependencies import get_language_model, get_settings from neuroagent.app.main import app diff --git a/tests/app/test_config.py b/tests/app/test_config.py index 459d457..38191f8 100644 --- a/tests/app/test_config.py +++ b/tests/app/test_config.py @@ -1,9 +1,10 @@ """Test config""" import pytest -from neuroagent.app.config import Settings from pydantic import ValidationError +from neuroagent.app.config import Settings + def test_required(monkeypatch, patch_required_env): settings = Settings() diff --git a/tests/app/test_dependencies.py b/tests/app/test_dependencies.py index 1acfa4b..cde38e3 100644 --- a/tests/app/test_dependencies.py +++ b/tests/app/test_dependencies.py @@ -11,6 +11,7 @@ from langchain_openai import ChatOpenAI from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver + from neuroagent.agents import SimpleAgent, SimpleChatAgent from neuroagent.app.dependencies import ( Settings, @@ -24,6 +25,7 @@ get_kg_morpho_feature_tool, get_language_model, get_literature_tool, + get_me_model_tool, get_morpho_tool, get_morphology_feature_tool, get_traces_tool, @@ -32,6 +34,7 @@ ) from neuroagent.tools import ( ElectrophysFeatureTool, + GetMEModelTool, GetMorphoTool, GetTracesTool, KGMorphoFeatureTool, @@ -116,6 +119,7 @@ def test_get_literature_tool(monkeypatch, patch_required_env): [get_traces_tool, True, "TRACE", GetTracesTool], [get_electrophys_feature_tool, False, None, ElectrophysFeatureTool], [get_morphology_feature_tool, False, None, MorphologyFeatureTool], + [get_me_model_tool, True, "ME_MODEL", GetMEModelTool], ), ) def test_get_tool( @@ -207,6 +211,9 @@ def test_get_agent(monkeypatch, patch_required_env): httpx_client=httpx_client, settings=settings, ) + me_model_tool = get_me_model_tool( + settings=settings, token=token, httpx_client=httpx_client + ) agent = get_agent( llm=language_model, @@ -218,6 +225,7 @@ def test_get_agent(monkeypatch, patch_required_env): electrophys_feature_tool=electrophys_feature_tool, traces_tool=traces_tool, settings=settings, + me_model_tool=me_model_tool, ) assert isinstance(agent, SimpleAgent) diff --git a/tests/app/test_main.py b/tests/app/test_main.py index 79fe9a8..3827ad8 100644 --- a/tests/app/test_main.py +++ b/tests/app/test_main.py @@ -2,6 +2,7 @@ from unittest.mock import patch from fastapi.testclient import TestClient + from neuroagent.app.dependencies import get_settings from neuroagent.app.main import app diff --git a/tests/app/test_middleware.py b/tests/app/test_middleware.py index 2707d4d..5a9bdcf 100644 --- a/tests/app/test_middleware.py +++ b/tests/app/test_middleware.py @@ -5,6 +5,7 @@ import pytest from fastapi.requests import Request from fastapi.responses import Response + from neuroagent.app.config import Settings from neuroagent.app.middleware import strip_path_prefix diff --git a/tests/conftest.py b/tests/conftest.py index d4c5f6d..c49e3a9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,13 +8,14 @@ from httpx import AsyncClient from langchain_core.language_models.fake_chat_models import GenericFakeChatModel from langchain_core.messages import AIMessage +from sqlalchemy import MetaData, create_engine +from sqlalchemy.exc import OperationalError +from sqlalchemy.orm import Session + 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 -from sqlalchemy import MetaData, create_engine -from sqlalchemy.exc import OperationalError -from sqlalchemy.orm import Session @pytest.fixture(name="app_client") diff --git a/tests/data/kg_me_model_output.json b/tests/data/kg_me_model_output.json new file mode 100644 index 0000000..6db0acd --- /dev/null +++ b/tests/data/kg_me_model_output.json @@ -0,0 +1,152 @@ +{ + "hits": { + "hits": [ + { + "sort": [ + 1716989796908 + ], + "_id": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/eeeeac3c-6bf1-47ed-ab97-460668eba2d2", + "_index": "nexus_search_711d6b8f-1285-42db-9259-b277dd687435_711d6b8f-1285-42db-9259-b277dd687435_21", + "_source": { + "@id": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/eeeeac3c-6bf1-47ed-ab97-460668eba2d2", + "@type": "https://neuroshapes.org/MEModel", + "analysisSuitable": false, + "brainRegion": { + "@id": "http://api.brain-map.org/api/v2/data/Structure/322", + "idLabel": "http://api.brain-map.org/api/v2/data/Structure/322|Primary somatosensory area", + "identifier": "http://api.brain-map.org/api/v2/data/Structure/322", + "label": "Primary somatosensory area" + }, + "createdAt": "2024-05-29T13:36:36.908Z", + "createdBy": "https://openbluebrain.com/api/nexus/v1/realms/bbp/users/antonel", + "deprecated": false, + "description": "My me-model", + "eType": { + "@id": "_:b0", + "identifier": "_:b0", + "label": "cNAC" + }, + "generation": {}, + "image": [ + { + "@id": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/2869d500-3c6e-4d15-bcae-b4ece55c5586", + "about": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/currentscape", + "identifier": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/2869d500-3c6e-4d15-bcae-b4ece55c5586" + }, + { + "@id": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/aea6c71b-41a9-4268-96e7-d0f9265693e1", + "about": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/parameters_distribution", + "identifier": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/aea6c71b-41a9-4268-96e7-d0f9265693e1" + }, + { + "@id": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/ea0b79b5-66bb-48bf-8933-9b43b56959d5", + "about": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/traces", + "identifier": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/ea0b79b5-66bb-48bf-8933-9b43b56959d5" + }, + { + "@id": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/ba3c8d1a-2839-4fdd-b13f-73950c84954a", + "about": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/currentscape", + "identifier": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/ba3c8d1a-2839-4fdd-b13f-73950c84954a" + }, + { + "@id": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/af301c80-0c51-4726-95b7-dc65e5e180fd", + "about": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/scores", + "identifier": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/af301c80-0c51-4726-95b7-dc65e5e180fd" + }, + { + "@id": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/798bab6a-4caa-460b-bac7-a6f08672f11b", + "about": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/currentscape", + "identifier": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/798bab6a-4caa-460b-bac7-a6f08672f11b" + }, + { + "@id": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/489673cc-27c1-46bf-ad6a-510ea124a4cc", + "about": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/currentscape", + "identifier": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/489673cc-27c1-46bf-ad6a-510ea124a4cc" + }, + { + "@id": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/3cb119a2-6a3a-4e75-99c5-53f2405d3e91", + "about": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/currentscape", + "identifier": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/3cb119a2-6a3a-4e75-99c5-53f2405d3e91" + }, + { + "@id": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/3819e25e-ace3-4809-9a96-964d5d9fb932", + "about": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/currentscape", + "identifier": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/3819e25e-ace3-4809-9a96-964d5d9fb932" + }, + { + "@id": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/33bcf9d3-b6db-4d2c-b3c6-3ea8a5932202", + "about": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/currentscape", + "identifier": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/33bcf9d3-b6db-4d2c-b3c6-3ea8a5932202" + }, + { + "@id": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/29dc72cb-bb9d-4966-bc3b-70e6ff65c6e3", + "about": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/currentscape", + "identifier": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/29dc72cb-bb9d-4966-bc3b-70e6ff65c6e3" + }, + { + "@id": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/e5e20e7a-87b5-47a2-abda-2d7212ab00e1", + "about": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/thumbnail", + "identifier": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/e5e20e7a-87b5-47a2-abda-2d7212ab00e1" + } + ], + "layer": [ + { + "@id": "http://purl.obolibrary.org/obo/UBERON_0005394", + "idLabel": "http://purl.obolibrary.org/obo/UBERON_0005394|layer 5", + "identifier": "http://purl.obolibrary.org/obo/UBERON_0005394", + "label": "layer 5" + } + ], + "mType": { + "@id": "_:b2", + "identifier": "_:b2", + "label": "L5_TPC:B" + }, + "memodel": { + "emodelResource": { + "@id": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/67acf101-12f6-4b7b-8a89-5237aadf94db", + "name": "EM__hipp_rat__CA1_int_cNAC_010710HP2_20190328163258__2" + }, + "neuronMorphology": { + "@id": "https://bbp.epfl.ch/data/bbp/mmb-point-neuron-framework-model/255e007e-a6d1-4fc9-b984-1e0221e39ea3", + "name": "ch150801A1" + }, + "validated": true + }, + "name": "My me-model", + "objectOfStudy": { + "@id": "http://bbp.epfl.ch/neurosciencegraph/taxonomies/objectsofstudy/singlecells", + "identifier": "http://bbp.epfl.ch/neurosciencegraph/taxonomies/objectsofstudy/singlecells", + "label": "Single Cell" + }, + "project": { + "@id": "https://openbluebrain.com/api/nexus/v1/projects/bbp/mmb-point-neuron-framework-model", + "identifier": "https://openbluebrain.com/api/nexus/v1/projects/bbp/mmb-point-neuron-framework-model", + "label": "bbp/mmb-point-neuron-framework-model" + }, + "seed": 2, + "subjectSpecies": { + "@id": "http://purl.obolibrary.org/obo/NCBITaxon_10116", + "identifier": "http://purl.obolibrary.org/obo/NCBITaxon_10116", + "label": "Rattus norvegicus" + }, + "updatedAt": "2024-07-09T10:53:01.208Z", + "updatedBy": "https://openbluebrain.com/api/nexus/v1/realms/bbp/users/ajaquier", + "_self": "https://openbluebrain.com/api/nexus/v1/resources/bbp/mmb-point-neuron-framework-model/_/https:%2F%2Fbbp.epfl.ch%2Fdata%2Fbbp%2Fmmb-point-neuron-framework-model%2Feeeeac3c-6bf1-47ed-ab97-460668eba2d2" + } + } + ], + "total": { + "relation": "eq", + "value": 1 + } + }, + "timed_out": false, + "took": 14, + "_shards": { + "failed": 0, + "skipped": 4, + "successful": 15, + "total": 15 + } +} \ No newline at end of file diff --git a/tests/test_cell_types.py b/tests/test_cell_types.py index 0596534..930ec71 100644 --- a/tests/test_cell_types.py +++ b/tests/test_cell_types.py @@ -3,6 +3,7 @@ from pathlib import Path import pytest + from neuroagent.cell_types import CellTypesMeta, get_celltypes_descendants CELL_TYPES_FILE = Path(__file__).parent / "data" / "kg_cell_types_hierarchy_test.json" diff --git a/tests/test_resolving.py b/tests/test_resolving.py index 4393b76..6d4f25a 100644 --- a/tests/test_resolving.py +++ b/tests/test_resolving.py @@ -1,5 +1,6 @@ import pytest from httpx import AsyncClient + from neuroagent.resolving import ( es_resolve, escape_punctuation, diff --git a/tests/test_utils.py b/tests/test_utils.py index 152232c..c43e33b 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -5,6 +5,7 @@ import pytest from httpx import AsyncClient + from neuroagent.schemas import KGMetadata from neuroagent.utils import ( RegionMeta, diff --git a/tests/tools/test_basic_tool.py b/tests/tools/test_basic_tool.py index fb9ab17..c29bc6c 100644 --- a/tests/tools/test_basic_tool.py +++ b/tests/tools/test_basic_tool.py @@ -4,9 +4,10 @@ from langchain_core.messages import AIMessage, HumanMessage from langchain_core.tools import ToolException from langgraph.prebuilt import create_react_agent -from neuroagent.tools.base_tool import BasicTool from pydantic import BaseModel +from neuroagent.tools.base_tool import BasicTool + class input_for_test(BaseModel): test_str: str diff --git a/tests/tools/test_electrophys_tool.py b/tests/tools/test_electrophys_tool.py index e7729fc..59e8763 100644 --- a/tests/tools/test_electrophys_tool.py +++ b/tests/tools/test_electrophys_tool.py @@ -6,6 +6,7 @@ import httpx import pytest from langchain_core.tools import ToolException + from neuroagent.tools import ElectrophysFeatureTool from neuroagent.tools.electrophys_tool import ( CALCULATED_FEATURES, diff --git a/tests/tools/test_get_me_model_tool.py b/tests/tools/test_get_me_model_tool.py new file mode 100644 index 0000000..2366b59 --- /dev/null +++ b/tests/tools/test_get_me_model_tool.py @@ -0,0 +1,159 @@ +"""Tests Get Morpho tool.""" + +import json +from pathlib import Path + +import httpx +import pytest +from langchain_core.tools import ToolException + +from neuroagent.tools.get_me_model_tool import GetMEModelTool, MEModelOutput + + +class TestGetMEModelTool: + @pytest.mark.asyncio + async def test_arun(self, httpx_mock, brain_region_json_path, tmp_path): + url = "http://fake_url" + json_path = ( + Path(__file__).resolve().parent.parent / "data" / "kg_me_model_output.json" + ) + with open(json_path) as f: + knowledge_graph_response = json.load(f) + + httpx_mock.add_response( + url=url, + json=knowledge_graph_response, + ) + tool = GetMEModelTool( + metadata={ + "url": url, + "search_size": 2, + "httpx_client": httpx.AsyncClient(), + "token": "fake_token", + "brainregion_path": brain_region_json_path, + "celltypes_path": tmp_path, + } + ) + response = await tool._arun( + brain_region_id="brain_region_id_link/549", + mtype_id="mtype_id_link/549", + etype_id="etype_id_link/549", + ) + assert isinstance(response, list) + assert len(response) == 1 + assert isinstance(response[0], MEModelOutput) + + @pytest.mark.asyncio + async def test_arun_errors(self, httpx_mock, brain_region_json_path, tmp_path): + url = "http://fake_url" + httpx_mock.add_response( + url=url, + json={}, + ) + + tool = GetMEModelTool( + metadata={ + "url": url, + "search_size": 2, + "httpx_client": httpx.AsyncClient(), + "token": "fake_token", + "brainregion_path": brain_region_json_path, + "celltypes_path": tmp_path, + } + ) + with pytest.raises(ToolException) as tool_exception: + _ = await tool._arun( + brain_region_id="brain_region_id_link/bad", + mtype_id="mtype_id_link/superbad", + etype_id="etype_id_link/superbad", + ) + + assert tool_exception.value.args[0] == "'hits'" + + +def test_create_query(): + url = "http://fake_url" + + tool = GetMEModelTool( + metadata={ + "url": url, + "search_size": 2, + "httpx_client": httpx.AsyncClient(), + "token": "fake_token", + } + ) + + # This should be a set, but passing a list here ensures that the test doesn;t rely on order. + brain_regions_ids = ["brain-region-id/68", "brain-region-id/131"] + mtype_id = "mtype-id/1234" + etype_id = "etype-id/1234" + + entire_query = tool.create_query( + brain_regions_ids=brain_regions_ids, mtype_ids={mtype_id}, etype_ids={etype_id} + ) + expected_query = { + "size": 2, + "track_total_hits": True, + "query": { + "bool": { + "must": [ + { + "bool": { + "should": [ + { + "term": { + "brainRegion.@id.keyword": "brain-region-id/68" + } + }, + { + "term": { + "brainRegion.@id.keyword": "brain-region-id/131" + } + }, + ] + } + }, + {"term": {"@type.keyword": ("https://neuroshapes.org/MEModel")}}, + {"term": {"deprecated": False}}, + {"bool": {"should": [{"match": {"mType.label": "mtype-id/1234"}}]}}, + {"bool": {"should": [{"match": {"eType.label": "etype-id/1234"}}]}}, + ] + } + }, + "sort": {"createdAt": {"order": "desc", "unmapped_type": "keyword"}}, + } + assert isinstance(entire_query, dict) + assert entire_query == expected_query + + # Case 2 with no mtype + entire_query1 = tool.create_query(brain_regions_ids=brain_regions_ids) + expected_query1 = { + "size": 2, + "track_total_hits": True, + "query": { + "bool": { + "must": [ + { + "bool": { + "should": [ + { + "term": { + "brainRegion.@id.keyword": "brain-region-id/68" + } + }, + { + "term": { + "brainRegion.@id.keyword": "brain-region-id/131" + } + }, + ] + } + }, + {"term": {"@type.keyword": ("https://neuroshapes.org/MEModel")}}, + {"term": {"deprecated": False}}, + ] + } + }, + "sort": {"createdAt": {"order": "desc", "unmapped_type": "keyword"}}, + } + assert entire_query1 == expected_query1 diff --git a/tests/tools/test_get_morpho_tool.py b/tests/tools/test_get_morpho_tool.py index d3ef4ab..87fdcbb 100644 --- a/tests/tools/test_get_morpho_tool.py +++ b/tests/tools/test_get_morpho_tool.py @@ -6,6 +6,7 @@ import httpx import pytest from langchain_core.tools import ToolException + from neuroagent.tools import GetMorphoTool from neuroagent.tools.get_morpho_tool import KnowledgeGraphOutput diff --git a/tests/tools/test_kg_morpho_features_tool.py b/tests/tools/test_kg_morpho_features_tool.py index 2f5473c..b346e26 100644 --- a/tests/tools/test_kg_morpho_features_tool.py +++ b/tests/tools/test_kg_morpho_features_tool.py @@ -6,6 +6,7 @@ import httpx import pytest from langchain_core.tools import ToolException + from neuroagent.tools import KGMorphoFeatureTool from neuroagent.tools.kg_morpho_features_tool import ( FeatRangeInput, diff --git a/tests/tools/test_literature_search_tool.py b/tests/tools/test_literature_search_tool.py index 0f58728..93adf4f 100644 --- a/tests/tools/test_literature_search_tool.py +++ b/tests/tools/test_literature_search_tool.py @@ -4,6 +4,7 @@ import httpx import pytest + from neuroagent.tools import LiteratureSearchTool from neuroagent.tools.literature_search_tool import ParagraphMetadata diff --git a/tests/tools/test_morphology_features_tool.py b/tests/tools/test_morphology_features_tool.py index 22943e6..92a311f 100644 --- a/tests/tools/test_morphology_features_tool.py +++ b/tests/tools/test_morphology_features_tool.py @@ -6,6 +6,7 @@ import httpx import pytest from langchain_core.tools import ToolException + from neuroagent.tools import MorphologyFeatureTool from neuroagent.tools.morphology_features_tool import MorphologyFeatureOutput diff --git a/tests/tools/test_resolve_br_tool.py b/tests/tools/test_resolve_br_tool.py index 697a6ef..31f893d 100644 --- a/tests/tools/test_resolve_br_tool.py +++ b/tests/tools/test_resolve_br_tool.py @@ -2,6 +2,7 @@ import pytest from httpx import AsyncClient + from neuroagent.tools import ResolveBrainRegionTool from neuroagent.tools.resolve_brain_region_tool import ( BRResolveOutput, diff --git a/tests/tools/test_traces_tool.py b/tests/tools/test_traces_tool.py index 0bae056..ad26bf4 100644 --- a/tests/tools/test_traces_tool.py +++ b/tests/tools/test_traces_tool.py @@ -6,6 +6,7 @@ import httpx import pytest from langchain_core.tools import ToolException + from neuroagent.tools import GetTracesTool from neuroagent.tools.traces_tool import TracesOutput