fabricatio 0.2.13.dev0__tar.gz → 0.2.13.dev2__tar.gz

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fabricatio
-Version: 0.2.13.dev0
+Version: 0.2.13.dev2
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: 3.12
@@ -165,6 +165,10 @@ max_tokens = 8192
     ```bash
     make test
     ```
+## TODO
+
+- Add an element based format strategy
+
 
 ## Contributing
 

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev2}/README.md

@@ -119,6 +119,10 @@ max_tokens = 8192
     ```bash
     make test
     ```
+## TODO
+
+- Add an element based format strategy
+
 
 ## Contributing
 

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev2}/examples/write_article/write_article.py

@@ -2,12 +2,20 @@
 
 import asyncio
 from pathlib import Path
+from typing import Optional
 
 import typer
 from fabricatio import Event, Role, WorkFlow, logger
-from fabricatio.actions.article import ExtractOutlineFromRaw, GenerateArticleProposal, GenerateInitialOutline
+from fabricatio.actions.article import (
+    ExtractOutlineFromRaw,
+    GenerateArticleProposal,
+    GenerateInitialOutline,
+    WriteChapterSummary,
+)
 from fabricatio.actions.article_rag import ArticleConsultRAG, WriteArticleContentRAG
 from fabricatio.actions.output import DumpFinalizedOutput, PersistentAll, RenderedDump
+from fabricatio.fs import safe_text_read
+from fabricatio.models.extra.article_main import Article
 from fabricatio.models.extra.article_outline import ArticleOutline
 from fabricatio.models.task import Task
 from fabricatio.utils import ok
@@ -18,10 +26,9 @@ from typer import Typer
 Role(
     name="Undergraduate Researcher",
     description="Write an outline for an article in typst format.",
-    llm_model="openai/deepseek-v3-250324",
-    llm_temperature=0.45,
+    llm_model="openai/qwen3-235b-a22b",
     # llm_api_endpoint=HttpUrl("https://dashscope.aliyuncs.com/compatible-mode/v1"),
-    llm_top_p=0.95,
+    llm_stream=True,
     llm_max_tokens=8191,
     llm_rpm=600,
     llm_tpm=900000,
@@ -41,7 +48,7 @@ Role(
                         threshold=0.58,
                         result_per_query=2,
                         extractor_model={"model": "openai/qwen-max"},
-                        query_model={"model": "openai/qwen-turbo", "temperature": 0.3, "top_p": 0.85},
+                        query_model={"model": "openai/qwen-turbo"},
                     )
                 ),
                 PersistentAll,
@@ -74,7 +81,12 @@ Role(
         Event.quick_instantiate(ns4 := "consult"): WorkFlow(
             name="Consult Article",
             description="Consult an article with given article outline. dump the outline to the given path. in typst format.",
-            steps=(ArticleConsultRAG(ref_q_model={"model":"openai/qwen-turbo"}).to_task_output(),),
+            steps=(ArticleConsultRAG(ref_q_model={"model": "openai/qwen-turbo"}).to_task_output(),),
+        ),
+        Event.quick_instantiate(ns5 := "chap-suma"): WorkFlow(
+            name="Chapter Summary",
+            description="Generate chapter summary based on given article outline. dump the outline to the given path. in typst format.",
+            steps=(WriteChapterSummary().to_task_output(),),
         ),
     },
 )
@@ -86,13 +98,12 @@ app = Typer()
 @app.command()
 def consult(
     collection_name: str = typer.Option("article_chunks", "-c", "--collection-name", help="Name of the collection."),
+    tei_endpoint: Optional[str] = typer.Option(None, "-t", "--tei-endpoint", help="TEI endpoint."),
 ) -> None:
     """Consult an article based on a given article outline."""
     _ = asyncio.run(
         Task(name="Answer Question")
-        .update_init_context(
-            collection_name=collection_name,
-        )
+        .update_init_context(collection_name=collection_name, tei_endpoint=tei_endpoint)
         .delegate(ns4)
     )
 
