Spaces:
Running
Running
File size: 4,477 Bytes
c2ea5ed 26d8ef2 5aa050e 26d8ef2 0b23e94 c2ea5ed 26d8ef2 c2ea5ed 290f82e c2ea5ed 5aa050e c2ea5ed 0b23e94 c2ea5ed 7bc750c c2ea5ed 7bc750c c2ea5ed 7bc750c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 |
#!/usr/bin/env python3
"""
Knowledge Graph Perturbation Tester using LiteLLM
This script loads a knowledge graph, tests relations for perturbation resistance,
and enriches the graph with test results using different LLM providers via LiteLLM.
All data is stored in the database, not in local files.
"""
# LiteLLM Monkey Patch to fix "Unsupported parameter" errors
from utils.fix_litellm_stop_param import patched_completion
import json
import argparse
import os
import time
import random
import logging
from typing import Dict, List, Any, Tuple, Optional, Union, Callable
import yaml
import sys
from datetime import datetime
import litellm
from utils.config import OPENAI_API_KEY
import uuid
# Configure logging for this module
logger = logging.getLogger(__name__)
# Import LiteLLM
from litellm import completion
# API Key setup - only set if value exists
# Only set environment variable if OPENAI_API_KEY is not None
if OPENAI_API_KEY:
os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
DEFAULT_MODEL = "gpt-5-mini"
from utils.config import LANGFUSE_PUBLIC_KEY, LANGFUSE_SECRET_KEY, LANGFUSE_AUTH, LANGFUSE_HOST
# Only set environment variables if they have values (not None)
if LANGFUSE_PUBLIC_KEY:
os.environ["LANGFUSE_PUBLIC_KEY"] = LANGFUSE_PUBLIC_KEY
if LANGFUSE_SECRET_KEY:
os.environ["LANGFUSE_SECRET_KEY"] = LANGFUSE_SECRET_KEY
if LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY:
os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = f"{LANGFUSE_HOST}/api/public/otel"
os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = f"Authorization=Basic {LANGFUSE_AUTH}"
# Commented out to reduce verbose telemetry output
# import openlit
# openlit.init()
# Only set environment variable if OPENAI_API_KEY is not None
if OPENAI_API_KEY:
os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
# (future) from .perturbation_types.rule_misunderstanding import RuleMisunderstandingPerturbationTester
# (future) from .perturbation_types.emotional_manipulation import EmotionalManipulationPerturbationTester
def load_litellm_config(model: str = "gpt-5-mini", api_key: str = None):
"""
Load LiteLLM config to route models to the correct provider.
This function is now simplified to only use environment variables or a passed key.
"""
try:
import litellm
litellm.set_verbose = False
if api_key:
litellm.api_key = api_key
logger.info("LiteLLM config loaded successfully.")
except ImportError:
logger.warning("LiteLLM not installed. LLM-based testing will not be available.")
except Exception as e:
logger.error(f"Error loading LiteLLM config: {str(e)}")
def run_knowledge_graph_tests(
testing_data: Dict[str, Any],
perturbation_types: List[str],
model: str = "gpt-5-mini",
max_relations: int = None,
judge_model: str = "gpt-5-mini",
openai_api_key: str = None,
progress_callback: Optional[Callable[[int, int, str], None]] = None,
**kwargs
) -> Dict[str, Any]:
"""
Run a suite of perturbation tests on a knowledge graph.
This is the main entry point for running tests.
"""
results = {}
# Load LiteLLM config
load_litellm_config(model, api_key=openai_api_key)
if "jailbreak" in perturbation_types:
from .perturbation_types.jailbreak import run_jailbreak_tests
jailbreak_results = run_jailbreak_tests(
testing_data,
model=model,
max_relations=max_relations,
judge_model=judge_model,
openai_api_key=openai_api_key,
progress_callback=progress_callback,
**kwargs,
)
results["jailbreak"] = jailbreak_results
if "counterfactual_bias" in perturbation_types:
from .perturbation_types.counterfactual_bias import run_counterfactual_bias_tests
cb_results = run_counterfactual_bias_tests(
testing_data,
model=model,
max_relations=max_relations,
judge_model=judge_model,
openai_api_key=openai_api_key,
progress_callback=progress_callback,
**kwargs,
)
results["counterfactual_bias"] = cb_results
# (future) Add other perturbation types here
# if "rule_misunderstanding" in perturbation_types:
# rm_results = run_rule_misunderstanding_tests(testing_data, model, **kwargs)
# results['rule_misunderstanding'] = rm_results
return results
if __name__ == '__main__':
sys.exit(0)
|