fix responses module

server/nlp/responses.py

@@ -4,9 +4,11 @@
 """
 
 import ast
+from typing import cast
 
 import numpy as np
 import openai
+from openai.types.chat import ChatCompletionMessageParam
 
 from server.config import OPENAI_API_KEY, OpenAIMessage, RedisDocument
 from server.nlp.embeddings import query_all
@@ -19,17 +21,12 @@
 def openai_response(thread: list[OpenAIMessage], sender: str) -> str:
     """Generate a response from OpenAI.
 
-    Parameters
-    ----------
-    thread: :obj:`list` of :obj:`OpenAIMessage`
-        previous email thread
-    sender: :obj:`str`
-        hacker email address
+    Args:
+        thread: previous email thread
+        sender: hacker email address
 
     Returns:
-    -------
-    str
-        email response
+        email response generated by
     """
     messages = [
         {
@@ -61,25 +58,25 @@ def openai_response(thread: list[OpenAIMessage], sender: str) -> str:
         }
     ]
 
+    messages = cast(list[ChatCompletionMessageParam], messages)
     response = openai.chat.completions.create(model=MODEL, messages=messages)
 
-    return response["choices"][0]["message"]["content"]
+    if response.choices[0].message.content is None:
+        return "openai unknown error"
+
+    return response.choices[0].message.content
 
 
 def openai_parse(email: str) -> list[str]:
     """Parse an email into questions using OpenAI.
 
-    Parameters
-    ----------
-    email : :obj:`str`
-        hacker email
+    Args:
+        email: hacker email
 
     Returns:
-    -------
-    :obj:`list` of :obj:`str`
-        parsed list of questions
+        list of questions parsed from the email
     """
-    response = openai.ChatCompletion.create(
+    response = openai.chat.completions.create(
         model=MODEL,
         messages=[
             {
@@ -92,7 +89,7 @@ def openai_parse(email: str) -> list[str]:
         ],
     )
     try:
-        questions = ast.literal_eval(response["choices"][0]["message"]["content"])
+        questions = ast.literal_eval(cast(str, response.choices[0].message.content))
         assert isinstance(questions, list)
         assert len(questions) > 0
         return questions
@@ -104,14 +101,10 @@ def openai_parse(email: str) -> list[str]:
 def confidence_metric(confidences: list[float]) -> float:
     """Compute confidence metric for a list of confidences.
 
-    Parameters
-    ----------
-    confidences : :obj:`list` of :obj:`float`
-        list of confidences
+    Args:
+        confidences: list of confidences
 
     Returns:
-    -------
-    float
         confidence metric
     """
     print("confidences", confidences)
@@ -123,20 +116,14 @@ def generate_context(
 ) -> tuple[list[OpenAIMessage], dict[str, list[RedisDocument]], float]:
     """Generate email context.
 
-    Parameters
-    ----------
-    email : :obj:`str`
-        hacker email
+    Args:
+        email: hacker email
 
     Returns:
-    -------
-    :obj:`list` of :obj:`OpenAIMessage`
-        list of contexts for all questions in email
-    :obj:`dict` of :obj:`[str, list[RedisDocument]]`
+        (list of contexts for all questions in email,
         dictionary mapping each question to list of context documents used to
-        answer question
-    :obj:`float`
-        confidence metric for all documents
+        answer question,
+        confidence metric for all documents)
     """
     questions = openai_parse(email)
     confidences = []
@@ -169,24 +156,16 @@ def generate_response(
 ) -> tuple[str, dict[str, list[RedisDocument]], float]:
     """Generate response to email.
 
-    Parameters
-    ----------
-    sender: :obj:`str`
-        hacker email address
-    email: :obj:`str`
-        newest incoming hacker email
-    thread : :obj:`list` of :obj:`OpenAIMessage`, optional
-        previous email thread
+    Args:
+        sender: hacker email address
+        email: newest incoming hacker email
+        thread : previous email thread
 
     Returns:
-    -------
-    str
-        email response
-    :obj:`dict` of :obj:`[str, list[RedisDocument]]`
+        (email response,
         dictionary mapping each question to list of context documents used to
-        answer question
-    float
-        confidence of response
+        answer question,
+        confidence of response)
     """
     if thread is None:
         thread = []