fabricatio 0.2.3.dev3__cp312-cp312-win_amd64.whl → 0.2.4.dev1__cp312-cp312-win_amd64.whl

fabricatio/models/usages.py

@@ -1,19 +1,18 @@
 """This module contains classes that manage the usage of language models and tools in tasks."""
 
 from asyncio import gather
-from typing import Callable, Dict, Iterable, List, Optional, Self, Set, Union, Unpack, overload
+from typing import Callable, Dict, Iterable, List, Optional, Self, Set, Type, Union, Unpack, overload
 
 import asyncstdlib
 import litellm
-import orjson
 from fabricatio._rust_instances import template_manager
 from fabricatio.config import configs
 from fabricatio.journal import logger
-from fabricatio.models.generic import Base, WithBriefing
+from fabricatio.models.generic import ScopedConfig, WithBriefing
 from fabricatio.models.kwargs_types import ChooseKwargs, EmbeddingKwargs, GenerateKwargs, LLMKwargs
 from fabricatio.models.task import Task
 from fabricatio.models.tool import Tool, ToolBox
-from fabricatio.models.utils import Messages, MilvusData
+from fabricatio.models.utils import Messages
 from fabricatio.parser import JsonCapture
 from litellm import stream_chunk_builder
 from litellm.types.utils import (
@@ -23,135 +22,16 @@ from litellm.types.utils import (
     StreamingChoices,
 )
 from litellm.utils import CustomStreamWrapper
-from pydantic import Field, HttpUrl, NonNegativeFloat, NonNegativeInt, PositiveInt, SecretStr
+from more_itertools import duplicates_everseen
+from pydantic import Field, NonNegativeInt, PositiveInt
 
 
-class LLMUsage(Base):
+class LLMUsage(ScopedConfig):
     """Class that manages LLM (Large Language Model) usage parameters and methods."""
 
-    llm_api_endpoint: Optional[HttpUrl] = None
-    """The OpenAI API endpoint."""
-
-    llm_api_key: Optional[SecretStr] = None
-    """The OpenAI API key."""
-
-    llm_timeout: Optional[PositiveInt] = None
-    """The timeout of the LLM model."""
-
-    llm_max_retries: Optional[PositiveInt] = None
-    """The maximum number of retries."""
-
-    llm_model: Optional[str] = None
-    """The LLM model name."""
-
-    llm_temperature: Optional[NonNegativeFloat] = None
-    """The temperature of the LLM model."""
-
-    llm_stop_sign: Optional[str | List[str]] = None
-    """The stop sign of the LLM model."""
-
-    llm_top_p: Optional[NonNegativeFloat] = None
-    """The top p of the LLM model."""
-
-    llm_generation_count: Optional[PositiveInt] = None
-    """The number of generations to generate."""
-
-    llm_stream: Optional[bool] = None
-    """Whether to stream the LLM model's response."""
-
-    llm_max_tokens: Optional[PositiveInt] = None
-    """The maximum number of tokens to generate."""
-
-    async def aembedding(
-        self,
-        input_text: List[str],
-        model: Optional[str] = None,
-        dimensions: Optional[int] = None,
-        timeout: Optional[PositiveInt] = None,
-        caching: Optional[bool] = False,
-    ) -> EmbeddingResponse:
-        """Asynchronously generates embeddings for the given input text.
-
-        Args:
-            input_text (List[str]): A list of strings to generate embeddings for.
-            model (Optional[str]): The model to use for embedding. Defaults to the instance's `llm_model` or the global configuration.
-            dimensions (Optional[int]): The dimensions of the embedding. Defaults to None.
-            timeout (Optional[PositiveInt]): The timeout for the embedding request. Defaults to the instance's `llm_timeout` or the global configuration.
-            caching (Optional[bool]): Whether to cache the embedding result. Defaults to False.
-
-
-        Returns:
-            EmbeddingResponse: The response containing the embeddings.
-        """
-        return await litellm.aembedding(
-            input=input_text,
-            caching=caching,
-            dimensions=dimensions,
-            model=model or self.llm_model or configs.llm.model,
-            timeout=timeout or self.llm_timeout or configs.llm.timeout,
-            api_key=self.llm_api_key.get_secret_value() if self.llm_api_key else configs.llm.api_key.get_secret_value(),
-            api_base=self.llm_api_endpoint.unicode_string().rstrip(
-                "/"
-            )  # seems embedding function takes no base_url end with a slash
-            if self.llm_api_endpoint
-            else configs.llm.api_endpoint.unicode_string().rstrip("/"),
-        )
-
-    @overload
-    async def vectorize(self, input_text: List[str], **kwargs: Unpack[EmbeddingKwargs]) -> List[List[float]]: ...
-    @overload
-    async def vectorize(self, input_text: str, **kwargs: Unpack[EmbeddingKwargs]) -> List[float]: ...
-
-    async def vectorize(
-        self, input_text: List[str] | str, **kwargs: Unpack[EmbeddingKwargs]
-    ) -> List[List[float]] | List[float]:
-        """Asynchronously generates vector embeddings for the given input text.
-
-        Args:
-            input_text (List[str] | str): A string or list of strings to generate embeddings for.
-            **kwargs (Unpack[EmbeddingKwargs]): Additional keyword arguments for embedding.
-
-        Returns:
-            List[List[float]] | List[float]: The generated embeddings.
-        """
-        if isinstance(input_text, str):
-            return (await self.aembedding([input_text], **kwargs)).data[0].get("embedding")
-
-        return [o.get("embedding") for o in (await self.aembedding(input_text, **kwargs)).data]
-
-    @overload
-    async def pack(
-        self, input_text: List[str], subject: Optional[str] = None, **kwargs: Unpack[EmbeddingKwargs]
-    ) -> List[MilvusData]: ...
-    @overload
-    async def pack(
-        self, input_text: str, subject: Optional[str] = None, **kwargs: Unpack[EmbeddingKwargs]
-    ) -> MilvusData: ...
-
-    async def pack(
-        self, input_text: List[str] | str, subject: Optional[str] = None, **kwargs: Unpack[EmbeddingKwargs]
-    ) -> List[MilvusData] | MilvusData:
-        """Asynchronously generates MilvusData objects for the given input text.
-
-        Args:
-            input_text (List[str] | str): A string or list of strings to generate embeddings for.
-            subject (Optional[str]): The subject of the input text. Defaults to None.
-            **kwargs (Unpack[EmbeddingKwargs]): Additional keyword arguments for embedding.
-
-        Returns:
-            List[MilvusData] | MilvusData: The generated MilvusData objects.
-        """
-        if isinstance(input_text, str):
-            return MilvusData(vector=await self.vectorize(input_text, **kwargs), text=input_text, subject=subject)
-        vecs = await self.vectorize(input_text, **kwargs)
-        return [
-            MilvusData(
-                vector=vec,
-                text=text,
-                subject=subject,
-            )
-            for text, vec in zip(input_text, vecs, strict=True)
-        ]
+    @classmethod
+    def _scoped_model(cls) -> Type["LLMUsage"]:
+        return LLMUsage
 
     async def aquery(
         self,
@@ -181,10 +61,8 @@ class LLMUsage(Base):
             stream=kwargs.get("stream") or self.llm_stream or configs.llm.stream,
             timeout=kwargs.get("timeout") or self.llm_timeout or configs.llm.timeout,
             max_retries=kwargs.get("max_retries") or self.llm_max_retries or configs.llm.max_retries,
-            api_key=self.llm_api_key.get_secret_value() if self.llm_api_key else configs.llm.api_key.get_secret_value(),
-            base_url=self.llm_api_endpoint.unicode_string()
-            if self.llm_api_endpoint
-            else configs.llm.api_endpoint.unicode_string(),
+            api_key=(self.llm_api_key or configs.llm.api_key).get_secret_value(),
+            base_url=(self.llm_api_endpoint or configs.llm.api_endpoint).unicode_string(),
         )
 
     async def ainvoke(
@@ -213,13 +91,13 @@ class LLMUsage(Base):
         if isinstance(resp, ModelResponse):
             return resp.choices
         if isinstance(resp, CustomStreamWrapper):
-            if configs.debug.streaming_visible:
-                chunks = []
-                async for chunk in resp:
-                    chunks.append(chunk)
-                    print(chunk.choices[0].delta.content or "", end="")  # noqa: T201
-                return stream_chunk_builder(chunks).choices
-            return stream_chunk_builder(await asyncstdlib.list()).choices
+            if not configs.debug.streaming_visible:
+                return stream_chunk_builder(await asyncstdlib.list()).choices
+            chunks = []
+            async for chunk in resp:
+                chunks.append(chunk)
+                print(chunk.choices[0].delta.content or "", end="")  # noqa: T201
+            return stream_chunk_builder(chunks).choices
         logger.critical(err := f"Unexpected response type: {type(resp)}")
         raise ValueError(err)
 
@@ -334,11 +212,10 @@ class LLMUsage(Base):
                     **kwargs,
                 )
             ) and (validated := validator(response)):
-                logger.debug(f"Successfully validated the response at {i}th attempt. response: \n{response}")
+                logger.debug(f"Successfully validated the response at {i}th attempt.")
                 return validated
-            logger.debug(f"Failed to validate the response at {i}th attempt. response: \n{response}")
-        logger.error(f"Failed to validate the response after {max_validations} attempts.")
-        raise ValueError("Failed to validate the response.")
+        logger.error(err := f"Failed to validate the response after {max_validations} attempts.")
+        raise ValueError(err)
 
     async def aask_validate_batch[T](
         self,
@@ -361,6 +238,26 @@ class LLMUsage(Base):
         """
         return await gather(*[self.aask_validate(question, validator, **kwargs) for question in questions])
 
+    async def aliststr(self, requirement: str, k: NonNegativeInt = 0, **kwargs: Unpack[GenerateKwargs]) -> List[str]:
+        """Asynchronously generates a list of strings based on a given requirement.
+
+        Args:
+            requirement (str): The requirement for the list of strings.
+            k (NonNegativeInt): The number of choices to select, 0 means infinite. Defaults to 0.
+            **kwargs (Unpack[GenerateKwargs]): Additional keyword arguments for the LLM usage.
+
+        Returns:
+            List[str]: The validated response as a list of strings.
+        """
+        return await self.aask_validate(
+            template_manager.render_template(
+                configs.templates.liststr_template,
+                {"requirement": requirement, "k": k},
+            ),
+            lambda resp: JsonCapture.validate_with(resp, target_type=list, elements_type=str, length=k),
+            **kwargs,
+        )
+
     async def achoose[T: WithBriefing](
         self,
         instruction: str,
@@ -384,28 +281,28 @@ class LLMUsage(Base):
             - Ensures response compliance through JSON parsing and format validation.
             - Relies on `aask_validate` to implement retry mechanisms with validation.
         """
+        if dup := duplicates_everseen(choices, key=lambda x: x.name):
+            logger.error(err := f"Redundant choices: {dup}")
+            raise ValueError(err)
         prompt = template_manager.render_template(
             configs.templates.make_choice_template,
             {
                 "instruction": instruction,
-                "options": [{"name": m.name, "briefing": m.briefing} for m in choices],
+                "options": [m.model_dump(include={"name", "briefing"}) for m in choices],
                 "k": k,
             },
         )
         names = {c.name for c in choices}
+
         logger.debug(f"Start choosing between {names} with prompt: \n{prompt}")
 
         def _validate(response: str) -> List[T] | None:
-            ret = JsonCapture.convert_with(response, orjson.loads)
-
-            if not isinstance(ret, List) or (0 < k != len(ret)):
-                logger.error(f"Incorrect Type or length of response: \n{ret}")
+            ret = JsonCapture.validate_with(response, target_type=List, elements_type=str, length=k)
+            if ret is None or set(ret) - names:
                 return None
-            if any(n not in names for n in ret):
-                logger.error(f"Invalid choice in response: \n{ret}")
-                return None
-
-            return [next(toolbox for toolbox in choices if toolbox.name == toolbox_str) for toolbox_str in ret]
+            return [
+                next(candidate for candidate in choices if candidate.name == candidate_name) for candidate_name in ret
+            ]
 
         return await self.aask_validate(
             question=prompt,
@@ -459,55 +356,91 @@ class LLMUsage(Base):
         Returns:
             bool: The judgment result (True or False) based on the AI's response.
         """
-
-        def _validate(response: str) -> bool | None:
-            ret = JsonCapture.convert_with(response, orjson.loads)
-            if not isinstance(ret, bool):
-                return None
-            return ret
-
         return await self.aask_validate(
             question=template_manager.render_template(
                 configs.templates.make_judgment_template,
                 {"prompt": prompt, "affirm_case": affirm_case, "deny_case": deny_case},
             ),
-            validator=_validate,
+            validator=lambda resp: JsonCapture.validate_with(resp, target_type=bool),
             **kwargs,
         )
 
-    def fallback_to(self, other: "LLMUsage") -> Self:
-        """Fallback to another instance's attribute values if the current instance's attributes are None.
+
+class EmbeddingUsage(LLMUsage):
+    """A class representing the embedding model."""
+
+    async def aembedding(
+        self,
+        input_text: List[str],
+        model: Optional[str] = None,
+        dimensions: Optional[int] = None,
+        timeout: Optional[PositiveInt] = None,
+        caching: Optional[bool] = False,
+    ) -> EmbeddingResponse:
+        """Asynchronously generates embeddings for the given input text.
 
         Args:
-            other (LLMUsage): Another instance from which to copy attribute values.
+            input_text (List[str]): A list of strings to generate embeddings for.
+            model (Optional[str]): The model to use for embedding. Defaults to the instance's `llm_model` or the global configuration.
+            dimensions (Optional[int]): The dimensions of the embedding output should have, which is used to validate the result. Defaults to None.
+            timeout (Optional[PositiveInt]): The timeout for the embedding request. Defaults to the instance's `llm_timeout` or the global configuration.
+            caching (Optional[bool]): Whether to cache the embedding result. Defaults to False.
+
 
         Returns:
-            Self: The current instance, allowing for method chaining.
+            EmbeddingResponse: The response containing the embeddings.
         """
-        # Iterate over the attribute names and copy values from 'other' to 'self' where applicable
-        # noinspection PydanticTypeChecker,PyTypeChecker
-        for attr_name in LLMUsage.model_fields:
-            # Copy the attribute value from 'other' to 'self' only if 'self' has None and 'other' has a non-None value
-            if getattr(self, attr_name) is None and (attr := getattr(other, attr_name)) is not None:
-                setattr(self, attr_name, attr)
-
-        # Return the current instance to allow for method chaining
-        return self
+        # check seq length
+        max_len = self.embedding_max_sequence_length or configs.embedding.max_sequence_length
+        if any(len(t) > max_len for t in input_text):
+            logger.error(err := f"Input text exceeds maximum sequence length {max_len}.")
+            raise ValueError(err)
+
+        return await litellm.aembedding(
+            input=input_text,
+            caching=caching or self.embedding_caching or configs.embedding.caching,
+            dimensions=dimensions or self.embedding_dimensions or configs.embedding.dimensions,
+            model=model or self.embedding_model or configs.embedding.model or self.llm_model or configs.llm.model,
+            timeout=timeout
+            or self.embedding_timeout
+            or configs.embedding.timeout
+            or self.llm_timeout
+            or configs.llm.timeout,
+            api_key=(
+                self.embedding_api_key or configs.embedding.api_key or self.llm_api_key or configs.llm.api_key
+            ).get_secret_value(),
+            api_base=(
+                self.embedding_api_endpoint
+                or configs.embedding.api_endpoint
+                or self.llm_api_endpoint
+                or configs.llm.api_endpoint
+            )
+            .unicode_string()
+            .rstrip("/"),
+            # seems embedding function takes no base_url end with a slash
+        )
+
+    @overload
+    async def vectorize(self, input_text: List[str], **kwargs: Unpack[EmbeddingKwargs]) -> List[List[float]]: ...
+    @overload
+    async def vectorize(self, input_text: str, **kwargs: Unpack[EmbeddingKwargs]) -> List[float]: ...
 
-    def hold_to(self, others: Union["LLMUsage", Iterable["LLMUsage"]]) -> Self:
-        """Hold to another instance's attribute values if the current instance's attributes are None.
+    async def vectorize(
+        self, input_text: List[str] | str, **kwargs: Unpack[EmbeddingKwargs]
+    ) -> List[List[float]] | List[float]:
+        """Asynchronously generates vector embeddings for the given input text.
 
         Args:
-            others (LLMUsage | Iterable[LLMUsage]): Another instance or iterable of instances from which to copy attribute values.
+            input_text (List[str] | str): A string or list of strings to generate embeddings for.
+            **kwargs (Unpack[EmbeddingKwargs]): Additional keyword arguments for embedding.
 
         Returns:
-            Self: The current instance, allowing for method chaining.
+            List[List[float]] | List[float]: The generated embeddings.
         """
-        for other in others:
-            # noinspection PyTypeChecker,PydanticTypeChecker
-            for attr_name in LLMUsage.model_fields:
-                if (attr := getattr(self, attr_name)) is not None and getattr(other, attr_name) is None:
-                    setattr(other, attr_name, attr)
+        if isinstance(input_text, str):
+            return (await self.aembedding([input_text], **kwargs)).data[0].get("embedding")
+
+        return [o.get("embedding") for o in (await self.aembedding(input_text, **kwargs)).data]
 
 
 class ToolBoxUsage(LLMUsage):

fabricatio/models/utils.py

@@ -1,5 +1,6 @@
 """A module containing utility classes for the models."""
 
+from enum import Enum
 from typing import Any, Dict, List, Literal, Optional, Self
 
 from pydantic import BaseModel, ConfigDict, Field
@@ -125,3 +126,21 @@ class MilvusData(BaseModel):
         """
         self.id = new_id
         return self
+
+
+class TaskStatus(Enum):
+    """An enumeration representing the status of a task.
+
+    Attributes:
+        Pending: The task is pending.
+        Running: The task is currently running.
+        Finished: The task has been successfully completed.
+        Failed: The task has failed.
+        Cancelled: The task has been cancelled.
+    """
+
+    Pending = "pending"
+    Running = "running"
+    Finished = "finished"
+    Failed = "failed"
+    Cancelled = "cancelled"

fabricatio/parser.py

@@ -1,9 +1,10 @@
 """A module to parse text using regular expressions."""
 
-from typing import Any, Callable, Self, Tuple
+from typing import Any, Callable, Optional, Self, Tuple, Type
 
+import orjson
 import regex
-from pydantic import BaseModel, ConfigDict, Field, PositiveInt, PrivateAttr
+from pydantic import BaseModel, ConfigDict, Field, PositiveInt, PrivateAttr, ValidationError
 from regex import Pattern, compile
 
 from fabricatio.journal import logger
@@ -27,11 +28,7 @@ class Capture(BaseModel):
     _compiled: Pattern = PrivateAttr()
 
     def model_post_init(self, __context: Any) -> None:
-        """Initialize the compiled regular expression pattern after the model is initialized.
-
-        Args:
-            __context (Any): The context in which the model is initialized.
-        """
+        """Initialize the compiled pattern."""
         self._compiled = compile(self.pattern, self.flags)
 
     def capture(self, text: str) -> Tuple[str, ...] | str | None:
@@ -70,10 +67,40 @@ class Capture(BaseModel):
             return None
         try:
             return convertor(cap)
-        except (ValueError, SyntaxError) as e:
+        except (ValueError, SyntaxError, ValidationError) as e:
             logger.error(f"Failed to convert text using {convertor.__name__} to convert.\nerror: {e}\n {cap}")
             return None
 
+    def validate_with[K, T, E](
+        self,
+        text: str,
+        target_type: Type[T],
+        elements_type: Optional[Type[E]] = None,
+        length: Optional[int] = None,
+        deserializer: Callable[[Tuple[str, ...]], K] | Callable[[str], K] = orjson.loads,
+    ) -> T | None:
+        """Validate the given text using the pattern.
+
+        Args:
+            text (str): The text to search the pattern in.
+            target_type (Type[T]): The expected type of the output, dict or list.
+            elements_type (Optional[Type[E]]): The expected type of the elements in the output dict keys or list elements.
+            length (Optional[int]): The expected length of the output, bool(length)==False means no length validation.
+            deserializer (Callable[[Tuple[str, ...]], K] | Callable[[str], K]): The function to deserialize the captured text.
+
+        Returns:
+            T | None: The validated text if the pattern is found and the output is of the expected type, otherwise None.
+        """
+        judges = [lambda output_obj: isinstance(output_obj, target_type)]
+        if elements_type:
+            judges.append(lambda output_obj: all(isinstance(e, elements_type) for e in output_obj))
+        if length:
+            judges.append(lambda output_obj: len(output_obj) == length)
+
+        if (out := self.convert_with(text, deserializer)) and all(j(out) for j in judges):
+            return out
+        return None
+
     @classmethod
     def capture_code_block(cls, language: str) -> Self:
         """Capture the first occurrence of a code block in the given text.