from smolagents import CodeAgent, DuckDuckGoSearchTool, tool, LiteLLMModel
import yaml
import traceback
import os
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

@tool
def search_news(topic: str) -> str:
    """
    주어진 주제에 대한 최신 뉴스를 검색하는 도구입니다.
    Args:
        topic: 검색할 뉴스 주제
    Returns:
        검색된 뉴스 정보
    """
    try:
        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}")
    except Exception as e:
        traceback.print_exc()
        return f"검색 중 오류 발생: {str(e)}"

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

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

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

        주제: {topic}

        검색 결과:
        {search_results}

        위 정보를 바탕으로 새롭고 통찰력 있는 기사를 작성해주세요."""

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        
        response = model(messages)
        if isinstance(response, dict) and 'content' in response:
            return response['content']
        elif hasattr(response, 'content'):
            return response.content
        else:
            return str(response)
    except Exception as e:
        traceback.print_exc()
        return f"기사 작성 중 오류 발생: {str(e)}"


final_answer = FinalAnswerTool()
model = LiteLLMModel(
    model_id="gemini/gemini-2.0-flash-exp", # 반드시 model_id 로 지정해야 함
    max_tokens=2096,
    temperature=0.7,
    api_key=os.getenv("GOOGLE_API_KEY")
)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[search_news, write_article, final_answer],
    max_steps=10,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

if __name__ == "__main__":
    try:
        GradioUI(agent).launch()
    except Exception as e:
        print(f"Error launching UI: {str(e)}")