resk-llm 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- resk-llm-0.1.0/LICENCE.txt +21 -0
- resk-llm-0.1.0/PKG-INFO +151 -0
- resk-llm-0.1.0/README.md +116 -0
- resk-llm-0.1.0/resk-llm/__init__.py +3 -0
- resk-llm-0.1.0/resk-llm/openai_protector.py +82 -0
- resk-llm-0.1.0/resk-llm/resk_context_manager.py +40 -0
- resk-llm-0.1.0/resk-llm/resk_models.py +124 -0
- resk-llm-0.1.0/resk-llm/resk_openai_tokens.py +26 -0
- resk-llm-0.1.0/resk-llm/resk_protector.py +42 -0
- resk-llm-0.1.0/resk_llm.egg-info/PKG-INFO +151 -0
- resk-llm-0.1.0/resk_llm.egg-info/SOURCES.txt +15 -0
- resk-llm-0.1.0/resk_llm.egg-info/dependency_links.txt +1 -0
- resk-llm-0.1.0/resk_llm.egg-info/requires.txt +1 -0
- resk-llm-0.1.0/resk_llm.egg-info/top_level.txt +1 -0
- resk-llm-0.1.0/setup.cfg +4 -0
- resk-llm-0.1.0/setup.py +29 -0
- resk-llm-0.1.0/tests/test_openai_protector.py.py +65 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c)
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
resk-llm-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,151 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: resk-llm
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Resk-LLM is a robust Python library designed to enhance security and manage context when interacting with OpenAI's language models. It provides a protective layer for API calls, safeguarding against common vulnerabilities and ensuring optimal performance.
|
|
5
|
+
Home-page: https://github.com/Resk-Security/Resk-LLM
|
|
6
|
+
Author: Resk
|
|
7
|
+
Author-email: nielzac@proton.me
|
|
8
|
+
License: MIT License
|
|
9
|
+
|
|
10
|
+
Copyright (c)
|
|
11
|
+
|
|
12
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
13
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
14
|
+
in the Software without restriction, including without limitation the rights
|
|
15
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
16
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
17
|
+
furnished to do so, subject to the following conditions:
|
|
18
|
+
|
|
19
|
+
The above copyright notice and this permission notice shall be included in all
|
|
20
|
+
copies or substantial portions of the Software.
|
|
21
|
+
|
|
22
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
23
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
24
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
25
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
26
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
27
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
28
|
+
SOFTWARE.
|
|
29
|
+
Classifier: Programming Language :: Python :: 3
|
|
30
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
31
|
+
Classifier: Operating System :: OS Independent
|
|
32
|
+
Requires-Python: >=3.6
|
|
33
|
+
Description-Content-Type: text/markdown
|
|
34
|
+
License-File: LICENCE.txt
|
|
35
|
+
|
|
36
|
+
# LLM Protector for OpenAI
|
|
37
|
+
|
|
38
|
+
LLM Protector is a Python library designed to secure calls to the OpenAI API by protecting against context overflow attacks, cleaning inputs, and managing a sliding context window.
|
|
39
|
+
|
|
40
|
+
## Installation
|
|
41
|
+
|
|
42
|
+
```
|
|
43
|
+
pip install resk ##soon avalailble
|
|
44
|
+
```
|
|
45
|
+
Instead use :
|
|
46
|
+
```
|
|
47
|
+
git clone
|
|
48
|
+
```
|
|
49
|
+
## Usage
|
|
50
|
+
|
|
51
|
+
```python
|
|
52
|
+
from openai import OpenAI
|
|
53
|
+
from llm_protector import OpenAIProtector
|
|
54
|
+
|
|
55
|
+
# Initialize the OpenAI client
|
|
56
|
+
client = OpenAI(api_key="your-api-key")
|
|
57
|
+
|
|
58
|
+
# Create a protector instance
|
|
59
|
+
protector = OpenAIProtector(model="gpt-4o-mini", preserved_prompts=2, reserved_tokens=1000)
|
|
60
|
+
|
|
61
|
+
# Example usage with OpenAI API
|
|
62
|
+
messages = [
|
|
63
|
+
{"role": "system", "content": "You are a helpful assistant."},
|
|
64
|
+
{"role": "user", "content": "Hello, can you help me?"},
|
|
65
|
+
{"role": "assistant", "content": "Of course! How can I assist you today?"},
|
|
66
|
+
{"role": "user", "content": "I'd like to learn more about Python programming."}
|
|
67
|
+
]
|
|
68
|
+
|
|
69
|
+
response = protector.protect_openai_call(
|
|
70
|
+
client.chat.completions.create,
|
|
71
|
+
model="gpt-4o-mini",
|
|
72
|
+
messages=messages
|
|
73
|
+
)
|
|
74
|
+
|
|
75
|
+
print(response.choices[0].message.content)
|
|
76
|
+
```
|
|
77
|
+
|
|
78
|
+
## Features
|
|
79
|
+
|
|
80
|
+
### Model Selection
|
|
81
|
+
|
|
82
|
+
You can specify the model when initializing the protector:
|
|
83
|
+
|
|
84
|
+
```python
|
|
85
|
+
protector = OpenAIProtector(model="gpt-4o-mini")
|
|
86
|
+
```
|
|
87
|
+
|
|
88
|
+
To get the list of available models:
|
|
89
|
+
|
|
90
|
+
```python
|
|
91
|
+
print(OpenAIProtector.get_available_models())
|
|
92
|
+
```
|
|
93
|
+
|
|
94
|
+
To get information about a specific model:
|
|
95
|
+
|
|
96
|
+
```python
|
|
97
|
+
print(OpenAIProtector.get_model_info("gpt-4o-mini"))
|
|
98
|
+
```
|
|
99
|
+
|
|
100
|
+
### Special Tokens
|
|
101
|
+
|
|
102
|
+
The protector uses a predefined list of OpenAI special tokens for input sanitization. You can view and modify this list if necessary.
|
|
103
|
+
|
|
104
|
+
To get the current list of special tokens:
|
|
105
|
+
|
|
106
|
+
```python
|
|
107
|
+
print(OpenAIProtector.get_special_tokens())
|
|
108
|
+
```
|
|
109
|
+
|
|
110
|
+
To update the list of special tokens:
|
|
111
|
+
|
|
112
|
+
```python
|
|
113
|
+
new_tokens = {
|
|
114
|
+
"general": ["<|endoftext|>", "<|fim_prefix|>"],
|
|
115
|
+
"chat": ["<|im_start|>user", "<|im_start|>assistant"]
|
|
116
|
+
}
|
|
117
|
+
OpenAIProtector.update_special_tokens(new_tokens)
|
|
118
|
+
```
|
|
119
|
+
|
|
120
|
+
### Context Management
|
|
121
|
+
|
|
122
|
+
The protector automatically manages a sliding context for long conversations. You can specify the number of prompts to preserve and the number of reserved tokens during initialization:
|
|
123
|
+
|
|
124
|
+
```python
|
|
125
|
+
protector = OpenAIProtector(model="gpt-4o-mini", preserved_prompts=2, reserved_tokens=1000)
|
|
126
|
+
```
|
|
127
|
+
|
|
128
|
+
- `preserved_prompts`: Number of messages at the beginning of the conversation to always preserve (default 2, typically the system prompt and the first user message).
|
|
129
|
+
- `reserved_tokens`: Number of tokens reserved for preserved messages and the last message (default 1000).
|
|
130
|
+
|
|
131
|
+
The protector:
|
|
132
|
+
- Cleans messages by removing non-UTF-8 characters and special tokens.
|
|
133
|
+
- Automatically closes open HTML tags.
|
|
134
|
+
- Escapes special HTML characters.
|
|
135
|
+
- Truncates messages that are too long if necessary.
|
|
136
|
+
- Removes the oldest messages to respect the model's context limit.
|
|
137
|
+
|
|
138
|
+
## Security
|
|
139
|
+
|
|
140
|
+
The protector automatically sanitizes all inputs and manages the context when using the `protect_openai_call` method. Always use this method for your OpenAI API calls to benefit from all protections.
|
|
141
|
+
|
|
142
|
+
```
|
|
143
|
+
|
|
144
|
+
## License
|
|
145
|
+
|
|
146
|
+
This project is licensed under a Custom Open Source License. See the [LICENSE](LICENSE) file for full details. Key points:
|
|
147
|
+
|
|
148
|
+
- Free for commercial and non-commercial use
|
|
149
|
+
- Modifications must be documented and attributed
|
|
150
|
+
- Paid versions based on this code are subject to royalties
|
|
151
|
+
- The original source must be cited in all derivative works
|
resk-llm-0.1.0/README.md
ADDED
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
# LLM Protector for OpenAI
|
|
2
|
+
|
|
3
|
+
LLM Protector is a Python library designed to secure calls to the OpenAI API by protecting against context overflow attacks, cleaning inputs, and managing a sliding context window.
|
|
4
|
+
|
|
5
|
+
## Installation
|
|
6
|
+
|
|
7
|
+
```
|
|
8
|
+
pip install resk ##soon avalailble
|
|
9
|
+
```
|
|
10
|
+
Instead use :
|
|
11
|
+
```
|
|
12
|
+
git clone
|
|
13
|
+
```
|
|
14
|
+
## Usage
|
|
15
|
+
|
|
16
|
+
```python
|
|
17
|
+
from openai import OpenAI
|
|
18
|
+
from llm_protector import OpenAIProtector
|
|
19
|
+
|
|
20
|
+
# Initialize the OpenAI client
|
|
21
|
+
client = OpenAI(api_key="your-api-key")
|
|
22
|
+
|
|
23
|
+
# Create a protector instance
|
|
24
|
+
protector = OpenAIProtector(model="gpt-4o-mini", preserved_prompts=2, reserved_tokens=1000)
|
|
25
|
+
|
|
26
|
+
# Example usage with OpenAI API
|
|
27
|
+
messages = [
|
|
28
|
+
{"role": "system", "content": "You are a helpful assistant."},
|
|
29
|
+
{"role": "user", "content": "Hello, can you help me?"},
|
|
30
|
+
{"role": "assistant", "content": "Of course! How can I assist you today?"},
|
|
31
|
+
{"role": "user", "content": "I'd like to learn more about Python programming."}
|
|
32
|
+
]
|
|
33
|
+
|
|
34
|
+
response = protector.protect_openai_call(
|
|
35
|
+
client.chat.completions.create,
|
|
36
|
+
model="gpt-4o-mini",
|
|
37
|
+
messages=messages
|
|
38
|
+
)
|
|
39
|
+
|
|
40
|
+
print(response.choices[0].message.content)
|
|
41
|
+
```
|
|
42
|
+
|
|
43
|
+
## Features
|
|
44
|
+
|
|
45
|
+
### Model Selection
|
|
46
|
+
|
|
47
|
+
You can specify the model when initializing the protector:
|
|
48
|
+
|
|
49
|
+
```python
|
|
50
|
+
protector = OpenAIProtector(model="gpt-4o-mini")
|
|
51
|
+
```
|
|
52
|
+
|
|
53
|
+
To get the list of available models:
|
|
54
|
+
|
|
55
|
+
```python
|
|
56
|
+
print(OpenAIProtector.get_available_models())
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
To get information about a specific model:
|
|
60
|
+
|
|
61
|
+
```python
|
|
62
|
+
print(OpenAIProtector.get_model_info("gpt-4o-mini"))
|
|
63
|
+
```
|
|
64
|
+
|
|
65
|
+
### Special Tokens
|
|
66
|
+
|
|
67
|
+
The protector uses a predefined list of OpenAI special tokens for input sanitization. You can view and modify this list if necessary.
|
|
68
|
+
|
|
69
|
+
To get the current list of special tokens:
|
|
70
|
+
|
|
71
|
+
```python
|
|
72
|
+
print(OpenAIProtector.get_special_tokens())
|
|
73
|
+
```
|
|
74
|
+
|
|
75
|
+
To update the list of special tokens:
|
|
76
|
+
|
|
77
|
+
```python
|
|
78
|
+
new_tokens = {
|
|
79
|
+
"general": ["<|endoftext|>", "<|fim_prefix|>"],
|
|
80
|
+
"chat": ["<|im_start|>user", "<|im_start|>assistant"]
|
|
81
|
+
}
|
|
82
|
+
OpenAIProtector.update_special_tokens(new_tokens)
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
### Context Management
|
|
86
|
+
|
|
87
|
+
The protector automatically manages a sliding context for long conversations. You can specify the number of prompts to preserve and the number of reserved tokens during initialization:
|
|
88
|
+
|
|
89
|
+
```python
|
|
90
|
+
protector = OpenAIProtector(model="gpt-4o-mini", preserved_prompts=2, reserved_tokens=1000)
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
- `preserved_prompts`: Number of messages at the beginning of the conversation to always preserve (default 2, typically the system prompt and the first user message).
|
|
94
|
+
- `reserved_tokens`: Number of tokens reserved for preserved messages and the last message (default 1000).
|
|
95
|
+
|
|
96
|
+
The protector:
|
|
97
|
+
- Cleans messages by removing non-UTF-8 characters and special tokens.
|
|
98
|
+
- Automatically closes open HTML tags.
|
|
99
|
+
- Escapes special HTML characters.
|
|
100
|
+
- Truncates messages that are too long if necessary.
|
|
101
|
+
- Removes the oldest messages to respect the model's context limit.
|
|
102
|
+
|
|
103
|
+
## Security
|
|
104
|
+
|
|
105
|
+
The protector automatically sanitizes all inputs and manages the context when using the `protect_openai_call` method. Always use this method for your OpenAI API calls to benefit from all protections.
|
|
106
|
+
|
|
107
|
+
```
|
|
108
|
+
|
|
109
|
+
## License
|
|
110
|
+
|
|
111
|
+
This project is licensed under a Custom Open Source License. See the [LICENSE](LICENSE) file for full details. Key points:
|
|
112
|
+
|
|
113
|
+
- Free for commercial and non-commercial use
|
|
114
|
+
- Modifications must be documented and attributed
|
|
115
|
+
- Paid versions based on this code are subject to royalties
|
|
116
|
+
- The original source must be cited in all derivative works
|
|
@@ -0,0 +1,82 @@
|
|
|
1
|
+
import re
|
|
2
|
+
import html
|
|
3
|
+
from typing import List, Dict, Any, Callable
|
|
4
|
+
|
|
5
|
+
from resk.resk_models import RESK_MODELS
|
|
6
|
+
from resk.resk_openai_tokens import OPENAI_SPECIAL_TOKENS
|
|
7
|
+
|
|
8
|
+
class OpenAIProtector:
|
|
9
|
+
def __init__(self, model: str = "gpt-4o", max_context_length: int = None, reserved_tokens: int = 1000, preserved_prompts: int = 2):
|
|
10
|
+
self.model = model
|
|
11
|
+
self.model_info = RESK_MODELS.get(model, RESK_MODELS["gpt-4o"])
|
|
12
|
+
self.max_context_length = max_context_length or self.model_info["context_window"]
|
|
13
|
+
self.max_output_tokens = self.model_info["max_output_tokens"]
|
|
14
|
+
self.reserved_tokens = reserved_tokens
|
|
15
|
+
self.preserved_prompts = preserved_prompts
|
|
16
|
+
self.special_tokens = set(OPENAI_SPECIAL_TOKENS["general"] + OPENAI_SPECIAL_TOKENS["chat"])
|
|
17
|
+
|
|
18
|
+
def sanitize_input(self, text: str) -> str:
|
|
19
|
+
text = text.encode('utf-8', errors='ignore').decode('utf-8')
|
|
20
|
+
for token in self.special_tokens:
|
|
21
|
+
text = text.replace(token, "")
|
|
22
|
+
text = self._close_html_tags(text)
|
|
23
|
+
return html.escape(text)
|
|
24
|
+
|
|
25
|
+
def _close_html_tags(self, text: str) -> str:
|
|
26
|
+
opened_tags = []
|
|
27
|
+
for match in re.finditer(r'<(\w+)[^>]*>', text):
|
|
28
|
+
tag = match.group(1)
|
|
29
|
+
if tag.lower() not in ['br', 'hr', 'img']:
|
|
30
|
+
opened_tags.append(tag)
|
|
31
|
+
for tag in reversed(opened_tags):
|
|
32
|
+
text += f'</{tag}>'
|
|
33
|
+
return text
|
|
34
|
+
|
|
35
|
+
def _truncate_text(self, text: str) -> str:
|
|
36
|
+
max_chars = self.max_context_length * 4 # Estimation grossière : 1 token ≈ 4 caractères
|
|
37
|
+
return text[:max_chars] if len(text) > max_chars else text
|
|
38
|
+
|
|
39
|
+
def manage_sliding_context(self, messages: List[Dict[str, str]]) -> List[Dict[str, str]]:
|
|
40
|
+
total_tokens = sum(len(msg['content'].split()) for msg in messages)
|
|
41
|
+
if total_tokens <= self.max_context_length:
|
|
42
|
+
return messages
|
|
43
|
+
|
|
44
|
+
preserved_messages = messages[:self.preserved_prompts]
|
|
45
|
+
remaining_messages = messages[self.preserved_prompts:]
|
|
46
|
+
|
|
47
|
+
available_tokens = self.max_context_length - self.reserved_tokens
|
|
48
|
+
preserved_tokens = sum(len(msg['content'].split()) for msg in preserved_messages)
|
|
49
|
+
available_tokens -= preserved_tokens
|
|
50
|
+
|
|
51
|
+
while remaining_messages and sum(len(msg['content'].split()) for msg in remaining_messages) > available_tokens:
|
|
52
|
+
remaining_messages.pop(0)
|
|
53
|
+
|
|
54
|
+
return preserved_messages + remaining_messages
|
|
55
|
+
|
|
56
|
+
def protect_openai_call(self, api_function: Callable, *args: Any, **kwargs: Any) -> Any:
|
|
57
|
+
sanitized_args = [self.sanitize_input(arg) if isinstance(arg, str) else arg for arg in args]
|
|
58
|
+
sanitized_kwargs = {k: self.sanitize_input(v) if isinstance(v, str) else v for k, v in kwargs.items()}
|
|
59
|
+
|
|
60
|
+
if 'messages' in sanitized_kwargs:
|
|
61
|
+
sanitized_kwargs['messages'] = self.manage_sliding_context(sanitized_kwargs['messages'])
|
|
62
|
+
|
|
63
|
+
sanitized_kwargs['model'] = self.model_info.get('version', self.model)
|
|
64
|
+
|
|
65
|
+
return api_function(*sanitized_args, **sanitized_kwargs)
|
|
66
|
+
|
|
67
|
+
@classmethod
|
|
68
|
+
def get_available_models(cls) -> List[str]:
|
|
69
|
+
return list(RESK_MODELS.keys())
|
|
70
|
+
|
|
71
|
+
@classmethod
|
|
72
|
+
def get_model_info(cls, model: str) -> Dict[str, Any]:
|
|
73
|
+
return RESK_MODELS.get(model, None)
|
|
74
|
+
|
|
75
|
+
@classmethod
|
|
76
|
+
def get_special_tokens(cls) -> Dict[str, List[str]]:
|
|
77
|
+
return OPENAI_SPECIAL_TOKENS
|
|
78
|
+
|
|
79
|
+
@classmethod
|
|
80
|
+
def update_special_tokens(cls, new_tokens: Dict[str, List[str]]) -> None:
|
|
81
|
+
global OPENAI_SPECIAL_TOKENS
|
|
82
|
+
OPENAI_SPECIAL_TOKENS = new_tokens
|
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
class ContextManager:
|
|
2
|
+
def __init__(self, model_info, preserved_prompts=2, reserved_tokens=1000):
|
|
3
|
+
self.max_context_length = model_info["context_window"]
|
|
4
|
+
self.preserved_prompts = preserved_prompts
|
|
5
|
+
self.reserved_tokens = reserved_tokens
|
|
6
|
+
|
|
7
|
+
def clean_message(self, message):
|
|
8
|
+
return ' '.join(message.split()) # Supprime les espaces multiples, tabulations et sauts de ligne
|
|
9
|
+
|
|
10
|
+
def truncate_message(self, message, available_tokens):
|
|
11
|
+
words = message.split()
|
|
12
|
+
return ' '.join(words[:available_tokens])
|
|
13
|
+
|
|
14
|
+
def manage_sliding_context(self, messages):
|
|
15
|
+
cleaned_messages = [{'role': msg['role'], 'content': self.clean_message(msg['content'])} for msg in messages]
|
|
16
|
+
total_tokens = sum(len(msg['content'].split()) for msg in cleaned_messages)
|
|
17
|
+
|
|
18
|
+
if total_tokens <= self.max_context_length:
|
|
19
|
+
return cleaned_messages
|
|
20
|
+
|
|
21
|
+
preserved_messages = cleaned_messages[:self.preserved_prompts]
|
|
22
|
+
remaining_messages = cleaned_messages[self.preserved_prompts:]
|
|
23
|
+
|
|
24
|
+
available_tokens = self.max_context_length - self.reserved_tokens
|
|
25
|
+
preserved_tokens = sum(len(msg['content'].split()) for msg in preserved_messages)
|
|
26
|
+
available_tokens -= preserved_tokens
|
|
27
|
+
|
|
28
|
+
while remaining_messages and available_tokens > 0:
|
|
29
|
+
message = remaining_messages.pop(0)
|
|
30
|
+
message_tokens = len(message['content'].split())
|
|
31
|
+
|
|
32
|
+
if message_tokens > available_tokens:
|
|
33
|
+
truncated_content = self.truncate_message(message['content'], available_tokens)
|
|
34
|
+
preserved_messages.append({'role': message['role'], 'content': truncated_content})
|
|
35
|
+
break
|
|
36
|
+
else:
|
|
37
|
+
preserved_messages.append(message)
|
|
38
|
+
available_tokens -= message_tokens
|
|
39
|
+
|
|
40
|
+
return preserved_messages
|
|
@@ -0,0 +1,124 @@
|
|
|
1
|
+
|
|
2
|
+
RESK_MODELS = {
|
|
3
|
+
# GPT-4o models
|
|
4
|
+
"gpt-4o": {
|
|
5
|
+
"description": "Our high-intelligence flagship model for complex, multi-step tasks. GPT-4o is cheaper and faster than GPT-4 Turbo.",
|
|
6
|
+
"context_window": 128000,
|
|
7
|
+
"max_output_tokens": 4096,
|
|
8
|
+
"training_data": "Up to Oct 2023",
|
|
9
|
+
"version": "gpt-4o-2024-05-13"
|
|
10
|
+
},
|
|
11
|
+
"gpt-4o-2024-05-13": {
|
|
12
|
+
"description": "gpt-4o currently points to this version.",
|
|
13
|
+
"context_window": 128000,
|
|
14
|
+
"max_output_tokens": 4096,
|
|
15
|
+
"training_data": "Up to Oct 2023"
|
|
16
|
+
},
|
|
17
|
+
"gpt-4o-2024-08-06": {
|
|
18
|
+
"description": "Latest snapshot that supports Structured Outputs",
|
|
19
|
+
"context_window": 128000,
|
|
20
|
+
"max_output_tokens": 16384,
|
|
21
|
+
"training_data": "Up to Oct 2023"
|
|
22
|
+
},
|
|
23
|
+
"chatgpt-4o-latest": {
|
|
24
|
+
"description": "Dynamic model continuously updated to the current version of GPT-4o in ChatGPT. Intended for research and evaluation.",
|
|
25
|
+
"context_window": 128000,
|
|
26
|
+
"max_output_tokens": 16384,
|
|
27
|
+
"training_data": "Up to Oct 2023"
|
|
28
|
+
},
|
|
29
|
+
|
|
30
|
+
# GPT-4o mini models
|
|
31
|
+
"gpt-4o-mini": {
|
|
32
|
+
"description": "Our affordable and intelligent small model for fast, lightweight tasks. GPT-4o mini is cheaper and more capable than GPT-3.5 Turbo.",
|
|
33
|
+
"context_window": 128000,
|
|
34
|
+
"max_output_tokens": 16384,
|
|
35
|
+
"training_data": "Up to Oct 2023",
|
|
36
|
+
"version": "gpt-4o-mini-2024-07-18"
|
|
37
|
+
},
|
|
38
|
+
"gpt-4o-mini-2024-07-18": {
|
|
39
|
+
"description": "gpt-4o-mini currently points to this version.",
|
|
40
|
+
"context_window": 128000,
|
|
41
|
+
"max_output_tokens": 16384,
|
|
42
|
+
"training_data": "Up to Oct 2023"
|
|
43
|
+
},
|
|
44
|
+
|
|
45
|
+
# GPT-4 Turbo and GPT-4 models
|
|
46
|
+
"gpt-4-turbo": {
|
|
47
|
+
"description": "The latest GPT-4 Turbo model with vision capabilities. Vision requests can now use JSON mode and function calling.",
|
|
48
|
+
"context_window": 128000,
|
|
49
|
+
"max_output_tokens": 4096,
|
|
50
|
+
"training_data": "Up to Dec 2023",
|
|
51
|
+
"version": "gpt-4-turbo-2024-04-09"
|
|
52
|
+
},
|
|
53
|
+
"gpt-4-turbo-2024-04-09": {
|
|
54
|
+
"description": "GPT-4 Turbo with Vision model. Vision requests can now use JSON mode and function calling.",
|
|
55
|
+
"context_window": 128000,
|
|
56
|
+
"max_output_tokens": 4096,
|
|
57
|
+
"training_data": "Up to Dec 2023"
|
|
58
|
+
},
|
|
59
|
+
"gpt-4-turbo-preview": {
|
|
60
|
+
"description": "GPT-4 Turbo preview model.",
|
|
61
|
+
"context_window": 128000,
|
|
62
|
+
"max_output_tokens": 4096,
|
|
63
|
+
"training_data": "Up to Dec 2023",
|
|
64
|
+
"version": "gpt-4-0125-preview"
|
|
65
|
+
},
|
|
66
|
+
"gpt-4-0125-preview": {
|
|
67
|
+
"description": "GPT-4 Turbo preview model intended to reduce cases of 'laziness' where the model doesn't complete a task.",
|
|
68
|
+
"context_window": 128000,
|
|
69
|
+
"max_output_tokens": 4096,
|
|
70
|
+
"training_data": "Up to Dec 2023"
|
|
71
|
+
},
|
|
72
|
+
"gpt-4-1106-preview": {
|
|
73
|
+
"description": "GPT-4 Turbo preview model featuring improved instruction following, JSON mode, reproducible outputs, parallel function calling, and more.",
|
|
74
|
+
"context_window": 128000,
|
|
75
|
+
"max_output_tokens": 4096,
|
|
76
|
+
"training_data": "Up to Apr 2023"
|
|
77
|
+
},
|
|
78
|
+
"gpt-4": {
|
|
79
|
+
"description": "Currently points to gpt-4-0613.",
|
|
80
|
+
"context_window": 8192,
|
|
81
|
+
"max_output_tokens": 8192,
|
|
82
|
+
"training_data": "Up to Sep 2021",
|
|
83
|
+
"version": "gpt-4-0613"
|
|
84
|
+
},
|
|
85
|
+
"gpt-4-0613": {
|
|
86
|
+
"description": "Snapshot of gpt-4 from June 13th 2023 with improved function calling support.",
|
|
87
|
+
"context_window": 8192,
|
|
88
|
+
"max_output_tokens": 8192,
|
|
89
|
+
"training_data": "Up to Sep 2021"
|
|
90
|
+
},
|
|
91
|
+
"gpt-4-0314": {
|
|
92
|
+
"description": "Legacy Snapshot of gpt-4 from March 14th 2023.",
|
|
93
|
+
"context_window": 8192,
|
|
94
|
+
"max_output_tokens": 8192,
|
|
95
|
+
"training_data": "Up to Sep 2021"
|
|
96
|
+
},
|
|
97
|
+
|
|
98
|
+
# GPT-3.5 Turbo models
|
|
99
|
+
"gpt-3.5-turbo": {
|
|
100
|
+
"description": "Currently points to gpt-3.5-turbo-0125.",
|
|
101
|
+
"context_window": 16385,
|
|
102
|
+
"max_output_tokens": 4096,
|
|
103
|
+
"training_data": "Up to Sep 2021",
|
|
104
|
+
"version": "gpt-3.5-turbo-0125"
|
|
105
|
+
},
|
|
106
|
+
"gpt-3.5-turbo-0125": {
|
|
107
|
+
"description": "The latest GPT-3.5 Turbo model with higher accuracy at responding in requested formats and a fix for a bug which caused a text encoding issue for non-English language function calls.",
|
|
108
|
+
"context_window": 16385,
|
|
109
|
+
"max_output_tokens": 4096,
|
|
110
|
+
"training_data": "Up to Sep 2021"
|
|
111
|
+
},
|
|
112
|
+
"gpt-3.5-turbo-1106": {
|
|
113
|
+
"description": "GPT-3.5 Turbo model with improved instruction following, JSON mode, reproducible outputs, parallel function calling, and more.",
|
|
114
|
+
"context_window": 16385,
|
|
115
|
+
"max_output_tokens": 4096,
|
|
116
|
+
"training_data": "Up to Sep 2021"
|
|
117
|
+
},
|
|
118
|
+
"gpt-3.5-turbo-instruct": {
|
|
119
|
+
"description": "Similar capabilities as GPT-3 era models. Compatible with legacy Completions endpoint and not Chat Completions.",
|
|
120
|
+
"context_window": 4096,
|
|
121
|
+
"max_output_tokens": 4096,
|
|
122
|
+
"training_data": "Up to Sep 2021"
|
|
123
|
+
}
|
|
124
|
+
}
|
|
@@ -0,0 +1,26 @@
|
|
|
1
|
+
# resk_openai_tokens.py
|
|
2
|
+
|
|
3
|
+
OPENAI_SPECIAL_TOKENS = {
|
|
4
|
+
"general": [
|
|
5
|
+
"<|endoftext|>",
|
|
6
|
+
"<|fim_prefix|>",
|
|
7
|
+
"<|fim_middle|>",
|
|
8
|
+
"<|fim_suffix|>",
|
|
9
|
+
"<|endofprompt|>",
|
|
10
|
+
"<s>",
|
|
11
|
+
"</s>",
|
|
12
|
+
"<|im_start|>",
|
|
13
|
+
"<|im_end|>",
|
|
14
|
+
"<|im_sep|>",
|
|
15
|
+
"[/INST]",
|
|
16
|
+
"SYS",
|
|
17
|
+
"[INST]",
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
],
|
|
21
|
+
"chat": [
|
|
22
|
+
"<|im_start|>user",
|
|
23
|
+
"<|im_start|>assistant",
|
|
24
|
+
"<|im_start|>system",
|
|
25
|
+
]
|
|
26
|
+
}
|
|
@@ -0,0 +1,42 @@
|
|
|
1
|
+
from resk_models import RESK_MODELS
|
|
2
|
+
from resk_openai_tokens import OPENAI_SPECIAL_TOKENS
|
|
3
|
+
from resk_context_manager import ContextManager
|
|
4
|
+
|
|
5
|
+
import html
|
|
6
|
+
|
|
7
|
+
class ReskProtector:
|
|
8
|
+
def __init__(self, model="gpt-4o", preserved_prompts=2, reserved_tokens=1000):
|
|
9
|
+
self.model = model
|
|
10
|
+
self.model_info = RESK_MODELS.get(model, RESK_MODELS["gpt-4o"])
|
|
11
|
+
self.special_tokens = set(OPENAI_SPECIAL_TOKENS["general"] + OPENAI_SPECIAL_TOKENS["chat"])
|
|
12
|
+
self.context_manager = ContextManager(self.model_info, preserved_prompts, reserved_tokens)
|
|
13
|
+
|
|
14
|
+
def sanitize_input(self, text):
|
|
15
|
+
text = text.encode('utf-8', errors='ignore').decode('utf-8')
|
|
16
|
+
for token in self.special_tokens:
|
|
17
|
+
text = text.replace(token, "")
|
|
18
|
+
text = self._close_html_tags(text)
|
|
19
|
+
text = html.escape(text)
|
|
20
|
+
return text
|
|
21
|
+
|
|
22
|
+
def _close_html_tags(self, text):
|
|
23
|
+
# Implémentation de la fermeture des balises HTML
|
|
24
|
+
pass
|
|
25
|
+
|
|
26
|
+
def protect_openai_call(self, api_function, *args, **kwargs):
|
|
27
|
+
sanitized_args = [self.sanitize_input(arg) if isinstance(arg, str) else arg for arg in args]
|
|
28
|
+
sanitized_kwargs = {k: self.sanitize_input(v) if isinstance(v, str) else v for k, v in kwargs.items()}
|
|
29
|
+
|
|
30
|
+
if 'messages' in sanitized_kwargs:
|
|
31
|
+
sanitized_kwargs['messages'] = self.context_manager.manage_sliding_context(sanitized_kwargs['messages'])
|
|
32
|
+
|
|
33
|
+
return api_function(*sanitized_args, **sanitized_kwargs)
|
|
34
|
+
|
|
35
|
+
@classmethod
|
|
36
|
+
def get_special_tokens(cls):
|
|
37
|
+
return OPENAI_SPECIAL_TOKENS
|
|
38
|
+
|
|
39
|
+
@classmethod
|
|
40
|
+
def update_special_tokens(cls, new_tokens):
|
|
41
|
+
global OPENAI_SPECIAL_TOKENS
|
|
42
|
+
OPENAI_SPECIAL_TOKENS = new_tokens
|
|
@@ -0,0 +1,151 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: resk-llm
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Resk-LLM is a robust Python library designed to enhance security and manage context when interacting with OpenAI's language models. It provides a protective layer for API calls, safeguarding against common vulnerabilities and ensuring optimal performance.
|
|
5
|
+
Home-page: https://github.com/Resk-Security/Resk-LLM
|
|
6
|
+
Author: Resk
|
|
7
|
+
Author-email: nielzac@proton.me
|
|
8
|
+
License: MIT License
|
|
9
|
+
|
|
10
|
+
Copyright (c)
|
|
11
|
+
|
|
12
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
13
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
14
|
+
in the Software without restriction, including without limitation the rights
|
|
15
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
16
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
17
|
+
furnished to do so, subject to the following conditions:
|
|
18
|
+
|
|
19
|
+
The above copyright notice and this permission notice shall be included in all
|
|
20
|
+
copies or substantial portions of the Software.
|
|
21
|
+
|
|
22
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
23
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
24
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
25
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
26
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
27
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
28
|
+
SOFTWARE.
|
|
29
|
+
Classifier: Programming Language :: Python :: 3
|
|
30
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
31
|
+
Classifier: Operating System :: OS Independent
|
|
32
|
+
Requires-Python: >=3.6
|
|
33
|
+
Description-Content-Type: text/markdown
|
|
34
|
+
License-File: LICENCE.txt
|
|
35
|
+
|
|
36
|
+
# LLM Protector for OpenAI
|
|
37
|
+
|
|
38
|
+
LLM Protector is a Python library designed to secure calls to the OpenAI API by protecting against context overflow attacks, cleaning inputs, and managing a sliding context window.
|
|
39
|
+
|
|
40
|
+
## Installation
|
|
41
|
+
|
|
42
|
+
```
|
|
43
|
+
pip install resk ##soon avalailble
|
|
44
|
+
```
|
|
45
|
+
Instead use :
|
|
46
|
+
```
|
|
47
|
+
git clone
|
|
48
|
+
```
|
|
49
|
+
## Usage
|
|
50
|
+
|
|
51
|
+
```python
|
|
52
|
+
from openai import OpenAI
|
|
53
|
+
from llm_protector import OpenAIProtector
|
|
54
|
+
|
|
55
|
+
# Initialize the OpenAI client
|
|
56
|
+
client = OpenAI(api_key="your-api-key")
|
|
57
|
+
|
|
58
|
+
# Create a protector instance
|
|
59
|
+
protector = OpenAIProtector(model="gpt-4o-mini", preserved_prompts=2, reserved_tokens=1000)
|
|
60
|
+
|
|
61
|
+
# Example usage with OpenAI API
|
|
62
|
+
messages = [
|
|
63
|
+
{"role": "system", "content": "You are a helpful assistant."},
|
|
64
|
+
{"role": "user", "content": "Hello, can you help me?"},
|
|
65
|
+
{"role": "assistant", "content": "Of course! How can I assist you today?"},
|
|
66
|
+
{"role": "user", "content": "I'd like to learn more about Python programming."}
|
|
67
|
+
]
|
|
68
|
+
|
|
69
|
+
response = protector.protect_openai_call(
|
|
70
|
+
client.chat.completions.create,
|
|
71
|
+
model="gpt-4o-mini",
|
|
72
|
+
messages=messages
|
|
73
|
+
)
|
|
74
|
+
|
|
75
|
+
print(response.choices[0].message.content)
|
|
76
|
+
```
|
|
77
|
+
|
|
78
|
+
## Features
|
|
79
|
+
|
|
80
|
+
### Model Selection
|
|
81
|
+
|
|
82
|
+
You can specify the model when initializing the protector:
|
|
83
|
+
|
|
84
|
+
```python
|
|
85
|
+
protector = OpenAIProtector(model="gpt-4o-mini")
|
|
86
|
+
```
|
|
87
|
+
|
|
88
|
+
To get the list of available models:
|
|
89
|
+
|
|
90
|
+
```python
|
|
91
|
+
print(OpenAIProtector.get_available_models())
|
|
92
|
+
```
|
|
93
|
+
|
|
94
|
+
To get information about a specific model:
|
|
95
|
+
|
|
96
|
+
```python
|
|
97
|
+
print(OpenAIProtector.get_model_info("gpt-4o-mini"))
|
|
98
|
+
```
|
|
99
|
+
|
|
100
|
+
### Special Tokens
|
|
101
|
+
|
|
102
|
+
The protector uses a predefined list of OpenAI special tokens for input sanitization. You can view and modify this list if necessary.
|
|
103
|
+
|
|
104
|
+
To get the current list of special tokens:
|
|
105
|
+
|
|
106
|
+
```python
|
|
107
|
+
print(OpenAIProtector.get_special_tokens())
|
|
108
|
+
```
|
|
109
|
+
|
|
110
|
+
To update the list of special tokens:
|
|
111
|
+
|
|
112
|
+
```python
|
|
113
|
+
new_tokens = {
|
|
114
|
+
"general": ["<|endoftext|>", "<|fim_prefix|>"],
|
|
115
|
+
"chat": ["<|im_start|>user", "<|im_start|>assistant"]
|
|
116
|
+
}
|
|
117
|
+
OpenAIProtector.update_special_tokens(new_tokens)
|
|
118
|
+
```
|
|
119
|
+
|
|
120
|
+
### Context Management
|
|
121
|
+
|
|
122
|
+
The protector automatically manages a sliding context for long conversations. You can specify the number of prompts to preserve and the number of reserved tokens during initialization:
|
|
123
|
+
|
|
124
|
+
```python
|
|
125
|
+
protector = OpenAIProtector(model="gpt-4o-mini", preserved_prompts=2, reserved_tokens=1000)
|
|
126
|
+
```
|
|
127
|
+
|
|
128
|
+
- `preserved_prompts`: Number of messages at the beginning of the conversation to always preserve (default 2, typically the system prompt and the first user message).
|
|
129
|
+
- `reserved_tokens`: Number of tokens reserved for preserved messages and the last message (default 1000).
|
|
130
|
+
|
|
131
|
+
The protector:
|
|
132
|
+
- Cleans messages by removing non-UTF-8 characters and special tokens.
|
|
133
|
+
- Automatically closes open HTML tags.
|
|
134
|
+
- Escapes special HTML characters.
|
|
135
|
+
- Truncates messages that are too long if necessary.
|
|
136
|
+
- Removes the oldest messages to respect the model's context limit.
|
|
137
|
+
|
|
138
|
+
## Security
|
|
139
|
+
|
|
140
|
+
The protector automatically sanitizes all inputs and manages the context when using the `protect_openai_call` method. Always use this method for your OpenAI API calls to benefit from all protections.
|
|
141
|
+
|
|
142
|
+
```
|
|
143
|
+
|
|
144
|
+
## License
|
|
145
|
+
|
|
146
|
+
This project is licensed under a Custom Open Source License. See the [LICENSE](LICENSE) file for full details. Key points:
|
|
147
|
+
|
|
148
|
+
- Free for commercial and non-commercial use
|
|
149
|
+
- Modifications must be documented and attributed
|
|
150
|
+
- Paid versions based on this code are subject to royalties
|
|
151
|
+
- The original source must be cited in all derivative works
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
LICENCE.txt
|
|
2
|
+
README.md
|
|
3
|
+
setup.py
|
|
4
|
+
resk-llm/__init__.py
|
|
5
|
+
resk-llm/openai_protector.py
|
|
6
|
+
resk-llm/resk_context_manager.py
|
|
7
|
+
resk-llm/resk_models.py
|
|
8
|
+
resk-llm/resk_openai_tokens.py
|
|
9
|
+
resk-llm/resk_protector.py
|
|
10
|
+
resk_llm.egg-info/PKG-INFO
|
|
11
|
+
resk_llm.egg-info/SOURCES.txt
|
|
12
|
+
resk_llm.egg-info/dependency_links.txt
|
|
13
|
+
resk_llm.egg-info/requires.txt
|
|
14
|
+
resk_llm.egg-info/top_level.txt
|
|
15
|
+
tests/test_openai_protector.py.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
openai
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
resk-llm
|
resk-llm-0.1.0/setup.cfg
ADDED
resk-llm-0.1.0/setup.py
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
from setuptools import setup, find_packages
|
|
2
|
+
|
|
3
|
+
with open("README.md", "r", encoding="utf-8") as fh:
|
|
4
|
+
long_description = fh.read()
|
|
5
|
+
|
|
6
|
+
with open("LICENCE.txt", "r", encoding="utf-8") as f:
|
|
7
|
+
license_text = f.read()
|
|
8
|
+
|
|
9
|
+
setup(
|
|
10
|
+
name="resk-llm",
|
|
11
|
+
version="0.1.0",
|
|
12
|
+
author="Resk",
|
|
13
|
+
author_email="nielzac@proton.me",
|
|
14
|
+
description="Resk-LLM is a robust Python library designed to enhance security and manage context when interacting with OpenAI's language models. It provides a protective layer for API calls, safeguarding against common vulnerabilities and ensuring optimal performance.",
|
|
15
|
+
long_description=long_description,
|
|
16
|
+
long_description_content_type="text/markdown",
|
|
17
|
+
url="https://github.com/Resk-Security/Resk-LLM",
|
|
18
|
+
packages=find_packages(),
|
|
19
|
+
classifiers=[
|
|
20
|
+
"Programming Language :: Python :: 3",
|
|
21
|
+
"License :: OSI Approved :: MIT License",
|
|
22
|
+
"Operating System :: OS Independent",
|
|
23
|
+
],
|
|
24
|
+
python_requires=">=3.6",
|
|
25
|
+
install_requires=[
|
|
26
|
+
"openai",
|
|
27
|
+
],
|
|
28
|
+
license=license_text,
|
|
29
|
+
)
|
|
@@ -0,0 +1,65 @@
|
|
|
1
|
+
import unittest
|
|
2
|
+
from unittest.mock import patch, MagicMock
|
|
3
|
+
from resk.openai_protector import OpenAIProtector
|
|
4
|
+
|
|
5
|
+
class TestOpenAIProtector(unittest.TestCase):
|
|
6
|
+
|
|
7
|
+
def setUp(self):
|
|
8
|
+
self.protector = OpenAIProtector(model="gpt-4o", preserved_prompts=2)
|
|
9
|
+
|
|
10
|
+
def test_sanitize_input(self):
|
|
11
|
+
input_text = "<script>alert('XSS')</script>Hello<endoftext>"
|
|
12
|
+
sanitized = self.protector.sanitize_input(input_text)
|
|
13
|
+
self.assertEqual(sanitized, "<script>alert('XSS')</script>Hello")
|
|
14
|
+
|
|
15
|
+
def test_close_html_tags(self):
|
|
16
|
+
input_text = "<p>Unclosed paragraph<div>Nested <b>bold"
|
|
17
|
+
closed = self.protector._close_html_tags(input_text)
|
|
18
|
+
self.assertEqual(closed, "<p>Unclosed paragraph<div>Nested <b>bold</b></div></p>")
|
|
19
|
+
|
|
20
|
+
def test_truncate_text(self):
|
|
21
|
+
long_text = "a" * (self.protector.max_context_length * 5)
|
|
22
|
+
truncated = self.protector._truncate_text(long_text)
|
|
23
|
+
self.assertEqual(len(truncated), self.protector.max_context_length * 4)
|
|
24
|
+
|
|
25
|
+
def test_manage_sliding_context(self):
|
|
26
|
+
messages = [
|
|
27
|
+
{"role": "system", "content": "You are a helpful assistant."},
|
|
28
|
+
{"role": "user", "content": "Hello"},
|
|
29
|
+
{"role": "assistant", "content": "Hi there!"},
|
|
30
|
+
{"role": "user", "content": "How are you?"},
|
|
31
|
+
{"role": "assistant", "content": "I'm doing well, thank you for asking!"},
|
|
32
|
+
]
|
|
33
|
+
managed = self.protector.manage_sliding_context(messages)
|
|
34
|
+
self.assertEqual(len(managed), 4) # 2 preserved + 2 most recent
|
|
35
|
+
self.assertEqual(managed[0], messages[0])
|
|
36
|
+
self.assertEqual(managed[1], messages[1])
|
|
37
|
+
self.assertEqual(managed[-2], messages[-2])
|
|
38
|
+
self.assertEqual(managed[-1], messages[-1])
|
|
39
|
+
|
|
40
|
+
@patch('openai.ChatCompletion.create')
|
|
41
|
+
def test_protect_openai_call(self, mock_create):
|
|
42
|
+
mock_response = MagicMock()
|
|
43
|
+
mock_response.choices = [MagicMock(message={"content": "Test response"})]
|
|
44
|
+
mock_create.return_value = mock_response
|
|
45
|
+
|
|
46
|
+
messages = [
|
|
47
|
+
{"role": "system", "content": "You are a helpful assistant."},
|
|
48
|
+
{"role": "user", "content": "Hello<endoftext>"},
|
|
49
|
+
]
|
|
50
|
+
|
|
51
|
+
response = self.protector.protect_openai_call(
|
|
52
|
+
mock_create,
|
|
53
|
+
model="gpt-4o",
|
|
54
|
+
messages=messages
|
|
55
|
+
)
|
|
56
|
+
|
|
57
|
+
self.assertEqual(response.choices[0].message["content"], "Test response")
|
|
58
|
+
mock_create.assert_called_once()
|
|
59
|
+
called_args = mock_create.call_args[1]
|
|
60
|
+
self.assertEqual(called_args['model'], "gpt-4o")
|
|
61
|
+
self.assertEqual(len(called_args['messages']), 2)
|
|
62
|
+
self.assertNotIn("<endoftext>", called_args['messages'][1]['content'])
|
|
63
|
+
|
|
64
|
+
if __name__ == '__main__':
|
|
65
|
+
unittest.main()
|