fabricatio 0.2.1.dev4__cp312-cp312-win_amd64.whl → 0.2.3__cp312-cp312-win_amd64.whl

fabricatio/models/generic.py

@@ -1,7 +1,7 @@
 """This module defines generic classes for models in the Fabricatio library."""
 
 from pathlib import Path
-from typing import Callable, List, Self
+from typing import Callable, Iterable, List, Optional, Self, Union, final
 
 import orjson
 from fabricatio._rust import blake3_hash
@@ -12,6 +12,11 @@ from pydantic import (
     BaseModel,
     ConfigDict,
     Field,
+    HttpUrl,
+    NonNegativeFloat,
+    PositiveFloat,
+    PositiveInt,
+    SecretStr,
 )
 
 
@@ -150,3 +155,104 @@ class WithDependency(Base):
                 for p in self.dependencies
             },
         )
+
+
+class ScopedConfig(Base):
+    """Class that manages a scoped configuration."""
+
+    llm_api_endpoint: Optional[HttpUrl] = None
+    """The OpenAI API endpoint."""
+
+    llm_api_key: Optional[SecretStr] = None
+    """The OpenAI API key."""
+
+    llm_timeout: Optional[PositiveInt] = None
+    """The timeout of the LLM model."""
+
+    llm_max_retries: Optional[PositiveInt] = None
+    """The maximum number of retries."""
+
+    llm_model: Optional[str] = None
+    """The LLM model name."""
+
+    llm_temperature: Optional[NonNegativeFloat] = None
+    """The temperature of the LLM model."""
+
+    llm_stop_sign: Optional[str | List[str]] = None
+    """The stop sign of the LLM model."""
+
+    llm_top_p: Optional[NonNegativeFloat] = None
+    """The top p of the LLM model."""
+
+    llm_generation_count: Optional[PositiveInt] = None
+    """The number of generations to generate."""
+
+    llm_stream: Optional[bool] = None
+    """Whether to stream the LLM model's response."""
+
+    llm_max_tokens: Optional[PositiveInt] = None
+    """The maximum number of tokens to generate."""
+
+    embedding_api_endpoint: Optional[HttpUrl] = None
+    """The OpenAI API endpoint."""
+
+    embedding_api_key: Optional[SecretStr] = None
+    """The OpenAI API key."""
+
+    embedding_timeout: Optional[PositiveInt] = None
+    """The timeout of the LLM model."""
+
+    embedding_model: Optional[str] = None
+    """The LLM model name."""
+
+    embedding_dimensions: Optional[PositiveInt] = None
+    """The dimensions of the embedding."""
+    embedding_caching: Optional[bool] = False
+    """Whether to cache the embedding result."""
+
+    milvus_uri: Optional[HttpUrl] = Field(default=None)
+    """The URI of the Milvus server."""
+    milvus_token: Optional[SecretStr] = Field(default=None)
+    """The token for the Milvus server."""
+    milvus_timeout: Optional[PositiveFloat] = Field(default=None)
+    """The timeout for the Milvus server."""
+    milvus_dimensions: Optional[PositiveInt] = Field(default=None)
+    """The dimensions of the Milvus server."""
+
+    @final
+    def fallback_to(self, other: "ScopedConfig") -> Self:
+        """Fallback to another instance's attribute values if the current instance's attributes are None.
+
+        Args:
+            other (LLMUsage): Another instance from which to copy attribute values.
+
+        Returns:
+            Self: The current instance, allowing for method chaining.
+        """
+        # Iterate over the attribute names and copy values from 'other' to 'self' where applicable
+        # noinspection PydanticTypeChecker,PyTypeChecker
+        for attr_name in ScopedConfig.model_fields:
+            # Copy the attribute value from 'other' to 'self' only if 'self' has None and 'other' has a non-None value
+            if getattr(self, attr_name) is None and (attr := getattr(other, attr_name)) is not None:
+                setattr(self, attr_name, attr)
+
+        # Return the current instance to allow for method chaining
+        return self
+
+    @final
+    def hold_to(self, others: Union["ScopedConfig", Iterable["ScopedConfig"]]) -> Self:
+        """Hold to another instance's attribute values if the current instance's attributes are None.
+
+        Args:
+            others (LLMUsage | Iterable[LLMUsage]): Another instance or iterable of instances from which to copy attribute values.
+
+        Returns:
+            Self: The current instance, allowing for method chaining.
+        """
+        if not isinstance(others, Iterable):
+            others = [others]
+        for other in others:
+            # noinspection PyTypeChecker,PydanticTypeChecker
+            for attr_name in ScopedConfig.model_fields:
+                if (attr := getattr(self, attr_name)) is not None and getattr(other, attr_name) is None:
+                    setattr(other, attr_name, attr)