@@ -193,5 +204,27 @@ def write(
     logger.success(f"The outline is saved in:\n{path}")
 
 
+@app.command()
+def suma(
+    article_path: Path = typer.Option(  # noqa: B008
+        Path("article.typ"), "-a", "--article-path", help="Path to the article file."
+    ),
+    dump_path: Path = typer.Option(Path("out.typ"), "-d", "--dump-path", help="Path to dump the final output."),  # noqa: B008
+) -> None:
+    """Write chap summary based on given article."""
+    path = ok(
+        asyncio.run(
+            Task(name="write an article")
+            .update_init_context(
+                article=Article.from_typst_code("article", body=safe_text_read(article_path)),
+                write_to=dump_path,
+            )
+            .delegate(ns5)
+        ),
+        "Failed to generate an article ",
+    )
+    logger.success(f"The outline is saved in:\n{path}")
+
+
 if __name__ == "__main__":
     app()

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev2}/pyproject.toml

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

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev2}/python/fabricatio/actions/article.py

@@ -8,10 +8,12 @@ from more_itertools import filter_map
 from pydantic import Field
 from rich import print as r_print
 
+from fabricatio import TEMPLATE_MANAGER
 from fabricatio.capabilities.censor import Censor
 from fabricatio.capabilities.extract import Extract
 from fabricatio.capabilities.propose import Propose
-from fabricatio.fs import safe_text_read
+from fabricatio.config import configs
+from fabricatio.fs import dump_text, safe_text_read
 from fabricatio.journal import logger
 from fabricatio.models.action import Action
 from fabricatio.models.extra.article_essence import ArticleEssence
@@ -21,6 +23,7 @@ from fabricatio.models.extra.article_proposal import ArticleProposal
 from fabricatio.models.extra.rule import RuleSet
 from fabricatio.models.kwargs_types import ValidateKwargs
 from fabricatio.models.task import Task
+from fabricatio.models.usages import LLMUsage
 from fabricatio.rust import BibManager, detect_language
 from fabricatio.utils import ok, wrapp_in_block
 
@@ -271,3 +274,48 @@ class LoadArticle(Action):
 
     async def _execute(self, article_outline: ArticleOutline, typst_code: str, **cxt) -> Article:
         return Article.from_mixed_source(article_outline, typst_code)
+
+
+class WriteChapterSummary(Action, LLMUsage):
+    """Write the chapter summary."""
+
+    output_key: str = "chapter_summaries"
+
+    paragraph_count: int = 1
+
+    summary_word_count: int = 200
+
+    summary_title: str = "Chapter Summary"
+    write_to: Optional[Path] = None
+
+    async def _execute(self, article: Article, write_to: Optional[Path] = None, **cxt) -> List[str]:
+        logger.info(";".join(a.title for a in article.chapters))
+
+        ret = [
+            f"== {self.summary_title}\n{raw}"
+            for raw in (
+                await self.aask(
+                    TEMPLATE_MANAGER.render_template(
+                        configs.templates.chap_summary_template,
+                        [
+                            {
+                                "chapter": a.to_typst_code(),
+                                "title": a.title,
+                                "language": a.language,
+                                "summary_word_count": self.summary_word_count,
+                                "paragraph_count": self.paragraph_count,
+                            }
+                            for a in article.chapters
+                        ],
+                    )
+                )
+            )
+        ]
+
+        if (to := (self.write_to or write_to)) is not None:
+            dump_text(
+                to,
+                "\n\n\n".join(f"//{a.title}\n\n{s}" for a, s in zip(article.chapters, ret, strict=True)),
+            )
+
+        return ret

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev2}/python/fabricatio/actions/article_rag.py

@@ -2,7 +2,7 @@
 
 from asyncio import gather
 from pathlib import Path
