fabricatio 0.2.2__tar.gz → 0.2.3.dev1__tar.gz

{fabricatio-0.2.2 → fabricatio-0.2.3.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fabricatio
-Version: 0.2.2
+Version: 0.2.3.dev1
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: 3.12

{fabricatio-0.2.2 → fabricatio-0.2.3.dev1}/examples/task_handle/handle_task.py

@@ -56,6 +56,30 @@ class WriteDocumentation(Action):
         return await self.aask(task_input.briefing, system_message=task_input.dependencies_prompt)
 
 
+class TestCancel(Action):
+    """Action that says hello to the world."""
+
+    name: str = "cancel"
+    description: str = "cancel the task"
+    output_key: str = "counter"
+
+    async def _execute(self, counter: int, **_) -> int:
+        logger.info(f"Counter: {counter}")
+        await asyncio.sleep(5)
+        counter += 1
+        return counter
+
+
+class WriteToOutput(Action):
+    """Action that says hello to the world."""
+
+    name: str = "write to output"
+    output_key: str = "task_output"
+
+    async def _execute(self, **_) -> str:
+        return "hi, this is the output"
+
+
 async def main() -> None:
     """Main function."""
     role = Role(
@@ -68,20 +92,32 @@ async def main() -> None:
             Event.instantiate_from("doc").push_wildcard().push("pending"): WorkFlow(
                 name="write documentation", steps=(WriteDocumentation, DumpText)
             ),
+            Event.instantiate_from("cancel_test").push_wildcard().push("pending"): WorkFlow(
+                name="cancel_test",
+                steps=(TestCancel, TestCancel, TestCancel, TestCancel, TestCancel, TestCancel, WriteToOutput),
+                extra_init_context={"counter": 0},
+            ),
         },
     )
 
-    prompt = "i want you to write a cli app implemented with python , which can calculate the sum to a given n, all write to a single file names `cli.py`, put it in `output` folder."
-
-    proposed_task = await role.propose(prompt)
-    path = await proposed_task.move_to("coding").delegate()
+    proposed_task = await role.propose(
+        "i want you to write a cli app implemented with python , which can calculate the sum to a given n, all write to a single file names `cli.py`, put it in `output` folder."
+    )
+    path = await proposed_task.delegate("coding")
     logger.success(f"Code Path: {path}")
 
-    proposed_task = await role.propose(f"write Readme.md file for the code, source file {path},save it in `README.md`")
-    proposed_task.override_dependencies([path])
-    doc = await proposed_task.move_to("doc").delegate()
+    proposed_task = await role.propose(
+        f"write Readme.md file for the code, source file {path},save it in `README.md`,which is in the `output` folder, too."
+    )
+    proposed_task.override_dependencies(path)
+    doc = await proposed_task.delegate("doc")
     logger.success(f"Documentation: \n{doc}")
 
+    proposed_task.publish("cancel_test")
+    await proposed_task.cancel()
+    out = await proposed_task.get_output()
+    logger.info(f"Canceled Task Output: {out}")
+
 
 if __name__ == "__main__":
     asyncio.run(main())

{fabricatio-0.2.2 → fabricatio-0.2.3.dev1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "fabricatio"
-version = "0.2.2"
+version = "0.2.3-dev1"
 description = "A LLM multi-agent framework."
 readme = "README.md"
 license = { file = "LICENSE" }

fabricatio-0.2.3.dev1/python/fabricatio/capabilities/rag.py

@@ -0,0 +1,46 @@
+"""A module for the RAG (Retrieval Augmented Generation) model."""
+
+from typing import Any, Dict, List, Optional, Self, Union
+
+from fabricatio.config import configs
+from fabricatio.models.utils import MilvusData
+
+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
+from pydantic import BaseModel, ConfigDict, PrivateAttr
+
+
+class Rag(BaseModel):
+    """A class representing the RAG (Retrieval Augmented Generation) model."""
+
+    model_config = ConfigDict(use_attribute_docstrings=True)
+
+    _client: MilvusClient = PrivateAttr(
+        default=MilvusClient(
+            uri=configs.rag.milvus_uri.unicode_string(),
+            token=configs.rag.milvus_token.get_secret_value(),
+            timeout=configs.rag.milvus_timeout,
+        ),
+    )
+    _target_collection: Optional[str] = PrivateAttr(default=None)
+
+    @property
+    def client(self) -> MilvusClient:
+        """The Milvus client."""
+        return self._client
+
+    def add_document[D: Union[Dict[str, Any] | MilvusData]](self, collection_name: str, data: D | List[D]) -> Self:
+        """Adds a document to the specified collection.
+
+        Args:
+            collection_name (str): The name of the collection.
+            data (dict): The data to be added to the collection.
+        """
+        if isinstance(data, MilvusData):
+            data = data.prepare_insertion()
+        if isinstance(data, list):
+            data = [d.prepare_insertion() if isinstance(d, MilvusData) else d for d in data]
+        self.client.insert(collection_name, data)
+        return self

{fabricatio-0.2.2 → fabricatio-0.2.3.dev1}/python/fabricatio/config.py

@@ -11,6 +11,7 @@ from pydantic import (
     FilePath,
     HttpUrl,
     NonNegativeFloat,
+    PositiveFloat,
     PositiveInt,
     SecretStr,
 )
@@ -207,6 +208,19 @@ class ToolBoxConfig(BaseModel):
     """The name of the module containing the data."""
 
 
+class RagConfig(BaseModel):
+    """RAG configuration class."""
+
+    model_config = ConfigDict(use_attribute_docstrings=True)
+
+    milvus_uri: HttpUrl = Field(default=HttpUrl("http://localhost:19530"))
+    """The URI of the Milvus server."""
+    milvus_timeout: Optional[PositiveFloat] = Field(default=None)
+    """The timeout of the Milvus server."""
+    milvus_token: Optional[SecretStr] = Field(default=None)
+    """The token of the Milvus server."""
+
+
 class Settings(BaseSettings):
     """Application settings class.
 
@@ -250,6 +264,9 @@ class Settings(BaseSettings):
     toolbox: ToolBoxConfig = Field(default_factory=ToolBoxConfig)
     """Toolbox Configuration"""
 
+    rag: RagConfig = Field(default_factory=RagConfig)
+    """RAG Configuration"""
+
     @classmethod
     def settings_customise_sources(
         cls,

{fabricatio-0.2.2 → fabricatio-0.2.3.dev1}/python/fabricatio/core.py

@@ -38,11 +38,11 @@ class Env(BaseModel):
 
     @overload
     def on[**P, R](
-            self,
-            event: str | Event,
-            func: Optional[Callable[P, R]] = None,
-            /,
-            ttl: int = -1,
+        self,
+        event: str | Event,
+        func: Optional[Callable[P, R]] = None,
+        /,
+        ttl: int = -1,
     ) -> Callable[[Callable[P, R]], Callable[P, R]]:
         """
         Registers an event listener with a specific function that listens indefinitely or for a specified number of times.
@@ -58,11 +58,11 @@ class Env(BaseModel):
         ...
 
     def on[**P, R](
-            self,
-            event: str | Event,
-            func: Optional[Callable[P, R]] = None,
-            /,
-            ttl=-1,
+        self,
+        event: str | Event,
+        func: Optional[Callable[P, R]] = None,
+        /,
+        ttl=-1,
     ) -> Callable[[Callable[P, R]], Callable[P, R]] | Self:
         """Registers an event listener with a specific function that listens indefinitely or for a specified number of times.
 
@@ -78,14 +78,13 @@ class Env(BaseModel):
             event = event.collapse()
         if func is None:
             return self._ee.on(event, ttl=ttl)
-
         self._ee.on(event, func, ttl=ttl)
         return self
 
     @overload
     def once[**P, R](
-            self,
-            event: str | Event,
+        self,
+        event: str | Event,
     ) -> Callable[[Callable[P, R]], Callable[P, R]]:
         """
         Registers an event listener that listens only once.
@@ -100,9 +99,9 @@ class Env(BaseModel):
 
     @overload
     def once[**P, R](
-            self,
-            event: str | Event,
-            func: Callable[[Callable[P, R]], Callable[P, R]],
+        self,
+        event: str | Event,
+        func: Callable[[Callable[P, R]], Callable[P, R]],
     ) -> Self:
         """
         Registers an event listener with a specific function that listens only once.
@@ -117,9 +116,9 @@ class Env(BaseModel):
         ...
 
     def once[**P, R](
-            self,
-            event: str | Event,
-            func: Optional[Callable[P, R]] = None,
+        self,
+        event: str | Event,
+        func: Optional[Callable[P, R]] = None,
     ) -> Callable[[Callable[P, R]], Callable[P, R]] | Self:
         """Registers an event listener with a specific function that listens only once.
 
@@ -163,5 +162,20 @@ class Env(BaseModel):
             event = event.collapse()
         return await self._ee.emit_async(event, *args, **kwargs)
 
+    def emit_future[**P](self, event: str | Event, *args: P.args, **kwargs: P.kwargs) -> None:
+        """Emits an event to all registered listeners and returns a future object.
+
+        Args:
+            event (str | Event): The event to emit.
+            *args: Positional arguments to pass to the listeners.
+            **kwargs: Keyword arguments to pass to the listeners.
+
+        Returns:
+            None: The future object.
+        """
+        if isinstance(event, Event):
+            event = event.collapse()
+        return self._ee.emit_future(event, *args, **kwargs)
+
 
 env = Env()

{fabricatio-0.2.2 → fabricatio-0.2.3.dev1}/python/fabricatio/models/action.py

@@ -2,7 +2,7 @@
 
 import traceback
 from abc import abstractmethod
-from asyncio import Queue
+from asyncio import Queue, create_task
 from typing import Any, Dict, Self, Tuple, Type, Union, Unpack
 
 from fabricatio.capabilities.rating import GiveRating
@@ -108,7 +108,11 @@ class WorkFlow(WithBriefing, ToolBoxUsage):
         try:
             for step in self._instances:
                 logger.debug(f"Executing step: {step.name}")
-                modified_ctx = await step.act(await self._context.get())
+                act_task = create_task(step.act(await self._context.get()))
+                if task.is_cancelled():
+                    act_task.cancel(f"Cancelled by task: {task.name}")
+                    break
+                modified_ctx = await act_task
                 await self._context.put(modified_ctx)
                 current_action = step.name
             logger.info(f"Finished executing workflow: {self.name}")

{fabricatio-0.2.2 → fabricatio-0.2.3.dev1}/python/fabricatio/models/kwargs_types.py

@@ -5,6 +5,15 @@ from typing import List, NotRequired, TypedDict
 from pydantic import NonNegativeFloat, NonNegativeInt, PositiveInt
 
 
+class EmbeddingKwargs(TypedDict):
+    """A type representing the keyword arguments for the embedding method."""
+
+    model: NotRequired[str]
+    dimensions: NotRequired[int]
+    timeout: NotRequired[PositiveInt]
+    caching: NotRequired[bool]
+
+
 class LLMKwargs(TypedDict):
     """A type representing the keyword arguments for the LLM (Large Language Model) usage."""
 

{fabricatio-0.2.2 → fabricatio-0.2.3.dev1}/python/fabricatio/models/task.py

@@ -60,7 +60,7 @@ class Task[T](WithBriefing, WithJsonExample, WithDependency):
     dependencies: List[str] = Field(default_factory=list)
     """A list of file paths, These file are needed to read or write to meet a specific requirement of this task, if it is not directly given out, it SHALL just be a empty list meaning `NOT ASSIGNED`"""
 
-    _output: Queue = PrivateAttr(default_factory=lambda: Queue(maxsize=1))
+    _output: Queue[T | None] = PrivateAttr(default_factory=Queue)
     """The output queue of the task."""
 
     _status: TaskStatus = PrivateAttr(default=TaskStatus.Pending)
@@ -131,7 +131,7 @@ class Task[T](WithBriefing, WithJsonExample, WithDependency):
             self.description = description
         return self
 
-    async def get_output(self) -> T:
+    async def get_output(self) -> T | None:
         """Get the output of the task.
 
         Returns:
@@ -232,6 +232,7 @@ class Task[T](WithBriefing, WithJsonExample, WithDependency):
         """
         logger.info(f"Cancelling task `{self.name}`")
         self._status = TaskStatus.Cancelled
+        await self._output.put(None)
         await env.emit_async(self.cancelled_label, self)
         return self
 
@@ -243,27 +244,38 @@ class Task[T](WithBriefing, WithJsonExample, WithDependency):
         """
         logger.info(f"Failing task `{self.name}`")
         self._status = TaskStatus.Failed
+        await self._output.put(None)
         await env.emit_async(self.failed_label, self)
         return self
 
-    async def publish(self) -> Self:
+    def publish(self, new_namespace: Optional[EventLike] = None) -> Self:
         """Publish the task to the event bus.
 
+        Args:
+            new_namespace(EventLike, optional): The new namespace to move the task to.
+
         Returns:
-            Task: The published instance of the `Task` class
+            Task: The published instance of the `Task` class.
         """
+        if new_namespace:
+            self.move_to(new_namespace)
         logger.info(f"Publishing task `{(label := self.pending_label)}`")
-        await env.emit_async(label, self)
+        env.emit_future(label, self)
         return self
 
-    async def delegate(self) -> T:
-        """Delegate the task to the event bus and wait for the output.
+    async def delegate(self, new_namespace: Optional[EventLike] = None) -> T | None:
+        """Delegate the task to the event.
+
+        Args:
+            new_namespace(EventLike, optional): The new namespace to move the task to.
 
         Returns:
-            T: The output of the task
+            T|None: The output of the task.
         """
+        if new_namespace:
+            self.move_to(new_namespace)
         logger.info(f"Delegating task `{(label := self.pending_label)}`")
-        await env.emit_async(label, self)
+        env.emit_future(label, self)
         return await self.get_output()
 
     @property
@@ -277,3 +289,43 @@ class Task[T](WithBriefing, WithJsonExample, WithDependency):
             configs.templates.task_briefing_template,
             self.model_dump(),
         )
+
+    def is_running(self) -> bool:
+        """Check if the task is running.
+
+        Returns:
+            bool: True if the task is running, False otherwise.
+        """
+        return self._status == TaskStatus.Running
+
+    def is_finished(self) -> bool:
+        """Check if the task is finished.
+
+        Returns:
+            bool: True if the task is finished, False otherwise.
+        """
+        return self._status == TaskStatus.Finished
+
+    def is_failed(self) -> bool:
+        """Check if the task is failed.
+
+        Returns:
+            bool: True if the task is failed, False otherwise.
+        """
+        return self._status == TaskStatus.Failed
+
+    def is_cancelled(self) -> bool:
+        """Check if the task is cancelled.
+
+        Returns:
+            bool: True if the task is cancelled, False otherwise.
+        """
+        return self._status == TaskStatus.Cancelled
+
+    def is_pending(self) -> bool:
+        """Check if the task is pending.
+
+        Returns:
+            bool: True if the task is pending, False otherwise.
+        """
+        return self._status == TaskStatus.Pending

{fabricatio-0.2.2 → fabricatio-0.2.3.dev1}/python/fabricatio/models/usages.py

@@ -10,14 +10,15 @@ 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.kwargs_types import ChooseKwargs, GenerateKwargs, LLMKwargs
+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
+from fabricatio.models.utils import Messages, MilvusData
 from fabricatio.parser import JsonCapture
 from litellm import stream_chunk_builder
 from litellm.types.utils import (
     Choices,
+    EmbeddingResponse,
     ModelResponse,
     StreamingChoices,
 )
