Commit ·
bcdc55d
1
Parent(s): c36efba
refactor(api): add LLM abstraction layer to replace hardcoded model strings
Browse filesReplace 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 +2 -0
- agent/agent.py +2 -1
- agent/agents/answer_extractor.py +3 -1
- agent/agents/websearchagent.py +2 -1
- agent/agents/websearchagents.py +5 -3
- agent/api/api.py +14 -0
- agent/tools/ocr_reader.py +2 -2
- requirements.txt +1 -0
.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=
|
| 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=
|
| 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=
|
| 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=
|
| 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=
|
| 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=
|
| 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
|
| 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 =
|
| 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
|