Update model_router.py

pull/1101/head
CI-DEV 1 week ago committed by GitHub
parent 9ebd54de57
commit a602dff844
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -4,7 +4,6 @@ from concurrent.futures import ThreadPoolExecutor
from typing import Optional
from pydantic import BaseModel, Field
from swarms.utils.function_caller_model import OpenAIFunctionCaller
from swarms.utils.any_to_str import any_to_str
from swarms.utils.formatter import formatter
from swarms.utils.litellm_wrapper import LiteLLM
@ -174,7 +173,7 @@ class ModelRouter:
temperature (float): Temperature parameter for model randomness
max_workers (int): Maximum concurrent workers for batch processing
model_output (ModelOutput): Pydantic model for structured outputs
model_caller (OpenAIFunctionCaller): Function calling interface
model_caller (LiteLLM): Function calling interface
"""
def __init__(
@ -210,11 +209,11 @@ class ModelRouter:
if self.max_workers == "auto":
self.max_workers = os.cpu_count()
self.model_caller = OpenAIFunctionCaller(
base_model=ModelOutput,
self.model_caller = LiteLLM(
model_name="gpt-4o",
response_format=ModelOutput,
temperature=self.temperature,
system_prompt=self.system_prompt,
api_key=api_key,
)
except Exception as e:
raise RuntimeError(

Loading…
Cancel
Save