@@ -61,6 +62,97 @@ class LLMUsage(Base):
     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)
+        ]
+
     async def aquery(
         self,
         messages: List[Dict[str, str]],

{fabricatio-0.2.2 → fabricatio-0.2.3.dev1}/python/fabricatio/models/utils.py

@@ -1,6 +1,6 @@
 """A module containing utility classes for the models."""
 
-from typing import Dict, List, Literal, Self
+from typing import Any, Dict, List, Literal, Optional, Self
 
 from pydantic import BaseModel, ConfigDict, Field
 
@@ -76,3 +76,52 @@ class Messages(list):
             list[dict]: A list of dictionaries representing the messages.
         """
         return [message.model_dump() for message in self]
+
+
+class MilvusData(BaseModel):
+    """A class representing data stored in Milvus."""
+
+    model_config = ConfigDict(use_attribute_docstrings=True)
+    id: Optional[int] = Field(default=None)
+    """The identifier of the data."""
+
+    vector: List[float]
+    """The vector representation of the data."""
+
+    text: str
+    """The text representation of the data."""
+
+    subject: Optional[str] = Field(default=None)
+    """A subject label that we use to demo metadata filtering later."""
+
+    def prepare_insertion(self) -> Dict[str, Any]:
+        """Prepares the data for insertion into Milvus.
+
+        Returns:
+            dict: A dictionary containing the data to be inserted into Milvus.
+        """
+        return self.model_dump(exclude_none=True)
+
+    def update_subject(self, new_subject: str) -> Self:
+        """Updates the subject label of the data.
+
+        Args:
+            new_subject (str): The new subject label.
+
+        Returns:
+            Self: The updated instance of MilvusData.
+        """
+        self.subject = new_subject
+        return self
+
+    def update_id(self, new_id: int) -> Self:
+        """Updates the identifier of the data.
+
+        Args:
+            new_id (int): The new identifier.
+
+        Returns:
+            Self: The updated instance of MilvusData.
+        """
+        self.id = new_id
+        return self