fabricatio 0.2.3.dev2__cp312-cp312-win_amd64.whl → 0.2.3.dev3__cp312-cp312-win_amd64.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,11 +1,14 @@
 """A module for the RAG (Retrieval Augmented Generation) model."""
 
+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
 
+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
@@ -14,67 +17,77 @@ 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
+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:
+    """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."""
 
-    _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)
+    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)
+    """The Milvus client used for the RAG model."""
 
     @property
     def client(self) -> MilvusClient:
-        """The Milvus client."""
+        """Return the Milvus client."""
         return self._client
 
-    def view(self, collection_name: str, create: bool = False) -> Self:
+    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.
             create (bool): Whether to create the collection if it does not exist.
         """
-        if create and self._client.has_collection(collection_name):
+        if create and collection_name and not self._client.has_collection(collection_name):
             self._client.create_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
@@ -92,7 +105,7 @@ 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)
+        self._client.insert(collection_name or self.safe_target_collection, data)
         return self
 
     def consume(
@@ -109,7 +122,7 @@ class Rag(LLMUsage):
             Self: The current instance, allowing for method chaining.
         """
         data = reader(Path(source))
-        self.add_document(data, collection_name or self.safe_viewing_collection)
+        self.add_document(data, collection_name or self.safe_target_collection)
         return self
 
     async def afetch_document(
@@ -132,7 +145,7 @@ 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,
             output_fields=desired_fields if isinstance(desired_fields, list) else [desired_fields],
             limit=result_per_query,
@@ -177,3 +190,36 @@ class Rag(LLMUsage):
             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.dev2.dist-info → fabricatio-0.2.3.dev3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fabricatio
-Version: 0.2.3.dev2
+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.dev2.dist-info → fabricatio-0.2.3.dev3.dist-info}/RECORD

@@ -1,13 +1,13 @@
-fabricatio-0.2.3.dev2.dist-info/METADATA,sha256=BlVLqYv59JMHHleZ0q386UhyragcQ3iZbMXonozLX3c,12339
-fabricatio-0.2.3.dev2.dist-info/WHEEL,sha256=tpW5AN9B-9qsM9WW2FXG2r193YXiqexDadpKp0A2daI,96
-fabricatio-0.2.3.dev2.dist-info/licenses/LICENSE,sha256=do7J7EiCGbq0QPbMAL_FqLYufXpHnCnXBOuqVPwSV8Y,1088
+fabricatio-0.2.3.dev3.dist-info/METADATA,sha256=_1gGVl7M89bTO5_yTp9V1Tc9GAAlgcB-QY5q-IhUUUQ,12296
+fabricatio-0.2.3.dev3.dist-info/WHEEL,sha256=tpW5AN9B-9qsM9WW2FXG2r193YXiqexDadpKp0A2daI,96
+fabricatio-0.2.3.dev3.dist-info/licenses/LICENSE,sha256=do7J7EiCGbq0QPbMAL_FqLYufXpHnCnXBOuqVPwSV8Y,1088
 fabricatio/actions/communication.py,sha256=NZxIIncKgJSDyBrqNebUtH_haqtxHa8ld2TZxT3CMdU,429
 fabricatio/actions/transmission.py,sha256=xpvKqbXqgpi1BWy-vUUvmd8NZ1GhRNfsYUBp-l2jLyk,862
 fabricatio/actions/__init__.py,sha256=eFmFVPQvtNgFynIXBVr3eP-vWQDWCPng60YY5LXvZgg,115
-fabricatio/capabilities/rag.py,sha256=WeUn-7wppyISbmNpm4flogzUZrCnrJ5iQYNjepZpdww,6954
+fabricatio/capabilities/rag.py,sha256=zUOxpX-MV1HVQ2iSjk2j5jMf0ZqAgqB42DbSBML5KGQ,9484
 fabricatio/capabilities/rating.py,sha256=zmTUvsUfxFgovRQzy4djL2zKRYTHmN6JY7A4lyT5uVQ,14907
 fabricatio/capabilities/task.py,sha256=d2xtrwQxXWI40UskQCR5YhHarY7ST0ppr8TjY12uWQE,5327
