Remove unused binary files and refactor main application structure to integrate FastAPI with new routing and utility functions.

chain/__init__.py

@@ -4,43 +4,19 @@ import json
 from datetime import datetime
 from venv import logger
 
-import torch
 from pymongo import errors
 from langchain_core.runnables.history import RunnableWithMessageHistory
+from langchain_core.output_parsers import PydanticOutputParser
 from langchain_core.messages import BaseMessage, message_to_dict
 from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain.chains.retrieval import create_retrieval_chain
 from langchain.prompts.chat import ChatPromptTemplate, MessagesPlaceholder
 from langchain_mongodb import MongoDBChatMessageHistory
-from langchain_huggingface import HuggingFacePipeline
-
-from models.llm import GPTModel, Phi4MiniONNXLLM, HuggingfaceModel
-
-# llm = GPTModel()
-# REPO_ID = "microsoft/Phi-4-mini-instruct-onnx"
-# SUBFOLDER = "cpu_and_mobile/cpu-int4-rtn-block-32-acc-level-4"
-# llm = Phi4MiniONNXLLM(REPO_ID, SUBFOLDER)
-
-# MODEL_NAME = "openai-community/gpt2"
-MODEL_NAME = "microsoft/phi-1_5"
-# llm = HuggingfaceModel(MODEL_NAME)
-
-hf_llm = HuggingFacePipeline.from_model_id(
-    model_id="microsoft/Phi-4",
-    task="text-generation",
-    pipeline_kwargs={
-        "max_new_tokens": 128,
-        "temperature": 0.3,
-        "top_k": 50,
-        "do_sample": True
-    },
-    model_kwargs={
-        "torch_dtype": "auto",
-        "device_map": torch.device("cuda" if torch.cuda.is_available() else "cpu"),
-        "max_memory": {0: "10GB"},
-        "use_cache": False
-    }
-)
+
+from schema import FollowUpQ
+from models.llm import GPTModel
+
+llm = GPTModel()
 
 SYS_PROMPT = """You are a knowledgeable financial professional. You can provide well elaborated and credible answers to user queries in economic and finance by referring to retrieved contexts.
             You should answer user queries strictly following the instructions below, and do not provide anything irrelevant. \n
@@ -108,7 +84,7 @@ def get_message_history(
     """
     return MessageHistory(
         session_id = session_id,
-        connection_string=str(mongo_url), database_name='emails')
+        connection_string=str(mongo_url), database_name='mailbox')
 
 class RAGChain(RunnableWithMessageHistory):
     """
@@ -130,3 +106,50 @@ class RAGChain(RunnableWithMessageHistory):
             history_messages_key="chat_history",
             output_messages_key="answer"
         )
+
+class FollowUpChain():
+    """
+    FollowUpQChain is a class to generate follow-up questions based on contexts and initial query.
+
+    Attributes:
+        parser (PydanticOutputParser): An instance of PydanticOutputParser to parse the output.
+        chain (Chain): A chain of prompts and models to generate follow-up questions.
+
+    Methods:
+        __init__():
+            Initializes the FollowUpQChain with a parser and a prompt chain.
+        
+        invoke(contexts, query):
+            Invokes the chain with the provided contexts and query to generate follow-up questions.
+
+                contexts (str): The contexts to be used for generating follow-up questions.
+                query (str): The initial query to be used for generating follow-up questions.
+    """
+    def __init__(self):
+        self.parser = PydanticOutputParser(pydantic_object=FollowUpQ)
+        prompt = ChatPromptTemplate.from_messages([
+                    ("system", "You are a professional commentator on current events.Your task\
+                      is to provide 3 follow-up questions based on contexts and initial query."),
+                    ("system", "contexts: {contexts}"),
+                    ("system", "initial query: {query}"),
+                    ("human", "Format instructions: {format_instructions}"),
+                    ("placeholder", "{agent_scratchpad}"),
+                ])
+        self.chain = prompt | llm | self.parser
+
+    def invoke(self, query, contexts):
+        """
+        Invokes the chain with the provided content and additional parameters.
+
+        Args:
+            content (str): The article content to be processed.
+
+        Returns:
+            The result of the chain invocation.
+        """
+        result = self.chain.invoke({
+            'contexts': contexts,
+            'format_instructions': self.parser.get_format_instructions(),
+            'query': query
+        })
+        return result.questions

main.py

@@ -1,20 +1,36 @@
-"""Module to run the mail collection process."""
-from dotenv import load_dotenv
-
-# from controllers import mail
-from chain import RAGChain
-from retriever import DocRetriever
-
-load_dotenv()
-
-if __name__ == "__main__":
-    # mail.collect()
-    # mail.get_documents()
-    req = {
-        "query": "Just give me an update?",
-    }
-    chain = RAGChain(DocRetriever(req=req))
-    result = chain.invoke({"input": req['query']},
-                       config={"configurable": {"session_id": "20250301"}})
-    print(result)
-    print(result.get("answer"))
+"""Module to handle the main FastAPI application and its endpoints."""
+import logging
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from router import main
+
+
+app = FastAPI(docs_url="/")
+
+app.include_router(main.router, tags=["content"])
+
+origins = [
+    "*"
+]
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials = True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+logging.basicConfig(
+    format='%(asctime)s - %(levelname)s - %(funcName)s - %(message)s')
+logging.getLogger().setLevel(logging.ERROR)
+
+
+@app.get("/_health")
+def health():
+    """
+    Returns the health status of the application.
+
+    :return: A string "OK" indicating the health status.
+    """
+    return "OK"

retriever/__init__.py

@@ -23,9 +23,9 @@ class DocRetriever(BaseRetriever):
                 list: A list of Document objects with relevant metadata.
     """
     retriever: VectorStoreRetriever = None
