@@ -52,8 +52,8 @@ async def get_next_action(model, messages, objective, session_id):
52
52
return "coming soon"
53
53
if model == "gemini-pro-vision" :
54
54
return call_gemini_pro_vision (messages , objective ), None
55
- if model == "llava" :
56
- operation = call_ollama_llava (messages )
55
+ if model == "llava" or model == "llava:13b" or "bakllava" or "llava-llama3" :
56
+ operation = call_ollama_llava (messages , model )
57
57
return operation , None
58
58
if model == "claude-3" :
59
59
operation = await call_claude_3_with_ocr (messages , objective , model )
@@ -453,9 +453,11 @@ async def call_gpt_4_vision_preview_labeled(messages, objective, model):
453
453
return call_gpt_4_vision_preview (messages )
454
454
455
455
456
- def call_ollama_llava (messages ):
456
+ def call_ollama_llava (messages , model ):
457
+ if model == "" :
458
+ model = "llava"
457
459
if config .verbose :
458
- print ("[call_ollama_llava]" )
460
+ print (f "[call_ollama_llava] model { model } " )
459
461
time .sleep (1 )
460
462
try :
461
463
screenshots_dir = "screenshots"
@@ -485,7 +487,7 @@ def call_ollama_llava(messages):
485
487
messages .append (vision_message )
486
488
487
489
response = ollama .chat (
488
- model = "llava" ,
490
+ model = model ,
489
491
messages = messages ,
490
492
)
491
493
@@ -527,7 +529,7 @@ def call_ollama_llava(messages):
527
529
)
528
530
if config .verbose :
529
531
traceback .print_exc ()
530
- return call_ollama_llava (messages )
532
+ return call_ollama_llava (messages , model )
531
533
532
534
533
535
async def call_claude_3_with_ocr (messages , objective , model ):
0 commit comments