Adding mypy

Makefile

@@ -13,6 +13,7 @@ lint:
 	poetry run ruff format --preview src/ tests/
 	poetry run black --preview --enable-unstable-feature=string_processing src/ tests/
 	poetry run ruff check --fix src/ tests/
+	poetry run mypy src/avalan
 
 test:
 	poetry sync --extras test

pyproject.toml

@@ -276,6 +276,74 @@ version = "1.3.10"
 [tool.poetry.group.dev.dependencies]
 ruff = "^0.11.11"
 black = "^25.1.0"
+mypy = "^1.14.0"
+
+[tool.mypy]
+python_version = "3.11"
+files = ["src/avalan"]
+plugins = ["pydantic.mypy"]
+strict = false
+warn_return_any = true
+warn_unused_ignores = false
+warn_redundant_casts = true
+warn_unused_configs = true
+show_error_codes = true
+show_column_numbers = true
+pretty = true
+check_untyped_defs = true
+disallow_untyped_defs = false
+disallow_incomplete_defs = false
+disallow_untyped_decorators = false
+
+[[tool.mypy.overrides]]
+module = [
+    "a2a.*",
+    "anthropic.*",
+    "bitsandbytes.*",
+    "boto3.*",
+    "botocore.*",
+    "diffusers.*",
+    "elasticsearch.*",
+    "faiss.*",
+    "google.*",
+    "huggingface_hub.*",
+    "imageio.*",
+    "jinja2.*",
+    "keyring.*",
+    "litellm.*",
+    "markitdown.*",
+    "markdownify.*",
+    "mcp.*",
+    "mlx.*",
+    "mlx_lm.*",
+    "numpy.*",
+    "openai.*",
+    "pandas.*",
+    "pgvector.*",
+    "PIL.*",
+    "playwright.*",
+    "psycopg.*",
+    "psycopg_pool.*",
+    "pydantic.*",
+    "RestrictedPython.*",
+    "rich.*",
+    "soundfile.*",
+    "sqlalchemy.*",
+    "sqlglot.*",
+    "sympy.*",
+    "tiktoken.*",
+    "torch.*",
+    "torchaudio.*",
+    "torchvision.*",
+    "tqdm.*",
+    "transformers.*",
+    "tree_sitter.*",
+    "tree_sitter_python.*",
+    "uvicorn.*",
+    "vllm.*",
+    "youtube_transcript_api.*",
+]
+ignore_missing_imports = true
 
 [tool.poetry-dynamic-versioning]
 enable = true

src/avalan/__init__.py

@@ -5,18 +5,18 @@
 from packaging.version import Version, parse
 
 
-def _config() -> dict:
+def _config() -> dict[str, str]:
     config = metadata("avalan")
     package_version = metadata_version("avalan")
     return {
-        "name": config["Name"],
+        "name": str(config["Name"]),
         "version": package_version,
-        "license": config["License"],
+        "license": str(config["License"]),
         "url": "https://avalan.ai",
     }
 
 
-config = _config()
+config: dict[str, str] = _config()
 
 
 def license() -> str:

src/avalan/agent/engine.py

@@ -76,7 +76,8 @@ async def input_token_count(self) -> int | None:
                 },
             )
         )
-        count = self._model.input_token_count(
+        assert hasattr(self._model, "input_token_count")
+        count: int | None = self._model.input_token_count(
             self._last_prompt[0],
             system_prompt=self._last_prompt[1],
             developer_prompt=self._last_prompt[2],
@@ -158,6 +159,7 @@ async def __call__(
                 },
             )
         )
