fabricatio 0.2.10.dev0__cp312-cp312-win_amd64.whl → 0.2.11__cp312-cp312-win_amd64.whl

fabricatio/models/extra/aricle_rag.py

@@ -0,0 +1,254 @@
+"""A Module containing the article rag models."""
+
+import re
+from pathlib import Path
+from typing import ClassVar, Dict, List, Optional, Self, Unpack
+
+from fabricatio.fs import safe_text_read
+from fabricatio.journal import logger
+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, is_chinese, split_into_chunks
+from fabricatio.utils import ok
+from more_itertools.recipes import flatten, unique
+from pydantic import Field
+
+
+class ArticleChunk(MilvusDataBase, AsPrompt):
+    """The chunk of an article."""
+
+    etc_word: ClassVar[str] = "等"
+    and_word: ClassVar[str] = "与"
+    _cite_number: Optional[int] = None
+
+    head_split: ClassVar[List[str]] = [
+        "引 言",
+        "引言",
+        "绪 论",
+        "绪论",
+        "前言",
+        "INTRODUCTION",
+        "Introduction",
+    ]
+    tail_split: ClassVar[List[str]] = [
+        "参 考 文 献",
+        "参  考  文  献",
+        "参考文献",
+        "REFERENCES",
+        "References",
+        "Bibliography",
+        "Reference",
+    ]
+    chunk: str
+    """The segment of the article"""
+    year: int
+    """The year of the article"""
+    authors: List[str] = Field(default_factory=list)
+    """The authors of the article"""
+    article_title: str
+    """The title of the article"""
+    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 cite_number(self) -> int:
+        """Get the cite number."""
+        return ok(self._cite_number, "cite number not set")
+
+    def _prepare_vectorization_inner(self) -> str:
+        return self.chunk
+
+    @classmethod
+    def from_file[P: str | Path](
+        cls, path: P | List[P], bib_mgr: BibManager, **kwargs: Unpack[ChunkKwargs]
+    ) -> List[Self]:
+        """Load the article chunks from the file."""
+        if isinstance(path, list):
+            result = list(flatten(cls._from_file_inner(p, bib_mgr, **kwargs) for p in path))
+            logger.debug(f"Number of chunks created from list of files: {len(result)}")
+            return result
+
+        return cls._from_file_inner(path, bib_mgr, **kwargs)
+
+    @classmethod
+    def _from_file_inner(cls, path: str | Path, bib_mgr: BibManager, **kwargs: Unpack[ChunkKwargs]) -> List[Self]:
+        path = Path(path)
+
+        title_seg = path.stem.split(" - ").pop()
+
+        key = (
+            bib_mgr.get_cite_key_by_title(title_seg)
+            or bib_mgr.get_cite_key_by_title_fuzzy(title_seg)
+            or bib_mgr.get_cite_key_fuzzy(path.stem)
+        )
+        if key is None:
+            logger.warning(f"no cite key found for {path.as_posix()}, skip.")
+            return []
+        authors = ok(bib_mgr.get_author_by_key(key), f"no author found for {key}")
+        year = ok(bib_mgr.get_year_by_key(key), f"no year found for {key}")
+        article_title = ok(bib_mgr.get_title_by_key(key), f"no title found for {key}")
+
+        result = [
+            cls(chunk=c, year=year, authors=authors, article_title=article_title, bibtex_cite_key=key)
+            for c in split_into_chunks(cls.purge_numeric_citation(cls.strip(safe_text_read(path))), **kwargs)
+        ]
+
+        logger.debug(f"Number of chunks created from file {path.as_posix()}: {len(result)}")
+        return result
+
+    @classmethod
+    def strip(cls, string: str) -> str:
+        """Strip the head and tail of the string."""
+        logger.debug(f"String length before strip: {(original := len(string))}")
+        for split in (s for s in cls.head_split if s in string):
+            logger.debug(f"Strip head using {split}")
+            parts = string.split(split)
+            string = split.join(parts[1:]) if len(parts) > 1 else parts[0]
+            break
+        logger.debug(
+            f"String length after head strip: {(stripped_len := len(string))}, decreased by {(d := original - stripped_len)}"
+        )
+        if not d:
+            logger.warning("No decrease at head strip, which is might be abnormal.")
+        for split in (s for s in cls.tail_split if s in string):
+            logger.debug(f"Strip tail using {split}")
+            parts = string.split(split)
+            string = split.join(parts[:-1]) if len(parts) > 1 else parts[0]
+            break
+        logger.debug(f"String length after tail strip: {len(string)}, decreased by {(d := stripped_len - len(string))}")
+        if not d:
+            logger.warning("No decrease at tail strip, which is might be abnormal.")
+
+        return string
+
+    def as_typst_cite(self) -> str:
+        """As typst cite."""
+        return f"#cite(<{self.bibtex_cite_key}>)"
+
+    @staticmethod
+    def purge_numeric_citation(string: str) -> str:
+        """Purge numeric citation."""
+        import re
+
+        return re.sub(r"\[[\d\s,\\~–-]+]", "", string)
+
+    @property
+    def auther_firstnames(self) -> List[str]:
+        """Get the first name of the authors."""
+        ret = []
+        for n in self.authors:
+            if is_chinese(n):
+                ret.append(n[0])
+            else:
+                ret.append(n.split()[-1])
+        return ret
+
+    def as_auther_seq(self) -> str:
+        """Get the auther sequence."""
+        match len(self.authors):
+            case 0:
+                raise ValueError("No authors found")
+            case 1:
+                return f"（{self.auther_firstnames[0]}，{self.year}）{self.as_typst_cite()}"
+            case 2:
+                return f"（{self.auther_firstnames[0]}{self.and_word}{self.auther_firstnames[1]}，{self.year}）{self.as_typst_cite()}"
+            case 3:
+                return f"（{self.auther_firstnames[0]}，{self.auther_firstnames[1]}{self.and_word}{self.auther_firstnames[2]}，{self.year}）{self.as_typst_cite()}"
+            case _:
+                return f"（{self.auther_firstnames[0]}，{self.auther_firstnames[1]}{self.and_word}{self.auther_firstnames[2]}{self.etc_word}，{self.year}）{self.as_typst_cite()}"
+
+    def update_cite_number(self, cite_number: int) -> Self:
+        """Update the cite number."""
+        self._cite_number = cite_number
+        return self
+
+
+class CitationManager(AsPrompt):
+    """Citation manager."""
+
+    article_chunks: List[ArticleChunk] = Field(default_factory=list)
+    """Article chunks."""
+
+    pat: str = r"(\[\[([\d\s,-]*)]])"
+    """Regex pattern to match citations."""
+    sep: str = ","
+    """Separator for citation numbers."""
+    abbr_sep: str = "-"
+    """Separator for abbreviated citation numbers."""
+
+    def update_chunks(
+        self, article_chunks: List[ArticleChunk], set_cite_number: bool = True, dedup: bool = True
+    ) -> Self:
+        """Update article chunks."""
+        self.article_chunks.clear()
+        self.article_chunks.extend(article_chunks)
+        if dedup:
+            self.article_chunks = list(unique(self.article_chunks, lambda c: blake3_hash(c.chunk.encode())))
+        if set_cite_number:
+            self.set_cite_number_all()
+        return self
+
+    def empty(self) -> Self:
+        """Empty the article chunks."""
+        self.article_chunks.clear()
+        return self
+
+    def add_chunks(self, article_chunks: List[ArticleChunk], set_cite_number: bool = True, dedup: bool = True) -> Self:
+        """Add article chunks."""
+        self.article_chunks.extend(article_chunks)
+        if dedup:
+            self.article_chunks = list(unique(self.article_chunks, lambda c: blake3_hash(c.chunk.encode())))
+        if set_cite_number:
+            self.set_cite_number_all()
+        return self
+
+    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)
+        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)}
+
+    def apply(self, string: str) -> str:
+        """Apply citation replacements to the input string."""
+        for origin, m in re.findall(self.pat, string):
+            logger.info(f"Matching citation: {m}")
+            notations = self.convert_to_numeric_notations(m)
+            logger.info(f"Citing Notations: {notations}")
+            citation_number_seq = list(flatten(self.decode_expr(n) for n in notations))
+            logger.info(f"Citation Number Sequence: {citation_number_seq}")
+            dedup = self.deduplicate_citation(citation_number_seq)
+            logger.info(f"Deduplicated Citation Number Sequence: {dedup}")
+            string = string.replace(origin, self.unpack_cite_seq(dedup))
+        return string
+
+    def decode_expr(self, string: str) -> List[int]:
+        """Decode citation expression into a list of integers."""
+        if self.abbr_sep in string:
+            start, end = string.split(self.abbr_sep)
+            return list(range(int(start), int(end) + 1))
+        return [int(string)]
+
+    def convert_to_numeric_notations(self, string: str) -> List[str]:
+        """Convert citation string into numeric notations."""
+        return [s.strip() for s in string.split(self.sep)]
+
+    def deduplicate_citation(self, citation_seq: List[int]) -> List[int]:
+        """Deduplicate citation sequence."""
+        chunk_seq = [a for a in self.article_chunks if a.cite_number in citation_seq]
+        deduped = unique(chunk_seq, lambda a: a.bibtex_cite_key)
+        return [a.cite_number for a in deduped]
+
+    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)

