diff --git a/src/neuroagent/tools/bluenaas_tool.py b/src/neuroagent/tools/bluenaas_tool.py index a9ecf82..b055e1b 100644 --- a/src/neuroagent/tools/bluenaas_tool.py +++ b/src/neuroagent/tools/bluenaas_tool.py @@ -1,6 +1,7 @@ """BlueNaaS single cell stimulation, simulation and synapse placement tool.""" -from typing import List, Literal, Optional -from pydantic import BaseModel, Field +from typing import List, Literal, Optional, Annotated +from pydantic import BaseModel, Field, PositiveInt +from annotated_types import Len import logging from typing import Any, Type @@ -10,42 +11,86 @@ logger = logging.getLogger(__name__) -class SynapseConfig(BaseModel): +class SynapseSimulationConfig(BaseModel): id: str delay: int duration: int - frequency: int - weightScalar: float + frequency: PositiveInt + weightScalar: int -class StimulusConfig(BaseModel): - stimulusType: Literal["current_clamp", "voltage_clamp"] - stimulusProtocol: Literal["ap_waveform", "idrest", "iv", "fire_pattern", "pos_cheops", "neg_cheops"] - amplitudes: List[float] +class SimulationStimulusConfig(BaseModel): + stimulusType: Literal["current_clamp", "voltage_clamp", "conductance"] + stimulusProtocol: Literal["ap_waveform", "idrest", "iv", "fire_pattern"] + amplitudes: Annotated[list[float], Len(min_length=1, max_length=15)] class CurrentInjectionConfig(BaseModel): injectTo: str - stimulus: StimulusConfig + stimulus: SimulationStimulusConfig -class RecordFromConfig(BaseModel): +class RecordingLocation(BaseModel): section: str offset: float -class ConditionsConfig(BaseModel): +class SimulationConditionsConfig(BaseModel): celsius: float vinit: float hypamp: float - max_time: int + max_time: float time_step: float seed: int +class SimulationWithSynapseBody(BaseModel): + directCurrentConfig: CurrentInjectionConfig + synapseConfigs: list[SynapseSimulationConfig] + + class InputBlueNaaS(BaseModel): - model_id: str - synapses: List[SynapseConfig] - currentInjection: CurrentInjectionConfig - recordFrom: List[RecordFromConfig] - conditions: ConditionsConfig - type: Literal["single-neuron-simulation"] - simulationDuration: int + """Inputs for the BlueNaaS single-neuron simulation.""" + + model_id: str = Field( + description=( + "ID of the neuron model to be used in the simulation. The model ID can be" + " fetched using the 'get-me-model-tool'." + ) + ) + synapses: list[SynapseSimulationConfig] = Field( + description=( + "List of synapse configurations. Each synapse configuration includes the" + " synapse ID, delay, duration, frequency, and weight scalar." + ) + ) + currentInjection: CurrentInjectionConfig = Field( + description=( + "Configuration for current injection. Includes the target section to inject" + " to and the stimulus configuration." + ) + ) + recordFrom: list[RecordingLocation] = Field( + description=( + "List of sections to record from during the simulation. Each record" + " configuration includes the section name and offset." + ) + ) + conditions: SimulationConditionsConfig = Field( + description=( + "Simulation conditions including temperature (celsius), initial voltage" + " (vinit, in mV), hyperpolarizing current (hypamp, in nA), maximum simulation time" + " (max_time, in ms), time step (time_step, in ms), and random seed (seed)." + ) + ) + #TODO: implement synaptome simulation + simulationType: Literal["single-neuron-simulation","synaptome-simulation"] = Field( + description=( + "Type of the simulation. it can be single neuron simulation for simulation" + " without synapse placement or synaptome-simulation to put synapses on the morphology." + ) + ) + simulationDuration: int = Field( + description=( + "Duration of the simulation in milliseconds." + ) + ) + class BlueNaaSOutput(BaseModel): status: str @@ -54,11 +99,11 @@ class BlueNaaSOutput(BaseModel): class BlueNaaSTool(BasicTool): - name: str = "blue-naas-tool" + name: str = "bluenaas-tool" description: str = """Runs a single-neuron simulation using the BlueNaaS service. Requires a 'model_id' which can be fetched using the 'get-me-model-tool'. - The input configuration should be provided by the user otherwise agent will probe the user - with the selected default values.""" + The input configuration should be provided by the user otherwise agent + will probe the user with the selected default values.""" metadata: dict[str, Any] args_schema: Type[BaseModel] = InputBlueNaaS