Spaces:
Build error
Build error
File size: 2,572 Bytes
5e56bcf | 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 | from upif import guard
class UpifOpenAI:
"""
Wrapper for OpenAI Client to automatically sanitize inputs and outputs.
"""
def __init__(self, client):
self.client = client
self.chat = self.Chat(client.chat)
class Chat:
def __init__(self, original_chat):
self.completions = self.Completions(original_chat.completions)
class Completions:
def __init__(self, original_completions):
self.create = self._create_wrapper(original_completions.create)
def _create_wrapper(self, original_create):
def wrapper(*args, **kwargs):
# 1. Scan Input (Messages)
messages = kwargs.get("messages", [])
for msg in messages:
if "content" in msg and isinstance(msg["content"], str):
processed = guard.process_input(msg["content"])
# Check failure (refusal)
if processed == guard.input_guard.refusal_message:
# Return a mocked response object that refuses
return self._mock_refusal(processed)
msg["content"] = processed
# 2. Call Original
response = original_create(*args, **kwargs)
# 3. Scan Output (Response Content)
# Handle object access (obj.choices[0].message.content)
try:
content = response.choices[0].message.content
if content:
safe_content = guard.process_output(content)
response.choices[0].message.content = safe_content
except Exception:
pass # Streaming or different structure
return response
return wrapper
def _mock_refusal(self, message):
"""Creates a fake OpenAI response object with the refusal message."""
from types import SimpleNamespace
return SimpleNamespace(
choices=[
SimpleNamespace(
message=SimpleNamespace(
content=message,
role="assistant"
)
)
]
)
|