fabricatio 0.2.9.dev4__cp312-cp312-win_amd64.whl → 0.2.10.dev0__cp312-cp312-win_amd64.whl

fabricatio/actions/article.py

@@ -11,7 +11,6 @@ from fabricatio.capabilities.propose import Propose
 from fabricatio.fs import safe_text_read
 from fabricatio.journal import logger
 from fabricatio.models.action import Action
-from fabricatio.models.extra.article_base import SubSectionBase
 from fabricatio.models.extra.article_essence import ArticleEssence
 from fabricatio.models.extra.article_main import Article
 from fabricatio.models.extra.article_outline import ArticleOutline
@@ -142,11 +141,17 @@ class GenerateInitialOutline(Action, Propose):
         article_proposal: ArticleProposal,
         **_,
     ) -> Optional[ArticleOutline]:
+        raw_outline = await self.aask(
+            f"{(article_proposal.as_prompt())}\n\nNote that you should use `{article_proposal.language}` to write the `ArticleOutline`\n"
+            f"Design each chapter of a proper and academic and ready for release manner.\n"
+            f"You Must make sure every chapter have sections, and every section have subsections.\n"
+            f"Make the chapter and sections and subsections bing divided into a specific enough article component.",
+        )
+
         return ok(
             await self.propose(
                 ArticleOutline,
-                f"{(article_proposal.as_prompt())}\n\nNote that you should use `{article_proposal.language}` to write the `ArticleOutline`\n"
-                f"You Must make sure every chapter have sections, and every section have subsections.",
+                f"{raw_outline}\n\n\n\noutline provided above is the outline i need to extract to a JSON,",
             ),
             "Could not generate the initial outline.",
         ).update_ref(article_proposal)
@@ -178,7 +183,7 @@ class FixIntrospectedErrors(Action, Censor):
                 await self.censor_obj(
                     article_outline,
                     ruleset=ok(intro_fix_ruleset or self.ruleset, "No ruleset provided"),
-                    reference=f"{article_outline.as_prompt()}\n # Fatal Error of the Original Article Outline\n{pack}",
+                    reference=f"{article_outline.display()}\n # Fatal Error of the Original Article Outline\n{pack}",
                 ),
                 "Could not correct the component.",
             ).update_ref(origin)
@@ -191,107 +196,6 @@ class FixIntrospectedErrors(Action, Censor):
         return article_outline
 
 
