見出し画像

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にしましたが、もちろん違うモデルでもできると思います

いいなと思ったら応援しよう!