├── LICENSE
├── README.md
├── analytics.py
├── reflection.py
└── wrapper.py
/LICENSE:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Copyright (c) 2023 Yohei Nakajima
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.
22 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # Wrapper's Delight: An Enhanced OpenAI Wrapper
2 | **`wrappers_delight`** is a Python wrapper built around OpenAI's ChatCompletion API. The main features of this wrapper are:
3 |
4 | * Automated logging to a NDJSON file for every interaction.
5 | * Analytics functions for visualizing model usage.
6 | * Parameter-based or AI-assisted querying of logs.
7 | * (Optional) Automated reflection and suggested improvements on every prompt.
8 |
9 |
10 | To get started, you'll need to clone this repository:
11 |
12 | ```
13 | git clone https://github.com/yoheinakajima/wrappers_delight.git
14 | cd wrappers_delight
15 | ```
16 | ## Prerequisites
17 | Ensure you have OpenAI's Python client installed:
18 |
19 | ```
20 | pip install openai
21 | ```
22 | # Usage
23 | ## Basic Chat Completion
24 | To use the wrapper for a basic chat with the model, load the wrapper and every prompt input and output will be stored:
25 | **Aside from line 2, this is a standard OpenAI ChatCompletion call.*
26 | ```
27 | import openai
28 | from wrappers_delight import ChatWrapper
29 |
30 | # Set your OpenAI API key
31 | openai.api_key = "YOUR_API_KEY"
32 |
33 | # Engage in a conversation with the model
34 | response = openai.ChatCompletion.create(
35 | model="gpt-3.5-turbo",
36 | messages=[
37 | {"role": "system", "content": "You are a helpful assistant."},
38 | {"role": "user", "content": "Tell me a fun fact."},
39 | ]
40 | )
41 |
42 | print(response.choices[0].message.content)
43 | ```
44 | ## With Custom Functions
45 | To register and use a custom function:
46 | **Aside from line 2, this is a standard OpenAI ChatCompletion call using functions.*
47 | ```
48 | import openai
49 | from wrappers_delight.wrapper import _enhanced_chat_completion
50 |
51 | # Set your OpenAI API key
52 | openai.api_key = "YOUR_OPENAI_KEY"
53 |
54 | completion = openai.ChatCompletion.create(
55 | model="gpt-3.5-turbo-0613",
56 | messages=[
57 | {
58 | "role": "user",
59 | "content": "I visited Tokyo, then moved to San Francisco, and finally settled in Toronto."
60 | }
61 | ],
62 | functions=[
63 | {
64 | "name": "extract_locations",
65 | "description": "Extract all locations mentioned in the text",
66 | "parameters": {
67 | "type": "object",
68 | "properties": {
69 | "locations": {
70 | "type": "array",
71 | "items": {
72 | "type": "object",
73 | "properties": {
74 | "name": {
75 | "type": "string",
76 | "description": "The name of the location"
77 | },
78 | "country_iso_alpha2": {
79 | "type": "string",
80 | "description": "The ISO alpha-2 code of the country where the location is situated"
81 | }
82 | },
83 | "required": ["name", "country_iso_alpha2"]
84 | }
85 | }
86 | },
87 | "required": ["locations"],
88 | },
89 | },
90 | ],
91 | function_call={"name": "extract_locations"}
92 | )
93 |
94 | response_data = completion.choices[0]['message']['function_call']['arguments']
95 | print(response_data)
96 | ```
97 | ## Logging
98 | All interactions with the model are automatically logged to log.ndjson. Each row in the NDJSON consists of the request parameters and the corresponding model response.
99 | ## Example Output
100 | The following are 2 example outputs in log.ndjson. The first is a standard ChatCompletion call, the second is a function call:
101 | ```
102 | {"timestamp": "2023-08-13 03:00:49", "response_time": 3.21, "params": {"model": "gpt-3.5-turbo", "messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Tell me a fun fact."}]}, "response": {"id": "chatcmpl-7mvdPu0H1QULvDZOUVJS6npdMslul", "object": "chat.completion", "created": 1691895647, "model": "gpt-3.5-turbo-0613", "choices": [{"index": 0, "message": {"role": "assistant", "content": "Sure! Here's a fun fact: Did you know that honey never spoils? Archaeologists have found pots of honey in ancient Egyptian tombs that are over 3,000 years old and still perfectly edible! Honey's low moisture content and acidic pH create an inhospitable environment for bacteria and other microorganisms, allowing it to last indefinitely."}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 23, "completion_tokens": 70, "total_tokens": 93}}, "total_tokens": 93, "model": "gpt-3.5-turbo"}
103 | {"timestamp": "2023-08-13 03:01:16", "response_time": 4.80, "params": {"model": "gpt-3.5-turbo-0613", "messages": [{"role": "user", "content": "I visited Tokyo, then moved to San Francisco, and finally settled in Toronto."}], "functions": [{"name": "extract_locations", "description": "Extract all locations mentioned in the text", "parameters": {"type": "object", "properties": {"locations": {"type": "array", "items": {"type": "object", "properties": {"name": {"type": "string", "description": "The name of the location"}, "country_iso_alpha2": {"type": "string", "description": "The ISO alpha-2 code of the country where the location is situated"}}, "required": ["name", "country_iso_alpha2"]}}}, "required": ["locations"]}}], "function_call": {"name": "extract_locations"}}, "response": {"id": "chatcmpl-7mvdqfScl0uQ2tfye1HdfcPEV5XYI", "object": "chat.completion", "created": 1691895674, "model": "gpt-3.5-turbo-0613", "choices": [{"index": 0, "message": {"role": "assistant", "content": null, "function_call": {"name": "extract_locations", "arguments": "{\n \"locations\": [\n {\n \"name\": \"Tokyo\",\n \"country_iso_alpha2\": \"JP\"\n },\n {\n \"name\": \"San Francisco\",\n \"country_iso_alpha2\": \"US\"\n },\n {\n \"name\": \"Toronto\",\n \"country_iso_alpha2\": \"CA\"\n }\n ]\n}"}}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 83, "completion_tokens": 74, "total_tokens": 157}}, "total_tokens": 157, "model": "gpt-3.5-turbo-0613"}
104 | ```
105 |
106 | # Analytics
107 | To run the analytics and visualize the model's usage, load the analytics:
108 | ```
109 | from wrappers_delight.analytics import plot_token_usage, plot_model_distribution, plot_response_time_distribution
110 |
111 | # Generate and display the analytics plots
112 | plot_token_usage()
113 | plot_model_distribution()
114 | plot_response_time()
115 | plot_response_time_distribution()
116 | ```
117 | # Query Logs
118 | To enable querying your logs, load the two query functions:
119 | ```
120 | from wrappers_delight.analytics import query_log, query_log_with_ai
121 | ```
122 | *query_log()* and *query_log_with_ai()* are functions that serve to retrieve specific entries from logs. While query_log directly queries the log based on various parameters, query_log_with_ai uses an AI model to further refine those results based on context or more complex requirements.
123 | ## Prerequisites
124 | Ensure you have the required libraries installed. This includes pandas, openai, and others depending on your needs.
125 | ## using query_log()
126 | Parameters:
127 | * columns: List of columns to be displayed in the result.
128 | * start_date: The starting date for the logs you want to retrieve.
129 | * end_date: The end date for the logs.
130 | * min_tokens: Minimum number of tokens in the response.
131 | * max_tokens: Maximum number of tokens in the response.
132 | * filter_by: Additional filter criteria.
133 | * sort_by: The column by which the results should be sorted (default is timestamp).
134 | * limit: Limit the number of rows in the result.
135 | * keyword: Search for a keyword in user message or response.
136 | * function_name: Filter by the function name.
137 | * model_version: Filter by model version.
138 |
139 | Sample Code:
140 |
141 | ```
142 | result = query_log(start_date="2023-08-01", end_date="2023-08-10", keyword="weather")
143 | print(result)
144 | ```
145 | ## using query_log_with_ai
146 | This function interprets natural language user queries and translates them into parameters to fetch appropriate log entries using the query_log function.
147 | ### Example Usage:
148 | User Query:
149 | ```
150 | result = query_log_with_ai("Show me the last 3 logs.")
151 | ```
152 | Expected Parameters for query_log:
153 | ```
154 | {
155 | "limit": 3,
156 | "sort_by": "timestamp"
157 | }
158 | ```
159 | User Query:
160 | ```
161 | result = query_log_with_ai("I'd like to see the top 5 logs that mention 'weather'.")
162 | ```
163 | Expected Parameters for query_log:
164 | ```
165 | {
166 | "keyword": "weather",
167 | "limit": 5
168 | }
169 | ```
170 | User Query:
171 | ```
172 | result = query_log_with_ai("Can you display the logs sorted by total tokens for the gpt-3.5-turbo model, but only show the timestamp, model, and total tokens columns?
173 | ")
174 | ```
175 | Expected Parameters for query_log:
176 | ```
177 | {
178 | "model_version": "gpt-3.5-turbo",
179 | "sort_by": "total_tokens",
180 | "columns": ["timestamp", "model", "total_tokens"]
181 | }
182 | ```
183 | User Query:
184 | ```
185 | result = query_log_with_ai("I want logs that used between 50 to 1000 tokens and only used the gpt-3.5-turbo model.")
186 | ```
187 | Expected Parameters for query_log:
188 | ```
189 | {
190 | "min_tokens": 50,
191 | "max_tokens": 1000,
192 | "filter_by": {
193 | "model": "gpt-3.5-turbo"
194 | }
195 | }
196 | ```
197 | User Query:
198 | ```
199 | result = query_log_with_ai("Show me the logs from August 12th to August 14th.")
200 | ```
201 | Expected Parameters for query_log:
202 | ```
203 | {
204 | "start_date": "2023-08-12",
205 | "end_date": "2023-08-14"
206 | }
207 | ```
208 | The *`query_log_with_ai()`* function will automatically use the generated parameters to call the *`query_log()`* function. The examples above were provides to help you understand the capabilities of *`query_log_with_ai()`* while also providing you with additional examples for the *`query_log()`* function.
209 |
210 | ## Automatic Storage of Log Queries
211 | * The results are stored in the *log_queries* directory with a unique filename.
212 | * Always ensure that sensitive data is properly handled and not exposed.
213 |
214 | # Enabling and Using the Reflection Feature
215 | The reflection capability provides insights into how the OpenAI model interpreted given prompts, evaluating its responses and suggesting potential improvements. This feature is integrated into the `wrapper.ChatWrapper`, making it easy to enable or disable.
216 |
217 | ## How to Enable:
218 | 1. Enable Reflection: Simply use the following line to enable the reflection feature for all subsequent ChatCompletion calls:
219 | ```
220 | wrapper.ChatWrapper.enable_reflection()
221 | ```
222 | This will automatically turn on reflection for all chat completion API calls made using the `ChatWrapper`.
223 |
224 | 2. Disable Reflection: If you wish to turn off the reflection feature:
225 | ```
226 | wrapper.ChatWrapper.disable_reflection()
227 | ```
228 | ## Storage of Reflections:
229 |
230 | Reflections are recorded and stored in the `prompt_reflections.ndjson` file. Each entry within this file represents a distinct reflection and is formatted as a newline-delimited JSON (NDJSON). This allows for easier appending of new reflections and streamlined parsing.
231 |
232 | The structure for each reflection entry is as follows:
233 |
234 | - **kwargs**: This contains the input prompt provided to the OpenAI model.
235 | - **reflection**:
236 | - **Analysis**: Insights into how the AI interpreted the prompt and the rationale behind its response.
237 | - **Is Satisfactory**: Indicates whether the AI believes its response was satisfactory. It can be `True` or `False`.
238 | - **Suggested Prompt**: If the AI feels the prompt could be clearer or more optimized, this provides a suggested improvement.
239 |
240 | # Wrappers Delight UI
241 |
242 | Looking for a user-friendly interface to visualize and interpret logs generated by `wrappers_delight`? Check out [`wrappers_delight_ui`](https://github.com/yoheinakajima/wrappers_delight_ui)!
243 |
244 | `wrappers_delight_ui` offers an intuitive way to see the reflections and prompt-response dynamics in action. Gain insights, optimize prompts, and elevate your understanding of how AI processes and responds to various prompts.
245 |
246 | [Explore Wrappers Delight UI on GitHub](https://github.com/yoheinakajima/wrappers_delight_ui)
247 |
248 | # Contributing
249 | If you'd like to contribute, please fork the repository and make changes as you'd like. Pull requests are warmly welcomed.
250 |
251 | # License
252 | Distributed under the MIT License. See LICENSE for more information.
253 |
--------------------------------------------------------------------------------
/analytics.py:
--------------------------------------------------------------------------------
1 | import os
2 | import openai
3 | from datetime import datetime
4 | import matplotlib.pyplot as plt
5 | import pandas as pd
6 | from .wrapper import _original_chat_completion
7 |
8 | LOG_FILE = 'log.ndjson'
9 |
10 | def calculate_cost(tokens):
11 | COST_PER_TOKEN = 0.01
12 | return tokens * COST_PER_TOKEN
13 |
14 | def plot_token_usage():
15 | data = pd.read_json(LOG_FILE, lines=True)
16 | dates = data['timestamp']
17 | token_counts = data['total_tokens']
18 |
19 | plt.plot(dates, token_counts)
20 | plt.xticks(rotation=45)
21 | plt.ylabel("Tokens")
22 | plt.xlabel("Timestamp")
23 | plt.title("Token Usage Over Time")
24 | plt.tight_layout()
25 | plt.savefig("token_usage.png")
26 | plt.cla()
27 |
28 | def plot_model_distribution():
29 | data = pd.read_json(LOG_FILE, lines=True)
30 | models = data['model'].value_counts().to_dict()
31 |
32 | plt.pie(models.values(), labels=models.keys(), autopct='%1.1f%%')
33 | plt.title("Model Usage Distribution")
34 | plt.savefig("model_distribution.png")
35 | plt.cla()
36 |
37 | def plot_response_time_distribution():
38 | data = pd.read_json(LOG_FILE, lines=True)
39 | dates = data['timestamp']
40 | response_times = data['response_time']
41 |
42 | plt.plot(dates, response_times)
43 | plt.xticks(rotation=45)
44 | plt.ylabel("Response Time (s)")
45 | plt.xlabel("Timestamp")
46 | plt.title("Response Time Over Time")
47 | plt.tight_layout()
48 | plt.savefig("response_time.png")
49 | plt.cla()
50 |
51 | def query_log(columns=None, start_date=None, end_date=None, min_tokens=None,
52 | max_tokens=None, filter_by=None, sort_by='timestamp', limit=None, keyword=None,
53 | function_name=None, model_version=None):
54 | """
55 | This function queries the log based on various parameters.
56 | """
57 | data = pd.read_json(LOG_FILE, lines=True)
58 |
59 | # Helper function to extract the user's message
60 | def extract_user_message(params):
61 | for msg in params['messages']:
62 | if msg['role'] == 'user' and 'content' in msg:
63 | return msg['content']
64 | return None
65 |
66 | # Extract the user's message
67 | data['user_message'] = data['params'].apply(extract_user_message)
68 |
69 | # Helper function to extract the assistant's response
70 | def extract_response_content(row):
71 | if 'function_call' in row:
72 | return str(row['function_call'])
73 | elif 'choices' in row and len(row['choices']) > 0:
74 | content = row['choices'][0]['message'].get('content')
75 | return content if content else None
76 | return None
77 |
78 | # Extract the assistant's response
79 | data['response_content'] = data['response'].apply(extract_response_content)
80 |
81 | # Combine the user's message and the assistant's response
82 | data['combined_content'] = data['user_message'].astype(str) + " | " + data['response_content'].astype(str)
83 |
84 | # Filter by date range
85 | if start_date:
86 | data = data[data['timestamp'] >= start_date]
87 | if end_date:
88 | data = data[data['timestamp'] <= end_date]
89 |
90 | # Filter by token range
91 | if min_tokens:
92 | data = data[data['total_tokens'] >= min_tokens]
93 | if max_tokens:
94 | data = data[data['total_tokens'] <= max_tokens]
95 |
96 | # Search for a keyword in combined_content
97 | if keyword:
98 | data = data[data['combined_content'].str.contains(keyword, case=False)]
99 |
100 | # Filter by function name or model version
101 | if function_name:
102 | data = data[data['function_name'] == function_name]
103 |
104 | if model_version:
105 | data = data[data['model'] == model_version]
106 |
107 | # Additional filter criteria
108 | if filter_by:
109 | for key, value in filter_by.items():
110 | data = data[data[key] == value]
111 |
112 | # Select specific columns
113 | if columns:
114 | data = data[columns]
115 |
116 | # Sort the data
117 | if sort_by:
118 | data = data.sort_values(by=sort_by, ascending=False)
119 |
120 | # Limit number of rows
121 | if limit:
122 | data = data.head(limit)
123 |
124 | # Use the create_file_name function to get the filename
125 | csv_filename = os.path.join("log_queries", create_file_name(start_date, end_date, keyword, model_version))
126 |
127 | if not os.path.exists('log_queries'):
128 | os.makedirs('log_queries')
129 |
130 | data.to_csv(csv_filename, index=False)
131 |
132 | return data
133 |
134 |
135 | def create_file_name(start_date=None, end_date=None, keyword=None, model_version=None):
136 | """
137 | Create a file name based on the given parameters.
138 |
139 | Parameters:
140 | - start_date (str): The start date of the log.
141 | - end_date (str): The end date of the log.
142 | - keyword (str): The keyword used in the query.
143 | - model_version (str): The model version used in the query.
144 |
145 | Returns:
146 | str: The constructed file name.
147 | """
148 | details = []
149 |
150 | if start_date and end_date:
151 | details.append(f"date_{start_date}_to_{end_date}")
152 | elif start_date:
153 | details.append(f"from_{start_date}")
154 | elif end_date:
155 | details.append(f"until_{end_date}")
156 |
157 | if keyword:
158 | details.append(f"keyword_{keyword}")
159 |
160 | if model_version:
161 | details.append(f"model_{model_version}")
162 |
163 | # Combining the details with underscores and limiting the length to ensure the filename isn't too long
164 | details_str = '_'.join(details)[:150] # Limiting to 150 characters
165 | timestamp_str = datetime.now().strftime('%Y%m%d_%H%M%S')
166 |
167 | return f"query_{timestamp_str}_{details_str}.csv"
168 |
169 |
170 | def query_log_with_ai(natural_language_input):
171 | # Define the initial session prompt
172 | messages = [
173 | {
174 | "role": "system",
175 | "content": """
176 | The function 'query_log' allows users to query a log of API requests based on various parameters.
177 | The available functionalities within 'query_log' are:
178 | 1. Filtering by specific date ranges using 'start_date' and 'end_date'.
179 | 2. Filtering by a range of token counts using 'min_tokens' and 'max_tokens'.
180 | 3. Searching for keywords in the prompt and response using 'keyword'.
181 | 4. Filtering by specific function names using 'function_name'.
182 | 5. Filtering by specific model versions using 'model_version'.
183 | 6. Adding custom filters through the 'filter_by' dictionary.
184 | 7. Specifying columns to return using 'columns'.
185 | 8. Sorting results by a specific column using 'sort_by' (default is 'timestamp').
186 | 9. Limiting the number of rows returned using 'limit'.
187 | 10. Always include .
188 | The goal is to turn natural language requests into specific queries for the 'query_log' function.
189 | """
190 | },
191 | {
192 | "role": "user",
193 | "content": "Show me the logs from August 12th to August 14th."
194 | },
195 | {
196 | "role": "assistant",
197 | "content": "{'start_date': '2023-08-12', 'end_date': '2023-08-14'}"
198 | },
199 | {
200 | "role": "user",
201 | "content": "I'd like to see the top 5 logs that mention 'weather'."
202 | },
203 | {
204 | "role": "assistant",
205 | "content": "{'keyword': 'weather', 'limit': 5}"
206 | },
207 | {
208 | "role": "user",
209 | "content": "Can you display the logs sorted by total tokens for the gpt-3.5-turbo model, but only show the timestamp, model, and total tokens columns?"
210 | },
211 | {
212 | "role": "assistant",
213 | "content": "{'model_version': 'gpt-3.5-turbo', 'sort_by': 'total_tokens', 'columns': ['timestamp', 'model', 'total_tokens']}"
214 | },
215 | {
216 | "role": "user",
217 | "content": "I want logs that used between 50 to 1000 tokens and only used the gpt-3.5-turbo model."
218 | },
219 | {
220 | "role": "assistant",
221 | "content": "{'min_tokens': 50, 'max_tokens': 1000, 'filter_by': {'model': 'gpt-3.5-turbo'}}"
222 | },
223 | {
224 | "role": "user",
225 | "content": "Give me up to 3 logs between August 10th and 15th, and only show their timestamp and response."
226 | },
227 | {
228 | "role": "assistant",
229 | "content": "{'start_date': '2023-08-10', 'end_date': '2023-08-15', 'columns': ['timestamp', 'response'], 'limit': 3}"
230 | },
231 | {
232 | "role": "user",
233 | "content": natural_language_input
234 | }
235 | ]
236 |
237 | # Send the messages to OpenAI's ChatCompletion API
238 | response = _original_chat_completion(
239 | model="gpt-3.5-turbo",
240 | messages=messages
241 | )
242 |
243 | # Extract the query from the assistant's most recent message
244 | query_str = response.choices[0].message['content'].strip()
245 |
246 | # Evaluate the query to get a Python dictionary
247 | query_dict = eval(query_str)
248 |
249 | # Use the query_log function to fetch the results
250 | results = query_log(**query_dict)
251 |
252 | return results
253 |
--------------------------------------------------------------------------------
/reflection.py:
--------------------------------------------------------------------------------
1 | import openai
2 | import json
3 | from .wrapper import _original_chat_completion
4 |
5 | def log_reflection_to_ndjson(kwargs, reflection_response):
6 | reflection_entry = {
7 | "kwargs": kwargs,
8 | "reflection": reflection_response
9 | }
10 |
11 | with open("prompt_reflections.ndjson", mode='a', encoding='utf-8') as file:
12 | file.write(json.dumps(reflection_entry) + '\n')
13 |
14 | def generate_prompt_reflection(kwargs, response):
15 | # Define the function to analyze and optimize prompts
16 | function_description = {
17 | "name": "analyze_and_optimize_prompt",
18 | "description": "Analyze the prompt and response to provide suggestions for optimization",
19 | "parameters": {
20 | "type": "object",
21 | "properties": {
22 | "analysis": {
23 | "type": "string",
24 | "description": "Analysis and possible optimizations for the prompt."
25 | },
26 | "is_satisfactory": {
27 | "type": "boolean",
28 | "description": "Indicates if the prompt was satisfactory."
29 | },
30 | "suggested_prompt": {
31 | "type": "string",
32 | "description": "A better version of the prompt, if applicable."
33 | }
34 | },
35 | "required": ["analysis", "is_satisfactory", "suggested_prompt"]
36 | },
37 | }
38 |
39 | function_call = {
40 | "name": "analyze_and_optimize_prompt",
41 | "arguments": {
42 | "prompt": kwargs.get("prompt", ""), # extract the user's prompt
43 | "response": response.get("choices", [{}])[0].get("text", "").strip()
44 | }
45 | }
46 |
47 | reflection_response = _original_chat_completion(
48 | model="gpt-3.5-turbo-16k",
49 | messages=[
50 | {
51 | "role": "system",
52 | "content": "You are an AI assistant specializing in providing advice on the use of large language models, specifically optimizing prompts for the ChatCompletion or Function Call endpoint. Analyze the user message, which will be a JSON array with data about the API call (both input and output), and analyze the user prompt (or function call), and provide suggested improvements."
53 | },
54 | {
55 | "role": "user",
56 | "content": json.dumps({
57 | "kwargs": kwargs,
58 | "response": response
59 | })
60 | }
61 | ],
62 | functions=[function_description],
63 | function_call=function_call
64 | )
65 |
66 | reflection_content = reflection_response["choices"][0]["message"]["function_call"]["arguments"]
67 |
68 | # Log the reflection content
69 | log_reflection_to_ndjson(kwargs, reflection_content)
70 |
71 | return reflection_content
72 |
--------------------------------------------------------------------------------
/wrapper.py:
--------------------------------------------------------------------------------
1 | import openai
2 | import json
3 | import datetime
4 | import time
5 |
6 | _original_chat_completion = openai.ChatCompletion.create
7 |
8 | class ChatWrapper:
9 | reflection_enabled = False
10 |
11 | @staticmethod
12 | def enable_reflection():
13 | ChatWrapper.reflection_enabled = True
14 |
15 | @staticmethod
16 | def disable_reflection():
17 | ChatWrapper.reflection_enabled = False
18 |
19 | @staticmethod
20 | def log_to_ndjson(params, task_label, response, response_time):
21 | timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
22 | model = params.get("model", "")
23 | total_tokens = response["usage"]["total_tokens"]
24 |
25 | log_entry = {
26 | "timestamp": timestamp,
27 | "response_time": round(response_time, 2),
28 | "params": params,
29 | "response": response,
30 | "total_tokens": total_tokens,
31 | "model": model,
32 | "taskLabel": task_label # Add taskLabel here
33 | }
34 |
35 | with open("log.ndjson", mode='a', encoding='utf-8') as file:
36 | file.write(json.dumps(log_entry) + '\n')
37 |
38 | @staticmethod
39 | def enhanced_chat_completion(*args, **kwargs):
40 | timestamp_start = time.perf_counter()
41 |
42 | # Extract and remove taskLabel from kwargs if it exists
43 | task_label = kwargs.pop("taskLabel", None)
44 |
45 | # Call the original OpenAI method
46 | response = _original_chat_completion(*args, **kwargs)
47 |
48 | timestamp_end = time.perf_counter()
49 | response_time = timestamp_end - timestamp_start
50 |
51 | # Log with taskLabel
52 | ChatWrapper.log_to_ndjson(kwargs, task_label, response, response_time)
53 |
54 | if ChatWrapper.reflection_enabled:
55 | from .reflection import generate_prompt_reflection
56 | generate_prompt_reflection(kwargs, response)
57 |
58 | return response
59 |
60 | # Override the ChatCompletion with the enhanced version
61 | openai.ChatCompletion.create = ChatWrapper.enhanced_chat_completion
62 |
--------------------------------------------------------------------------------