-class FixIllegalReferences(Action, Censor):
-    """Fix illegal references in the article outline."""
-
-    output_key: str = "illegal_references_fixed_outline"
-    """The key of the output data."""
-
-    ruleset: Optional[RuleSet] = None
-    """Ruleset to use to fix the illegal references."""
-    max_error_count: Optional[int] = None
-    """The maximum number of errors to fix."""
-
-    async def _execute(
-        self,
-        article_outline: ArticleOutline,
-        ref_fix_ruleset: Optional[RuleSet] = None,
-        **_,
-    ) -> Optional[ArticleOutline]:
-        counter = 0
-        while pack := article_outline.find_illegal_ref(gather_identical=True):
-            logger.info(f"Found {counter}th illegal references")
-            ref_seq, err = ok(pack)
-            logger.warning(f"Found illegal referring error: {err}")
-            new = ok(
-                await self.censor_obj(
-                    ref_seq[0],
-                    ruleset=ok(ref_fix_ruleset or self.ruleset, "No ruleset provided"),
-                    reference=f"{article_outline.as_prompt()}\n# Some Basic errors found that need to be fixed\n{err}",
-                ),
-                "Could not correct the component",
-            )
-            for r in ref_seq:
-                r.update_from(new)
-            if self.max_error_count and counter > self.max_error_count:
-                logger.warning("Max error count reached, stopping.")
-                break
-            counter += 1
-
-        return article_outline
-
-
-class TweakOutlineForwardRef(Action, Censor):
-    """Tweak the forward references in the article outline.
-
-    Ensures that the conclusions of the current chapter effectively support the analysis of subsequent chapters.
-    """
-
-    output_key: str = "article_outline_fw_ref_checked"
-    ruleset: Optional[RuleSet] = None
-    """Ruleset to use to fix the illegal references."""
-
-    async def _execute(
-        self, article_outline: ArticleOutline, ref_twk_ruleset: Optional[RuleSet] = None, **cxt
-    ) -> ArticleOutline:
-        return await self._inner(
-            article_outline,
-            ruleset=ok(ref_twk_ruleset or self.ruleset, "No ruleset provided"),
-            field_name="support_to",
-        )
-
-    async def _inner(self, article_outline: ArticleOutline, ruleset: RuleSet, field_name: str) -> ArticleOutline:
-        await gather(
-            *[self._loop(a[-1], article_outline, field_name, ruleset) for a in article_outline.iter_subsections()],
-        )
-
-        return article_outline
-
-    async def _loop(
-        self, a: SubSectionBase, article_outline: ArticleOutline, field_name: str, ruleset: RuleSet
-    ) -> None:
-        if judge := await self.evidently_judge(
-            f"{article_outline.as_prompt()}\n\n{a.display()}\n"
-            f"Does the `{a.__class__.__name__}`'s `{field_name}` field need to be extended or tweaked?"
-        ):
-            await self.censor_obj_inplace(
-                a,
-                ruleset=ruleset,
-                reference=f"{article_outline.as_prompt()}\n"
-                f"The Article component titled `{a.title}` whose `{field_name}` field needs to be extended or tweaked.\n"
-                f"# Judgement\n{judge.display()}",
-            )
-
-
-class TweakOutlineBackwardRef(TweakOutlineForwardRef):
-    """Tweak the backward references in the article outline.
-
-    Ensures that the prerequisites of the current chapter are correctly referenced in the `depend_on` field.
-    """
-
-    output_key: str = "article_outline_bw_ref_checked"
-    ruleset: Optional[RuleSet] = None
-
-    async def _execute(
-        self, article_outline: ArticleOutline, ref_twk_ruleset: Optional[RuleSet] = None, **cxt
-    ) -> ArticleOutline:
-        return await self._inner(
-            article_outline,
-            ruleset=ok(ref_twk_ruleset or self.ruleset, "No ruleset provided"),
-            field_name="depend_on",
-        )
-
-
 class GenerateArticle(Action, Censor):
     """Generate the article based on the outline."""
 

fabricatio/capabilities/check.py

@@ -104,7 +104,8 @@ class Check(AdvancedJudge, Propose):
             - Proposes Improvement only when violation is confirmed
         """
         if judge := await self.evidently_judge(
-            f"# Content to exam\n{input_text}\n\n# Rule Must to follow\n{rule.display()}\nDoes `Content to exam` provided above violate the `Rule Must to follow` provided above?",
+            f"# Content to exam\n{input_text}\n\n# Rule Must to follow\n{rule.display()}\nDoes `Content to exam` provided above violate the `{rule.name}` provided above?"
+            f"should I take some measure to fix that violation? true for I do need, false for I don't need.",
             **override_kwargs(kwargs, default=None),
         ):
             logger.info(f"Rule `{rule.name}` violated: \n{judge.display()}")

fabricatio/capabilities/rag.py

@@ -3,28 +3,22 @@
 try:
     from pymilvus import MilvusClient
 except ImportError as e:
-    raise RuntimeError("pymilvus is not installed. Have you installed `fabricatio[rag]` instead of `fabricatio`?") from e
+    raise RuntimeError(
+        "pymilvus is not installed. Have you installed `fabricatio[rag]` instead of `fabricatio`?"
+    ) from e
 from functools import lru_cache
 from operator import itemgetter
-from os import PathLike
-from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Self, Union, Unpack, cast, overload
+from typing import List, Optional, Self, Type, Unpack
 
 from more_itertools.recipes import flatten, unique
 from pydantic import Field, PrivateAttr
 
 from fabricatio.config import configs
 from fabricatio.journal import logger
-from fabricatio.models.kwargs_types import (
-    ChooseKwargs,
-    CollectionConfigKwargs,
-    EmbeddingKwargs,
-    FetchKwargs,
-    LLMKwargs,
-    RetrievalKwargs,
-)
+from fabricatio.models.adv_kwargs_types import CollectionConfigKwargs, FetchKwargs
+from fabricatio.models.extra.rag import MilvusDataBase
+from fabricatio.models.kwargs_types import ChooseKwargs
 from fabricatio.models.usages import EmbeddingUsage
-from fabricatio.models.utils import MilvusData
 from fabricatio.rust_instances import TEMPLATE_MANAGER
 from fabricatio.utils import ok
 
@@ -78,40 +72,6 @@ class RAG(EmbeddingUsage):
             raise RuntimeError("Client is not initialized. Have you called `self.init_client()`?")
         return self
 
-    @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)
-        ]
-
     def view(
         self, collection_name: Optional[str], create: bool = False, **kwargs: Unpack[CollectionConfigKwargs]
     ) -> Self:
@@ -152,29 +112,27 @@ class RAG(EmbeddingUsage):
         Returns:
             str: The name of the collection being viewed.
         """
-        if self.target_collection is None:
-            raise RuntimeError("No collection is being viewed. Have you called `self.view()`?")
-        return self.target_collection
+        return ok(self.target_collection, "No collection is being viewed. Have you called `self.view()`?")
 
-    def add_document[D: Union[Dict[str, Any], MilvusData]](
-        self, data: D | List[D], collection_name: Optional[str] = None, flush: bool = False
+    async def add_document[D: MilvusDataBase](
+        self, data: List[D] | D, collection_name: Optional[str] = None, flush: bool = False
     ) -> Self:
         """Adds a document to the specified collection.
 
         Args:
-            data (Union[Dict[str, Any], MilvusData] | List[Union[Dict[str, Any], MilvusData]]): The data to be added to the collection.
+            data (Union[Dict[str, Any], MilvusDataBase] | List[Union[Dict[str, Any], MilvusDataBase]]): The data to be added to the collection.
             collection_name (Optional[str]): The name of the collection. If not provided, the currently viewed collection is used.
             flush (bool): Whether to flush the collection after insertion.
 
         Returns:
             Self: The current instance, allowing for method chaining.
         """
-        if isinstance(data, MilvusData):
-            prepared_data = data.prepare_insertion()
-        elif isinstance(data, list):
-            prepared_data = [d.prepare_insertion() if isinstance(d, MilvusData) else d for d in data]
-        else:
-            raise TypeError(f"Expected MilvusData or list of MilvusData, got {type(data)}")
+        if isinstance(data, MilvusDataBase):
+            data = [data]
+
+        data_vec = await self.vectorize([d.to_vectorize for d in data])
+        prepared_data = [d.prepare_insertion(vec) for d, vec in zip(data, data_vec, strict=True)]
+
         c_name = collection_name or self.safe_target_collection
         self.check_client().client.insert(c_name, prepared_data)
 
@@ -183,84 +141,33 @@ class RAG(EmbeddingUsage):
             self.client.flush(c_name)
         return self
 
-    async def consume_file(
-        self,
-        source: List[PathLike] | PathLike,
-        reader: Callable[[PathLike], str] = lambda path: Path(path).read_text(encoding="utf-8"),
-        collection_name: Optional[str] = None,
-    ) -> Self:
-        """Consume a file and add its content to the collection.
-
-        Args:
-            source (PathLike): The path to the file to be consumed.
-            reader (Callable[[PathLike], MilvusData]): The reader function to read the file.
-            collection_name (Optional[str]): The name of the collection. If not provided, the currently viewed collection is used.
-
-        Returns:
-            Self: The current instance, allowing for method chaining.
-        """
-        if not isinstance(source, list):
-            source = [source]
-        return await self.consume_string([reader(s) for s in source], collection_name)
-
-    async def consume_string(self, text: List[str] | str, collection_name: Optional[str] = None) -> Self:
-        """Consume a string and add it to the collection.
-
-        Args:
-            text (List[str] | str): The text to be added to the collection.
-            collection_name (Optional[str]): The name of the collection. If not provided, the currently viewed collection is used.
-
-        Returns:
-            Self: The current instance, allowing for method chaining.
-        """
-        self.add_document(await self.pack(text), collection_name or self.safe_target_collection, flush=True)
-        return self
-
-    @overload
-    async def afetch_document[V: (int, str, float, bytes)](
+    async def afetch_document[D: MilvusDataBase](
         self,
         vecs: List[List[float]],
-        desired_fields: List[str],
+        document_model: Type[D],
         collection_name: Optional[str] = None,
         similarity_threshold: float = 0.37,
         result_per_query: int = 10,
-    ) -> List[Dict[str, V]]: ...
-
-    @overload
-    async def afetch_document[V: (int, str, float, bytes)](
-        self,
-        vecs: List[List[float]],
-        desired_fields: str,
-        collection_name: Optional[str] = None,
-        similarity_threshold: float = 0.37,
-        result_per_query: int = 10,
-    ) -> List[V]: ...
-    async def afetch_document[V: (int, str, float, bytes)](
-        self,
-        vecs: List[List[float]],
-        desired_fields: List[str] | str,
-        collection_name: Optional[str] = None,
-        similarity_threshold: float = 0.37,
-        result_per_query: int = 10,
-    ) -> List[Dict[str, Any]] | List[V]:
-        """Fetch data from the collection.
+    ) -> List[D]:
+        """Asynchronously fetches documents from a Milvus database based on input vectors.
 
         Args:
-            vecs (List[List[float]]): The vectors to search for.
-            desired_fields (List[str] | str): The fields to retrieve.
-            collection_name (Optional[str]): The name of the collection. If not provided, the currently viewed collection is used.
-            similarity_threshold (float): The threshold for similarity, only results above this threshold will be returned.
-            result_per_query (int): The number of results to return per query.
+           vecs (List[List[float]]): A list of vectors to search for in the database.
+           document_model (Type[D]): The model class used to convert fetched data into document objects.
+           collection_name (Optional[str]): The name of the collection to search within.
+                                             If None, the currently viewed collection is used.
+           similarity_threshold (float): The similarity threshold for vector search. Defaults to 0.37.
+           result_per_query (int): The maximum number of results to return per query. Defaults to 10.
 
         Returns:
-            List[Dict[str, Any]] | List[Any]: The retrieved data.
+           List[D]: A list of document objects created from the fetched data.
         """
         # Step 1: Search for vectors
         search_results = self.check_client().client.search(
             collection_name or self.safe_target_collection,
             vecs,
             search_params={"radius": similarity_threshold},
-            output_fields=desired_fields if isinstance(desired_fields, list) else [desired_fields],
+            output_fields=list(document_model.model_fields),
             limit=result_per_query,
         )
 
@@ -270,20 +177,20 @@ class RAG(EmbeddingUsage):
         # Step 3: Sort by distance (descending)
         sorted_results = sorted(unique_results, key=itemgetter("distance"), reverse=True)
 
-        logger.debug(f"Searched similarities: {[t['distance'] for t in sorted_results]}")
+        logger.debug(
+            f"Fetched {len(sorted_results)} document,searched similarities: {[t['distance'] for t in sorted_results]}"
+        )
         # Step 4: Extract the entities
         resp = [result["entity"] for result in sorted_results]
 
-        if isinstance(desired_fields, list):
-            return resp
-        return [r.get(desired_fields) for r in resp]  # extract the single field as list
+        return document_model.from_sequence(resp)
 
-    async def aretrieve(
+    async def aretrieve[D: MilvusDataBase](
         self,
         query: List[str] | str,
         final_limit: int = 20,
-        **kwargs: Unpack[FetchKwargs],
-    ) -> List[str]:
+        **kwargs: Unpack[FetchKwargs[D]],
+    ) -> List[D]:
         """Retrieve data from the collection.
 
         Args:
@@ -292,82 +199,17 @@ class RAG(EmbeddingUsage):
             **kwargs (Unpack[FetchKwargs]): Additional keyword arguments for retrieval.
 
         Returns:
-            List[str]: A list of strings containing the retrieved data.
+            List[D]: A list of document objects created from the retrieved data.
         """
         if isinstance(query, str):
             query = [query]
-        return cast(
-            "List[str]",
+        return (
             await self.afetch_document(
                 vecs=(await self.vectorize(query)),
-                desired_fields="text",
                 **kwargs,
-            ),
+            )
         )[:final_limit]
 
-    async def aretrieve_compact(
-        self,
-        query: List[str] | str,
-        **kwargs: Unpack[RetrievalKwargs],
-    ) -> str:
-        """Retrieve data from the collection and format it for display.
-
-        Args:
-            query (List[str] | str): The query to be used for retrieval.
-            **kwargs (Unpack[RetrievalKwargs]): Additional keyword arguments for retrieval.
-
-        Returns:
-            str: A formatted string containing the retrieved data.
-        """
-        return TEMPLATE_MANAGER.render_template(
-            configs.templates.retrieved_display_template, {"docs": (await self.aretrieve(query, **kwargs))}
-        )
-
-    async def aask_retrieved(
-        self,
-        question: str,
-        query: Optional[List[str] | str] = None,
-        collection_name: Optional[str] = None,
-        extra_system_message: str = "",
-        result_per_query: int = 10,
-        final_limit: int = 20,
-        similarity_threshold: float = 0.37,
-        **kwargs: Unpack[LLMKwargs],
-    ) -> str:
-        """Asks a question by retrieving relevant documents based on the provided query.
-
-        This method performs document retrieval using the given query, then asks the
-        specified question using the retrieved documents as context.
-
-        Args:
-            question (str): The question to be asked.
-            query (List[str] | str): The query or list of queries used for document retrieval.
-            collection_name (Optional[str]): The name of the collection to retrieve documents from.
-                                              If not provided, the currently viewed collection is used.
-            extra_system_message (str): An additional system message to be included in the prompt.
-            result_per_query (int): The number of results to return per query. Default is 10.
-            final_limit (int): The maximum number of retrieved documents to consider. Default is 20.
-            similarity_threshold (float): The threshold for similarity, only results above this threshold will be returned.
-            **kwargs (Unpack[LLMKwargs]): Additional keyword arguments passed to the underlying `aask` method.
-
-        Returns:
-            str: A string response generated after asking with the context of retrieved documents.
-        """
-        rendered = await self.aretrieve_compact(
-            query or question,
-            final_limit=final_limit,
-            collection_name=collection_name,
-            result_per_query=result_per_query,
-            similarity_threshold=similarity_threshold,
-        )
-
-        logger.debug(f"Retrieved Documents: \n{rendered}")
-        return await self.aask(
-            question,
-            f"{rendered}\n\n{extra_system_message}",
-            **kwargs,
-        )
-
     async def arefined_query(self, question: List[str] | str, **kwargs: Unpack[ChooseKwargs]) -> Optional[List[str]]:
         """Refines the given question using a template.
 
@@ -385,38 +227,3 @@ class RAG(EmbeddingUsage):
             ),
             **kwargs,
         )
