Add type hints where appropriate

Author: TomasKorbar

agents/metrics_middleware.py

@@ -11,7 +11,7 @@
 
 
 class MetricsMiddleware(RunMiddlewareProtocol):
-    def __init__(self):
+    def __init__(self) -> None:
         self.start_time: datetime | None = None
         self.end_time: datetime | None = None
         self.tool_calls: int = 0
@@ -28,10 +28,10 @@ def bind(self, ctx: RunContext) -> None:
             EmitterOptions(is_blocking=True, priority=1),
         )
 
-    async def _on_run_context_start(self, event: RunContextStartEvent, meta: EventMeta):
+    async def _on_run_context_start(self, event: RunContextStartEvent, meta: EventMeta) -> None:
         self.start_time = datetime.now()
 
-    async def _on_run_context_finish(self, event: RunContextFinishEvent, meta: EventMeta):
+    async def _on_run_context_finish(self, event: RunContextFinishEvent, meta: EventMeta) -> None:
         self.end_time = datetime.now()
 
     @property
@@ -40,5 +40,5 @@ def duration(self) -> float:
             return (self.end_time - self.start_time).total_seconds()
         return 0
 
-    def get_metrics(self) -> dict:
+    def get_metrics(self) -> dict[str, float]:
         return {"duration": self.duration}

agents/tests/e2e/conftest.py

@@ -1,7 +1,13 @@
+from typing import Generator
+
+
 import pytest
 
+
 @pytest.hookimpl(wrapper=True)
-def pytest_terminal_summary(terminalreporter: pytest.TerminalReporter, exitstatus, config: pytest.Config):
+def pytest_terminal_summary(
+    terminalreporter: pytest.TerminalReporter, exitstatus, config: pytest.Config
+) -> Generator:
     yield
     metrics = config.stash.get("metrics", None)
 

agents/tests/e2e/test_triage.py

@@ -9,9 +9,9 @@
 
 
 class TriageAgentTestCase:
-    def __init__(self, input, expected_output):
-        self.input = input
-        self.expected_output = expected_output
+    def __init__(self, input: str, expected_output: TriageOutputSchema):
+        self.input: str = input
+        self.expected_output: TriageOutputSchema = expected_output
         self.metrics: dict = None
 
     async def run(self) -> TriageState: