bayesian-llm-guard 1.0.0__tar.gz

bayesian_llm_guard-1.0.0/LICENSE

@@ -0,0 +1,17 @@
+Apache License
+Version 2.0, January 2004
+http://www.apache.org/licenses/
+
+Copyright 2026 XAIDeep.com
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

bayesian_llm_guard-1.0.0/PKG-INFO

@@ -0,0 +1,35 @@
+Metadata-Version: 2.4
+Name: bayesian-llm-guard
+Version: 1.0.0
+Summary: Epistemic Uncertainty Estimation & Guardrails for Agentic LLMs and RAG Systems
+License: Apache-2.0
+License-File: LICENSE
+Author: AI Research Team
+Author-email: research@xaideep.com
+Requires-Python: >=3.10,<4.0
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python :: 3.15
+Requires-Dist: numpy (>=1.24.0,<2.0.0)
+Requires-Dist: pydantic (>=2.0.0,<3.0.0)
+Requires-Dist: pydantic-settings (>=2.0.0,<3.0.0)
+Requires-Dist: torch (>=2.0.0)
+Description-Content-Type: text/markdown
+
+# bayesian-llm-guard 🛡️
+
+**bayesian-llm-guard** is a lightweight, production-ready Python library that bridges Bayesian Epistemic Uncertainty Estimation with Agentic LLM guardrails to mitigate hallucinations in RAG systems.
+
+## Installation
+
+_(Coming soon to PyPI)_
+
+## License
+
+Apache 2.0
+

bayesian_llm_guard-1.0.0/README.md

@@ -0,0 +1,11 @@
+# bayesian-llm-guard 🛡️
+
+**bayesian-llm-guard** is a lightweight, production-ready Python library that bridges Bayesian Epistemic Uncertainty Estimation with Agentic LLM guardrails to mitigate hallucinations in RAG systems.
+
+## Installation
+
+_(Coming soon to PyPI)_
+
+## License
+
+Apache 2.0

bayesian_llm_guard-1.0.0/bayesian_llm_guard/__init__.py

@@ -0,0 +1,21 @@
+"""
+bayesian-llm-guard
+Epistemic Uncertainty Estimation & Guardrails for Agentic LLMs.
+"""
+
+from .core import (
+    UQGuardConfig,
+    UQEngine,
+    uq_guard,
+    UQGuardException,
+    AgentSelfCorrectionLoop
+)
+
+__version__ = "0.1.0"
+__all__ = [
+    "UQGuardConfig",
+    "UQEngine",
+    "uq_guard",
+    "UQGuardException",
+    "AgentSelfCorrectionLoop"
+]

bayesian_llm_guard-1.0.0/bayesian_llm_guard/core/__init__.py

@@ -0,0 +1,12 @@
+from .config import UQGuardConfig
+from .uq_engine import UQEngine
+from .guardrail import uq_guard, UQGuardException
+from .agent_loop import AgentSelfCorrectionLoop
+
+__all__ = [
+    "UQGuardConfig",
+    "UQEngine",
+    "uq_guard",
+    "UQGuardException",
+    "AgentSelfCorrectionLoop"
+]

bayesian_llm_guard-1.0.0/bayesian_llm_guard/core/agent_loop.py

