!pip install openai
!pip install gradio
#!pip install arxiv

import gradio as gr
from openai import OpenAI
import json
import arxiv  # 사전 설치 필요: !pip install arxiv

# OpenAI 클라이언트 초기화
client = OpenAI()


# 논문 검색 및 첫 논문 요약 번역 함수
def search_paper(query):
    results = arxiv.Search(
        query=query,
        max_results=3,
        sort_by=arxiv.SortCriterion.Relevance
    )

    papers = []
    abstracts = []
    for r in results.results():
        papers.append(f"📄 {r.title}\n🧑‍🔬 저자: {', '.join([a.name for a in r.authors])}\n📅 출판일: {r.published.date()}\n🔗 링크: {r.pdf_url}")
        abstracts.append(r.summary)

    if not papers:
        return "❌ 해당 주제에 대한 논문을 찾을 수 없습니다."

    # 첫 번째 논문 요약 번역 요청
    abstract_summary = abstracts[0]
    summary_translation = translate_to_korean(abstract_summary)

    result_text = "\n\n".join(papers)
    result_text += f"\n\n📘 첫 번째 논문 요약 (한글 번역):\n{summary_translation}"
    return result_text

# 초록 번역 함수
def translate_to_korean(abstract_text):
    completion = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "아래 영문 초록을 자연스러운 한국어로 요약 번역해줘."},
            {"role": "user", "content": abstract_text}
        ]
    )
    return completion.choices[0].message.content.strip()

# OpenAI Tool 등록
tools = [
    {
        "type": "function",
        "function": {
            "name": "search_paper",
            "description": "주어진 키워드로 arXiv에서 논문을 검색합니다.",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "논문 검색 키워드이며 반드시 영어 단어나 문장으로 변환해줘",
                    },
                },
                "required": ["query"]
            },
        },
    },
]

# 응답 생성 함수
def generate_response(prompt_text: str):
    completion = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "사용자의 요청에 따라 논문 검색에 필요한 키워드를 영어로 번역하여 추출한 후 search_paper 함수를 호출해줘. 그 외에는 일반적인 답변을 해줘."},
            {"role": "user", "content": prompt_text}
        ],
        tools=tools,
    )

    response = completion.choices[0].message

    if response.tool_calls:
        tool_call = response.tool_calls[0]
        function_name = tool_call.function.name
        arguments_json = json.loads(tool_call.function.arguments)

        if function_name == "search_paper":
            query = arguments_json.get('query')
            return search_paper(query)

        else:
            return "⚠️ 지원되지 않는 함수입니다."

    elif response.function_call:
        return str(response.function_call)

    else:
        return response.content

# Gradio 인터페이스
def chatgpt_interface():
    with gr.Blocks() as demo:
        gr.Markdown("## 💬 ChatGPT API 데모 (날씨 + 논문 검색 + 요약 번역)")

        with gr.Row():
            with gr.Column():
                input_text = gr.Textbox(label="질문을 입력하세요 (예: 양자 컴퓨팅 논문 검색)")
                output_text = gr.Textbox(label="AI 응답", interactive=False, lines=18)

                input_text.submit(generate_response, inputs=input_text, outputs=output_text)

        demo.launch()

# 실행
chatgpt_interface()