-
-    async def aask_refined(
-        self,
-        question: str,
-        collection_name: Optional[str] = None,
-        extra_system_message: str = "",
-        result_per_query: int = 10,
-        final_limit: int = 20,
-        similarity_threshold: float = 0.37,
-        **kwargs: Unpack[LLMKwargs],
-    ) -> str:
-        """Asks a question using a refined query based on the provided question.
-
-        Args:
-            question (str): The question to be asked.
-            collection_name (Optional[str]): The name of the collection to retrieve documents from.
-            extra_system_message (str): An additional system message to be included in the prompt.
-            result_per_query (int): The number of results to return per query. Default is 10.
-            final_limit (int): The maximum number of retrieved documents to consider. Default is 20.
-            similarity_threshold (float): The threshold for similarity, only results above this threshold will be returned.
-            **kwargs (Unpack[LLMKwargs]): Additional keyword arguments passed to the underlying `aask` method.
-
-        Returns:
-            str: A string response generated after asking with the refined question.
-        """
-        return await self.aask_retrieved(
-            question,
-            await self.arefined_query(question, **kwargs),
-            collection_name=collection_name,
-            extra_system_message=extra_system_message,
-            result_per_query=result_per_query,
-            final_limit=final_limit,
-            similarity_threshold=similarity_threshold,
-            **kwargs,
-        )

