fix typing

Author: Samoed

autointent/generation/_generator.py

@@ -67,11 +67,11 @@ def get_chat_completion(self, messages: list[Message]) -> str:
             messages: List of messages to send to the model.
         """
         response = self.client.chat.completions.create(
-            messages=messages,  # type: ignore[arg-type]
+            messages=messages,  # type: ignore[call-overload]
             model=self.model_name,
             **self.generation_params,
         )
-        return response.choices[0].message.content  # type: ignore[return-value]
+        return response.choices[0].message.content  # type: ignore[no-any-return]
 
     async def get_chat_completion_async(self, messages: list[Message]) -> str:
         """Prompt LLM and return its answer asynchronously.
@@ -80,11 +80,15 @@ async def get_chat_completion_async(self, messages: list[Message]) -> str:
             messages: List of messages to send to the model.
         """
         response = await self.async_client.chat.completions.create(
-            messages=messages,  # type: ignore[arg-type]
+            messages=messages,  # type: ignore[call-overload]
             model=self.model_name,
             **self.generation_params,
         )
-        return response.choices[0].message.content  # type: ignore[return-value]
+
+        if response is None or not response.choices:
+            msg = "No response received from the model."
+            raise RuntimeError(msg)
+        return response.choices[0].message.content  # type: ignore[no-any-return]
 
     def _create_retry_messages(self, error_message: str, raw: str | None) -> list[Message]:
         """Create a follow-up message for retry with error details and schema."""
@@ -128,7 +132,7 @@ async def _get_structured_output_openai_async(
                 model=self.model_name,
                 messages=messages,  # type: ignore[arg-type]
                 response_format=output_model,
-                **self.generation_params,
+                **self.generation_params,  # type: ignore[arg-type]
             )
             raw = response.choices[0].message.content
             res = response.choices[0].message.parsed
@@ -154,7 +158,7 @@ async def _get_structured_output_vllm_async(
             json_schema = output_model.model_json_schema()
             response = await self.async_client.chat.completions.create(
                 model=self.model_name,
-                messages=messages,  # type: ignore[arg-type]
+                messages=messages,  # type: ignore[call-overload]
                 extra_body={"guided_json": json_schema},
                 **self.generation_params,
             )
@@ -245,7 +249,7 @@ def _get_structured_output_openai_sync(
                 model=self.model_name,
                 messages=messages,  # type: ignore[arg-type]
                 response_format=output_model,
-                **self.generation_params,
+                **self.generation_params,  # type: ignore[arg-type]
             )
             raw = response.choices[0].message.content
             res = response.choices[0].message.parsed
@@ -271,7 +275,7 @@ def _get_structured_output_vllm_sync(
             json_schema = output_model.model_json_schema()
             response = self.client.chat.completions.create(
                 model=self.model_name,
-                messages=messages,  # type: ignore[arg-type]
+                messages=messages,  # type: ignore[call-overload]
                 extra_body={"guided_json": json_schema},
                 **self.generation_params,
             )

autointent/generation/intents/_description_generation.py

@@ -60,15 +60,10 @@ async def create_intent_description(
     intent_name = intent_name if intent_name is not None else ""
     utterances = random.sample(utterances, min(5, len(utterances)))
 
-    result = await client.get_chat_completion_async(
+    return await client.get_chat_completion_async(
         messages=prompt.to_messages(intent_name, utterances),
     )
 
-    if not isinstance(result, str):
-        error_text = f"Unexpected response type: expected str, got {type(result).__name__}"
-        raise TypeError(error_text)
-    return result
-
 
 async def generate_intent_descriptions(
     client: Generator,

tests/generation/intents/test_description_generation.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from unittest.mock import AsyncMock, Mock, patch
+from unittest.mock import AsyncMock, patch
 
 import pytest
 
@@ -243,8 +243,8 @@ async def test_generate_intent_descriptions_empty_utterances_patterns():
     mock_create.assert_called_once_with(
         messages=[
             {
-                 "role": "system",
-                 "content": prompt.system_text,
+                "role": "system",
+                "content": prompt.system_text,
             },
             {
                 "role": "user",