fabricatio 0.2.3.dev3__cp312-cp312-manylinux_2_34_x86_64.whl → 0.2.4.dev1__cp312-cp312-manylinux_2_34_x86_64.whl

fabricatio/__init__.py

@@ -3,11 +3,13 @@
 from importlib.util import find_spec
 
 from fabricatio._rust_instances import template_manager
+from fabricatio.actions import ExtractArticleEssence
 from fabricatio.core import env
 from fabricatio.fs import magika
 from fabricatio.journal import logger
 from fabricatio.models.action import Action, WorkFlow
 from fabricatio.models.events import Event
+from fabricatio.models.extra import ArticleEssence
 from fabricatio.models.role import Role
 from fabricatio.models.task import Task
 from fabricatio.models.tool import ToolBox
@@ -17,9 +19,11 @@ from fabricatio.toolboxes import arithmetic_toolbox, basic_toolboxes, fs_toolbox
 
 __all__ = [
     "Action",
+    "ArticleEssence",
     "Capture",
     "CodeBlockCapture",
     "Event",
+    "ExtractArticleEssence",
     "JsonCapture",
     "Message",
     "Messages",
@@ -40,6 +44,6 @@ __all__ = [
 
 
 if find_spec("pymilvus"):
-    from fabricatio.capabilities.rag import Rag
+    from fabricatio.capabilities.rag import RAG
 
-    __all__ += ["Rag"]
+    __all__ += ["RAG"]

fabricatio/actions/__init__.py

@@ -1,5 +1,5 @@
 """module for actions."""
 
-from fabricatio.actions.transmission import PublishTask
+from fabricatio.actions.article import ExtractArticleEssence
 
-__all__ = ["PublishTask"]
+__all__ = ["ExtractArticleEssence"]

fabricatio/actions/article.py

@@ -0,0 +1,44 @@
+"""Actions for transmitting tasks to targets."""
+
+from os import PathLike
+from pathlib import Path
+from typing import Callable, List
+
+from fabricatio.journal import logger
+from fabricatio.models.action import Action
+from fabricatio.models.extra import ArticleEssence
+from fabricatio.models.task import Task
+
+
+class ExtractArticleEssence(Action):
+    """Extract the essence of article(s)."""
+
+    name: str = "extract article essence"
+    """The name of the action."""
+    description: str = "Extract the essence of an article. output as json"
+    """The description of the action."""
+
+    output_key: str = "article_essence"
+    """The key of the output data."""
+
+    async def _execute[P: PathLike | str](
+        self,
+        task_input: Task,
+        reader: Callable[[P], str] = lambda p: Path(p).read_text(encoding="utf-8"),
+        **_,
+    ) -> List[ArticleEssence]:
+        if not await self.ajudge(
+            f"= Task\n{task_input.briefing}\n\n\n= Role\n{self.briefing}",
+            affirm_case="The task does not violate the role, and could be approved since the file dependencies are specified.",
+            deny_case="The task does violate the role, and could not be approved.",
+        ):
+            logger.info(err := "Task not approved.")
+            raise RuntimeError(err)
+
+        # trim the references
+        contents = ["References".join(c.split("References")[:-1]) for c in map(reader, task_input.dependencies)]
+        return await self.propose(
+            ArticleEssence,
+            contents,
+            system_message=f"# your personal briefing: \n{self.briefing}",
+        )

fabricatio/capabilities/propose.py

@@ -0,0 +1,55 @@
+"""A module for the task capabilities of the Fabricatio library."""
+
+from typing import List, Type, Unpack, overload
+
+from fabricatio.models.generic import ProposedAble
+from fabricatio.models.kwargs_types import GenerateKwargs
+from fabricatio.models.usages import LLMUsage
+
+
+class Propose[M: ProposedAble](LLMUsage):
+    """A class that proposes an Obj based on a prompt."""
+
+    @overload
+    async def propose(
+        self,
+        cls: Type[M],
+        prompt: List[str],
+        **kwargs: Unpack[GenerateKwargs],
+    ) -> List[M]: ...
+
+    @overload
+    async def propose(
+        self,
+        cls: Type[M],
+        prompt: str,
+        **kwargs: Unpack[GenerateKwargs],
+    ) -> M: ...
+
+    async def propose(
+        self,
+        cls: Type[M],
+        prompt: List[str] | str,
+        **kwargs: Unpack[GenerateKwargs],
+    ) -> List[M] | M:
+        """Asynchronously proposes a task based on a given prompt and parameters.
+
+        Parameters:
+            cls: The class type of the task to be proposed.
+            prompt: The prompt text for proposing a task, which is a string that must be provided.
+            **kwargs: The keyword arguments for the LLM (Large Language Model) usage.
+
+        Returns:
+            A Task object based on the proposal result.
+        """
+        if isinstance(prompt, str):
+            return await self.aask_validate(
+                question=cls.create_json_prompt(prompt),
+                validator=cls.instantiate_from_string,
+                **kwargs,
+            )
+        return await self.aask_validate_batch(
+            questions=[cls.create_json_prompt(p) for p in prompt],
+            validator=cls.instantiate_from_string,
+            **kwargs,
+        )

fabricatio/capabilities/rag.py

@@ -1,71 +1,114 @@
 """A module for the RAG (Retrieval Augmented Generation) model."""
 
+try:
+    from pymilvus import MilvusClient
+except ImportError as e:
+    raise RuntimeError("pymilvus is not installed. Have you installed `fabricatio[rag]` instead of `fabricatio`") from e
 from functools import lru_cache
 from operator import itemgetter
 from os import PathLike
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Self, Union, Unpack
+from typing import Any, Callable, Dict, List, Optional, Self, Union, Unpack, overload
 
-from fabricatio import template_manager
+from fabricatio._rust_instances import template_manager
 from fabricatio.config import configs
-from fabricatio.models.kwargs_types import LLMKwargs
-from fabricatio.models.usages import LLMUsage
+from fabricatio.journal import logger
+from fabricatio.models.kwargs_types import CollectionSimpleConfigKwargs, EmbeddingKwargs, FetchKwargs, LLMKwargs
+from fabricatio.models.usages import EmbeddingUsage
 from fabricatio.models.utils import MilvusData
 from more_itertools.recipes import flatten
-
-try:
-    from pymilvus import MilvusClient
-except ImportError as e:
-    raise RuntimeError("pymilvus is not installed. Have you installed `fabricatio[rag]` instead of `fabricatio`") from e
 from pydantic import Field, PrivateAttr
 
 
 @lru_cache(maxsize=None)
-def create_client(
-    uri: Optional[str] = None, token: Optional[str] = None, timeout: Optional[float] = None
-) -> MilvusClient:
+def create_client(uri: str, token: str = "", timeout: Optional[float] = None) -> MilvusClient:
     """Create a Milvus client."""
     return MilvusClient(
-        uri=uri or configs.rag.milvus_uri.unicode_string(),
-        token=token or configs.rag.milvus_token.get_secret_value() if configs.rag.milvus_token else "",
-        timeout=timeout or configs.rag.milvus_timeout,
+        uri=uri,
+        token=token,
+        timeout=timeout,
     )
 
 
-class Rag(LLMUsage):
+class RAG(EmbeddingUsage):
     """A class representing the RAG (Retrieval Augmented Generation) model."""
 
-    milvus_uri: Optional[str] = Field(default=None, frozen=True)
-    """The URI of the Milvus server."""
-    milvus_token: Optional[str] = Field(default=None, frozen=True)
-    """The token for the Milvus server."""
-    milvus_timeout: Optional[float] = Field(default=None, frozen=True)
-    """The timeout for the Milvus server."""
     target_collection: Optional[str] = Field(default=None)
     """The name of the collection being viewed."""
 
-    _client: MilvusClient = PrivateAttr(None)
+    _client: Optional[MilvusClient] = PrivateAttr(None)
     """The Milvus client used for the RAG model."""
 
     @property
     def client(self) -> MilvusClient:
         """Return the Milvus client."""
+        if self._client is None:
+            raise RuntimeError("Client is not initialized. Have you called `self.init_client()`?")
         return self._client
 
-    def model_post_init(self, __context: Any) -> None:
-        """Initialize the RAG model by creating the collection if it does not exist."""
-        self._client = create_client(self.milvus_uri, self.milvus_token, self.milvus_timeout)
-        self.view(self.target_collection, create=True)
+    def init_client(
+        self,
+        milvus_uri: Optional[str] = None,
+        milvus_token: Optional[str] = None,
+        milvus_timeout: Optional[float] = None,
+    ) -> Self:
+        """Initialize the Milvus client."""
+        self._client = create_client(
+            uri=milvus_uri or (self.milvus_uri or configs.rag.milvus_uri).unicode_string(),
+            token=milvus_token
+            or (token.get_secret_value() if (token := (self.milvus_token or configs.rag.milvus_token)) else ""),
+            timeout=milvus_timeout or self.milvus_timeout,
+        )
+        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.
 
-    def view(self, collection_name: Optional[str], create: bool = False) -> Self:
+        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[CollectionSimpleConfigKwargs]
+    ) -> Self:
         """View the specified collection.
 
         Args:
             collection_name (str): The name of the collection.
             create (bool): Whether to create the collection if it does not exist.
+            **kwargs (Unpack[CollectionSimpleConfigKwargs]): Additional keyword arguments for collection configuration.
         """
         if create and collection_name and not self._client.has_collection(collection_name):
-            self._client.create_collection(collection_name)
+            kwargs["dimension"] = kwargs.get("dimension") or self.milvus_dimensions or configs.rag.milvus_dimensions
+            self._client.create_collection(collection_name, auto_id=True, **kwargs)
+            logger.info(f"Creating collection {collection_name}")
 
         self.target_collection = collection_name
         return self
@@ -90,13 +133,14 @@ class Rag(LLMUsage):
         return self.target_collection
 
     def add_document[D: Union[Dict[str, Any], MilvusData]](
-        self, data: D | List[D], collection_name: Optional[str] = None
+        self, data: D | List[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.
             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.
@@ -105,11 +149,19 @@ class Rag(LLMUsage):
             data = data.prepare_insertion()
         if isinstance(data, list):
             data = [d.prepare_insertion() if isinstance(d, MilvusData) else d for d in data]
-        self._client.insert(collection_name or self.safe_target_collection, data)
+        c_name = collection_name or self.safe_target_collection
+        self._client.insert(c_name, data)
+
+        if flush:
+            logger.debug(f"Flushing collection {c_name}")
+            self._client.flush(c_name)
         return self
 
-    def consume(
-        self, source: PathLike, reader: Callable[[PathLike], MilvusData], collection_name: Optional[str] = None
+    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.
 
@@ -121,8 +173,21 @@ class Rag(LLMUsage):
         Returns:
             Self: The current instance, allowing for method chaining.
         """
-        data = reader(Path(source))
-        self.add_document(data, collection_name or self.safe_target_collection)
+        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
 
     async def afetch_document(
@@ -130,6 +195,7 @@ class Rag(LLMUsage):
         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[Any]:
         """Fetch data from the collection.
@@ -138,6 +204,7 @@ class Rag(LLMUsage):
             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.
 
         Returns:
@@ -147,6 +214,7 @@ class Rag(LLMUsage):
         search_results = self._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],
             limit=result_per_query,
         )
@@ -157,6 +225,7 @@ class Rag(LLMUsage):
         # Step 3: Sort by distance (descending)
         sorted_results = sorted(flattened_results, key=itemgetter("distance"), reverse=True)
 
+        logger.debug(f"Searched similarities: {[t['distance'] for t in sorted_results]}")
         # Step 4: Extract the entities
         resp = [result["entity"] for result in sorted_results]
 
@@ -168,27 +237,29 @@ class Rag(LLMUsage):
         self,
         query: List[str] | str,
         collection_name: Optional[str] = None,
-        result_per_query: int = 10,
         final_limit: int = 20,
+        **kwargs: Unpack[FetchKwargs],
     ) -> List[str]:
         """Retrieve data from the collection.
 
         Args:
             query (List[str] | str): The query to be used for retrieval.
             collection_name (Optional[str]): The name of the collection. If not provided, the currently viewed collection is used.
-            result_per_query (int): The number of results to be returned per query.
             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.
         """
         if isinstance(query, str):
             query = [query]
-        return await self.afetch_document(
-            vecs=(await self.vectorize(query)),
-            desired_fields="text",
-            collection_name=collection_name,
-            result_per_query=result_per_query,
+        return (
+            await self.afetch_document(
+                vecs=(await self.vectorize(query)),
+                desired_fields="text",
+                collection_name=collection_name,
+                **kwargs,
+            )
         )[:final_limit]
 
     async def aask_retrieved(
@@ -196,8 +267,10 @@ class Rag(LLMUsage):
         question: str | List[str],
         query: List[str] | 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 by retrieving relevant documents based on the provided query.
@@ -210,16 +283,28 @@ class Rag(LLMUsage):
             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.
         """
-        docs = await self.aretrieve(query, collection_name, result_per_query, final_limit)
+        docs = await self.aretrieve(
+            query,
+            collection_name,
+            final_limit,
+            result_per_query=result_per_query,
+            similarity_threshold=similarity_threshold,
+        )
+
+        rendered = template_manager.render_template(configs.templates.retrieved_display_template, {"docs": docs[::-1]})
+
+        logger.debug(f"Retrieved Documents: \n{rendered}")
         return await self.aask(
             question,
-            template_manager.render_template(configs.templates.retrieved_display_template, {"docs": docs}),
+            f"{rendered}\n\n{extra_system_message}",
             **kwargs,
         )

fabricatio/capabilities/rating.py

@@ -131,8 +131,7 @@ class GiveRating(WithBriefing, LLMUsage):
 
         def _validator(response: str) -> Dict[str, str] | None:
             if (
-                (json_data := JsonCapture.convert_with(response, orjson.loads)) is not None
-                and isinstance(json_data, dict)
+                (json_data := JsonCapture.validate_with(response, target_type=dict, elements_type=str)) is not None
                 and json_data.keys() == criteria
                 and all(isinstance(v, str) for v in json_data.values())
             ):
@@ -173,11 +172,10 @@ class GiveRating(WithBriefing, LLMUsage):
 
         def _validator(response: str) -> Set[str] | None:
             if (
-                (json_data := JsonCapture.convert_with(response, orjson.loads)) is not None
-                and isinstance(json_data, list)
-                and all(isinstance(v, str) for v in json_data)
-                and (criteria_count == 0 or len(json_data) == criteria_count)
-            ):
+                json_data := JsonCapture.validate_with(
+                    response, target_type=list, elements_type=str, length=criteria_count
+                )
+            ) is not None:
                 return set(json_data)
             return None
 
@@ -219,27 +217,6 @@ class GiveRating(WithBriefing, LLMUsage):
         Returns:
             Set[str]: A set of drafted rating criteria.
         """
-
-        def _reasons_validator(response: str) -> List[str] | None:
-            if (
-                (json_data := JsonCapture.convert_with(response, orjson.loads)) is not None
-                and isinstance(json_data, list)
-                and all(isinstance(v, str) for v in json_data)
-                and len(json_data) == reasons_count
-            ):
-                return json_data
-            return None
-
-        def _criteria_validator(response: str) -> Set[str] | None:
-            if (
-                (json_data := JsonCapture.convert_with(response, orjson.loads)) is not None
-                and isinstance(json_data, list)
-                and all(isinstance(v, str) for v in json_data)
-                and len(json_data) == criteria_count
-            ):
-                return set(json_data)
-            return None
-
         kwargs = GenerateKwargs(system_message=f"# your personal briefing: \n{self.briefing}", **kwargs)
         # extract reasons from the comparison of ordered pairs of extracted from examples
         reasons = flatten(
@@ -256,7 +233,9 @@ class GiveRating(WithBriefing, LLMUsage):
                     )
                     for pair in (permutations(examples, 2))
                 ],
-                validator=_reasons_validator,
+                validator=lambda resp: JsonCapture.validate_with(
+                    resp, target_type=list, elements_type=str, length=reasons_count
+                ),
                 **kwargs,
             )
         )
