From 0f5ecafc57f34eecd33cebf17426c06b7bda4f3d Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Sun, 14 Apr 2024 16:51:13 -0400 Subject: [PATCH] fix: api usage issue --- backend/apps/ollama/main.py | 28 ++++++++++++++++++++++++---- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/backend/apps/ollama/main.py b/backend/apps/ollama/main.py index 5e19a8e36..7ed41eb2e 100644 --- a/backend/apps/ollama/main.py +++ b/backend/apps/ollama/main.py @@ -612,7 +612,12 @@ async def generate_embeddings( user=Depends(get_current_user), ): if url_idx == None: - if form_data.model in app.state.MODELS: + model = form_data.model + + if ":" not in model: + model = f"{model}:latest" + + if model in app.state.MODELS: url_idx = random.choice(app.state.MODELS[form_data.model]["urls"]) else: raise HTTPException( @@ -672,7 +677,12 @@ async def generate_completion( ): if url_idx == None: - if form_data.model in app.state.MODELS: + model = form_data.model + + if ":" not in model: + model = f"{model}:latest" + + if model in app.state.MODELS: url_idx = random.choice(app.state.MODELS[form_data.model]["urls"]) else: raise HTTPException( @@ -770,7 +780,12 @@ async def generate_chat_completion( ): if url_idx == None: - if form_data.model in app.state.MODELS: + model = form_data.model + + if ":" not in model: + model = f"{model}:latest" + + if model in app.state.MODELS: url_idx = random.choice(app.state.MODELS[form_data.model]["urls"]) else: raise HTTPException( @@ -874,7 +889,12 @@ async def generate_openai_chat_completion( ): if url_idx == None: - if form_data.model in app.state.MODELS: + model = form_data.model + + if ":" not in model: + model = f"{model}:latest" + + if model in app.state.MODELS: url_idx = random.choice(app.state.MODELS[form_data.model]["urls"]) else: raise HTTPException(