from langchain_anthropic import ChatAnthropic
from pydantic import BaseModel, Field
from typing import List
import pandas as pd
from dotenv import load_dotenv
import os
import warnings
import json
# Ignorar todos los warnings
warnings.filterwarnings("ignore")

load_dotenv()

llm = ChatAnthropic(model="claude-3-haiku-20240307", temperature=0)

# Pydantic class definitions


class Title(BaseModel):
    text: str = Field(description="Title of the notification")

class Message(BaseModel):
    content: str = Field(description="Main content or message of the notification")

class Number(BaseModel):
    value: int = Field(description="Relevant number or identifier within the notification")

class NotificationAnalysis(BaseModel):
    title: List[Title] = Field(default=[], description="Title of the push notification")
    message: List[Message] = Field(default=[], description="Main message or content")
    number: List[Number] = Field(default=[], description="Relevant number or ID within the notification")

# Function to process notification text
def extract_from_notification(llm, notification_text):
    prompt = f"""
    Given the following push notification text, structure the content into JSON format under the sections "title", "message", and "number":
    - Extract the title of the notification and add it to "title".
    - Extract the main message or content and add it to "message".
    - If there is any relevant number or identifier, add it to "number".
    - Do not include information that is not directly related to title, message, or relevant number.

    Notification to process:
    {notification_text}
    
    Expected format:
    {{
        "title": [{{"text": "Title of the notification"}}],
        "message": [{{"content": "Main message or content"}}],
        "number": [{{"value": 123}}]
    }}
    """
    # Set up the LLM with structured output according to the NotificationAnalysis class
    structured_llm = llm.with_structured_output(NotificationAnalysis)

    # Invoke the structured model and obtain the response
    response = structured_llm.invoke(prompt)
    response = response.dict()


    return response


notification_text = """
New update available: Version 2.0 has been released with improved features. 
Your device ID is 12345. Tap to update now.
"""

# Ejecuta la función con el modelo Anthropic real
response = extract_from_notification(llm, notification_text)

# Muestra la respuesta
print(response)
print(type(response))