Skip to content

GitLab

Commit 0f5ecafc authored by Timothy J. Baek's avatar Timothy J. Baek
Browse files

fix: api usage issue

parent 98c16776
Hide whitespace changes
Inline Side-by-side
Showing with 24 additions and 4 deletions
backend/apps/ollama/main.py
View file @ 0f5ecafc
...@@ -612,7 +612,12 @@ async def generate_embeddings( ...@@ -612,7 +612,12 @@ async def generate_embeddings(
user=Depends(get_current_user), user=Depends(get_current_user),
): ):
if url_idx == None: 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"]) url_idx = random.choice(app.state.MODELS[form_data.model]["urls"])
else: else:
raise HTTPException( raise HTTPException(
...@@ -672,7 +677,12 @@ async def generate_completion( ...@@ -672,7 +677,12 @@ async def generate_completion(
): ):
if url_idx == None: 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"]) url_idx = random.choice(app.state.MODELS[form_data.model]["urls"])
else: else:
raise HTTPException( raise HTTPException(
...@@ -770,7 +780,12 @@ async def generate_chat_completion( ...@@ -770,7 +780,12 @@ async def generate_chat_completion(
): ):
if url_idx == None: 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"]) url_idx = random.choice(app.state.MODELS[form_data.model]["urls"])
else: else:
raise HTTPException( raise HTTPException(
...@@ -874,7 +889,12 @@ async def generate_openai_chat_completion( ...@@ -874,7 +889,12 @@ async def generate_openai_chat_completion(
): ):
if url_idx == None: 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"]) url_idx = random.choice(app.state.MODELS[form_data.model]["urls"])
else: else:
raise HTTPException( raise HTTPException(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment