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,29 @@ 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 summarizing papers for your readers.\n"
+            "Instructions:\n"
+            "respond with fewer than 100 words\n"
+            "start your response with 'This collection of papers'"
+        ),
+        "user": f"Summarize these research papers:\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

@@ -1,3 +1,5 @@
+import re
+
 from tests.context import docs
 from tests.context import welcome as w
 
@@ -15,3 +17,11 @@ 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 isinstance(prompt["system"], str)
+    assert re.search(r"'''\* this is abstract\n\* second abstract'''", prompt["user"])