fabricatio 0.2.13.dev0__tar.gz → 0.2.13.dev1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fabricatio
-Version: 0.2.13.dev0
+Version: 0.2.13.dev1
 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.dev1}/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.dev1}/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,7 +26,7 @@ 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_model="openai/qwen-plus",
     llm_temperature=0.45,
     # llm_api_endpoint=HttpUrl("https://dashscope.aliyuncs.com/compatible-mode/v1"),
     llm_top_p=0.95,
@@ -74,7 +82,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 +99,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 +205,26 @@ 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:
+    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.dev1}/pyproject.toml

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

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev1}/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.dev1}/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
 
@@ -257,6 +257,7 @@ class WriteArticleContentRAG(Action, RAG, Extract):
 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."""
@@ -270,6 +271,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 +302,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.dev1}/python/fabricatio/capabilities/rag.py

@@ -143,21 +143,25 @@ 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,
     ) -> 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.
 
         Returns:
            List[D]: A list of document objects created from the fetched data.
@@ -165,15 +169,34 @@ 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),
             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])
+                retrieved_id.update(res["id"] for res in g)
+                rank_scores = await reranker.arerank(q, [m.prepare_vectorization() for m in models])
+                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 +228,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.dev1}/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.dev1}/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.dev1}/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,5 @@ 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]

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev1}/python/fabricatio/models/extra/article_base.py

@@ -2,8 +2,9 @@
 
 from abc import ABC
 from enum import StrEnum
-from typing import Generator, List, Optional, Self, Tuple
+from typing import ClassVar, Generator, List, Optional, Self, Tuple, Type
 
+from fabricatio.fs.readers import extract_sections
 from fabricatio.models.generic import (
     AsPrompt,
     Described,
@@ -105,12 +106,6 @@ class ArticleOutlineBase(
         self.description = other.description
         return self
 
-    def display_metadata(self) -> str:
-        """Displays the metadata of the current instance."""
-        return self.model_dump_json(
-            indent=1, include={"title", "writing_aim", "description", "support_to", "depend_on"}
-        )
-
     def update_from_inner(self, other: Self) -> Self:
         """Updates the current instance with the attributes of another instance."""
         return self.update_metadata(other)
@@ -140,6 +135,8 @@ class SectionBase[T: SubSectionBase](ArticleOutlineBase):
     subsections: List[T]
     """Subsections of the section. Contains at least one subsection. You can also add more as needed."""
 
+    child_type: ClassVar[Type[SubSectionBase]]
+
     def to_typst_code(self) -> str:
         """Converts the section into a Typst formatted code snippet.
 
@@ -148,6 +145,17 @@ class SectionBase[T: SubSectionBase](ArticleOutlineBase):
         """
         return f"== {super().to_typst_code()}" + "\n\n".join(subsec.to_typst_code() for subsec in self.subsections)
 
+    @classmethod
+    def from_typst_code(cls, title: str, body: str, **kwargs) -> Self:
+        """Creates an Article object from the given Typst code."""
+        return super().from_typst_code(
+            title,
+            body,
+            subsections=[
+                cls.child_type.from_typst_code(*pack) for pack in extract_sections(body, level=3, section_char="=")
+            ],
+        )
+
     def resolve_update_conflict(self, other: Self) -> str:
         """Resolve update errors in the article outline."""
         out = ""
@@ -186,11 +194,23 @@ class ChapterBase[T: SectionBase](ArticleOutlineBase):
 
     sections: List[T]
     """Sections of the chapter. Contains at least one section. You can also add more as needed."""
+    child_type: ClassVar[Type[SectionBase]]
 
     def to_typst_code(self) -> str:
         """Converts the chapter into a Typst formatted code snippet for rendering."""
         return f"= {super().to_typst_code()}" + "\n\n".join(sec.to_typst_code() for sec in self.sections)
 
+    @classmethod
+    def from_typst_code(cls, title: str, body: str, **kwargs) -> Self:
+        """Creates an Article object from the given Typst code."""
+        return super().from_typst_code(
+            title,
+            body,
+            sections=[
+                cls.child_type.from_typst_code(*pack) for pack in extract_sections(body, level=2, section_char="=")
+            ],
+        )
+
     def resolve_update_conflict(self, other: Self) -> str:
         """Resolve update errors in the article outline."""
         out = ""
@@ -238,6 +258,19 @@ class ArticleBase[T: ChapterBase](FinalizedDumpAble, AsPrompt, FromTypstCode, To
     chapters: List[T]
     """Chapters of the article. Contains at least one chapter. You can also add more as needed."""
 
+    child_type: ClassVar[Type[ChapterBase]]
+
+    @classmethod
+    def from_typst_code(cls, title: str, body: str, **kwargs) -> Self:
+        """Generates an article from the given Typst code."""
+        return super().from_typst_code(
+            title,
+            body,
+            chapters=[
+                cls.child_type.from_typst_code(*pack) for pack in extract_sections(body, level=1, section_char="=")
+            ],
+        )
+
     def iter_dfs_rev(
         self,
     ) -> Generator[ArticleOutlineBase, None, None]:

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev1}/python/fabricatio/models/extra/article_main.py

@@ -1,9 +1,8 @@
 """ArticleBase and ArticleSubsection classes for managing hierarchical document components."""
 
-from typing import Dict, Generator, List, Self, Tuple, override
+from typing import ClassVar, Dict, Generator, List, Self, Tuple, Type, 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 (
     ArticleBase,
@@ -52,6 +51,10 @@ class Paragraph(SketchedAble, WordCount, Described):
         """Create a Paragraph object from the given content."""
         return cls(elaboration="", aims=[], expected_word_count=word_count(content), content=content)
 
+    @property
+    def exact_wordcount(self) -> int:
+        return word_count(self.content)
+
 
 class ArticleParagraphSequencePatch(SequencePatch[Paragraph]):
     """Patch for `Paragraph` list of `ArticleSubsection`."""
@@ -115,31 +118,13 @@ class ArticleSubsection(SubSectionBase):
 class ArticleSection(SectionBase[ArticleSubsection]):
     """Atomic argumentative unit with high-level specificity."""
 
-    @classmethod
-    def from_typst_code(cls, title: str, body: str, **kwargs) -> Self:
-        """Creates an Article object from the given Typst code."""
-        return super().from_typst_code(
-            title,
-            body,
-            subsections=[
-                ArticleSubsection.from_typst_code(*pack) for pack in extract_sections(body, level=3, section_char="=")
-            ],
-        )
+    child_type: ClassVar[Type[SubSectionBase]] = ArticleSubsection
 
 
 class ArticleChapter(ChapterBase[ArticleSection]):
     """Thematic progression implementing research function."""
 
-    @classmethod
-    def from_typst_code(cls, title: str, body: str, **kwargs) -> Self:
-        """Creates an Article object from the given Typst code."""
-        return super().from_typst_code(
-            title,
-            body,
-            sections=[
-                ArticleSection.from_typst_code(*pack) for pack in extract_sections(body, level=2, section_char="=")
-            ],
-        )
+    child_type: ClassVar[Type[SectionBase]] = ArticleSection
 
 
 class Article(
@@ -153,6 +138,8 @@ class Article(
     aiming to provide a comprehensive model for academic papers.
     """
 
