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

fabricatio/__init__.py

@@ -1,5 +1,7 @@
 """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.core import env
 from fabricatio.fs import magika
@@ -35,3 +37,9 @@ __all__ = [
     "task_toolbox",
     "template_manager",
 ]
+
+
+if find_spec("pymilvus"):
+    from fabricatio.capabilities.rag import Rag
+
+    __all__ += ["Rag"]

fabricatio/capabilities/rag.py

@@ -1,46 +1,225 @@
 """A module for the RAG (Retrieval Augmented Generation) model."""
 
-from typing import Any, Dict, List, Optional, Self, Union
+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 fabricatio import template_manager
 from fabricatio.config import configs
+from fabricatio.models.kwargs_types import LLMKwargs
+from fabricatio.models.usages import LLMUsage
 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 BaseModel, ConfigDict, PrivateAttr
+from pydantic import Field, PrivateAttr
 
 
-class Rag(BaseModel):
+@lru_cache(maxsize=None)
+def create_client(
+    uri: Optional[str] = None, token: Optional[str] = None, 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,
+    )
+
+
+class Rag(LLMUsage):
     """A class representing the RAG (Retrieval Augmented Generation) model."""
 
-    model_config = ConfigDict(use_attribute_docstrings=True)
+    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(
-        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)
+    _client: MilvusClient = PrivateAttr(None)
+    """The Milvus client used for the RAG model."""
 
     @property
     def client(self) -> MilvusClient:
-        """The Milvus client."""
+        """Return the Milvus client."""
         return self._client
 
-    def add_document[D: Union[Dict[str, Any] | MilvusData]](self, collection_name: str, data: D | List[D]) -> Self:
-        """Adds a document to the specified collection.
+    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 view(self, collection_name: Optional[str], create: bool = False) -> Self:
+        """View the specified collection.
 
         Args:
             collection_name (str): The name of the collection.
-            data (dict): The data to be added to the collection.
+            create (bool): Whether to create the collection if it does not exist.
+        """
+        if create and collection_name and not self._client.has_collection(collection_name):
+            self._client.create_collection(collection_name)
+
+        self.target_collection = collection_name
+        return self
+
+    def quit_viewing(self) -> Self:
+        """Quit the current view.
+
+        Returns:
+            Self: The current instance, allowing for method chaining.
+        """
+        return self.view(None)
+
+    @property
+    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:
+            raise RuntimeError("No collection is being viewed. Have you called `self.view()`?")
+        return self.target_collection
+
+    def add_document[D: Union[Dict[str, Any], MilvusData]](
+        self, data: D | List[D], collection_name: Optional[str] = None
+    ) -> 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.
+
+        Returns:
+            Self: The current instance, allowing for method chaining.
         """
         if isinstance(data, MilvusData):
             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, data)
+        self._client.insert(collection_name or self.safe_target_collection, data)
         return self
+
+    def consume(
+        self, source: PathLike, reader: Callable[[PathLike], MilvusData], 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.
+        """
+        data = reader(Path(source))
+        self.add_document(data, collection_name or self.safe_target_collection)
+        return self
+
+    async def afetch_document(
+        self,
+        vecs: List[List[float]],
+        desired_fields: List[str] | str,
+        collection_name: Optional[str] = None,
+        result_per_query: int = 10,
+    ) -> List[Dict[str, Any]] | List[Any]:
+        """Fetch data from the collection.
+
+        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.
+            result_per_query (int): The number of results to return per query.
+
+        Returns:
+            List[Dict[str, Any]] | List[Any]: The retrieved data.
+        """
+        # Step 1: Search for vectors
+        search_results = self._client.search(
+            collection_name or self.safe_target_collection,
+            vecs,
+            output_fields=desired_fields if isinstance(desired_fields, list) else [desired_fields],
+            limit=result_per_query,
+        )
+
+        # Step 2: Flatten the search results
+        flattened_results = flatten(search_results)
+
+        # Step 3: Sort by distance (descending)
+        sorted_results = sorted(flattened_results, key=itemgetter("distance"), reverse=True)
+
+        # 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]
+
+    async def aretrieve(
+        self,
+        query: List[str] | str,
+        collection_name: Optional[str] = None,
+        result_per_query: int = 10,
+        final_limit: int = 20,
+    ) -> 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.
+
+        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,
+        )[:final_limit]
+
+    async def aask_retrieved(
+        self,
+        question: str | List[str],
+        query: List[str] | str,
+        collection_name: Optional[str] = None,
+        result_per_query: int = 10,
+        final_limit: int = 20,
+        **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.
+            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.
+            **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)
+        return await self.aask(
+            question,
+            template_manager.render_template(configs.templates.retrieved_display_template, {"docs": docs}),
+            **kwargs,
+        )

