fabricatio 0.2.9.dev4__cp312-cp312-win_amd64.whl → 0.2.10.dev1__cp312-cp312-win_amd64.whl

fabricatio/capabilities/rag.py

@@ -3,28 +3,22 @@
 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
+    raise RuntimeError(
+        "pymilvus is not installed. Have you installed `fabricatio[rag]` instead of `fabricatio`?"
+    ) from e
 from functools import lru_cache
 from operator import itemgetter
-from os import PathLike
-from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Self, Union, Unpack, cast, overload
+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.kwargs_types import (
-    ChooseKwargs,
-    CollectionConfigKwargs,
-    EmbeddingKwargs,
-    FetchKwargs,
-    LLMKwargs,
-    RetrievalKwargs,
-)
+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.models.utils import MilvusData
 from fabricatio.rust_instances import TEMPLATE_MANAGER
 from fabricatio.utils import ok
 
@@ -78,40 +72,6 @@ class RAG(EmbeddingUsage):
             raise RuntimeError("Client is not initialized. Have you called `self.init_client()`?")
         return self
 
-    @overload
-    async def pack(
-        self, input_text: List[str], subject: Optional[str] = None, **kwargs: Unpack[EmbeddingKwargs]
-    ) -> List[MilvusData]: ...
-    @overload
-    async def pack(
-        self, input_text: str, subject: Optional[str] = None, **kwargs: Unpack[EmbeddingKwargs]
-    ) -> MilvusData: ...
-
-    async def pack(
-        self, input_text: List[str] | str, subject: Optional[str] = None, **kwargs: Unpack[EmbeddingKwargs]
-    ) -> List[MilvusData] | MilvusData:
-        """Asynchronously generates MilvusData objects for the given input text.
-
-        Args:
-            input_text (List[str] | str): A string or list of strings to generate embeddings for.
-            subject (Optional[str]): The subject of the input text. Defaults to None.
-            **kwargs (Unpack[EmbeddingKwargs]): Additional keyword arguments for embedding.
-
-        Returns:
-            List[MilvusData] | MilvusData: The generated MilvusData objects.
-        """
-        if isinstance(input_text, str):
-            return MilvusData(vector=await self.vectorize(input_text, **kwargs), text=input_text, subject=subject)
-        vecs = await self.vectorize(input_text, **kwargs)
-        return [
-            MilvusData(
-                vector=vec,
-                text=text,
-                subject=subject,
-            )
-            for text, vec in zip(input_text, vecs, strict=True)
-        ]
-
     def view(
         self, collection_name: Optional[str], create: bool = False, **kwargs: Unpack[CollectionConfigKwargs]
     ) -> Self:
@@ -152,29 +112,27 @@ class RAG(EmbeddingUsage):
         Returns:
             str: The name of the collection being viewed.
         """
-        if self.target_collection is None:
-            raise RuntimeError("No collection is being viewed. Have you called `self.view()`?")
-        return self.target_collection
+        return ok(self.target_collection, "No collection is being viewed. Have you called `self.view()`?")
 
-    def add_document[D: Union[Dict[str, Any], MilvusData]](
-        self, data: D | List[D], collection_name: Optional[str] = None, flush: bool = False
+    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], MilvusData] | List[Union[Dict[str, Any], MilvusData]]): The data to be added to the collection.
+            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, MilvusData):
-            prepared_data = data.prepare_insertion()
-        elif isinstance(data, list):
-            prepared_data = [d.prepare_insertion() if isinstance(d, MilvusData) else d for d in data]
-        else:
-            raise TypeError(f"Expected MilvusData or list of MilvusData, got {type(data)}")
+        if isinstance(data, MilvusDataBase):
+            data = [data]
+
+        data_vec = await self.vectorize([d.prepare_vectorization() 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)
 
@@ -183,84 +141,33 @@ class RAG(EmbeddingUsage):
             self.client.flush(c_name)
         return self
 
-    async def consume_file(
-        self,
-        source: List[PathLike] | PathLike,
-        reader: Callable[[PathLike], str] = lambda path: Path(path).read_text(encoding="utf-8"),
-        collection_name: Optional[str] = None,
-    ) -> Self:
-        """Consume a file and add its content to the collection.
-
-        Args:
-            source (PathLike): The path to the file to be consumed.
-            reader (Callable[[PathLike], MilvusData]): The reader function to read the file.
-            collection_name (Optional[str]): The name of the collection. If not provided, the currently viewed collection is used.
-
-        Returns:
-            Self: The current instance, allowing for method chaining.
-        """
-        if not isinstance(source, list):
-            source = [source]
-        return await self.consume_string([reader(s) for s in source], collection_name)
-
-    async def consume_string(self, text: List[str] | str, collection_name: Optional[str] = None) -> Self:
-        """Consume a string and add it to the collection.
-
-        Args:
-            text (List[str] | str): The text to be added to the collection.
-            collection_name (Optional[str]): The name of the collection. If not provided, the currently viewed collection is used.
-
-        Returns:
-            Self: The current instance, allowing for method chaining.
-        """
-        self.add_document(await self.pack(text), collection_name or self.safe_target_collection, flush=True)
-        return self
-
-    @overload
-    async def afetch_document[V: (int, str, float, bytes)](
+    async def afetch_document[D: MilvusDataBase](
         self,
         vecs: List[List[float]],
-        desired_fields: List[str],
+        document_model: Type[D],
         collection_name: Optional[str] = None,
         similarity_threshold: float = 0.37,
         result_per_query: int = 10,
-    ) -> List[Dict[str, V]]: ...
-
-    @overload
-    async def afetch_document[V: (int, str, float, bytes)](
-        self,
-        vecs: List[List[float]],
-        desired_fields: str,
-        collection_name: Optional[str] = None,
-        similarity_threshold: float = 0.37,
-        result_per_query: int = 10,
-    ) -> List[V]: ...
-    async def afetch_document[V: (int, str, float, bytes)](
-        self,
-        vecs: List[List[float]],
-        desired_fields: List[str] | str,
-        collection_name: Optional[str] = None,
-        similarity_threshold: float = 0.37,
-        result_per_query: int = 10,
-    ) -> List[Dict[str, Any]] | List[V]:
-        """Fetch data from the collection.
+    ) -> List[D]:
+        """Asynchronously fetches documents from a Milvus database based on input vectors.
 
         Args:
-            vecs (List[List[float]]): The vectors to search for.
-            desired_fields (List[str] | str): The fields to retrieve.
-            collection_name (Optional[str]): The name of the collection. If not provided, the currently viewed collection is used.
-            similarity_threshold (float): The threshold for similarity, only results above this threshold will be returned.
-            result_per_query (int): The number of results to return per query.
+           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[Dict[str, Any]] | List[Any]: The retrieved data.
+           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=desired_fields if isinstance(desired_fields, list) else [desired_fields],
+            output_fields=list(document_model.model_fields),
             limit=result_per_query,
         )
 
@@ -270,104 +177,42 @@ class RAG(EmbeddingUsage):
         # Step 3: Sort by distance (descending)
         sorted_results = sorted(unique_results, key=itemgetter("distance"), reverse=True)
 
-        logger.debug(f"Searched similarities: {[t['distance'] for t in sorted_results]}")
+        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]
 
-        if isinstance(desired_fields, list):
-            return resp
-        return [r.get(desired_fields) for r in resp]  # extract the single field as list
+        return document_model.from_sequence(resp)
 
-    async def aretrieve(
+    async def aretrieve[D: MilvusDataBase](
         self,
         query: List[str] | str,
+        document_model: Type[D],
         final_limit: int = 20,
         **kwargs: Unpack[FetchKwargs],
-    ) -> List[str]:
+    ) -> List[D]:
         """Retrieve data from the collection.
 
         Args:
             query (List[str] | str): The query to be used for retrieval.
+            document_model (Type[D]): The model class used to convert retrieved data into document objects.
             final_limit (int): The final limit on the number of results to return.
             **kwargs (Unpack[FetchKwargs]): Additional keyword arguments for retrieval.
 
         Returns:
