fabricatio 0.2.9.dev4__tar.gz → 0.2.10.dev0__tar.gz

{fabricatio-0.2.9.dev4 → fabricatio-0.2.10.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fabricatio
-Version: 0.2.9.dev4
+Version: 0.2.10.dev0
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: 3.12

fabricatio-0.2.10.dev0/examples/write_article/write_article.py

@@ -0,0 +1,241 @@
+"""Example of using the library."""
+
+import asyncio
+from pathlib import Path
+from typing import List
+
+from fabricatio import Event, Role, WorkFlow, logger
+from fabricatio.actions.article import (
+    FixIntrospectedErrors,
+    GenerateArticle,
+    GenerateArticleProposal,
+    GenerateInitialOutline,
+)
+from fabricatio.actions.article_rag import TweakArticleRAG
+from fabricatio.actions.output import DumpFinalizedOutput, GatherAsList, PersistentAll, RetrieveFromLatest
+from fabricatio.actions.rules import DraftRuleSet, GatherRuleset
+from fabricatio.models.action import Action
+from fabricatio.models.extra.article_main import Article
+from fabricatio.models.extra.article_outline import ArticleOutline
+from fabricatio.models.extra.article_proposal import ArticleProposal
+from fabricatio.models.extra.rule import RuleSet
+from fabricatio.models.task import Task
+from fabricatio.utils import ok
+
+
+class Connect(Action):
+    """Connect the article with the article_outline and article_proposal."""
+
+    output_key: str = "article"
+    """Connect the article with the article_outline and article_proposal."""
+
+    async def _execute(
+        self,
+        article_briefing: str,
+        article_proposal: ArticleProposal,
+        article_outline: ArticleOutline,
+        article: Article,
+        **cxt,
+    ) -> Article:
+        """Connect the article with the article_outline and article_proposal."""
+        return article.update_ref(article_outline.update_ref(article_proposal.update_ref(article_briefing)))
+
+
+async def main(article: bool, rule: bool = False, fintune: bool = False) -> None:
+    """Main function."""
+    Role(
+        name="Undergraduate Researcher",
+        description="Write an outline for an article in typst format.",
+        llm_model="openai/deepseek-v3-250324",
+        llm_temperature=1.3,
+        llm_top_p=1.0,
+        llm_max_tokens=8191,
+        llm_rpm=1000,
+        llm_tpm=5000000,
+        registry={
+            Event.quick_instantiate(ns := "article"): WorkFlow(
+                name="Generate Article Outline",
+                description="Generate an outline for an article. dump the outline to the given path. in typst format.",
+                steps=(
+                    RetrieveFromLatest(
+                        retrieve_cls=RuleSet,
+                        load_path="persistent_ruleset/outline_ruleset",
+                        output_key="outline_ruleset",
+                    ),
+                    RetrieveFromLatest(
+                        retrieve_cls=RuleSet,
+                        load_path="persistent_ruleset/dep_ref_ruleset",
+                        output_key="dep_ref_ruleset",
+                    ),
+                    RetrieveFromLatest(
+                        retrieve_cls=RuleSet,
+                        load_path="persistent_ruleset/rev_dep_ref_ruleset",
+                        output_key="rev_dep_ref_ruleset",
+                    ),
+                    RetrieveFromLatest(
+                        retrieve_cls=RuleSet, load_path="persistent_ruleset/cite_ruleset", output_key="cite_ruleset"
+                    ),
+                    RetrieveFromLatest(
+                        retrieve_cls=RuleSet, load_path="persistent_ruleset/para_ruleset", output_key="para_ruleset"
+                    ),
+                    RetrieveFromLatest(
+                        retrieve_cls=RuleSet, load_path="persistent_ruleset/ref_ruleset", output_key="ref_ruleset"
+                    ),
+                    RetrieveFromLatest(
+                        retrieve_cls=RuleSet, load_path="persistent_ruleset/lang_ruleset", output_key="lang_ruleset"
+                    ),
+                    GenerateArticleProposal(llm_temperature=1.3),
+                    GenerateInitialOutline(
+                        output_key="article_outline",
+                    ),
+                    GatherRuleset(output_key="intro_fix_ruleset", to_gather=["para_ruleset"]),
+                    FixIntrospectedErrors(
+                        output_key="article_outline",
+                    ),
+                    PersistentAll,
+                    GatherRuleset(output_key="ref_fix_ruleset", to_gather=["ref_ruleset"]),
+                    GatherRuleset(output_key="ref_twk_ruleset", to_gather=["dep_ref_ruleset", "ref_ruleset"]),
+                    GatherRuleset(output_key="article_gen_ruleset", to_gather=["para_ruleset"]),
+                    GenerateArticle(
+                        output_key="article",
+                    ),
+                    PersistentAll,
+                    GatherRuleset(
+                        output_key="twk_rag_ruleset", to_gather=["para_ruleset", "cite_ruleset", "lang_ruleset"]
+                    ),
+                    TweakArticleRAG(output_key="to_dump"),
+                    DumpFinalizedOutput(output_key="task_output"),
+                    PersistentAll,
+                ),
+            ).update_init_context(
+                article_briefing=Path("./article_briefing.txt").read_text(),
+                dump_path="out.typ",
+                persist_dir="persistent",
+                collection_name="article_essence_0324",
+            ),
+            Event.quick_instantiate(finetune := "article_finetune"): WorkFlow(
+                name="Generate Article Outline",
+                description="Generate an outline for an article. dump the outline to the given path. in typst format.",
+                steps=(
+                    RetrieveFromLatest(
+                        retrieve_cls=RuleSet,
+                        load_path="persistent_ruleset/outline_ruleset",
+                        output_key="outline_ruleset",
+                    ),
+                    RetrieveFromLatest(
+                        retrieve_cls=RuleSet,
+                        load_path="persistent_ruleset/dep_ref_ruleset",
+                        output_key="dep_ref_ruleset",
+                    ),
+                    RetrieveFromLatest(
+                        retrieve_cls=RuleSet,
+                        load_path="persistent_ruleset/rev_dep_ref_ruleset",
+                        output_key="rev_dep_ref_ruleset",
+                    ),
+                    RetrieveFromLatest(
+                        retrieve_cls=RuleSet, load_path="persistent_ruleset/cite_ruleset", output_key="cite_ruleset"
+                    ),
+                    RetrieveFromLatest(
+                        retrieve_cls=RuleSet, load_path="persistent_ruleset/para_ruleset", output_key="para_ruleset"
+                    ),
+                    RetrieveFromLatest(
+                        retrieve_cls=RuleSet, load_path="persistent_ruleset/ref_ruleset", output_key="ref_ruleset"
+                    ),
+                    RetrieveFromLatest(
+                        retrieve_cls=ArticleProposal,
+                        load_path="persistent/article_proposal",
+                        output_key="article_proposal",
+                    ),
+                    RetrieveFromLatest(
+                        retrieve_cls=ArticleOutline,
+                        load_path="persistent/article_outline",
+                        output_key="article_outline",
+                    ),
+                    RetrieveFromLatest(retrieve_cls=Article, load_path="persistent/article", output_key="article"),
+                    Connect,
+                    GatherRuleset(output_key="intro_fix_ruleset", to_gather=["para_ruleset"]),
+                    GatherRuleset(output_key="ref_fix_ruleset", to_gather=["ref_ruleset"]),
+                    GatherRuleset(output_key="article_gen_ruleset", to_gather=["para_ruleset"]),
+                    GatherRuleset(output_key="twk_rag_ruleset", to_gather=["para_ruleset", "cite_ruleset"]),
+                    TweakArticleRAG(
+                        output_key="to_dump",
+                    ),
+                    DumpFinalizedOutput(output_key="task_output"),
+                    PersistentAll,
+                ),
+            ).update_init_context(
+                article_briefing=Path("./article_briefing.txt").read_text(),
+                dump_path="out_fix.typ",
+                persist_dir="persistent_fix",
+                collection_name="article_essence_0324",
+            ),
+            Event.quick_instantiate(rule_ns := "rule"): WorkFlow(
+                name="Generate Draft Rule Set",
+                description="Generate a draft rule set for the article.",
+                llm_model="openai/deepseek-v3-250324",
+                llm_stream=False,
+                steps=(
+                    # 精简后的para_ruleset规则
+                    DraftRuleSet(
+                        ruleset_requirement="如果`paragraphs`字段为空列表，那么你就需要按照`expected_word_count`来为章节补充内容",
+                        output_key="para_ruleset",
+                        rule_count=1,
+                    ),
+                    # 精简后的cite_ruleset规则
+                    DraftRuleSet(
+                        ruleset_requirement="1. 参考文献引用格式：(作者等, 年份)#cite(<bibtex_key>)\n"
+                        "2. #cite()必须用尖括号包裹单个BibTeX键，多引用需重复使用",
+                        output_key="cite_ruleset",
+                        rule_count=1,
+                    ),
+                    # 新增中文检测规则集
+                    DraftRuleSet(
+                        ruleset_requirement="1. 所有标题和正文内容必须使用中文,如果不为中文你需要翻译过来\n"
+                        "2. 术语和专业词汇需使用中文表述,英文缩写第一次出现的时候需要在其后面‘()’来辅助说明",
+                        output_key="lang_ruleset",
+                        rule_count=1,
+                    ),
+                    # 其他规则集保持原有结构但简化内容
+                    DraftRuleSet(
+                        ruleset_requirement="章节的`depend_on`字段的`ArticleRef`只能引用当前章节之前的元素。\n",
+                        output_key="dep_ref_ruleset",
+                        rule_count=1,
+                    ),
+                    DraftRuleSet(
+                        ruleset_requirement="章节的`support_to`字段的`ArticleRef`只能引用当前章节之后的元素。\n",
+                        output_key="rev_dep_ref_ruleset",
+                        rule_count=1,
+                    ),
+                    DraftRuleSet(
+                        ruleset_requirement="ArticleRef必须指向已定义元素",
+                        output_key="ref_ruleset",
+                        rule_count=1,
+                    ),
+                    DraftRuleSet(
+                        ruleset_requirement="标题使用学术术语",
+                        output_key="outline_ruleset",
+                        rule_count=1,
+                    ),
+                    GatherAsList(gather_suffix="ruleset").to_task_output(),
+                    PersistentAll(persist_dir="persistent_ruleset"),
+                ),
+            ),
+        },
+    )
+
+    if rule:
+        draft_rule_task: Task[List[RuleSet]] = Task(name="draft a rule set")
+        rule_set = ok(await draft_rule_task.delegate(rule_ns), "Failed to generate ruleset")
+        logger.success(f"Ruleset:\n{len(rule_set)}")
+    if article:
+        proposed_task = Task(name="write an article")
+        path = ok(await proposed_task.delegate(ns), "Failed to generate ruleset")
+        logger.success(f"The outline is saved in:\n{path}")
+    if fintune:
+        proposed_task = Task(name="finetune an article")
+        path = ok(await proposed_task.delegate(finetune), "Failed to generate ruleset")
+        logger.success(f"The outline is saved in:\n{path}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main(True, False, False))

{fabricatio-0.2.9.dev4 → fabricatio-0.2.10.dev0}/pyproject.toml

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

{fabricatio-0.2.9.dev4 → fabricatio-0.2.10.dev0}/python/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-0.2.9.dev4 → fabricatio-0.2.10.dev0}/python/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-0.2.10.dev0/python/fabricatio/capabilities/rag.py

@@ -0,0 +1,229 @@
+"""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 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.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.rust_instances import TEMPLATE_MANAGER
+from fabricatio.utils import ok
+
+
+@lru_cache(maxsize=None)
+def create_client(uri: str, token: str = "", timeout: Optional[float] = None) -> MilvusClient:
+    """Create a Milvus client."""
+    return MilvusClient(
+        uri=uri,
+        token=token,
+        timeout=timeout,
+    )
+
+
+class RAG(EmbeddingUsage):
+    """A class representing the RAG (Retrieval Augmented Generation) model."""
+
+    target_collection: Optional[str] = Field(default=None)
+    """The name of the collection being viewed."""
+
+    _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 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 ok(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 or configs.rag.milvus_timeout,
+        )
+        return self
+
+    def check_client(self, init: bool = True) -> Self:
+        """Check if the client is initialized, and if not, initialize it."""
+        if self._client is None and init:
+            return self.init_client()
+        if self._client is None and not init:
+            raise RuntimeError("Client is not initialized. Have you called `self.init_client()`?")
+        return self
+
+    def view(
+        self, collection_name: Optional[str], create: bool = False, **kwargs: Unpack[CollectionConfigKwargs]
+    ) -> 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[CollectionConfigKwargs]): Additional keyword arguments for collection configuration.
+        """
+        if create and collection_name and not self.check_client().client.has_collection(collection_name):
+            kwargs["dimension"] = ok(
+                kwargs.get("dimension")
+                or self.milvus_dimensions
+                or configs.rag.milvus_dimensions
+                or self.embedding_dimensions
+                or configs.embedding.dimensions,
+                "`dimension` is not set at any level.",
+            )
+            self.client.create_collection(collection_name, auto_id=True, **kwargs)
+            logger.info(f"Creating collection {collection_name}")
+
+        self.target_collection = collection_name
+        return self
+
+    def quit_viewing(self) -> Self:
+        """Quit the current view.
+
+        Returns:
+            Self: The current instance, allowing for method chaining.
+        """
+        return self.view(None)
+
+    @property
+    def safe_target_collection(self) -> str:
+        """Get the name of the collection being viewed, raise an error if not viewing any collection.
+
+        Returns:
+            str: The name of the collection being viewed.
+        """
+        return ok(self.target_collection, "No collection is being viewed. Have you called `self.view()`?")
+
+    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], 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, 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)
+
+        if flush:
+            logger.debug(f"Flushing collection {c_name}")
+            self.client.flush(c_name)
+        return self
+
+    async def afetch_document[D: MilvusDataBase](
+        self,
+        vecs: List[List[float]],
+        document_model: Type[D],
+        collection_name: Optional[str] = None,
+        similarity_threshold: float = 0.37,
+        result_per_query: int = 10,
+    ) -> List[D]:
+        """Asynchronously fetches documents from a Milvus database based on input vectors.
+
+        Args:
+           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[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=list(document_model.model_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(unique_results, key=itemgetter("distance"), reverse=True)
+
+        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]
+
+        return document_model.from_sequence(resp)
+
+    async def aretrieve[D: MilvusDataBase](
+        self,
+        query: List[str] | str,
+        final_limit: int = 20,
+        **kwargs: Unpack[FetchKwargs[D]],
+    ) -> List[D]:
+        """Retrieve data from the collection.
+
+        Args:
+            query (List[str] | str): The query to be used for retrieval.
+            final_limit (int): The final limit on the number of results to return.
+            **kwargs (Unpack[FetchKwargs]): Additional keyword arguments for retrieval.
+
+        Returns:
+            List[D]: A list of document objects created from the retrieved data.
+        """
+        if isinstance(query, str):
+            query = [query]
+        return (
+            await self.afetch_document(
+                vecs=(await self.vectorize(query)),
+                **kwargs,
+            )
+        )[:final_limit]
+
+    async def arefined_query(self, question: List[str] | str, **kwargs: Unpack[ChooseKwargs]) -> Optional[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.alist_str(
+            TEMPLATE_MANAGER.render_template(
+                configs.templates.refined_query_template,
+                {"question": [question] if isinstance(question, str) else question},
+            ),
+            **kwargs,
+        )

fabricatio-0.2.10.dev0/python/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"