Create tool to get news using NewsAPI

app/ai/agent.py

@@ -14,6 +14,7 @@
 from langchain_community.utilities.openweathermap import OpenWeatherMapAPIWrapper
 from langchain.tools import WikipediaQueryRun
 from langchain_community.utilities import WikipediaAPIWrapper
+from .tools import news
 
 langchain.debug = True
 
@@ -62,7 +63,14 @@ def setup_agent(session_id: str, model: str) -> AgentExecutor:
 
     weather = OpenWeatherMapAPIWrapper(openweathermap_api_key=settings.OPENWEATHERAPP_API_KEY)
     wikipedia = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())
+    news_tool = news.NewsTool()
+
     tools = [
+        Tool(
+            name="News",
+            func=news_tool.run,
+            description="Useful for when you need to answer questions about current news"
+        ),
         Tool(
             name="Wikipedia",
             func=wikipedia.run,

app/ai/tools/news.py

@@ -0,0 +1,43 @@
+from typing import Optional
+
+from langchain.callbacks.manager import (
+    AsyncCallbackManagerForToolRun,
+    CallbackManagerForToolRun,
+)
+from langchain_core.tools import BaseTool
+from newsapi import NewsApiClient
+from config.settings import settings
+
+
+newsapi = NewsApiClient(api_key=settings.NEWS_API)
+
+
+def get_top_headlines_formatted() -> str:
+    # Get top headlines
+    result = ""
+
+    top_headlines = newsapi.get_top_headlines()
+
+    articles = top_headlines["articles"]
+    for idx, article in enumerate(articles):
+        title = f"{idx} + {article['title']} + '\n' "
+        result += title
+        result += f"Description: {article["description"]} + '\n'"
+
+    return result
+
+
+class NewsTool(BaseTool):
+    name = "news_tool"
+    description = "Useful for when you need to answer questions about current news"
+
+    def _run(
+            self, run_manager: Optional[CallbackManagerForToolRun] = None
+    ) -> str:
+        return get_top_headlines_formatted()
+
+    async def _arun(
+            self, run_manager: Optional[AsyncCallbackManagerForToolRun] = None
+    ) -> str:
+        """Use the tool asynchronously."""
+        raise NotImplementedError("custom_search does not support async")

app/config/settings.py

@@ -51,5 +51,7 @@ class Settings(BaseSettings):
 
     OPENWEATHERAPP_API_KEY: Optional[str] = os.getenv("OPENWEATHERAPP_API_KEY")
 
+    NEWS_API: Optional[str] = os.getenv("NEWS_API")
+
 
 settings = Settings()