Spaces:
Running
Running
| #!/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) | |