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.
@@ -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.
@@ -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,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,3 @@
1
+ from .openai_protector import OpenAIProtector
2
+
3
+ __version__ = "0.1.0"
@@ -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
+ openai
@@ -0,0 +1 @@
1
+ resk-llm
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -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, "&lt;script&gt;alert('XSS')&lt;/script&gt;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()