Queer European MD passionate about IT

Mission903Solutions.py 6.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167
  1. from openai import OpenAI
  2. import tiktoken
  3. import json
  4. from datetime import datetime
  5. import os
  6. import streamlit as st
  7. DEFAULT_API_KEY = os.environ.get("TOGETHER_API_KEY")
  8. DEFAULT_BASE_URL = "https://api.together.xyz/v1"
  9. DEFAULT_MODEL = "meta-llama/Llama-3-8b-chat-hf"
  10. DEFAULT_TEMPERATURE = 0.7
  11. DEFAULT_MAX_TOKENS = 512
  12. DEFAULT_TOKEN_BUDGET = 4096
  13. class ConversationManager:
  14. def __init__(self, api_key=None, base_url=None, model=None, history_file=None, temperature=None, max_tokens=None, token_budget=None):
  15. if not api_key:
  16. api_key = DEFAULT_API_KEY
  17. if not base_url:
  18. base_url = DEFAULT_BASE_URL
  19. self.client = OpenAI(
  20. api_key=api_key,
  21. base_url=base_url
  22. )
  23. self.model = model if model else DEFAULT_MODEL
  24. self.temperature = temperature if temperature else DEFAULT_TEMPERATURE
  25. self.max_tokens = max_tokens if max_tokens else DEFAULT_MAX_TOKENS
  26. self.token_budget = token_budget if token_budget else DEFAULT_TOKEN_BUDGET
  27. self.system_messages = {
  28. "sassy_assistant": "You are a sassy assistant that is fed up with answering questions.",
  29. "angry_assistant": "You are an angry assistant that likes yelling in all caps.",
  30. "thoughtful_assistant": "You are a thoughtful assistant, always ready to dig deeper. You ask clarifying questions to ensure understanding and approach problems with a step-by-step methodology.",
  31. "custom": "Enter your custom system message here."
  32. }
  33. self.system_message = self.system_messages["sassy_assistant"] # Default persona
  34. self.conversation_history = [{"role": "system", "content": self.system_message}]
  35. def count_tokens(self, text):
  36. try:
  37. encoding = tiktoken.encoding_for_model(self.model)
  38. except KeyError:
  39. encoding = tiktoken.get_encoding("cl100k_base")
  40. tokens = encoding.encode(text)
  41. return len(tokens)
  42. def total_tokens_used(self):
  43. try:
  44. return sum(self.count_tokens(message['content']) for message in self.conversation_history)
  45. except Exception as e:
  46. print(f"An unexpected error occurred while calculating the total tokens used: {e}")
  47. return None
  48. def enforce_token_budget(self):
  49. try:
  50. while self.total_tokens_used() > self.token_budget:
  51. if len(self.conversation_history) <= 1:
  52. break
  53. self.conversation_history.pop(1)
  54. except Exception as e:
  55. print(f"An unexpected error occurred while enforcing the token budget: {e}")
  56. def set_persona(self, persona):
  57. if persona in self.system_messages:
  58. self.system_message = self.system_messages[persona]
  59. self.update_system_message_in_history()
  60. else:
  61. raise ValueError(f"Unknown persona: {persona}. Available personas are: {list(self.system_messages.keys())}")
  62. def set_custom_system_message(self, custom_message):
  63. if not custom_message:
  64. raise ValueError("Custom message cannot be empty.")
  65. self.system_messages['custom'] = custom_message
  66. self.set_persona('custom')
  67. def update_system_message_in_history(self):
  68. try:
  69. if self.conversation_history and self.conversation_history[0]["role"] == "system":
  70. self.conversation_history[0]["content"] = self.system_message
  71. else:
  72. self.conversation_history.insert(0, {"role": "system", "content": self.system_message})
  73. except Exception as e:
  74. print(f"An unexpected error occurred while updating the system message in the conversation history: {e}")
  75. def chat_completion(self, prompt, temperature=None, max_tokens=None, model=None):
  76. temperature = temperature if temperature is not None else self.temperature
  77. max_tokens = max_tokens if max_tokens is not None else self.max_tokens
  78. model = model if model is not None else self.model
  79. self.conversation_history.append({"role": "user", "content": prompt})
  80. self.enforce_token_budget()
  81. try:
  82. response = self.client.chat.completions.create(
  83. model=model,
  84. messages=self.conversation_history,
  85. temperature=temperature,
  86. max_tokens=max_tokens,
  87. )
  88. except Exception as e:
  89. print(f"An error occurred while generating a response: {e}")
  90. return None
  91. ai_response = response.choices[0].message.content
  92. self.conversation_history.append({"role": "assistant", "content": ai_response})
  93. return ai_response
  94. def reset_conversation_history(self):
  95. self.conversation_history = [{"role": "system", "content": self.system_message}]
  96. ### Streamlit code ###
  97. st.title("Sassy Chatbot :face_with_rolling_eyes:")
  98. # Sidebar
  99. st.sidebar.header("Options")
  100. # Initialize the ConversationManager object
  101. if 'chat_manager' not in st.session_state:
  102. st.session_state['chat_manager'] = ConversationManager()
  103. chat_manager = st.session_state['chat_manager']
  104. # Set the token budget, max tokens per message, and temperature with sliders
  105. max_tokens_per_message = st.sidebar.slider("Max Tokens Per Message", min_value=10, max_value=500, value=50)
  106. temperature = st.sidebar.slider("Temperature", min_value=0.0, max_value=1.0, value=0.7, step=0.01)
  107. # Select and set system message with a selectbox
  108. system_message = st.sidebar.selectbox("System message", ['Sassy', 'Angry', 'Thoughtful', 'Custom'])
  109. if system_message == 'Sassy':
  110. chat_manager.set_persona('sassy_assistant')
  111. elif system_message == 'Angry':
  112. chat_manager.set_persona('angry_assistant')
  113. elif system_message == 'Thoughtful':
  114. chat_manager.set_persona('thoughtful_assistant')
  115. # Open text area for custom system message if "Custom" is selected
  116. elif system_message == 'Custom':
  117. custom_message = st.sidebar.text_area("Custom system message")
  118. if st.sidebar.button("Set custom system message"):
  119. chat_manager.set_custom_system_message(custom_message)
  120. if st.sidebar.button("Reset conversation history", on_click=chat_manager.reset_conversation_history):
  121. st.session_state['conversation_history'] = chat_manager.conversation_history
  122. if 'conversation_history' not in st.session_state:
  123. st.session_state['conversation_history'] = chat_manager.conversation_history
  124. conversation_history = st.session_state['conversation_history']
  125. # Chat input from the user
  126. user_input = st.chat_input("Write a message")
  127. # Call the chat manager to get a response from the AI. Uses settings from the sidebar.
  128. if user_input:
  129. response = chat_manager.chat_completion(user_input, temperature=temperature, max_tokens=max_tokens_per_message)
  130. # Display the conversation history
  131. for message in conversation_history:
  132. if message["role"] != "system":
  133. with st.chat_message(message["role"]):
  134. st.write(message["content"])