@@ -0,0 +1,79 @@
+"""
+Agent Loop Module for bayesian-llm-guard.
+Implements an autonomous self-correction loop to retry and refine 
+queries when high uncertainty (hallucination) is detected.
+"""
+
+import logging
+from typing import Any, Callable, Dict
+
+from bayesian_llm_guard.core.guardrail import UQGuardException
+
+logger = logging.getLogger("bayesian_agent_loop")
+
+
+class AgentSelfCorrectionLoop:
+    """
+    Manages iterative query refinement and re-execution for RAG pipelines.
+    Automatically catches uncertainty exceptions and refines the user's prompt.
+    """
+
+    def __init__(
+        self, 
+        max_retries: int = 3, 
+        query_refiner: Callable[[str], str] = lambda q: q + " (provide more specific context)"
+    ) -> None:
+        """
+        Args:
+            max_retries (int): Maximum number of times to retry the pipeline.
+            query_refiner (Callable): A function that modifies the failed query to improve retrieval.
+        """
+        self.max_retries = max_retries
+        self.query_refiner = query_refiner
+
+    def execute_with_fallback(
+        self, 
+        rag_pipeline_func: Callable[..., Dict[str, Any]], 
+        initial_query: str, 
+        **kwargs: Any
+    ) -> Dict[str, Any]:
+        """
+        Executes the provided RAG function. If it triggers the UQ Guardrail, 
+        the query is refined and retried automatically.
+
+        Args:
+            rag_pipeline_func (Callable): The decorated RAG function to execute.
+            initial_query (str): The original user prompt.
+            
+        Returns:
+            Dict[str, Any]: The safe, validated result from the pipeline.
+            
+        Raises:
+            RuntimeError: If the maximum number of retries is exhausted.
+        """
+        current_query = initial_query
+        attempt = 0
+
+        while attempt < self.max_retries:
+            try:
+                logger.info(f"Agent Loop Attempt {attempt + 1}/{self.max_retries}. Query: '{current_query}'")
+                
+                # Attempt to execute the pipeline
+                return rag_pipeline_func(query=current_query, **kwargs)
+                
+            except UQGuardException as exception:
+                attempt += 1
+                logger.warning(f"Guardrail intercepted the response: {exception}")
+                
+                if attempt < self.max_retries:
+                    logger.info("Refining query and retrying...")
+                    current_query = self.query_refiner(current_query)
+
+        # If the loop finishes without returning, the model could not produce a safe answer
+        error_msg = (
+            f"Agent loop failed after {self.max_retries} attempts. "
+            "The model consistently exhibited high epistemic uncertainty. "
+            "Please check your knowledge base or refine the original prompt."
+        )
+        logger.error(error_msg)
+        raise RuntimeError(error_msg)

bayesian_llm_guard-1.0.0/bayesian_llm_guard/core/config.py

@@ -0,0 +1,38 @@
+"""
+Configuration Module for bayesian-llm-guard.
+Manages type-safe settings and environment variables.
+"""
+
+from pydantic import Field
+from pydantic_settings import BaseSettings
+
+
+class UQGuardConfig(BaseSettings):
+    """
+    Core configuration for the Bayesian Guardrail.
+    Values can be overridden using environment variables with the 'UQ_' prefix.
+    Example: export UQ_THRESHOLD=0.03
+    """
+    
+    threshold: float = Field(
+        default=0.05,
+        ge=0.0,
+        le=1.0,
+        description="Maximum allowed epistemic variance. Values above this trigger the guardrail."
+    )
+    
+    mc_samples: int = Field(
+        default=30,
+        ge=5,
+        le=100,
+        description="Number of Monte Carlo forward passes for uncertainty estimation."
+    )
+    
+    max_retries: int = Field(
+        default=3,
+        ge=0,
+        description="Maximum number of self-correction attempts when high uncertainty is detected."
+    )
+
+    class Config:
+        env_prefix = "UQ_"

bayesian_llm_guard-1.0.0/bayesian_llm_guard/core/guardrail.py

