Selaa lähdekoodia

fix: api usage issue

Timothy J. Baek 1 vuosi sitten
vanhempi
commit
0f5ecafc57
1 muutettua tiedostoa jossa 24 lisäystä ja 4 poistoa
  1. 24 4
      backend/apps/ollama/main.py

+ 24 - 4
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(