fabricatio/config.py

@@ -176,6 +176,9 @@ class TemplateConfig(BaseModel):
     draft_rating_weights_klee_template: str = Field(default="draft_rating_weights_klee")
     """The name of the draft rating weights klee template which will be used to draft rating weights with Klee method."""
 
+    retrieved_display_template: str = Field(default="retrieved_display")
+    """The name of the retrieved display template which will be used to display retrieved documents."""
+
 
 class MagikaConfig(BaseModel):
     """Magika configuration class."""

fabricatio/models/task.py

@@ -46,19 +46,19 @@ class Task[T](WithBriefing, WithJsonExample, WithDependency):
     """
 
     name: str = Field(...)
-    """The name of the task, which should be a concise and descriptive name."""
+    """The name of the task, which should be concise and descriptive."""
 
     description: str = Field(default="")
-    """The description of the task, which should provide every details and noting about the task if provided, obeying the CEFR level rule and 5W1H rule."""
+    """A detailed explanation of the task that includes all necessary information. Should be clear and answer what, why, when, where, who, and how questions."""
 
-    goal: List[str] = Field(default=[])
-    """The goal of the task, a list of strings. The goal should be a concise and clear statement of what the task is intended to achieve, goal SHALL NOT be too broad or too narrow."""
+    goals: List[str] = Field(default=[])
+    """A list of objectives that the task aims to accomplish. Each goal should be clear and specific. Complex tasks should be broken into multiple smaller goals."""
 
     namespace: List[str] = Field(default_factory=list)
-    """The namespace of the task, a list of namespace segment, as string, if it is not directly given out, it SHALL just be a empty list meaning `NOT ASSIGNED`"""
+    """A list of string segments that identify the task's location in the system. If not specified, defaults to an empty list."""
 
     dependencies: List[str] = Field(default_factory=list)
-    """A list of file paths, These file are needed to read or write to meet a specific requirement of this task, if it is not directly given out, it SHALL just be a empty list meaning `NOT ASSIGNED`"""
+    """A list of file paths that are needed (either reading or writing) to complete this task. If not specified, defaults to an empty list."""
 
     _output: Queue[T | None] = PrivateAttr(default_factory=Queue)
     """The output queue of the task."""
@@ -113,7 +113,7 @@ class Task[T](WithBriefing, WithJsonExample, WithDependency):
         Returns:
             Task: A new instance of the `Task` class.
         """
-        return cls(name=name, goal=goal, description=description)
+        return cls(name=name, goals=goal, description=description)
 
     def update_task(self, goal: Optional[List[str] | str] = None, description: Optional[str] = None) -> Self:
         """Update the goal and description of the task.