fabricatio/constants.py

@@ -0,0 +1,20 @@
+"""A module containing constants used throughout the library."""
+from enum import StrEnum
+
+
+class TaskStatus(StrEnum):
+    """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/decorators.py

@@ -2,6 +2,7 @@
 
 from asyncio import iscoroutinefunction
 from functools import wraps
+from importlib.util import find_spec
 from inspect import signature
 from shutil import which
 from types import ModuleType
@@ -209,3 +210,25 @@ def logging_exec_time[**P, R](func: Callable[P, R]) -> Callable[P, R]:
         return result
 
     return _wrapper
+
+
+def precheck_package[**P, R](package_name: str, msg: str) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """Check if a package exists in the current environment.
+
+    Args:
+        package_name (str): The name of the package to check.
+        msg (str): The message to display if the package is not found.
+
+    Returns:
+        bool: True if the package exists, False otherwise.
+    """
+
+    def _wrapper(func: Callable[P, R]) -> Callable[P, R]:
+        def _inner(*args: P.args, **kwargs: P.kwargs) -> R:
+            if find_spec(package_name):
+                return func(*args, **kwargs)
+            raise RuntimeError(msg)
+
+        return _inner
+
+    return _wrapper

fabricatio/models/adv_kwargs_types.py

@@ -1,5 +1,10 @@
 """A module containing kwargs types for content correction and checking operations."""
+
+from importlib.util import find_spec
+from typing import Required, Type, TypedDict
+
 from fabricatio.models.extra.problem import Improvement
+from fabricatio.models.extra.rag import MilvusDataBase
 from fabricatio.models.extra.rule import RuleSet
 from fabricatio.models.generic import SketchedAble
 from fabricatio.models.kwargs_types import ReferencedKwargs
@@ -23,3 +28,40 @@ class CheckKwargs(ReferencedKwargs[Improvement], total=False):
     """
 
     ruleset: RuleSet
