veadk-python 0.2.15__py3-none-any.whl → 0.2.17__py3-none-any.whl

veadk/evaluation/types.py

@@ -17,6 +17,25 @@ from dataclasses import dataclass
 
 @dataclass
 class EvalResultCaseData:
+    """Holds data for a single evaluation case result.
+
+    This dataclass stores input, outputs, score, and status for one test case.
+    Used in evaluation reporting and metrics export.
+
+    Attributes:
+        id (str): Unique ID of the case.
+        input (str): User input for the case.
+        actual_output (str): Actual agent response.
+        expected_output (str): Expected agent response.
+        score (str): Score as string from evaluation.
+        reason (str): Reason for the score.
+        status (str): Status like 'PASSED' or 'FAILURE'.
+        latency (str): Latency in milliseconds as string.
+
+    Note:
+        Score and latency are strings for compatibility with external systems.
+    """
+
     id: str
     input: str
     actual_output: str
@@ -29,5 +48,18 @@ class EvalResultCaseData:
 
 @dataclass
 class EvalResultMetadata:
+    """Stores metadata about the evaluation run.
+
+    This dataclass captures model information for the evaluation.
+    Used in reporting and tracing.
+
+    Attributes:
+        tested_model (str): Name of the model being tested.
+        judge_model (str): Name of the judge model used.
+
+    Note:
+        Simple structure for quick metadata access.
+    """
+
     tested_model: str
     judge_model: str

veadk/evaluation/utils/prometheus.py

@@ -22,6 +22,23 @@ from veadk.evaluation.types import EvalResultCaseData, EvalResultMetadata
 
 
 class PrometheusPushgatewayConfig:
