You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
86 lines
2.9 KiB
Python
86 lines
2.9 KiB
Python
from typing import Dict, Any, List, Optional
|
|
|
|
import vertexai
|
|
from google.api_core.exceptions import NotFound
|
|
from vertexai.generative_models import FunctionDeclaration, GenerativeModel, Tool
|
|
|
|
from app.core.settings import settings
|
|
from app.models.tool_model import ToolDefinition
|
|
|
|
|
|
class LLMService:
|
|
|
|
def __init__(self):
|
|
"""Inicializa o cliente Vertex AI e define modelos de fallback."""
|
|
vertexai.init(
|
|
project=settings.google_project_id,
|
|
location=settings.google_location,
|
|
)
|
|
|
|
configured = settings.vertex_model_name.strip()
|
|
fallback_models = ["gemini-2.5-flash", "gemini-2.0-flash-001", "gemini-1.5-pro"]
|
|
self.model_names = [configured] + [m for m in fallback_models if m != configured]
|
|
|
|
def build_vertex_tools(self, tools: List[ToolDefinition]) -> Optional[List[Tool]]:
|
|
"""Converte tools internas para o formato esperado pelo Vertex AI."""
|
|
# Vertex espera uma lista de Tool, com function_declarations agrupadas em um unico Tool.
|
|
if not tools:
|
|
return None
|
|
|
|
function_declarations = [
|
|
FunctionDeclaration(
|
|
name=tool.name,
|
|
description=tool.description,
|
|
parameters=tool.parameters,
|
|
)
|
|
for tool in tools
|
|
]
|
|
|
|
return [Tool(function_declarations=function_declarations)]
|
|
|
|
async def generate_response(
|
|
self,
|
|
message: str,
|
|
tools: List[ToolDefinition],
|
|
history: List[Dict[str, Any]] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Gera resposta textual ou chamada de tool a partir da mensagem do usuario."""
|
|
vertex_tools = self.build_vertex_tools(tools)
|
|
|
|
response = None
|
|
last_error = None
|
|
|
|
for model_name in self.model_names:
|
|
try:
|
|
model = GenerativeModel(model_name)
|
|
chat = model.start_chat(history=history or [])
|
|
send_kwargs = {"tools": vertex_tools} if vertex_tools else {}
|
|
response = chat.send_message(message, **send_kwargs)
|
|
break
|
|
except NotFound as err:
|
|
last_error = err
|
|
continue
|
|
|
|
if response is None:
|
|
if last_error:
|
|
raise RuntimeError(
|
|
f"Nenhum modelo Vertex disponivel. Verifique VERTEX_MODEL_NAME e acesso no projeto. Erro: {last_error}"
|
|
) from last_error
|
|
raise RuntimeError("Falha ao gerar resposta no Vertex AI.")
|
|
|
|
part = response.candidates[0].content.parts[0]
|
|
|
|
if part.function_call:
|
|
return {
|
|
"response": None,
|
|
"tool_call": {
|
|
"name": part.function_call.name,
|
|
"arguments": dict(part.function_call.args),
|
|
},
|
|
}
|
|
|
|
return {
|
|
"response": response.text,
|
|
"tool_call": None,
|
|
}
|