Mahbodez/knee_report_checklist
1
1import colorama2from colorama import Fore, Style3import openai4from tenacity import retry, stop_after_attempt, wait_fixed5import json6import os7import tiktoken8import functools as ft9import time10 11JSON_TEMPLATE = """12{question}13The required key(s) are: {keys}.14Only and only respond with the key(s) and value(s) mentioned above.15Your answer in valid JSON format:\n16"""17 18MODEL_COST_DICT = {19 "gpt-3.5-turbo": {20 "input": 0.0015,21 "output": 0.002,22 },23 "gpt-4": {24 "input": 0.03,25 "output": 0.06,26 },27}28 29 30def set_api_key(key=None):31 """Sets the OpenAI API key."""32 if key is None:33 key = os.environ.get("OPENAI_API_KEY")34 openai.api_key = key35 36 37def num_tokens_from_string(string: str, encoding_name: str) -> int:38 """Returns the number of tokens in a text string."""39 encoding = tiktoken.get_encoding(encoding_name)40 num_tokens = len(encoding.encode(string))41 return num_tokens42 43 44def num_tokens_from_messages(messages: list[dict], model="gpt-3.5-turbo-0613"):45 """Returns the number of tokens used by a list of messages."""46 try:47 encoding = tiktoken.encoding_for_model(model)48 except KeyError:49 encoding = tiktoken.get_encoding("cl100k_base")50 if model == "gpt-3.5-turbo-0613": # note: future models may deviate from this51 num_tokens = 052 for message in messages:53 num_tokens += (54 4 # every message follows <im_start>{role/name}\n{content}<im_end>\n55 )56 for key, value in message.items():57 num_tokens += len(encoding.encode(value))58 if key == "name": # if there's a name, the role is omitted59 num_tokens += -1 # role is always required and always 1 token60 num_tokens += 2 # every reply is primed with <im_start>assistant61 return num_tokens62 else:63 raise NotImplementedError(64 f"""num_tokens_from_messages() is not presently implemented for model {model}.65 See https://github.com/openai/openai-python/blob/main/chatml.md for information on how messages are converted to tokens."""66 )67 68 69@retry(stop=stop_after_attempt(3), wait=wait_fixed(2))70def chat(messages: list[dict], model="gpt-3.5-turbo", temperature=0.0):71 response = openai.ChatCompletion().create(72 model=model,73 messages=messages,74 temperature=temperature,75 )76 return response["choices"][0]["message"]["content"]77 78 79def make_message(role: str, content: str) -> dict:80 return {81 "role": role,82 "content": content,83 }84 85 86def make_prompt(template: str, **kwargs):87 return template.format(**kwargs)88 89 90def unravel_messages(messages: list[dict]) -> list[str]:91 """Returns a string representation of a list of messages."""92 return [f"{message['role']}: {message['content']}" for message in messages]93 94 95class LLM:96 def __init__(self, model="gpt-3.5-turbo", temperature=0.0):97 self.model = model98 self.temperature = temperature99 self.token_counter = 0100 self.cost = 0.0101 102 @retry(stop=stop_after_attempt(3), wait=wait_fixed(2))103 def chat(self, messages: list[dict]):104 response = openai.ChatCompletion().create(105 model=self.model,106 messages=messages,107 temperature=self.temperature,108 )109 self.token_counter += int(response["usage"]["total_tokens"])110 self.cost += (111 response["usage"]["prompt_tokens"]112 / 1000113 * MODEL_COST_DICT[self.model]["input"]114 + response["usage"]["completion_tokens"]115 / 1000116 * MODEL_COST_DICT[self.model]["output"]117 )118 return response["choices"][0]["message"]["content"]119 120 def reset(self):121 self.token_counter = 0122 self.cost = 0.0123 124 def __call__(self, messages: list[dict]):125 return self.chat(messages)126 127 128class SummaryMemory:129 """130 A class that manages a memory of messages and automatically summarizes them when the maximum token limit is reached.131 132 Attributes:133 max_token_limit (int): The maximum number of tokens allowed in the memory before summarization occurs.134 messages (list[dict]): A list of messages in the memory.135 model (str): The name of the GPT model to use for chat completion.136 ai_role (str): The role of the AI in the conversation.137 human_role (str): The role of the human in the conversation.138 auto_summarize (bool): Whether to automatically summarize the messages when the maximum token limit is reached.139 """140 141 # ...142 summary_template = "Summarize the following messages into a paragraph and replace '{user}' with '{human_role}', and '{assistant}' with '{ai_role}':\n{messages}"143 144 def __init__(145 self,146 system_prompt="",147 max_token_limit=4000,148 model="gpt-3.5-turbo",149 ai_role="answer",150 human_role="question/exam",151 auto_summarize=False,152 ):153 self.max_token_limit = max_token_limit154 self.messages: list[dict] = []155 self.model = model156 self.ai_role = ai_role157 self.human_role = human_role158 self.auto_summarize = auto_summarize159 self.system_prompt = system_prompt160 self.reset()161 162 def reset(self):163 self.messages = [self.system_prompt]164 165 def remove_last(self):166 if len(self.messages) > 1: # don't remove the system prompt167 self.messages.pop()168 169 def remove(170 self, index: int171 ): # don't remove the system prompt and start counting from 1172 if index > 0 and index < len(self.messages):173 self.messages.pop(index)174 175 def replace(self, index: int, message: dict):176 if index > 0 and index < len(self.messages):177 self.messages[index] = message178 179 def change_system_prompt(self, new_prompt: str):180 self.system_prompt = new_prompt181 self.messages[0] = new_prompt182 183 def remove_first(self):184 # dont remove the system prompt185 if len(self.messages) > 1:186 self.messages.pop(1) # remove the first message after the system prompt187 188 def append(self, message: dict):189 total_tokens = num_tokens_from_messages(self.messages + [message])190 191 while (192 self.auto_summarize and total_tokens > self.max_token_limit193 ): # keep summarizing until we're under the limit194 self.summarize()195 total_tokens = num_tokens_from_messages(self.messages + [message])196 197 self.messages.append(message)198 199 def summarize(self):200 prompt = make_prompt(201 self.summary_template,202 user="user",203 human_role=self.human_role,204 assistant="assistant",205 ai_role=self.ai_role,206 messages="\n".join(207 unravel_messages(self.messages[1:])208 ), # don't include the system prompt209 )210 summary = chat(211 messages=[make_message("user", prompt)],212 model=self.model,213 )214 self.reset()215 self.append(make_message("user", summary))216 217 def get_messages(self):218 return self.messages[1:] # don't include the system prompt219 220 def get_unraveled_messages(self):221 return unravel_messages(self.messages[1:])222 223 224class MemoryBuffer:225 """226 A class that manages a buffer of messages and clips them to a maximum token limit.227 228 Attributes:229 max_token_limit (int): The maximum number of tokens allowed in the buffer.230 messages (list[dict]): A list of messages in the buffer.231 """232 233 def __init__(234 self,235 system_prompt,236 max_token_limit=1000,237 ):238 """239 Initializes a new instance of the MemoryBuffer class.240 241 Args:242 max_token_limit (int, optional): The maximum number of tokens allowed in the buffer. Defaults to 1000.243 """244 self.max_token_limit = max_token_limit245 self.messages = []246 self.system_prompt = system_prompt247 self.reset()248 249 def reset(self):250 """251 Resets the buffer by clearing all messages.252 """253 self.messages = [self.system_prompt]254 255 def add(self, message: dict):256 """257 Adds a message to the buffer and clips the buffer to the maximum token limit.258 259 Args:260 message (dict): The message to add to the buffer.261 """262 total_tokens = num_tokens_from_messages(self.messages + [message])263 if total_tokens > self.max_token_limit:264 # clip the messages to the max token limit265 # from the end of the list266 # remove messages from the beginning of the list267 # until the total number of tokens is less than the max token limit268 while total_tokens > self.max_token_limit:269 self.messages = self.messages[1:]270 total_tokens = num_tokens_from_messages(self.messages + [message])271 self.messages.append(message)272 273 def remove(self, message: dict):274 """275 Removes a message from the buffer.276 277 Args:278 message (dict): The message to remove from the buffer.279 """280 if message in self.messages:281 self.messages.remove(message)282 283 def remove_last(self):284 """285 Removes the last message from the buffer.286 """287 if len(self.messages) > 0:288 self.messages.pop()289 290 def remove_first(self):291 """292 Removes the first message from the buffer.293 """294 if len(self.messages) > 0:295 self.messages.pop(0)296 297 298def json2dict(string: str) -> dict:299 """Returns a dictionary of variables from a string containing JSON."""300 try:301 return json.loads(string)302 except json.decoder.JSONDecodeError:303 print("Error: JSONDecodeError")304 return {}305 306 307def print_help(num_nodes, color):308 """309 Prints the help message for the AI assistant.310 """311 colorama.init()312 print(color + "The AI assistant presents a clinical case and asks for a diagnosis.")313 print(314 color + "You need to explore the case by asking questions to the AI assistant."315 )316 print(317 color318 + "You have to ask questions in a logical order, conforming to the clinical guidelines."319 )320 print(321 color322 + "You need to minimize the number of jump between subjects, while covering as many subjects as possible."323 )324 print(color + f"there are a total of {num_nodes} visitable nodes in the tree")325 print(326 color327 + "you have to explore the tree as much as possible while avoiding jumps and travelling excessively."328 )329 print(Style.RESET_ALL)330 331 332def make_question(template=JSON_TEMPLATE, role="user", **kwargs) -> dict:333 prompt = make_prompt(template=template, **kwargs)334 message = make_message(role, prompt)335 return message336 337 338# a debugging decorator and use functools to preserve the function name and docstring339# the decorator gets DEBUG as an argument to turn on or off debugging340def debug(DEBUG, print_func, measure_time=True):341 def decorator(func):342 @ft.wraps(func)343 def wrapper(*args, **kwargs):344 if DEBUG:345 print_func(f"\nCalling {func.__name__}")346 if measure_time and DEBUG:347 start = time.time()348 result = func(*args, **kwargs)349 if measure_time and DEBUG:350 end = time.time()351 print_func(f"Elapsed time: {end - start:.2f}s")352 if DEBUG:353 print_func(f"Returning {func.__name__}")354 return result355 356 return wrapper357 358 return decorator359 360 361# to use the decorator, add @debug(DEBUG) above the function definition362