update description generator (#237)

* update description generator

* fix test

* fix typing

* fix tests

Author: Samoed

autointent/generation/_generator.py

@@ -5,7 +5,7 @@
 import os
 from pathlib import Path
 from textwrap import dedent
-from typing import Any, ClassVar, Literal, TypedDict, TypeVar
+from typing import Any, Literal, TypedDict, TypeVar
 
 import openai
 from dotenv import load_dotenv
@@ -57,27 +57,21 @@ class Generator:
 
     _dump_data_filename = "init_params.json"
 
-    _default_generation_params: ClassVar[dict[str, Any]] = {
-        "max_tokens": 150,
-        "n": 1,
-        "stop": None,
-        "temperature": 0.7,
-    }
-    """Default generation parameters for API requests."""
-
     def __init__(
         self,
         base_url: str | None = None,
         model_name: str | None = None,
         use_cache: bool = True,
-        **generation_params: Any,  # noqa: ANN401
+        client_params: dict[str, Any] | None = None,
+        **generation_params: dict[str, Any],
     ) -> None:
         """Initialize the Generator with API configuration.
 
         Args:
             base_url: OpenAI API compatible server URL.
             model_name: Name of the language model to use.
             use_cache: Whether to use caching for structured outputs.
+            client_params: Additional parameters for client.
             **generation_params: Additional generation parameters to override defaults passed to OpenAI completions API.
         """
         base_url = base_url or os.getenv("OPENAI_BASE_URL")
@@ -91,27 +85,23 @@ def __init__(
         self.base_url = base_url
         self.use_cache = use_cache
 
-        self.client = openai.OpenAI(base_url=base_url)
-        self.async_client = openai.AsyncOpenAI(base_url=base_url)
+        self.client = openai.OpenAI(base_url=base_url, **(client_params or {}))
+        self.async_client = openai.AsyncOpenAI(base_url=base_url, **(client_params or {}))
+        self.generation_params = generation_params
         self.cache = StructuredOutputCache(use_cache=use_cache)
 
-        self.generation_params = {
-            **self._default_generation_params,
-            **generation_params,
-        }  #  https://stackoverflow.com/a/65539348
-
     def get_chat_completion(self, messages: list[Message]) -> str:
         """Prompt LLM and return its answer.
 
         Args:
             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.
@@ -120,11 +110,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."""
@@ -168,7 +162,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
@@ -194,12 +188,12 @@ 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,
             )
             raw = response.choices[0].message.content
-            res = output_model.model_validate_json(raw)  # type: ignore[arg-type]
+            res = output_model.model_validate_json(raw)
         except (ValidationError, ValueError) as e:
             msg = f"Failed to obtain structured output for model {self.model_name} and messages {messages}: {e!s}"
             logger.warning(msg)
@@ -252,6 +246,10 @@ async def get_structured_output_async(
             current_messages.extend(self._create_retry_messages(error, raw))
 
         if res is None:
+            msg = (
+                f"Failed to generate valid structured output after {max_retries + 1} attempts.\n"
+                f"Messages: {current_messages}"
+            )
             logger.exception(msg)
             raise RetriesExceededError(max_retries=max_retries, messages=current_messages)
 
@@ -281,7 +279,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
@@ -307,12 +305,12 @@ 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,
             )
             raw = response.choices[0].message.content
-            res = output_model.model_validate_json(raw)  # type: ignore[arg-type]
+            res = output_model.model_validate_json(raw)
         except (ValidationError, ValueError) as e:
             msg = f"Failed to obtain structured output for model {self.model_name} and messages {messages}: {e!s}"
             logger.warning(msg)
@@ -365,6 +363,7 @@ def get_structured_output_sync(
             current_messages.extend(self._create_retry_messages(error, raw))
 
         if res is None:
+            msg = "Structured output returned None but no error was caught."
             logger.exception(msg)
             raise RetriesExceededError(max_retries=max_retries, messages=current_messages)
 

autointent/generation/chat_templates/_intent_descriptions.py

@@ -2,77 +2,63 @@
 
 from pydantic import BaseModel, field_validator
 
-PROMPT_DESCRIPTION = """
+from autointent.generation.chat_templates import Message, Role
+
+PROMPT_DESCRIPTION_SYSTEM = """
 Your task is to write a description of the intent.
 
-You are given the name of the intent, user intentions related to it, and
-regular expressions that match user utterances. The description should be:
+You are given the name of the intent, user intentions related to it. The description should be:
 1) In declarative form.
 2) No more than one sentence.
-3) In the language in which the utterances or regular expressions are written.
+3) In the language in which the utterances.
 
 Remember:
 - Respond with just the description, no extra details.
-- Keep in mind that either the names, user queries, or regex patterns may not be provided.
+- Keep in mind that either the names or user queries may not be provided.
 
 For example:
 
+Input:
 name:
 activate_my_card
 user utterances:
-Please help me with my card. It won't activate.
-I tried but am unable to activate my card.
-I want to start using my card.
-regex patterns:
-(activate.*card)|(start.*using.*card)
-description:
+- Please help me with my card. It won't activate.
+- I tried but am unable to activate my card.
+- I want to start using my card.
+
+Output:
 User wants to activate his card.
 
+Input:
 name:
 beneficiary_not_allowed
 user utterances:
 
-regex patterns:
-(not.*allowed.*beneficiary)|(cannot.*add.*beneficiary)
-description:
+Output:
 User wants to know why his beneficiary is not allowed.
