Lucas-C-R commited on
Commit
3f77561
·
1 Parent(s): 19ef247

refactor: simplify LLM model for retriever agent

Browse files
Files changed (1) hide show
  1. services/agent_services.py +11 -7
services/agent_services.py CHANGED
@@ -30,12 +30,16 @@ def load_prompt(name: str) -> str:
30
  return prompts[name]
31
 
32
 
33
- def create_llm(model: Literal["groq", "openai"] = "openai") -> BaseChatModel:
34
- return (
35
- ChatGroq(model="qwen-qwq-32b", temperature=0)
36
- if model == "groq"
37
- else ChatOpenAI(model="gpt-4.1", temperature=0)
38
- )
 
 
 
 
39
 
40
 
41
  def create_agent(
@@ -78,7 +82,7 @@ def create_workflow() -> Callable:
78
  llm = create_llm()
79
 
80
  retriever_agent = create_agent(
81
- llm=llm,
82
  tools=[retriever_tool],
83
  prompt_name="retriever_prompt",
84
  name="retriever_agent",
 
30
  return prompts[name]
31
 
32
 
33
+ def create_llm(
34
+ model: Literal["groq", "openai", "openai-nano"] = "openai",
35
+ ) -> BaseChatModel:
36
+ match (model):
37
+ case "groq":
38
+ return ChatGroq(model="qwen-qwq-32b", temperature=0)
39
+ case "openai":
40
+ return ChatOpenAI(model="gpt-4.1", temperature=0)
41
+ case "openai-nano":
42
+ return ChatOpenAI(model="gpt-4.1-nano", temperature=0)
43
 
44
 
45
  def create_agent(
 
82
  llm = create_llm()
83
 
84
  retriever_agent = create_agent(
85
+ llm=create_llm("openai-nano"),
86
  tools=[retriever_tool],
87
  prompt_name="retriever_prompt",
88
  name="retriever_agent",