import uuid
from typing import TypedDict
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from langchain_openai import ChatOpenAI
import os
import re
from dotenv import load_dotenv

load_dotenv()

class MainState(TypedDict):
    question: HumanMessage | str| None
    scratchpad: list[AIMessage | ToolMessage] | None
    answer: AIMessage | None
    next_node: str | None
    history: list[HumanMessage | AIMessage]

llm = ChatOpenAI(
    model="sabiazinho-3",
    temperature=0,
    api_key=os.environ.get('MARITACA_API_KEY'), # Insira sua chave aqui
    base_url="https://chat.maritaca.ai/api",
)

def generate_uuid():
    return str(uuid.uuid4())

def post_process(message: AIMessage) -> AIMessage:
    matches = re.findall(r"\[SOT\](.*)\[EOT\]", message.content, re.DOTALL)
    matches = matches[0] if matches else None

    if matches:
        return AIMessage(content='', additional_kwargs={'custom_tool_call': matches} ,tool_call_id=generate_uuid())
    return AIMessage(content=message.content)