From ade4686b5e7efdc263fe8bbad507f6751228f88e Mon Sep 17 00:00:00 2001
From: CounterFire2023 <136581895+CounterFire2023@users.noreply.github.com>
Date: Wed, 18 Dec 2024 11:58:04 +0800
Subject: [PATCH] use mongodb for store chat history

---
 .env                  |  5 +++-
 .vscode/settings.json |  2 +-
 env/agent2.yaml       |  2 +-
 env/requirements.txt  |  3 ++-
 llm/agent_util.py     |  2 +-
 llm/db_history.py     | 59 +++++++++++++++++++++++++++++++++++++++++++
 llm/ollama.py         |  2 +-
 7 files changed, 69 insertions(+), 6 deletions(-)
 create mode 100644 llm/db_history.py

diff --git a/.env b/.env
index f0cbd71..5d04709 100644
--- a/.env
+++ b/.env
@@ -1,4 +1,7 @@
 TAVILY_API_KEY=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
 OLLAMA_BOT_TOKEN=7965398310:AAF2QDGXRRwDta0eYC7lFSvRDLW7L-SQH_E
 OLLAMA_BOT_TOKEN=7802269927:AAHXso9DZe5_LgQ4kN9J5jEAgJpMNQVaYI0
-OLLAMA_BOT_CHAT_IDS=5015834404
\ No newline at end of file
+OLLAMA_BOT_CHAT_IDS=5015834404
+# DB_MAIN=mongodb://localhost:27017
+DB_MAIN=mongodb://192.168.100.22:27017
+DB_NAME=aicq_dev
\ No newline at end of file
diff --git a/.vscode/settings.json b/.vscode/settings.json
index c9f7a98..7849afd 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,3 +1,3 @@
 {
-  "python.venvPath": "/Users/zhl/miniconda/envs/agent2/bin/python"
+  "python.venvPath": "/Users/zhl/miniconda/envs/agent/bin/python",
 }
\ No newline at end of file
diff --git a/env/agent2.yaml b/env/agent2.yaml
index 9b19e17..7560c80 100644
--- a/env/agent2.yaml
+++ b/env/agent2.yaml
@@ -1,4 +1,4 @@
-name: agent2
+name: agent
 channels:
   - defaults
   - conda-forge
diff --git a/env/requirements.txt b/env/requirements.txt
index da04f82..7db0c55 100644
--- a/env/requirements.txt
+++ b/env/requirements.txt
@@ -10,4 +10,5 @@ langchain_chroma==0.1.4
 python-dotenv
 python-telegram-bot
 chromadb==0.4.22
-sentence_transformers
\ No newline at end of file
+sentence_transformers
+pymongo==4.10.1
\ No newline at end of file
diff --git a/llm/agent_util.py b/llm/agent_util.py
index d7ea63e..5313789 100644
--- a/llm/agent_util.py
+++ b/llm/agent_util.py
@@ -152,7 +152,7 @@ def config_agent(llama3, llama3_json):
         # TODO:: 根据context特定的内容生产答案
         # check if context is not None and not empty
         if order_info is not None:
-            return {"generation": order_info}
+            return {"generation": "order_info::"+order_info}
         else:
             generation = generate_chain.invoke(
                 {"context": context, "question": question, "history": state["history"]})
diff --git a/llm/db_history.py b/llm/db_history.py
new file mode 100644
index 0000000..12a5f15
--- /dev/null
+++ b/llm/db_history.py
@@ -0,0 +1,59 @@
+from pymongo import MongoClient
+import os
+from datetime import datetime
+
+message_template = """
+User: {question}
+
+Answer: {answer}
+"""
+
+COLLECTION_NAME = 'chat_history'
+
+class ChatHistory:
+    def __init__(self) -> None:
+        self.history = {}
+        db_url = os.environ.get('DB_MAIN')
+        db_name = os.environ.get('DB_NAME')
+        self.db_client = MongoClient(db_url)
+        self.db = self.db_client[db_name]
+        self.collection = self.db[COLLECTION_NAME]
+        self.collection.create_index('user')
+
+    def append(self, question: str, answer: str, uid: str):
+        new_message = message_template.format(
+            question=question,
+            answer=answer
+        )
+        if uid not in self.history:
+            self.history[uid] = []
+        self.history[uid].append(new_message)
+        self.collection.insert_one({
+            'user': uid,
+            'question': question,
+            'answer': answer,
+            'created_at': int(datetime.now().timestamp())
+        })
+
+    def delete(self, uid: str):
+        if uid in self.history:
+            del self.history[uid]
+        self.collection.delete_many({'user': uid})
+    
+
+    def get_context(self, question: str, uid: str):
+        records = self.collection.find({'user': uid}).sort("_id", -1).limit(10)
+        # transform records into list of message_template format
+        documents = []
+        for record in records:
+            documents.append(message_template.format(
+                question=record['question'],
+                answer=record['answer']
+            ))
+        context = '\n'.join([d for d in documents])
+        return context
+    
+    def get_last_few(self, uid: str):
+        if uid not in self.history:
+            return ''
+        return '\n'.join(self.history[uid][-3:])
\ No newline at end of file
diff --git a/llm/ollama.py b/llm/ollama.py
index ce91d04..eeb3931 100644
--- a/llm/ollama.py
+++ b/llm/ollama.py
@@ -2,7 +2,7 @@ from ollama import Client
 # from langchain.llms.ollama import Ollama
 # from langchain.llms import Ollama
 from langchain_community.chat_models import ChatOllama
-from .chat_history import ChatHistory
+from .db_history import ChatHistory
 from .agent_util import config_agent
 import os