-            List[str]: A list of strings containing the retrieved data.
+            List[D]: A list of document objects created from the retrieved data.
         """
         if isinstance(query, str):
             query = [query]
-        return cast(
-            "List[str]",
+        return (
             await self.afetch_document(
                 vecs=(await self.vectorize(query)),
-                desired_fields="text",
+                document_model=document_model,
                 **kwargs,
-            ),
+            )
         )[:final_limit]
 
-    async def aretrieve_compact(
-        self,
-        query: List[str] | str,
-        **kwargs: Unpack[RetrievalKwargs],
-    ) -> str:
-        """Retrieve data from the collection and format it for display.
-
-        Args:
-            query (List[str] | str): The query to be used for retrieval.
-            **kwargs (Unpack[RetrievalKwargs]): Additional keyword arguments for retrieval.
-
-        Returns:
-            str: A formatted string containing the retrieved data.
-        """
-        return TEMPLATE_MANAGER.render_template(
-            configs.templates.retrieved_display_template, {"docs": (await self.aretrieve(query, **kwargs))}
-        )
-
-    async def aask_retrieved(
-        self,
-        question: str,
-        query: Optional[List[str] | str] = None,
-        collection_name: Optional[str] = None,
-        extra_system_message: str = "",
-        result_per_query: int = 10,
-        final_limit: int = 20,
-        similarity_threshold: float = 0.37,
-        **kwargs: Unpack[LLMKwargs],
-    ) -> str:
-        """Asks a question by retrieving relevant documents based on the provided query.
-
-        This method performs document retrieval using the given query, then asks the
-        specified question using the retrieved documents as context.
-
-        Args:
-            question (str): The question to be asked.
-            query (List[str] | str): The query or list of queries used for document retrieval.
-            collection_name (Optional[str]): The name of the collection to retrieve documents from.
-                                              If not provided, the currently viewed collection is used.
-            extra_system_message (str): An additional system message to be included in the prompt.
-            result_per_query (int): The number of results to return per query. Default is 10.
-            final_limit (int): The maximum number of retrieved documents to consider. Default is 20.
-            similarity_threshold (float): The threshold for similarity, only results above this threshold will be returned.
-            **kwargs (Unpack[LLMKwargs]): Additional keyword arguments passed to the underlying `aask` method.
-
-        Returns:
-            str: A string response generated after asking with the context of retrieved documents.
-        """
-        rendered = await self.aretrieve_compact(
-            query or question,
-            final_limit=final_limit,
-            collection_name=collection_name,
-            result_per_query=result_per_query,
-            similarity_threshold=similarity_threshold,
-        )
-
-        logger.debug(f"Retrieved Documents: \n{rendered}")
-        return await self.aask(
-            question,
-            f"{rendered}\n\n{extra_system_message}",
-            **kwargs,
-        )
-
     async def arefined_query(self, question: List[str] | str, **kwargs: Unpack[ChooseKwargs]) -> Optional[List[str]]:
         """Refines the given question using a template.
 
@@ -385,38 +230,3 @@ class RAG(EmbeddingUsage):
             ),
             **kwargs,
         )
-
-    async def aask_refined(
-        self,
-        question: str,
-        collection_name: Optional[str] = None,
-        extra_system_message: str = "",
-        result_per_query: int = 10,
-        final_limit: int = 20,
-        similarity_threshold: float = 0.37,
-        **kwargs: Unpack[LLMKwargs],
-    ) -> str:
-        """Asks a question using a refined query based on the provided question.
-
-        Args:
-            question (str): The question to be asked.
-            collection_name (Optional[str]): The name of the collection to retrieve documents from.
-            extra_system_message (str): An additional system message to be included in the prompt.
-            result_per_query (int): The number of results to return per query. Default is 10.
-            final_limit (int): The maximum number of retrieved documents to consider. Default is 20.
-            similarity_threshold (float): The threshold for similarity, only results above this threshold will be returned.
-            **kwargs (Unpack[LLMKwargs]): Additional keyword arguments passed to the underlying `aask` method.
-
-        Returns:
-            str: A string response generated after asking with the refined question.
-        """
-        return await self.aask_retrieved(
-            question,
-            await self.arefined_query(question, **kwargs),
-            collection_name=collection_name,
-            extra_system_message=extra_system_message,
-            result_per_query=result_per_query,
-            final_limit=final_limit,
-            similarity_threshold=similarity_threshold,
-            **kwargs,
-        )

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"

fabricatio/decorators.py

@@ -2,6 +2,7 @@
 
 from asyncio import iscoroutinefunction
 from functools import wraps
+from importlib.util import find_spec
 from inspect import signature
 from shutil import which
 from types import ModuleType
@@ -209,3 +210,25 @@ def logging_exec_time[**P, R](func: Callable[P, R]) -> Callable[P, R]:
         return result
 
     return _wrapper
+
+
+def precheck_package[**P, R](package_name: str, msg: str) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """Check if a package exists in the current environment.
+
+    Args:
+        package_name (str): The name of the package to check.
+        msg (str): The message to display if the package is not found.
+
+    Returns:
+        bool: True if the package exists, False otherwise.
+    """
+
+    def _wrapper(func: Callable[P, R]) -> Callable[P, R]:
+        def _inner(*args: P.args, **kwargs: P.kwargs) -> R:
+            if find_spec(package_name):
+                return func(*args, **kwargs)
+            raise RuntimeError(msg)
+
+        return _inner
+
+    return _wrapper

