diff --git a/.changeset/silver-masks-speak.md b/.changeset/silver-masks-speak.md new file mode 100644 index 0000000000000000000000000000000000000000..704befc9a67d7c620933e16b0514f7d3d55c3b69 --- /dev/null +++ b/.changeset/silver-masks-speak.md @@ -0,0 +1,5 @@ +--- +"create-llama": patch +--- + +Remove hard-coded Gemini supported models diff --git a/templates/types/streaming/fastapi/app/settings.py b/templates/types/streaming/fastapi/app/settings.py index 3f138b10085ea2471874574f626f2566232661e8..87c591af46e5eacf5360fe98c3eb2fa90f95ccb3 100644 --- a/templates/types/streaming/fastapi/app/settings.py +++ b/templates/types/streaming/fastapi/app/settings.py @@ -115,18 +115,8 @@ def init_gemini(): from llama_index.llms.gemini import Gemini from llama_index.embeddings.gemini import GeminiEmbedding - model_map: Dict[str, str] = { - "gemini-1.5-pro-latest": "models/gemini-1.5-pro-latest", - "gemini-pro": "models/gemini-pro", - "gemini-pro-vision": "models/gemini-pro-vision", - } + model_name = f"models/{os.getenv('MODEL')}" + embed_model_name = f"models/{os.getenv('EMBEDDING_MODEL')}" - embed_model_map: Dict[str, str] = { - "embedding-001": "models/embedding-001", - "text-embedding-004": "models/text-embedding-004", - } - - Settings.llm = Gemini(model=model_map[os.getenv("MODEL")]) - Settings.embed_model = GeminiEmbedding( - model_name=embed_model_map[os.getenv("EMBEDDING_MODEL")] - ) + Settings.llm = Gemini(model=model_name) + Settings.embed_model = GeminiEmbedding(model_name=embed_model_name)