fabricatio/models/kwargs_types.py

@@ -5,6 +5,29 @@ from typing import List, NotRequired, TypedDict
 from pydantic import NonNegativeFloat, NonNegativeInt, PositiveInt
 
 
+class CollectionSimpleConfigKwargs(TypedDict):
+    """A type representing the configuration for a collection."""
+
+    dimension: NotRequired[int]
+    timeout: NotRequired[float]
+
+
+class FetchKwargs(TypedDict):
+    """A type representing the keyword arguments for the fetch method."""
+
+    similarity_threshold: NotRequired[float]
+    result_per_query: NotRequired[int]
+
+
+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/task.py

@@ -46,21 +46,21 @@ 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 = PrivateAttr(default_factory=lambda: Queue(maxsize=1))
+    _output: Queue[T | None] = PrivateAttr(default_factory=Queue)
     """The output queue of the task."""
 
     _status: TaskStatus = PrivateAttr(default=TaskStatus.Pending)
@@ -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,12 +126,12 @@ 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
 
-    async def get_output(self) -> T:
+    async def get_output(self) -> T | None:
         """Get the output of the task.
 
         Returns:
@@ -232,6 +232,7 @@ class Task[T](WithBriefing, WithJsonExample, WithDependency):
         """
         logger.info(f"Cancelling task `{self.name}`")
         self._status = TaskStatus.Cancelled
+        await self._output.put(None)
         await env.emit_async(self.cancelled_label, self)
         return self
 
@@ -243,27 +244,38 @@ class Task[T](WithBriefing, WithJsonExample, WithDependency):
         """
         logger.info(f"Failing task `{self.name}`")
         self._status = TaskStatus.Failed
+        await self._output.put(None)
         await env.emit_async(self.failed_label, self)
         return self
 
-    async def publish(self) -> Self:
+    def publish(self, new_namespace: Optional[EventLike] = None) -> Self:
         """Publish the task to the event bus.
 
+        Args:
+            new_namespace(EventLike, optional): The new namespace to move the task to.
+
         Returns:
-            Task: The published instance of the `Task` class
+            Task: The published instance of the `Task` class.
         """
+        if new_namespace:
+            self.move_to(new_namespace)
         logger.info(f"Publishing task `{(label := self.pending_label)}`")
-        await env.emit_async(label, self)
+        env.emit_future(label, self)
         return self
 
-    async def delegate(self) -> T:
-        """Delegate the task to the event bus and wait for the output.
+    async def delegate(self, new_namespace: Optional[EventLike] = None) -> T | None:
+        """Delegate the task to the event.
+
+        Args:
+            new_namespace(EventLike, optional): The new namespace to move the task to.
 
         Returns:
-            T: The output of the task
+            T|None: The output of the task.
         """
+        if new_namespace:
+            self.move_to(new_namespace)
         logger.info(f"Delegating task `{(label := self.pending_label)}`")
-        await env.emit_async(label, self)
+        env.emit_future(label, self)
         return await self.get_output()
 
     @property
@@ -277,3 +289,43 @@ class Task[T](WithBriefing, WithJsonExample, WithDependency):
             configs.templates.task_briefing_template,
             self.model_dump(),
         )
+
+    def is_running(self) -> bool:
+        """Check if the task is running.
+
+        Returns:
+            bool: True if the task is running, False otherwise.
+        """
+        return self._status == TaskStatus.Running
+
+    def is_finished(self) -> bool:
+        """Check if the task is finished.
+
+        Returns:
+            bool: True if the task is finished, False otherwise.
+        """
+        return self._status == TaskStatus.Finished
+
+    def is_failed(self) -> bool:
+        """Check if the task is failed.
+
+        Returns:
+            bool: True if the task is failed, False otherwise.
+        """
+        return self._status == TaskStatus.Failed
+
+    def is_cancelled(self) -> bool:
+        """Check if the task is cancelled.
+
+        Returns:
+            bool: True if the task is cancelled, False otherwise.
+        """
+        return self._status == TaskStatus.Cancelled
+
+    def is_pending(self) -> bool:
+        """Check if the task is pending.
+
+        Returns:
+            bool: True if the task is pending, False otherwise.
+        """
+        return self._status == TaskStatus.Pending

fabricatio/models/usages.py

@@ -1,7 +1,7 @@
 """This module contains classes that manage the usage of language models and tools in tasks."""
 
 from asyncio import gather
-from typing import Callable, Dict, Iterable, List, Optional, Self, Set, Union, Unpack, overload
+from typing import Callable, Dict, Iterable, List, Optional, Self, Set, Type, Union, Unpack, overload
 
 import asyncstdlib
 import litellm
@@ -9,8 +9,8 @@ import orjson
 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.generic import ScopedConfig, WithBriefing
+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
@@ -18,48 +18,20 @@ from fabricatio.parser import JsonCapture
 from litellm import stream_chunk_builder
 from litellm.types.utils import (
     Choices,
+    EmbeddingResponse,
     ModelResponse,
     StreamingChoices,
 )
 from litellm.utils import CustomStreamWrapper
-from pydantic import Field, HttpUrl, NonNegativeFloat, NonNegativeInt, PositiveInt, SecretStr
+from pydantic import Field, NonNegativeInt, PositiveInt
 
 
-class LLMUsage(Base):
+class LLMUsage(ScopedConfig):
     """Class that manages LLM (Large Language Model) usage parameters and methods."""
 
-    llm_api_endpoint: Optional[HttpUrl] = None
-    """The OpenAI API endpoint."""
-
-    llm_api_key: Optional[SecretStr] = None
-    """The OpenAI API key."""
-
-    llm_timeout: Optional[PositiveInt] = None
-    """The timeout of the LLM model."""
-
-    llm_max_retries: Optional[PositiveInt] = None
-    """The maximum number of retries."""
-
-    llm_model: Optional[str] = None
-    """The LLM model name."""
-
-    llm_temperature: Optional[NonNegativeFloat] = None
-    """The temperature of the LLM model."""
-
-    llm_stop_sign: Optional[str | List[str]] = None
-    """The stop sign of the LLM model."""
-
-    llm_top_p: Optional[NonNegativeFloat] = None
-    """The top p of the LLM model."""
-
-    llm_generation_count: Optional[PositiveInt] = None
-    """The number of generations to generate."""
-
-    llm_stream: Optional[bool] = None
-    """Whether to stream the LLM model's response."""
-
-    llm_max_tokens: Optional[PositiveInt] = None
-    """The maximum number of tokens to generate."""
+    @classmethod
+    def _scoped_model(cls) -> Type["LLMUsage"]:
+        return LLMUsage
 
     async def aquery(
         self,
@@ -89,10 +61,8 @@ class LLMUsage(Base):
             stream=kwargs.get("stream") or self.llm_stream or configs.llm.stream,
             timeout=kwargs.get("timeout") or self.llm_timeout or configs.llm.timeout,
             max_retries=kwargs.get("max_retries") or self.llm_max_retries or configs.llm.max_retries,
-            api_key=self.llm_api_key.get_secret_value() if self.llm_api_key else configs.llm.api_key.get_secret_value(),
-            base_url=self.llm_api_endpoint.unicode_string()
-            if self.llm_api_endpoint
-            else configs.llm.api_endpoint.unicode_string(),
+            api_key=(self.llm_api_key or configs.llm.api_key).get_secret_value(),
+            base_url=(self.llm_api_endpoint or configs.llm.api_endpoint).unicode_string(),
         )
 
     async def ainvoke(
@@ -121,13 +91,13 @@ class LLMUsage(Base):
         if isinstance(resp, ModelResponse):
             return resp.choices
         if isinstance(resp, CustomStreamWrapper):
-            if configs.debug.streaming_visible:
-                chunks = []
-                async for chunk in resp:
-                    chunks.append(chunk)
-                    print(chunk.choices[0].delta.content or "", end="")  # noqa: T201
-                return stream_chunk_builder(chunks).choices
-            return stream_chunk_builder(await asyncstdlib.list()).choices
+            if not configs.debug.streaming_visible:
+                return stream_chunk_builder(await asyncstdlib.list()).choices
+            chunks = []
+            async for chunk in resp:
+                chunks.append(chunk)
+                print(chunk.choices[0].delta.content or "", end="")  # noqa: T201
+            return stream_chunk_builder(chunks).choices
         logger.critical(err := f"Unexpected response type: {type(resp)}")
         raise ValueError(err)
 
@@ -383,39 +353,76 @@ class LLMUsage(Base):
             **kwargs,
         )
 
