mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-07-06 03:37:49 +02:00
feat(add executor):
This commit is contained in:
Alexander Myasoedov
@@ -0,0 +1,12 @@
|
||||
"""Advanced concurrent execution package for security scanning."""
|
||||
|
||||
from agentic_security.executor.rate_limiter import TokenBucketRateLimiter
|
||||
from agentic_security.executor.circuit_breaker import CircuitBreaker
|
||||
from agentic_security.executor.concurrent import ConcurrentExecutor, ExecutorMetrics
|
||||
|
||||
__all__ = [
|
||||
"TokenBucketRateLimiter",
|
||||
"CircuitBreaker",
|
||||
"ConcurrentExecutor",
|
||||
"ExecutorMetrics",
|
||||
]
|
||||
@@ -0,0 +1,109 @@
|
||||
"""Circuit breaker pattern for fault tolerance."""
|
||||
|
||||
import time
|
||||
from typing import Literal
|
||||
|
||||
|
||||
CircuitState = Literal["closed", "open", "half_open"]
|
||||
|
||||
|
||||
class CircuitBreaker:
|
||||
"""Circuit breaker to prevent cascading failures.
|
||||
|
||||
Implements the circuit breaker pattern with three states:
|
||||
- closed: Normal operation, requests pass through
|
||||
- open: Failure threshold exceeded, requests fail fast
|
||||
- half_open: Recovery attempt, limited requests allowed
|
||||
|
||||
Example:
|
||||
>>> breaker = CircuitBreaker(failure_threshold=0.5, recovery_timeout=30)
|
||||
>>> if breaker.is_open():
|
||||
... raise Exception("Circuit breaker is open")
|
||||
>>> try:
|
||||
... result = make_request()
|
||||
... breaker.record_success()
|
||||
>>> except Exception:
|
||||
... breaker.record_failure()
|
||||
"""
|
||||
|
||||
def __init__(self, failure_threshold: float = 0.5, recovery_timeout: int = 30):
|
||||
"""Initialize circuit breaker.
|
||||
|
||||
Args:
|
||||
failure_threshold: Failure rate (0.0-1.0) that triggers open state
|
||||
recovery_timeout: Seconds to wait before attempting recovery
|
||||
"""
|
||||
self.failure_threshold = failure_threshold
|
||||
self.recovery_timeout = recovery_timeout
|
||||
self.failures = 0
|
||||
self.successes = 0
|
||||
self.state: CircuitState = "closed"
|
||||
self.last_failure_time: float | None = None
|
||||
|
||||
def record_success(self):
|
||||
"""Record a successful request."""
|
||||
self.successes += 1
|
||||
|
||||
# If in half_open state and we have enough successes, close the circuit
|
||||
if self.state == "half_open" and self.successes >= 3:
|
||||
self.state = "closed"
|
||||
self.failures = 0
|
||||
self.successes = 0
|
||||
|
||||
def record_failure(self):
|
||||
"""Record a failed request."""
|
||||
self.failures += 1
|
||||
self.last_failure_time = time.monotonic()
|
||||
|
||||
total = self.failures + self.successes
|
||||
|
||||
# Need minimum sample size before opening circuit
|
||||
if total >= 10:
|
||||
failure_rate = self.failures / total
|
||||
if failure_rate >= self.failure_threshold:
|
||||
self.state = "open"
|
||||
|
||||
def is_open(self) -> bool:
|
||||
"""Check if circuit breaker is open.
|
||||
|
||||
Returns:
|
||||
bool: True if circuit is open and requests should be blocked
|
||||
"""
|
||||
if self.state == "open":
|
||||
# Check if we should attempt recovery
|
||||
if self.last_failure_time is not None:
|
||||
if time.monotonic() - self.last_failure_time > self.recovery_timeout:
|
||||
self.state = "half_open"
|
||||
# Reset counters for half-open state
|
||||
self.failures = 0
|
||||
self.successes = 0
|
||||
return False
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def get_state(self) -> CircuitState:
|
||||
"""Get current circuit breaker state.
|
||||
|
||||
Returns:
|
||||
CircuitState: Current state (closed, open, or half_open)
|
||||
"""
|
||||
return self.state
|
||||
|
||||
def get_failure_rate(self) -> float:
|
||||
"""Get current failure rate.
|
||||
|
||||
Returns:
|
||||
float: Failure rate (0.0-1.0), or 0.0 if no requests recorded
|
||||
"""
|
||||
total = self.failures + self.successes
|
||||
if total == 0:
|
||||
return 0.0
|
||||
return self.failures / total
|
||||
|
||||
def reset(self):
|
||||
"""Reset circuit breaker to initial state."""
|
||||
self.failures = 0
|
||||
self.successes = 0
|
||||
self.state = "closed"
|
||||
self.last_failure_time = None
|
||||
@@ -0,0 +1,232 @@
|
||||
"""Concurrent executor with rate limiting and circuit breaking."""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from agentic_security.executor.rate_limiter import TokenBucketRateLimiter
|
||||
from agentic_security.executor.circuit_breaker import CircuitBreaker
|
||||
from agentic_security.logutils import logger
|
||||
from agentic_security.probe_actor.state import FuzzerState
|
||||
|
||||
|
||||
class ExecutorMetrics:
|
||||
"""Track executor performance metrics."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize metrics tracking."""
|
||||
self.successful_requests = 0
|
||||
self.failed_requests = 0
|
||||
self.total_latency = 0.0
|
||||
self.latencies: list[float] = []
|
||||
|
||||
def record_success(self, latency: float):
|
||||
"""Record a successful request.
|
||||
|
||||
Args:
|
||||
latency: Request latency in seconds
|
||||
"""
|
||||
self.successful_requests += 1
|
||||
self.total_latency += latency
|
||||
self.latencies.append(latency)
|
||||
|
||||
def record_failure(self):
|
||||
"""Record a failed request."""
|
||||
self.failed_requests += 1
|
||||
|
||||
def get_stats(self) -> dict[str, Any]:
|
||||
"""Get current statistics.
|
||||
|
||||
Returns:
|
||||
dict: Statistics including total requests, success rate, latency metrics
|
||||
"""
|
||||
total_requests = self.successful_requests + self.failed_requests
|
||||
|
||||
if total_requests == 0:
|
||||
return {
|
||||
"total_requests": 0,
|
||||
"success_rate": 0.0,
|
||||
"avg_latency_ms": 0.0,
|
||||
"p95_latency_ms": 0.0,
|
||||
}
|
||||
|
||||
success_rate = self.successful_requests / total_requests
|
||||
avg_latency_ms = (
|
||||
(self.total_latency / self.successful_requests * 1000)
|
||||
if self.successful_requests > 0
|
||||
else 0.0
|
||||
)
|
||||
|
||||
# Calculate p95 latency
|
||||
if self.latencies:
|
||||
sorted_latencies = sorted(self.latencies)
|
||||
p95_index = int(len(sorted_latencies) * 0.95)
|
||||
p95_latency_ms = sorted_latencies[p95_index] * 1000 if p95_index < len(sorted_latencies) else 0.0
|
||||
else:
|
||||
p95_latency_ms = 0.0
|
||||
|
||||
return {
|
||||
"total_requests": total_requests,
|
||||
"successful_requests": self.successful_requests,
|
||||
"failed_requests": self.failed_requests,
|
||||
"success_rate": success_rate,
|
||||
"avg_latency_ms": avg_latency_ms,
|
||||
"p95_latency_ms": p95_latency_ms,
|
||||
}
|
||||
|
||||
|
||||
class ConcurrentExecutor:
|
||||
"""Enhanced concurrent executor with rate limiting and circuit breaking.
|
||||
|
||||
Provides advanced concurrency control for security scanning with:
|
||||
- Token bucket rate limiting
|
||||
- Circuit breaker for fault tolerance
|
||||
- Metrics collection
|
||||
- Semaphore-based concurrency limits
|
||||
|
||||
Example:
|
||||
>>> executor = ConcurrentExecutor(max_concurrent=20, rate_limit=10, burst=5)
|
||||
>>> tokens, failures = await executor.execute_batch(
|
||||
... request_factory, prompts, "module_name", fuzzer_state
|
||||
... )
|
||||
>>> print(executor.metrics.get_stats())
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_concurrent: int = 50,
|
||||
rate_limit: float = 100,
|
||||
burst: int = 20,
|
||||
failure_threshold: float = 0.5,
|
||||
recovery_timeout: int = 30,
|
||||
):
|
||||
"""Initialize concurrent executor.
|
||||
|
||||
Args:
|
||||
max_concurrent: Maximum number of concurrent requests
|
||||
rate_limit: Requests per second limit
|
||||
burst: Maximum burst size for rate limiter
|
||||
failure_threshold: Failure rate that triggers circuit breaker
|
||||
recovery_timeout: Seconds before attempting circuit recovery
|
||||
"""
|
||||
self.semaphore = asyncio.Semaphore(max_concurrent)
|
||||
self.rate_limiter = TokenBucketRateLimiter(rate_limit, burst)
|
||||
self.circuit_breaker = CircuitBreaker(failure_threshold, recovery_timeout)
|
||||
self.metrics = ExecutorMetrics()
|
||||
|
||||
logger.info(
|
||||
f"ConcurrentExecutor initialized: max_concurrent={max_concurrent}, "
|
||||
f"rate_limit={rate_limit}/s, burst={burst}"
|
||||
)
|
||||
|
||||
async def execute_batch(
|
||||
self,
|
||||
request_factory,
|
||||
prompts: list[str],
|
||||
module_name: str,
|
||||
fuzzer_state: FuzzerState,
|
||||
) -> tuple[int, int]:
|
||||
"""Execute a batch of prompts with rate limiting and circuit breaking.
|
||||
|
||||
This is compatible with the existing process_prompt_batch signature.
|
||||
|
||||
Args:
|
||||
request_factory: Request factory with fn() method
|
||||
prompts: List of prompts to process
|
||||
module_name: Name of the module being scanned
|
||||
fuzzer_state: State tracking object
|
||||
|
||||
Returns:
|
||||
tuple[int, int]: (total_tokens, failures)
|
||||
"""
|
||||
tasks = [
|
||||
self._execute_single(request_factory, prompt, module_name, fuzzer_state)
|
||||
for prompt in prompts
|
||||
]
|
||||
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
# Aggregate results
|
||||
total_tokens = 0
|
||||
failures = 0
|
||||
|
||||
for result in results:
|
||||
if isinstance(result, Exception):
|
||||
failures += 1
|
||||
logger.error(f"Task failed with exception: {result}")
|
||||
else:
|
||||
tokens, refused = result
|
||||
total_tokens += tokens
|
||||
if refused:
|
||||
failures += 1
|
||||
|
||||
return total_tokens, failures
|
||||
|
||||
async def _execute_single(
|
||||
self,
|
||||
request_factory,
|
||||
prompt: str,
|
||||
module_name: str,
|
||||
fuzzer_state: FuzzerState,
|
||||
) -> tuple[int, bool]:
|
||||
"""Execute a single prompt with rate limiting and circuit breaking.
|
||||
|
||||
Args:
|
||||
request_factory: Request factory with fn() method
|
||||
prompt: Prompt to process
|
||||
module_name: Name of the module being scanned
|
||||
fuzzer_state: State tracking object
|
||||
|
||||
Returns:
|
||||
tuple[int, bool]: (tokens, refused)
|
||||
|
||||
Raises:
|
||||
Exception: If circuit breaker is open
|
||||
"""
|
||||
# Rate limiting
|
||||
await self.rate_limiter.acquire()
|
||||
|
||||
# Circuit breaker check
|
||||
if self.circuit_breaker.is_open():
|
||||
self.metrics.record_failure()
|
||||
raise Exception("Circuit breaker is open - too many failures")
|
||||
|
||||
# Concurrency control
|
||||
async with self.semaphore:
|
||||
start_time = time.monotonic()
|
||||
|
||||
try:
|
||||
# Import here to avoid circular dependency
|
||||
from agentic_security.probe_actor.fuzzer import process_prompt
|
||||
|
||||
tokens = 0 # Initial token count for this prompt
|
||||
result = await process_prompt(
|
||||
request_factory, prompt, tokens, module_name, fuzzer_state
|
||||
)
|
||||
|
||||
# Record success
|
||||
self.circuit_breaker.record_success()
|
||||
latency = time.monotonic() - start_time
|
||||
self.metrics.record_success(latency)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
# Record failure
|
||||
self.circuit_breaker.record_failure()
|
||||
self.metrics.record_failure()
|
||||
logger.error(f"Error executing prompt: {e}")
|
||||
raise
|
||||
|
||||
def get_metrics(self) -> dict[str, Any]:
|
||||
"""Get current executor metrics.
|
||||
|
||||
Returns:
|
||||
dict: Metrics including request stats, latency, and circuit breaker state
|
||||
"""
|
||||
stats = self.metrics.get_stats()
|
||||
stats["circuit_breaker_state"] = self.circuit_breaker.get_state()
|
||||
stats["circuit_breaker_failure_rate"] = self.circuit_breaker.get_failure_rate()
|
||||
stats["available_tokens"] = self.rate_limiter.get_available_tokens()
|
||||
|
||||
return stats
|
||||
@@ -0,0 +1,63 @@
|
||||
"""Token bucket rate limiter for controlling request rate."""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
|
||||
class TokenBucketRateLimiter:
|
||||
"""Token bucket rate limiter with configurable rate and burst capacity.
|
||||
|
||||
This implements the token bucket algorithm where tokens are added at a fixed
|
||||
rate and consumed for each request. Supports bursting up to the bucket capacity.
|
||||
|
||||
Example:
|
||||
>>> limiter = TokenBucketRateLimiter(rate=10, burst=20)
|
||||
>>> await limiter.acquire() # Will wait if no tokens available
|
||||
"""
|
||||
|
||||
def __init__(self, rate: float, burst: int):
|
||||
"""Initialize rate limiter.
|
||||
|
||||
Args:
|
||||
rate: Tokens added per second (requests/sec)
|
||||
burst: Maximum bucket capacity (max concurrent burst)
|
||||
"""
|
||||
self.rate = rate
|
||||
self.burst = burst
|
||||
self.tokens = float(burst)
|
||||
self.last_update = time.monotonic()
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def acquire(self):
|
||||
"""Acquire a token, waiting if necessary.
|
||||
|
||||
This method will block until a token is available.
|
||||
"""
|
||||
async with self._lock:
|
||||
now = time.monotonic()
|
||||
elapsed = now - self.last_update
|
||||
|
||||
# Add tokens based on elapsed time
|
||||
self.tokens = min(self.burst, self.tokens + elapsed * self.rate)
|
||||
self.last_update = now
|
||||
|
||||
if self.tokens >= 1:
|
||||
# Token available, consume it
|
||||
self.tokens -= 1
|
||||
return
|
||||
|
||||
# Need to wait for next token
|
||||
wait_time = (1 - self.tokens) / self.rate
|
||||
await asyncio.sleep(wait_time)
|
||||
self.tokens = 0
|
||||
self.last_update = time.monotonic()
|
||||
|
||||
def get_available_tokens(self) -> float:
|
||||
"""Get current number of available tokens (non-blocking).
|
||||
|
||||
Returns:
|
||||
float: Number of tokens currently available
|
||||
"""
|
||||
now = time.monotonic()
|
||||
elapsed = now - self.last_update
|
||||
return min(self.burst, self.tokens + elapsed * self.rate)
|
||||
@@ -475,3 +475,47 @@ def prepare_prompts(
|
||||
datasets.append(load_csv(name))
|
||||
|
||||
return datasets
|
||||
|
||||
|
||||
async def prepare_prompts_unified(configs: list) -> list[ProbeDataset]:
|
||||
"""Prepare datasets using unified loader configuration.
|
||||
|
||||
This is an alternative to prepare_prompts() that uses the UnifiedDatasetLoader
|
||||
for streamlined configuration and merging of multiple sources.
|
||||
|
||||
Args:
|
||||
configs: List of InputSourceConfig objects or dicts
|
||||
|
||||
Returns:
|
||||
list[ProbeDataset]: List containing the merged dataset
|
||||
|
||||
Example:
|
||||
>>> from agentic_security.probe_data.unified_loader import InputSourceConfig
|
||||
>>> configs = [
|
||||
... InputSourceConfig(
|
||||
... source_type="huggingface",
|
||||
... dataset_name="deepset/prompt-injections",
|
||||
... enabled=True,
|
||||
... weight=1.0
|
||||
... )
|
||||
... ]
|
||||
>>> datasets = await prepare_prompts_unified(configs)
|
||||
"""
|
||||
from agentic_security.probe_data.unified_loader import (
|
||||
UnifiedDatasetLoader,
|
||||
InputSourceConfig,
|
||||
)
|
||||
|
||||
# Convert dicts to InputSourceConfig if needed
|
||||
config_objects = []
|
||||
for config in configs:
|
||||
if isinstance(config, dict):
|
||||
config_objects.append(InputSourceConfig(**config))
|
||||
else:
|
||||
config_objects.append(config)
|
||||
|
||||
loader = UnifiedDatasetLoader(config_objects)
|
||||
merged_dataset = await loader.load_all()
|
||||
|
||||
# Return as list for compatibility with existing code
|
||||
return [merged_dataset] if merged_dataset.prompts else []
|
||||
|
||||
@@ -0,0 +1,250 @@
|
||||
"""Unified dataset loader for CSV, HuggingFace, and proxy sources."""
|
||||
|
||||
from typing import Any, Literal, Optional
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from agentic_security.logutils import logger
|
||||
from agentic_security.probe_data.data import (
|
||||
load_dataset_generic,
|
||||
load_csv,
|
||||
create_probe_dataset,
|
||||
)
|
||||
from agentic_security.probe_data.models import ProbeDataset
|
||||
|
||||
|
||||
class InputSourceConfig(BaseModel):
|
||||
"""Configuration for a single input source."""
|
||||
|
||||
source_type: Literal["csv", "huggingface", "proxy"] = Field(
|
||||
description="Type of input source"
|
||||
)
|
||||
enabled: bool = Field(default=True, description="Whether this source is enabled")
|
||||
dataset_name: str = Field(description="Name/identifier of the dataset")
|
||||
weight: float = Field(default=1.0, ge=0.0, description="Sampling weight for merging")
|
||||
|
||||
# CSV-specific fields
|
||||
path: Optional[str] = Field(
|
||||
default=None, description="File path for CSV sources"
|
||||
)
|
||||
prompt_column: Optional[str] = Field(
|
||||
default="prompt", description="Column name containing prompts"
|
||||
)
|
||||
|
||||
# HuggingFace-specific fields
|
||||
split: Optional[str] = Field(
|
||||
default="train", description="Dataset split to load (train/test/validation)"
|
||||
)
|
||||
max_samples: Optional[int] = Field(
|
||||
default=None, ge=1, description="Maximum number of samples to load"
|
||||
)
|
||||
|
||||
# URL for custom sources
|
||||
url: Optional[str] = Field(
|
||||
default=None, description="URL for remote CSV files"
|
||||
)
|
||||
|
||||
|
||||
class UnifiedDatasetLoader:
|
||||
"""Loads and merges datasets from multiple sources."""
|
||||
|
||||
def __init__(self, configs: list[InputSourceConfig]):
|
||||
"""Initialize with list of input source configurations.
|
||||
|
||||
Args:
|
||||
configs: List of InputSourceConfig objects defining data sources
|
||||
"""
|
||||
self.configs = configs
|
||||
logger.info(f"Initialized UnifiedDatasetLoader with {len(configs)} sources")
|
||||
|
||||
async def load_all(self) -> ProbeDataset:
|
||||
"""Load all enabled sources and merge into a single dataset.
|
||||
|
||||
Returns:
|
||||
ProbeDataset: Merged dataset from all enabled sources
|
||||
"""
|
||||
datasets = []
|
||||
|
||||
for config in self.configs:
|
||||
if not config.enabled:
|
||||
logger.debug(f"Skipping disabled source: {config.dataset_name}")
|
||||
continue
|
||||
|
||||
try:
|
||||
dataset = await self._load_single(config)
|
||||
if dataset and dataset.prompts:
|
||||
datasets.append((dataset, config.weight))
|
||||
logger.info(
|
||||
f"Loaded {len(dataset.prompts)} prompts from {config.dataset_name} "
|
||||
f"(weight={config.weight})"
|
||||
)
|
||||
else:
|
||||
logger.warning(f"No prompts loaded from {config.dataset_name}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading {config.dataset_name}: {e}")
|
||||
|
||||
if not datasets:
|
||||
logger.warning("No datasets loaded successfully")
|
||||
return create_probe_dataset("unified_empty", [], {"sources": []})
|
||||
|
||||
return self._merge_weighted(datasets)
|
||||
|
||||
async def _load_single(self, config: InputSourceConfig) -> ProbeDataset:
|
||||
"""Load a single dataset based on its configuration.
|
||||
|
||||
Args:
|
||||
config: Configuration for the source to load
|
||||
|
||||
Returns:
|
||||
ProbeDataset: Loaded dataset
|
||||
"""
|
||||
if config.source_type == "csv":
|
||||
return self._load_csv_source(config)
|
||||
elif config.source_type == "huggingface":
|
||||
return self._load_huggingface_source(config)
|
||||
elif config.source_type == "proxy":
|
||||
return self._load_proxy_source(config)
|
||||
else:
|
||||
raise ValueError(f"Unknown source type: {config.source_type}")
|
||||
|
||||
def _load_csv_source(self, config: InputSourceConfig) -> ProbeDataset:
|
||||
"""Load dataset from CSV file.
|
||||
|
||||
Args:
|
||||
config: CSV source configuration
|
||||
|
||||
Returns:
|
||||
ProbeDataset: Dataset loaded from CSV
|
||||
"""
|
||||
if config.path:
|
||||
# Local CSV file
|
||||
logger.info(f"Loading CSV from path: {config.path}")
|
||||
dataset = load_csv(config.path)
|
||||
elif config.url:
|
||||
# Remote CSV file
|
||||
logger.info(f"Loading CSV from URL: {config.url}")
|
||||
mappings = {config.prompt_column: "prompt"} if config.prompt_column else None
|
||||
dataset = load_dataset_generic(
|
||||
name=config.dataset_name,
|
||||
url=config.url,
|
||||
mappings=mappings,
|
||||
metadata={"source_type": "csv", "url": config.url}
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"CSV source {config.dataset_name} requires either path or url")
|
||||
|
||||
# Apply max_samples limit if specified
|
||||
if config.max_samples and len(dataset.prompts) > config.max_samples:
|
||||
logger.info(
|
||||
f"Limiting {config.dataset_name} from {len(dataset.prompts)} "
|
||||
f"to {config.max_samples} samples"
|
||||
)
|
||||
dataset.prompts = dataset.prompts[:config.max_samples]
|
||||
|
||||
return dataset
|
||||
|
||||
def _load_huggingface_source(self, config: InputSourceConfig) -> ProbeDataset:
|
||||
"""Load dataset from HuggingFace.
|
||||
|
||||
Args:
|
||||
config: HuggingFace source configuration
|
||||
|
||||
Returns:
|
||||
ProbeDataset: Dataset loaded from HuggingFace
|
||||
"""
|
||||
logger.info(
|
||||
f"Loading HuggingFace dataset: {config.dataset_name} "
|
||||
f"(split={config.split})"
|
||||
)
|
||||
|
||||
# Build column mappings
|
||||
mappings = None
|
||||
if config.prompt_column and config.prompt_column != "prompt":
|
||||
mappings = {config.prompt_column: "prompt"}
|
||||
|
||||
dataset = load_dataset_generic(
|
||||
name=config.dataset_name,
|
||||
mappings=mappings,
|
||||
metadata={
|
||||
"source_type": "huggingface",
|
||||
"split": config.split,
|
||||
}
|
||||
)
|
||||
|
||||
# Apply max_samples limit if specified
|
||||
if config.max_samples and len(dataset.prompts) > config.max_samples:
|
||||
logger.info(
|
||||
f"Limiting {config.dataset_name} from {len(dataset.prompts)} "
|
||||
f"to {config.max_samples} samples"
|
||||
)
|
||||
dataset.prompts = dataset.prompts[:config.max_samples]
|
||||
|
||||
return dataset
|
||||
|
||||
def _load_proxy_source(self, config: InputSourceConfig) -> ProbeDataset:
|
||||
"""Load dataset from proxy queue (placeholder for PoC).
|
||||
|
||||
Args:
|
||||
config: Proxy source configuration
|
||||
|
||||
Returns:
|
||||
ProbeDataset: Empty dataset (proxy integration not implemented in PoC)
|
||||
"""
|
||||
logger.warning(
|
||||
f"Proxy source {config.dataset_name} not implemented in PoC - returning empty dataset"
|
||||
)
|
||||
return create_probe_dataset(
|
||||
config.dataset_name,
|
||||
[],
|
||||
{"source_type": "proxy", "status": "not_implemented"}
|
||||
)
|
||||
|
||||
def _merge_weighted(
|
||||
self, datasets: list[tuple[ProbeDataset, float]]
|
||||
) -> ProbeDataset:
|
||||
"""Merge multiple datasets with weighted sampling.
|
||||
|
||||
For PoC, this implements simple concatenation with optional weighting.
|
||||
Production version would implement proper stratified sampling.
|
||||
|
||||
Args:
|
||||
datasets: List of (ProbeDataset, weight) tuples
|
||||
|
||||
Returns:
|
||||
ProbeDataset: Merged dataset
|
||||
"""
|
||||
if not datasets:
|
||||
return create_probe_dataset("unified_empty", [], {"sources": []})
|
||||
|
||||
# For PoC: simple concatenation, repeat prompts based on weight
|
||||
all_prompts = []
|
||||
source_names = []
|
||||
total_tokens = 0
|
||||
|
||||
for dataset, weight in datasets:
|
||||
source_names.append(dataset.dataset_name)
|
||||
|
||||
# Calculate how many times to include this dataset based on weight
|
||||
# Weight of 1.0 = include once, 2.0 = include twice, etc.
|
||||
repeat_count = max(1, int(weight))
|
||||
|
||||
for _ in range(repeat_count):
|
||||
all_prompts.extend(dataset.prompts)
|
||||
|
||||
total_tokens += dataset.tokens * repeat_count
|
||||
|
||||
logger.info(
|
||||
f"Merged {len(datasets)} datasets into {len(all_prompts)} total prompts "
|
||||
f"from sources: {source_names}"
|
||||
)
|
||||
|
||||
return ProbeDataset(
|
||||
dataset_name="unified",
|
||||
metadata={
|
||||
"sources": source_names,
|
||||
"source_count": len(datasets),
|
||||
"weights": {ds.dataset_name: w for ds, w in datasets},
|
||||
},
|
||||
prompts=all_prompts,
|
||||
tokens=total_tokens,
|
||||
approx_cost=0.0,
|
||||
)
|
||||
Reference in New Issue
Block a user