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

fabricatio/capabilities/rag.py

@@ -0,0 +1,179 @@
+"""A module for the RAG (Retrieval Augmented Generation) model."""
+
+from operator import itemgetter
+from os import PathLike
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Self, Union
+
+from fabricatio.config import configs
+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 PrivateAttr
+
+
+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)
+
+    @property
+    def client(self) -> MilvusClient:
+        """The Milvus client."""
+        return self._client
+
+    def view(self, collection_name: 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):
+            self._client.create_collection(collection_name)
+
+        self._target_collection = collection_name
+        return self
+
+    def quit_view(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
+
+    @property
+    def safe_viewing_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 or self.safe_viewing_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_viewing_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_viewing_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]

fabricatio/config.py

@@ -11,6 +11,7 @@ from pydantic import (
     FilePath,
     HttpUrl,
     NonNegativeFloat,
+    PositiveFloat,
     PositiveInt,
     SecretStr,
 )
@@ -207,6 +208,19 @@ 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."""
+
+
 class Settings(BaseSettings):
     """Application settings class.
 
@@ -250,6 +264,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/models/kwargs_types.py

@@ -5,6 +5,15 @@ from typing import List, NotRequired, TypedDict
 from pydantic import NonNegativeFloat, NonNegativeInt, PositiveInt
 
 
+class EmbeddingKwargs(TypedDict):
+    """A type representing the keyword arguments for the embedding method."""
+
+    model: NotRequired[str]
+    dimensions: NotRequired[int]
+    timeout: NotRequired[PositiveInt]
+    caching: NotRequired[bool]
+
+
 class LLMKwargs(TypedDict):
     """A type representing the keyword arguments for the LLM (Large Language Model) usage."""
 

fabricatio/models/usages.py

@@ -10,14 +10,15 @@ from fabricatio._rust_instances import template_manager
 from fabricatio.config import configs
 from fabricatio.journal import logger
 from fabricatio.models.generic import Base, WithBriefing
-from fabricatio.models.kwargs_types import ChooseKwargs, GenerateKwargs, LLMKwargs
+from fabricatio.models.kwargs_types import ChooseKwargs, EmbeddingKwargs, GenerateKwargs, LLMKwargs
 from fabricatio.models.task import Task
 from fabricatio.models.tool import Tool, ToolBox
-from fabricatio.models.utils import Messages
+from fabricatio.models.utils import Messages, MilvusData
 from fabricatio.parser import JsonCapture
 from litellm import stream_chunk_builder
 from litellm.types.utils import (
     Choices,
+    EmbeddingResponse,
     ModelResponse,
     StreamingChoices,
 )
@@ -61,6 +62,97 @@ class LLMUsage(Base):
     llm_max_tokens: Optional[PositiveInt] = None
     """The maximum number of tokens to generate."""
 