-    k: int = 5
+    k: int = 3
 
-    def __init__(self, req, k: int = 2) -> None:
+    def __init__(self, req, k: int = 3) -> None:
         super().__init__()
         # _filter={}
         # if req.site != []:
@@ -52,7 +52,7 @@ class DocRetriever(BaseRetriever):
                     metadata = {
                         "content": doc.page_content,
                         # "id": doc.metadata['id'],
-                        # "title": doc.metadata['title'],
+                        "title": doc.metadata['subject'],
                         # "site": doc.metadata['site'],
                         # "link": doc.metadata['link'],
                         # "publishDate": doc.metadata['publishDate'].strftime('%Y-%m-%d'),

router/main.py

@@ -0,0 +1,76 @@
+"""Module for defining the main routes of the API."""
+from fastapi import APIRouter
+from fastapi.responses import StreamingResponse
+from schema import ReqData
+from utils import generate
+
+router = APIRouter()
+
+@router.post("/stream")
+async def stream(query: ReqData):
+    """
+    Handles streaming of data based on the provided query.
+
+    Args:
+        query (ReqData): The request data containing the query parameters.
+
+    Returns:
+        StreamingResponse: A streaming response with generated data with type 'text/event-stream'.
+    """
+    return StreamingResponse(generate(query), media_type='text/event-stream')
+
+# # @router.post("/followup")
+# # def follow_up(req: ReqFollowUp):
+# #     """
+# #     Handles the follow-up POST request.
+
+# #     Args:
+# #         req (ReqFollowUp): The request object containing follow-up data.
+
+# #     Returns:
+# #         Response: The response from the follow-up processing function.
+# #     """
+# #     return followup(req)
+
+# @router.post("/chat/history")
+# def retrieve_history(chat_history: ChatHistory):
+#     """
+#     Endpoint to retrieve chat history.
+
+#     This endpoint handles POST requests to the "/chat/history" URL. It accepts a 
+#     ChatHistory object as input and returns the chat history.
+
+#     Args:
+#         chat_history (ChatHistory): The chat history object containing the details 
+#         of the chat to be retrieved.
+
+#     Returns:
+#         The chat history retrieved by the retrieve_chat_history function.
+#     """
+#     return get_chat_history(chat_history)
+
+# @router.post("/chat/session")
+# def retrieve_session(chat_session: ChatSession):
+#     """
+#     Retrieve a chat session.
+
+#     Args:
+#         chat_session (ChatSession): The chat session to retrieve.
+
+#     Returns:
+#         ChatSession: The retrieved chat session.
+#     """
+#     return get_chat_session(chat_session)
+
+# @router.post("/chat/history/clear")
+# def clear_history(chat_history: ChatHistory):
+#     """
+#     Clears the chat history.
+
+#     Args:
+#         chat_history (ChatHistory): The chat history object to be cleared.
+
+#     Returns:
+#         The result of the clear_chat_history function.
+#     """
+#     return clear_chat(chat_history)

schema/__init__.py

@@ -0,0 +1,60 @@
+"""Module containing the data models for the application."""
+from typing import Optional, List
+from pydantic import BaseModel, Field
+
+class ReqData(BaseModel):
+    """
+    RequestData is a Pydantic model that represents the data structure for a request.
+
+    Attributes:
+        query (str): The query string provided by the user.
+        chat_id (str): The unique identifier for the chat session.
+        user_id (str): The unique identifier for the user.
+        web (Optional[bool]): A flag indicating if the request is from the web. Defaults to False.
+    """
+    query: str
+    id: Optional[List[str]] = []
+    site: Optional[List[str]] = []
+    chat_id: str
+    user_id: str
+    web: Optional[bool] = False
+
+class ReqFollowUp(BaseModel):
+    """
+    RequestFollowUp is a Pydantic model that represents a request for follow-up.
+
+    Attributes:
+        query (str): The query string that needs follow-up.
+        contexts (list[str]): A list of context strings related to the query.
+    """
+    query: str
+    contexts: list[str]
+
+class FollowUpQ(BaseModel):
+    """
+    FollowUpQ model to represent a follow-up question based on context information.
+
+    Attributes:
+        question (list[str]): A list of follow-up questions based on context information.
+    """
+    questions: list[str] = Field(..., description="3 Follow up questions based on context.")
+
+class ChatHistory(BaseModel):
+    """
+    ChatHistory model representing a chat session.
+
+    Attributes:
+        chat_id (str): The unique identifier for the chat session.
+        user_id (str): The unique identifier for the user.
+    """
+    chat_id: str
+    user_id: str
+
+class ChatSession(BaseModel):
+    """
+    ChatSession model representing a chat session.
+
+    Attributes:
+        user_id (str): The unique identifier for the user.
+    """
+    user_id: str

utils.py

@@ -0,0 +1,36 @@
+"""Module containing utility functions for the chatbot application."""
+import json
+from chain import RAGChain, FollowUpChain
+from schema import ReqData
+from retriever import DocRetriever
+
+followUpChain = FollowUpChain()
+
+async def generate(req: ReqData):
+    """
+    Asynchronously generates responses based on the provided request data.
+
+    This function uses different processing chains depending on the `web` attribute of the request.
+    It streams chunks of data and yields server-sent events (SSE) for answers and contexts.
+    Additionally, it generates follow-up questions and updates citations.
+
+    Args:
+        req (ReqData): Request data containing user and chat info, query, and other parameters.
+
+    Yields:
+        str: Server-sent events (SSE) for answers, contexts, and follow-up questions in JSON format.
+    """
+    chain = RAGChain(DocRetriever(req=req))
+    session_id = "/".join([req.user_id, req.chat_id])
+    contexts = []
+    for chunk in chain.stream({"input": req.query},
+                                   config={"configurable": {"session_id": session_id}}):
+        if 'answer' in chunk:
+            yield "event: answer\n"
+            yield f"data: {json.dumps(chunk)}\n\n"
+        elif 'context' in chunk:
+            for context in chunk['context']:
+                yield "event: context\n"
+                yield f"data: {json.dumps({'context': context.metadata})}\n\n"
+    yield "event: questions\n"
+    yield f"data: {json.dumps({'questions': followUpChain.invoke(req.query, contexts)})}\n\n"