ollamaのチャットでmem0を適用する
import os
import time
from datetime import datetime
from mem0 import Memory
import ollama
# pip install chromadb
# openaiはembeddingで使う
os.environ['OPENAI_API_KEY'] = "************************************************"
system_content = f"""<<SYS>>あなたは優秀なAIアシスタントです。質問に答えてください。
現在時刻は{datetime.now()}です。<</SYS>>"""
config = {
"vector_store": {
"provider": "chroma",
"config": {
"collection_name": "chat",
"path": "chatdb",
}
}
}
# llama-3.1:latest
model = "llama3.1:latest"
class PersonalAssistant:
def __init__(self):
self.client = ollama.Client()
self.memory = Memory.from_config(config)
self.messages = [{"role": "system", "content": system_content}]
def ask_question(self, question, user_id):
# Fetch previous related memories
previous_memories = self.search_memories(question, user_id=user_id)
prompt = question
if previous_memories:
prompt = f"User input: {question}\n Previous memories: {previous_memories}"
self.messages.append({"role": "user", "content": prompt})
# Generate response using llama3.1:latest
stream = self.client.chat(
model=model,
messages=self.messages,
stream=True,
)
answer = ""
for chunk in stream:
tok = chunk['message']['content']
answer += tok
print(tok, end='')
if chunk['done'] == 'ture':
break
print() # 改行
self.messages.append({"role": "assistant", "content": answer})
# Store the question in memory
self.memory.add(question, user_id=user_id)
return answer
def get_memories(self, user_id):
memories = self.memory.get_all(user_id=user_id)
return [m["memory"] for m in memories]
def search_memories(self, query, user_id):
memories = self.memory.search(query, user_id=user_id)
return [m["memory"] for m in memories]
# Usage example
user_id = "user_123"
ai_assistant = PersonalAssistant()
def main():
while True:
question = input("You: ")
if question.lower() in ['']:
print("exit...")
break
ai_assistant.ask_question(question, user_id=user_id)
memories = ai_assistant.get_memories(user_id=user_id)
print("Memories:")
for memory in memories:
print(f"- {memory}")
print("-----")
time.sleep(0.3)
if __name__ == "__main__":
main()
llama-cpp-pythonでもできないことはないのですが、mem0ai公式がollamaを採用しているのでollamaのコードにします
PCにollamaがインストールされていることが前提になります
openaiのapi_keyはembeddingで使うので必要です。embeddingまでローカルで済ます方向には来ています。
今回はllama3.1-latestにしましたが、もちろん違うモデルでもできると思います