+    async def aembedding(
+        self,
+        input_text: List[str],
+        model: Optional[str] = None,
+        dimensions: Optional[int] = None,
+        timeout: Optional[PositiveInt] = None,
+        caching: Optional[bool] = False,
+    ) -> EmbeddingResponse:
+        """Asynchronously generates embeddings for the given input text.
+
+        Args:
+            input_text (List[str]): A list of strings to generate embeddings for.
+            model (Optional[str]): The model to use for embedding. Defaults to the instance's `llm_model` or the global configuration.
+            dimensions (Optional[int]): The dimensions of the embedding. Defaults to None.
+            timeout (Optional[PositiveInt]): The timeout for the embedding request. Defaults to the instance's `llm_timeout` or the global configuration.
+            caching (Optional[bool]): Whether to cache the embedding result. Defaults to False.
+
+
+        Returns:
+            EmbeddingResponse: The response containing the embeddings.
+        """
+        return await litellm.aembedding(
+            input=input_text,
+            caching=caching,
+            dimensions=dimensions,
+            model=model or self.llm_model or configs.llm.model,
+            timeout=timeout or self.llm_timeout or configs.llm.timeout,
+            api_key=self.llm_api_key.get_secret_value() if self.llm_api_key else configs.llm.api_key.get_secret_value(),
+            api_base=self.llm_api_endpoint.unicode_string().rstrip(
+                "/"
+            )  # seems embedding function takes no base_url end with a slash
+            if self.llm_api_endpoint
+            else configs.llm.api_endpoint.unicode_string().rstrip("/"),
+        )
+
+    @overload
+    async def vectorize(self, input_text: List[str], **kwargs: Unpack[EmbeddingKwargs]) -> List[List[float]]: ...
+    @overload
+    async def vectorize(self, input_text: str, **kwargs: Unpack[EmbeddingKwargs]) -> List[float]: ...
+
+    async def vectorize(
+        self, input_text: List[str] | str, **kwargs: Unpack[EmbeddingKwargs]
+    ) -> List[List[float]] | List[float]:
+        """Asynchronously generates vector embeddings for the given input text.
+
+        Args:
+            input_text (List[str] | str): A string or list of strings to generate embeddings for.
+            **kwargs (Unpack[EmbeddingKwargs]): Additional keyword arguments for embedding.
+
+        Returns:
+            List[List[float]] | List[float]: The generated embeddings.
+        """
+        if isinstance(input_text, str):
+            return (await self.aembedding([input_text], **kwargs)).data[0].get("embedding")
+
+        return [o.get("embedding") for o in (await self.aembedding(input_text, **kwargs)).data]
+
+    @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)
+        ]
+
     async def aquery(
         self,
         messages: List[Dict[str, str]],

fabricatio/models/utils.py

@@ -1,6 +1,6 @@
 """A module containing utility classes for the models."""
 
-from typing import Dict, List, Literal, Self
+from typing import Any, Dict, List, Literal, Optional, Self
 
 from pydantic import BaseModel, ConfigDict, Field
 
@@ -76,3 +76,52 @@ class Messages(list):
             list[dict]: A list of dictionaries representing the messages.
         """
         return [message.model_dump() for message in self]
+
+
+class MilvusData(BaseModel):
+    """A class representing data stored in Milvus."""
+
+    model_config = ConfigDict(use_attribute_docstrings=True)
+    id: Optional[int] = Field(default=None)
+    """The identifier of the data."""
+
+    vector: List[float]
+    """The vector representation of the data."""
+
+    text: str
+    """The text representation of the data."""
+
+    subject: Optional[str] = Field(default=None)
+    """A subject label that we use to demo metadata filtering later."""
+
+    def prepare_insertion(self) -> Dict[str, Any]:
+        """Prepares the data for insertion into Milvus.
+
+        Returns:
+            dict: A dictionary containing the data to be inserted into Milvus.
+        """
+        return self.model_dump(exclude_none=True)
+
+    def update_subject(self, new_subject: str) -> Self:
+        """Updates the subject label of the data.
+
+        Args:
+            new_subject (str): The new subject label.
+
+        Returns:
+            Self: The updated instance of MilvusData.
+        """
+        self.subject = new_subject
+        return self
+
+    def update_id(self, new_id: int) -> Self:
+        """Updates the identifier of the data.
+
+        Args:
+            new_id (int): The new identifier.
+
+        Returns:
+            Self: The updated instance of MilvusData.
+        """
+        self.id = new_id
+        return self

{fabricatio-0.2.3.dev0.dist-info → fabricatio-0.2.3.dev2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fabricatio
-Version: 0.2.3.dev0
+Version: 0.2.3.dev2
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: 3.12

{fabricatio-0.2.3.dev0.dist-info → fabricatio-0.2.3.dev2.dist-info}/RECORD

@@ -1,12 +1,13 @@
-fabricatio-0.2.3.dev0.dist-info/METADATA,sha256=AbSeMTWK-eEnD_dWakSjEUsP-XE9mBHA-8_1JK3EYr4,12339
-fabricatio-0.2.3.dev0.dist-info/WHEEL,sha256=tpW5AN9B-9qsM9WW2FXG2r193YXiqexDadpKp0A2daI,96
-fabricatio-0.2.3.dev0.dist-info/licenses/LICENSE,sha256=do7J7EiCGbq0QPbMAL_FqLYufXpHnCnXBOuqVPwSV8Y,1088
+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/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/rating.py,sha256=zmTUvsUfxFgovRQzy4djL2zKRYTHmN6JY7A4lyT5uVQ,14907
 fabricatio/capabilities/task.py,sha256=d2xtrwQxXWI40UskQCR5YhHarY7ST0ppr8TjY12uWQE,5327
-fabricatio/config.py,sha256=wzaaUHZZMRCYc37M_M4qKuLOYtwdEjYtyG77-AGkqCg,11467
+fabricatio/config.py,sha256=6_cpyGu87NmuPC2Kc0fq_cwOogH3fjXTILiDY7poWi4,12041
 fabricatio/core.py,sha256=VQ_JKgUGIy2gZ8xsTBZCdr_IP7wC5aPg0_bsOmjQ588,6458
 fabricatio/decorators.py,sha256=uzsP4tFKQNjDHBkofsjjoJA0IUAaYOtt6YVedoyOqlo,6551
 fabricatio/fs/curd.py,sha256=faMstgGUiQ4k2AW3OXfvvWWTldTtKXco7QINYaMjmyA,3981
@@ -16,12 +17,12 @@ fabricatio/journal.py,sha256=siqimKF0M_QaaOCMxtjr_BJVNyUIAQWILzE9Q4T6-7c,781
 fabricatio/models/action.py,sha256=NpklAVUHYO5JIY9YLwYowZ-U8R9CFf5aC10DhLF7gxQ,5924
 fabricatio/models/events.py,sha256=mrihNEFgQ5o7qFWja1z_qX8dnaTLwPBoJdVlzxQV5oM,2719
 fabricatio/models/generic.py,sha256=WEjZ96rTyBjaBjkM6e8E4Pg_Naot4xWRvGJteqBiCCI,5133
-fabricatio/models/kwargs_types.py,sha256=nTtD3wzSpCg-NlrJ43yW6lmfeWzD2V_XGMPlL5mXzyc,1147
+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/tool.py,sha256=WTFnpF6xZ1nJbmIOonLsGQcM-kkDCeZiAFqyil9xg2U,6988
-fabricatio/models/usages.py,sha256=iLxas1gE7MA55ZtQJJ-qu3W6JP5KLjPgmNKqNYIF6yU,23972
-fabricatio/models/utils.py,sha256=i_kpcQpct04mQFk1nbcVGV-pl1YThWu4Qk3wbewzKkc,2535
+fabricatio/models/usages.py,sha256=bzsTDrAekiQyIwKeWds5YdgsXk8qiZkD7OZdno1Q_Ck,28213
+fabricatio/models/utils.py,sha256=mXea76bd4r2jy_zx74GM4t5kCvkMu0JTOaw_VGvTCxk,3952
 fabricatio/parser.py,sha256=uLabsvF07wRKW1PoTGuGEENCx3P4mhmuO8JkmOEkKko,3522
 fabricatio/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 fabricatio/toolboxes/arithmetic.py,sha256=WLqhY-Pikv11Y_0SGajwZx3WhsLNpHKf9drzAqOf_nY,1369
@@ -31,6 +32,6 @@ fabricatio/toolboxes/__init__.py,sha256=b13KmASO8q5fBLwew964fn9oH86ER5g-S1PgA4fZ
 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=MMgVw-oe5Yw0aBgcHHDX4FExLxDD0Xx1mnDmzDM7Wcc,1266176
-fabricatio-0.2.3.dev0.data/scripts/tdown.exe,sha256=AqyvPpwXz-hImUl93bB9nQO67AtojXG53SDN-Ltor8M,3397120
-fabricatio-0.2.3.dev0.dist-info/RECORD,,
+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,,