updated token counts and context windows explicitly for each model

app.py

@@ -11,8 +11,6 @@ from llama_index.core import (
 from llama_parse import LlamaParse
 from streamlit_pdf_viewer import pdf_viewer
 
-MAX_OUTPUT_TOKENS = 2048  
-
 class MistralTokens:
     """
     Returns tokens for MistralAI models.
@@ -71,11 +69,11 @@ def main():
 
         # Select LLM
         if provider == 'google':
-            llm_list = ['gemini-1.0-pro', 'gemini-1.5-flash', 'gemini-1.5-pro', 'aqa']
+            llm_list = ['gemini-1.0-pro', 'gemini-1.5-flash', 'gemini-1.5-pro']
         elif provider == 'huggingface':
             llm_list = []
         elif provider == 'mistralai':
-            llm_list = ["mistral-small-latest", "mistral-large-latest", "open-mistral-nemo-latest"]
+            llm_list = ["mistral-large-latest", "open-mistral-nemo-latest"]
         elif provider == 'openai':
             llm_list = ['gpt-3.5-turbo', 'gpt-4', 'gpt-4-turbo', 'gpt-4o', 'gpt-4o-mini']
         else:
@@ -127,75 +125,100 @@ def main():
             if provider == 'google':
                 from llama_index.llms.gemini import Gemini
                 from llama_index.embeddings.gemini import GeminiEmbedding
-
+                max_output_tokens = 8192 # https://firebase.google.com/docs/vertex-ai/gemini-models
+                
                 os.environ['GOOGLE_API_KEY'] = str(llm_key)
                 Settings.llm = Gemini(
                     model=f"models/{llm_name}",
                     token=os.environ.get("GOOGLE_API_KEY"),
                     temperature=temperature,
-                    max_tokens=MAX_OUTPUT_TOKENS
+                    max_tokens=max_output_tokens
                 )
                 Settings.tokenizer = GeminiTokens(llm_name) 
-                Settings.num_output = MAX_OUTPUT_TOKENS
+                Settings.num_output = max_output_tokens
                 Settings.embed_model = GeminiEmbedding(
                     model_name="models/text-embedding-004", api_key=os.environ.get("GOOGLE_API_KEY") #, title="this is a document"
                 )
-                Settings.context_window = 4096 
+                if llm_name == 'gemini-1.0-pro':
+                    total_token_limit = 32760
+                else:
+                    total_token_limit = 1e6
+                Settings.context_window = total_token_limit -  max_output_tokens # Gemini counts total tokens
             elif provider == 'huggingface':
                 if llm_name is not None and embed_name is not None:
                     from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI 
                     from llama_index.embeddings.huggingface import HuggingFaceInferenceAPIEmbedding
                     from transformers import AutoTokenizer
 
+                    max_output_tokens = 2048 # Just a generic value
+
                     os.environ['HFTOKEN'] = str(llm_key)
                     Settings.llm = HuggingFaceInferenceAPI(
                         model_name=llm_name, 
                         token=os.environ.get("HFTOKEN"),
                         temperature=temperature,
-                        max_tokens=MAX_OUTPUT_TOKENS
+                        max_tokens=max_output_tokens
                     )
                     Settings.tokenizer = AutoTokenizer.from_pretrained(
                         llm_name,
                         token=os.environ.get("HFTOKEN"),
                     )
-                    Settings.num_output = MAX_OUTPUT_TOKENS
+                    Settings.num_output = max_output_tokens
                     Settings.embed_model = HuggingFaceInferenceAPIEmbedding(
                         model_name=embed_name
                     )
-                    Settings.context_window = 4096 
+                    Settings.context_window = 4096 # Just a generic value
             elif provider == 'mistralai':
                 from llama_index.llms.mistralai import MistralAI
                 from llama_index.embeddings.mistralai import MistralAIEmbedding
+                max_output_tokens = 8192 # Based on internet consensus since this is not well documented
                 
                 os.environ['MISTRAL_API_KEY'] = str(llm_key)
                 Settings.llm = MistralAI(
                     model=llm_name, 
                     temperature=temperature,
-                    max_tokens=MAX_OUTPUT_TOKENS,
+                    max_tokens=max_output_tokens,
                     random_seed=42,
                     safe_mode=True
                 )
                 Settings.tokenizer = MistralTokens(llm_name)
-                Settings.num_output = MAX_OUTPUT_TOKENS
+                Settings.num_output = max_output_tokens
                 Settings.embed_model = MistralAIEmbedding(
                     model_name="mistral-embed", 
                     api_key=os.environ.get("MISTRAL_API_KEY")
                 )
-                Settings.context_window = 4096 # max possible
+                Settings.context_window = 128000 # 128k for flagship models - doesn't seem to count input tokens
             elif provider == 'openai':
                 from llama_index.llms.openai import OpenAI
                 from llama_index.embeddings.openai import OpenAIEmbedding
 
+                # https://platform.openai.com/docs/models/gpt-4-turbo-and-gpt-4
+                if llm_name == 'gpt-3.5-turbo':
+                    max_output_tokens = 4096
+                    context_window = 16385
+                elif llm_name == 'gpt-4' :
+                    max_output_tokens = 8192
+                    context_window = 8192
+                elif llm_name == 'gpt-4-turbo'
+                    max_output_tokens = 4096
+                    context_window = 128000
+                elif llm_name == 'gpt-4o':
+                    max_output_tokens = 4096
+                    context_window = 128000
+                elif llm_name == 'gpt-4o-mini':
+                    max_output_tokens = 16384
+                    context_window = 128000
+
                 os.environ["OPENAI_API_KEY"] = str(llm_key)
                 Settings.llm = OpenAI(
                     model=llm_name, 
                     temperature=temperature,
-                    max_tokens=MAX_OUTPUT_TOKENS
+                    max_tokens=max_output_tokens
                 )
                 Settings.tokenizer = tiktoken.encoding_for_model(llm_name).encode
-                Settings.num_output = MAX_OUTPUT_TOKENS
+                Settings.num_output = max_output_tokens
                 Settings.embed_model = OpenAIEmbedding()
-                Settings.context_window = 4096 # max possible
+                Settings.context_window = context_window 
             else:
                 raise NotImplementedError(f"{provider} is not supported yet")