Update final_answer.py

final_answer.py

@@ -44,10 +44,13 @@ def check_reasoning(final_answer, agent_memory):
 
 def ensure_formatting(final_answer, agent_memory):
     # Ensure the final answer is formatted correctly
-    model_name = 'granite3.3:8b'
+    model_name = 'cogito:14b'
     # Initialize the chat model
-    model = LiteLLMModel(model_id=f'ollama_chat/{model_name}',
-                             flatten_messages_as_text=True)
+    multimodal_model = LiteLLMModel(
+            model_id=model_name,
+            provider="ollama",  # ADD THIS
+            api_key="ollama"    # ADD THIS (ollama doesn't need real key but LiteLLM might expect it)
+        )
     prompt = f"""
         Here is a user-given task and the agent steps: {agent_memory.get_succinct_steps()}. Now here is the FINAL ANSWER that was given: 
         {final_answer}