-from typing import List, Optional
+from typing import ClassVar, List, Optional
 
 from pydantic import Field, PositiveInt
 
@@ -21,17 +21,17 @@ from fabricatio.models.extra.article_outline import ArticleOutline
 from fabricatio.models.extra.rule import RuleSet
 from fabricatio.models.kwargs_types import ChooseKwargs, LLMKwargs
 from fabricatio.rust import convert_to_block_formula, convert_to_inline_formula
-from fabricatio.utils import ask_retain, ok
+from fabricatio.utils import ok
 
 TYPST_CITE_USAGE = (
-    "citation number is REQUIRED to cite any reference!,for example in Auther Pattern: 'Doe et al.[[1]], Jack et al.[[2]]' or in Sentence Suffix Sattern: 'Global requirement is incresing[[1]].'\n"
-    "Everything is build upon the typst language, which is similar to latex, \n"
+    "citation number is REQUIRED to cite any reference!'\n"
     "Legal citing syntax examples(seperated by |): [[1]]|[[1,2]]|[[1-3]]|[[12,13-15]]|[[1-3,5-7]]\n"
     "Illegal citing syntax examples(seperated by |): [[1],[2],[3]]|[[1],[1-2]]\n"
+    "You SHALL not cite a single reference more than once!"
     "Those reference mark shall not be omitted during the extraction\n"
     "It's recommended to cite multiple references that supports your conclusion at a time.\n"
     "Wrap inline expression with '\\(' and '\\)',like '\\(>5m\\)' '\\(89%\\)', and wrap block equation with '\\[' and '\\]'.\n"
-    "In addition to that, you can add a label outside the block equation which can be used as a cross reference identifier, the label is a string wrapped in `<` and `>` like `<energy-release-rate-equation>`.Note that the label string should be a summarizing title for the equation being labeled.\n"
+    "In addition to that, you can add a label outside the block equation which can be used as a cross reference identifier, the label is a string wrapped in `<` and `>` like `<energy-release-rate-equation>`.Note that the label string should be a summarizing title for the equation being labeled and should never be written within the formula block.\n"
     "you can refer to that label by using the syntax with prefix of `@eqt:`, which indicate that this notation is citing a label from the equations. For example ' @eqt:energy-release-rate-equation ' DO remember that the notation shall have both suffixed and prefixed space char which enable the compiler to distinguish the notation from the plaintext."
     "Below is two usage example:\n"
     "```typst\n"
@@ -47,6 +47,7 @@ TYPST_CITE_USAGE = (
 class WriteArticleContentRAG(Action, RAG, Extract):
     """Write an article based on the provided outline."""
 
+    ctx_override: ClassVar[bool] = True
     search_increment_multiplier: float = 1.6
     """The increment multiplier of the search increment."""
     ref_limit: int = 35
@@ -63,6 +64,7 @@ class WriteArticleContentRAG(Action, RAG, Extract):
     """The number of results to be returned per query."""
     req: str = TYPST_CITE_USAGE
     """The req of the write article content."""
+    tei_endpoint: Optional[str] = None
 
     async def _execute(
         self,
@@ -108,16 +110,7 @@ class WriteArticleContentRAG(Action, RAG, Extract):
 
         while not await confirm("Accept this version and continue?").ask_async():
             if inst := await text("Search for more refs for additional spec.").ask_async():
-                await self.search_database(
-                    article,
-                    article_outline,
-                    chap,
-                    sec,
-                    subsec,
-                    cm,
-                    supervisor=True,
-                    extra_instruction=inst,
-                )
+                await self.search_database(article, article_outline, chap, sec, subsec, cm, extra_instruction=inst)
 
             if instruction := await text("Enter the instructions to improve").ask_async():
                 raw = await self.write_raw(article, article_outline, chap, sec, subsec, cm, instruction)
@@ -200,7 +193,6 @@ class WriteArticleContentRAG(Action, RAG, Extract):
         subsec: ArticleSubsection,
         cm: CitationManager,
         extra_instruction: str = "",
-        supervisor: bool = False,
     ) -> None:
         """Search database for related references."""
         search_req = (
@@ -208,61 +200,31 @@ class WriteArticleContentRAG(Action, RAG, Extract):
             f"More specifically, i m witting the Chapter `{chap.title}` >> Section `{sec.title}` >> Subsection `{subsec.title}`.\n"
             f"I need to search related references to build up the content of the subsec mentioned above, which is `{subsec.title}`.\n"
             f"provide 10~16 queries as possible, to get best result!\n"
-            f"You should provide both English version and chinese version of the refined queries!\n{extra_instruction}\n"
+            f"You should provide both English version and chinese version of the refined queries!\n{extra_instruction}"
         )
 
-        ref_q = ok(
-            await self.arefined_query(
-                search_req,
-                **self.query_model,
-            ),
-            "Failed to refine query.",
-        )
-
-        if supervisor:
-            ref_q = await ask_retain(ref_q)
-        ret = await self.aretrieve(
-            ref_q,
-            ArticleChunk,
-            max_accepted=self.ref_limit,
-            result_per_query=self.result_per_query,
-            similarity_threshold=self.threshold,
-        )
-
-        cm.add_chunks(ok(ret))
-        ref_q = await self.arefined_query(
-            f"{cm.as_prompt()}\n\nAbove is the retrieved references in the first RAG, now we need to perform the second RAG.\n\n{search_req}",
-            **self.query_model,
-        )
-
-        if ref_q is None:
-            logger.warning("Second refine query is None, skipping.")
-            return
-        if supervisor:
-            ref_q = await ask_retain(ref_q)
-
-        ret = await self.aretrieve(
-            ref_q,
-            ArticleChunk,
-            max_accepted=int(self.ref_limit * self.search_increment_multiplier),
-            result_per_query=int(self.result_per_query * self.search_increment_multiplier),
-            similarity_threshold=self.threshold,
+        await self.clued_search(
+            search_req,
+            cm,
+            refinery_kwargs=self.ref_q_model,
+            expand_multiplier=self.search_increment_multiplier,
+            base_accepted=self.ref_limit,
+            result_per_query=self.ref_per_q,
+            similarity_threshold=self.similarity_threshold,
+            tei_endpoint=self.tei_endpoint,
         )
-        if ret is None:
-            logger.warning("Second retrieve is None, skipping.")
-            return
-        cm.add_chunks(ret)
 
 
 class ArticleConsultRAG(Action, AdvancedRAG):
     """Write an article based on the provided outline."""
 
+    ctx_override: ClassVar[bool] = True
     output_key: str = "consult_count"
     search_increment_multiplier: float = 1.6
     """The multiplier to increase the limit of references to retrieve per query."""
-    ref_limit: int = 20
+    ref_limit: int = 26
     """The final limit of references."""
-    ref_per_q: int = 3
+    ref_per_q: int = 13
     """The limit of references to retrieve per query."""
     similarity_threshold: float = 0.62
     """The similarity threshold of references to retrieve."""
@@ -270,6 +232,7 @@ class ArticleConsultRAG(Action, AdvancedRAG):
     """The model to use for refining query."""
     req: str = TYPST_CITE_USAGE
     """The request for the rag model."""
+    tei_endpoint: Optional[str] = None
 
     @precheck_package(
         "questionary", "`questionary` is required for supervisor mode, please install it by `fabricatio[qa]`"
@@ -300,6 +263,7 @@ class ArticleConsultRAG(Action, AdvancedRAG):
                 base_accepted=self.ref_limit,
                 result_per_query=self.ref_per_q,
                 similarity_threshold=self.similarity_threshold,
+                tei_endpoint=self.tei_endpoint,
             )
 
             ret = await self.aask(f"{cm.as_prompt()}\n{self.req}\n{req}")

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev2}/python/fabricatio/capabilities/advanced_rag.py

@@ -7,6 +7,7 @@ from fabricatio.journal import logger
 from fabricatio.models.adv_kwargs_types import FetchKwargs
 from fabricatio.models.extra.aricle_rag import ArticleChunk, CitationManager
 from fabricatio.models.kwargs_types import ChooseKwargs
+from fabricatio.utils import fallback_kwargs
 
 
 class AdvancedRAG(RAG):
@@ -40,10 +41,13 @@ class AdvancedRAG(RAG):
                 f"\n\n{requirement}",
                 **refinery_kwargs,
             )
+
             if ref_q is None:
                 logger.error(f"At round [{i}/{max_round}] search, failed to refine the query, exit.")
                 return cm
-            refs = await self.aretrieve(ref_q, ArticleChunk, base_accepted, **kwargs)
+            refs = await self.aretrieve(
+                ref_q, ArticleChunk, base_accepted, **fallback_kwargs(kwargs, filter_expr=cm.as_milvus_filter_expr())
+            )
 
             if (max_capacity := max_capacity - len(refs)) < 0:
                 cm.add_chunks(refs[0:max_capacity])

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev2}/python/fabricatio/capabilities/rag.py

