document cli compiles

Author: azliu0

server/cli.py

@@ -10,6 +10,7 @@
     increment_response_count,
     thread_emails_to_openai_messages,
 )
+from server.models.document import Document
 from server.models.email import Email
 from server.models.response import Response
 from server.models.thread import Thread
@@ -28,8 +29,9 @@ def email():
     thread = Thread()
     db.session.add(thread)
     db.session.commit()
+
     email = Email(
-        date=datetime.datetime.now(datetime.timezone.utc),
+        date=datetime.datetime.now(datetime.timezone.utc),  # type: ignore
         sender="azliu@mit.edu",
         subject=subject,
         body=body,
@@ -46,14 +48,26 @@ def email():
         questions, document_ids, document_confidences = document_data(documents)
         db.session.add(email)
         db.session.commit()
+
+        documents = []
+        for i, doc_ids_question in enumerate(document_ids):
+            documents.append([])
+            for doc_id in doc_ids_question:
+                document = db.session.execute(
+                    db.select(Document).where(Document.id == doc_id)
+                ).scalar()
+                if document:
+                    documents[i].append(document)
+
         r = Response(
             openai_res,
             questions,
-            document_ids,
+            documents,
             document_confidences,
             confidence,
             email.id,
         )
+
         db.session.add(r)
         thread.last_email = email.id
         thread.resolved = False

server/controllers/emails.py

@@ -6,6 +6,7 @@
 import os
 import re
 from datetime import datetime, timezone
+from typing import List
 
 import boto3
 from apiflask import APIBlueprint
@@ -33,17 +34,13 @@
 emails = APIBlueprint("emails", __name__, url_prefix="/emails", tag="Emails")
 
 
-def thread_emails_to_openai_messages(thread_emails: list[Email]) -> list[OpenAIMessage]:
+def thread_emails_to_openai_messages(thread_emails: List[Email]) -> List[OpenAIMessage]:
     """Converts list of email to openai messages.
 
-    Parameters:
-    ----------
-    thread_email_ids : :obj:`list` of :obj:`Email`
-        list of email ids
+    Args:
+        thread_emails: list of emails
 
     Returns:
-    -------
-    :obj:`list` of :obj:`OpenAIMessage`
         list of openai messages
     """
     openai_messages = []
@@ -54,11 +51,11 @@ def thread_emails_to_openai_messages(thread_emails: list[Email]) -> list[OpenAIM
 
 
 def document_data(
-    documents: dict[str, list[RedisDocument]],
-) -> tuple[list[str], list[list[int]], list[list[float]]]:
+    documents: dict[str, List[RedisDocument]],
+) -> tuple[List[str], List[List[int]], List[List[float]]]:
     """Process raw openai document output.
 
-    Parameters
+    Args:
     ----------
     documents : :obj:`list` of :obj:`dict`
         raw openai document output
@@ -88,14 +85,13 @@ def document_data(
     return questions, doc_ids, doc_confidences
 
 
-def increment_response_count(document_ids: list[list[int]]):
+def increment_response_count(document_ids: List[List[int]]):
     """Increment response count for documents.
 
-    Parameters
-    ----------
-    document_ids : :obj:`list` of :obj:`list` of :obj:`int`
-        list of document ids. each element in the list is a list of document ids used to
-        answer a specific question
+    Args:
+        document_ids : :obj:`list` of :obj:`list` of :obj:`int`
+            list of document ids. each element in the list is a list of document ids
+            used to answer a specific question
     """
     for doc_ids_question in document_ids:
         for doc_id in doc_ids_question:
@@ -105,14 +101,13 @@ def increment_response_count(document_ids: list[list[int]]):
             db.session.commit()
 
 
-def decrement_response_count(document_ids: list[list[int]]):
+def decrement_response_count(document_ids: List[List[int]]):
     """Decrement response count for documents.
 
-    Parameters
-    ----------
-    document_ids : :obj:`list` of :obj:`list` of :obj:`int`
-        list of document ids. each element in the list is a list of document ids used to
-        answer a specific question
+    Args:
+        document_ids : :obj:`list` of :obj:`list` of :obj:`int`
+            list of document ids. each element in the list is a list of document ids
+            used to answer a specific question
     """
     for doc_ids_question in document_ids:
         for doc_id in doc_ids_question:

server/models/email.py

@@ -44,8 +44,8 @@ class Email(db.Model):
 
     is_reply: Mapped[bool] = mapped_column(nullable=False)
 
-    thread_id: Mapped[int] = mapped_column(ForeignKey("Threads.id"), nullable=False)
-    thread: Mapped[Thread] = relationship(back_populates="emails")
+    thread_id: Mapped[str] = mapped_column(ForeignKey("Threads.id"), nullable=False)
+    thread: Mapped[Thread] = relationship(back_populates="emails", init=False)
 
     def map(self):
         """Map the email to a dictionary."""

server/models/response.py

@@ -44,7 +44,7 @@ class Response(db.Model):
 
     confidence: Mapped[float] = mapped_column(nullable=False)
 
-    email_id: Mapped[int] = mapped_column(ForeignKey("Emails.id", ondelete="CASCADE"))
+    email_id: Mapped[str] = mapped_column(ForeignKey("Emails.id", ondelete="CASCADE"))
     email: Mapped[Email] = relationship(
         back_populates="response", init=False, single_parent=True
     )

server/models/thread.py

@@ -26,7 +26,7 @@ class Thread(db.Model):
     __tablename__ = "Threads"
 
     id: Mapped[str] = mapped_column(primary_key=True, init=False)
-    last_email: Mapped[Optional[str]] = mapped_column(nullable=True)
+    last_email: Mapped[Optional[str]] = mapped_column(nullable=True, init=False)
     resolved: Mapped[bool] = mapped_column(nullable=False, default=False)
 
     emails: Mapped[List[Email]] = relationship(