fabricatio 0.2.1.dev4__cp312-cp312-win_amd64.whl → 0.2.3__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

@@ -0,0 +1,310 @@
+"""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, overload
+
+from fabricatio._rust_instances import template_manager
+from fabricatio.config import configs
+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
+
+
+@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(EmbeddingUsage):
+    """A class representing the RAG (Retrieval Augmented Generation) model."""
+
+    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:
+        """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 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 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
+        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, 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.
+        """
+        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]
+        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
+
+    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.
+
+        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.
+        """
+        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(
+        self,
+        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.
+
+        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.
+            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:
+            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,
+            search_params={"radius": similarity_threshold},
+            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)
+
+        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]
+
+        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,
+        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.
+            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,
+                **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/rating.py

@@ -12,7 +12,7 @@ from fabricatio.models.generic import WithBriefing
 from fabricatio.models.kwargs_types import GenerateKwargs, ValidateKwargs
 from fabricatio.models.usages import LLMUsage
 from fabricatio.parser import JsonCapture
-from more_itertools import flatten
+from more_itertools import flatten, windowed
 from pydantic import NonNegativeInt, PositiveInt
 
 
@@ -275,3 +275,81 @@ class GiveRating(WithBriefing, LLMUsage):
             validator=_criteria_validator,
             **kwargs,
         )
+
+    async def drafting_rating_weights_klee(
+        self,
+        topic: str,
+        criteria: Set[str],
+        **kwargs: Unpack[ValidateKwargs],
+    ) -> Dict[str, float]:
+        """Drafts rating weights for a given topic and criteria using the Klee method.
+
+        Args:
+            topic (str): The topic for the rating weights.
+            criteria (Set[str]): A set of criteria for the rating weights.
+            **kwargs (Unpack[ValidateKwargs]): Additional keyword arguments for the LLM usage.
+
+        Returns:
+            Dict[str, float]: A dictionary representing the drafted rating weights for each criterion.
+        """
+        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)
+
+        # get the importance multiplier indicating how important is second criterion compared to the first one
+        relative_weights = await self.aask_validate_batch(
+            questions=[
+                template_manager.render_template(
+                    configs.templates.draft_rating_weights_klee_template,
+                    {
+                        "topic": topic,
+                        "first": pair[0],
+                        "second": pair[1],
+                    },
+                )
+                for pair in windows
+            ],
+            validator=_validator,
+            **GenerateKwargs(system_message=f"# your personal briefing: \n{self.briefing}", **kwargs),
+        )
+        weights = [1]
+        for rw in relative_weights:
+            weights.append(weights[-1] * rw)
+        total = sum(weights)
+        return dict(zip(criteria, [w / total for w in weights], strict=True))
+
+    async def composite_score(
+        self,
+        topic: str,
+        to_rate: List[str],
+        reasons_count: PositiveInt = 2,
+        criteria_count: PositiveInt = 5,
+        **kwargs: Unpack[ValidateKwargs],
+    ) -> List[float]:
+        """Calculates the composite scores for a list of items based on a given topic and criteria.
+
+        Args:
+            topic (str): The topic for the rating.
+            to_rate (List[str]): A list of strings to be rated.
+            reasons_count (PositiveInt, optional): The number of reasons to extract from each pair of examples. Defaults to 2.
+            criteria_count (PositiveInt, optional): The number of criteria to draft. Defaults to 5.
+            **kwargs (Unpack[ValidateKwargs]): Additional keyword arguments for the LLM usage.
+
+        Returns:
+            List[float]: A list of composite scores for the items.
+        """
+        criteria = await self.draft_rating_criteria_from_examples(
+            topic, to_rate, reasons_count, criteria_count, **kwargs
+        )
+        weights = await self.drafting_rating_weights_klee(topic, criteria, **kwargs)
+        logger.info(f"Criteria: {criteria}\nWeights: {weights}")
+        ratings_seq = await self.rate(to_rate, topic, criteria, **kwargs)
+
+        return [sum(ratings[c] * weights[c] for c in criteria) for ratings in ratings_seq]

fabricatio/config.py

@@ -11,6 +11,7 @@ from pydantic import (
     FilePath,
     HttpUrl,
     NonNegativeFloat,
+    PositiveFloat,
     PositiveInt,
     SecretStr,
 )
@@ -79,6 +80,30 @@ class LLMConfig(BaseModel):
     """The maximum number of tokens to generate. Set to 8192 as per request."""
 
 
+class EmbeddingConfig(BaseModel):
+    """Embedding configuration class."""
+
+    model_config = ConfigDict(use_attribute_docstrings=True)
+
+    model: str = Field(default="text-embedding-ada-002")
+    """The embedding model name. """
+
+    dimensions: Optional[PositiveInt] = Field(default=None)
+    """The dimensions of the embedding. Default is None."""
+
+    timeout: Optional[PositiveInt] = Field(default=None)
+    """The timeout of the embedding model in seconds. Default is 300 seconds as per request."""
+
+    caching: bool = Field(default=False)
+    """Whether to cache the embedding. Default is True."""
+
+    api_endpoint: Optional[HttpUrl] = None
+    """The OpenAI API endpoint."""
+
+    api_key: Optional[SecretStr] = None
+    """The OpenAI API key."""
+
+
 class PymitterConfig(BaseModel):
     """Pymitter configuration class.
 
@@ -172,6 +197,12 @@ class TemplateConfig(BaseModel):
     extract_criteria_from_reasons_template: str = Field(default="extract_criteria_from_reasons")
     """The name of the extract criteria from reasons template which will be used to extract criteria from reasons."""
 
+    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."""
@@ -204,6 +235,21 @@ class ToolBoxConfig(BaseModel):
     """The name of the module containing the data."""
 
 
+class RagConfig(BaseModel):
+    """RAG configuration class."""
+
+    model_config = ConfigDict(use_attribute_docstrings=True)
+
+    milvus_uri: HttpUrl = Field(default=HttpUrl("http://localhost:19530"))
+    """The URI of the Milvus server."""
+    milvus_timeout: Optional[PositiveFloat] = Field(default=None)
+    """The timeout of the Milvus server."""
+    milvus_token: Optional[SecretStr] = Field(default=None)
+    """The token of the Milvus server."""
+    milvus_dimensions: Optional[PositiveInt] = Field(default=None)
+    """The dimensions of the Milvus server."""
+
+
 class Settings(BaseSettings):
     """Application settings class.
 
@@ -229,6 +275,9 @@ class Settings(BaseSettings):
     llm: LLMConfig = Field(default_factory=LLMConfig)
     """LLM Configuration"""
 
+    embedding: EmbeddingConfig = Field(default_factory=EmbeddingConfig)
+    """Embedding Configuration"""
+
     debug: DebugConfig = Field(default_factory=DebugConfig)
     """Debug Configuration"""
 
@@ -247,6 +296,9 @@ class Settings(BaseSettings):
     toolbox: ToolBoxConfig = Field(default_factory=ToolBoxConfig)
     """Toolbox Configuration"""
 
+    rag: RagConfig = Field(default_factory=RagConfig)
+    """RAG Configuration"""
+
     @classmethod
     def settings_customise_sources(
         cls,

fabricatio/core.py

@@ -38,11 +38,11 @@ class Env(BaseModel):
 
     @overload
     def on[**P, R](
-            self,
-            event: str | Event,
-            func: Optional[Callable[P, R]] = None,
-            /,
-            ttl: int = -1,
+        self,
+        event: str | Event,
+        func: Optional[Callable[P, R]] = None,
+        /,
+        ttl: int = -1,
     ) -> Callable[[Callable[P, R]], Callable[P, R]]:
         """
         Registers an event listener with a specific function that listens indefinitely or for a specified number of times.
@@ -58,11 +58,11 @@ class Env(BaseModel):
         ...
 
     def on[**P, R](
-            self,
-            event: str | Event,
-            func: Optional[Callable[P, R]] = None,
-            /,
-            ttl=-1,
+        self,
+        event: str | Event,
+        func: Optional[Callable[P, R]] = None,
+        /,
+        ttl=-1,
     ) -> Callable[[Callable[P, R]], Callable[P, R]] | Self:
         """Registers an event listener with a specific function that listens indefinitely or for a specified number of times.
 
@@ -78,14 +78,13 @@ class Env(BaseModel):
             event = event.collapse()
         if func is None:
             return self._ee.on(event, ttl=ttl)
-
         self._ee.on(event, func, ttl=ttl)
         return self
 
     @overload
     def once[**P, R](
-            self,
-            event: str | Event,
+        self,
+        event: str | Event,
     ) -> Callable[[Callable[P, R]], Callable[P, R]]:
         """
         Registers an event listener that listens only once.
@@ -100,9 +99,9 @@ class Env(BaseModel):
 
     @overload
     def once[**P, R](
-            self,
-            event: str | Event,
-            func: Callable[[Callable[P, R]], Callable[P, R]],
+        self,
+        event: str | Event,
+        func: Callable[[Callable[P, R]], Callable[P, R]],
     ) -> Self:
         """
         Registers an event listener with a specific function that listens only once.
@@ -117,9 +116,9 @@ class Env(BaseModel):
         ...
 
     def once[**P, R](
-            self,
-            event: str | Event,
-            func: Optional[Callable[P, R]] = None,
+        self,
+        event: str | Event,
+        func: Optional[Callable[P, R]] = None,
     ) -> Callable[[Callable[P, R]], Callable[P, R]] | Self:
         """Registers an event listener with a specific function that listens only once.
 
@@ -163,5 +162,20 @@ class Env(BaseModel):
             event = event.collapse()
         return await self._ee.emit_async(event, *args, **kwargs)
 
+    def emit_future[**P](self, event: str | Event, *args: P.args, **kwargs: P.kwargs) -> None:
+        """Emits an event to all registered listeners and returns a future object.
+
+        Args:
+            event (str | Event): The event to emit.
+            *args: Positional arguments to pass to the listeners.
+            **kwargs: Keyword arguments to pass to the listeners.
+
+        Returns:
+            None: The future object.
+        """
+        if isinstance(event, Event):
+            event = event.collapse()
+        return self._ee.emit_future(event, *args, **kwargs)
+
 
 env = Env()

fabricatio/models/action.py

@@ -2,7 +2,7 @@
 
 import traceback
 from abc import abstractmethod
-from asyncio import Queue
+from asyncio import Queue, create_task
 from typing import Any, Dict, Self, Tuple, Type, Union, Unpack
 
 from fabricatio.capabilities.rating import GiveRating
@@ -108,7 +108,11 @@ class WorkFlow(WithBriefing, ToolBoxUsage):
         try:
             for step in self._instances:
                 logger.debug(f"Executing step: {step.name}")
-                modified_ctx = await step.act(await self._context.get())
+                act_task = create_task(step.act(await self._context.get()))
+                if task.is_cancelled():
+                    act_task.cancel(f"Cancelled by task: {task.name}")
+                    break
+                modified_ctx = await act_task
                 await self._context.put(modified_ctx)
                 current_action = step.name
             logger.info(f"Finished executing workflow: {self.name}")