@@ -143,21 +143,27 @@ class RAG(EmbeddingUsage):
 
     async def afetch_document[D: MilvusDataBase](
         self,
-        vecs: List[List[float]],
+        query: List[str],
         document_model: Type[D],
         collection_name: Optional[str] = None,
         similarity_threshold: float = 0.37,
         result_per_query: int = 10,
+        tei_endpoint: Optional[str] = None,
+        reranker_threshold: float = 0.7,
+        filter_expr: str = "",
     ) -> 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.
+           query (List[str]): 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.
+           tei_endpoint (str): the endpoint of the TEI api.
+           reranker_threshold (float): The threshold used to filtered low relativity document.
+            filter_expr (str): filter_expression parsed into pymilvus search.
 
         Returns:
            List[D]: A list of document objects created from the fetched data.
@@ -165,15 +171,38 @@ class RAG(EmbeddingUsage):
         # Step 1: Search for vectors
         search_results = self.check_client().client.search(
             collection_name or self.safe_target_collection,
-            vecs,
+            await self.vectorize(query),
             search_params={"radius": similarity_threshold},
             output_fields=list(document_model.model_fields),
+            filter=filter_expr,
             limit=result_per_query,
         )
+        if tei_endpoint is not None:
+            from fabricatio.utils import RerankerAPI
+
+            reranker = RerankerAPI(base_url=tei_endpoint)
+
+            retrieved_id = set()
+            raw_result = []
+
+            for q, g in zip(query, search_results, strict=True):
+                models = document_model.from_sequence([res["entity"] for res in g if res["id"] not in retrieved_id])
+                logger.debug(f"Retrived {len(g)} raw document, filtered out {len(models)}.")
+                retrieved_id.update(res["id"] for res in g)
+                if not models:
+                    continue
+                rank_scores = await reranker.arerank(q, [m.prepare_vectorization() for m in models], truncate=True)
+                raw_result.extend(
+                    (models[s["index"]], s["score"]) for s in rank_scores if s["score"] > reranker_threshold
+                )
+
+            raw_result_sorted = sorted(raw_result, key=lambda x: x[1], reverse=True)
+            return [r[0] for r in raw_result_sorted]
 
         # 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)
 
