from langchain_ollama import ChatOllama
from langchain_ibm import ChatWatsonx
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
from langchain_core.output_parsers import (
    StrOutputParser,
    JsonOutputParser,
    PydanticOutputParser,
)
from pydantic import BaseModel, Field
from typing import Literal
from dotenv import load_dotenv
import os
import gradio as gr

qwen_llm = ChatOllama(model="qwen3.5:4b", temperature=0)
exaone_llm = ChatOllama(model="exaone3.5:2.4b", temperature=0)


class NewsResult(BaseModel):
    title: str
    date: str = Field(description="YYYY-MM-DD 형식, 없는 경우 '없음'")
    category: Literal["정치", "경제", "사회", "문화", "스포츠", "IT", "국제", "기타"]
    keywords: list[str] = Field(description="핵심 키워드 3개 이내")


pydantic_parser = PydanticOutputParser(pydantic_object=NewsResult)
template = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "당신은 뉴스 분석 전문가입니다. 기사에서 아래 정보를 추출하세요. {format_instructions}",
        ),
        ("human", "{article}"),
    ]
).partial(format_instructions=pydantic_parser.get_format_instructions())


chain = template | exaone_llm | pydantic_parser


def news_input(texts):
    # ===을 기준으로 기사 분리
    articles = [article for article in texts.split("===") if article.strip()]

    response = chain.batch([{"article": article} for article in articles])
    return "\n\n".join(str(item) for item in response)


app = gr.Interface(
    news_input,
    inputs=[
        gr.Textbox(
            label="기사",
            placeholder="여러 기사 입력 시 구분자로 ===을 사용하세요",
            lines=20,
        ),
    ],
    outputs=[
        gr.Textbox(label="요약", lines=20),
    ],
    title="✒ 뉴스 분석 전문가",
    description="뉴스 기사에서 정보를 추출합니다.",
)

app.launch()