AlexTrinityBlock commited on
Commit
bcdc55d
·
1 Parent(s): c36efba

refactor(api): add LLM abstraction layer to replace hardcoded model strings

Browse files

Replace hardcoded "google_genai:gemini-3-flash-preview" model references across
all agent modules (supervisor, websearch, answer extractor, ocr reader) with
configurable get_llm() function from agent.api.api module

.env.example CHANGED
@@ -1,2 +1,4 @@
1
  GOOGLE_API_KEY=your-google-api-key-here
2
  TAVILY_API_KEY=your-tavily-api-key-here
 
 
 
1
  GOOGLE_API_KEY=your-google-api-key-here
2
  TAVILY_API_KEY=your-tavily-api-key-here
3
+ MINIMAX_API_KEY=your-minimax-api-key-here
4
+ MINIMAX_GROUP_ID=your-minimax-group-id-here
agent/agent.py CHANGED
@@ -3,6 +3,7 @@ from dotenv import load_dotenv
3
  from colorama import Fore, Style # type: ignore[import]
4
  from langchain.agents import create_agent
5
  from langchain_core.messages import HumanMessage
 
6
  from agent.tools.math_solver import math_solver
7
  from agent.tools.file_downloader import file_downloader
8
  from agent.tools.ocr_reader import ocr_reader
@@ -20,7 +21,7 @@ load_dotenv()
20
  def supervisor_agent():
21
  """Return a supervisor agent instance with math_solver and websearch_agent."""