@@ -0,0 +1,81 @@
+"""
+Guardrail Module for bayesian-llm-guard.
+Provides the @uq_guard decorator to intercept LLM/RAG pipelines
+and evaluate epistemic uncertainty before returning results.
+"""
+
+import logging
+from functools import wraps
+from typing import Any, Callable, Dict, Optional
+
+from bayesian_llm_guard.core.uq_engine import UQEngine
+from bayesian_llm_guard.core.config import UQGuardConfig
+
+# Setup basic logging for the guardrail
+logger = logging.getLogger("bayesian_guardrail")
+logging.basicConfig(level=logging.INFO)
+
+
+class UQGuardException(Exception):
+    """
+    Raised when the model's epistemic uncertainty exceeds the defined safety threshold.
+    Indicates a potential hallucination.
+    """
+    pass
+
+
+def uq_guard(
+    uq_engine: UQEngine,
+    config: Optional[UQGuardConfig] = None,
+    fallback_action: Optional[Callable[..., Any]] = None,
+) -> Callable[..., Callable[..., Any]]:
+    """
+    A decorator that wraps LLM/RAG execution functions.
+    It intercepts the returned tensor features, calculates the variance using the UQEngine,
+    and blocks the response if the model is too uncertain.
+
+    Args:
+        uq_engine (UQEngine): The instantiated uncertainty quantification engine.
+        config (UQGuardConfig, optional): Configuration containing the safety threshold.
+        fallback_action (Callable, optional): A function to execute if the threshold is breached.
+    """
+    conf = config or UQGuardConfig()
+
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Dict[str, Any]:
+            # Execute the underlying LLM or RAG pipeline
+            result = func(*args, **kwargs)
+            
+            # The wrapped function must return a dictionary containing 'tensor_features'
+            if isinstance(result, dict) and "tensor_features" in result:
+                features = result["tensor_features"]
+                
+                # Calculate the maximum epistemic variance
+                max_variance = uq_engine.estimate_variance(features)
+
+                logger.info(
+                    f"Calculated Epistemic Variance: {max_variance:.4f} "
+                    f"(Threshold: {conf.threshold:.4f})"
+                )
+
+                # Check if the uncertainty violates the safety threshold
+                if max_variance > conf.threshold:
+                    logger.warning("Uncertainty threshold exceeded! High risk of hallucination.")
+                    
+                    if fallback_action:
+                        logger.info("Executing defined fallback action.")
+                        return fallback_action(*args, **kwargs, uncertainty=max_variance)
+                    
+                    # If no fallback is defined, raise an exception to stop execution
+                    raise UQGuardException(
+                        f"Safety block triggered: Uncertainty score {max_variance:.4f} "
+                        f"is higher than the allowed threshold of {conf.threshold:.4f}."
+                    )
+                
+                # Inject the variance metrics into the final result for downstream analytics
+                result["uq_metrics"] = {"epistemic_variance": max_variance}
+            
+            return result
+        return wrapper
+    return decorator

bayesian_llm_guard-1.0.0/bayesian_llm_guard/core/uq_engine.py

@@ -0,0 +1,46 @@
+"""
+Uncertainty Quantification (UQ) Engine.
+Implements Epistemic Uncertainty estimation using Monte Carlo Dropout.
+Optimized for high-throughput GPU inference using Tensor Batching.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class UQEngine:
+    """
+    Executes vectorized stochastic forward passes through a PyTorch model
+    to calculate epistemic variance.
+    """
+    
+    def __init__(self, model: nn.Module, mc_samples: int = 30) -> None:
+        self.model = model
+        self.mc_samples = mc_samples
+
+    def _enable_dropout(self) -> None:
+        """Forces dropout layers to remain active during evaluation mode."""
+        for module in self.model.modules():
+            if isinstance(module, nn.Dropout):
+                module.train()
+
+    def estimate_variance(self, input_tensor: torch.Tensor) -> float:
+        """
+        Calculates the maximum epistemic variance using vectorized Monte Carlo samples.
+        """
+        self._enable_dropout()
+        
+        with torch.no_grad():
+            # ENTERPRISE BATCHING: Duplicate the tensor along the batch dimension.
+            # If input is (1, 768), repeated_tensor becomes (30, 768).
+            # This allows the GPU to process all Monte Carlo samples in parallel.
+            repeat_dims = [self.mc_samples] + [1] * (input_tensor.dim() - 1)
+            repeated_tensor = input_tensor.repeat(*repeat_dims)
+            
+            # Single highly-parallel forward pass
+            predictions = self.model(repeated_tensor)
+            
+        # Calculate statistical variance across the batch dimension (dim=0)
+        variances = torch.var(predictions, dim=0, unbiased=True)
+        
+        return float(variances.max().item())

bayesian_llm_guard-1.0.0/pyproject.toml

@@ -0,0 +1,30 @@
+[build-system]
+requires = ["poetry-core>=1.0.0"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry]
+name = "bayesian-llm-guard"
+version = "1.0.0"
+description = "Epistemic Uncertainty Estimation & Guardrails for Agentic LLMs and RAG Systems"
+readme = "README.md"
+authors = ["AI Research Team <research@xaideep.com>"]
+license = "Apache-2.0"
+
+[tool.poetry.dependencies]
+python = "^3.10"
+torch = ">=2.0.0"
+numpy = "^1.24.0"
+pydantic = "^2.0.0"
+pydantic-settings = "^2.0.0"
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^7.3.0"
+black = "^23.3.0"
+isort = "^5.12.0"
+
+[tool.semantic_release]
+version_toml = ["pyproject.toml:tool.poetry.version"]
+branch = "main"
+build_command = "pip install poetry && poetry build"
+commit_message = "chore(release): v{version} [skip ci]"
+commit_author = "github-actions[bot] <actions@github.com>"