fabricatio 0.2.6__cp39-cp39-win_amd64.whl

fabricatio/models/generic.py

@@ -0,0 +1,406 @@
+"""This module defines generic classes for models in the Fabricatio library."""
+
+from abc import abstractmethod
+from pathlib import Path
+from typing import Any, Callable, Dict, Iterable, List, Optional, Self, Union, final, overload
+
+import orjson
+from fabricatio._rust import blake3_hash
+from fabricatio._rust_instances import TEMPLATE_MANAGER
+from fabricatio.config import configs
+from fabricatio.fs.readers import MAGIKA, safe_text_read
+from fabricatio.journal import logger
+from fabricatio.parser import JsonCapture
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    HttpUrl,
+    NonNegativeFloat,
+    PositiveFloat,
+    PositiveInt,
+    SecretStr,
+)
+
+
+class Base(BaseModel):
+    """Base class for all models with Pydantic configuration."""
+
+    model_config = ConfigDict(use_attribute_docstrings=True)
+
+
+class Display(Base):
+    """Class that provides a method to display the model in a formatted JSON string."""
+
+    def display(self) -> str:
+        """Display the model in a formatted JSON string.
+
+        Returns:
+            str: The formatted JSON string of the model.
+        """
+        return self.model_dump_json(indent=1)
+
+    def compact(self) -> str:
+        """Display the model in a compact JSON string.
+
+        Returns:
+            str: The compact JSON string of the model.
+        """
+        return self.model_dump_json()
+
+
+class Named(Base):
+    """Class that includes a name attribute."""
+
+    name: str = Field(frozen=True)
+    """The name of the object."""
+
+
+class Described(Base):
+    """Class that includes a description attribute."""
+
+    description: str = Field(default="", frozen=True)
+    """The description of the object."""
+
+
+class WithBriefing(Named, Described):
+    """Class that provides a briefing based on the name and description."""
+
+    @property
+    def briefing(self) -> str:
+        """Get the briefing of the object.
+
+        Returns:
+            str: The briefing of the object.
+        """
+        return f"{self.name}: {self.description}" if self.description else self.name
+
+    def prepend[D: Dict[str, Any]](self, kwargs: D) -> D:
+        """Prepend the briefing to the system message in the kwargs.
+
+        Args:
+            kwargs (Dict[str, Any]): The keyword arguments to modify.
+
+        Returns:
+            Dict[str, Any]: The modified keyword arguments.
+        """
+        kwargs["system_message"] = f"# your personal briefing: \n{self.briefing}\n" + kwargs.get("system_message", "")
+        return kwargs
+
+
+class WithFormatedJsonSchema(Base):
+    """Class that provides a formatted JSON schema of the model."""
+
+    @classmethod
+    def formated_json_schema(cls) -> str:
+        """Get the JSON schema of the model in a formatted string.
+
+        Returns:
+            str: The JSON schema of the model in a formatted string.
+        """
+        return orjson.dumps(
+            cls.model_json_schema(),
+            option=orjson.OPT_INDENT_2 | orjson.OPT_SORT_KEYS,
+        ).decode()
+
+
+class CreateJsonObjPrompt(WithFormatedJsonSchema):
+    """Class that provides a prompt for creating a JSON object."""
+
+    @classmethod
+    @overload
+    def create_json_prompt(cls, requirement: List[str]) -> List[str]: ...
+
+    @classmethod
+    @overload
+    def create_json_prompt(cls, requirement: str) -> str: ...
+
+    @classmethod
+    def create_json_prompt(cls, requirement: str | List[str]) -> str | List[str]:
+        """Create the prompt for creating a JSON object with given requirement.
+
+        Args:
+            requirement (str): The requirement for the JSON object.
+
+        Returns:
+            str: The prompt for creating a JSON object with given requirement.
+        """
+        if isinstance(requirement, str):
+            return TEMPLATE_MANAGER.render_template(
+                configs.templates.create_json_obj_template,
+                {"requirement": requirement, "json_schema": cls.formated_json_schema()},
+            )
+        return [
+            TEMPLATE_MANAGER.render_template(
+                configs.templates.create_json_obj_template,
+                {"requirement": r, "json_schema": cls.formated_json_schema()},
+            )
+            for r in requirement
+        ]
+
+
+class InstantiateFromString(Base):
+    """Class that provides a method to instantiate the class from a string."""
+
+    @classmethod
+    def instantiate_from_string(cls, string: str) -> Self | None:
+        """Instantiate the class from a string.
+
+        Args:
+            string (str): The string to instantiate the class from.
+
+        Returns:
+            Self | None: The instance of the class or None if the string is not valid.
+        """
+        return JsonCapture.convert_with(string, cls.model_validate_json)
+
+
+class ProposedAble(CreateJsonObjPrompt, InstantiateFromString):
+    """Class that provides a method to propose a JSON object based on the requirement."""
+
+    pass
+
+
+class FinalizedDumpAble(Base):
+    """Class that provides a method to finalize the dump of the object."""
+
+    @abstractmethod
+    def finalized_dump(self) -> str:
+        """Finalize the dump of the object.
+
+        Returns:
+            str: The finalized dump of the object.
+        """
+
+    def finalized_dump_to(self, path: str | Path) -> Self:
+        """Finalize the dump of the object to a file.
+
+        Args:
+            path (str | Path): The path to save the finalized dump.
+
+        Returns:
+            Self: The current instance of the object.
+        """
+        Path(path).write_text(self.finalized_dump(), encoding="utf-8")
+        return self
+
+
+class WithDependency(Base):
+    """Class that manages file dependencies."""
+
+    dependencies: List[str] = Field(default_factory=list)
+    """The file dependencies which is needed to read or write to meet a specific requirement, a list of file paths."""
+
+    def add_dependency[P: str | Path](self, dependency: P | List[P]) -> Self:
+        """Add a file dependency to the task.
+
+        Args:
+            dependency (str | Path | List[str | Path]): The file dependency to add to the task.
+
+        Returns:
+            Self: The current instance of the task.
+        """
+        if not isinstance(dependency, list):
+            dependency = [dependency]
+        self.dependencies.extend(Path(d).as_posix() for d in dependency)
+        return self
+
+    def remove_dependency[P: str | Path](self, dependency: P | List[P]) -> Self:
+        """Remove a file dependency from the task.
+
+        Args:
+            dependency (str | Path | List[str | Path]): The file dependency to remove from the task.
+
+        Returns:
+            Self: The current instance of the task.
+        """
+        if not isinstance(dependency, list):
+            dependency = [dependency]
+        for d in dependency:
+            self.dependencies.remove(Path(d).as_posix())
+        return self
+
+    def clear_dependencies(self) -> Self:
+        """Clear all file dependencies from the task.
+
+        Returns:
+            Self: The current instance of the task.
+        """
+        self.dependencies.clear()
+        return self
+
+    def override_dependencies[P: str | Path](self, dependencies: List[P] | P) -> Self:
+        """Override the file dependencies of the task.
+
+        Args:
+            dependencies (List[str | Path] | str | Path): The file dependencies to override the task's dependencies.
+
+        Returns:
+            Self: The current instance of the task.
+        """
+        return self.clear_dependencies().add_dependency(dependencies)
+
+    def pop_dependence[T](self, idx: int = -1, reader: Callable[[str], T] = safe_text_read) -> T:
+        """Pop the file dependencies from the task.
+
+        Returns:
+            str: The popped file dependency
+        """
+        return reader(self.dependencies.pop(idx))
+
+    @property
+    def dependencies_prompt(self) -> str:
+        """Generate a prompt for the task based on the file dependencies.
+
+        Returns:
+            str: The generated prompt for the task.
+        """
+        return TEMPLATE_MANAGER.render_template(
+            configs.templates.dependencies_template,
+            {
+                (pth := Path(p)).name: {
+                    "path": pth.as_posix(),
+                    "exists": pth.exists(),
+                    "description": (identity := MAGIKA.identify_path(pth)).output.description,
+                    "size": f"{pth.stat().st_size / (1024 * 1024) if pth.exists() and pth.is_file() else 0:.3f} MB",
+                    "content": (text := safe_text_read(pth)),
+                    "lines": len(text.splitlines()),
+                    "language": identity.output.ct_label,
+                    "checksum": blake3_hash(pth.read_bytes()) if pth.exists() and pth.is_file() else "unknown",
+                }
+                for p in self.dependencies
+            },
+        )
+
+
+class PrepareVectorization(Base):
+    """Class that prepares the vectorization of the model."""
+
+    @abstractmethod
+    def _prepare_vectorization_inner(self) -> str:
+        """Prepare the vectorization of the model."""
+
+    def prepare_vectorization(self, max_length: Optional[int] = None) -> str:
+        """Prepare the vectorization of the model.
+
+        Returns:
+            str: The prepared vectorization of the model.
+        """
+        max_length = max_length or configs.embedding.max_sequence_length
+        chunk = self._prepare_vectorization_inner()
+        if len(chunk) > max_length:
+            logger.error(err := f"Chunk exceeds maximum sequence length {max_length}.")
+            raise ValueError(err)
+
+        return chunk
+
+
+class ScopedConfig(Base):
+    """Class that manages a scoped configuration."""
+
+    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."""
+
+    llm_tpm: Optional[PositiveInt] = None
+    """The tokens per minute of the LLM model."""
+
+    llm_rpm: Optional[PositiveInt] = None
+    """The requests per minute of the LLM model."""
+
+    embedding_api_endpoint: Optional[HttpUrl] = None
+    """The OpenAI API endpoint."""
+
+    embedding_api_key: Optional[SecretStr] = None
+    """The OpenAI API key."""
+
+    embedding_timeout: Optional[PositiveInt] = None
+    """The timeout of the LLM model."""
+
+    embedding_model: Optional[str] = None
+    """The LLM model name."""
+
+    embedding_max_sequence_length: Optional[PositiveInt] = None
+    """The maximum sequence length."""
+
+    embedding_dimensions: Optional[PositiveInt] = None
+    """The dimensions of the embedding."""
+    embedding_caching: Optional[bool] = False
+    """Whether to cache the embedding result."""
+
+    milvus_uri: Optional[HttpUrl] = Field(default=None)
+    """The URI of the Milvus server."""
+    milvus_token: Optional[SecretStr] = Field(default=None)
+    """The token for the Milvus server."""
+    milvus_timeout: Optional[PositiveFloat] = Field(default=None)
+    """The timeout for the Milvus server."""
+    milvus_dimensions: Optional[PositiveInt] = Field(default=None)
+    """The dimensions of the Milvus server."""
+
+    @final
+    def fallback_to(self, other: "ScopedConfig") -> Self:
+        """Fallback to another instance's attribute values if the current instance's attributes are None.
+
+        Args:
+            other (LLMUsage): Another instance from which to copy attribute values.
+
+        Returns:
+            Self: The current instance, allowing for method chaining.
+        """
+        # Iterate over the attribute names and copy values from 'other' to 'self' where applicable
+        # noinspection PydanticTypeChecker,PyTypeChecker
+        for attr_name in ScopedConfig.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
+
+    @final
+    def hold_to(self, others: Union["ScopedConfig", Iterable["ScopedConfig"]]) -> Self:
+        """Hold to another instance's attribute values if the current instance's attributes are None.
+
+        Args:
+            others (LLMUsage | Iterable[LLMUsage]): Another instance or iterable of instances from which to copy attribute values.
+
+        Returns:
+            Self: The current instance, allowing for method chaining.
+        """
+        if not isinstance(others, Iterable):
+            others = [others]
+        for other in others:
+            # noinspection PyTypeChecker,PydanticTypeChecker
+            for attr_name in ScopedConfig.model_fields:
+                if (attr := getattr(self, attr_name)) is not None and getattr(other, attr_name) is None:
+                    setattr(other, attr_name, attr)
+        return self

fabricatio/models/kwargs_types.py

@@ -0,0 +1,169 @@
+"""This module contains the types for the keyword arguments of the methods in the models module."""
+
+from importlib.util import find_spec
+from typing import Any, Required, TypedDict
+
+from litellm.caching.caching import CacheMode
+from litellm.types.caching import CachingSupportedCallTypes
+
+if find_spec("pymilvus"):
+    from pymilvus import CollectionSchema
+    from pymilvus.milvus_client import IndexParams
+
+    class CollectionConfigKwargs(TypedDict, total=False):
+        """Configuration parameters for a vector collection.
+
+        These arguments are typically used when configuring connections to vector databases.
+        """
+
+        dimension: int | None
+        primary_field_name: str
+        id_type: str
+        vector_field_name: str
+        metric_type: str
+        timeout: float | None
+        schema: CollectionSchema | None
+        index_params: IndexParams | None
+
+
+class FetchKwargs(TypedDict, total=False):
+    """Arguments for fetching data from vector collections.
+
+    Controls how data is retrieved from vector databases, including filtering
+    and result limiting parameters.
+    """
+
+    collection_name: str | None
+    similarity_threshold: float
+    result_per_query: int
+
+
+class EmbeddingKwargs(TypedDict, total=False):
+    """Configuration parameters for text embedding operations.
+
+    These settings control the behavior of embedding models that convert text
+    to vector representations.
+    """
+
+    model: str
+    dimensions: int
+    timeout: int
+    caching: bool
+
+
+class LLMKwargs(TypedDict, total=False):
+    """Configuration parameters for language model inference.
+
+    These arguments control the behavior of large language model calls,
+    including generation parameters and caching options.
+    """
+
+    model: str
+    temperature: float
+    stop: str | list[str]
+    top_p: float
+    max_tokens: int
+    stream: bool
+    timeout: int
+    max_retries: int
+    no_cache: bool  # if the req uses cache in this call
+    no_store: bool  # If store the response of this call to cache
+    cache_ttl: int  # how long the stored cache is alive, in seconds
+    s_maxage: int  # max accepted age of cached response, in seconds
+
+
+class GenerateKwargs(LLMKwargs, total=False):
+    """Arguments for content generation operations.
+
+    Extends LLMKwargs with additional parameters specific to generation tasks,
+    such as the number of generated items and the system message.
+    """
+
+    system_message: str
+
+
+class ValidateKwargs[T](GenerateKwargs, total=False):
+    """Arguments for content validation operations.
+
+    Extends LLMKwargs with additional parameters specific to validation tasks,
+    such as limiting the number of validation attempts.
+    """
+
+    default: T
+    max_validations: int
+    co_extractor: GenerateKwargs
+
+
+# noinspection PyTypedDict
+class ReviewKwargs[T](ValidateKwargs[T], total=False):
+    """Arguments for content review operations.
+
+    Extends GenerateKwargs with parameters for evaluating content against
+    specific topics and review criteria.
+    """
+
+    topic: Required[str]
+    criteria: set[str]
+
+
+class CorrectKwargs[T](ReviewKwargs[T], total=False):
+    """Arguments for content correction operations.
+
+    Extends GenerateKwargs with parameters for correcting content based on
+    specific criteria and templates.
+    """
+
+    reference: str
+    supervisor_check: bool
+
+
+# noinspection PyTypedDict
+class ChooseKwargs[T](ValidateKwargs[T], total=False):
+    """Arguments for selection operations.
+
+    Extends GenerateKwargs with parameters for selecting among options,
+    such as the number of items to choose.
+    """
+
+    k: int
+
+
+class CacheKwargs(TypedDict, total=False):
+    """Configuration parameters for the caching system.
+
+    These arguments control the behavior of various caching backends,
+    including in-memory, Redis, S3, and vector database caching options.
+    """
+
+    mode: CacheMode  # when default_on cache is always on, when default_off cache is opt in
+    host: str
+    port: str
+    password: str
+    namespace: str
+    ttl: float
+    default_in_memory_ttl: float
+    default_in_redis_ttl: float
+    similarity_threshold: float
+    supported_call_types: list[CachingSupportedCallTypes]
+    # s3 Bucket, boto3 configuration
+    s3_bucket_name: str
+    s3_region_name: str
+    s3_api_version: str
+    s3_use_ssl: bool
+    s3_verify: bool | str
+    s3_endpoint_url: str
+    s3_aws_access_key_id: str
+    s3_aws_secret_access_key: str
+    s3_aws_session_token: str
+    s3_config: Any
+    s3_path: str
+    redis_semantic_cache_use_async: bool
+    redis_semantic_cache_embedding_model: str
+    redis_flush_size: int
+    redis_startup_nodes: list
+    disk_cache_dir: Any
+    qdrant_api_base: str
+    qdrant_api_key: str
+    qdrant_collection_name: str
+    qdrant_quantization_config: str
+    qdrant_semantic_cache_embedding_model: str

fabricatio/models/role.py

@@ -0,0 +1,72 @@
+"""Module that contains the Role class for managing workflows and their event registrations."""
+
+from typing import Any, Self, Set
+
+from fabricatio.capabilities.correct import Correct
+from fabricatio.capabilities.task import HandleTask, ProposeTask
+from fabricatio.core import env
+from fabricatio.journal import logger
+from fabricatio.models.action import WorkFlow
+from fabricatio.models.events import Event
+from fabricatio.models.tool import ToolBox
+from pydantic import Field
+
+
+class Role(ProposeTask, HandleTask, Correct):
+    """Class that represents a role with a registry of events and workflows.
+
+    A Role serves as a container for workflows, managing their registration to events
+    and providing them with shared configuration like tools and personality.
+
+    Attributes:
+        registry: Mapping of events to workflows that handle them
+        toolboxes: Set of toolboxes available to this role and its workflows
+    """
+
+    registry: dict[Event | str, WorkFlow] = Field(default_factory=dict)
+    """The registry of events and workflows."""
+
+    toolboxes: Set[ToolBox] = Field(default_factory=set)
+    """Collection of tools available to this role."""
+
+    def model_post_init(self, __context: Any) -> None:
+        """Initialize the role by resolving configurations and registering workflows.
+
+        Args:
+            __context: The context used for initialization
+        """
+        self.resolve_configuration().register_workflows()
+
+    def register_workflows(self) -> Self:
+        """Register each workflow in the registry to its corresponding event in the event bus.
+
+        Returns:
+            Self: The role instance for method chaining
+        """
+        for event, workflow in self.registry.items():
+            logger.debug(
+                f"Registering workflow: `{workflow.name}` for event: `{Event.instantiate_from(event).collapse()}`"
+            )
+            env.on(event, workflow.serve)
+        return self
+
+    def resolve_configuration(self) -> Self:
+        """Apply role-level configuration to all workflows in the registry.
+
+        This includes setting up fallback configurations, injecting personality traits,
+        and providing tool access to workflows and their steps.
+
+        Returns:
+            Self: The role instance for method chaining
+        """
+        for workflow in self.registry.values():
+            logger.debug(f"Resolving config for workflow: `{workflow.name}`")
+            (
+                workflow.fallback_to(self)
+                .steps_fallback_to_self()
+                .inject_personality(self.briefing)
+                .supply_tools_from(self)
+                .steps_supply_tools_from_self()
+            )
+
+        return self