LSWenda/plugins/zhishiku_rtst.py

from langchain.embeddings import HuggingFaceEmbeddings
import sentence_transformers
import numpy as np
import re,os
from plugins.common import settings,allowCROS
from plugins.common import error_helper 
from plugins.common import success_print 
if settings.librarys.rtst.backend=="Annoy":
    from langchain.vectorstores.annoy import Annoy as Vectorstore
else:
    from langchain.vectorstores.faiss import FAISS as Vectorstore
divider='\n'

if not os.path.exists('memory'):
    os.mkdir('memory')
cunnrent_setting=settings.librarys.rtst
def get_doc_by_id(id,memory_name):
    return vectorstores[memory_name].docstore.search(vectorstores[memory_name].index_to_docstore_id[id])

def process_strings(A, C, B):
    # find the longest common suffix of A and prefix of B
    common = ""
    for i in range(1, min(len(A), len(B)) + 1):
        if A[-i:] == B[:i]:
            common = A[-i:]
    # if there is a common substring, replace one of them with C and concatenate
    if common:
        return A[:-len(common)] + C + B
    # otherwise, just return A + B
    else:
        return A + B
    
def get_title_by_doc(doc):
    return re.sub('【.+】', '', doc.metadata['source'])
def get_doc(id,score,step,memory_name):
    doc = get_doc_by_id(id,memory_name)
    final_content=doc.page_content
    # print("文段分数：",score,[doc.page_content])
    if step > 0:
        for i in range(1, step+1):
            try:
                doc_before=get_doc_by_id(id-i,memory_name)
                if get_title_by_doc(doc_before)==get_title_by_doc(doc):
                    final_content=process_strings(doc_before.page_content,divider,final_content)
                    # print("上文分数：",score,doc.page_content)
            except:
                pass
            try:
                doc_after=get_doc_by_id(id+i,memory_name)
                if get_title_by_doc(doc_after)==get_title_by_doc(doc):
                    final_content=process_strings(final_content,divider,doc_after.page_content)
            except:
                pass
    if doc.metadata['source'].endswith(".pdf") or doc.metadata['source'].endswith(".txt"):
        title=f"[{doc.metadata['source']}](/txt/{doc.metadata['source']})"
    else:
        title=doc.metadata['source']
    return {'title': title,'content':re.sub(r'\n+', "\n", final_content),"score":int(score)}
def find(s,step = 0,memory_name="default"):
    try:
        embedding = get_vectorstore(memory_name).embedding_function(s)
        scores, indices = vectorstores[memory_name].index.search(np.array([embedding], dtype=np.float32), int(cunnrent_setting.count))
        docs = []
        for j, i in enumerate(indices[0]):
            if i == -1:
                continue
            if scores[0][j]>260:continue
            docs.append(get_doc(i,scores[0][j],step,memory_name))

        return docs
    except Exception as e:
        print(e)
        return []
model_path=cunnrent_setting.model_path

try:
    if model_path.startswith("http"):#"http://127.0.0.1:3000/"
        from langchain.embeddings import OpenAIEmbeddings
        import os
        os.environ["OPENAI_API_TYPE"] = "open_ai"
        os.environ["OPENAI_API_BASE"] = model_path
        os.environ["OPENAI_API_KEY"] = "your OpenAI key"

        from langchain.embeddings.openai import OpenAIEmbeddings
        embeddings = OpenAIEmbeddings(
            deployment="text-embedding-ada-002",
            model="text-embedding-ada-002"
        )
    else:
        from langchain.embeddings import HuggingFaceEmbeddings
        embeddings = HuggingFaceEmbeddings(model_name='')
        embeddings.client = sentence_transformers.SentenceTransformer(
            model_path, device="cuda")
except Exception as e:
    error_helper("embedding加载失败",
                 r"https://github.com/l15y/wenda")
    raise e
vectorstores={}
def get_vectorstore(memory_name):
    try:
        return vectorstores[memory_name]
    except Exception  as e:
        try:
            vectorstores[memory_name] = Vectorstore.load_local(
                'memory/'+memory_name, embeddings=embeddings)
            return vectorstores[memory_name]
        except Exception  as e:
            success_print("没有读取到RTST记忆区%s，将新建。"%memory_name)
    return None

from langchain.docstore.document import Document
from langchain.text_splitter import CharacterTextSplitter
from bottle import route, response, request, static_file, hook
import bottle
@route('/upload_rtst_zhishiku', method=("POST","OPTIONS"))
def upload_zhishiku():
    allowCROS()
    try:
        data = request.json
        title=data.get("title")
        memory_name=data.get("memory_name")
        data = data.get("txt")
        # data = re.sub(r'！', "！\n", data)
        # data = re.sub(r'：', "：\n", data)
        # data = re.sub(r'。', "。\n", data)
        data = re.sub(r"\n\s*\n", "\n", data)
        data = re.sub(r'\r', "\n", data)
        data = re.sub(r'\n\n', "\n", data)
        docs=[Document(page_content=data, metadata={"source":title })]
        print(docs)
        
        if hasattr(settings.librarys.rtst,"size") and hasattr(settings.librarys.rtst,"overlap"):
            text_splitter = CharacterTextSplitter(
                chunk_size=int(settings.librarys.rtst.size), chunk_overlap=int(settings.librarys.rtst.overlap), separator='\n')
        else:
            text_splitter = CharacterTextSplitter(
                chunk_size=20, chunk_overlap=0, separator='\n')
        doc_texts = text_splitter.split_documents(docs)

        texts = [d.page_content for d in doc_texts]
        metadatas = [d.metadata for d in doc_texts]
        vectorstore_new = Vectorstore.from_texts(texts, embeddings, metadatas=metadatas)
        vectorstore=get_vectorstore(memory_name)
        if vectorstore is None:
            vectorstores[memory_name]=vectorstore_new
        else:
            vectorstores[memory_name].merge_from(vectorstore_new)
        return '成功'
    except Exception as e:
        return str(e)
@route('/save_rtst_zhishiku', method=("POST","OPTIONS"))
def save_zhishiku():
    allowCROS()
    try:
        data = request.json
        memory_name=data.get("memory_name")
        vectorstores[memory_name].save_local('memory/'+memory_name)
        return "保存成功"
    except Exception as e:
        return str(e)
import json
@route('/find_rtst_in_memory', method=("POST","OPTIONS"))
def api_find():
    allowCROS()
    try:
        data = request.json
        prompt = data.get('prompt')
        step = data.get('step')
        memory_name=data.get("memory_name")
        if step is None:
            step = int(settings.library.step)
        return json.dumps(find(prompt,int(step),memory_name))
    except Exception as e:
        return str(e)
    
@route('/list_rtst_in_disk', method=("POST","OPTIONS"))
def api_find():
    allowCROS()
    return json.dumps(os.listdir('memory'))
    
@route('/del_rtst_in_memory', method=("POST","OPTIONS"))
def api_find():
    allowCROS()
    try:
        data = request.json
        memory_name=data.get("memory_name")
        del vectorstores[memory_name]
    except Exception as e:
        return str(e)

@route('/save_news', method=("POST","OPTIONS"))
def save_news():
    allowCROS()
    try:
        data = request.json
        if not data:
            return 'no data'
        title = data.get('title')
        txt = data.get('txt')
        cut_file = f"txt/{title}.txt"
        with open(cut_file, 'w', encoding='utf-8') as f:
            f.write(txt)
            f.close()
        return 'success'
    except Exception as e:
        return(e)