fabricatio 0.2.3.dev3__cp312-cp312-manylinux_2_34_x86_64.whl → 0.2.4__cp312-cp312-manylinux_2_34_x86_64.whl

fabricatio/__init__.py

@@ -3,23 +3,32 @@
 from importlib.util import find_spec
 
 from fabricatio._rust_instances import template_manager
+from fabricatio.actions.article import ExtractArticleEssence, GenerateArticleProposal, GenerateOutline
+from fabricatio.actions.output import DumpFinalizedOutput
 from fabricatio.core import env
-from fabricatio.fs import magika
+from fabricatio.fs import magika, safe_json_read, safe_text_read
 from fabricatio.journal import logger
 from fabricatio.models.action import Action, WorkFlow
 from fabricatio.models.events import Event
+from fabricatio.models.extra import ArticleEssence
 from fabricatio.models.role import Role
 from fabricatio.models.task import Task
 from fabricatio.models.tool import ToolBox
 from fabricatio.models.utils import Message, Messages
 from fabricatio.parser import Capture, CodeBlockCapture, JsonCapture, PythonCapture
-from fabricatio.toolboxes import arithmetic_toolbox, basic_toolboxes, fs_toolbox, task_toolbox
+from fabricatio.toolboxes import arithmetic_toolbox, basic_toolboxes, fs_toolbox
+from fabricatio.workflows.articles import WriteOutlineWorkFlow
 
 __all__ = [
     "Action",
+    "ArticleEssence",
     "Capture",
     "CodeBlockCapture",
+    "DumpFinalizedOutput",
     "Event",
+    "ExtractArticleEssence",
+    "GenerateArticleProposal",
+    "GenerateOutline",
     "JsonCapture",
     "Message",
     "Messages",
@@ -28,18 +37,22 @@ __all__ = [
     "Task",
     "ToolBox",
     "WorkFlow",
+    "WriteOutlineWorkFlow",
     "arithmetic_toolbox",
     "basic_toolboxes",
     "env",
     "fs_toolbox",
     "logger",
     "magika",
-    "task_toolbox",
+    "safe_json_read",
+    "safe_text_read",
     "template_manager",
 ]
 
 
 if find_spec("pymilvus"):
-    from fabricatio.capabilities.rag import Rag
+    from fabricatio.actions.rag import InjectToDB
+    from fabricatio.capabilities.rag import RAG
+    from fabricatio.workflows.rag import StoreArticle
 
-    __all__ += ["Rag"]
+    __all__ += ["RAG", "InjectToDB", "StoreArticle"]

fabricatio/actions/article.py

@@ -0,0 +1,81 @@
+"""Actions for transmitting tasks to targets."""
+
+from os import PathLike
+from pathlib import Path
+from typing import Callable, List
+
+from fabricatio.fs import safe_text_read
+from fabricatio.journal import logger
+from fabricatio.models.action import Action
+from fabricatio.models.extra import ArticleEssence, ArticleOutline, ArticleProposal
+from fabricatio.models.task import Task
+
+
+class ExtractArticleEssence(Action):
+    """Extract the essence of article(s) in text format from the paths specified in the task dependencies.
+
+    Notes:
+        This action is designed to extract vital information from articles with Markdown format, which is pure text, and
+        which is converted from pdf files using `magic-pdf` from the `MinerU` project, see https://github.com/opendatalab/MinerU
+    """
+
+    output_key: str = "article_essence"
+    """The key of the output data."""
+
+    async def _execute[P: PathLike | str](
+        self,
+        task_input: Task,
+        reader: Callable[[P], str] = lambda p: Path(p).read_text(encoding="utf-8"),
+        **_,
+    ) -> List[ArticleEssence]:
+        if not task_input.dependencies:
+            logger.info(err := "Task not approved, since no dependencies are provided.")
+            raise RuntimeError(err)
+
+        # trim the references
+        contents = ["References".join(c.split("References")[:-1]) for c in map(reader, task_input.dependencies)]
+        return await self.propose(
+            ArticleEssence,
+            contents,
+            system_message=f"# your personal briefing: \n{self.briefing}",
+        )
+
+
+class GenerateArticleProposal(Action):
+    """Generate an outline for the article based on the extracted essence."""
+
+    output_key: str = "article_proposal"
+    """The key of the output data."""
+
+    async def _execute(
+        self,
+        task_input: Task,
+        **_,
+    ) -> ArticleProposal:
+        input_path = await self.awhich_pathstr(
+            f"{task_input.briefing}\nExtract the path of file, which contains the article briefing that I need to read."
+        )
+
+        return await self.propose(
+            ArticleProposal,
+            safe_text_read(input_path),
+            system_message=f"# your personal briefing: \n{self.briefing}",
+        )
+
+
+class GenerateOutline(Action):
+    """Generate the article based on the outline."""
+
+    output_key: str = "article"
+    """The key of the output data."""
+
+    async def _execute(
+        self,
+        article_proposal: ArticleProposal,
+        **_,
+    ) -> ArticleOutline:
+        return await self.propose(
+            ArticleOutline,
+            article_proposal.display(),
+            system_message=f"# your personal briefing: \n{self.briefing}",
+        )

fabricatio/actions/output.py

@@ -0,0 +1,21 @@
+"""Dump the finalized output to a file."""
+
+from typing import Unpack
+
+from fabricatio.models.action import Action
+from fabricatio.models.generic import FinalizedDumpAble
+from fabricatio.models.task import Task
+
+
+class DumpFinalizedOutput(Action):
+    """Dump the finalized output to a file."""
+
+    output_key: str = "dump_path"
+
+    async def _execute(self, task_input: Task, to_dump: FinalizedDumpAble, **cxt: Unpack) -> str:
+        dump_path = await self.awhich_pathstr(
+            f"{task_input.briefing}\n\nExtract a single path of the file, to which I will dump the data."
+        )
+
+        to_dump.finalized_dump_to(dump_path)
+        return dump_path

fabricatio/actions/rag.py

@@ -0,0 +1,25 @@
+"""Inject data into the database."""
+
+from typing import List, Optional, Unpack
+
+from fabricatio.capabilities.rag import RAG
+from fabricatio.models.action import Action
+from fabricatio.models.generic import PrepareVectorization
+
+
+class InjectToDB(Action, RAG):
+    """Inject data into the database."""
+
+    output_key: str = "collection_name"
+
+    async def _execute[T: PrepareVectorization](
+        self, to_inject: T | List[T], collection_name: Optional[str] = "my_collection", **cxt: Unpack
+    ) -> str:
+        if not isinstance(to_inject, list):
+            to_inject = [to_inject]
+
+        await self.view(collection_name, create=True).consume_string(
+            [t.prepare_vectorization(self.embedding_max_sequence_length) for t in to_inject],
+        )
+
+        return collection_name

fabricatio/capabilities/propose.py

@@ -0,0 +1,55 @@
+"""A module for the task capabilities of the Fabricatio library."""
+
+from typing import List, Type, Unpack, overload
+
+from fabricatio.models.generic import ProposedAble
+from fabricatio.models.kwargs_types import GenerateKwargs
+from fabricatio.models.usages import LLMUsage
+
+
+class Propose[M: ProposedAble](LLMUsage):
+    """A class that proposes an Obj based on a prompt."""
+
+    @overload
+    async def propose(
+        self,
+        cls: Type[M],
+        prompt: List[str],
+        **kwargs: Unpack[GenerateKwargs],
+    ) -> List[M]: ...
+
+    @overload
+    async def propose(
+        self,
+        cls: Type[M],
+        prompt: str,
+        **kwargs: Unpack[GenerateKwargs],
+    ) -> M: ...
+
+    async def propose(
+        self,
+        cls: Type[M],
+        prompt: List[str] | str,
+        **kwargs: Unpack[GenerateKwargs],
+    ) -> List[M] | M:
+        """Asynchronously proposes a task based on a given prompt and parameters.
+
+        Parameters:
+            cls: The class type of the task to be proposed.
+            prompt: The prompt text for proposing a task, which is a string that must be provided.
+            **kwargs: The keyword arguments for the LLM (Large Language Model) usage.
+
+        Returns:
+            A Task object based on the proposal result.
+        """
+        if isinstance(prompt, str):
+            return await self.aask_validate(
+                question=cls.create_json_prompt(prompt),
+                validator=cls.instantiate_from_string,
+                **kwargs,
+            )
+        return await self.aask_validate_batch(
+            questions=[cls.create_json_prompt(p) for p in prompt],
+            validator=cls.instantiate_from_string,
+            **kwargs,
+        )

fabricatio/capabilities/rag.py

@@ -1,71 +1,120 @@
 """A module for the RAG (Retrieval Augmented Generation) model."""
 
+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 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
+from typing import Any, Callable, Dict, List, Optional, Self, Union, Unpack, overload
 
-from fabricatio import template_manager
+from fabricatio._rust_instances import template_manager
 from fabricatio.config import configs
-from fabricatio.models.kwargs_types import LLMKwargs
-from fabricatio.models.usages import LLMUsage
+from fabricatio.journal import logger
+from fabricatio.models.kwargs_types import (
+    ChooseKwargs,
+    CollectionSimpleConfigKwargs,
+    EmbeddingKwargs,
+    FetchKwargs,
+    LLMKwargs,
+)
+from fabricatio.models.usages import EmbeddingUsage
 from fabricatio.models.utils import MilvusData
-from more_itertools.recipes import flatten
-
-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 more_itertools.recipes import flatten, unique
 from pydantic import Field, PrivateAttr
 
 
 @lru_cache(maxsize=None)
-def create_client(
-    uri: Optional[str] = None, token: Optional[str] = None, timeout: Optional[float] = None
-) -> MilvusClient:
+def create_client(uri: str, token: str = "", timeout: Optional[float] = None) -> MilvusClient:
     """Create a Milvus client."""
     return MilvusClient(
-        uri=uri or configs.rag.milvus_uri.unicode_string(),
-        token=token or configs.rag.milvus_token.get_secret_value() if configs.rag.milvus_token else "",
-        timeout=timeout or configs.rag.milvus_timeout,
+        uri=uri,
+        token=token,
+        timeout=timeout,
     )
 
 
-class Rag(LLMUsage):
+class RAG(EmbeddingUsage):
     """A class representing the RAG (Retrieval Augmented Generation) model."""
 
-    milvus_uri: Optional[str] = Field(default=None, frozen=True)
-    """The URI of the Milvus server."""
-    milvus_token: Optional[str] = Field(default=None, frozen=True)
-    """The token for the Milvus server."""
-    milvus_timeout: Optional[float] = Field(default=None, frozen=True)
-    """The timeout for the Milvus server."""
     target_collection: Optional[str] = Field(default=None)
     """The name of the collection being viewed."""
 
-    _client: MilvusClient = PrivateAttr(None)
+    _client: Optional[MilvusClient] = PrivateAttr(None)
     """The Milvus client used for the RAG model."""
 
     @property
     def client(self) -> MilvusClient:
         """Return the Milvus client."""
+        if self._client is None:
+            raise RuntimeError("Client is not initialized. Have you called `self.init_client()`?")
         return self._client
 
-    def model_post_init(self, __context: Any) -> None:
-        """Initialize the RAG model by creating the collection if it does not exist."""
-        self._client = create_client(self.milvus_uri, self.milvus_token, self.milvus_timeout)
-        self.view(self.target_collection, create=True)
+    def init_client(
+        self,
+        milvus_uri: Optional[str] = None,
+        milvus_token: Optional[str] = None,
+        milvus_timeout: Optional[float] = None,
+    ) -> Self:
+        """Initialize the Milvus client."""
+        self._client = create_client(
+            uri=milvus_uri or (self.milvus_uri or configs.rag.milvus_uri).unicode_string(),
+            token=milvus_token
+            or (token.get_secret_value() if (token := (self.milvus_token or configs.rag.milvus_token)) else ""),
+            timeout=milvus_timeout or self.milvus_timeout,
+        )
+        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: ...
 
-    def view(self, collection_name: Optional[str], create: bool = False) -> Self:
+    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[CollectionSimpleConfigKwargs]
+    ) -> Self:
         """View the specified collection.
 
         Args:
             collection_name (str): The name of the collection.
             create (bool): Whether to create the collection if it does not exist.
+            **kwargs (Unpack[CollectionSimpleConfigKwargs]): Additional keyword arguments for collection configuration.
         """
         if create and collection_name and not self._client.has_collection(collection_name):
