fabricatio 0.2.3.dev2__cp312-cp312-win_amd64.whl → 0.2.4.dev0__cp312-cp312-win_amd64.whl

fabricatio/__init__.py

@@ -1,6 +1,9 @@
 """Fabricatio is a Python library for building llm app using event-based agent structure."""
 
+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
@@ -18,6 +21,7 @@ __all__ = [
     "Capture",
     "CodeBlockCapture",
     "Event",
+    "ExtractArticleEssence",
     "JsonCapture",
     "Message",
     "Messages",
@@ -35,3 +39,9 @@ __all__ = [
     "task_toolbox",
     "template_manager",
 ]
+
+
+if find_spec("pymilvus"):
+    from fabricatio.capabilities.rag import 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,127 @@
+"""Actions for transmitting tasks to targets."""
+
+from os import PathLike
+from pathlib import Path
+from typing import Callable, List
+
+from pydantic import BaseModel, Field
+from pydantic.config import ConfigDict
+
+from fabricatio.journal import logger
+from fabricatio.models.action import Action
+from fabricatio.models.generic import ProposedAble
+from fabricatio.models.task import Task
+
+
+class Equation(BaseModel):
+    """Structured representation of mathematical equations (including their physical or conceptual meanings)."""
+
+    model_config = ConfigDict(use_attribute_docstrings=True)
+
+    description: str = Field(...)
+    """A concise explanation of the equation's meaning, purpose, and relevance in the context of the research."""
+
+    latex_code: str = Field(...)
+    """The LaTeX code used to represent the equation in a publication-ready format."""
+
+
+class Figure(BaseModel):
+    """Structured representation of figures (including their academic significance and explanatory captions)."""
+
+    model_config = ConfigDict(use_attribute_docstrings=True)
+
+    description: str = Field(...)
+    """A detailed explanation of the figure's content and its role in conveying key insights."""
+
+    figure_caption: str = Field(...)
+    """The caption accompanying the figure, summarizing its main points and academic value."""
+
+
+class ArticleEssence(ProposedAble):
+    """Structured representation of the core elements of an academic paper(providing a comprehensive digital profile of the paper's essential information)."""
+
+    # Basic Metadata
+    title: str = Field(...)
+    """The full title of the paper, including any subtitles if applicable."""
+
+    authors: List[str] = Field(default_factory=list)
+    """A list of the paper's authors, typically in the order of contribution."""
+
+    keywords: List[str] = Field(default_factory=list)
+    """A list of keywords that summarize the paper's focus and facilitate indexing."""
+
+    publication_year: int = Field(None)
+    """The year in which the paper was published."""
+
+    # Core Content Elements
+    domain: List[str] = Field(default_factory=list)
+    """The research domains or fields addressed by the paper (e.g., ['Natural Language Processing', 'Computer Vision'])."""
+
+    abstract: str = Field(...)
+    """A structured abstract that outlines the research problem, methodology, and conclusions in three distinct sections."""
+
+    core_contributions: List[str] = Field(default_factory=list)
+    """Key academic contributions that distinguish the paper from prior work in the field."""
+
+    technical_novelty: List[str] = Field(default_factory=list)
+    """Specific technical innovations introduced by the research, listed as individual points."""
+
+    # Academic Achievements Showcase
+    highlighted_equations: List[Equation] = Field(default_factory=list)
+    """Core mathematical equations that represent breakthroughs in the field, accompanied by explanations of their physical or conceptual significance."""
+
+    highlighted_algorithms: List[str] = Field(default_factory=list)
+    """Pseudocode for key algorithms, annotated to highlight innovative components."""
+
+    highlighted_figures: List[Figure] = Field(default_factory=list)
+    """Critical diagrams or illustrations, each accompanied by a caption explaining their academic importance."""
+
+    highlighted_tables: List[str] = Field(default_factory=list)
+    """Important data tables, annotated to indicate statistical significance or other notable findings."""
+
+    # Academic Discussion Dimensions
+    research_problem: str = Field("")
+    """A clearly defined research question or problem addressed by the study."""
+
+    limitations: List[str] = Field(default_factory=list)
+    """An analysis of the methodological or experimental limitations of the research."""
+
+    future_work: List[str] = Field(default_factory=list)
+    """Suggestions for potential directions or topics for follow-up studies."""
+
+    impact_analysis: str = Field("")
+    """An assessment of the paper's potential influence on the development of the field."""
+
+
+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,89 +1,146 @@
 """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
+from typing import Any, Callable, Dict, List, Optional, Self, Union, Unpack, overload
 
+from fabricatio._rust_instances import template_manager
 from fabricatio.config import configs
-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
+from pydantic import Field, PrivateAttr
 
-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 PrivateAttr
+
+@lru_cache(maxsize=None)
+def create_client(uri: str, token: str = "", timeout: Optional[float] = None) -> MilvusClient:
+    """Create a Milvus client."""
+    return MilvusClient(
+        uri=uri,
+        token=token,
+        timeout=timeout,
+    )
 
 
-class Rag(LLMUsage):
+class RAG(EmbeddingUsage):
     """A class representing the RAG (Retrieval Augmented Generation) model."""
 
-    _client: MilvusClient = PrivateAttr(
-        default=MilvusClient(
-            uri=configs.rag.milvus_uri.unicode_string(),
-            token=configs.rag.milvus_token.get_secret_value(),
-            timeout=configs.rag.milvus_timeout,
-        ),
-    )
-    _target_collection: Optional[str] = PrivateAttr(default=None)
+    target_collection: Optional[str] = Field(default=None)
+    """The name of the collection being viewed."""
+
+    _client: Optional[MilvusClient] = PrivateAttr(None)
+    """The Milvus client used for the RAG model."""
 
     @property
     def client(self) -> MilvusClient:
-        """The Milvus client."""
+        """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 view(self, collection_name: str, create: bool = False) -> Self:
+    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.
+
+        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 self._client.has_collection(collection_name):
-            self._client.create_collection(collection_name)
+        if create and collection_name and not self._client.has_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
+        self.target_collection = collection_name
         return self
 
-    def quit_view(self) -> Self:
+    def quit_viewing(self) -> Self:
         """Quit the current view.
 
         Returns:
             Self: The current instance, allowing for method chaining.
         """
-        self._target_collection = None
-        return self
-
-    @property
-    def viewing_collection(self) -> Optional[str]:
-        """Get the name of the collection being viewed.
-
-        Returns:
-            Optional[str]: The name of the collection being viewed.
-        """
-        return self._target_collection
+        return self.view(None)
 
     @property
-    def safe_viewing_collection(self) -> str:
+    def safe_target_collection(self) -> str:
         """Get the name of the collection being viewed, raise an error if not viewing any collection.
 
         Returns:
             str: The name of the collection being viewed.
         """
-        if self._target_collection is None:
+        if self.target_collection is None:
             raise RuntimeError("No collection is being viewed. Have you called `self.view()`?")
-        return self._target_collection
+        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.
@@ -92,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_viewing_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.
 
@@ -108,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_viewing_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(
@@ -117,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.
@@ -125,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:
@@ -132,8 +212,9 @@ class Rag(LLMUsage):
         """
         # Step 1: Search for vectors
         search_results = self._client.search(
-            collection_name or self.safe_viewing_collection,
+            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,
         )
@@ -144,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]
 
@@ -155,25 +237,74 @@ 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(
+        self,
+        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.
+
+        This method performs document retrieval using the given query, then asks the
+        specified question using the retrieved documents as context.
+
+        Args:
+            question (str | List[str]): The question or list of questions 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.
+        """
+        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,
+            f"{rendered}\n\n{extra_system_message}",
+            **kwargs,
+        )

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):