fix tests

Author: BoBer78

backend/src/neuroagent/app/config.py

@@ -99,14 +99,6 @@ class SettingsKGMorpho(BaseModel):
     model_config = ConfigDict(frozen=True)
 
 
-class SettingsGetMorpho(BaseModel):
-    """Get Morpho settings."""
-
-    search_size: int = 10
-
-    model_config = ConfigDict(frozen=True)
-
-
 class SettingsGetMEModel(BaseModel):
     """Get ME Model settings."""
 
@@ -163,7 +155,6 @@ class SettingsTools(BaseModel):
     obi_one: SettingsObiOne = SettingsObiOne()
     bluenaas: SettingsBlueNaaS = SettingsBlueNaaS()
     entitycore: SettingsEntityCore = SettingsEntityCore()
-    morpho: SettingsGetMorpho = SettingsGetMorpho()
     trace: SettingsTrace = SettingsTrace()
     kg_morpho_features: SettingsKGMorpho = SettingsKGMorpho()
     me_model: SettingsGetMEModel = SettingsGetMEModel()

backend/src/neuroagent/app/dependencies.py

@@ -35,7 +35,6 @@
     MorphoMetricsTool,
     PlotGeneratorTool,
     ResolveBrainRegionTool,
-    ResolveETypeTool,
     ResolveMtypeTool,
     SCSGetAllTool,
     SCSGetOneTool,
@@ -181,7 +180,6 @@ def get_tool_list() -> list[type[BaseTool]]:
         KGMorphoFeatureTool,
         ResolveBrainRegionTool,
         ResolveMtypeTool,
-        ResolveETypeTool,
         MorphoMetricsTool,
         GetTracesTool,
         PlotGeneratorTool,