fabricatio/models/adv_kwargs_types.py

@@ -1,4 +1,8 @@
 """A module containing kwargs types for content correction and checking operations."""
+
+from importlib.util import find_spec
+from typing import NotRequired, TypedDict
+
 from fabricatio.models.extra.problem import Improvement
 from fabricatio.models.extra.rule import RuleSet
 from fabricatio.models.generic import SketchedAble
@@ -23,3 +27,34 @@ class CheckKwargs(ReferencedKwargs[Improvement], total=False):
     """
 
     ruleset: RuleSet
+
+
+if find_spec("pymilvus"):
+    from pymilvus import CollectionSchema
+    from pymilvus.milvus_client import IndexParams
+
+    class CollectionConfigKwargs(TypedDict, total=False):
+        """Configuration parameters for a vector collection.
+
+        These arguments are typically used when configuring connections to vector databases.
+        """
+
+        dimension: int | None
+        primary_field_name: str
+        id_type: str
+        vector_field_name: str
+        metric_type: str
+        timeout: float | None
+        schema: CollectionSchema | None
+        index_params: IndexParams | None
+
+    class FetchKwargs(TypedDict):
+        """Arguments for fetching data from vector collections.
+
+        Controls how data is retrieved from vector databases, including filtering
+        and result limiting parameters.
+        """
+
+        collection_name: NotRequired[str | None]
+        similarity_threshold: NotRequired[float]
+        result_per_query: NotRequired[int]

fabricatio/models/events.py

@@ -3,7 +3,7 @@
 from typing import List, Self, Union
 
 from fabricatio.config import configs
-from fabricatio.models.utils import TaskStatus
+from fabricatio.constants import TaskStatus
 from pydantic import BaseModel, ConfigDict, Field
 
 type EventLike = Union[str, List[str], "Event"]
@@ -77,23 +77,23 @@ class Event(BaseModel):
 
     def push_pending(self) -> Self:
         """Push a pending segment to the event."""
-        return self.push(TaskStatus.Pending.value)
+        return self.push(TaskStatus.Pending)
 
     def push_running(self) -> Self:
         """Push a running segment to the event."""
-        return self.push(TaskStatus.Running.value)
+        return self.push(TaskStatus.Running)
 
     def push_finished(self) -> Self:
         """Push a finished segment to the event."""
-        return self.push(TaskStatus.Finished.value)
+        return self.push(TaskStatus.Finished)
 
     def push_failed(self) -> Self:
         """Push a failed segment to the event."""
-        return self.push(TaskStatus.Failed.value)
+        return self.push(TaskStatus.Failed)
 
     def push_cancelled(self) -> Self:
         """Push a cancelled segment to the event."""
-        return self.push(TaskStatus.Cancelled.value)
+        return self.push(TaskStatus.Cancelled)
 
     def pop(self) -> str:
         """Pop a segment from the event."""

fabricatio/models/extra/advanced_judge.py

@@ -12,7 +12,7 @@ class JudgeMent(SketchedAble):
     """
 
     issue_to_judge: str
-    """The issue to be judged, true for affirmation, false for denial."""
+    """The issue to be judged, including the original question and context"""
 
     deny_evidence: List[str]
     """List of clues supporting the denial."""
@@ -21,7 +21,7 @@ class JudgeMent(SketchedAble):
     """List of clues supporting the affirmation."""
 
     final_judgement: bool
-    """The final judgment made according to all extracted clues."""
+    """The final judgment made according to all extracted clues. true for the `issue_to_judge` is correct and false for incorrect."""
 
     def __bool__(self) -> bool:
         """Return the final judgment value.

fabricatio/models/extra/aricle_rag.py

@@ -0,0 +1,120 @@
+"""A Module containing the article rag models."""
+
+from pathlib import Path
+from typing import ClassVar, Dict, List, 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, split_into_chunks
+from fabricatio.utils import ok, wrapp_in_block
+from more_itertools.recipes import flatten
+from pydantic import Field
+
+
+class ArticleChunk(MilvusDataBase, AsPrompt):
+    """The chunk of an article."""
+
+    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 {
+            self.article_title: f"{wrapp_in_block(self.chunk, 'Referring Content')}\n"
+            f"Authors: {';'.join(self.authors)}\n"
+            f"Published Year: {self.year}\n"
+            f"Bibtex Key: {self.bibtex_cite_key}\n",
+        }
+
+    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.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