conversation history

Author: EdmundKorley

.env.example

@@ -1 +1,8 @@
+# Copy this file to .env and replace the values with your own
 OPENAI_API_KEY=XXXXXXXXXX
+DATABASE_URL=postgres://llm_agent:llm_password@localhost:5432/llm_db
+DB_NAME=llm_db
+DB_USER=llm_agent
+DB_PASSWORD=llm_password
+DB_HOST=localhost
+DB_PORT=5432

Pipfile

@@ -11,6 +11,8 @@ tiktoken = "*"
 djangorestframework = "*"
 pypdf = "*"
 openai = "*"
+python-dotenv = "*"
+psycopg = "*"
 
 [dev-packages]
 

Pipfile.lock

@@ -1,5 +1,8 @@
 import os
 import django
+import secrets
+import string
+
 from rest_framework import status
 from rest_framework.decorators import api_view
 from rest_framework.response import Response
@@ -16,12 +19,16 @@
 
 @api_view(["POST"])
 def create_chat(request):
-    question = request.data.get("question").strip()
+    prompt = request.data.get("prompt").strip()
+    session_id = (request.data.get("session_id") or generate_short_id()).strip()
 
-    answer = run_chat_chain(question)
+    answer = run_chat_chain(prompt=prompt, session_id=session_id)
 
     return Response(
-        {"answer": answer},
+        {
+            "answer": answer,
+            "session_id": session_id,
+        },
         status=status.HTTP_201_CREATED,
     )
 
@@ -36,4 +43,9 @@ def create_embeddings(request):
     return Response(
         f"Created {document_file_name} index",
         status=status.HTTP_201_CREATED,
-    )
+    )
+
+
+def generate_short_id(length=6):
+    alphanumeric = string.ascii_letters + string.digits
+    return ''.join(secrets.choice(alphanumeric) for _ in range(length))

llm/api.py

@@ -0,0 +1,6 @@
+from django.apps import AppConfig
+
+
+class LlmConfig(AppConfig):
+    default_auto_field = "django.db.models.BigAutoField"
+    name = "llm"

llm/apps.py

@@ -1,22 +1,25 @@
-from llm.chains import config, embeddings
+from llm.chains import config, embeddings, memory
 
 from langchain.chat_models import ChatOpenAI
 from langchain.chains import RetrievalQA
 
 
-def run_chat_chain(prompt: str):
+def run_chat_chain(prompt: str, session_id: str):
     llm = get_llm()
     retriever = embeddings.get_retriever()
-    qa_chain = RetrievalQA.from_chain_type(
-        llm=llm, retriever=retriever, chain_type="stuff", return_source_documents=True
+    chain_memory = memory.conversation_history(session_id=session_id)
+    chain = RetrievalQA.from_chain_type(
+        llm=llm,
+        retriever=retriever,
+        chain_type="stuff",
+        return_source_documents=True,
+        memory=chain_memory,
     )
-    result = qa_chain({"query": prompt})
-
-    print(result)
+    result = chain({"query": prompt})
 
     return result
 
 
 def get_llm():
     llm = ChatOpenAI(model_name="gpt-3.5-turbo-16k", temperature=config.BOT_TEMPATURE)
-    return llm
+    return llm

llm/chains/chat.py

@@ -0,0 +1,19 @@
+import os
+
+from langchain.memory import PostgresChatMessageHistory, ConversationBufferMemory
+
+
+def conversation_history(session_id: str):
+    history = PostgresChatMessageHistory(
+      connection_string=os.getenv("DATABASE_URL"),
+      session_id=session_id,
+    )
+    memory = ConversationBufferMemory(
+      memory_key="chat_history",
+      return_messages=True,
+      input_key="query",
+      output_key="result",
+      chat_memory=history
+    )
+    return memory
+

llm/chains/memory.py

@@ -0,0 +1,25 @@
+# Generated by Django 4.2.4 on 2023-08-30 20:46
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    initial = True
+
+    dependencies = [
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='MessageStore',
+            fields=[
+                ('id', models.AutoField(primary_key=True, serialize=False)),
+                ('session_id', models.TextField()),
+                ('message', models.JSONField()),
+            ],
+            options={
+                'db_table': 'message_store',
+            },
+        ),
+    ]

llm/migrations/0001_initial.py

@@ -0,0 +1,12 @@
+from django.db import models
+
+class MessageStore(models.Model):
+    """
+    See Langchain source code for the table schema we match here: https://github.com/langchain-ai/langchain/blob/7fa82900cb15d9c41099ad7dbb8aaa66941f6905/libs/langchain/langchain/memory/chat_message_histories/postgres.py#L39-L42
+    """
+    id = models.AutoField(primary_key=True)
+    session_id = models.TextField()
+    message = models.JSONField()
+
+    class Meta:
+        db_table = "message_store"

llm/migrations/__init__.py

@@ -9,9 +9,12 @@
 For the full list of settings and their values, see
 https://docs.djangoproject.com/en/4.2/ref/settings/
 """
-
+import os
 from pathlib import Path
 
+from dotenv import load_dotenv
+load_dotenv()
+
 # Build paths inside the project like this: BASE_DIR / 'subdir'.
 BASE_DIR = Path(__file__).resolve().parent.parent
 
@@ -31,6 +34,7 @@
 # Application definition
 
 INSTALLED_APPS = [
+    'llm.apps.LlmConfig',
     'django.contrib.admin',
     'django.contrib.auth',
     'django.contrib.contenttypes',
@@ -75,8 +79,12 @@
 
 DATABASES = {
     'default': {
-        'ENGINE': 'django.db.backends.sqlite3',
-        'NAME': BASE_DIR / 'db.sqlite3',
+        'ENGINE': 'django.db.backends.postgresql_psycopg2',
+        'NAME': os.getenv("DB_NAME"),
+        'USER': os.getenv("DB_USER"),
+        'PASSWORD': os.getenv("DB_PASSWORD"),
+        'HOST': os.getenv("DB_HOST"),
+        'PORT': os.getenv("DB_PORT"),
     }
 }