app.py

from smolagents import CodeAgent, ToolCallingAgent, DuckDuckGoSearchTool, LiteLLMModel, tool
from typing import Optional
import os
from dotenv import load_dotenv
import gradio as gr

# .env 파일에서 환경 변수를 로드
load_dotenv()

# Gemini AI 모델 초기화
model = LiteLLMModel(
    model_id="gemini/gemini-2.0-flash-exp",
    api_key=os.getenv("GOOGLE_API_KEY")
)

@tool
def search_news(topic: str) -> str:
    """
    주어진 주제에 대한 최신 뉴스를 검색하는 도구입니다.
    Args:
        topic: 검색할 뉴스 주제
    Returns:
        검색된 뉴스 정보
    """
    search_agent = CodeAgent(
        tools=[DuckDuckGoSearchTool()],
        model=model,
        additional_authorized_imports=["requests", "bs4"]
    )
    search_query = f"최신 뉴스 {topic}"
    return search_agent.run(f"Search for: {search_query}")

@tool
def write_article(topic: str, search_results: str) -> str:
    """
    검색 결과를 바탕으로 새로운 기사를 작성하는 도구입니다.
    Args:
        topic: 기사 주제
        search_results: 검색된 뉴스 정보
    Returns:
        작성된 기사
    """
    NEWS_WRITER_PROMPT = """
    당신은 전문 기자입니다. 주어진 주제에 대해 검색된 정보를 바탕으로 새로운 기사를 한국어로 작성해야 합니다.
    다음 형식으로 기사를 작성해주세요:

    1. 제목 (흥미롭고 눈에 띄는 제목)
    2. 요약 (핵심 내용을 2-3문장으로 요약)
    3. 본문 (상세한 내용을 단락으로 구분하여 작성)
    4. 결론 (기사의 의의나 향후 전망)

    기사는 객관적이고 사실에 기반하여 작성해야 하며, 검색된 정보를 재구성하여 새로운 시각으로 작성해주세요.
    """
    
    writer_agent = ToolCallingAgent(
        model=model,
        system_prompt=NEWS_WRITER_PROMPT
    )
    
    prompt = f"""
    다음 주제와 검색 결과를 바탕으로 새로운 기사를 한국어로 작성해주세요:

    주제: {topic}

    검색 결과:
    {search_results}

    위 정보를 바탕으로 새롭고 통찰력 있는 기사를 한국어로 작성해주세요.
    """
    return writer_agent.run(prompt)

class NewsWriterGradioUI:
    def __init__(self):
        self.agent = CodeAgent(
            tools=[search_news, write_article],
            model=model
        )
    
    def generate_article(self, topic: str) -> str:
        try:
            prompt = f"""
            1. '{topic}' 주제에 대해 search_news 도구를 사용하여 최신 뉴스를 검색하세요.
            2. 검색된 결과를 바탕으로 write_article 도구를 사용하여 새로운 기사를 작성하세요.
            """
            result = self.agent.run(prompt)
            return result
        except Exception as e:
            return f"에러가 발생했습니다: {str(e)}"
    
    def launch(self):
        interface = gr.Interface(
            fn=self.generate_article,
            inputs=gr.Textbox(label="뉴스 주제를 입력하세요"),
            outputs=gr.Textbox(label="생성된 기사"),
            title="AI 기자",
            description="주제를 입력하면 관련 뉴스를 검색하고 새로운 기사를 작성합니다."
        )
        interface.launch()

def main():
    ui = NewsWriterGradioUI()
    ui.launch()

if __name__ == "__main__":
    main()