fabricatio/models/extra/article_base.py

@@ -18,6 +18,7 @@ from fabricatio.models.generic import (
     Titled,
     WordCount,
 )
+from fabricatio.rust import comment
 from pydantic import Field
 
 
@@ -29,11 +30,9 @@ class ReferringType(StrEnum):
     SUBSECTION = "subsection"
 
 
-
 type RefKey = Tuple[str, Optional[str], Optional[str]]
 
 
-
 class ArticleMetaData(SketchedAble, Described, WordCount, Titled, Language):
     """Metadata for an article component."""
 
@@ -47,7 +46,16 @@ class ArticleMetaData(SketchedAble, Described, WordCount, Titled, Language):
     aims: List[str]
     """List of writing aims of the research component in academic style."""
 
-
+    @property
+    def typst_metadata_comment(self) -> str:
+        """Generates a comment for the metadata of the article component."""
+        return (
+            (comment(f"Desc:\n{self.description}\n") if self.description else "")
+            + (comment(f"Aims:\n{'\n  '.join(self.aims)}\n") if self.aims else "")
+            + (comment(f"Expected Word Count:{self.expected_word_count}") if self.expected_word_count else "")
+            if self.expected_word_count
+            else ""
+        )
 
 
 class ArticleOutlineBase(
@@ -92,7 +100,7 @@ class SubSectionBase(ArticleOutlineBase):
 
     def to_typst_code(self) -> str:
         """Converts the component into a Typst code snippet for rendering."""
-        return f"=== {self.title}\n"
+        return f"=== {self.title}\n{self.typst_metadata_comment}\n"
 
     def introspect(self) -> str:
         """Introspects the article subsection outline."""
@@ -117,7 +125,9 @@ class SectionBase[T: SubSectionBase](ArticleOutlineBase):
         Returns:
             str: The formatted Typst code snippet.
         """
-        return f"== {self.title}\n" + "\n\n".join(subsec.to_typst_code() for subsec in self.subsections)
+        return f"== {self.title}\n{self.typst_metadata_comment}\n" + "\n\n".join(
+            subsec.to_typst_code() for subsec in self.subsections
+        )
 
     def resolve_update_conflict(self, other: Self) -> str:
         """Resolve update errors in the article outline."""
@@ -160,7 +170,9 @@ class ChapterBase[T: SectionBase](ArticleOutlineBase):
 
     def to_typst_code(self) -> str:
         """Converts the chapter into a Typst formatted code snippet for rendering."""
-        return f"= {self.title}\n" + "\n\n".join(sec.to_typst_code() for sec in self.sections)
+        return f"= {self.title}\n{self.typst_metadata_comment}\n" + "\n\n".join(
+            sec.to_typst_code() for sec in self.sections
+        )
 
     def resolve_update_conflict(self, other: Self) -> str:
         """Resolve update errors in the article outline."""
@@ -302,4 +314,41 @@ class ArticleBase[T: ChapterBase](FinalizedDumpAble, AsPrompt, WordCount, Descri
             === Implementation Details
             == Evaluation Protocol
         """
-        return "\n\n".join(a.to_typst_code() for a in self.chapters)
+        return (
+            comment(
+                f"Title:{self.title}\n"
+                + (f"Desc:\n{self.description}\n" if self.description else "")
+                + f"Word Count:{self.expected_word_count}"
+                if self.expected_word_count
+                else ""
+            )
+            + "\n\n"
+            + "\n\n".join(a.to_typst_code() for a in self.chapters)
+        )
+
+    def avg_chap_wordcount[S](self:S) -> S:
+        """Set all chap have same word count sum up to be `self.expected_word_count`."""
+        avg = int(self.expected_word_count / len(self.chapters))
+        for c in self.chapters:
+            c.expected_word_count = avg
+        return self
+
+    def avg_sec_wordcount[S](self:S) -> S:
+        """Set all sec have same word count sum up to be `self.expected_word_count`."""
+        for c in self.chapters:
+            avg = int(c.expected_word_count / len(c.sections))
+            for s in c.sections:
+                s.expected_word_count = avg
+        return self
+
+    def avg_subsec_wordcount[S](self:S) -> S:
+        """Set all subsec have same word count sum up to be `self.expected_word_count`."""
+        for _, s in self.iter_sections():
+            avg = int(s.expected_word_count / len(s.subsections))
+            for ss in s.subsections:
+                ss.expected_word_count = avg
+        return self
+
+    def avg_wordcount_recursive(self) -> Self:
+        """Set all chap, sec, subsec have same word count sum up to be `self.expected_word_count`."""
+        return self.avg_chap_wordcount().avg_sec_wordcount().avg_sec_wordcount()

fabricatio/models/extra/article_essence.py

@@ -1,8 +1,9 @@
 """ArticleEssence: Semantic fingerprint of academic paper for structured analysis."""
 
-from typing import List, Self
+from typing import List
 
-from fabricatio.models.generic import Display, PersistentAble, ProposedAble, Vectorizable
+from fabricatio.models.extra.rag import MilvusDataBase
+from fabricatio.models.generic import PersistentAble, SketchedAble
 from pydantic import BaseModel
 
 
@@ -54,7 +55,7 @@ class Highlightings(BaseModel):
     """
 
 
-class ArticleEssence(ProposedAble, Display, PersistentAble, Vectorizable):
+class ArticleEssence(SketchedAble, PersistentAble, MilvusDataBase):
     """Structured representation of a scientific article's core elements in its original language."""
 
     language: str
@@ -93,7 +94,7 @@ class ArticleEssence(ProposedAble, Display, PersistentAble, Vectorizable):
     bibtex_cite_key: str
     """Bibtex cite key of the original article."""
 
-    def update_cite_key(self, new_cite_key: str) -> Self:
-        """Update the bibtex_cite_key of the article."""
-        self.bibtex_cite_key = new_cite_key
-        return self
+    def _prepare_vectorization_inner(self) -> str:
+        return self.compact()
+
+

fabricatio/models/extra/article_main.py

@@ -2,6 +2,7 @@
 
 from typing import Dict, Generator, List, Self, Tuple, override
 
+from fabricatio.decorators import precheck_package
 from fabricatio.fs.readers import extract_sections
 from fabricatio.journal import logger
 from fabricatio.models.extra.article_base import (
@@ -11,11 +12,15 @@ from fabricatio.models.extra.article_base import (
     SubSectionBase,
 )
 from fabricatio.models.extra.article_outline import (
+    ArticleChapterOutline,
     ArticleOutline,
+    ArticleSectionOutline,
+    ArticleSubsectionOutline,
 )
 from fabricatio.models.generic import Described, PersistentAble, SequencePatch, SketchedAble, WithRef, WordCount
-from fabricatio.rust import word_count
-from pydantic import Field
+from fabricatio.rust import convert_all_block_tex, convert_all_inline_tex, word_count
+from fabricatio.utils import fallback_kwargs
+from pydantic import Field, NonNegativeInt
 
 PARAGRAPH_SEP = "// - - -"
 
@@ -23,6 +28,9 @@ PARAGRAPH_SEP = "// - - -"
 class Paragraph(SketchedAble, WordCount, Described):
     """Structured academic paragraph blueprint for controlled content generation."""
 
+    expected_word_count: NonNegativeInt = 0
+    """The expected word count of this paragraph, 0 means not specified"""
+
     description: str = Field(
         alias="elaboration",
         description=Described.model_fields["description"].description,
@@ -85,7 +93,7 @@ class ArticleSubsection(SubSectionBase):
         Returns:
             str: Typst code snippet for rendering.
         """
-        return f"=== {self.title}\n" + f"\n{PARAGRAPH_SEP}\n".join(p.content for p in self.paragraphs)
+        return super().to_typst_code() + f"\n\n{PARAGRAPH_SEP}\n\n".join(p.content for p in self.paragraphs)
 
     @classmethod
     def from_typst_code(cls, title: str, body: str) -> Self:
@@ -153,10 +161,74 @@ class Article(
             "Original Article": self.display(),
         }
 
+    def convert_tex(self) -> Self:
+        """Convert tex to typst code."""
+        for _, _, subsec in self.iter_subsections():
+            for p in subsec.paragraphs:
+                p.content = convert_all_inline_tex(p.content)
+                p.content = convert_all_block_tex(p.content)
+        return self
+
+    def fix_wrapper(self) -> Self:
+        """Fix wrapper."""
+        for _, _, subsec in self.iter_subsections():
+            for p in subsec.paragraphs:
+                p.content = (
+                    p.content.replace(r" \( ", "$")
+                    .replace(r" \) ", "$")
+                    .replace("\\[\n", "$$\n")
+                    .replace("\n\\]", "\n$$")
+                )
+        return self
+
     @override
     def iter_subsections(self) -> Generator[Tuple[ArticleChapter, ArticleSection, ArticleSubsection], None, None]:
         return super().iter_subsections()  # pyright: ignore [reportReturnType]
 
+    def extrac_outline(self) -> ArticleOutline:
+        """Extract outline from article."""
+        # Create an empty list to hold chapter outlines
+        chapters = []
+
+        # Iterate through each chapter in the article
+        for chapter in self.chapters:
+            # Create an empty list to hold section outlines
+            sections = []
+
+            # Iterate through each section in the chapter
+            for section in chapter.sections:
+                # Create an empty list to hold subsection outlines
+                subsections = []
+
+                # Iterate through each subsection in the section
+                for subsection in section.subsections:
+                    # Create a subsection outline and add it to the list
+                    subsections.append(
+                        ArticleSubsectionOutline(**subsection.model_dump(exclude={"paragraphs"}, by_alias=True))
+                    )
+
+                # Create a section outline and add it to the list
+                sections.append(
+                    ArticleSectionOutline(
+                        **section.model_dump(exclude={"subsections"}, by_alias=True),
+                        subsections=subsections,
+                    )
+                )
+
+            # Create a chapter outline and add it to the list
+            chapters.append(
+                ArticleChapterOutline(
+                    **chapter.model_dump(exclude={"sections"}, by_alias=True),
+                    sections=sections,
+                )
+            )
+
+        # Create and return the article outline
+        return ArticleOutline(
+            **self.model_dump(exclude={"chapters"}, by_alias=True),
+            chapters=chapters,
+        )
+
     @classmethod
     def from_outline(cls, outline: ArticleOutline) -> "Article":
         """Generates an article from the given outline.
@@ -194,13 +266,37 @@ class Article(
         return article
 
     @classmethod
-    def from_typst_code(cls, title: str, body: str) -> Self:
+    def from_typst_code(cls, title: str, body: str, **kwargs) -> Self:
         """Generates an article from the given Typst code."""
         return cls(
             chapters=[
                 ArticleChapter.from_typst_code(*pack) for pack in extract_sections(body, level=1, section_char="=")
             ],
             heading=title,
-            expected_word_count=word_count(body),
-            abstract="",
+            **fallback_kwargs(
+                kwargs,
+                expected_word_count=word_count(body),
+                abstract="",
+            ),
         )
+
+    @classmethod
+    def from_mixed_source(cls, article_outline: ArticleOutline, typst_code: str) -> Self:
+        """Generates an article from the given outline and Typst code."""
+        self = cls.from_typst_code(article_outline.title, typst_code)
+        self.expected_word_count = article_outline.expected_word_count
+        self.description = article_outline.description
+        for a, o in zip(self.iter_dfs(), article_outline.iter_dfs(), strict=True):
+            a.update_metadata(o)
+        return self.update_ref(article_outline)
+
+    @precheck_package(
+        "questionary", "'questionary' is required to run this function. Have you installed `fabricatio[qa]`?."
+    )
+    async def edit_titles(self) -> Self:
+        """Edits the titles of the article."""
+        from questionary import text
+
+        for a in self.iter_dfs():
+            a.title = await text(f"Edit `{a.title}`.", default=a.title).ask_async() or a.title
+        return self

fabricatio/models/extra/problem.py

@@ -7,7 +7,6 @@ from fabricatio.journal import logger
 from fabricatio.models.generic import SketchedAble, WithBriefing
 from fabricatio.utils import ask_edit
 from pydantic import Field
-from questionary import Choice, checkbox, text
 from rich import print as r_print
 
 
@@ -74,6 +73,9 @@ class ProblemSolutions(SketchedAble):
         return len(self.solutions) > 0
 
     async def edit_problem(self) -> Self:
+        """Interactively edit the problem description."""
+        from questionary import text
+
         """Interactively edit the problem description."""
         self.problem = Problem.model_validate_strings(
             await text("Please edit the problem below:", default=self.problem.display()).ask_async()
@@ -127,6 +129,8 @@ class Improvement(SketchedAble):
         Returns:
             Self: The current instance with filtered problems and solutions.
         """
+        from questionary import Choice, checkbox
+
         # Choose the problems to retain
         chosen_ones: List[ProblemSolutions] = await checkbox(
             "Please choose the problems you want to retain.(Default: retain all)",