File size: 10,693 Bytes
76962bf 749fa40 76962bf 749fa40 76962bf 749fa40 b1198f0 749fa40 84ae02f 749fa40 b1198f0 749fa40 b1198f0 749fa40 b1198f0 76962bf 749fa40 76962bf fb01a6c b1198f0 76962bf fb01a6c 749fa40 b1198f0 749fa40 76962bf fb01a6c 749fa40 76962bf fb01a6c 749fa40 fb01a6c 749fa40 fb01a6c 749fa40 fb01a6c 749fa40 76962bf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 | import os
import time
import json
import asyncio
import requests
from dotenv import load_dotenv
from langchain_core.messages import AIMessage, SystemMessage, HumanMessage
from src.utils.logger import setup_logger
logger = setup_logger("ModelManager")
load_dotenv()
class ReqModel:
def __init__(self, model: str, temperature: float, base_url: str, api_key: str, headers: dict):
self.model = model
self.temperature = temperature
self.base_url = base_url
self.api_key = api_key
self.headers = headers
self.bound_tools = None
def bind_tools(self, tools, **kwargs):
new_model = ReqModel(self.model, self.temperature, self.base_url, self.api_key, self.headers)
new_model.bound_tools = tools
return new_model
def _convert_messages(self, messages):
req_msgs = []
for m in messages:
if isinstance(m, SystemMessage):
req_msgs.append({"role": "system", "content": m.content})
elif isinstance(m, HumanMessage):
req_msgs.append({"role": "user", "content": m.content})
elif isinstance(m, AIMessage):
req_msgs.append({"role": "assistant", "content": m.content})
elif isinstance(m, dict) and "role" in m and "content" in m:
req_msgs.append(m)
else:
req_msgs.append({"role": "user", "content": str(getattr(m, 'content', m))})
return req_msgs
def _format_tools(self):
if not self.bound_tools:
return None
tools_list = []
for tool in self.bound_tools:
if hasattr(tool, "name") and hasattr(tool, "description") and hasattr(tool, "args_schema"):
tools_list.append({
"type": "function",
"function": {
"name": tool.name,
"description": tool.description,
"parameters": tool.args_schema.schema() if tool.args_schema else {"type": "object", "properties": {}}
}
})
return tools_list
def _make_request(self, messages, config=None, **kwargs):
url = f"{self.base_url}/chat/completions"
headers = {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json"
}
headers.update(self.headers)
payload = {
"model": self.model,
"messages": self._convert_messages(messages),
"temperature": self.temperature,
}
formatted_tools = self._format_tools()
if formatted_tools:
payload["tools"] = formatted_tools
response = requests.post(url, headers=headers, json=payload)
response.raise_for_status()
data = response.json()
message = data["choices"][0]["message"]
content = message.get("content", "")
ai_message = AIMessage(content=content if content else "")
if "tool_calls" in message and message["tool_calls"]:
tool_calls = []
for tc in message["tool_calls"]:
try:
args = json.loads(tc["function"]["arguments"])
except Exception:
args = {}
tool_calls.append({
"name": tc["function"]["name"],
"args": args,
"id": tc["id"]
})
ai_message.additional_kwargs["tool_calls"] = message["tool_calls"]
ai_message.tool_calls = tool_calls
ai_message.response_metadata = {"token_usage": data.get("usage", {})}
return ai_message
def invoke(self, messages, config=None, **kwargs):
return self._make_request(messages, config, **kwargs)
async def ainvoke(self, messages, config=None, **kwargs):
loop = asyncio.get_event_loop()
return await loop.run_in_executor(None, lambda: self._make_request(messages, config, **kwargs))
def stream(self, messages, config=None, **kwargs):
yield self._make_request(messages, config, **kwargs)
async def astream(self, messages, config=None, **kwargs):
loop = asyncio.get_event_loop()
response = await loop.run_in_executor(None, lambda: self._make_request(messages, config, **kwargs))
yield response
class RateLimitFallbackWrapper:
def __init__(self, main_llm, fallback_llms):
self.main_llm = main_llm
self.fallback_llms = fallback_llms
self.bound_tools = None
def bind_tools(self, tools, **kwargs):
new_main = self.main_llm.bind_tools(tools, **kwargs)
new_falls = [llm.bind_tools(tools, **kwargs) for llm in self.fallback_llms]
new_wrapper = RateLimitFallbackWrapper(new_main, new_falls)
new_wrapper.bound_tools = tools
return new_wrapper
async def ainvoke(self, messages, config=None, **kwargs):
try:
return await self.main_llm.ainvoke(messages, config=config, **kwargs)
except Exception as e:
logger.warning(f"LLM Error with main model: {e}. Attempting fallbacks immediately.")
for idx, fb_llm in enumerate(self.fallback_llms):
try:
logger.info(f"Trying fallback model {idx+1} [Model: {fb_llm.model}]")
return await fb_llm.ainvoke(messages, config=config, **kwargs)
except Exception as fb_e:
logger.warning(f"Fallback {idx+1} failed: {fb_e}")
raise RuntimeError("All models (main and fallbacks) failed.") from None
def invoke(self, messages, config=None, **kwargs):
try:
return self.main_llm.invoke(messages, config=config, **kwargs)
except Exception as e:
logger.warning(f"LLM Error with main model: {e}. Attempting fallbacks immediately.")
for idx, fb_llm in enumerate(self.fallback_llms):
try:
logger.info(f"Trying fallback model {idx+1}")
return fb_llm.invoke(messages, config=config, **kwargs)
except Exception as fb_e:
logger.warning(f"Fallback {idx+1} failed: {fb_e}")
raise RuntimeError("All models failed synchronously.") from None
def stream(self, messages, config=None, **kwargs):
try:
yield from self.main_llm.stream(messages, config=config, **kwargs)
return
except Exception as e:
logger.warning(f"LLM stream error with main model: {e}. Attempting fallbacks immediately.")
for idx, fb_llm in enumerate(self.fallback_llms):
try:
logger.info(f"Trying fallback model {idx+1} for streaming")
yield from fb_llm.stream(messages, config=config, **kwargs)
return
except Exception as fb_e:
logger.warning(f"Fallback {idx+1} streaming failed: {fb_e}")
raise RuntimeError("All models failed while streaming.")
async def astream(self, messages, config=None, **kwargs):
try:
async for chunk in self.main_llm.astream(messages, config=config, **kwargs):
yield chunk
return
except Exception as e:
logger.warning(f"LLM async stream error with main model: {e}. Attempting fallbacks immediately.")
for idx, fb_llm in enumerate(self.fallback_llms):
try:
logger.info(f"Trying fallback model {idx+1} for async streaming")
async for chunk in fb_llm.astream(messages, config=config, **kwargs):
yield chunk
return
except Exception as fb_e:
logger.warning(f"Fallback {idx+1} async streaming failed: {fb_e}")
raise RuntimeError("All models failed while async streaming.")
class ModelManager:
def __init__(self, model_name: str = "google/gemma-4-26b-a4b-it:free"):
self.provider = os.getenv("MODEL_PROVIDER", "openrouter").lower()
self.model_name = os.getenv("OPENROUTER_MODEL_NAME", model_name)
def _get_openrouter_api_keys(self):
primary_api_key = os.getenv("OPENROUTER_API_KEY")
secondary_api_key = os.getenv("OPENROUTER_SECONDARY_API_KEY")
if not primary_api_key and secondary_api_key:
logger.warning("OPENROUTER_API_KEY missing; using OPENROUTER_SECONDARY_API_KEY as the active key.")
primary_api_key = secondary_api_key
secondary_api_key = None
if not primary_api_key:
raise EnvironmentError(
"OpenRouter requires OPENROUTER_API_KEY or OPENROUTER_SECONDARY_API_KEY in the environment."
)
return primary_api_key, secondary_api_key
def get_llm(self, temperature: float = 0, model_name: str = None):
model = model_name or self.model_name
logger.info(f"Initializing LLM: Provider={self.provider}, Model={model}")
primary_api_key, secondary_api_key = self._get_openrouter_api_keys()
base_url = "https://openrouter.ai/api/v1"
main_llm = ReqModel(
model=model,
temperature=temperature,
base_url=base_url,
api_key=primary_api_key,
headers={
"HTTP-Referer": "https://github.com/Sudharshan-3904/dmChatbot",
"X-Title": "Medical AI Chatbot"
}
)
fallback_llms = []
if secondary_api_key:
fallback_llms.append(
ReqModel(
model=model,
temperature=temperature,
base_url=base_url,
api_key=secondary_api_key,
headers={
"HTTP-Referer": "https://github.com/Sudharshan-3904/dmChatbot",
"X-Title": "Medical AI Chatbot"
}
)
)
fallback_models = [
"google/gemma-4-26b-a4b-it:free",
"google/gemma-4-31b-it:free",
"openai/gpt-oss-20b:free"
]
fallback_llms.extend([
ReqModel(
model=m,
temperature=temperature,
base_url=base_url,
api_key=primary_api_key,
headers={
"HTTP-Referer": "https://github.com/Sudharshan-3904/dmChatbot",
"X-Title": "Medical AI Chatbot"
}
) for m in fallback_models if m != model
])
return RateLimitFallbackWrapper(main_llm, fallback_llms)
model_manager = ModelManager()
|