support testing malicious prompts with no guidelines

This commit is contained in:
Adam Wilson
2025-06-28 12:18:35 -06:00
parent 036d36bf4f
commit cb1be6746f
18 changed files with 447 additions and 76 deletions
@@ -4,10 +4,13 @@ from src.text_generation.adapters.embedding_model import EmbeddingModel
from src.text_generation.adapters.text_generation_foundation_model import TextGenerationFoundationModel
from src.text_generation.entrypoints.http_api_controller import HttpApiController
from src.text_generation.entrypoints.server import RestApiServer
from src.text_generation.services.logging.file_logging_service import FileLoggingService
from src.text_generation.services.logging.json_web_traffic_logging_service import JSONWebTrafficLoggingService
from src.text_generation.services.nlp.semantic_similarity_service import SemanticSimilarityService
from src.text_generation.services.nlp.text_generation_response_service import TextGenerationResponseService
from src.text_generation.services.nlp.retrieval_augmented_generation_response_service import RetrievalAugmentedGenerationResponseService
from src.text_generation.services.guardrails.generated_text_guardrail_service import GeneratedTextGuardrailService
from src.text_generation.services.guidelines.rag_guidelines_service import RetrievalAugmentedGenerationGuidelinesService
from src.text_generation.services.utilities.response_processing_service import ResponseProcessingService
class DependencyInjectionContainer(containers.DeclarativeContainer):
@@ -16,8 +19,7 @@ class DependencyInjectionContainer(containers.DeclarativeContainer):
config = providers.Configuration(yaml_files=['config.yml'])
logging_service = providers.Singleton(
FileLoggingService,
filename='test.log'
JSONWebTrafficLoggingService
)
foundation_model = providers.Singleton(
@@ -28,9 +30,21 @@ class DependencyInjectionContainer(containers.DeclarativeContainer):
EmbeddingModel
)
rag_guidelines_service = providers.Factory(
RetrievalAugmentedGenerationGuidelinesService,
embedding_model=embedding_model
)
response_processing_service = providers.Factory(
ResponseProcessingService
)
rag_response_service = providers.Factory(
RetrievalAugmentedGenerationResponseService,
embedding_model=embedding_model
foundation_model=foundation_model,
embedding_model=embedding_model,
rag_guidelines_service=rag_guidelines_service,
response_processing_service=response_processing_service
)
# add / implement guidelines svc
@@ -41,9 +55,14 @@ class DependencyInjectionContainer(containers.DeclarativeContainer):
'another example'
]
semantic_similarity_service = providers.Factory(
SemanticSimilarityService,
embedding_model=embedding_model
)
generated_text_guardrail_service = providers.Factory(
GeneratedTextGuardrailService,
embedding_model=embedding_model,
semantic_similarity_service=semantic_similarity_service,
comparison_texts=comparison_texts
)
@@ -62,7 +81,6 @@ class DependencyInjectionContainer(containers.DeclarativeContainer):
rest_api_server = providers.Factory(
RestApiServer,
logging_service=logging_service,
listening_port=9999, # config.server.port,
api_controller=api_controller
)
+3
View File
@@ -0,0 +1,3 @@
class Average:
def from_list_of_floats(self, floats: list[float]) -> float:
return sum(floats) / len(floats)
@@ -1,30 +1,24 @@
import json
import traceback
from src.text_generation.services.nlp.text_generation_response_service import TextGenerationResponseService
from src.text_generation.services.nlp.retrieval_augmented_generation_response_service import RetrievalAugmentedGenerationResponseService
from src.text_generation.services.logging.file_logging_service import FileLoggingService
from src.text_generation.services.guardrails.generated_text_guardrail_service import GeneratedTextGuardrailService
from src.text_generation.services.logging.abstract_web_traffic_logging_service import AbstractWebTrafficLoggingService
from src.text_generation.services.nlp.abstract_language_model_response_service import AbstractLanguageModelResponseService
from src.text_generation.services.guardrails.abstract_generated_text_guardrail_service import AbstractGeneratedTextGuardrailService
class HttpApiController:
def __init__(
self,
logging_service: FileLoggingService,
text_generation_response_service: TextGenerationResponseService,
rag_response_service: RetrievalAugmentedGenerationResponseService,
generated_text_guardrail_service: GeneratedTextGuardrailService
logging_service: AbstractWebTrafficLoggingService,
text_generation_response_service: AbstractLanguageModelResponseService,
rag_response_service: AbstractLanguageModelResponseService,
generated_text_guardrail_service: AbstractGeneratedTextGuardrailService
):
self.logger = logging_service.logger
# TODO: temp debug
self.original_info = self.logger.info
self.logger.info = self.debug_info
self.logging_service = logging_service
self.text_generation_response_service = text_generation_response_service
self.rag_response_service = rag_response_service
self.generated_text_guardrail_service = generated_text_guardrail_service
self.routes = {}
self.register_routes()
@@ -39,23 +33,15 @@ class HttpApiController:
def register_routes(self):
"""Register all API routes"""
self.routes[('GET', '/')] = self.health_check
self.routes[('POST', '/api/conversations')] = self.handle_conversations
self.routes[('POST', '/api/rag_conversations')] = self.handle_conversations_with_rag
def __http_415_notsupported(self, env, start_response):
response_headers = [('Content-Type', 'application/json')]
start_response('415 Unsupported Media Type', response_headers)
return [json.dumps({'error': 'Unsupported Content-Type'}).encode('utf-8')]
self.routes[('POST', '/api/completions')] = self.handle_conversations
self.routes[('POST', '/api/completions/rag-guided')] = self.handle_conversations_with_rag
def format_response(self, data):
"""Format response data as JSON with 'response' key"""
response_data = {'response': data}
try:
response_body = json.dumps(response_data).encode('utf-8')
except:
# If serialization fails, convert data to string first
response_body = json.dumps({'response': str(data)}).encode('utf-8')
return response_body
@@ -66,7 +52,7 @@ class HttpApiController:
return [response_body]
def handle_conversations(self, env, start_response):
"""Handle POST requests to /api/conversations"""
"""POST /api/completions"""
try:
request_body_size = int(env.get('CONTENT_LENGTH', 0))
except ValueError:
@@ -89,11 +75,11 @@ class HttpApiController:
http_status_code = 200 # make enum
response_headers = [('Content-Type', 'application/json'), ('Content-Length', str(len(response_body)))]
start_response(f'{http_status_code} OK', response_headers)
self.logger.info('non-RAG response: request body %s | status: %s | response: %s', request_body, http_status_code, response_body)
self.logging_service.log_request_response(request=prompt, response=response_text)
return [response_body]
def handle_conversations_with_rag(self, env, start_response):
"""Handle POST requests to /api/rag_conversations with RAG functionality"""
"""POST /api/completions/rag-guided"""
try:
request_body_size = int(env.get('CONTENT_LENGTH', 0))
except ValueError:
@@ -116,7 +102,7 @@ class HttpApiController:
http_status_code = 200 # make enum
response_headers = [('Content-Type', 'application/json'), ('Content-Length', str(len(response_body)))]
start_response(f'{http_status_code} OK', response_headers)
self.logger.info('RAG response: request body %s | status: %s | response: %s', request_body, http_status_code, response_body)
self.logging_service.log_request_response(request=prompt, response=response_text)
return [response_body]
def _http_200_ok(self, env, start_response):
+2 -5
View File
@@ -1,18 +1,15 @@
from wsgiref.simple_server import make_server
from src.text_generation.entrypoints.http_api_controller import HttpApiController
from src.text_generation.services.logging.file_logging_service import FileLoggingService
class RestApiServer:
def __init__(
self,
listening_port: int,
logging_service: FileLoggingService,
listening_port: int,
api_controller: HttpApiController
):
self.listening_port = listening_port
self.logger = logging_service.logger
self.api_controller = api_controller
def listen(self):
@@ -21,4 +18,4 @@ class RestApiServer:
print(f'listening on port {self.listening_port}...')
wsgi_srv.serve_forever()
except Exception as e:
self.logger.debug(e)
print(e)
@@ -0,0 +1,12 @@
import abc
from typing import Any, Dict, List
class AbstractWebTrafficLoggingService(abc.ABC):
@abc.abstractmethod
def log_request_response(self, request: str, response: str):
raise NotImplementedError
@abc.abstractmethod
def get_logs(self) -> List[Dict[str, Any]]:
raise NotImplementedError
@@ -1,4 +1,6 @@
import logging
from src.text_generation.services.logging.abstract_logging_service import AbstractLoggingService
@@ -0,0 +1,51 @@
import calendar
import json
import os
import threading
import time
from datetime import datetime
from typing import Any, Dict, List
from src.text_generation.services.logging.abstract_web_traffic_logging_service import AbstractWebTrafficLoggingService
class JSONWebTrafficLoggingService(AbstractWebTrafficLoggingService):
def __init__(self):
self._lock = threading.Lock()
timestamp = calendar.timegm(time.gmtime())
self.log_file_path = f"http_logs_{timestamp}.json"
self._ensure_log_file_exists()
def _ensure_log_file_exists(self):
if not os.path.exists(self.log_file_path):
with open(self.log_file_path, 'w') as f:
json.dump([], f)
def _read_logs(self) -> List[Dict[str, Any]]:
try:
with open(self.log_file_path, 'r') as f:
return json.load(f)
except (json.JSONDecodeError, FileNotFoundError):
return []
def _write_logs(self, logs: List[Dict[str, Any]]):
with open(self.log_file_path, 'w') as f:
json.dump(logs, f, indent=2, ensure_ascii=False)
def log_request_response(
self,
request: str,
response: str):
with self._lock:
logs = self._read_logs()
log_entry = {
"request": request,
"response": response,
"timestamp": datetime.now().isoformat()
}
logs.append(log_entry)
self._write_logs(logs)
def get_logs(self) -> List[Dict[str, Any]]:
with self._lock:
return self._read_logs()