-            self._client.create_collection(collection_name)
+            kwargs["dimension"] = kwargs.get("dimension") or self.milvus_dimensions or configs.rag.milvus_dimensions
+            self._client.create_collection(collection_name, auto_id=True, **kwargs)
+            logger.info(f"Creating collection {collection_name}")
 
         self.target_collection = collection_name
         return self
@@ -90,13 +139,14 @@ class Rag(LLMUsage):
         return self.target_collection
 
     def add_document[D: Union[Dict[str, Any], MilvusData]](
-        self, data: D | List[D], collection_name: Optional[str] = None
+        self, data: D | List[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.
             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.
@@ -105,11 +155,19 @@ class Rag(LLMUsage):
             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 or self.safe_target_collection, data)
+        c_name = collection_name or self.safe_target_collection
+        self._client.insert(c_name, data)
+
+        if flush:
+            logger.debug(f"Flushing collection {c_name}")
+            self._client.flush(c_name)
         return self
 
-    def consume(
-        self, source: PathLike, reader: Callable[[PathLike], MilvusData], collection_name: Optional[str] = None
+    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.
 
@@ -121,8 +179,21 @@ class Rag(LLMUsage):
         Returns:
             Self: The current instance, allowing for method chaining.
         """
-        data = reader(Path(source))
-        self.add_document(data, collection_name or self.safe_target_collection)
+        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
 
     async def afetch_document(
@@ -130,6 +201,7 @@ class Rag(LLMUsage):
         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[Any]:
         """Fetch data from the collection.
@@ -138,6 +210,7 @@ class Rag(LLMUsage):
             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.
 
         Returns:
@@ -147,16 +220,18 @@ class Rag(LLMUsage):
         search_results = self._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],
             limit=result_per_query,
         )
 
         # Step 2: Flatten the search results
         flattened_results = flatten(search_results)
-
+        unique_results = unique(flattened_results, key=itemgetter("id"))
         # Step 3: Sort by distance (descending)
-        sorted_results = sorted(flattened_results, key=itemgetter("distance"), reverse=True)
+        sorted_results = sorted(unique_results, key=itemgetter("distance"), reverse=True)
 
+        logger.debug(f"Searched similarities: {[t['distance'] for t in sorted_results]}")
         # Step 4: Extract the entities
         resp = [result["entity"] for result in sorted_results]
 
@@ -167,37 +242,38 @@ class Rag(LLMUsage):
     async def aretrieve(
         self,
         query: List[str] | str,
-        collection_name: Optional[str] = None,
-        result_per_query: int = 10,
         final_limit: int = 20,
+        **kwargs: Unpack[FetchKwargs],
     ) -> List[str]:
         """Retrieve data from the collection.
 
         Args:
             query (List[str] | str): The query to be used for retrieval.
-            collection_name (Optional[str]): The name of the collection. If not provided, the currently viewed collection is used.
-            result_per_query (int): The number of results to be returned per query.
             final_limit (int): The final limit on the number of results to return.
+            **kwargs (Unpack[FetchKwargs]): Additional keyword arguments for retrieval.
 
         Returns:
             List[str]: A list of strings containing the retrieved data.
         """
         if isinstance(query, str):
             query = [query]
-        return await self.afetch_document(
-            vecs=(await self.vectorize(query)),
-            desired_fields="text",
-            collection_name=collection_name,
-            result_per_query=result_per_query,
+        return (
+            await self.afetch_document(
+                vecs=(await self.vectorize(query)),
+                desired_fields="text",
+                **kwargs,
+            )
         )[:final_limit]
 
     async def aask_retrieved(
         self,
-        question: str | List[str],
-        query: List[str] | str,
+        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.
@@ -206,20 +282,85 @@ class Rag(LLMUsage):
         specified question using the retrieved documents as context.
 
         Args:
-            question (str | List[str]): The question or list of questions to be asked.
+            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.
         """
-        docs = await self.aretrieve(query, collection_name, result_per_query, final_limit)
+        docs = await self.aretrieve(
+            query or question,
+            final_limit,
+            collection_name=collection_name,
+            result_per_query=result_per_query,
+            similarity_threshold=similarity_threshold,
+        )
+
+        rendered = template_manager.render_template(configs.templates.retrieved_display_template, {"docs": docs[::-1]})
+
+        logger.debug(f"Retrieved Documents: \n{rendered}")
         return await self.aask(
             question,
-            template_manager.render_template(configs.templates.retrieved_display_template, {"docs": docs}),
+            f"{rendered}\n\n{extra_system_message}",
+            **kwargs,
+        )
+
+    async def arefined_query(self, question: List[str] | str, **kwargs: Unpack[ChooseKwargs]) -> List[str]:
+        """Refines the given question using a template.
+
+        Args:
+            question (List[str] | str): The question to be refined.
+            **kwargs (Unpack[ChooseKwargs]): Additional keyword arguments for the refinement process.
+
+        Returns:
+            List[str]: A list of refined questions.
+        """
+        return await self.aliststr(
+            template_manager.render_template(
+                configs.templates.refined_query_template,
+                {"question": [question] if isinstance(question, str) else question},
+            ),
+            **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,
         )