scaffolded welcome.summarize_prompt()

mind_palace/welcome.py

@@ -1,6 +1,7 @@
 from typing import List
 
 from docs import Section
+from llama_index.llms import ChatMessage, OpenAI
 
 
 def parse_abstracts(nodes) -> List[str]:
@@ -12,9 +13,24 @@ def parse_abstracts(nodes) -> List[str]:
     ]
 
 
-def summarize(texts):
-    # TODO:
-    return
+def summarize_prompt(abstracts: List[str]):
+    bullet_points = "\n".join([f"* {text}" for text in abstracts])
+    return {
+        "system": "You are a science journalist",
+        "user": f"Summarize these research paper into a single, short paragraph:\n'''{bullet_points}''''",
+    }
+
+
+def summarize(gpt_model, texts: List[str]):
+    prompt = summarize_prompt(texts)
+    messages = [
+        ChatMessage(role="system", content=prompt["system"]),
+        ChatMessage(role="user", content=prompt["user"]),
+    ]
+    resp = OpenAI(model=gpt_model).chat(messages)
+
+    # TODO: parse response
+    return resp
 
 
 def get_welcome_message(nodes):

tests/unit/test_welcome.py

@@ -15,3 +15,14 @@ def test_parse_abstracts():
     assert len(abstracts) == 1
     assert isinstance(abstracts, list)
     assert all(isinstance(abstract, str) for abstract in abstracts)
+
+
+def test_summarize_prompt():
+    abstracts = ["this is abstract", "second abstract"]
+    prompt = w.summarize_prompt(abstracts)
+    assert isinstance(prompt, dict)
+    assert prompt["system"] == "You are a science journalist"
+    assert (
+        prompt["user"]
+        == "Summarize these research paper into a single, short paragraph:\n'''* this is abstract\n* second abstract''''"
+    )