Merge pull request #1 from CoToYo/db_update

Tiny code structure updates.

controllers/NewsController.py

@@ -1,27 +1,27 @@
-import services.NewsService as NewsService
+from services.NewsService import NewsService
 
-"""
-return news without considering keywords
-"""
 
+class NewsController:
+    def __init__(self) -> None:
+        self.news_service = NewsService()
 
-def getNews():
-    return NewsService.getNews()
+    """
+    return news without considering keywords
+    """
 
+    def getNews(self):
+        return self.news_service.getNews()
 
-"""
-return news based on certain keywords
-"""
+    """
+    return news based on certain keywords
+    """
 
+    def getNewsWithKeywords(self, user_keywords):
+        return self.news_service.getNewsWithKeywords(user_keywords)
 
-def getNewsWithKeywords(user_keywords):
-    return NewsService.getNewsWithKeywords(user_keywords)
+    """
+    deal requests with wrong route
+    """
 
-
-"""
-deal requests with wrong route
-"""
-
-
-def notFound(error):
-    return NewsService.notFound(error)
+    def notFound(self, error):
+        return self.news_service.notFound(error)

db_update/Update.py

@@ -52,7 +52,7 @@
 Notice: "id" should be removed before the news is inserted into the database.
 """
 
-# db has limited storage space(512mb), clean up the old news everytime
+# db has limited storage space(512mb), clean up old news before fetching new news.
 collection.delete_many({})
 
 # Insert all feteched news into the database.

models/NewsModel.py

@@ -1,7 +1,14 @@
 from config.Database import client
 
-db = client.CybernewsDB
-
 """
 TODO: Read news from database.
 """
+
+
+class CybernewsDB:
+    def __init__(self):
+        self.client = client
+        self.db = self.client["CybernewsDB"]
+
+    def get_collections(self, collection_name):
+        return self.db.get_collection(collection_name)

routes/NewsRoutes.py

@@ -1,7 +1,8 @@
 from flask import *
-from controllers.NewsController import *
+from controllers.NewsController import NewsController
 
 routes = Blueprint("routes", __name__)
+news_controller = NewsController()
 
 """
 return news without considering keywords
@@ -10,7 +11,7 @@
 
 @routes.route("/news", methods=["GET"])
 def getNews_route():
-    return getNews()
+    return news_controller.getNews()
 
 
 """
@@ -22,7 +23,7 @@ def getNews_route():
 def getNewsWithKeywords_route():
     # get list of keywords as argument from User's request
     user_keywords = request.args.getlist("keywords")
-    return getNewsWithKeywords(user_keywords)
+    return news_controller.getNewsWithKeywords(user_keywords)
 
 
 """
@@ -32,4 +33,4 @@ def getNewsWithKeywords_route():
 
 @routes.errorhandler(404)
 def notFound_route(error):
-    notFound(error)
+    news_controller.notFound(error)

services/NewsService.py

@@ -3,100 +3,95 @@
 from langchain.llms import OpenAI
 from langchain.chat_models import ChatOpenAI
 from langchain.prompts import PromptTemplate