-fabricatio/config.py,sha256=6_cpyGu87NmuPC2Kc0fq_cwOogH3fjXTILiDY7poWi4,12041
+fabricatio/config.py,sha256=9ZgPn1GJMgDSGt31Xt6oPD10s2o-ssxbGHBc5PD_drs,12222
 fabricatio/core.py,sha256=VQ_JKgUGIy2gZ8xsTBZCdr_IP7wC5aPg0_bsOmjQ588,6458
 fabricatio/decorators.py,sha256=uzsP4tFKQNjDHBkofsjjoJA0IUAaYOtt6YVedoyOqlo,6551
 fabricatio/fs/curd.py,sha256=faMstgGUiQ4k2AW3OXfvvWWTldTtKXco7QINYaMjmyA,3981
@@ -19,7 +19,7 @@ fabricatio/models/events.py,sha256=mrihNEFgQ5o7qFWja1z_qX8dnaTLwPBoJdVlzxQV5oM,2
 fabricatio/models/generic.py,sha256=WEjZ96rTyBjaBjkM6e8E4Pg_Naot4xWRvGJteqBiCCI,5133
 fabricatio/models/kwargs_types.py,sha256=a-e7rdMZJi8xTBL_RLmTC9OPzI-Js7rlS689PR03VvA,1401
 fabricatio/models/role.py,sha256=gYvleTeKUGDUNKPAC5B0EPMLC4jZ4vHsFHmHiVXkU6c,1830
-fabricatio/models/task.py,sha256=ip6VeOV7vgXqhiQFOCjVl3hzc6lgdhfyvxbBuSz2-C0,11529
+fabricatio/models/task.py,sha256=M6jeDFE3jX6cNV9bdOwhjHqgBHI3FKtFLWcmlqhYgcs,11419
 fabricatio/models/tool.py,sha256=WTFnpF6xZ1nJbmIOonLsGQcM-kkDCeZiAFqyil9xg2U,6988
 fabricatio/models/usages.py,sha256=bzsTDrAekiQyIwKeWds5YdgsXk8qiZkD7OZdno1Q_Ck,28213
 fabricatio/models/utils.py,sha256=mXea76bd4r2jy_zx74GM4t5kCvkMu0JTOaw_VGvTCxk,3952
@@ -31,7 +31,7 @@ fabricatio/toolboxes/task.py,sha256=kU4a501awIDV7GwNDuSlK3_Ym-5OhCp5sS-insTmUmQ,
 fabricatio/toolboxes/__init__.py,sha256=b13KmASO8q5fBLwew964fn9oH86ER5g-S1PgA4fZ_xs,482
 fabricatio/_rust.pyi,sha256=0wCqtwWkVxxoqprvk8T27T8QYKIAKHS7xgsmdMNjQKc,1756
 fabricatio/_rust_instances.py,sha256=dl0-yZ4UvT5g20tQgnPJpmqtkjFGXNG_YK4eLfi_ugQ,279
-fabricatio/__init__.py,sha256=opIrN8lGyT-h2If4Qez0bRuWBa3uIT9GsM9CZy7_XJ0,1100
-fabricatio/_rust.cp312-win_amd64.pyd,sha256=wJkACEnysxtUfMBOwQ3851hSYNy1HOtY9vSIepLzEFg,1272320
-fabricatio-0.2.3.dev2.data/scripts/tdown.exe,sha256=A9eIQbGxwDhpQ0gYqtEHA4Ms-LakyYkSkjiNnccvfYY,3398144
-fabricatio-0.2.3.dev2.dist-info/RECORD,,
+fabricatio/__init__.py,sha256=5MAnDTbECGJUBwv2_Hg40zCuWgPVTFR34wzu9EqrwXc,1246
+fabricatio/_rust.cp312-win_amd64.pyd,sha256=EedusrQbVZW1_gQ8NhTCcyNdhopSOEJOKAxU9hB6AWU,1264128
+fabricatio-0.2.3.dev3.data/scripts/tdown.exe,sha256=19y_6dFzSkW1wYvplySjLMz6KqUgV6-nZzl2smJkyXQ,3397632
+fabricatio-0.2.3.dev3.dist-info/RECORD,,