unify token counting

app.py

@@ -305,7 +305,7 @@ def render_response(
         st.info(response_text)
 
     usage_mod.display_api_usage(
-        response["aimessage"].response_metadata, model_info, provider, tag=tag
+        response["aimessage"], model_info, provider, tag=tag
     )
     doc_format_mod.render_retrieved_chunks(response["docs"], tag=tag)
 

sidebar_mod.py

@@ -8,13 +8,18 @@ def render_outreach_links():
     hf_url = "https://huggingface.co/hyperdemocracy"
     pc_url = "https://www.pinecone.io/blog/serverless"
     together_url = "https://www.together.ai/"
+    google_gemini_url = "https://ai.google.dev/gemini-api"
+    anthropic_url = "https://www.anthropic.com/api"
+    openai_url = "https://platform.openai.com/docs/overview"
     st.subheader(":brain: About [hyperdemocracy](https://hyperdemocracy.us)")
     st.subheader(f":world_map: Visualize [nomic atlas]({nomic_url})")
     st.subheader(f":hugging_face: Raw [huggingface datasets]({hf_url})")
     st.subheader(f":evergreen_tree: Index [pinecone serverless]({pc_url})")
     st.subheader(f":pancakes: Inference [together.ai]({together_url})")
-
-
+    st.subheader(f":eyeglasses: Inference [google-gemini]({google_gemini_url})")
+    st.subheader(f":hut: Inference [anthropic]({anthropic_url})")
+    st.subheader(f":sparkles: Inference [openai]({openai_url})")
+    
 def render_sidebar():
     with st.container(border=True):
         render_outreach_links()

usage_mod.py

@@ -1,9 +1,10 @@
 import streamlit as st
+from langchain_core.messages import AIMessage
 
 
-def get_openai_token_usage(response_metadata: dict, model_info: dict):
-    input_tokens = response_metadata["token_usage"]["prompt_tokens"]
-    output_tokens = response_metadata["token_usage"]["completion_tokens"]
+def get_openai_token_usage(aimessage: AIMessage, model_info: dict):
+    input_tokens = aimessage.usage_metadata["input_tokens"]
+    output_tokens = aimessage.usage_metadata["output_tokens"]
     cost = (
         input_tokens * 1e-6 * model_info["cost"]["pmi"]
         + output_tokens * 1e-6 * model_info["cost"]["pmo"]
@@ -15,9 +16,9 @@ def get_openai_token_usage(response_metadata: dict, model_info: dict):
     }
 
 
-def get_anthropic_token_usage(response_metadata: dict, model_info: dict):
-    input_tokens = response_metadata["usage"]["input_tokens"]
-    output_tokens = response_metadata["usage"]["output_tokens"]
+def get_anthropic_token_usage(aimessage: AIMessage, model_info: dict):
+    input_tokens = aimessage.usage_metadata["input_tokens"]
+    output_tokens = aimessage.usage_metadata["output_tokens"]
     cost = (
         input_tokens * 1e-6 * model_info["cost"]["pmi"]
         + output_tokens * 1e-6 * model_info["cost"]["pmo"]
@@ -29,9 +30,9 @@ def get_anthropic_token_usage(response_metadata: dict, model_info: dict):
     }
 
 
-def get_together_token_usage(response_metadata: dict, model_info: dict):
-    input_tokens = response_metadata["token_usage"]["prompt_tokens"]
-    output_tokens = response_metadata["token_usage"]["completion_tokens"]
+def get_together_token_usage(aimessage: AIMessage, model_info: dict):
+    input_tokens = aimessage.usage_metadata["input_tokens"]
+    output_tokens = aimessage.usage_metadata["output_tokens"]
     cost = (
         input_tokens * 1e-6 * model_info["cost"]["pmi"]
         + output_tokens * 1e-6 * model_info["cost"]["pmo"]
@@ -43,9 +44,9 @@ def get_together_token_usage(response_metadata: dict, model_info: dict):
     }
 
 
-def get_google_token_usage(response_metadata: dict, model_info: dict):
-    input_tokens = 0
-    output_tokens = 0
+def get_google_token_usage(aimessage: AIMessage, model_info: dict):
+    input_tokens = aimessage.usage_metadata["input_tokens"]
+    output_tokens = aimessage.usage_metadata["output_tokens"]
     cost = (
         input_tokens * 1e-6 * model_info["cost"]["pmi"]
         + output_tokens * 1e-6 * model_info["cost"]["pmo"]
@@ -57,29 +58,29 @@ def get_google_token_usage(response_metadata: dict, model_info: dict):
     }
 
 
-def get_token_usage(response_metadata: dict, model_info: dict, provider: str):
+def get_token_usage(aimessage: AIMessage, model_info: dict, provider: str):
     match provider:
         case "OpenAI":
-            return get_openai_token_usage(response_metadata, model_info)
+            return get_openai_token_usage(aimessage, model_info)
         case "Anthropic":
-            return get_anthropic_token_usage(response_metadata, model_info)
+            return get_anthropic_token_usage(aimessage, model_info)
         case "Together":
-            return get_together_token_usage(response_metadata, model_info)
+            return get_together_token_usage(aimessage, model_info)
         case "Google":
-            return get_google_token_usage(response_metadata, model_info)
+            return get_google_token_usage(aimessage, model_info)
         case _:
             raise ValueError()
 
 
 def display_api_usage(
-    response_metadata: dict, model_info: dict, provider: str, tag: str | None = None
+    aimessage: AIMessage, model_info: dict, provider: str, tag: str | None = None
 ):
     with st.container(border=True):
         if tag is None:
             st.write("API Usage")
         else:
             st.write(f"API Usage ({tag})")
-        token_usage = get_token_usage(response_metadata, model_info, provider)
+        token_usage = get_token_usage(aimessage, model_info, provider)
         col1, col2, col3 = st.columns(3)
         with col1:
             st.metric("Input Tokens", token_usage["input_tokens"])
@@ -87,5 +88,6 @@ def display_api_usage(
             st.metric("Output Tokens", token_usage["output_tokens"])
         with col3:
             st.metric("Cost", f"${token_usage['cost']:.4f}")
-        with st.expander("Response Metadata"):
-            st.warning(response_metadata)
+        with st.expander("AIMessage Metadata"):
+            dd = {key: val for key, val in aimessage.dict().items() if key != "content"}
+            st.write(dd)