fix

.gitignore

@@ -0,0 +1,6 @@
+db.json
+__pycache__/embedding.cpython-312.pyc
+__pycache__/retriever.cpython-312.pyc
+files/input/SenseNebula AIS产品培训_20220727.pdf
+knowledge_bases/ceshi .faiss
+knowledge_bases/ceshi .pkl

Dockerfile

@@ -0,0 +1,25 @@
+# 使用官方的 Python 基础镜像
+FROM python:3.12-slim
+
+# 设置工作目录
+WORKDIR /app
+
+# 复制依赖文件
+COPY requirements.txt .
+
+ENV PIP_NO_CACHE_DIR=off
+# 设置环境变量
+ENV PYTHONUNBUFFERED=1
+
+# 安装 Python 依赖包
+RUN pip install --upgrade pip
+RUN pip install --no-cache-dir -r requirements.txt
+
+# 复制项目文件
+COPY . .
+
+# 暴露端口（如果需要）
+EXPOSE 7860
+
+# 运行项目
+CMD ["python", "app.py"]

app.py

@@ -0,0 +1,199 @@
+
+import gradio as gr
+import numpy as np
+from gradio_image_prompter import ImagePrompter
+import time
+from pathlib import Path
+from retriever import knowledgeBase
+import llm
+
+current_file_path = Path(__file__).resolve()
+absolute_path = (current_file_path.parent / "files" / "input").resolve()
+
+components = {}
+
+params = {
+    "algo_type": None,
+    "input_image":None
+}
+
+
+def gradio(*keys):
+    if len(keys) == 1 and type(keys[0]) in [list, tuple]:
+        keys = keys[0]
+
+    return [components[k] for k in keys]
+
+def create_ui():
+    with gr.Blocks() as demo:
+        with gr.Tab("知识库"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    with gr.Group():
+                        components["db_view"] = gr.Dataframe(
+                                                    headers=["列表"],
+                                                    datatype=["str"],
+                                                    row_count=2,
+                                                    col_count=(1, "fixed"),
+                                                    interactive=False
+                        )
+                        components["file_expr"] = gr.FileExplorer(
+                            scale=1,
+                            value=[],
+                            file_count="single",
+                            root=absolute_path,
+                            # ignore_glob="**/__init__.py",
+                            elem_id="file_expr",
+                        )
+                with gr.Column(scale=2):
+                    with gr.Row():
+                        with gr.Column(scale=2):
+                            components["db_name"] = gr.Textbox(label="名称", info="请输入库名称", lines=1, value="")
+                        with gr.Column(scale=2):
+                            components["db_submit_btn"] = gr.Button(value="提交")
+                    components["file_upload"] = gr.File(elem_id='file_upload',file_count='multiple',label='文档上传', file_types=[".pdf", ".doc", '.docx', '.json', '.csv'])
+            with gr.Row():
+                with gr.Column(scale=2):
+                    components["db_input"] = gr.Textbox(label="关键词", lines=1, value="")
+                with gr.Column(scale=1):
+                    components["db_test_select"] = gr.Dropdown(knowledgeBase.get_bases(),multiselect=True, label="知识库选择")
+                with gr.Column(scale=1):
+                    components["dbtest_submit_btn"] = gr.Button(value="检索")
+            with gr.Row():
+                with gr.Group():
+                    components["db_search_result"] = gr.JSON(label="检索结果")
+
+        with gr.Tab("问答"):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    with gr.Group():
+                        components["chatbot"] = gr.Chatbot(
+                                            [(None,"你好，有什么需要帮助的？")],
+                                            elem_id="chatbot",
+                                            bubble_full_width=False,
+                                            height=600
+                            )
+                        components["chat_input"] = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message or upload file...", show_label=False)
+                        components["db_select"] = gr.CheckboxGroup(knowledgeBase.get_bases(),label="知识库", info="可选择1个或多个知识库")
+        create_event_handlers()
+        demo.load(init,None,gradio("db_view","db_select","db_test_select"))
+    return demo
+
+def init():
+    db_list = knowledgeBase.get_bases()
+    db_df_list = knowledgeBase.get_df_bases()
+    return db_df_list,gr.CheckboxGroup(db_list,label="知识库", info="可选择1个或多个知识库"),gr.Dropdown(db_list,multiselect=True, label="知识库选择")
+
+def create_event_handlers():
+
+    components["db_submit_btn"].click(
+        file_handler,gradio('file_upload','db_name'),gradio("db_view",'db_select',"db_test_select")
+    )
+
+    components["chat_input"].submit(
+        do_llm_request, gradio("chatbot", "chat_input"), gradio("chatbot", "chat_input")
+    ).then(
+        do_llm_response, gradio("chatbot","db_select"), gradio("chatbot"), api_name="bot_response"
+    ).then(
+        lambda: gr.MultimodalTextbox(interactive=True), None, gradio('chat_input')
+    )
+
+    # components["chatbot"].like(print_like_dislike, None, None)
+
+    components['dbtest_submit_btn'].click(
+        do_search, gradio('db_test_select','db_input'), gradio('db_search_result')
+    )
+
+    components['db_view'].select(
+        db_expr, gradio('db_view'), gradio('file_expr')
+    )
+
+def print_like_dislike(x: gr.LikeData):
+    print(x.index, x.value, x.liked)
+
+def do_llm_request(history, message):
+    for x in message["files"]:
+        history.append(((x,), None))
+    if message["text"] is not None:
+        history.append((message["text"], None))
+    return history, gr.MultimodalTextbox(value=None, interactive=False)
+
+def do_llm_response(history,selected_dbs):
+    print("do_llm_response:",history,selected_dbs)
+    user_input = history[-1][0]
+    prompt = ""
+    quote = ""
+    if len(selected_dbs) > 0:
+        knowledge = knowledgeBase.retrieve_documents(selected_dbs,user_input)
+        print("do_llm_response context:",knowledge)
+        prompt = f'''
+背景1：{knowledge[0]["content"]}
+背景2：{knowledge[1]["content"]}
+背景3：{knowledge[2]["content"]}
+基于以上事实回答问题：{user_input}
+        '''
+
+        quote = f'''
+> 文档：{knowledge[0]["meta"]["source"]}，页码：{knowledge[0]["meta"]["page"]}
+> 文档：{knowledge[1]["meta"]["source"]}，页码：{knowledge[1]["meta"]["page"]}
+> 文档：{knowledge[2]["meta"]["source"]}，页码：{knowledge[2]["meta"]["page"]}
+'''
+    else:
+        prompt = user_input
+    
+    history[-1][1] = ""
+    if llm_client is None:
+        gr.Warning("请先设置大模型")
+        response = "模型参数未设置"
+    else:
+        print("do_llm_response prompt:",prompt)
+        response = llm_client(prompt)
+        response = response.removeprefix(prompt)
+        response += quote
+
+    for character in response:
+        history[-1][1] += character
+        time.sleep(0.01)
+        yield history
+
+
+llm_client = llm.baidu_client
+
+
+def file_handler(file_objs,name):
+    import shutil
+    import os
+    
+    print("file_obj:",file_objs)
+    
+    os.makedirs(os.path.dirname("./files/input/"), exist_ok=True)
+
+    for idx, file in enumerate(file_objs):
+        print(file)
+        file_path = "./files/input/" +  os.path.basename(file.name)
+        if not os.path.exists(file_path):
+            shutil.move(file.name,"./files/input/")
+        
+        knowledgeBase.add_documents_to_kb(name,[file_path])
+
+    dbs = knowledgeBase.get_bases()
+    dfs = knowledgeBase.get_df_bases()
+    return dfs,gr.CheckboxGroup(dbs,label="知识库", info="可选择1个或多个知识库"),gr.Dropdown(dbs,multiselect=True, label="知识库选择")
+
+def db_expr(selected_index: gr.SelectData, dataframe_origin):
+    print("db_expr",selected_index.index)
+    
+    dbname = dataframe_origin.iloc[selected_index.index[0],selected_index.index[1]]
+    print("db_expr",dbname)
+    
+    return knowledgeBase.get_db_files(dbname)
+
+def do_search(selected_dbs,user_input):
+    print("do_search:",selected_dbs,user_input)
+    context = knowledgeBase.retrieve_documents(selected_dbs,user_input)
+    return context
+
+if __name__ == "__main__":
+    demo = create_ui()
+    # demo.launch(server_name="10.151.124.137")
+    demo.launch()

config.py

@@ -0,0 +1,8 @@
+
+wenxin_ak = ""
+wenxin_sk = ""
+
+tongyi_ak = ""
+tongyi_sk = ""
+
+hg_token = ""

embedding.py

@@ -0,0 +1,60 @@
+from transformers import AutoModel, AutoTokenizer
+from typing import Any, List, Mapping, Optional,Union
+from langchain.callbacks.manager import (
+    CallbackManagerForLLMRun
+)
+from langchain_core.embeddings import Embeddings
+import torch
+
+class Embedding(Embeddings):
+
+    def __init__(self,**kwargs):
+        self.model=AutoModel.from_pretrained('BAAI/bge-small-zh-v1.5')
+        self.tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-small-zh-v1.5')
+        self.model.eval()
+       
+    @property
+    def _llm_type(self) -> str:
+        return "BAAI/bge-small-zh-v1.5"
+    
+    @property
+    def model_name(self) -> str:
+        return "embedding"
+    
+    def _call(
+        self,
+        prompt: List[str],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        encoded_input = self.tokenizer(prompt, padding=True, truncation=True, return_tensors='pt')
+
+        with torch.no_grad():
+            model_output = self.model(**encoded_input)
+            # Perform pooling. In this case, cls pooling.
+            sentence_embeddings = model_output[0][:, 0]
+            print(sentence_embeddings.shape)
+        # normalize embeddings
+        sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
+        return sentence_embeddings.numpy()
+
+    @property
+    def _identifying_params(self) -> Mapping[str, Any]:
+        """Get the identifying parameters."""
+        return {"model_path": self.model_path}
+    
+    def embed_documents(self, texts) -> List[List[float]]:
+        # Embed a list of documents
+        embeddings = []
+        print("embed_documents:",len(texts),type(texts))
+        embedding = self._call(texts)
+        for row in embedding:
+            embeddings.append(row)
+        # print("embed_documents: shape",embeddings.shape)
+        return embeddings
+    
+    def embed_query(self, text) -> List[float]:
+        # Embed a single query
+        embedding = self._call([text])
+        return embedding[0]

llm.py

@@ -0,0 +1,73 @@
+import requests
+import json
+from http import HTTPStatus
+from dashscope import Application
+import config
+
+def baidu_client(input):
+    
+    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-lite-8k?access_token=" + get_access_token()
+    
+    payload = json.dumps({
+        "temperature": 0.95,
+        "top_p": 0.7,
+        "penalty_score": 1,
+        "messages": [
+            {
+                "role": "user",
+                "content": input
+            }
+        ],
+        "system": ""
+    })
+    headers = {
+        'Content-Type': 'application/json'
+    }
+    
+    response = requests.request("POST", url, headers=headers, data=payload)
+    
+    print("baidu_client",response.text)
+    return response.json()["result"]
+    
+
+def get_access_token():
+    """
+    使用 AK，SK 生成鉴权签名（Access Token）
+    :return: access_token，或是None(如果错误)
+    """
+    url = "https://aip.baidubce.com/oauth/2.0/token"
+    params = {"grant_type": "client_credentials", "client_id": config.wenxin_ak, "client_secret": config.wenxin_sk}
+    return str(requests.post(url, params=params).json().get("access_token"))
+
+
+def qwen_agent_app(input):
+    response = Application.call(app_id=config.tongyi_ak,
+                                prompt=input,
+                                api_key=config.tongyi_sk,
+                                )
+
+    if response.status_code != HTTPStatus.OK:
+        print('request_id=%s, code=%s, message=%s\n' % (response.request_id, response.status_code, response.message))
+        return ""
+    else:
+        print('request_id=%s\n output=%s\n usage=%s\n' % (response.request_id, response.output, response.usage))
+        return response.output["text"]
+    
+
+def hg_client(input):
+
+    import requests
+    API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3"
+    headers = {"Authorization": f"Bearer {config.hg_token}"}
+
+    def query(payload):
+        response = requests.post(API_URL, headers=headers, json=payload)
+        return response.json()
+        
+    output = query({
+        "inputs": input,
+    })
+    print(output)
+    if len(output) >0:
+        return output[0]['generated_text']
+    return ""

makefile

@@ -0,0 +1,54 @@
+# 定义变量
+IMAGE_NAME=guojingneo/rag-app
+DOCKERFILE_PATH=Dockerfile
+CONTAINER_NAME=rag-app-container
+PORT=7860
+
+# 获取 Git 提交 ID
+COMMIT_ID := $(shell git rev-parse --short HEAD)
+
+# 默认目标
+.PHONY: all
+all: build
+
+# 构建 Docker 镜像
+.PHONY: build
+build:
+	docker build -t $(IMAGE_NAME):$(COMMIT_ID) -f $(DOCKERFILE_PATH) .
+
+# 运行 Docker 容器
+.PHONY: run
+run:
+	docker run -d --name $(CONTAINER_NAME) -p $(PORT):$(PORT) $(IMAGE_NAME):$(COMMIT_ID)
+
+# 停止并删除容器
+.PHONY: stop
+stop:
+	docker stop $(CONTAINER_NAME) || true
+	docker rm $(CONTAINER_NAME) || true
+
+# 推送 Docker 镜像到注册表
+.PHONY: push
+push:
+	docker push $(IMAGE_NAME):$(COMMIT_ID)
+
+# 清理未使用的 Docker 镜像和容器
+.PHONY: clean
+clean:
+	docker system prune -f
+
+# 打包镜像并推送
+.PHONY: package
+package: build push
+
+# 显示帮助信息
+.PHONY: help
+help:
+	@echo "使用方法:"
+	@echo "  make build      构建 Docker 镜像"
+	@echo "  make run        运行 Docker 容器"
+	@echo "  make stop       停止并删除容器"
+	@echo "  make push       推送 Docker 镜像到注册表"
+	@echo "  make clean      清理未使用的 Docker 镜像和容器"
+	@echo "  make package    构建并推送 Docker 镜像"
+	@echo "  make help       显示帮助信息"

requirements.txt

@@ -0,0 +1,6 @@
+faiss-cpu==1.8.0
+pypdf==4.2.0
+langchain==0.2.5
+langchain-community==0.2.5
+transformers==4.32.1
+dashscope==1.20.0

retriever.py

@@ -0,0 +1,170 @@
+from langchain_community.vectorstores import FAISS
+from langchain_community.document_loaders import TextLoader, JSONLoader, PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.docstore.in_memory import InMemoryDocstore
+import faiss
+import os
+import glob
+import json
+from typing import Any,List,Dict
+from embedding import Embedding
+
+
+class KnowledgeBaseManager:
+    def __init__(self, base_path="./knowledge_bases", embedding_dim=512, batch_size=16):
+        self.base_path = base_path
+        self.embedding_dim = embedding_dim
+        self.batch_size = batch_size
+        self.embeddings = Embedding()
+        self.knowledge_bases: Dict[str, FAISS] = {}
+        self.db_files_map: Dict[str, list] = {}
+        os.makedirs(self.base_path, exist_ok=True)
+
+        faiss_files = glob.glob(os.path.join(base_path, '*.faiss'))
+        # 获取不带后缀的名称
+        file_names_without_extension = [os.path.splitext(os.path.basename(file))[0] for file in faiss_files]
+        for name in file_names_without_extension:
+            self.load_knowledge_base(name)
+
+
+    def create_knowledge_base(self, name: str):
+        index = faiss.IndexFlatL2(self.embedding_dim)
+        kb = FAISS(self.embeddings, index, InMemoryDocstore(), {})
+        if name in self.knowledge_bases:
+            print(f"Knowledge base '{name}' already exists.")
+            return
+        
+        self.knowledge_bases[name] = kb
+        self.db_files_map[name] = []
+        self.save_knowledge_base(name)
+        print(f"Knowledge base '{name}' created.")
+
+    def delete_knowledge_base(self, name: str):
+        if name in self.knowledge_bases:
+            del self.knowledge_bases[name]
+            del self.db_files_map[name]
+            os.remove(os.path.join(self.base_path, f"{name}.faiss"))
+            print(f"Knowledge base '{name}' deleted.")
+        else:
+            print(f"Knowledge base '{name}' does not exist.")
+
+    def load_knowledge_base(self, name: str):
+        kb_path = os.path.join(self.base_path, f"{name}.faiss")
+        if os.path.exists(kb_path):
+            self.knowledge_bases[name] = FAISS.load_local(self.base_path, self.embeddings, name, allow_dangerous_deserialization=True)
+            # 加载文件中的数据
+            try:
+                with open('db.json', 'r+') as f:
+                    self.db_files_map = json.load(f)
+            except FileNotFoundError:
+                # 如果文件不存在，则创建一个空的文件并初始化 self.db_files_map
+                with open('db.json', 'w+') as f:
+                    self.db_files_map = {}
+                    json.dump(self.db_files_map, f)
+            print(f"Knowledge base '{name}' loaded.")
+        else:
+            print(f"Knowledge base '{name}' does not exist.")
+
+    def save_knowledge_base(self, name: str):
+        if name in self.knowledge_bases:
+            self.knowledge_bases[name].save_local(self.base_path, name)
+            with open('db.json', 'w') as f:
+                json.dump(self.db_files_map, f)
+            print(f"Knowledge base '{name}' saved.")
+        else:
+            print(f"Knowledge base '{name}' does not exist.")
+
+    # Document(page_content = '渠道版', metadata = {
+	# 'source': './files/input/PS004.pdf',
+	# 'page': 0
+    # }), Document(page_content = '2/20.', metadata = {
+    #     'source': './files/input/PS004.pdf',
+    #     'page': 1
+    # })
+    def add_documents_to_kb(self, name: str, file_paths: List[str]):
+        if name not in self.knowledge_bases:
+            print(f"Knowledge base '{name}' does not exist.")
+            self.create_knowledge_base(name)
+        
+        kb = self.knowledge_bases[name]
+        self.db_files_map[name].extend([os.path.basename(file_path) for file_path in file_paths])
+        documents = self.load_documents(file_paths)
+        print(f"Loaded {len(documents)} documents.")
+        print(documents)
+        pages = self.split_documents(documents)
+        print(f"Split documents into {len(pages)} pages.")
+        # print(pages)
+        
+        doc_ids = []
+        for i in range(0, len(pages), self.batch_size):
+            batch = pages[i:i+self.batch_size]
+            doc_ids.extend(kb.add_documents(batch))
+        
+        self.save_knowledge_base(name)
+        return doc_ids
+
+    def load_documents(self, file_paths: List[str]):
+        documents = []
+        for file_path in file_paths:
+            loader = self.get_loader(file_path)
+            documents.extend(loader.load())
+        return documents
+
+    def get_loader(self, file_path: str):
+        if file_path.endswith('.txt'):
+            return TextLoader(file_path)
+        elif file_path.endswith('.json'):
+            return JSONLoader(file_path)
+        elif file_path.endswith('.pdf'):
+            return PyPDFLoader(file_path)
+        else:
+            raise ValueError("Unsupported file format")
+
+    def split_documents(self, documents):
+        text_splitter = RecursiveCharacterTextSplitter(separators=[
+                                                    "\n\n",
+                                                    "\n",
+                                                    " ",
+                                                    ".",
+                                                    ",",
+                                                    "\u200b",  # Zero-width space
+                                                    "\uff0c",  # Fullwidth comma
+                                                    "\u3001",  # Ideographic comma
+                                                    "\uff0e",  # Fullwidth full stop
+                                                    "\u3002",  # Ideographic full stop
+                                                    "",
+                                                ],
+                                                chunk_size=512, chunk_overlap=0)
+        return text_splitter.split_documents(documents)
+
+    def retrieve_documents(self, names: List[str], query: str):
+        results = []
+        for name in names:
+            if name not in self.knowledge_bases:
+                print(f"Knowledge base '{name}' does not exist.")
+                continue
+            
+            retriever = self.knowledge_bases[name].as_retriever(
+                search_type="mmr",
+                search_kwargs={"score_threshold": 0.5, "k": 3}
+            )
+            docs = retriever.get_relevant_documents(query)
+            results.extend([{"name": name, "content": doc.page_content,"meta": doc.metadata} for doc in docs])
+            
+        
+        return results
+    
+    def get_db_files(self,name):
+        data = self.db_files_map[name]
+        return data
+    
+    def get_bases(self):
+        data = self.knowledge_bases.keys()
+        return list(data)
+    
+    def get_df_bases(self):
+        import pandas as pd
+        data = self.knowledge_bases.keys()
+        return pd.DataFrame(list(data), columns=['列表'])
+
+knowledgeBase = KnowledgeBaseManager()