@@ -320,15 +318,14 @@ def get_context_variables(
         "knowledge_graph_url": settings.knowledge_graph.url,
         "literature_search_url": settings.tools.literature.url,
         "me_model_search_size": settings.tools.me_model.search_size,
-        "morpho_search_size": settings.tools.morpho.search_size,
         "obi_one_url": settings.tools.obi_one.url,
         "openai_client": openai_client,
         "project_id": thread.project_id,
         "retriever_k": settings.tools.literature.retriever_k,
         "s3_client": s3_client,
         "starting_agent": starting_agent,
-        "thread_id": thread.thread_id,
         "tavily_api_key": settings.tools.web_search.tavily_api_key,
+        "thread_id": thread.thread_id,
         "token": token,
         "trace_search_size": settings.tools.trace.search_size,
         "use_reranker": settings.tools.literature.use_reranker,
@@ -348,12 +345,12 @@ def get_healthcheck_variables(
     correct service.
     """
     return {
+        "bluenaas_url": settings.tools.bluenaas.url.rstrip("/") + "/",
+        "entitycore_url": settings.tools.entitycore.url.rstrip("/") + "/",
         "httpx_client": httpx_client,
-        "literature_search_url": settings.tools.literature.url.rstrip("/") + "/",
         "knowledge_graph_url": settings.knowledge_graph.base_url.rstrip("/") + "/",
-        "bluenaas_url": settings.tools.bluenaas.url.rstrip("/") + "/",
+        "literature_search_url": settings.tools.literature.url.rstrip("/") + "/",
         "obi_one_url": settings.tools.obi_one.url.rstrip("/") + "/",
-        "entitycore_url": settings.tools.entitycore.url.rstrip("/") + "/",
     }
 
 

backend/src/neuroagent/tools/__init__.py

@@ -19,7 +19,6 @@
 from neuroagent.tools.morphology_viewer import MorphologyViewerTool
 from neuroagent.tools.now import NowTool
 from neuroagent.tools.resolve_brain_region_tool import ResolveBrainRegionTool
-from neuroagent.tools.resolve_etypes_tool import ResolveETypeTool
 from neuroagent.tools.resolve_mtypes_tool import ResolveMtypeTool
 from neuroagent.tools.traces_tool import GetTracesTool
 from neuroagent.tools.weather import WeatherTool
@@ -45,7 +44,6 @@
     "RandomPlotGeneratorTool",
     "ResolveBrainRegionTool",
     "ResolveMtypeTool",
-    "ResolveETypeTool",
     "WeatherTool",
     "WebSearchTool",
 ]

backend/src/neuroagent/tools/get_morpho_tool.py

@@ -1,11 +1,14 @@
 """Get Morpho tool."""
 
 import logging
-from typing import Any, ClassVar
+from typing import ClassVar
 
 from httpx import AsyncClient
 from pydantic import BaseModel, Field
 
+from neuroagent.tools.autogenerated_types.entitycore.models import (
+    ListResponseReconstructionMorphologyRead,
+)
 from neuroagent.tools.base_tool import BaseMetadata, BaseTool
 
 logger = logging.getLogger(__name__)
@@ -15,12 +18,13 @@ class GetMorphoInput(BaseModel):
     """Inputs of the knowledge graph API."""
 
     brain_region_id: str = Field(
-        description="ID of the brain region of interest in UUID format. To find the ID use the resolve-entity-tool first."
+        description="ID of the brain region of interest in UUID format. To find the ID use the resolve-brain-region-tool first."
     )
+    page_size: int = Field(default=10, description="Number of results per page.")
     page: int = Field(default=1, description="Page number for pagination.")
     mtype_id: str | None = Field(
         default=None,
-        description="ID of the M-type of interest. To find the ID use the resolve-entity-tool first.",
+        description="ID of the M-type of interest. To find the ID use the resolve-mtype-tool first.",
     )
 
 
@@ -29,7 +33,8 @@ class GetMorphoMetadata(BaseMetadata):
 
     entitycore_url: str
     token: str
-    morpho_search_size: int
+    vlab_id: str | None
+    project_id: str | None
 
 
 class MtypeOutput(BaseModel):
@@ -40,29 +45,6 @@ class MtypeOutput(BaseModel):
     alt_label: str | None
 
 
-class MorphologieOutput(BaseModel):
-    """Output schema for the knowledge graph API."""
-
-    morphology_id: str
-    morphology_name: str | None
-    morphology_description: str | None
-    mtype: list[MtypeOutput] | None
-
-    brain_region_id: str
-    brain_region_name: str | None
-
-    subject_species_name: str | None
-
-
-class GetMorphoToolOutput(BaseModel):
-    """Output schema for the Morpho tool."""
-
-    morphologies: list[MorphologieOutput]
-    current_page: int
-    page_size: int
-    total_items_found: int
-
-
 class GetMorphoTool(BaseTool):
     """Class defining the Get Morpho logic."""
 
@@ -92,7 +74,7 @@ class GetMorphoTool(BaseTool):
     metadata: GetMorphoMetadata
     input_schema: GetMorphoInput
 
-    async def arun(self) -> GetMorphoToolOutput:
+    async def arun(self) -> ListResponseReconstructionMorphologyRead:
         """From a brain region ID, extract morphologies.
 
         Returns
@@ -105,9 +87,22 @@ async def arun(self) -> GetMorphoToolOutput:
 
         response = await self.metadata.httpx_client.get(
             url=self.metadata.entitycore_url + "/reconstruction-morphology",
-            headers={"Authorization": f"Bearer {self.metadata.token}"},
+            headers={
+                "Authorization": f"Bearer {self.metadata.token}",
+                **(
+                    {"virtual-lab-id": self.metadata.vlab_id}
+                    if self.metadata.vlab_id is not None
+                    else {}
+                ),
+                **(
+                    {"project-id": self.metadata.project_id}
+                    if self.metadata.project_id is not None
+                    else {}
+                ),
+            },
             params={
-                "page_size": self.metadata.morpho_search_size,
+                "page_size": self.input_schema.page_size,
+                "page": self.input_schema.page,
                 "within_brain_region_hierachy_id": "e3e70682-c209-4cac-a29f-6fbed82c07cd",  # TEMP for mouse brain
                 "within_brain_region_brain_region_id": self.input_schema.brain_region_id,
                 "within_brain_region_ascendants": False,
@@ -118,48 +113,11 @@ async def arun(self) -> GetMorphoToolOutput:
                 ),
             },
         )
-
-        return self._process_output(response.json())
-
-    @staticmethod
-    def _process_output(output: Any) -> GetMorphoToolOutput:
-        """Process output to fit the KnowledgeGraphOutput pydantic class defined above.
-
-        Parameters
-        ----------
-        output
-            Raw output of the arun method, which comes from the KG
-
-        Returns
-        -------
-            list of KGMorphoFeatureOutput to describe the morphology and its metadata.
-        """
-        formatted_output = [
-            MorphologieOutput(
-                morphology_id=res["id"],
-                morphology_name=res.get("name"),
-                morphology_description=res.get("description"),
-                mtype=[
-                    MtypeOutput(
-                        mtype_id=mtype["id"],
-                        pref_label=mtype.get("pref_label"),
-                        alt_label=mtype.get("alt_label"),
-                    )
-                    for mtype in res.get("mtypes")
-                ],
-                brain_region_id=res["brain_region"]["id"],
-                brain_region_name=res["brain_region"].get("Name"),
-                subject_species_name=res["species"].get("name"),
+        if response.status_code != 200:
+            raise ValueError(
+                f"The morphology endpoint returned a non 200 response code. Error: {response.text}"
             )
-            for res in output["data"]
-        ]
-        pagination_data = output["pagination"]
-        return GetMorphoToolOutput(
-            morphologies=formatted_output,
-            current_page=pagination_data["page"],
-            page_size=pagination_data["page_size"],
-            total_items_found=pagination_data["total_items"],
-        )
+        return ListResponseReconstructionMorphologyRead(**response.json())
 
     @classmethod
     async def is_online(cls, *, httpx_client: AsyncClient, entitycore_url: str) -> bool:

backend/src/neuroagent/tools/resolve_brain_region_tool.py

@@ -68,6 +68,7 @@ async def arun(
             url=self.metadata.entitycore_url + "/brain-region",
             headers={"Authorization": f"Bearer {self.metadata.token}"},
             params={
+                "hierarchy_id": "e3e70682-c209-4cac-a29f-6fbed82c07cd",
                 "page_size": 500,
                 "name__ilike": self.input_schema.brain_region,
             },

backend/src/neuroagent/tools/resolve_etypes_tool.py

@@ -69,7 +69,7 @@ async def arun(
             url=self.metadata.entitycore_url + "/mtype",
             headers={"Authorization": f"Bearer {self.metadata.token}"},
             params={
-                "page_size": 100,
+                "page_size": 1,
                 "pref_label": self.input_schema.mtype,
             },
         )

backend/src/neuroagent/tools/resolve_mtypes_tool.py

@@ -29,11 +29,9 @@ def test_no_settings():
 
 def test_setup_tools(monkeypatch, patch_required_env):
     monkeypatch.setenv("NEUROAGENT_TOOLS__TRACE__SEARCH_SIZE", "20")
-    monkeypatch.setenv("NEUROAGENT_TOOLS__MORPHO__SEARCH_SIZE", "20")
     monkeypatch.setenv("NEUROAGENT_TOOLS__KG_MORPHO_FEATURES__SEARCH_SIZE", "20")
 
     settings = Settings()
 
-    assert settings.tools.morpho.search_size == 20
     assert settings.tools.trace.search_size == 20
     assert settings.tools.kg_morpho_features.search_size == 20