+        assert context.input is not None
         output = await self._run(context, context.input, **run_args)
         await self._event_manager.trigger(
             Event(
@@ -246,11 +248,18 @@ async def _run(
                 if previous_message:
                     await self.sync_message(previous_message)
 
-            for current_message in input_value:
+            for current_item in input_value:
+                current_message = (
+                    current_item
+                    if isinstance(current_item, Message)
+                    else Message(role=MessageRole.USER, content=current_item)
+                )
                 await self.sync_message(current_message)
 
             # Make recent memory the new model input
-            input_value = [rm.message for rm in self._memory.recent_messages]
+            recent = self._memory.recent_messages
+            assert recent is not None
+            input_value = [rm.message for rm in recent]
 
         # Have model generate output from input
 
@@ -289,7 +298,7 @@ async def _run(
             tool=self._tool,
             context=context,
         )
-        output = await self._model_manager(model_task)
+        output: TextGenerationResponse = await self._model_manager(model_task)
         await self._event_manager.trigger(
             Event(
                 type=EventType.MODEL_EXECUTE_AFTER,
@@ -344,6 +353,7 @@ async def sync_message(self, message: Message) -> None:
                 },
             )
         )
+        assert self._model.model_id is not None
         await self._memory.append_message(
             EngineMessage(
                 agent_id=self._id,

src/avalan/agent/loader.py

@@ -63,6 +63,7 @@ class OrchestratorLoader:
     )
     _OPENAI_RESPONSES_ALIASES = frozenset({"response", "responses"})
 
+    _event_manager: EventManager
     _hub: HuggingfaceHub
     _logger: Logger
     _participant_id: UUID
@@ -75,12 +76,19 @@ def __init__(
         logger: Logger,
         participant_id: UUID,
         stack: AsyncExitStack,
+        event_manager: EventManager | None = None,
     ) -> None:
+        self._event_manager = event_manager or EventManager()
         self._hub = hub
         self._logger = logger
         self._participant_id = participant_id
         self._stack = stack
 
+    @property
+    def event_manager(self) -> EventManager:
+        """Return the event manager instance."""
+        return self._event_manager
+
     @staticmethod
     def parse_permanent_store_value(
         value: str,
@@ -372,7 +380,7 @@ async def from_file(
                 agent_config=agent_config,
                 uri=uri,
                 engine_config=engine_config,
-                tools=enable_tools,
+                tools=enable_tools or [],
                 call_options=call_options,
                 template_vars=template_vars,
                 memory_permanent_message=memory_permanent_message,
@@ -495,7 +503,7 @@ async def from_settings(
 
         _l("Loading event manager")
 
-        event_manager = EventManager()
+        event_manager = self._event_manager
         if settings.log_events:
 
             def _log_event(event: Event) -> None:
@@ -607,6 +615,7 @@ def _log_event(event: Event) -> None:
 
         assert settings.agent_id
 
+        agent: Orchestrator
         if settings.orchestrator_type == "json":
             assert settings.json_config is not None
             agent = self._load_json_orchestrator(

src/avalan/agent/orchestrator/__init__.py

@@ -6,6 +6,7 @@
     Message,
     MessageContentText,
     MessageRole,
+    TransformerEngineSettings,
 )
 from ...entities import Modality as Modality
 from ...event import Event, EventType
@@ -14,18 +15,24 @@
 from ...model.call import ModelCallContext
 from ...model.engine import Engine
 from ...model.manager import ModelManager
+from ...model.response.text import (
+    TextGenerationResponse as TextGenerationResponse,
+)
 from ...tool.manager import ToolManager
 from .. import (
     AgentOperation,
-    EngineEnvironment,
     InputType,
     NoOperationAvailableException,
     Specification,
 )
+from .. import (
+    EngineEnvironment as EngineEnvironment,
+)
 from ..engine import EngineAgent
 from ..renderer import Renderer, TemplateEngineAgent
 from .response.orchestrator_response import OrchestratorResponse
 
+from collections.abc import Callable, Coroutine
 from contextlib import ExitStack
 from dataclasses import asdict
 from json import dumps
@@ -46,7 +53,7 @@ class Orchestrator:
     _memory: MemoryManager
     _tool: ToolManager
     _event_manager: EventManager
-    _engine_agents: dict[EngineEnvironment, EngineAgent] = {}
+    _engine_agents: dict[str, EngineAgent] = {}
     _engines_stack: ExitStack = ExitStack()
     _operation_step: int | None = None
     _model_ids: set[str] = set()
@@ -108,7 +115,9 @@ def id(self) -> UUID:
         return self._id
 
     @property
-    def input_token_count(self) -> int | None:
+    def input_token_count(
+        self,
+    ) -> Callable[[], Coroutine[Any, Any, int | None]] | None:
         return (
             self._last_engine_agent.input_token_count
             if self._last_engine_agent
@@ -243,7 +252,7 @@ async def __call__(self, input: Input, **kwargs) -> OrchestratorResponse:
 
         return OrchestratorResponse(
             messages,
-            result,
+            result,  # type: ignore[arg-type]
             engine_agent,
             operation,
             engine_args,
@@ -264,7 +273,11 @@ async def __aenter__(self):
             environment = operation.environment
             environment_hash = dumps(asdict(environment))
             if environment_hash not in self._engine_agents:
+                assert environment.engine_uri.model_id is not None
                 model_ids.append(environment.engine_uri.model_id)
+                assert isinstance(
+                    environment.settings, TransformerEngineSettings
+                )
                 engine = self._model_manager.load_engine(
                     environment.engine_uri,
                     environment.settings,
@@ -359,18 +372,22 @@ def _input_messages(
                     else message.content
                 )
                 render_vars.update({"input": message_content})
-                content = (
-                    self._renderer(self._user_template, **render_vars)
-                    if self._user_template
-                    else self._renderer.from_string(
+                if self._user_template:
+                    content = self._renderer(
+                        self._user_template, **render_vars
+                    )
+                else:
+                    assert self._user is not None
+                    content = self._renderer.from_string(
                         self._user, template_vars=render_vars
                     )
-                )
                 message = Message(role=message.role, content=content)
 
                 if isinstance(input, list):
-                    input[-1] = message
+                    input[-1] = message  # type: ignore[call-overload]
                 else:
                     input = message
 
-        return input
+        # The return type must be Message | list[Message] per the signature
+        assert isinstance(input, (Message, list))
+        return input  # type: ignore[return-value]

src/avalan/agent/orchestrator/orchestrators/default.py

@@ -1,4 +1,10 @@
-from ....agent import AgentOperation, EngineEnvironment, Goal, Specification
+from ....agent import (
+    AgentOperation,
+    EngineEnvironment,
+    Goal,
+    Role,
+    Specification,
+)
 from ....agent.orchestrator import Orchestrator
 from ....entities import EngineUri, Modality, TransformerEngineSettings
 from ....event.manager import EventManager
@@ -47,7 +53,7 @@ def __init__(
             )
         else:
             specification = Specification(
-                role=role,
+                role=Role(persona=[role]) if role else None,
                 goal=(
                     Goal(task=task, instructions=[instructions])
                     if task and instructions
@@ -66,7 +72,8 @@ def __init__(
             AgentOperation(
                 specification=specification,
                 environment=EngineEnvironment(
-                    engine_uri=engine_uri, settings=settings
+                    engine_uri=engine_uri,
+                    settings=settings or TransformerEngineSettings(),
                 ),
                 modality=Modality.TEXT_GENERATION,
             ),