-
-"""
-Create a LLM object
-"""
-OPENAI_API_KEY = dotenv_values(".env").get("OPENAI_API_KEY")
-llm = OpenAI(temperature=0.9, openai_api_key=OPENAI_API_KEY)
-
-"""
-Return news without considering keywords
-"""
-
-
-def getNews():
-    # fixed format for answer
-    answer_format = "news content (data source, date)"
-
-    # add prompt
-    prompt = PromptTemplate(
-        input_variables=["format"],
-        template="""
-        Please give me a list of the most recent cybersecurity news.
-        Indicate data source and date.
-        Return the answer with the format of '{format}'.
-        """,
-    )
-    news = llm(prompt.format(format=answer_format))
-    print(news)
-
-    # convert news data into JSON format
-    news_JSON = toJSON(news)
-
-    return news_JSON
-
-
-"""
-return news based on certain keywords
-"""
-
-
-def getNewsWithKeywords(user_keywords):
-    # fixed format for answer
-    answer_format = "news content (data source, date)"
-
-    # add prompt
-    prompt = PromptTemplate(
-        input_variables=["keywords", "format"],
-        template="""
-        Please give me a list of the most recent cybersecurity news with keywords of {keywords}.
-        Indicate data source and date.
-        Return the answer with the format of '{format}'.
-        """,
-    )
-    news = llm(prompt.format(keywords=user_keywords, format=answer_format))
-    print(news)
-
-    # convert news data into JSON format
-    news_JSON = toJSON(news)
-
-    return news_JSON
-
-
-"""
-deal requests with wrong route
-"""
-
-
-def notFound(error):
-    return jsonify({"error": "Not Found"}), 404
-
-
-"""
-Convert news given by OpenAI API into JSON format.
-"""
-
-
-def toJSON(data: str):
-    news_list = data.split("\n")
-    news_list_json = []
-
-    for item in news_list:
-        # Avoid dirty data
-        if len(item) == 0:
-            continue
-        # Split the string at the first occurrence of '('
-        title, remaining = item.split("(", 1)
-
-        # Extract the source by splitting at ',' and removing leading/trailing whitespace
-        source = remaining.split(",")[0].strip()
-
-        # Extract the date by splitting at ',' and removing leading/trailing whitespace
-        date = remaining.split(",")[1].strip().rstrip(")")
-
-        # Create a dictionary for each news item and append it to the news_list
-        news_item = {"title": title.strip(), "source": source, "date": date}
-        news_list_json.append(news_item)
-
-    return jsonify(news_list_json)
+from models.NewsModel import CybernewsDB
+
+
+class NewsService:
+    def __init__(self) -> None:
+        self.OPENAI_API_KEY = dotenv_values(".env").get("OPENAI_API_KEY")
+        self.llm = OpenAI(temperature=0.9, openai_api_key=self.OPENAI_API_KEY)
+        self.db = CybernewsDB()
+
+    """
+    Return news without considering keywords
+    """
+
+    def getNews(self):
+        # fixed format for answer
+        answer_format = "news content (data source, date)"
+
+        # add prompt
+        prompt = PromptTemplate(
+            input_variables=["format"],
+            template="""
+            Please give me a list of the most recent cybersecurity news.
+            Indicate data source and date.
+            Return the answer with the format of '{format}'.
+            """,
+        )
+        news = self.llm(prompt.format(format=answer_format))
+        print(news)
+
+        # convert news data into JSON format
+        news_JSON = self.toJSON(news)
+
+        return news_JSON
+
+    """
+    return news based on certain keywords
+    """
+
+    def getNewsWithKeywords(self, user_keywords):
+        # fixed format for answer
+        answer_format = "news content (data source, date)"
+
+        # add prompt
+        prompt = PromptTemplate(
+            input_variables=["keywords", "format"],
+            template="""
+            Please give me a list of the most recent cybersecurity news with keywords of {keywords}.
+            Indicate data source and date.
+            Return the answer with the format of '{format}'.
+            """,
+        )
+        news = self.llm(prompt.format(keywords=user_keywords, format=answer_format))
+        print(news)
+
+        # convert news data into JSON format
+        news_JSON = self.toJSON(news)
+
+        return news_JSON
+
+    """
+    deal requests with wrong route
+    """
+
+    def notFound(self, error):
+        return jsonify({"error": error}), 404
+
+    """
+    Convert news given by OpenAI API into JSON format.
+    """
+
+    def toJSON(self, data: str):
+        news_list = data.split("\n")
+        news_list_json = []
+
+        for item in news_list:
+            # Avoid dirty data
+            if len(item) == 0:
+                continue
+            # Split the string at the first occurrence of '('
+            title, remaining = item.split("(", 1)
+
+            # Extract the source by splitting at ',' and removing leading/trailing whitespace
+            source = remaining.split(",")[0].strip()
+
+            # Extract the date by splitting at ',' and removing leading/trailing whitespace
+            date = remaining.split(",")[1].strip().rstrip(")")
+
+            # Create a dictionary for each news item and append it to the news_list
+            news_item = {"title": title.strip(), "source": source, "date": date}
+            news_list_json.append(news_item)
+
+        return jsonify(news_list_json)