+
+
+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[D: MilvusDataBase](TypedDict, total=False):
+        """Arguments for fetching data from vector collections.
+
+        Controls how data is retrieved from vector databases, including filtering
+        and result limiting parameters.
+        """
+
+        document_model: Required[Type[D]]
+        collection_name: str | None
+        similarity_threshold: float
+        result_per_query: int
+
+    class RetrievalKwargs(FetchKwargs, total=False):
+        """Arguments for retrieval operations."""
+
+        final_limit: int

fabricatio/models/events.py

@@ -3,7 +3,7 @@
 from typing import List, Self, Union
 
 from fabricatio.config import configs
-from fabricatio.models.utils import TaskStatus
+from fabricatio.constants import TaskStatus
 from pydantic import BaseModel, ConfigDict, Field
 
 type EventLike = Union[str, List[str], "Event"]
@@ -77,23 +77,23 @@ class Event(BaseModel):
 
     def push_pending(self) -> Self:
         """Push a pending segment to the event."""
-        return self.push(TaskStatus.Pending.value)
+        return self.push(TaskStatus.Pending)
 
     def push_running(self) -> Self:
         """Push a running segment to the event."""
-        return self.push(TaskStatus.Running.value)
+        return self.push(TaskStatus.Running)
 
     def push_finished(self) -> Self:
         """Push a finished segment to the event."""
-        return self.push(TaskStatus.Finished.value)
+        return self.push(TaskStatus.Finished)
 
     def push_failed(self) -> Self:
         """Push a failed segment to the event."""
-        return self.push(TaskStatus.Failed.value)
+        return self.push(TaskStatus.Failed)
 
     def push_cancelled(self) -> Self:
         """Push a cancelled segment to the event."""
-        return self.push(TaskStatus.Cancelled.value)
+        return self.push(TaskStatus.Cancelled)
 
     def pop(self) -> str:
         """Pop a segment from the event."""