22
  return create_agent(
23
- model="google_genai:gemini-3-flash-preview",
24
  # tools=[math_solver, websearch_agent, web_search_agents],
25
  tools=[math_solver, web_search_agents, file_downloader, ocr_reader, list_files, http_get],
26
  system_prompt=(
 
3
  from colorama import Fore, Style # type: ignore[import]
4
  from langchain.agents import create_agent
5
  from langchain_core.messages import HumanMessage
6
+ from agent.api.api import get_llm
7
  from agent.tools.math_solver import math_solver
8
  from agent.tools.file_downloader import file_downloader
9
  from agent.tools.ocr_reader import ocr_reader
 
21
  def supervisor_agent():
22
  """Return a supervisor agent instance with math_solver and websearch_agent."""
23
  return create_agent(
24
+ model=get_llm(),
25
  # tools=[math_solver, websearch_agent, web_search_agents],
26
  tools=[math_solver, web_search_agents, file_downloader, ocr_reader, list_files, http_get],
27
  system_prompt=(
agent/agents/answer_extractor.py CHANGED
@@ -2,6 +2,8 @@ from colorama import Fore, Style # type: ignore[import]
2
  from langchain.agents import create_agent
3
  from pydantic import BaseModel, Field
4
 
 
 
5
 
6
  class ExtractedAnswer(BaseModel):
7
  """Structured output for answer extraction."""
@@ -19,7 +21,7 @@ def extract_answer(raw_answer: str, question: str) -> str:
19
  print(f"{Fore.CYAN}[AnswerExtractor] Extracting...{Style.RESET_ALL}")
20
 
21
  agent = create_agent(
22
- model="google_genai:gemini-3-flash-preview",
23
  response_format=ExtractedAnswer,
24
  system_prompt="""\
25
  You are part of a system that demands extreme precision.
 
2
  from langchain.agents import create_agent
3
  from pydantic import BaseModel, Field
4
 
5
+ from agent.api.api import get_llm
6
+
7
 
8
  class ExtractedAnswer(BaseModel):
9
  """Structured output for answer extraction."""
 
21
  print(f"{Fore.CYAN}[AnswerExtractor] Extracting...{Style.RESET_ALL}")
22
 
23
  agent = create_agent(
24
+ model=get_llm(),
25
  response_format=ExtractedAnswer,
26
  system_prompt="""\
27
  You are part of a system that demands extreme precision.
agent/agents/websearchagent.py CHANGED
@@ -3,6 +3,7 @@ from colorama import Fore, Style # type: ignore[import]
3
  from langchain_core.tools import tool
4
  from langchain.agents import create_agent
5
  from langgraph.errors import GraphRecursionError
 
6
  from agent.tools.search import web_search
7
 
8
 
@@ -26,7 +27,7 @@ def websearch_agent(query: str) -> str:
26
  """
27
  print(f"{Fore.YELLOW}[SupervisorAgent -> WebSearchAgent] {query}{Style.RESET_ALL}")
28
  base_agent = create_agent(
29
- model="google_genai:gemini-3-flash-preview",
30
  tools=[web_search],
31
  system_prompt=(
32
  f"Current time is: {datetime.now(timezone.utc).isoformat()}. "
 
3
  from langchain_core.tools import tool
4
  from langchain.agents import create_agent
5
  from langgraph.errors import GraphRecursionError
6
+ from agent.api.api import get_llm
7
  from agent.tools.search import web_search
8
 
9
 
 
27
  """
28
  print(f"{Fore.YELLOW}[SupervisorAgent -> WebSearchAgent] {query}{Style.RESET_ALL}")
29
  base_agent = create_agent(
30
+ model=get_llm(),
31
  tools=[web_search],
32
  system_prompt=(
33
  f"Current time is: {datetime.now(timezone.utc).isoformat()}. "
agent/agents/websearchagents.py CHANGED
@@ -6,6 +6,8 @@ from colorama import Fore, Style # type: ignore[import]
6
  from langchain.agents import create_agent
7
  from langchain_core.tools import tool
8
  from pydantic import BaseModel, Field
 
 
9
  from tavily import TavilyClient # type: ignore[import]
10
 
11
 
@@ -45,7 +47,7 @@ def expand_queries(origin_question: str, query: str) -> list[str]:
45
  print(f"{Fore.CYAN}[QueryExpander] Expanding: {query}{Style.RESET_ALL}")
46
 
47
  agent = create_agent(
48
- model="google_genai:gemini-3-flash-preview",
49
  response_format=ExpandedQueries,
50
  system_prompt=(
51
  f"Current time is: {datetime.now(timezone.utc).isoformat()}. "
@@ -157,7 +159,7 @@ def subagent(origin_question: str, query: str) -> str:
157
  print(f"{color}[SubAgent] Investigating ({len(prompt)} chars)...{Style.RESET_ALL}")
158
 
159
  agent = create_agent(
160
- model="google_genai:gemini-3-flash-preview",
161
  system_prompt=(
162
  f"Current time is: {datetime.now(timezone.utc).isoformat()}. "
163
  "You are a research analyst. You are given a web page's full content "
@@ -194,7 +196,7 @@ def combine_result_agent(origin_question: str, query: str) -> str:
194
  )
195
 
196
  agent = create_agent(
197
- model="google_genai:gemini-3-flash-preview",
198
  system_prompt=(
199
  f"Current time is: {datetime.now(timezone.utc).isoformat()}. "
200
  "You are a research synthesizer. You receive findings from multiple "
 
6
  from langchain.agents import create_agent
7
  from langchain_core.tools import tool
8
  from pydantic import BaseModel, Field
9
+
10
+ from agent.api.api import get_llm
11
  from tavily import TavilyClient # type: ignore[import]
12
 
13
 
 
47
  print(f"{Fore.CYAN}[QueryExpander] Expanding: {query}{Style.RESET_ALL}")
48
 
49
  agent = create_agent(
50
+ model=get_llm(),
51
  response_format=ExpandedQueries,
52
  system_prompt=(
53
  f"Current time is: {datetime.now(timezone.utc).isoformat()}. "
 
159
  print(f"{color}[SubAgent] Investigating ({len(prompt)} chars)...{Style.RESET_ALL}")
160
 
161
  agent = create_agent(
162
+ model=get_llm(),
163
  system_prompt=(
164
  f"Current time is: {datetime.now(timezone.utc).isoformat()}. "
165
  "You are a research analyst. You are given a web page's full content "
 
196
  )
197
 
198
  agent = create_agent(
199
+ model=get_llm(),
200
  system_prompt=(
201
  f"Current time is: {datetime.now(timezone.utc).isoformat()}. "
202
  "You are a research synthesizer. You receive findings from multiple "
agent/api/api.py ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # from langchain_google_genai import ChatGoogleGenerativeAI
2
+
3
+ from langchain_community.chat_models import MiniMaxChat
4
+
5
+ _llm = None
6
+
7
+
8
+ def get_llm():
9
+ """Return a shared LLM instance."""
10
+ global _llm
11
+ if _llm is None:
12
+ # _llm = ChatGoogleGenerativeAI(model="gemini-3-flash-preview")
13
+ _llm = MiniMaxChat(model="m2.5-highspeed")
14
+ return _llm
agent/tools/ocr_reader.py CHANGED
@@ -4,7 +4,7 @@ from pathlib import Path
4
  from colorama import Fore, Style # type: ignore[import]
5
  from langchain_core.messages import HumanMessage
6
  from langchain_core.tools import tool
7
- from langchain_google_genai import ChatGoogleGenerativeAI
8
 
9
  # Workspace directory where downloaded files are stored
10
  WORKSPACE_DIR = Path(__file__).resolve().parents[2] / "workspace"
@@ -80,7 +80,7 @@ def ocr_reader(filename: str, query: str = "") -> str:
80
  base64_image = _encode_image(str(path))
81
  mime_type = _guess_mime_type(str(path))
82
 
83
- llm = ChatGoogleGenerativeAI(model="gemini-3-flash-preview")
84
 
85
  prompt = f"{SYSTEM_PROMPT}\n\nUser question: {query}" if query else SYSTEM_PROMPT
86
 
 
4
  from colorama import Fore, Style # type: ignore[import]
5
  from langchain_core.messages import HumanMessage
6
  from langchain_core.tools import tool
7
+ from agent.api.api import get_llm
8
 
9
  # Workspace directory where downloaded files are stored
10
  WORKSPACE_DIR = Path(__file__).resolve().parents[2] / "workspace"
 
80
  base64_image = _encode_image(str(path))
81
  mime_type = _guess_mime_type(str(path))
82
 
83
+ llm = get_llm()
84
 
85
  prompt = f"{SYSTEM_PROMPT}\n\nUser question: {query}" if query else SYSTEM_PROMPT
86
 
requirements.txt CHANGED
@@ -5,6 +5,7 @@ python-dotenv==1.2.2
5
  langchain
6
  langchain-core
7
  langchain-google-genai
 
8
  langgraph
9
  tavily-python
10
  sympy
 
5
  langchain
6
  langchain-core
7
  langchain-google-genai
8
+ langchain-community
9
  langgraph
10
  tavily-python
11
  sympy