Refactor LLM selection to use specific models for reasoning, parsing, and balanced tasks; improve code clarity.

src/gaia_solving_agent/agent.py

@@ -15,11 +15,12 @@ from gaia_solving_agent.prompts import PLANING_PROMPT, FORMAT_ANSWER
 from gaia_solving_agent.tools import tavily_search_web, wikipedia_tool_spec
 
 # Choice of the model
-model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-# model_name = "deepseek-ai/DeepSeek-R1-0528"
-# model_name = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"  # For VLM needs
+cheap_model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+balanced_model_name = "meta-llama/Meta-Llama-3.1-70B-Instruct"
+reasoning_model_name = "deepseek-ai/DeepSeek-R1-0528"
+vlm_model_name = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"  # For VLM needs
 
-def get_llm(model_name=model_name):
+def get_llm(model_name=cheap_model_name):
     return NebiusLLM(
         model=model_name,
         api_key=NEBIUS_API_KEY,
@@ -45,7 +46,7 @@ class AnswerEvent(Event):
 class GaiaWorkflow(Workflow):
     @step
     async def setup(self, ev: StartEvent) -> QueryEvent:
-        llm = get_llm()
+        llm = get_llm(reasoning_model_name)
         prompt_template = RichPromptTemplate(PLANING_PROMPT)
         file_extension = Path(ev.additional_file_path).suffix if ev.additional_file_path else ""
         plan = llm.complete(prompt_template.format(
@@ -75,7 +76,7 @@ class GaiaWorkflow(Workflow):
 
     @step
     async def parse_answer(self, ev: AnswerEvent) -> StopEvent:
-        llm = get_llm()
+        llm = get_llm(balanced_model_name)
         prompt_template = RichPromptTemplate(FORMAT_ANSWER)
         pattern = r"<Question> :\s*(.*)[\n$]"
         search = re.search(pattern, ev.plan)
@@ -145,7 +146,7 @@ gaia_solving_agent = FunctionAgent(
         *simple_web_page_reader_toolspec.to_tool_list(),
         *RequestsToolSpec().to_tool_list(),
     ],
-    llm=get_llm(),
+    llm=get_llm(balanced_model_name),
     system_prompt="""
     You are a helpful assistant that uses tools to browse additional information and resources on the web to answer questions.
     """,