@@ -126,7 +126,7 @@ class Task[T](WithBriefing, WithJsonExample, WithDependency):
             Task: The updated instance of the `Task` class.
         """
         if goal:
-            self.goal = goal if isinstance(goal, list) else [goal]
+            self.goals = goal if isinstance(goal, list) else [goal]
         if description:
             self.description = description
         return self

{fabricatio-0.2.3.dev1.dist-info → fabricatio-0.2.3.dev3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fabricatio
-Version: 0.2.3.dev1
+Version: 0.2.3.dev3
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: 3.12
@@ -98,32 +98,32 @@ from fabricatio import Action, Role, Task, logger
 
 
 class Hello(Action):
-    """Action that says hello."""
+   """Action that says hello."""
 
-    name: str = "hello"
-    output_key: str = "task_output"
+   name: str = "hello"
+   output_key: str = "task_output"
 
-    async def _execute(self, task_input: Task[str], **_) -> Any:
-        ret = "Hello fabricatio!"
-        logger.info("executing talk action")
-        return ret
+   async def _execute(self, task_input: Task[str], **_) -> Any:
+      ret = "Hello fabricatio!"
+      logger.info("executing talk action")
+      return ret
 
 
 async def main() -> None:
-    """Main function."""
-    role = Role(
-        name="talker", 
-        description="talker role", 
-        registry={Task.pending_label: WorkFlow(name="talk", steps=(Hello,))}
-    )
+   """Main function."""
+   role = Role(
+      name="talker",
+      description="talker role",
+      registry={Task.pending_label: WorkFlow(name="talk", steps=(Hello,))}
+   )
 
-    task = Task(name="say hello", goal="say hello", description="say hello to the world")
-    result = await task.delegate()
-    logger.success(f"Result: {result}")
+   task = Task(name="say hello", goals="say hello", description="say hello to the world")
+   result = await task.delegate()
+   logger.success(f"Result: {result}")
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+   asyncio.run(main())
 ```
 
 #### Writing and Dumping Code
@@ -311,17 +311,18 @@ from fabricatio.models.task import Task
 
 toolbox_usage = ToolBoxUsage()
 
+
 async def handle_security_vulnerabilities():
-    task = Task(
-        name="Security Check",
-        goal=["Identify security vulnerabilities"],
-        description="Perform a thorough security review on the project.",
-        dependencies=["./src/main.py"]
-    )
-    
-    vulnerabilities = await toolbox_usage.gather_tools_fine_grind(task)
-    for vulnerability in vulnerabilities:
-        print(f"Found vulnerability: {vulnerability.name}")
+   task = Task(
+      name="Security Check",
+      goals=["Identify security vulnerabilities"],
+      description="Perform a thorough security review on the project.",
+      dependencies=["./src/main.py"]
+   )
+
+   vulnerabilities = await toolbox_usage.gather_tools_fine_grind(task)
+   for vulnerability in vulnerabilities:
+      print(f"Found vulnerability: {vulnerability.name}")
 ```
 
 #### Managing CTF Challenges
@@ -334,19 +335,22 @@ from fabricatio.models.task import Task
 
 toolbox_usage = ToolBoxUsage()
 
+
 async def solve_ctf_challenge(challenge_name: str, challenge_description: str, files: list[str]):
-    task = Task(
-        name=challenge_name,
-        goal=[f"Solve {challenge_name} challenge"],
-        description=challenge_description,
-        dependencies=files
-    )
-    
-    solution = await toolbox_usage.gather_tools_fine_grind(task)
-    print(f"Challenge Solved: {solution}")
+   task = Task(
+      name=challenge_name,
+      goals=[f"Solve {challenge_name} challenge"],
+      description=challenge_description,
+      dependencies=files
+   )
+
+   solution = await toolbox_usage.gather_tools_fine_grind(task)
+   print(f"Challenge Solved: {solution}")
+
 
 if __name__ == "__main__":
-    asyncio.run(solve_ctf_challenge("Binary Exploitation", "CTF Binary Exploitation Challenge", ["./challenges/binary_exploit"]))
+   asyncio.run(
+      solve_ctf_challenge("Binary Exploitation", "CTF Binary Exploitation Challenge", ["./challenges/binary_exploit"]))
 ```
 
 ### Configuration

{fabricatio-0.2.3.dev1.dist-info → fabricatio-0.2.3.dev3.dist-info}/RECORD

@@ -1,6 +1,6 @@
-fabricatio-0.2.3.dev1.dist-info/METADATA,sha256=0NOXlwnHxtpH1XjaFx4xx_1wSQGLOX4XzsodmXBBTgk,11960
-fabricatio-0.2.3.dev1.dist-info/WHEEL,sha256=RIvmwLDYujv60MYBx2jxyP4vdn1DD7X0kBgz1TQvZuc,108
-fabricatio-0.2.3.dev1.dist-info/licenses/LICENSE,sha256=yDZaTLnOi03bi3Dk6f5IjhLUc5old2yOsihHWU0z-i0,1067
+fabricatio-0.2.3.dev3.dist-info/METADATA,sha256=-3639SKLzCTOvpu4nz4VKm00nP4zt-Ot3pz4_e6PG7o,11913
+fabricatio-0.2.3.dev3.dist-info/WHEEL,sha256=RIvmwLDYujv60MYBx2jxyP4vdn1DD7X0kBgz1TQvZuc,108
+fabricatio-0.2.3.dev3.dist-info/licenses/LICENSE,sha256=yDZaTLnOi03bi3Dk6f5IjhLUc5old2yOsihHWU0z-i0,1067
 fabricatio/decorators.py,sha256=cJHsxxbnMhc4SzPl4454CPLuDP3H0qbTrzV_U2rLPrs,6372
 fabricatio/core.py,sha256=MaEKZ6DDmbdScAY-7F1gwGA6fr7ADX6Mz5rNVi2msFA,6277
 fabricatio/models/generic.py,sha256=4j6DRNkHLNhx8U6aijf0Cz9HSahtEPQBh22i6d_PGTs,4981
@@ -10,7 +10,7 @@ fabricatio/models/kwargs_types.py,sha256=litzaTVzzLjOpddAVdLfsuys_rFoqTyv0etCh9y
 fabricatio/models/utils.py,sha256=ZmVHsfMbIWLGQoZSJLYInDZkQvmpmS5RGBTkJXveNHc,3825
 fabricatio/models/usages.py,sha256=0pyQBahL2lp_Bb_nSyxrW44cm3Me1Btxk3c7YVK32WI,27554
 fabricatio/models/events.py,sha256=TkMSJYSzRgurbfY2knWIc8gYw7rmaWqB7KeYzy5jeWU,2637
-fabricatio/models/task.py,sha256=tPeX9i-dPpxQfrkX7xR2ESFlIXUTg2z-TqOVx5_qUCU,11198
+fabricatio/models/task.py,sha256=U63ULoPkOXx9UyqYqdGmALBGxn9lAidXbSxIZZXlbVM,11088
 fabricatio/models/action.py,sha256=006Jjfr9p-EBErboAn29yvyolIGi5mbJ_DA5og0Uh64,5780
 fabricatio/toolboxes/fs.py,sha256=oZjGOmtN8ZbkMXxCMatqvdPavVXyQ87AEDc1lx9jimY,483
 fabricatio/toolboxes/__init__.py,sha256=iEBsH0aRPLccIe6eyZ6iH5GdK8yYNQTnCwLzw4pWXAA,465
@@ -20,18 +20,18 @@ fabricatio/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 fabricatio/fs/readers.py,sha256=lqMrM5Nt3l6QJmPDoBno2PSaN2BbFwfUjBqaghrbK3A,1002
 fabricatio/fs/curd.py,sha256=g74Pc2WPSkPlCXX0cWj2jyiDYTfkLwpwM480eyCkmc8,3851
 fabricatio/fs/__init__.py,sha256=bYE9r8uR0dtknzbg_YaGv_6Wwa27ntkQt0Tl7Kb3HFI,117
-fabricatio/config.py,sha256=7lmhzEkKi8eSiVBaHjjV0cEyGCanAUmFXYOrOGn7Y7c,11742
+fabricatio/config.py,sha256=X4oDFVKPycSGuFYrn9cVURalU3fmUSRXBIF7bfpONZo,11920
 fabricatio/journal.py,sha256=bzxZay48ZWI0VIkkDXm4Wc_Cc9lBQYa2VGx3Hxy_PtA,753
-fabricatio/__init__.py,sha256=h0FgwSAdI2yhxYiuQi46uFvzqkK5LfJLt5XsHZ9rWIo,1063
+fabricatio/__init__.py,sha256=DNPOP3s4_5KswbEyFojSEjMp3JuckcsGz8OVGfRaUhY,1201
 fabricatio/actions/communication.py,sha256=MrL_revv1zm1VXqjDMKZkg4PHmyMxYThAmmw2wF791k,414
 fabricatio/actions/transmission.py,sha256=WM-kssiqzTLZuZSiqbRvBaC2qynd9kjrKXya4f9Zcb0,839
 fabricatio/actions/__init__.py,sha256=eLa_5ACZ-FqdrLtOfCHk5nQBxzhIs1kgMIXWmkm2P8Y,110
 fabricatio/_rust_instances.py,sha256=JAtO-vL8ihvduf1SHLNf0w7ZSVGCJeIv6zZ9Ekyy1hY,271
 fabricatio/parser.py,sha256=Q4laV79Svggl09UKa3lAw6NYPuDz1sP2F-6U4tsSvQw,3429
-fabricatio/capabilities/rag.py,sha256=EN1plVCKUPylbe2rilvCBNmf0F13zFmOck9vLUkz3lQ,1671
+fabricatio/capabilities/rag.py,sha256=prL7RL3RtB64lXj0q8TUyksZ39MrEXnWXDgaACQryRQ,9259
 fabricatio/capabilities/rating.py,sha256=KLJlaufG74lHJ5hHI_tjhQgjWjHhO_38Q1ILW2cY9fA,14552
 fabricatio/capabilities/task.py,sha256=UnkeON3VCADXUBtUmQbMwl9TEwK0yQSED3EglqMMHQ4,5196
 fabricatio/_rust.pyi,sha256=clhcURuiB9zlFo4m3VyoWQ8Xs4tvg6KNHXpF-ok9h4o,1703
-fabricatio/_rust.cpython-312-x86_64-linux-gnu.so,sha256=RwfP3a0Y0W82yUIHmH-SuNI6ZD_4X1dmu_QoREHvhnw,1336568
-fabricatio-0.2.3.dev1.data/scripts/tdown,sha256=sYeLflc31k_aVFo7_y-QEBe3yfz6uN0PqJ0cLpyMkOU,4569184
-fabricatio-0.2.3.dev1.dist-info/RECORD,,
+fabricatio/_rust.cpython-312-x86_64-linux-gnu.so,sha256=8sQa5DyKEnGWd4_JacuJksUI8m0pjyGKoF66BzOXmNA,1336728
+fabricatio-0.2.3.dev3.data/scripts/tdown,sha256=KnHYVZuHJisg6DR_zhOUGujgx25BZ9Jsw3HykXNrOlM,4578648
+fabricatio-0.2.3.dev3.dist-info/RECORD,,