@@ -205,15 +234,18 @@ class RAG(EmbeddingUsage):
         """
         if isinstance(query, str):
             query = [query]
+
         return (
             await self.afetch_document(
-                vecs=(await self.vectorize(query)),
+                query=query,
                 document_model=document_model,
                 **kwargs,
             )
         )[:max_accepted]
 
-    async def arefined_query(self, question: List[str] | str, **kwargs: Unpack[ChooseKwargs[Optional[List[str]]]]) -> Optional[List[str]]:
+    async def arefined_query(
+        self, question: List[str] | str, **kwargs: Unpack[ChooseKwargs[Optional[List[str]]]]
+    ) -> Optional[List[str]]:
         """Refines the given question using a template.
 
         Args:

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev2}/python/fabricatio/config.py

@@ -86,10 +86,12 @@ class LLMConfig(BaseModel):
 
     tpm: Optional[PositiveInt] = Field(default=1000000)
     """The rate limit of the LLM model in tokens per minute. None means not checked."""
-    presence_penalty:Optional[PositiveFloat]=None
+    presence_penalty: Optional[PositiveFloat] = None
     """The presence penalty of the LLM model."""
-    frequency_penalty:Optional[PositiveFloat]=None
+    frequency_penalty: Optional[PositiveFloat] = None
     """The frequency penalty of the LLM model."""
+
+
 class EmbeddingConfig(BaseModel):
     """Embedding configuration class."""
 
@@ -252,10 +254,13 @@ class TemplateConfig(BaseModel):
     rule_requirement_template: str = Field(default="rule_requirement")
     """The name of the rule requirement template which will be used to generate a rule requirement."""
 
-
     extract_template: str = Field(default="extract")
     """The name of the extract template which will be used to extract model from string."""
 
+    chap_summary_template: str = Field(default="chap_summary")
+    """The name of the chap summary template which will be used to generate a chapter summary."""
+
+
 class MagikaConfig(BaseModel):
     """Magika configuration class."""
 

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev2}/python/fabricatio/models/action.py

@@ -12,7 +12,7 @@ Classes:
 import traceback
 from abc import abstractmethod
 from asyncio import Queue, create_task
-from typing import Any, Dict, Self, Sequence, Tuple, Type, Union, final
+from typing import Any, ClassVar, Dict, Self, Sequence, Tuple, Type, Union, final
 
 from fabricatio.journal import logger
 from fabricatio.models.generic import WithBriefing
@@ -33,6 +33,9 @@ class Action(WithBriefing):
     a specific operation and can modify the shared context data.
     """
 
+    ctx_override: ClassVar[bool] = False
+    """Whether to override the instance attr by the context variable."""
+
     name: str = Field(default="")
     """The name of the action."""
 
@@ -157,6 +160,15 @@ class WorkFlow(WithBriefing, ToolBoxUsage):
             action.personality = personality
         return self
 
+    def override_action_variable(self, action: Action, ctx: Dict[str, Any]) -> Self:
+        """Override action variable with context values."""
+        if action.ctx_override:
+            for k, v in ctx.items():
+                if hasattr(action, k):
+                    setattr(action, k, v)
+
+        return self
+
     async def serve(self, task: Task) -> None:
         """Execute workflow to complete given task.
 
@@ -178,11 +190,12 @@ class WorkFlow(WithBriefing, ToolBoxUsage):
         try:
             # Process each action in sequence
             for i, step in enumerate(self._instances):
-                current_action = step.name
-                logger.info(f"Executing step [{i}] >> {current_action}")
+                logger.info(f"Executing step [{i}] >> {(current_action := step.name)}")
 
                 # Get current context and execute action
                 context = await self._context.get()
+
+                self.override_action_variable(step, context)
                 act_task = create_task(step.act(context))
                 # Handle task cancellation
                 if task.is_cancelled():

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev2}/python/fabricatio/models/adv_kwargs_types.py

@@ -1,7 +1,7 @@
 """A module containing kwargs types for content correction and checking operations."""
 
 from importlib.util import find_spec
-from typing import NotRequired, TypedDict
+from typing import NotRequired, Optional, TypedDict
 
 from fabricatio.models.extra.problem import Improvement
 from fabricatio.models.extra.rule import RuleSet
@@ -58,3 +58,6 @@ if find_spec("pymilvus"):
         collection_name: NotRequired[str | None]
         similarity_threshold: NotRequired[float]
         result_per_query: NotRequired[int]
