Refactor mail handling: update collect function to use MailReqData model and remove unused code

app/controllers/mail.py

@@ -6,7 +6,6 @@ from datetime import datetime, timedelta
 from venv import logger
 from ics import Calendar
 
-# import pandas as pd
 from langchain_core.documents import Document
 from langchain_community.document_loaders import (
     PyPDFLoader,
@@ -16,7 +15,6 @@ from langchain_community.document_loaders import (
 )
 
 from models.db import vectorstore
-# from models.mails import build_gmail_service
 
 SCOPES = ["https://www.googleapis.com/auth/gmail.readonly"]
 EMAIL_PATTERN = r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}"
@@ -78,6 +76,7 @@ def list_emails(service, messages):
     for message in messages:
         msg = service.users().messages().get(userId="me", id=message["id"], format="full").execute()
         metadata = {}
+        logger.info("vectorstore.index_to_docstore_id: %s", vectorstore.index_to_docstore_id)
         if msg["id"] in vectorstore.index_to_docstore_id:
             logger.info("Email already exists in the database.")
             continue

app/models/chroma/__init__.py

@@ -1,27 +0,0 @@
-# """Module for the Vector Database."""
-# from langchain_chroma import Chroma
-# from models.llm import EmbeddingsModel
-
-# vectorstore = Chroma(
-#     embedding_function=EmbeddingsModel("all-MiniLM-L6-v2"),
-#     collection_name="email",
-#     persist_directory="models/chroma/data"
-# )
-
-# # def create_or_get_collection(collection_name: str):
-# #     """
-# #     Creates a new collection or gets an existing collection from the Vector Database.
-
-# #     Args:
-# #         collection_name (str): The name of the collection.
-
-# #     Returns:
-# #         chromadb.Collection: The collection associated with the provided name.
-# #     """
-# #     chroma_client = chromadb.PersistentClient(path="models/chroma/data")
-# #     collection = chroma_client.get_or_create_collection(collection_name)
-# #     # try:
-# #     #     collection = chroma_client.create_collection(collection_name)
-# #     # except chromadb.errors.UniqueConstraintError:
-# #     #     collection = chroma_client.get_collection(collection_name)
-# #     return collection

app/models/mails/__init__.py

@@ -3,7 +3,6 @@ import os.path
 import pickle
 
 from google.auth.transport.requests import Request
-# from google.oauth2.credentials import Credentials
 from google_auth_oauthlib.flow import InstalledAppFlow
 from googleapiclient.discovery import build
 

app/playground/app.py

@@ -1,44 +0,0 @@
-"""Streamlit app example."""
-import logging
-import uuid
-import streamlit as st
-
-from chain import RAGChain
-from retriever import DocRetriever
-from controllers import mail
-
-logging.basicConfig(
-    format='%(asctime)s - %(levelname)s - %(funcName)s - %(message)s')
-logging.getLogger().setLevel(logging.ERROR)
-
-with st.sidebar:
-    st.header("Controls")
-    if st.button("Collect Data"):
-        result = mail.collect()
-        with st.chat_message("assistant"):
-            response_content = st.markdown(result)
-
-if 'chat_id' not in st.session_state:
-    st.session_state.chat_id = str(uuid.uuid4())
-    st.session_state.user_id = str(uuid.uuid4())
-
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
-if prompt := st.chat_input("What is up?"):
-
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    with st.chat_message("user"):
-        st.markdown(prompt)
-    req = {"query": prompt}
-    chain = RAGChain(DocRetriever(req=req))
-
-    result = chain.invoke({"input": req['query']},
-                    config={"configurable": {"session_id": st.session_state.chat_id}})
-    with st.chat_message("assistant"):
-        response_content = st.markdown(result['answer'])
-    st.session_state.messages.append({"role": "assistant", "content": result['answer']})

app/playground/phi-4-mini-instruct.py

@@ -1,23 +0,0 @@
-from langchain_huggingface import HuggingFacePipeline
-
-# Define the model ID
-model_id = "gpt2"
-model_id = "microsoft/Phi-4-mini-instruct"
-model_id = "Qwen/Qwen2.5-7B-Instruct"
-model_id = "microsoft/Phi-3-small-8k-instruct"
-
-# Create a pipeline for text generation
-llm = HuggingFacePipeline.from_model_id(
-    model_id=model_id,
-    task="text-generation",
-    device=-1,
-    # trust_remote_code=True,
-    pipeline_kwargs={
-        "max_new_tokens": 256,
-        "top_k": 50
-    },
-)
-
-# Use the model to generate text
-response = llm.invoke("Hello, how are you?")
-print(response)

app/playground/test.py

@@ -1,5 +0,0 @@
-from controllers import mail
-
-if __name__ == "__main__":
-    mail.collect()
-    # mail.get_documents()

app/router/mail.py

@@ -8,10 +8,12 @@ from controllers import mail
 from google.oauth2.credentials import Credentials
 from googleapiclient.discovery import build
 
+from schema import MailReqData
+
 router = APIRouter(prefix="/mail", tags=["mail"])
 
 @router.post("")
-def collect(email: str, request: Request):
+def collect(query: MailReqData, request: Request):
     """
     Handles the chat POST request.
 
@@ -22,8 +24,8 @@ def collect(email: str, request: Request):
         str: The generated response from the chat function.
     """
     try:
-        if os.path.exists(f"cache/{email}.pickle"):
-            with open(f"cache/{email}.pickle", "rb") as token:
+        if os.path.exists(f"cache/{query.email}.pickle"):
+            with open(f"cache/{query.email}.pickle", "rb") as token:
                 credentials = pickle.load(token)
         else:
             cred_dict = request.state.session.get("credential")
@@ -36,7 +38,7 @@ def collect(email: str, request: Request):
                 scopes=cred_dict["scopes"],
             )
         mailservice = build("gmail", "v1", credentials=credentials)
-        mail.collect(mailservice)
+        mail.collect(mailservice, query.query)
         return JSONResponse(content={"message": "Mail collected successfully."})
     except Exception as e:
         return JSONResponse(content={"error": str(e)}, status_code=500)

app/schema/__init__.py

@@ -19,6 +19,17 @@ class ReqData(BaseModel):
     user_id: str
     web: Optional[bool] = False
 
+class MailReqData(BaseModel):
+    """
+    MailReqData is a data model representing the structure of a mail request.
+
+    Attributes:
+        email (str): The email address of the sender.
+        query (str): The query or message content sent by the user.
+    """
+    email: str
+    query: str
+
 class ReqFollowUp(BaseModel):
     """
     RequestFollowUp is a Pydantic model that represents a request for follow-up.