use mongodb for store chat history

This commit is contained in:
CounterFire2023 2024-12-18 11:58:04 +08:00
parent ee55f9f23f
commit ade4686b5e
7 changed files with 69 additions and 6 deletions

5
.env
View File

@ -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
OLLAMA_BOT_CHAT_IDS=5015834404
# DB_MAIN=mongodb://localhost:27017
DB_MAIN=mongodb://192.168.100.22:27017
DB_NAME=aicq_dev

View File

@ -1,3 +1,3 @@
{
"python.venvPath": "/Users/zhl/miniconda/envs/agent2/bin/python"
"python.venvPath": "/Users/zhl/miniconda/envs/agent/bin/python",
}

2
env/agent2.yaml vendored
View File

@ -1,4 +1,4 @@
name: agent2
name: agent
channels:
- defaults
- conda-forge

View File

@ -10,4 +10,5 @@ langchain_chroma==0.1.4
python-dotenv
python-telegram-bot
chromadb==0.4.22
sentence_transformers
sentence_transformers
pymongo==4.10.1

View File

@ -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"]})

59
llm/db_history.py Normal file
View File

@ -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:])

View File

@ -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