+        tei_endpoint: NotRequired[Optional[str]]
+        reranker_threshold: NotRequired[float]
+        filter_expr: NotRequired[str]

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev2}/python/fabricatio/models/extra/aricle_rag.py

@@ -1,6 +1,7 @@
 """A Module containing the article rag models."""
 
 import re
+from itertools import groupby
 from pathlib import Path
 from typing import ClassVar, Dict, List, Optional, Self, Unpack
 
@@ -10,12 +11,13 @@ from fabricatio.models.extra.rag import MilvusDataBase
 from fabricatio.models.generic import AsPrompt
 from fabricatio.models.kwargs_types import ChunkKwargs
 from fabricatio.rust import BibManager, blake3_hash, split_into_chunks
-from fabricatio.utils import ok
+from fabricatio.utils import ok, wrapp_in_block
+from more_itertools.more import first
 from more_itertools.recipes import flatten, unique
 from pydantic import Field
 
 
-class ArticleChunk(MilvusDataBase, AsPrompt):
+class ArticleChunk(MilvusDataBase):
     """The chunk of an article."""
 
     etc_word: ClassVar[str] = "等"
@@ -51,10 +53,9 @@ class ArticleChunk(MilvusDataBase, AsPrompt):
     bibtex_cite_key: str
     """The bibtex cite key of the article"""
 
-    def _as_prompt_inner(self) -> Dict[str, str]:
-        return {
-            f"[[{ok(self._cite_number, 'You need to update cite number first.')}]] reference `{self.article_title}` from {self.as_auther_seq()}": self.chunk
-        }
+    @property
+    def reference_header(self) -> str:
+        return f"[[{ok(self._cite_number, 'You need to update cite number first.')}]] reference `{self.article_title}` from {self.as_auther_seq()}"
 
     @property
     def cite_number(self) -> int:
@@ -204,13 +205,23 @@ class CitationManager(AsPrompt):
 
     def set_cite_number_all(self) -> Self:
         """Set citation numbers for all article chunks."""
-        for i, a in enumerate(self.article_chunks, 1):
-            a.update_cite_number(i)
+        number_mapping = {a.bibtex_cite_key: 0 for a in self.article_chunks}
+
+        for i, k in enumerate(number_mapping.keys()):
+            number_mapping[k] = i
+
+        for a in self.article_chunks:
+            a.update_cite_number(number_mapping[a.bibtex_cite_key])
         return self
 
     def _as_prompt_inner(self) -> Dict[str, str]:
         """Generate prompt inner representation."""
-        return {"References": "\n".join(r.as_prompt() for r in self.article_chunks)}
+        seg = []
+        for k, g in groupby(self.article_chunks, key=lambda a: a.bibtex_cite_key):
+            g = list(g)
+            logger.debug(f"Group [{k}]: {len(g)}")
+            seg.append(wrapp_in_block("\n\n".join(a.chunk for a in g), first(g).reference_header))
+        return {"References": "\n".join(seg)}
 
     def apply(self, string: str) -> str:
         """Apply citation replacements to the input string."""
@@ -261,5 +272,9 @@ class CitationManager(AsPrompt):
 
     def unpack_cite_seq(self, citation_seq: List[int]) -> str:
         """Unpack citation sequence into a string."""
-        chunk_seq = [a for a in self.article_chunks if a.cite_number in citation_seq]
-        return "".join(a.as_typst_cite() for a in chunk_seq)
+        chunk_seq = {a.bibtex_cite_key: a for a in self.article_chunks if a.cite_number in citation_seq}
+        return "".join(a.as_typst_cite() for a in chunk_seq.values())
+
+    def as_milvus_filter_expr(self, blacklist: bool = True) -> str:
+        if blacklist:
+            return " and ".join(f'bibtex_cite_key != "{a.bibtex_cite_key}"' for a in self.article_chunks)