Add token usage tracking for openai and fix token usage tracking for anthropic.

app.py

@@ -61,6 +61,19 @@ def anthropic_streamlit_streamer(stream, llm):
     """
     for event in stream:
         if hasattr(event, "type"):
+            # Count input token usage.
+            if event.type == "message_start":
+                st.session_state["input_token_usage"][
+                    llm
+                ] += event.message.usage.input_tokens
+                st.session_state["output_token_usage"][
+                    llm
+                ] += event.message.usage.output_tokens
+
+            # Count output token usage.
+            if event.type == "message_delta":
+                st.session_state["output_token_usage"][llm] += event.usage.output_tokens
+
             # Handle content blocks
             if event.type == "content_block_delta" and hasattr(event, "delta"):
                 # Extract text delta from the event
@@ -68,18 +81,6 @@ def anthropic_streamlit_streamer(stream, llm):
                 if text_delta:
                     yield text_delta
 
-            # Count input token usage.
-            if event.type == "message_start":
-                input_token_usage = event["usage"]["input_tokens"]
-                output_token_usage = event["usage"]["output_tokens"]
-                st.session_state["input_token_usage"][llm] += input_token_usage
-                st.session_state["output_token_usage"][llm] += output_token_usage
-
-            # Count output token usage.
-            if event.type == "message_delta":
-                output_token_usage = event["usage"]["output_tokens"]
-                st.session_state["output_token_usage"][llm] += output_token_usage
-
             # Handle message completion events (optional if needed)
             elif event.type == "message_stop":
                 break  # End of message, stop streaming
@@ -101,6 +102,17 @@ def google_streamlit_streamer(stream):
         yield chunk.text
 
 
+def openai_streamlit_streamer(stream, llm):
+    # https://platform.openai.com/docs/api-reference/streaming
+    for event in stream:
+        if event.usage:
+            st.session_state["input_token_usage"][llm] += event.usage.prompt_tokens
+            st.session_state["output_token_usage"][llm] += event.usage.completion_tokens
+        if event.choices:
+            if event.choices[0].delta.content:
+                yield event.choices[0].delta.content
+
+
 def together_streamlit_streamer(stream, llm):
     # https://docs.together.ai/docs/chat-overview#streaming-responses
     for chunk in stream:
@@ -111,21 +123,6 @@ def together_streamlit_streamer(stream, llm):
         yield chunk.choices[0].delta.content
 
 
-def llm_streamlit_streamer(stream, llm):
-    if llm.startswith("anthropic"):
-        print(f"Using Anthropic streamer for llm: {llm}")
-        return anthropic_streamlit_streamer(stream, llm)
-    elif llm.startswith("vertex"):
-        print(f"Using Vertex streamer for llm: {llm}")
-        return google_streamlit_streamer(stream)
-    elif llm.startswith("together"):
-        print(f"Using Together streamer for llm: {llm}")
-        return together_streamlit_streamer(stream, llm)
-    else:
-        print(f"Using OpenAI streamer for llm: {llm}")
-        return openai_streamlit_streamer(stream, llm)
-
-
 # Helper functions for LLM council and aggregator selection
 def llm_council_selector():
     selected_council = st.radio(
@@ -144,6 +141,7 @@ def get_openai_response(model_name, prompt):
         model=model_name,
         messages=[{"role": "user", "content": prompt}],
         stream=True,
+        stream_options={"include_usage": True},
     )
 
 
@@ -175,7 +173,9 @@ def get_llm_response_stream(model_identifier, prompt):
     """Returns a streamlit-friendly stream of response tokens from the LLM."""
     provider, model_name = model_identifier.split("://")
     if provider == "openai":
-        return get_openai_response(model_name, prompt)
+        return openai_streamlit_streamer(
+            get_openai_response(model_name, prompt), model_identifier
+        )
     elif provider == "anthropic":
         return anthropic_streamlit_streamer(
             get_anthropic_response(model_name, prompt), model_identifier
@@ -360,6 +360,13 @@ def parse_judging_responses(
         ],
         response_format=DirectAssessmentJudgingResponse,
     )
+    # Track token usage.
+    st.session_state["input_token_usage"][
+        "gpt-4o-mini"
+    ] += completion.usage.prompt_tokens
+    st.session_state["output_token_usage"][
+        "gpt-4o-mini"
+    ] += completion.usage.completion_tokens
     return completion.choices[0].message.parsed
 
 
@@ -582,12 +589,12 @@ def st_render_responses(user_prompt):
         else:
             message_placeholder = st.empty()
             aggregator_stream = get_llm_response_stream(
-                selected_aggregator, aggregator_prompt
+                st.session_state.selected_aggregator, aggregator_prompt
             )
             if aggregator_stream:
-                st.session_state.responses[get_aggregator_key(selected_aggregator)] = (
-                    message_placeholder.write_stream(aggregator_stream)
-                )
+                st.session_state.responses[
+                    get_aggregator_key(st.session_state.selected_aggregator)
+                ] = message_placeholder.write_stream(aggregator_stream)
 
     st.session_state.responses_collected = True
 
@@ -825,8 +832,6 @@ def main():
                 )
 
         if submit_button:
-            st.markdown("#### Responses")
-
             # Udpate state.
             st.session_state.selected_models = selected_models
             st.session_state.selected_aggregator = selected_aggregator

constants.py

@@ -9,7 +9,8 @@ if os.getenv("DEBUG_MODE") == "True":
         "Smalls": [
             "together://meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
             "together://meta-llama/Llama-3.2-3B-Instruct-Turbo",
-            # "anthropic://claude-3-haiku-20240307",
+            "anthropic://claude-3-haiku-20240307",
+            "openai://gpt-4o-mini",
         ],
         "Flagships": [
             "together://meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",