Update app.py

app.py

@@ -1,17 +1,9 @@
-from smolagents import CodeAgent, ToolCallingAgent, DuckDuckGoSearchTool, LiteLLMModel, tool
-from typing import Optional
+from smolagents import CodeAgent, DuckDuckGoSearchTool, tool, LiteLLMModel
+import yaml
+import traceback
 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")
-)
+from tools.final_answer import FinalAnswerTool
+from Gradio_UI import GradioUI
 
 @tool
 def search_news(topic: str) -> str:
@@ -22,13 +14,17 @@ def search_news(topic: str) -> str:
     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}")
+    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:
@@ -40,66 +36,60 @@ def write_article(topic: str, search_results: str) -> str:
     Returns:
         작성된 기사
     """
-    NEWS_WRITER_PROMPT = """
-    당신은 전문 기자입니다. 주어진 주제에 대해 검색된 정보를 바탕으로 새로운 기사를 한국어로 작성해야 합니다.
-    다음 형식으로 기사를 작성해주세요:
+    try:
+        system_prompt = """
+        당신은 전문 기자입니다. 주어진 주제에 대해 검색된 정보를 바탕으로 새로운 기사를 한국어로 작성해야 합니다.
+        다음 형식으로 기사를 작성해주세요:
 
-    1. 제목 (흥미롭고 눈에 띄는 제목)
-    2. 요약 (핵심 내용을 2-3문장으로 요약)
-    3. 본문 (상세한 내용을 단락으로 구분하여 작성)
-    4. 결론 (기사의 의의나 향후 전망)
+        1. 제목 (흥미롭고 눈에 띄는 제목)
+        2. 요약 (핵심 내용을 2-3문장으로 요약)
+        3. 본문 (상세한 내용을 단락으로 구분하여 작성)
+        4. 결론 (기사의 의의나 향후 전망)
 
-    기사는 객관적이고 사실에 기반하여 작성해야 하며, 검색된 정보를 재구성하여 새로운 시각으로 작성해주세요.
-    """
-    
-    writer_agent = ToolCallingAgent(
-        model=model,
-        system_prompt=NEWS_WRITER_PROMPT
-    )
-    
-    prompt = f"""
-    다음 주제와 검색 결과를 바탕으로 새로운 기사를 한국어로 작성해주세요:
+        기사는 객관적이고 사실에 기반하여 작성해야 하며, 검색된 정보를 재구성하여 새로운 시각으로 작성해주세요.
+        """
+        
+        prompt = f"""
+        다음 주제와 검색 결과를 바탕으로 새로운 기사를 한국어로 작성해주세요:
 
-    주제: {topic}
+        주제: {topic}
 
-    검색 결과:
-    {search_results}
+        검색 결과:
+        {search_results}
 
-    위 정보를 바탕으로 새롭고 통찰력 있는 기사를 한국어로 작성해주세요.
-    """
-    return writer_agent.run(prompt)
+        위 정보를 바탕으로 새롭고 통찰력 있는 기사를 5문단으로 작성해주세요.
+        """
+        return model.generate(system_prompt + prompt)
+    except Exception as e:
+        traceback.print_exc()
+        return f"기사 작성 중 오류 발생: {str(e)}"
 
-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()
+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=6,
+    verbosity_level=1,
+    grammar=None,
+    planning_interval=None,
+    name=None,
+    description=None,
+    prompt_templates=prompt_templates
+)
 
 if __name__ == "__main__":
-    main() 
+    try:
+        GradioUI(agent).launch()
+    except Exception as e:
+        print(f"Error launching UI: {str(e)}")