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.
68 lines
1.9 KiB
Python
68 lines
1.9 KiB
Python
from typing import Dict, Any, List
|
|
import vertexai
|
|
from vertexai.generative_models import GenerativeModel, Tool, FunctionDeclaration
|
|
from app.core.settings import settings
|
|
from app.models.tool_model import ToolDefinition
|
|
|
|
|
|
class LLMService:
|
|
|
|
def __init__(self):
|
|
vertexai.init(
|
|
project=settings.google_project_id,
|
|
location=settings.google_location
|
|
)
|
|
|
|
self.model = GenerativeModel("gemini-1.5-pro")
|
|
|
|
def build_vertex_tools(self, tools: List[ToolDefinition]): # Converte as Tools internas (ToolDefinition) para o formato que o Vertex AI entende.
|
|
|
|
vertex_tools = []
|
|
|
|
for tool in tools:
|
|
vertex_tools.append(
|
|
Tool(
|
|
function_declarations=[
|
|
FunctionDeclaration(
|
|
name=tool.name,
|
|
description=tool.description,
|
|
parameters=tool.parameters
|
|
)
|
|
]
|
|
)
|
|
)
|
|
|
|
return vertex_tools
|
|
|
|
async def generate_response(
|
|
self,
|
|
message: str,
|
|
tools: List[ToolDefinition],
|
|
history: List[Dict[str, Any]] = None
|
|
) -> Dict[str, Any]:
|
|
|
|
vertex_tools = self.build_vertex_tools(tools) # Convertendo tools para formato do Vertex
|
|
|
|
chat = self.model.start_chat(
|
|
history=history or [],
|
|
tools=vertex_tools
|
|
)
|
|
|
|
response = chat.send_message(message)
|
|
|
|
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
|
|
}
|