add count costs

ai_talks/chat.py

@@ -61,6 +61,10 @@ if "input_kind" not in st.session_state:
     st.session_state.input_kind = st.session_state.locale.input_kind_1
 if "seed" not in st.session_state:
     st.session_state.seed = randrange(10**3)  # noqa: S311
+if "costs" not in st.session_state:
+    st.session_state.costs = []
+if "total_tokens" not in st.session_state:
+    st.session_state.total_tokens = []
 
 
 def main() -> None:

ai_talks/src/utils/conversation.py

@@ -11,9 +11,6 @@ from .agi.chat_gpt import create_gpt_completion
 from .stt import show_voice_input
 from .tts import show_audio_player
 
-phind.cf_clearance = st.secrets.api_credentials.phind_cf_clearance
-phind.user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/112.0.0.0 Safari/537.36"  # noqa: E501
-
 
 def clear_chat() -> None:
     st.session_state.generated = []
@@ -21,6 +18,8 @@ def clear_chat() -> None:
     st.session_state.messages = []
     st.session_state.user_text = ""
     st.session_state.seed = randrange(10**8)  # noqa: S311
+    st.session_state.costs = []
+    st.session_state.total_tokens = []
 
 
 def show_text_input() -> None:
@@ -60,12 +59,30 @@ def show_chat(ai_content: str, user_text: str) -> None:
             message(st.session_state.past[i], is_user=True, key=str(i) + "_user", seed=st.session_state.seed)
             message("", key=str(i), seed=st.session_state.seed)
             st.markdown(st.session_state.generated[i])
+            st.caption(f"""
+                {st.session_state.locale.tokens_count}{st.session_state.total_tokens[i]} |
+                {st.session_state.locale.message_cost}{st.session_state.costs[i]:.5f}$
+            """, help=f"{st.session_state.locale.total_cost}{sum(st.session_state.costs):.5f}$")
+
+
+def calc_cost(usage: dict) -> None:
+    total_tokens = usage.get("total_tokens")
+    prompt_tokens = usage.get("prompt_tokens")
+    completion_tokens = usage.get("completion_tokens")
+    st.session_state.total_tokens.append(total_tokens)
+    # pricing logic: https://openai.com/pricing#language-models
+    if st.session_state.model == "gpt-3.5-turbo":
+        cost = total_tokens * 0.002 / 1000
+    else:
+        cost = (prompt_tokens * 0.03 + completion_tokens * 0.06) / 1000
+    st.session_state.costs.append(cost)
 
 
 def show_gpt_conversation() -> None:
     try:
         completion = create_gpt_completion(st.session_state.model, st.session_state.messages)
         ai_content = completion.get("choices")[0].get("message").get("content")
+        calc_cost(completion.get("usage"))
         st.session_state.messages.append({"role": "assistant", "content": ai_content})
         if ai_content:
             show_chat(ai_content, st.session_state.user_text)
@@ -93,6 +110,9 @@ def show_bard_conversation() -> None:
 
 
 def phind_get_answer(question: str):
+    phind.cf_clearance = st.secrets.api_credentials.phind_cf_clearance
+    phind.user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/112.0.0.0 Safari/537.36"  # noqa: E501
+
     try:
         result = phind.Completion.create(
             model="gpt-4",

ai_talks/src/utils/lang.py

@@ -38,6 +38,9 @@ class Locale:
     footer_channel: str
     responsibility_denial: str
     donates_info: str
+    tokens_count: str
+    message_cost: str
+    total_cost: str
     empty_api_handler: str
 
 
@@ -84,6 +87,9 @@ en = Locale(
         This allows you to provide access to communication with AI for all users.
         Support us for joint development and interaction with the intelligence of the future!
     """,
+    tokens_count="Tokens count: ",
+    message_cost="Message cost: ",
+    total_cost="Total cost of conversation: ",
     empty_api_handler=f"""
         API key not found. Create `.streamlit/secrets.toml` with your API key.
         See [README.md]({README_URL}) for instructions or use the original [AI Talks]({AI_TALKS_URL}).
@@ -133,6 +139,9 @@ ru = Locale(
         Это позволяет обеспечить доступ к общению с ИИ для всех желающих пользователей.
         Поддержите нас для совместного развития и взаимодействия с интеллектом будущего!
     """,
+    tokens_count="Количество токенов: ",
+    message_cost="Cтоимость сообщения: ",
+    total_cost="Общая стоимость разговора: ",
     empty_api_handler=f"""
         Ключ API не найден. Создайте `.streamlit/secrets.toml` с вашим ключом API.
         Инструкции см. в [README.md]({README_URL}) или используйте оригинальный [AI Talks]({AI_TALKS_URL}).