+    child_type: ClassVar[Type[ChapterBase]] = ArticleChapter
+
     def _as_prompt_inner(self) -> Dict[str, str]:
         return {
             "Original Article Briefing": self.referenced.referenced.referenced,
@@ -261,17 +248,6 @@ class Article(
             article.chapters.append(article_chapter)
         return article
 
-    @classmethod
-    def from_typst_code(cls, title: str, body: str, **kwargs) -> Self:
-        """Generates an article from the given Typst code."""
-        return super().from_typst_code(
-            title,
-            body,
-            chapters=[
-                ArticleChapter.from_typst_code(*pack) for pack in extract_sections(body, level=1, section_char="=")
-            ],
-        )
-
     @classmethod
     def from_mixed_source(cls, article_outline: ArticleOutline, typst_code: str) -> Self:
         """Generates an article from the given outline and Typst code."""
@@ -292,3 +268,15 @@ class Article(
         for a in self.iter_dfs():
             a.title = await text(f"Edit `{a.title}`.", default=a.title).ask_async() or a.title
         return self
+
+    def check_short_paragraphs(self, threshold: int = 60) -> str:
+        """Checks for short paragraphs in the article."""
+        err = []
+        for chap, sec, subsec in self.iter_subsections():
+            for i, p in enumerate(subsec.paragraphs):
+                if p.exact_wordcount <= threshold:
+                    err.append(
+                        f"{chap.title}->{sec.title}->{subsec.title}-> Paragraph [{i}] is too short, {p.exact_wordcount} words."
+                    )
+
+        return "\n".join(err)

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev1}/python/fabricatio/models/generic.py

@@ -312,11 +312,11 @@ class Language(Base):
     @property
     def language(self) -> str:
         """Get the language of the object."""
-        if isinstance(self, Described):
+        if isinstance(self, Described) and self.description:
             return detect_language(self.description)
-        if isinstance(self, Titled):
+        if isinstance(self, Titled) and self.title:
             return detect_language(self.title)
-        if isinstance(self, Named):
+        if isinstance(self, Named) and self.name:
             return detect_language(self.name)
 
         return detect_language(self.model_dump_json(by_alias=True))

{fabricatio-0.2.13.dev0 → fabricatio-0.2.13.dev1}/python/fabricatio/models/kwargs_types.py

@@ -1,6 +1,6 @@
 """This module contains the types for the keyword arguments of the methods in the models module."""
 
-from typing import Any, Dict, List, NotRequired, Optional, Required, TypedDict
+from typing import Any, Dict, List, Literal, NotRequired, Optional, Required, TypedDict
 
 from litellm.caching.caching import CacheMode
 from litellm.types.caching import CachingSupportedCallTypes
@@ -164,3 +164,11 @@ class CacheKwargs(TypedDict, total=False):
     qdrant_collection_name: str
     qdrant_quantization_config: str
     qdrant_semantic_cache_embedding_model: str
+
+
+class RerankOptions(TypedDict, total=False):
+    """Optional keyword arguments for the rerank method."""
+
+    raw_scores: bool
+    truncate: bool
+    truncation_direction: Literal["Left", "Right"]