@@ -272,7 +251,9 @@ class GiveRating(WithBriefing, LLMUsage):
                     },
                 )
             ),
-            validator=_criteria_validator,
+            validator=lambda resp: set(out)
+            if (out := JsonCapture.validate_with(resp, target_type=list, elements_type=str, length=criteria_count))
+            else None,
             **kwargs,
         )
 
@@ -295,11 +276,6 @@ class GiveRating(WithBriefing, LLMUsage):
         if len(criteria) < 2:  # noqa: PLR2004
             raise ValueError("At least two criteria are required to draft rating weights")
 
-        def _validator(resp: str) -> float | None:
-            if (cap := JsonCapture.convert_with(resp, orjson.loads)) is not None and isinstance(cap, float):
-                return cap
-            return None
-
         criteria = list(criteria)  # freeze the order
         windows = windowed(criteria, 2)
 
@@ -316,7 +292,7 @@ class GiveRating(WithBriefing, LLMUsage):
                 )
                 for pair in windows
             ],
-            validator=_validator,
+            validator=lambda resp: JsonCapture.validate_with(resp, target_type=float),
             **GenerateKwargs(system_message=f"# your personal briefing: \n{self.briefing}", **kwargs),
         )
         weights = [1]

fabricatio/capabilities/task.py

@@ -5,21 +5,21 @@ from typing import Any, Dict, List, Optional, Tuple, Unpack
 
 import orjson
 from fabricatio._rust_instances import template_manager
+from fabricatio.capabilities.propose import Propose
 from fabricatio.config import configs
 from fabricatio.models.generic import WithBriefing
 from fabricatio.models.kwargs_types import ChooseKwargs, ValidateKwargs
 from fabricatio.models.task import Task
 from fabricatio.models.tool import Tool, ToolExecutor
-from fabricatio.models.usages import LLMUsage, ToolBoxUsage
+from fabricatio.models.usages import ToolBoxUsage
 from fabricatio.parser import JsonCapture, PythonCapture
 from loguru import logger
-from pydantic import ValidationError
 
 
-class ProposeTask(WithBriefing, LLMUsage):
+class ProposeTask(WithBriefing, Propose):
     """A class that proposes a task based on a prompt."""
 
-    async def propose[T](
+    async def propose_task[T](
         self,
         prompt: str,
         **kwargs: Unpack[ValidateKwargs],
@@ -34,27 +34,10 @@ class ProposeTask(WithBriefing, LLMUsage):
             A Task object based on the proposal result.
         """
         if not prompt:
-            err = f"{self.name}: Prompt must be provided."
-            logger.error(err)
+            logger.error(err := f"{self.name}: Prompt must be provided.")
             raise ValueError(err)
 
-        def _validate_json(response: str) -> None | Task:
-            try:
-                cap = JsonCapture.capture(response)
-                logger.debug(f"Response: \n{response}")
-                logger.info(f"Captured JSON: \n{cap}")
-                return Task.model_validate_json(cap)
-            except ValidationError as e:
-                logger.error(f"Failed to parse task from JSON: {e}")
-                return None
-
-        template_data = {"prompt": prompt, "json_example": Task.json_example()}
-        return await self.aask_validate(
-            question=template_manager.render_template(configs.templates.propose_task_template, template_data),
-            validator=_validate_json,
-            system_message=f"# your personal briefing: \n{self.briefing}",
-            **kwargs,
-        )
+        return await self.propose(Task, prompt, system_message=f"# your personal briefing: \n{self.briefing}", **kwargs)
 
 
 class HandleTask(WithBriefing, ToolBoxUsage):