-
-name:
-vacation_registration
-user utterances:
-как оформить отпуск
-в какие даты надо оформить отпуск
-как запланировать отпуск
-regex patterns:
-
-description:
-Пользователь спрашивает про оформление отпуска.
-
+"""
+PROMPT_DESCRIPTION_USER = """
 name:
 {intent_name}
 user utterances:
 {user_utterances}
-regex patterns:
-{regex_patterns}
-description:
-
 """
 
 
 class PromptDescription(BaseModel):
     """Prompt description configuration."""
 
-    text: str = PROMPT_DESCRIPTION
+    system_text: str = PROMPT_DESCRIPTION_SYSTEM
+    user_text: str = PROMPT_DESCRIPTION_USER
     """
     The template for the prompt to generate descriptions for intents.
     Should include placeholders for {intent_name} and {user_utterances}.
     - `{intent_name}` will be replaced with the name of the intent.
     - `{user_utterances}` will be replaced with the user utterances related to the intent.
-    - (optionally) `{regex_patterns}` will be replaced with the regular expressions that match user utterances.
     """
 
     @classmethod
-    @field_validator("text")
+    @field_validator("user_text")
     def check_valid_prompt(cls, value: str) -> str:
         """Validate the prompt description template.
 
@@ -89,3 +75,13 @@ def check_valid_prompt(cls, value: str) -> str:
             )
             raise ValueError(text_error)
         return value
+
+    def to_messages(self, intent_name: str | None, utterances: list[str]) -> list[Message]:
+        user_message_content = self.user_text.format(
+            intent_name=intent_name,
+            user_utterances="\n - ".join(utterances),
+        )
+        return [
+            Message(role=Role.SYSTEM, content=self.system_text),
+            Message(role=Role.USER, content=user_message_content),
+        ]

autointent/generation/intents/_description_generation.py

@@ -9,9 +9,8 @@
 import random
 from collections import defaultdict
 
-from openai import AsyncOpenAI
-
 from autointent import Dataset
+from autointent.generation import Generator
 from autointent.generation.chat_templates import PromptDescription
 from autointent.schemas import Intent, Sample
 
@@ -41,55 +40,36 @@ def group_utterances_by_label(samples: list[Sample]) -> dict[int, list[str]]:
 
 
 async def create_intent_description(
-    client: AsyncOpenAI,
+    client: Generator,
     intent_name: str | None,
     utterances: list[str],
-    regex_patterns: list[str],
     prompt: PromptDescription,
-    model_name: str,
 ) -> str:
     """Generate a description for a specific intent using an OpenAI model.
 
     Args:
         client: OpenAI client instance for model communication.
         intent_name: Name of the intent to describe (empty string if None).
         utterances: Example utterances related to the intent.
-        regex_patterns: Regular expression patterns associated with the intent.
         prompt: Template for model prompt with placeholders for intent_name,
                user_utterances, and regex_patterns.
-        model_name: Identifier of the OpenAI model to use.
 
     Raises:
         TypeError: If the model response is not a string.
     """
     intent_name = intent_name if intent_name is not None else ""
     utterances = random.sample(utterances, min(5, len(utterances)))
-    regex_patterns = random.sample(regex_patterns, min(3, len(regex_patterns)))
 
-    content = prompt.text.format(
-        intent_name=intent_name,
-        user_utterances="\n".join(utterances),
-        regex_patterns="\n".join(regex_patterns),
-    )
-    chat_completion = await client.chat.completions.create(
-        messages=[{"role": "user", "content": content}],
-        model=model_name,
-        temperature=0.2,
+    return await client.get_chat_completion_async(
+        messages=prompt.to_messages(intent_name, utterances),
     )
-    result = chat_completion.choices[0].message.content
-
-    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: AsyncOpenAI,
+    client: Generator,
     intent_utterances: dict[int, list[str]],
     intents: list[Intent],
     prompt: PromptDescription,
-    model_name: str,
 ) -> list[Intent]:
     """Generate descriptions for multiple intents using an OpenAI model.
 
@@ -99,22 +79,18 @@ async def generate_intent_descriptions(
         intents: List of intents needing descriptions.
         prompt: Template for model prompt with placeholders for intent_name,
                user_utterances, and regex_patterns.
-        model_name: Name of the OpenAI model to use.
     """
     tasks = []
     for intent in intents:
         if intent.description is not None:
             continue
         utterances = intent_utterances.get(intent.id, [])
-        regex_patterns = intent.regex_full_match + intent.regex_partial_match
         task = asyncio.create_task(
             create_intent_description(
                 client=client,
                 intent_name=intent.name,
                 utterances=utterances,
-                regex_patterns=regex_patterns,
                 prompt=prompt,
-                model_name=model_name,
             ),
         )
         tasks.append((intent, task))
@@ -127,8 +103,7 @@ async def generate_intent_descriptions(
 
 def generate_descriptions(
     dataset: Dataset,
-    client: AsyncOpenAI,
-    model_name: str,
+    client: Generator,
     prompt: PromptDescription | None = None,
 ) -> Dataset:
     """Add LLM-generated text descriptions to dataset's intents.
@@ -138,7 +113,6 @@ def generate_descriptions(
         client: OpenAI client for generating descriptions.
         prompt: Template for model prompt with placeholders for intent_name,
                user_utterances, and regex_patterns.
-        model_name: OpenAI model identifier for generating descriptions.
 
     See :ref:`intent_description_generation` tutorial.
     """
@@ -149,6 +123,6 @@ def generate_descriptions(
     if prompt is None:
         prompt = PromptDescription()
     dataset.intents = asyncio.run(
-        generate_intent_descriptions(client, intent_utterances, dataset.intents, prompt, model_name),
+        generate_intent_descriptions(client, intent_utterances, dataset.intents, prompt),
     )
     return dataset