+    """Configures connection to Prometheus Pushgateway for metrics export.
+
+    This class holds settings for pushing evaluation metrics to Prometheus.
+    It uses environment variables for default values.
+
+    Attributes:
+        url (str): URL of the Prometheus Pushgateway endpoint.
+            Defaults to OBSERVABILITY_PROMETHEUS_PUSHGATEWAY_URL environment variable.
+        username (str): Username for authentication.
+            Defaults to OBSERVABILITY_PROMETHEUS_USERNAME environment variable.
+        password (str): Password for authentication.
+            Defaults to OBSERVABILITY_PROMETHEUS_PASSWORD environment variable.
+
+    Note:
+        All fields are optional and use environment variables if not provided.
+    """
+
     url: str = Field(
         default_factory=lambda: getenv(
             "OBSERVABILITY_PROMETHEUS_PUSHGATEWAY_URL",
@@ -87,6 +104,26 @@ def post_pushgateway(
     registry: CollectorRegistry,
     grouping_key: dict[str, str] | None = None,
 ):
+    """Pushes metrics to Prometheus Pushgateway with authentication.
+
+    This function sends collected metrics to the specified Pushgateway URL.
+    It uses basic authentication and optional grouping keys.
+
+    Args:
+        pushgateway_url (str): URL of the Pushgateway endpoint.
+        username (str): Authentication username.
+        password (str): Authentication password.
+        job_name (str): Name of the job for metrics labeling.
+        registry (CollectorRegistry): Registry containing metrics to push.
+        grouping_key (dict[str, str] | None): Optional key-value pairs for grouping.
+
+    Raises:
+        Exception: If push operation fails due to network or auth issues.
+
+    Note:
+        Authentication handler is created internally using provided credentials.
+    """
+
     def auth_handler(url, method, timeout, headers, data):
         return basic_auth_handler(
             url, method, timeout, headers, data, username, password
@@ -114,6 +151,30 @@ def push_to_prometheus(
     username: str = "",
     password: str = "",
 ):
+    """Sets and pushes evaluation metrics to Prometheus.
+
+    This function updates gauge metrics with evaluation results and pushes them.
+    It handles counts, thresholds, and specific data labels.
+
+    Args:
+        test_name (str): Name of the test for grouping.
+        test_cases_total (int): Total number of test cases.
+        test_cases_failure (int): Number of failed test cases.
+        test_cases_pass (int): Number of passed test cases.
+        test_data_list (list[EvalResultCaseData]): List of case data for labeling.
+        eval_data (EvalResultMetadata): Metadata for evaluation.
+        case_threshold (float): Threshold value for cases. Defaults to 0.5.
+        diff_threshold (float): Diff threshold value. Defaults to 0.2.
+        url (str): Pushgateway URL. Defaults to empty.
+        username (str): Auth username. Defaults to empty.
+        password (str): Auth password. Defaults to empty.
+
+    Returns:
+        None: Metrics are set and pushed directly.
+
+    Raises:
+        ValueError: If required data is invalid.
+    """
     test_cases_total_metric.set(test_cases_total)
     test_cases_failure_metric.set(test_cases_failure)
     test_cases_pass_metric.set(test_cases_pass)

veadk/knowledgebase/backends/base_backend.py

@@ -18,8 +18,21 @@ from pydantic import BaseModel
 
 
 class BaseKnowledgebaseBackend(ABC, BaseModel):
+    """Base backend for knowledgebase.
+
+    Attributes:
+        index (str): Index or collection name of the vector storage.
+
+    Examples:
+        You can implement your own knowledgebase backend.
+
+        ```python
+        class CustomKnowledgebaseBackend(BaseKnowledgebaseBackend):
+            pass
+        ```
+    """
+
     index: str
-    """Index or collection name of the vector storage."""
 
     @abstractmethod
     def precheck_index_naming(self) -> None:

veadk/knowledgebase/backends/in_memory_backend.py

@@ -24,10 +24,19 @@ from veadk.knowledgebase.backends.utils import get_llama_index_splitter
 
 
 class InMemoryKnowledgeBackend(BaseKnowledgebaseBackend):
+    """A in-memory implementation backend for knowledgebase.
+
+    In-memory backend stores embedded text in a vector storage from Llama-index.
+
+    Attributes:
+        embedding_config (EmbeddingModelConfig):
+            Embedding config for text embedding and search.
+            Embedding config contains embedding model name and the corresponding dim.
+    """
+
     embedding_config: NormalEmbeddingModelConfig | EmbeddingModelConfig = Field(
         default_factory=EmbeddingModelConfig
     )
-    """Embedding model configs"""
 
     def model_post_init(self, __context: Any) -> None:
         self._embed_model = OpenAILikeEmbedding(

veadk/knowledgebase/backends/opensearch_backend.py

@@ -43,6 +43,32 @@ except ImportError:
 
 
 class OpensearchKnowledgeBackend(BaseKnowledgebaseBackend):
+    """Opensearch-based backend for knowledgebase.
+
+    Opensearch backend stores embedded text in a opensearch database by Llama-index.
+
+    Attributes:
+        opensearch_config (OpensearchConfig):
+            Opensearch database configurations.
+            Mainly contains opensearch host, port, username, password, etc.
+        embedding_config (EmbeddingModelConfig):
+            Embedding config for text embedding and search.
+            Embedding config contains embedding model name and the corresponding dim.
+
+    Examples:
+        Init a knowledgebase based on opensearch backend.
+
+        ```python
+        knowledgebase = Knowledgebase(backend="opensearch")
+        ```
+
+        With more configurations:
+
+        ```python
+        ...
+        ```
+    """
+
     opensearch_config: OpensearchConfig = Field(default_factory=OpensearchConfig)
     """Opensearch client configs"""
 

veadk/knowledgebase/backends/redis_backend.py

@@ -44,13 +44,40 @@ except ImportError:
 
 
 class RedisKnowledgeBackend(BaseKnowledgebaseBackend):
+    """Redis based backend for knowledgebase.
+
+    Redis backend stores embedded text in a redis database by Llama-index.
+
+    Attributes:
+        redis_config (RedisConfig):
+            Redis database configurations.
+            Mainly contains redis database host, port, etc.
+        embedding_config (EmbeddingModelConfig):
+            Embedding configurations for text embedding and search.
+            Embedding config contains embedding model name and the corresponding dim.
+
+    Notes:
+        Please ensure that your redis database supports Redisaearch stack.
+
+    Examples:
+        Init a knowledgebase based on redis backend.
+
+        ```python
+        knowledgebase = Knowledgebase(backend="redis")
+        ```
+
+        With more configurations:
+
+        ```python
+        ...
+        ```
+    """
+
     redis_config: RedisConfig = Field(default_factory=RedisConfig)
-    """Redis client configs"""
 
     embedding_config: EmbeddingModelConfig | NormalEmbeddingModelConfig = Field(
         default_factory=EmbeddingModelConfig
     )
-    """Embedding model configs"""
 
     def model_post_init(self, __context: Any) -> None:
         # We will use `from_url` to init Redis client once the

veadk/knowledgebase/backends/vikingdb_knowledge_backend.py

@@ -59,24 +59,62 @@ def get_files_in_directory(directory: str):
 
 
 class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
+    """Volcengine Viking DB knowledgebase backend.
+
+    Volcegnine Viking DB knowledgebase provides powerful knowledgebase storage and search.
+
+    Attributes:
+        volcengine_access_key (str | None):
+            Access key for Volcengine. Loaded automatically from the
+            `VOLCENGINE_ACCESS_KEY` environment variable if not provided.
+
+        volcengine_secret_key (str | None):
+            Secret key for Volcengine. Loaded automatically from the
+            `VOLCENGINE_SECRET_KEY` environment variable if not provided.
+
+        session_token (str):
+            Optional session token for temporary credentials. Defaults to an empty string.
+
+        volcengine_project (str):
+            VikingDB knowledgebase project name in the Volcengine console platform.
+            Defaults to `"default"`.
+
+        region (str): Region of the VikingDB knowledgebase. Defaults to `"cn-beijing"`.
+
+        tos_config (TOSConfig | NormalTOSConfig):
+            TOS configuration used for uploading files to TOS (Volcengine’s Object Storage).
+            Defaults to a new instance of `TOSConfig`.
+
+    Notes:
+        Please make sure that you have created a bucket in your TOS.
+
+    Examples:
+        Init a knowledgebase based on VikingDB knowledgebase backend.
+
+        ```python
+        knowledgebase = Knowledgebase(backend="redis")
+        ```
+
+        With more configurations:
+
+        ```python
+        ...
+        ```
+    """
+
     volcengine_access_key: str | None = Field(
         default_factory=lambda: os.getenv("VOLCENGINE_ACCESS_KEY")
     )
-
     volcengine_secret_key: str | None = Field(
         default_factory=lambda: os.getenv("VOLCENGINE_SECRET_KEY")
     )
-
     session_token: str = ""
 
     volcengine_project: str = "default"
-    """VikingDB knowledgebase project in Volcengine console platform. Default by `default`"""
 
     region: str = "cn-beijing"
-    """VikingDB knowledgebase region"""
 
     tos_config: TOSConfig | NormalTOSConfig = Field(default_factory=TOSConfig)
-    """TOS config, used to upload files to TOS"""
 
     def model_post_init(self, __context: Any) -> None:
         self.precheck_index_naming()

veadk/knowledgebase/knowledgebase.py

@@ -56,6 +56,110 @@ def _get_backend_cls(backend: str) -> type[BaseKnowledgebaseBackend]:
 
 
 class KnowledgeBase(BaseModel):
+    """A knowledge base for storing user-related information.
+
+    This class represents a knowledge base used to store and retrieve user-specific data.
+    It supports multiple backend options, including in-memory, OpenSearch, Redis, and Volcengine's
+    VikingDB. The knowledge base allows for efficient document retrieval based on similarity,
+    with the ability to configure backend-specific settings.
+
+    Attributes:
+        name (str): The name of the knowledge base. Default is "user_knowledgebase".
+        description (str): A description of the knowledge base. Default is "This knowledgebase stores some user-related information."
+        backend (Union[Literal["local", "opensearch", "viking", "redis"], BaseKnowledgebaseBackend]):
+            The type of backend to use for storing and querying the knowledge base. Supported options include:
+            - 'local' for in-memory storage (data is lost when the program exits).
+            - 'opensearch' for OpenSearch (requires OpenSearch cluster).
+            - 'viking' for Volcengine VikingDB (requires VikingDB service).
+            - 'redis' for Redis with vector search capability (requires Redis).
+            Default is 'local'.
+        backend_config (dict): Configuration dictionary for the selected backend.
+        top_k (int): The number of top similar documents to retrieve during a search. Default is 10.
+        app_name (str): The name of the application associated with the knowledge base. If index is not provided, this value will be set to `index`.
+        index (str): The name of the knowledge base index.
+
+    Notes:
+        Please ensure that you have set the embedding-related configurations in environment variables.
+
+    Examples:
+        ### Simple backend
+
+        Create a local knowledgebase:
+
+        ```python
+        from veadk import Agent, Runner
+        from veadk.knowledgebase.knowledgebase import KnowledgeBase
+        from veadk.memory.short_term_memory import ShortTermMemory
+
+        app_name = "veadk_playground_app"
+        user_id = "veadk_playground_user"
+        session_id = "veadk_playground_session"
+
+
+        knowledgebase = KnowledgeBase(backend="opensearch", app_name=app_name)
+        knowledgebase.add_from_files(files=[knowledgebase_file])
+
+        agent = Agent(knowledgebase=knowledgebase)
+
+        runner = Runner(
+            agent=agent,
+            short_term_memory=ShortTermMemory(),
+            app_name=app_name,
+            user_id=user_id,
+        )
+
+        response = await runner.run(
+            messages="Tell me the secret of green.", session_id=session_id
+        )
+        print(response)
+        ```
+
+        ### Initialize knowledgebase with metadata
+
+        ```python
+        from veadk.knowledgebase import KnowledgeBase
+
+        knowledgebase = KnowledgeBase(
+            name="user_data",
+            description="A knowledgebase contains user hobbies.",
+            index="my_app",
+        )
+        ```
+
+        ### Initialize knowledgebase with backend instance
+
+        ```python
+        import veadk.config  # noqa
+
+        from veadk.knowledgebase import KnowledgeBase
+        from veadk.knowledgebase.backends.in_memory_backend import InMemoryKnowledgeBackend
+
+        backend = InMemoryKnowledgeBackend(
+            index="my_app",
+            embedding_config=...,
+        )
+
+        knowledgebase = KnowledgeBase(
+            name="user_data",
+            description="A knowledgebase contains user hobbies.",
+            backend=backend,
+        )
+        ```
+
+        ### Initialize knowledgebase with backend config
+
+        ```python
+        from veadk.knowledgebase import KnowledgeBase
+
+        knowledgebase = KnowledgeBase(
+            name="user_data",
+            description="A knowledgebase contains user hobbies.",
+            backend="local",
+            backend_config={"index": "user_app"},
+        )
+        ```
+    """
+
     name: str = "user_knowledgebase"
 
     description: str = "This knowledgebase stores some user-related information."
@@ -63,23 +167,14 @@ class KnowledgeBase(BaseModel):
     backend: Union[
         Literal["local", "opensearch", "viking", "redis"], BaseKnowledgebaseBackend
     ] = "local"
-    """Knowledgebase backend type. Supported backends are:
-    - local: In-memory knowledgebase, data will be lost when the program exits.
-    - opensearch: OpenSearch knowledgebase, requires an OpenSearch cluster.
-    - viking: Volcengine VikingDB knowledgebase, requires VikingDB service.
-    - redis: Redis knowledgebase, requires Redis with vector search capability.
-    Default is `local`."""
 
     backend_config: dict = Field(default_factory=dict)
-    """Configuration for the backend"""
 
     top_k: int = 10
-    """Number of top similar documents to retrieve during search"""
 
     app_name: str = ""
 
     index: str = ""
-    """The name of the knowledgebase index. If not provided, it will be generated based on the `app_name`."""
 
     def model_post_init(self, __context: Any) -> None:
         if isinstance(self.backend, BaseKnowledgebaseBackend):
@@ -108,15 +203,81 @@ class KnowledgeBase(BaseModel):
         )
 
     def add_from_directory(self, directory: str, **kwargs) -> bool:
-        """Add knowledge from file path to knowledgebase"""
+        """Add knowledge from file path to knowledgebase.
+
+        Add the files in the directory to knowledgebase backend.
+
+        Args:
+            directory (str): The directory path that needs to store.
+
+        Returns:
+            bool: True if successfully store the knowledgebase, False otherwise.
+
+        Examples:
+            Store a directory to knowledgebase:
+
+            ```python
+            knowledgebase = Knowledgebase(backend="local")
+
+            if knowledgebase.add_from_directory("./knowledgebase"):
+                # add successfully
+                ...
+            else:
+                raise RuntimeError("Uploaded directory failed.")
+            ```
+        """
         return self._backend.add_from_directory(directory=directory, **kwargs)
 
     def add_from_files(self, files: list[str], **kwargs) -> bool:
-        """Add knowledge (e.g, documents, strings, ...) to knowledgebase"""
+        """Add knowledge files to knowledgebase.
+
+        Add a list of files to knowledgebase backend.
+
+        Args:
+            files (str): The list of files.
+
+        Returns:
+            bool: True if successfully store the knowledgebase, False otherwise.
+
+        Examples:
+            Store files to knowledgebase:
+
+            ```python
+            knowledgebase = Knowledgebase(backend="local")
+
+            if knowledgebase.add_from_files("./knowledgebase"):
+                # add successfully
+                ...
+            else:
+                raise RuntimeError("Uploaded files failed.")
+            ```
+        """
         return self._backend.add_from_files(files=files, **kwargs)
 
     def add_from_text(self, text: str | list[str], **kwargs) -> bool:
-        """Add knowledge from text to knowledgebase"""
+        """Add a piece of text or a list of text to knowledgebase.
+
+        The `text` can be a string or a list of string. The text will be embedded and stored by the corresponding backend.
+
+        Args:
+            text (str | list[str]): The text string or a list of text strings.
+
+        Returns:
+            bool: True if successfully store the knowledgebase, False otherwise.
+
+        Examples:
+            Store a string or a list of string to knowledgebase:
+
+            ```python
+            knowledgebase = Knowledgebase(backend="local")
+
+            if knowledgebase.add_from_text("./knowledgebase"):
+                # add successfully
+                ...
+            else:
+                raise RuntimeError("Uploaded text failed.")
+            ```
+        """
         return self._backend.add_from_text(text=text, **kwargs)
 
     def search(self, query: str, top_k: int = 0, **kwargs) -> list[KnowledgebaseEntry]: