fabricatio 0.2.0.dev4__cp312-cp312-win_amd64.whl

fabricatio/models/task.py

@@ -0,0 +1,283 @@
+"""This module defines the `Task` class, which represents a task with a status and output.
+
+It includes methods to manage the task's lifecycle, such as starting, finishing, cancelling, and failing the task.
+"""
+
+from asyncio import Queue
+from enum import Enum
+from typing import Any, List, Optional, Self
+
+from fabricatio.core import env
+from fabricatio.journal import logger
+from fabricatio.models.events import Event, EventLike
+from fabricatio.models.generic import LLMUsage, WithBriefing, WithDependency, WithJsonExample
+from fabricatio.parser import JsonCapture
+from pydantic import Field, PrivateAttr, ValidationError
+
+
+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"
+
+
+class Task[T](WithBriefing, WithJsonExample, WithDependency):
+    """A class representing a task with a status and output.
+
+    Attributes:
+        name (str): The name of the task.
+        description (str): The description of the task.
+        goal (str): The goal of the task.
+        dependencies (List[str]): The file dependencies of the task, a list of file paths.
+        namespace (List[str]): The namespace of the task, a list of namespace segment, as string.
+    """
+
+    name: str = Field(...)
+    """The name of the task."""
+
+    description: str = Field(default="")
+    """The description of the task."""
+
+    goal: str = Field(default="")
+    """The goal of the task."""
+
+    namespace: List[str] = Field(default_factory=list)
+    """The namespace of the task, a list of namespace segment, as string."""
+
+    _output: Queue = PrivateAttr(default_factory=lambda: Queue(maxsize=1))
+    """The output queue of the task."""
+
+    _status: TaskStatus = PrivateAttr(default=TaskStatus.Pending)
+    """The status of the task."""
+
+    _namespace: Event = PrivateAttr(default_factory=Event)
+    """The namespace of the task as an event, which is generated from the namespace list."""
+
+    def model_post_init(self, __context: Any) -> None:
+        """Initialize the task with a namespace event."""
+        self._namespace.segments.extend(self.namespace)
+
+    def move_to(self, new_namespace: EventLike) -> Self:
+        """Move the task to a new namespace.
+
+        Args:
+            new_namespace (List[str]): The new namespace to move the task to.
+
+        Returns:
+            Task: The moved instance of the `Task` class.
+        """
+        self.namespace = new_namespace
+        self._namespace.clear().concat(new_namespace)
+        return self
+
+    @classmethod
+    def simple_task(cls, name: str, goal: str, description: str) -> Self:
+        """Create a simple task with a name, goal, and description.
+
+        Args:
+            name (str): The name of the task.
+            goal (str): The goal of the task.
+            description (str): The description of the task.
+
+        Returns:
+            Task: A new instance of the `Task` class.
+        """
+        return cls(name=name, goal=goal, description=description)
+
+    def update_task(self, goal: Optional[str] = None, description: Optional[str] = None) -> Self:
+        """Update the goal and description of the task.
+
+        Args:
+            goal (str, optional): The new goal of the task.
+            description (str, optional): The new description of the task.
+
+        Returns:
+            Task: The updated instance of the `Task` class.
+        """
+        if goal:
+            self.goal = goal
+        if description:
+            self.description = description
+        return self
+
+    async def get_output(self) -> T:
+        """Get the output of the task.
+
+        Returns:
+            T: The output of the task.
+        """
+        logger.debug(f"Getting output for task {self.name}")
+        return await self._output.get()
+
+    def status_label(self, status: TaskStatus) -> str:
+        """Return a formatted status label for the task.
+
+        Args:
+            status (TaskStatus): The status of the task.
+
+        Returns:
+            str: The formatted status label.
+        """
+        return self._namespace.derive(self.name).push(status.value).collapse()
+
+    @property
+    def pending_label(self) -> str:
+        """Return the pending status label for the task.
+
+        Returns:
+            str: The pending status label.
+        """
+        return self.status_label(TaskStatus.Pending)
+
+    @property
+    def running_label(self) -> str:
+        """Return the running status label for the task.
+
+        Returns:
+            str: The running status label.
+        """
+        return self.status_label(TaskStatus.Running)
+
+    @property
+    def finished_label(self) -> str:
+        """Return the finished status label for the task.
+
+        Returns:
+            str: The finished status label.
+        """
+        return self.status_label(TaskStatus.Finished)
+
+    @property
+    def failed_label(self) -> str:
+        """Return the failed status label for the task.
+
+        Returns:
+            str: The failed status label.
+        """
+        return self.status_label(TaskStatus.Failed)
+
+    @property
+    def cancelled_label(self) -> str:
+        """Return the cancelled status label for the task.
+
+        Returns:
+            str: The cancelled status label.
+        """
+        return self.status_label(TaskStatus.Cancelled)
+
+    async def finish(self, output: T) -> Self:
+        """Mark the task as finished and set the output.
+
+        Args:
+            output (T): The output of the task.
+
+        Returns:
+            Task: The finished instance of the `Task` class.
+        """
+        logger.info(f"Finishing task {self.name}")
+        self._status = TaskStatus.Finished
+        await self._output.put(output)
+        logger.debug(f"Output set for task {self.name}")
+        await env.emit_async(self.finished_label, self)
+        logger.debug(f"Emitted finished event for task {self.name}")
+        return self
+
+    async def start(self) -> Self:
+        """Mark the task as running.
+
+        Returns:
+            Task: The running instance of the `Task` class.
+        """
+        logger.info(f"Starting task {self.name}")
+        self._status = TaskStatus.Running
+        await env.emit_async(self.running_label, self)
+        return self
+
+    async def cancel(self) -> Self:
+        """Mark the task as cancelled.
+
+        Returns:
+            Task: The cancelled instance of the `Task` class.
+        """
+        self._status = TaskStatus.Cancelled
+        await env.emit_async(self.cancelled_label, self)
+        return self
+
+    async def fail(self) -> Self:
+        """Mark the task as failed.
+
+        Returns:
+            Task: The failed instance of the `Task` class.
+        """
+        logger.error(f"Task {self.name} failed")
+        self._status = TaskStatus.Failed
+        await env.emit_async(self.failed_label, self)
+        return self
+
+    async def publish(self) -> Self:
+        """Publish the task to the event bus.
+
+        Returns:
+            Task: The published instance of the `Task` class
+        """
+        logger.info(f"Publishing task {self.name}")
+        await env.emit_async(self.pending_label, self)
+        return self
+
+    async def delegate(self) -> T:
+        """Delegate the task to the event bus and wait for the output.
+
+        Returns:
+            T: The output of the task
+        """
+        logger.info(f"Delegating task {self.name}")
+        await env.emit_async(self.pending_label, self)
+        return await self.get_output()
+
+    @property
+    def briefing(self) -> str:
+        """Return a briefing of the task including its goal.
+
+        Returns:
+            str: The briefing of the task.
+        """
+        return f"{super().briefing}\n{self.goal}"
+
+
+class ProposeTask(LLMUsage, WithBriefing):
+    """A class that proposes a task based on a prompt."""
+
+    async def propose(self, prompt: str) -> Task:
+        """Propose a task based on the provided prompt."""
+        assert prompt, "Prompt must be provided."
+
+        def _validate_json(response: str) -> None | Task:
+            try:
+                cap = JsonCapture.capture(response)
+                logger.debug(f"Response: \n{response}")
+                logger.info(f"Captured JSON: \n{cap[0]}")
+                return Task.model_validate_json(cap[0] if cap else response)
+            except ValidationError as e:
+                logger.error(f"Failed to parse task from JSON: {e}")
+                return None
+
+        return await self.aask_validate(
+            f"{prompt} \n\nBased on requirement above, "
+            f"you need to construct a task to satisfy that requirement in JSON format "
+            f"written like this: \n\n```json\n{Task.json_example()}\n```\n\n"
+            f"No extra explanation needed. ",
+            _validate_json,
+            system_message=f"# your personal briefing: \n{self.briefing}",
+        )

fabricatio/models/tool.py

@@ -0,0 +1,100 @@
+"""A module for defining tools and toolboxes."""
+
+from inspect import getfullargspec, signature
+from typing import Any, Callable, List, Self
+
+from fabricatio.models.generic import WithBriefing
+from pydantic import Field
+
+
+class Tool[**P, R](WithBriefing):
+    """A class representing a tool with a callable source function."""
+
+    name: str = Field(default="")
+    """The name of the tool."""
+
+    description: str = Field(default="")
+    """The description of the tool."""
+
+    source: Callable[P, R]
+    """The source function of the tool."""
+
+    def model_post_init(self, __context: Any) -> None:
+        """Initialize the tool with a name and a source function."""
+        self.name = self.name or self.source.__name__
+        assert self.name, "The tool must have a name."
+        self.description = self.description or self.source.__doc__ or ""
+
+    def invoke(self, *args: P.args, **kwargs: P.kwargs) -> R:
+        """Invoke the tool's source function with the provided arguments."""
+        return self.source(*args, **kwargs)
+
+    @property
+    def briefing(self) -> str:
+        """Return a brief description of the tool.
+
+        Returns:
+            str: A brief description of the tool.
+        """
+        source_signature = str(signature(self.source))
+        # 获取源函数的返回类型
+        return_annotation = getfullargspec(self.source).annotations.get("return", "None")
+        return f"{self.name}{source_signature} -> {return_annotation}\n{self.description}"
+
+
+class ToolBox(WithBriefing):
+    """A class representing a collection of tools."""
+
+    tools: List[Tool] = Field(default_factory=list)
+    """A list of tools in the toolbox."""
+
+    def collect_tool[**P, R](self, func: Callable[P, R]) -> Callable[P, R]:
+        """Add a callable function to the toolbox as a tool.
+
+        Args:
+            func (Callable[P, R]): The function to be added as a tool.
+
+        Returns:
+            Callable[P, R]: The added function.
+        """
+        self.tools.append(Tool(source=func))
+        return func
+
+    def add_tool[**P, R](self, func: Callable[P, R]) -> Self:
+        """Add a callable function to the toolbox as a tool.
+
+        Args:
+            func (Callable): The function to be added as a tool.
+
+        Returns:
+            Self: The current instance of the toolbox.
+        """
+        self.tools.append(Tool(source=func))
+        return self
+
+    @property
+    def briefing(self) -> str:
+        """Return a brief description of the toolbox.
+
+        Returns:
+            str: A brief description of the toolbox.
+        """
+        list_out = "\n\n".join([f"- {tool.briefing}" for tool in self.tools])
+        toc = f"## {self.name}: {self.description}\n## {len(self.tools)} tools available:"
+        return f"{toc}\n\n{list_out}"
+
+    def get[**P, R](self, name: str) -> Tool[P, R]:
+        """Invoke a tool by name with the provided arguments.
+
+        Args:
+            name (str): The name of the tool to invoke.
+
+        Returns:
+            Tool: The tool instance with the specified name.
+
+        Raises:
+            AssertionError: If no tool with the specified name is found.
+        """
+        tool = next((tool for tool in self.tools if tool.name == name), None)
+        assert tool, f"No tool named {name} found."
+        return tool

fabricatio/models/utils.py

@@ -0,0 +1,78 @@
+"""A module containing utility classes for the models."""
+
+from typing import Dict, List, Literal, Self
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class Message(BaseModel):
+    """A class representing a message."""
+
+    model_config = ConfigDict(use_attribute_docstrings=True)
+    role: Literal["user", "system", "assistant"] = Field(default="user")
+    """
+    Who is sending the message.
+    """
+    content: str = Field(default="")
+    """
+    The content of the message.
+    """
+
+
+class Messages(list):
+    """A list of messages."""
+
+    def add_message(self, role: Literal["user", "system", "assistant"], content: str) -> Self:
+        """Adds a message to the list with the specified role and content.
+
+        Args:
+            role (Literal["user", "system", "assistant"]): The role of the message sender.
+            content (str): The content of the message.
+
+        Returns:
+            Self: The current instance of Messages to allow method chaining.
+        """
+        if content:
+            self.append(Message(role=role, content=content))
+        return self
+
+    def add_user_message(self, content: str) -> Self:
+        """Adds a user message to the list with the specified content.
+
+        Args:
+            content (str): The content of the user message.
+
+        Returns:
+            Self: The current instance of Messages to allow method chaining.
+        """
+        return self.add_message("user", content)
+
+    def add_system_message(self, content: str) -> Self:
+        """Adds a system message to the list with the specified content.
+
+        Args:
+            content (str): The content of the system message.
+
+        Returns:
+            Self: The current instance of Messages to allow method chaining.
+        """
+        return self.add_message("system", content)
+
+    def add_assistant_message(self, content: str) -> Self:
+        """Adds an assistant message to the list with the specified content.
+
+        Args:
+            content (str): The content of the assistant message.
+
+        Returns:
+            Self: The current instance of Messages to allow method chaining.
+        """
+        return self.add_message("assistant", content)
+
+    def as_list(self) -> List[Dict[str, str]]:
+        """Converts the messages to a list of dictionaries.
+
+        Returns:
+            list[dict]: A list of dictionaries representing the messages.
+        """
+        return [message.model_dump() for message in self]

fabricatio/parser.py

@@ -0,0 +1,69 @@
+"""A module to parse text using regular expressions."""
+
+from typing import Any, Self, Tuple
+
+import regex
+from pydantic import Field, PositiveInt, PrivateAttr
+from regex import Pattern, compile
+
+from fabricatio.models.generic import Base
+
+
+class Capture(Base):
+    """A class to capture patterns in text using regular expressions.
+
+    Attributes:
+        pattern (str): The regular expression pattern to search for.
+        _compiled (Pattern): The compiled regular expression pattern.
+    """
+
+    target_groups: Tuple[int, ...] = Field(default_factory=tuple)
+    """The target groups to capture from the pattern."""
+    pattern: str = Field(frozen=True)
+    """The regular expression pattern to search for."""
+    flags: PositiveInt = Field(default=regex.DOTALL | regex.MULTILINE | regex.IGNORECASE, frozen=True)
+    """The flags to use when compiling the regular expression pattern."""
+    _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.
+        """
+        self._compiled = compile(self.pattern, self.flags)
+
+    def capture(self, text: str) -> Tuple[str, ...] | None:
+        """Capture the first occurrence of the pattern in the given text.
+
+        Args:
+            text (str): The text to search the pattern in.
+
+        Returns:
+            str | None: The captured text if the pattern is found, otherwise None.
+
+        """
+        match = self._compiled.search(text)
+        if match is None:
+            return None
+
+        if self.target_groups:
+            return tuple(match.group(g) for g in self.target_groups)
+        return (match.group(),)
+
+    @classmethod
+    def capture_code_block(cls, language: str) -> Self:
+        """Capture the first occurrence of a code block in the given text.
+
+        Args:
+            language (str): The text containing the code block.
+
+        Returns:
+            Self: The instance of the class with the captured code block.
+        """
+        return cls(pattern=f"```{language}\n(.*?)\n```", target_groups=(1,))
+
+
+JsonCapture = Capture.capture_code_block("json")
+PythonCapture = Capture.capture_code_block("python")
+CodeBlockCapture = Capture.capture_code_block("")

fabricatio/templates.py

@@ -0,0 +1,41 @@
+"""A module that manages templates for code generation."""
+
+from typing import Any, Dict, List, Self
+
+from pydantic import BaseModel, ConfigDict, DirectoryPath, Field, FilePath, PrivateAttr
+
+from fabricatio.config import configs
+from fabricatio.journal import logger
+
+
+class TemplateManager(BaseModel):
+    """A class that manages templates for code generation."""
+
+    model_config = ConfigDict(use_attribute_docstrings=True)
+    templates_dir: List[DirectoryPath] = Field(default_factory=lambda: list(configs.code2prompt.template_dir))
+    """The directories containing the templates. first element has the highest override priority."""
+    _discovered_templates: Dict[str, FilePath] = PrivateAttr(default_factory=dict)
+
+    def model_post_init(self, __context: Any) -> None:
+        """Post-initialization method for the model."""
+        self.discover_templates()
+
+    def discover_templates(self) -> Self:
+        """Discover the templates in the template directories."""
+        discovered = [
+            f
+            for d in self.templates_dir[::-1]
+            for f in d.rglob(f"*{configs.code2prompt.template_suffix}", case_sensitive=False)
+            if f.is_file()
+        ]
+
+        self._discovered_templates = {f.stem: f for f in discovered}
+        logger.info(f"Discovered {len(self._discovered_templates)} templates.")
+        return self
+
+    def get_template(self, name: str) -> FilePath | None:
+        """Get the template with the specified name."""
+        return self._discovered_templates.get(name, None)
+
+
+templates_manager = TemplateManager()

fabricatio/toolboxes/__init__.py

@@ -0,0 +1,7 @@
+"""Contains the built-in toolboxes for the Fabricatio package."""
+
+from fabricatio.toolboxes.task import TaskToolBox
+
+__all__ = [
+    "TaskToolBox",
+]

fabricatio/toolboxes/task.py

@@ -0,0 +1,4 @@
+from fabricatio.models.task import Task
+from fabricatio.models.tool import ToolBox
+
+TaskToolBox = ToolBox(name="TaskToolBox", description="A toolbox for tasks management.").add_tool(Task.simple_task)