update tool

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