fabricatio 0.2.6.dev2__cp312-cp312-win_amd64.whl → 0.2.7.dev2__cp312-cp312-win_amd64.whl

fabricatio/models/generic.py

@@ -10,6 +10,7 @@ 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.models.utils import ok
 from fabricatio.parser import JsonCapture
 from pydantic import (
     BaseModel,
@@ -19,6 +20,7 @@ from pydantic import (
     NonNegativeFloat,
     PositiveFloat,
     PositiveInt,
+    PrivateAttr,
     SecretStr,
 )
 
@@ -63,6 +65,44 @@ class Described(Base):
     """The description of the object."""
 
 
+class AsPrompt(Base):
+    """Class that provides a method to generate a prompt from the model."""
+
+    @final
+    def as_prompt(self) -> str:
+        """Generate a prompt from the model.
+
+        Returns:
+            str: The generated prompt.
+        """
+        return TEMPLATE_MANAGER.render_template(
+            configs.templates.as_prompt_template,
+            self._as_prompt_inner(),
+        )
+
+    @abstractmethod
+    def _as_prompt_inner(self) -> Dict[str, str]: ...
+
+
+class WithRef[T](Base):
+    """Class that provides a reference to another object."""
+
+    _reference: Optional[T] = PrivateAttr(None)
+
+    @property
+    def referenced(self) -> T:
+        """Get the referenced object."""
+        return ok(self._reference, "_reference is None")
+
+    def update_ref[S](self: S, reference: T | S) -> S:  # noqa: PYI019
+        """Update the reference of the object."""
+        if isinstance(reference, self.__class__):
+            self._reference = reference.referenced
+        else:
+            self._reference = reference
+        return self
+
+
 class WithBriefing(Named, Described):
     """Class that provides a briefing based on the name and description."""
 
@@ -75,17 +115,26 @@ class WithBriefing(Named, Described):
         """
         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.
+    def _prepend_inner(self, input_text: str) -> str:
+        return f"# your personal briefing: \n{self.briefing}\n{input_text}"
+
+    def prepend_sys_msg[D: (Dict[str, Any], str)](self, system_msg_like: D = "") -> Dict[str, Any]:
+        """Prepend the system message with the briefing.
 
         Args:
-            kwargs (Dict[str, Any]): The keyword arguments to modify.
+            system_msg_like (str | dict): The system message or a dictionary containing the system message.
 
         Returns:
-            Dict[str, Any]: The modified keyword arguments.
+            dict: The system message with the briefing prepended.
         """
-        kwargs["system_message"] = f"# your personal briefing: \n{self.briefing}\n" + kwargs.get("system_message", "")
-        return kwargs
+        match system_msg_like:
+            case dict(d):
+                d["system_message"] = self._prepend_inner(d.get("system_message", ""))
+                return d
+            case str(s):
+                return {"system_message": self._prepend_inner(s)}
+            case _:
+                raise TypeError(f"{system_msg_like} is not a dict or str")
 
 
 class WithFormatedJsonSchema(Base):
@@ -158,19 +207,17 @@ class InstantiateFromString(Base):
 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.
         """
+        return self.model_dump_json()
 
     def finalized_dump_to(self, path: str | Path) -> Self:
         """Finalize the dump of the object to a file.
@@ -185,6 +232,10 @@ class FinalizedDumpAble(Base):
         return self
 
 
+class CensoredAble(ProposedAble, FinalizedDumpAble):
+    """Class that provides a method to censor the object."""
+
+
 class WithDependency(Base):
     """Class that manages file dependencies."""
 

fabricatio/models/kwargs_types.py

@@ -1,19 +1,29 @@
 """This module contains the types for the keyword arguments of the methods in the models module."""
 
-from typing import Any, TypedDict
+from importlib.util import find_spec
+from typing import Any, Dict, 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 CollectionSimpleConfigKwargs(TypedDict, total=False):
-    """Configuration parameters for a vector collection.
+    class CollectionConfigKwargs(TypedDict, total=False):
+        """Configuration parameters for a vector collection.
 
-    These arguments are typically used when configuring connections to vector databases.
-    """
+        These arguments are typically used when configuring connections to vector databases.
+        """
 
-    dimension: int | None
-    timeout: float
+        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):
@@ -28,6 +38,12 @@ class FetchKwargs(TypedDict, total=False):
     result_per_query: int
 
 
+class RetrievalKwargs(FetchKwargs, total=False):
+    """Arguments for retrieval operations."""
+
+    final_limit: int
+
+
 class EmbeddingKwargs(TypedDict, total=False):
     """Configuration parameters for text embedding operations.
 
@@ -81,18 +97,25 @@ class ValidateKwargs[T](GenerateKwargs, total=False):
 
     default: T
     max_validations: int
+    co_extractor: GenerateKwargs
+
+
+class ReviewInnerKwargs[T](ValidateKwargs[T], total=False):
+    """Arguments for content review operations."""
+
+    criteria: set[str]
 
 
 # noinspection PyTypedDict
-class ReviewKwargs[T](ValidateKwargs[T], total=False):
+class ReviewKwargs[T](ReviewInnerKwargs[T], total=False):
     """Arguments for content review operations.
 
     Extends GenerateKwargs with parameters for evaluating content against
     specific topics and review criteria.
     """
 
-    topic: str
-    criteria: set[str]
+    rating_manual: Dict[str, str]
+    topic: Required[str]
 
 
 class CorrectKwargs[T](ReviewKwargs[T], total=False):
@@ -106,6 +129,13 @@ class CorrectKwargs[T](ReviewKwargs[T], total=False):
     supervisor_check: bool
 
 
+class CensoredCorrectKwargs[T](ReviewInnerKwargs[T], total=False):
+    """Arguments for content censorship operations."""
+
+    reference: str
+    supervisor_check: bool
+
+
 # noinspection PyTypedDict
 class ChooseKwargs[T](ValidateKwargs[T], total=False):
     """Arguments for selection operations.

fabricatio/models/role.py

@@ -1,4 +1,4 @@
-"""Module that contains the Role class."""
+"""Module that contains the Role class for managing workflows and their event registrations."""
 
 from typing import Any, Self, Set
 
@@ -13,19 +13,36 @@ from pydantic import Field
 
 
 class Role(ProposeTask, HandleTask, Correct):
-    """Class that represents a role with a registry of events and workflows."""
+    """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."""
+    """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:
-        """Register the workflows in the role to the event bus."""
+        """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 the workflows in the role to the event bus."""
+        """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()}`"
@@ -34,7 +51,14 @@ class Role(ProposeTask, HandleTask, Correct):
         return self
 
     def resolve_configuration(self) -> Self:
-        """Resolve the configuration of the role."""
+        """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}`")
             (

fabricatio/models/tool.py

@@ -136,7 +136,9 @@ class ToolExecutor(BaseModel):
 
     def inject_tools[M: ModuleType](self, module: Optional[M] = None) -> M:
         """Inject the tools into the provided module or default."""
-        module = module or cast(M, module_from_spec(spec=ModuleSpec(name=configs.toolbox.tool_module_name, loader=None)))
+        module = module or cast(
+            "M", module_from_spec(spec=ModuleSpec(name=configs.toolbox.tool_module_name, loader=None))
+        )
         for tool in self.candidates:
             logger.debug(f"Injecting tool: {tool.name}")
             setattr(module, tool.name, tool.invoke)
@@ -144,7 +146,9 @@ class ToolExecutor(BaseModel):
 
     def inject_data[M: ModuleType](self, module: Optional[M] = None) -> M:
         """Inject the data into the provided module or default."""
-        module = module or cast(M,module_from_spec(spec=ModuleSpec(name=configs.toolbox.data_module_name, loader=None)))
+        module = module or cast(
+            'M', module_from_spec(spec=ModuleSpec(name=configs.toolbox.data_module_name, loader=None))
+        )
         for key, value in self.data.items():
             logger.debug(f"Injecting data: {key}")
             setattr(module, key, value)