40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
import json
|
|
from services.ollama_service import OllamaService
|
|
from services.external_model_service import ExternalModelService
|
|
|
|
class ModelRouter:
|
|
def __init__(self):
|
|
self.ollama = OllamaService()
|
|
self.external = ExternalModelService()
|
|
|
|
with open("models.json", "r") as f:
|
|
self.models = json.load(f)
|
|
|
|
def pick_model(self, intent: str, mode: str):
|
|
# regras simples (pode ser expandido)
|
|
if intent == "document" and mode == "advogado":
|
|
return self.models["local"]["llama3.1-70b"]
|
|
|
|
if intent == "document" and mode == "corporativo":
|
|
return self.models["local"]["qwen2.5-7b"]
|
|
|
|
if intent == "search":
|
|
return self.models["external"]["claude-3.5-sonnet"]
|
|
|
|
if intent == "monitoring":
|
|
return self.models["external"]["mistral-large"]
|
|
|
|
# fallback
|
|
return self.models["local"]["llama3.1-8b"]
|
|
|
|
def generate(self, intent: str, mode: str, prompt: str):
|
|
model = self.pick_model(intent, mode)
|
|
|
|
if model["provider"] == "ollama":
|
|
return self.ollama.generate(prompt)
|
|
|
|
if model["provider"] in ["openai", "anthropic", "mistral"]:
|
|
return self.external.generate(model["provider"], model["model"], prompt)
|
|
|
|
return "Erro: modelo não encontrado."
|