fabricatio 0.2.6.dev2__cp312-cp312-win_amd64.whl → 0.2.7.dev3__cp312-cp312-win_amd64.whl

fabricatio/capabilities/rag.py

@@ -15,13 +15,14 @@ from fabricatio.config import configs
 from fabricatio.journal import logger
 from fabricatio.models.kwargs_types import (
     ChooseKwargs,
-    CollectionSimpleConfigKwargs,
+    CollectionConfigKwargs,
     EmbeddingKwargs,
     FetchKwargs,
     LLMKwargs,
+    RetrievalKwargs,
 )
 from fabricatio.models.usages import EmbeddingUsage
-from fabricatio.models.utils import MilvusData
+from fabricatio.models.utils import MilvusData, ok
 from more_itertools.recipes import flatten, unique
 from pydantic import Field, PrivateAttr
 
@@ -60,13 +61,21 @@ class RAG(EmbeddingUsage):
     ) -> Self:
         """Initialize the Milvus client."""
         self._client = create_client(
-            uri=milvus_uri or (self.milvus_uri or configs.rag.milvus_uri).unicode_string(),
+            uri=milvus_uri or ok(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,
+            timeout=milvus_timeout or self.milvus_timeout or configs.rag.milvus_timeout,
         )
         return self
 
+    def check_client(self, init: bool = True) -> Self:
+        """Check if the client is initialized, and if not, initialize it."""
+        if self._client is None and init:
+            return self.init_client()
+        if self._client is None and not init:
+            raise RuntimeError("Client is not initialized. Have you called `self.init_client()`?")
+        return self
+
     @overload
     async def pack(
         self, input_text: List[str], subject: Optional[str] = None, **kwargs: Unpack[EmbeddingKwargs]
@@ -102,17 +111,24 @@ class RAG(EmbeddingUsage):
         ]
 
     def view(
-        self, collection_name: Optional[str], create: bool = False, **kwargs: Unpack[CollectionSimpleConfigKwargs]
+        self, collection_name: Optional[str], create: bool = False, **kwargs: Unpack[CollectionConfigKwargs]
     ) -> 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.
+            **kwargs (Unpack[CollectionConfigKwargs]): Additional keyword arguments for collection configuration.
         """
-        if create and collection_name and self.client.has_collection(collection_name):
-            kwargs["dimension"] = kwargs.get("dimension") or self.milvus_dimensions or configs.rag.milvus_dimensions
+        if create and collection_name and not self.check_client().client.has_collection(collection_name):
+            kwargs["dimension"] = ok(
+                kwargs.get("dimension")
+                or self.milvus_dimensions
+                or configs.rag.milvus_dimensions
+                or self.embedding_dimensions
+                or configs.embedding.dimensions,
+                "`dimension` is not set at any level.",
+            )
             self.client.create_collection(collection_name, auto_id=True, **kwargs)
             logger.info(f"Creating collection {collection_name}")
 
@@ -158,7 +174,7 @@ class RAG(EmbeddingUsage):
         else:
             raise TypeError(f"Expected MilvusData or list of MilvusData, got {type(data)}")
         c_name = collection_name or self.safe_target_collection
-        self.client.insert(c_name, prepared_data)
+        self.check_client().client.insert(c_name, prepared_data)
 
         if flush:
             logger.debug(f"Flushing collection {c_name}")
@@ -198,6 +214,25 @@ class RAG(EmbeddingUsage):
         self.add_document(await self.pack(text), collection_name or self.safe_target_collection, flush=True)
         return self
 
+    @overload
+    async def afetch_document[V: (int, str, float, bytes)](
+        self,
+        vecs: List[List[float]],
+        desired_fields: List[str],
+        collection_name: Optional[str] = None,
+        similarity_threshold: float = 0.37,
+        result_per_query: int = 10,
+    ) -> List[Dict[str, V]]: ...
+
+    @overload
+    async def afetch_document[V: (int, str, float, bytes)](
+        self,
+        vecs: List[List[float]],
+        desired_fields: str,
+        collection_name: Optional[str] = None,
+        similarity_threshold: float = 0.37,
+        result_per_query: int = 10,
+    ) -> List[V]: ...
     async def afetch_document[V: (int, str, float, bytes)](
         self,
         vecs: List[List[float]],
@@ -219,7 +254,7 @@ class RAG(EmbeddingUsage):
             List[Dict[str, Any]] | List[Any]: The retrieved data.
         """
         # Step 1: Search for vectors
-        search_results = self.client.search(
+        search_results = self.check_client().client.search(
             collection_name or self.safe_target_collection,
             vecs,
             search_params={"radius": similarity_threshold},
@@ -260,7 +295,7 @@ class RAG(EmbeddingUsage):
         if isinstance(query, str):
             query = [query]
         return cast(
-            List[str],
+            "List[str]",
             await self.afetch_document(
                 vecs=(await self.vectorize(query)),
                 desired_fields="text",
@@ -268,6 +303,24 @@ class RAG(EmbeddingUsage):
             ),
         )[:final_limit]
 
+    async def aretrieve_compact(
+        self,
+        query: List[str] | str,
+        **kwargs: Unpack[RetrievalKwargs],
+    ) -> str:
+        """Retrieve data from the collection and format it for display.
+
+        Args:
+            query (List[str] | str): The query to be used for retrieval.
+            **kwargs (Unpack[RetrievalKwargs]): Additional keyword arguments for retrieval.
+
+        Returns:
+            str: A formatted string containing the retrieved data.
+        """
+        return TEMPLATE_MANAGER.render_template(
+            configs.templates.retrieved_display_template, {"docs": (await self.aretrieve(query, **kwargs))}
+        )
+
     async def aask_retrieved(
         self,
         question: str,
@@ -298,16 +351,14 @@ class RAG(EmbeddingUsage):
         Returns:
             str: A string response generated after asking with the context of retrieved documents.
         """
-        docs = await self.aretrieve(
+        rendered = await self.aretrieve_compact(
             query or question,
-            final_limit,
+            final_limit=final_limit,
             collection_name=collection_name,
             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,
@@ -315,7 +366,7 @@ class RAG(EmbeddingUsage):
             **kwargs,
         )
 
-    async def arefined_query(self, question: List[str] | str, **kwargs: Unpack[ChooseKwargs]) -> List[str]:
+    async def arefined_query(self, question: List[str] | str, **kwargs: Unpack[ChooseKwargs]) -> Optional[List[str]]:
         """Refines the given question using a template.
 
         Args:

fabricatio/capabilities/rating.py

@@ -10,6 +10,7 @@ from fabricatio.journal import logger
 from fabricatio.models.generic import WithBriefing
 from fabricatio.models.kwargs_types import ValidateKwargs
 from fabricatio.models.usages import LLMUsage
+from fabricatio.models.utils import override_kwargs
 from fabricatio.parser import JsonCapture
 from more_itertools import flatten, windowed
 from pydantic import NonNegativeInt, PositiveInt
@@ -126,13 +127,13 @@ class GiveRating(WithBriefing, LLMUsage):
         return await self.rate_fine_grind(to_rate, manual, score_range, **kwargs)
 
     async def draft_rating_manual(
-        self, topic: str, criteria: Set[str], **kwargs: Unpack[ValidateKwargs[Dict[str, str]]]
+        self, topic: str, criteria: Optional[Set[str]] = None, **kwargs: Unpack[ValidateKwargs[Dict[str, str]]]
     ) -> Optional[Dict[str, str]]:
         """Drafts a rating manual based on a topic and dimensions.
 
         Args:
             topic (str): The topic for the rating manual.
-            criteria (Set[str]): A set of dimensions for the rating manual.
+            criteria (Optional[Set[str]], optional): A set of criteria for the rating manual. If not specified, then this method will draft the criteria automatically.
             **kwargs (Unpack[ValidateKwargs]): Additional keyword arguments for the LLM usage.
 
         Returns:
@@ -148,6 +149,14 @@ class GiveRating(WithBriefing, LLMUsage):
                 return json_data
             return None
 
+        criteria = criteria or await self.draft_rating_criteria(
+            topic, **self.prepend_sys_msg(override_kwargs(dict(kwargs), default=None))
+        )
+
+        if criteria is None:
+            logger.error(f"Failed to draft rating criteria for topic {topic}")
+            return None
+
         return await self.aask_validate(
             question=(
                 TEMPLATE_MANAGER.render_template(
@@ -159,7 +168,7 @@ class GiveRating(WithBriefing, LLMUsage):
                 )
             ),
             validator=_validator,
-            **self.prepend(kwargs),
+            **self.prepend_sys_msg(kwargs),
         )
 
     async def draft_rating_criteria(
@@ -191,7 +200,7 @@ class GiveRating(WithBriefing, LLMUsage):
             validator=lambda resp: set(out)
             if (out := JsonCapture.validate_with(resp, list, str, criteria_count)) is not None
             else out,
-            **self.prepend(kwargs),
+            **self.prepend_sys_msg(kwargs),
         )
 
     async def draft_rating_criteria_from_examples(
@@ -244,7 +253,7 @@ class GiveRating(WithBriefing, LLMUsage):
                 validator=lambda resp: JsonCapture.validate_with(
                     resp, target_type=list, elements_type=str, length=reasons_count
                 ),
-                **self.prepend(kwargs),
+                **self.prepend_sys_msg(kwargs),
             )
         )
         # extract certain mount of criteria from reasons according to their importance and frequency
@@ -301,7 +310,7 @@ class GiveRating(WithBriefing, LLMUsage):
                 for pair in windows
             ],
             validator=lambda resp: JsonCapture.validate_with(resp, target_type=float),
-            **self.prepend(kwargs),
+            **self.prepend_sys_msg(kwargs),
         )
         weights = [1]
         for rw in relative_weights:

fabricatio/capabilities/review.py

@@ -1,6 +1,6 @@
 """A module that provides functionality to rate tasks based on a rating manual and score range."""
 
-from typing import List, Optional, Self, Set, Unpack, cast
+from typing import Dict, List, Optional, Self, Set, Unpack, cast
 
 from fabricatio._rust_instances import TEMPLATE_MANAGER
 from fabricatio.capabilities.propose import Propose
@@ -121,7 +121,7 @@ class ReviewResult[T](ProposedAble, Display):
             ReviewResult[K]: The current instance with updated reference type.
         """
         self._ref = ref  # pyright: ignore [reportAttributeAccessIssue]
-        return cast(ReviewResult[K], self)
+        return cast("ReviewResult[K]", self)
 
     def deref(self) -> T:
         """Retrieve the referenced object that was reviewed.
@@ -200,13 +200,14 @@ class Review(GiveRating, Propose):
             ReviewResult[Task[T]]: A review result containing identified problems and proposed solutions,
                 with a reference to the original task.
         """
-        return cast(ReviewResult[Task[T]], await self.review_obj(task, **kwargs))
+        return cast("ReviewResult[Task[T]]", await self.review_obj(task, **kwargs))
 
     async def review_string(
         self,
         input_text: str,
         topic: str,
         criteria: Optional[Set[str]] = None,
+        rating_manual: Optional[Dict[str, str]] = None,
         **kwargs: Unpack[ValidateKwargs[ReviewResult[str]]],
     ) -> ReviewResult[str]:
         """Review a string based on specified topic and criteria.
@@ -219,6 +220,7 @@ class Review(GiveRating, Propose):
             topic (str): The subject topic for the review criteria.
             criteria (Optional[Set[str]], optional): A set of criteria for the review.
                 If not provided, criteria will be drafted automatically. Defaults to None.
+            rating_manual (Optional[Dict[str,str]], optional): A dictionary of rating criteria and their corresponding scores.
             **kwargs (Unpack[ValidateKwargs]): Additional keyword arguments for the LLM usage.
 
         Returns:
@@ -227,12 +229,13 @@ class Review(GiveRating, Propose):
         """
         default = None
         if "default" in kwargs:
+            # this `default` is the default for the `propose` method
             default = kwargs.pop("default")
 
         criteria = criteria or (await self.draft_rating_criteria(topic, **kwargs))
         if not criteria:
             raise ValueError("No criteria provided for review.")
-        manual = await self.draft_rating_manual(topic, criteria, **kwargs)
+        manual = rating_manual or await self.draft_rating_manual(topic, criteria, **kwargs)
 
         if default is not None:
             kwargs["default"] = default

fabricatio/capabilities/task.py

@@ -23,7 +23,7 @@ class ProposeTask(WithBriefing, Propose):
         self,
         prompt: str,
         **kwargs: Unpack[ValidateKwargs[Task[T]]],
-    ) -> Task[T]:
+    ) -> Optional[Task[T]]:
         """Asynchronously proposes a task based on a given prompt and parameters.
 
         Parameters:
@@ -37,7 +37,7 @@ class ProposeTask(WithBriefing, Propose):
             logger.error(err := f"{self.name}: Prompt must be provided.")
             raise ValueError(err)
 
-        return await self.propose(Task, prompt, **self.prepend(cast(Dict[str, Any], kwargs)))
+        return await self.propose(Task, prompt, **self.prepend_sys_msg(cast("Dict[str, Any]", kwargs)))
 
 
 class HandleTask(WithBriefing, ToolBoxUsage):
@@ -81,10 +81,10 @@ class HandleTask(WithBriefing, ToolBoxUsage):
         return await self.aask_validate(
             question=q,
             validator=_validator,
-            **self.prepend(cast(Dict[str, Any], kwargs)),
+            **self.prepend_sys_msg(cast("Dict[str, Any]", kwargs)),
         )
 
-    async def handle_fin_grind(
+    async def handle_fine_grind(
         self,
         task: Task,
         data: Dict[str, Any],
@@ -110,4 +110,4 @@ class HandleTask(WithBriefing, ToolBoxUsage):
 
     async def handle(self, task: Task, data: Dict[str, Any], **kwargs: Unpack[ValidateKwargs]) -> Optional[Tuple]:
         """Asynchronously handles a task based on a given task object and parameters."""
-        return await self.handle_fin_grind(task, data, **kwargs)
+        return await self.handle_fine_grind(task, data, **kwargs)

fabricatio/config.py

@@ -48,37 +48,37 @@ class LLMConfig(BaseModel):
     """
 
     model_config = ConfigDict(use_attribute_docstrings=True)
-    api_endpoint: HttpUrl = Field(default=HttpUrl("https://api.openai.com"))
+    api_endpoint: Optional[HttpUrl] = Field(default=HttpUrl("https://api.openai.com"))
     """OpenAI API Endpoint."""
 
-    api_key: SecretStr = Field(default=SecretStr(""))
+    api_key: Optional[SecretStr] = Field(default=SecretStr("sk-setyourkey"))
     """OpenAI API key. Empty by default for security reasons, should be set before use."""
 
-    timeout: PositiveInt = Field(default=300)
+    timeout: Optional[PositiveInt] = Field(default=300)
     """The timeout of the LLM model in seconds. Default is 300 seconds as per request."""
 
-    max_retries: PositiveInt = Field(default=3)
+    max_retries: Optional[PositiveInt] = Field(default=3)
     """The maximum number of retries. Default is 3 retries."""
 
-    model: str = Field(default="gpt-3.5-turbo")
+    model: Optional[str] = Field(default="gpt-3.5-turbo")
     """The LLM model name. Set to 'gpt-3.5-turbo' as per request."""
 
-    temperature: NonNegativeFloat = Field(default=1.0)
+    temperature: Optional[NonNegativeFloat] = Field(default=1.0)
     """The temperature of the LLM model. Controls randomness in generation. Set to 1.0 as per request."""
 
-    stop_sign: str | List[str] = Field(default_factory=lambda: ["\n\n\n", "User:"])
+    stop_sign: Optional[str | List[str]] = Field(default=None)
     """The stop sign of the LLM model. No default stop sign specified."""
 
-    top_p: NonNegativeFloat = Field(default=0.35)
+    top_p: Optional[NonNegativeFloat] = Field(default=0.35)
     """The top p of the LLM model. Controls diversity via nucleus sampling. Set to 0.35 as per request."""
 
-    generation_count: PositiveInt = Field(default=1)
+    generation_count: Optional[PositiveInt] = Field(default=1)
     """The number of generations to generate. Default is 1."""
 
-    stream: bool = Field(default=False)
+    stream: Optional[bool] = Field(default=False)
     """Whether to stream the LLM model's response. Default is False."""
 
-    max_tokens: PositiveInt = Field(default=8192)
+    max_tokens: Optional[PositiveInt] = Field(default=None)
     """The maximum number of tokens to generate. Set to 8192 as per request."""
 
     rpm: Optional[PositiveInt] = Field(default=100)
@@ -93,7 +93,7 @@ class EmbeddingConfig(BaseModel):
 
     model_config = ConfigDict(use_attribute_docstrings=True)
 
-    model: str = Field(default="text-embedding-ada-002")
+    model: Optional[str] = Field(default="text-embedding-ada-002")
     """The embedding model name. """
 
     dimensions: Optional[PositiveInt] = Field(default=None)
@@ -102,10 +102,10 @@ class EmbeddingConfig(BaseModel):
     timeout: Optional[PositiveInt] = Field(default=None)
     """The timeout of the embedding model in seconds."""
 
-    max_sequence_length: PositiveInt = Field(default=8192)
+    max_sequence_length: Optional[PositiveInt] = Field(default=8192)
     """The maximum sequence length of the embedding model. Default is 8192 as per request."""
 
-    caching: bool = Field(default=False)
+    caching: Optional[bool] = Field(default=False)
     """Whether to cache the embedding. Default is False."""
 
     api_endpoint: Optional[HttpUrl] = None
@@ -148,13 +148,13 @@ class DebugConfig(BaseModel):
     log_level: Literal["DEBUG", "INFO", "SUCCESS", "WARNING", "ERROR", "CRITICAL"] = Field(default="INFO")
     """The log level of the application."""
 
-    log_file: FilePath = Field(default=Path(rf"{ROAMING_DIR}\fabricatio.log"))
+    log_file: FilePath = Field(default=Path(rf"{ROAMING_DIR}\fabricatio.log"), frozen=True)
     """The log file of the application."""
 
-    rotation: int = Field(default=1)
+    rotation: int = Field(default=1, frozen=True)
     """The rotation of the log file. in weeks."""
 
-    retention: int = Field(default=2)
+    retention: int = Field(default=2, frozen=True)
     """The retention of the log file. in weeks."""
 
     streaming_visible: bool = Field(default=False)
@@ -232,6 +232,12 @@ class TemplateConfig(BaseModel):
     correct_template: str = Field(default="correct")
     """The name of the correct template which will be used to correct a string."""
 
+    co_validation_template: str = Field(default="co_validation")
+    """The name of the co-validation template which will be used to co-validate a string."""
+
+    as_prompt_template: str = Field(default="as_prompt")
+    """The name of the as prompt template which will be used to convert a string to a prompt."""
+
 
 class MagikaConfig(BaseModel):
     """Magika configuration class."""
@@ -272,9 +278,9 @@ class RagConfig(BaseModel):
 
     model_config = ConfigDict(use_attribute_docstrings=True)
 
-    milvus_uri: HttpUrl = Field(default=HttpUrl("http://localhost:19530"))
+    milvus_uri: Optional[HttpUrl] = Field(default=HttpUrl("http://localhost:19530"))
     """The URI of the Milvus server."""
-    milvus_timeout: Optional[PositiveFloat] = Field(default=None)
+    milvus_timeout: Optional[PositiveFloat] = Field(default=30.0)
     """The timeout of the Milvus server."""
     milvus_token: Optional[SecretStr] = Field(default=None)
     """The token of the Milvus server."""
@@ -300,11 +306,13 @@ class RoutingConfig(BaseModel):
 
     model_config = ConfigDict(use_attribute_docstrings=True)
 
-    allowed_fails: Optional[int] = 1
+    max_parallel_requests: Optional[int] = 60
+    """The maximum number of parallel requests. None means not checked."""
+    allowed_fails: Optional[int] = 3
     """The number of allowed fails before the routing is considered failed."""
     retry_after: int = 15
     """The time in seconds to wait before retrying the routing after a fail."""
-    cooldown_time: Optional[int] = 120
+    cooldown_time: Optional[int] = 30
     """The time in seconds to wait before retrying the routing after a cooldown."""
 
 

fabricatio/decorators.py

@@ -177,3 +177,35 @@ def use_temp_module[**P, R](modules: ModuleType | List[ModuleType]) -> Callable[
         return _wrapper
 
     return _decorator
+
+
+def logging_exec_time[**P, R](func: Callable[P, R]) -> Callable[P, R]:
+    """Decorator to log the execution time of a function.
+
+    Args:
+        func (Callable): The function to be executed
+
+    Returns:
+        Callable: A decorator that wraps the function to log the execution time.
+    """
+    from time import time
+
+    if iscoroutinefunction(func):
+
+        @wraps(func)
+        async def _async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            start_time = time()
+            result = await func(*args, **kwargs)
+            logger.debug(f"Execution time of `{func.__name__}`: {time() - start_time:.2f} s")
+            return result
+
+        return _async_wrapper
+
+    @wraps(func)
+    def _wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        start_time = time()
+        result = func(*args, **kwargs)
+        logger.debug(f"Execution time of {func.__name__}: {(time() - start_time) * 1000:.2f} ms")
+        return result
+
+    return _wrapper