new: added support for reasoning models (closes #61)

Author: evilsocket

docs/concepts.md

@@ -23,6 +23,10 @@ using:
   - shell
 ```
 
+#### Enable Reasoning
+
+For models supporting reasoning, you can add a `reasoning` field to enable it, with a value that can either be `low`, `medium` or `high`.
+
 ### Prompt Interpolation and Variables
 Nerve supports [Jinja2](https://jinja.palletsprojects.com/) templating for dynamic prompt construction. You can:
 - Inject command line arguments (`{{ target }}`)

docs/index.md

@@ -70,6 +70,10 @@ Run it with:
 nerve run new-agent --url cnn.com
 ```
 
+#### Enable Reasoning
+
+For models supporting reasoning, you can add a `reasoning` field to enable it, with a value that can either be `low`, `medium` or `high`.
+
 ### Prompting & Variables
 Supports [Jinja2](https://jinja.palletsprojects.com/) templating. You can:
 - Include files: `{% include 'filename.md' %}`

nerve/generation/__init__.py

@@ -4,14 +4,16 @@
 import typing as t
 from abc import ABC, abstractmethod
 
+from litellm import ConfigDict
+from pydantic import BaseModel
 from loguru import logger
 
 from nerve.models import Usage
 from nerve.runtime import state
 from nerve.tools.protocol import get_tool_response, get_tool_schema
 
 
-class WindowStrategy(ABC):
+class WindowStrategy(ABC, BaseModel):
     @abstractmethod
     async def get_window(self, history: list[dict[str, t.Any]]) -> list[dict[str, t.Any]]:
         pass
@@ -20,23 +22,28 @@ async def get_window(self, history: list[dict[str, t.Any]]) -> list[dict[str, t.
     def __str__(self) -> str:
         pass
 
+class GenerationConfig(BaseModel):
+    generator_id: str
+    reasoning_effort: str | None = None
+    window_strategy: WindowStrategy
+    tools: list[t.Callable[..., t.Any]] | None = None
+
 
 class Engine(ABC):
     def __init__(
         self,
-        generator_id: str,
-        window_strategy: WindowStrategy,
-        tools: list[t.Callable[..., t.Any]] | None = None,
+        config: GenerationConfig,
     ):
-        self.generator_id = generator_id
+        self.config = config
+        self.generator_id = config.generator_id
         self.generator_params: dict[str, t.Any] = {}
 
         self._parse_generator_params()
 
         self.history: list[dict[str, t.Any]] = []
-        self.window_strategy = window_strategy
+        self.window_strategy = config.window_strategy
 
-        self.tools = {fn.__name__: fn for fn in (tools or [])}
+        self.tools = {fn.__name__: fn for fn in (config.tools or [])}
         self.tools_schemas = []
         for tool_name, tool_fn in self.tools.items():
             if not tool_fn.__doc__:

nerve/generation/litellm.py

@@ -7,7 +7,7 @@
 import litellm
 from loguru import logger
 
-from nerve.generation import Engine, WindowStrategy
+from nerve.generation import Engine, GenerationConfig, WindowStrategy
 from nerve.generation.conversation import SlidingWindowStrategy
 from nerve.generation.ollama import OllamaGlue
 from nerve.models import Usage
@@ -34,11 +34,9 @@ def _convert_to_serializable(obj: t.Any) -> t.Any:
 class LiteLLMEngine(Engine):
     def __init__(
         self,
-        generator_id: str,
-        window_strategy: WindowStrategy,
-        tools: list[t.Callable[..., t.Any]] | None = None,
+        config: GenerationConfig,
     ):
-        super().__init__(generator_id, window_strategy, tools)
+        super().__init__(config)
 
         # until this is not fixed, ollama needs special treatment: https://github.com/BerriAI/litellm/issues/6353
         self.is_ollama = "ollama" in self.generator_id
@@ -60,13 +58,20 @@ async def _litellm_generate(
             logger.debug(f"litellm.conversation: {json.dumps(conversation, indent=2)}")
 
             # litellm.set_verbose = True
+
+            #if self.config.reasoning_effort:
+            #    # if the model does not support reasoning, avoid raising litellm.UnsupportedParamsError
+            #    # by dropping the unsupported parameter
+            #    litellm.drop_params = True
+
             response = litellm.completion(
                 model=self.generator_id,
                 messages=conversation,
                 tools=tools_schema,
                 tool_choice="auto" if tools_schema else None,
                 verbose=False,
                 api_base=self.api_base,
+                reasoning_effort=self.config.reasoning_effort,
                 **self.generator_params,
             )
 

nerve/models.py

@@ -74,7 +74,7 @@ class Argument(BaseModel):
     tool: str | None = None
 
 
-def _check_required_version(required: str | None) -> str | None:
+def _validate_required_version(required: str | None) -> str | None:
     if required:
         from packaging.requirements import Requirement
 
@@ -98,6 +98,12 @@ def _check_required_version(required: str | None) -> str | None:
     return required
 
 
+def _validate_reasoning_effort(effort: str | None) -> str | None:
+    if effort not in [None, "low", "medium", "high"]:
+        raise ValueError(f"invalid reasoning effort: {effort}")
+    return effort
+
+
 class Configuration(BaseModel):
     """
     Configuration for an agent determining its "identity", task and capabilities.
@@ -127,13 +133,14 @@ class Limits(BaseModel):
 
     # legacy field used to detect if the user is loading a legacy file
     system_prompt: str | None = Field(default=None, exclude=True)
-
     # optional generator
     generator: str | None = None
+    # thinking effort for models supporting reasoning
+    reasoning: t.Annotated[str | None, AfterValidator(_validate_reasoning_effort)] = None
     # optional agent description
     description: str = ""
     # optional nerve version requirement
-    requires: t.Annotated[str | None, AfterValidator(_check_required_version)] = None
+    requires: t.Annotated[str | None, AfterValidator(_validate_required_version)] = None
     # used for versioning the agents
     version: str = "1.0.0"
     # the system prompt, the agent identity

nerve/runtime/agent.py

@@ -4,7 +4,7 @@
 from loguru import logger
 
 import nerve.runtime.state as state
-from nerve.generation import Engine, WindowStrategy
+from nerve.generation import Engine, GenerationConfig, WindowStrategy
 from nerve.generation.conversation import FullHistoryStrategy
 from nerve.generation.litellm import LiteLLMEngine
 from nerve.models import Configuration, Tool, Usage
@@ -77,10 +77,17 @@ async def create(
             configuration=configuration,
         )
 
+        engine_config = GenerationConfig(
+            generator_id=configuration.generator,
+            window_strategy=window_strategy,
+            tools=runtime.tools,
+            reasoning_effort=configuration.reasoning,
+        )
+
         return cls(
             runtime=runtime,
             configuration=configuration,
-            generation_engine=LiteLLMEngine(configuration.generator, window_strategy, runtime.tools),
+            generation_engine=LiteLLMEngine(engine_config),
             conv_window_strategy=window_strategy,
         )
 

nerve/runtime/logging.py

@@ -95,9 +95,14 @@ def log_event_to_terminal(event: Event) -> None:
             data["agent"] = DictWrapper(data["agent"])
 
         generator = data["agent"].runtime.generator
+        reasoning = data["agent"].configuration.reasoning
         name = data["agent"].runtime.name
         version = data["agent"].configuration.version
         tools = len(data["agent"].runtime.tools)
+
+        if reasoning:
+            generator = f"{generator} (reasoning={reasoning})"
+
         logger.info(f"🤖 {generator} | {name} v{version} with {tools} tools")
 
     elif event.name == "before_tool_called":