-    def fallback_to(self, other: "LLMUsage") -> Self:
-        """Fallback to another instance's attribute values if the current instance's attributes are None.
+
+class EmbeddingUsage(LLMUsage):
+    """A class representing the embedding model."""
+
+    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:
-            other (LLMUsage): Another instance from which to copy attribute values.
+            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 output should have, which is used to validate the result. 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:
-            Self: The current instance, allowing for method chaining.
+            EmbeddingResponse: The response containing the embeddings.
         """
-        # Iterate over the attribute names and copy values from 'other' to 'self' where applicable
-        # noinspection PydanticTypeChecker,PyTypeChecker
-        for attr_name in LLMUsage.model_fields:
-            # Copy the attribute value from 'other' to 'self' only if 'self' has None and 'other' has a non-None value
-            if getattr(self, attr_name) is None and (attr := getattr(other, attr_name)) is not None:
-                setattr(self, attr_name, attr)
-
-        # Return the current instance to allow for method chaining
-        return self
+        return await litellm.aembedding(
+            input=input_text,
+            caching=caching or self.embedding_caching or configs.embedding.caching,
+            dimensions=dimensions or self.embedding_dimensions or configs.embedding.dimensions,
+            model=model or self.embedding_model or configs.embedding.model or self.llm_model or configs.llm.model,
+            timeout=timeout
+            or self.embedding_timeout
+            or configs.embedding.timeout
+            or self.llm_timeout
+            or configs.llm.timeout,
+            api_key=(
+                self.embedding_api_key or configs.embedding.api_key or self.llm_api_key or configs.llm.api_key
+            ).get_secret_value(),
+            api_base=(
+                self.embedding_api_endpoint
+                or configs.embedding.api_endpoint
+                or self.llm_api_endpoint
+                or configs.llm.api_endpoint
+            )
+            .unicode_string()
+            .rstrip("/"),
+            # seems embedding function takes no base_url end with a slash
+        )
 
-    def hold_to(self, others: Union["LLMUsage", Iterable["LLMUsage"]]) -> Self:
-        """Hold to another instance's attribute values if the current instance's attributes are None.
+    @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:
-            others (LLMUsage | Iterable[LLMUsage]): Another instance or iterable of instances from which to copy attribute values.
+            input_text (List[str] | str): A string or list of strings to generate embeddings for.
+            **kwargs (Unpack[EmbeddingKwargs]): Additional keyword arguments for embedding.
 
         Returns:
-            Self: The current instance, allowing for method chaining.
+            List[List[float]] | List[float]: The generated embeddings.
         """
-        for other in others:
-            # noinspection PyTypeChecker,PydanticTypeChecker
-            for attr_name in LLMUsage.model_fields:
-                if (attr := getattr(self, attr_name)) is not None and getattr(other, attr_name) is None:
-                    setattr(other, attr_name, attr)
+        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]
 
 
 class ToolBoxUsage(LLMUsage):

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