From a26c51014bf771f3491fe3d29ad39504967a54fb Mon Sep 17 00:00:00 2001
From: Kenneth Reitz <me@kennethreitz.org>
Date: Wed, 6 Nov 2024 08:08:43 -0500
Subject: [PATCH] Refactor enhanced_context.py and update requirements.txt

---
 examples/enhanced_context.py | 99 ++++++++++++++++++++++++++++++++++++
 examples/requirements.txt    |  3 ++
 2 files changed, 102 insertions(+)
 create mode 100644 examples/enhanced_context.py

diff --git a/examples/enhanced_context.py b/examples/enhanced_context.py
new file mode 100644
index 0000000..c475785
--- /dev/null
+++ b/examples/enhanced_context.py
@@ -0,0 +1,99 @@
+from datetime import datetime
+
+import spacy
+import sqlite3
+
+from _context import simplemind as sm
+
+
+class EnhancedContextPlugin(sm.BasePlugin):
+    def __init__(self):
+        super().__init__()
+        # Initialize NLP model and memory database
+        object.__setattr__(self, "nlp", spacy.load("en_core_web_sm"))
+        object.__setattr__(self, "conn", sqlite3.connect(":memory:"))
+
+        self.init_db()
+
+    def init_db(self):
+        # Create a table to store entities and their last mention time
+        with self.conn:
+            self.conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS memory (
+                    entity TEXT PRIMARY KEY,
+                    last_mentioned TIMESTAMP
+                )
+            """
+            )
+
+    def store_entity(self, entity):
+        # Store or update entity mention time
+        print(f"Storing entity in memory: {entity}")
+        with self.conn:
+            self.conn.execute(
+                """
+                INSERT OR REPLACE INTO memory (entity, last_mentioned)
+                VALUES (?, ?)
+            """,
+                (entity, datetime.now()),
+            )
+
+    def retrieve_recent_entities(self):
+        # Retrieve entities mentioned in the last 7 days
+        cur = self.conn.cursor()
+        cur.execute(
+            """
+            SELECT entity FROM memory
+            WHERE last_mentioned >= datetime('now', '-7 days')
+        """
+        )
+        return [row[0] for row in cur.fetchall()]
+
+    def extract_entities(self, text):
+        # Extract entities (people, places, organizations) from text
+        doc = self.nlp(text)
+        return [
+            ent.text
+            for ent in doc.ents
+            if ent.label_ in {"PERSON", "ORG", "GPE", "NORP"}
+        ]
+
+    def pre_send_hook(self, conversation: sm.Conversation):
+        # Process the latest user message
+        last_message = conversation.get_last_message(role="user")
+        if last_message:
+            # Extract entities and store in memory
+            entities = self.extract_entities(last_message.text)
+            print(f"Extracted entities: {entities}")
+            for entity in entities:
+                self.store_entity(entity)
+
+            # Retrieve recent entities for context
+            recent_entities = self.retrieve_recent_entities()
+            if recent_entities:
+                print(f"Recent entities found: {recent_entities}")
+                context_message = f"Here are some topics recently discussed: {', '.join(recent_entities)}. Feel free to bring them up if relevant."
+                conversation.add_message(role="system", text=context_message)
+
+
+# Create a conversation and add the plugin
+conversation = sm.create_conversation(llm_model="gpt-4o-mini", llm_provider="openai")
+conversation.add_plugin(EnhancedContextPlugin())
+
+
+# Replace the single message test with an interactive chat loop
+def chat():
+    print("Welcome to the enhanced context chat! Type 'quit' to exit.")
+    while True:
+        user_input = input("\nYou: ").strip()
+        if user_input.lower() in ["quit", "exit"]:
+            break
+
+        conversation.add_message(role="user", text=user_input)
+        response = conversation.send()
+        print(f"\nAssistant: {response.text!r}")
+
+
+if __name__ == "__main__":
+    chat()
diff --git a/examples/requirements.txt b/examples/requirements.txt
index 476de75..628a084 100644
--- a/examples/requirements.txt
+++ b/examples/requirements.txt
@@ -1,6 +1,9 @@
+# python -m spacy download en_core_web_sm
+
 numpy
 openai
 pydantic
 faiss-cpu
 rich
 nltk
+spacy