import numpy as np
from ollama import embeddings
class Kb:
    def __init__(self, filepath):
        # 讀取文件內(nèi)容
        content = self.read_file(filepath)
        # print(content)
        # 讀取拆分好的數(shù)組
        self.chunks = self.split_content(content)
        # print(chunks)
        # for chunk in chunks:
        #     print(chunk)
        #     print('=' * 10)
        # 轉(zhuǎn)換成向量
        self.embeds = self.get_embeddings(self.chunks)
    # 讀取文件
    def read_file(self, filepath):
        with open(filepath, 'r', encoding='utf-8') as f:
            content = f.read()
        return content
    # 拆分知識庫
    @staticmethod
    def split_content(content):
        chunks = content.split('# ')
        # 過濾掉空塊
        chunks = [chunk.strip() for chunk in chunks if chunk.strip()]
        return chunks
    # 字符串轉(zhuǎn)向量(embeddings)
    def get_embedding(self, chunk):
        # milkey/m3e    0.642084887746903
        # bge-m3    0.6073383067378445
        # nomic-embed-text  完全找不到
        res = embeddings(model='milkey/m3e', prompt=chunk)
        # print(chunk)
        # print(res)
        # print(res['embedding'])
        return res['embedding']
    def get_embeddings(self, chunks):
        embeds = []
        for chunk in chunks:
            embed = self.get_embedding(chunk)
            embeds.append(embed)
        return np.array(embeds)
    # 查詢相似性向量
    def search(self, text):
        print(text)
        max_similarity = 0
        max_similarity_index = 0
        ask_embed = self.get_embedding(text)
        for kb_embed_index, kb_embed in enumerate(self.embeds):
            similarity = self.similarity(kb_embed, ask_embed)
            # print(similarity)
            # print(self.chunks[kb_embed_index])
            if similarity > max_similarity:
                max_similarity = similarity
                max_similarity_index = kb_embed_index
        print(max_similarity)
        print(self.chunks[max_similarity_index])
        # print(self.embeds[max_similarity_index])
        # 返回查到的相關(guān)文本
        return self.chunks[max_similarity_index]
    # 相似度
    @staticmethod
    def similarity(A, B):
        # 計算點積
        dot_product = np.dot(A, B)
        # 計算范數(shù)
        norm_A = np.linalg.norm(A)
        norm_B = np.linalg.norm(B)
        # 計算余弦相似度
        cosine_sim = dot_product / (norm_A * norm_B)
        return cosine_sim

from kb import Kb
from ollama import chat, Message
class Rag:
    def __init__(self, model, kb_filepath):
        self.kb_filepath = kb_filepath
        self.kb = Kb(kb_filepath)
        self.model = model
        self.prompt_template = """
        基于：%s
        回答：%s
        """
    def chat(self, message):
        # 用戶消息檢索相關(guān)上下文
        context = self.kb.search(message)
        # print(context)
        # prompt = self.prompt_template % (context, message)
        prompt = '請基于以下內(nèi)容回答問題：\n' + context
        response = chat(self.model, [Message(role='system', content=prompt), Message(role='user', content=message)])
        return response['message']

from rag import Rag
rag = Rag('deepseek-r1:14b', '私人知識庫.txt')
msg = rag.chat('請介紹下劉芳')
print(msg)

    def stream_chat(self, message):
        context = self.kb.search(message)
        prompt = '請基于以下內(nèi)容回答問題：\n' + context
        response = chat(self.model, [Message(role='system', content=prompt), Message(role='user', content=message)], stream=True)
        # 遍歷流式響應(yīng)
        for chunk in response:
            print(chunk['message']['content'], end='', flush=True)  # 實時打印輸出
        print()  # 輸出完成后換行

from rag import Rag
rag = Rag('deepseek-r1:14b', '私人知識庫.txt')
rag.stream_chat('請介紹下劉芳')