from .model_config import *
from .server_config import *

VERSION = "v0.0.1"
configs/__init__.py
codefuse-ai/codefuse-chatbot
from .base_chat import Chat
from .knowledge_chat import KnowledgeChat
from .llm_chat import LLMChat
from .search_chat import SearchChat

__all__ = [
    "Chat", "KnowledgeChat", "LLMChat", "SearchChat"
]

dev_opsgpt/chat/__init__.py
codefuse-ai/codefuse-chatbot
from fastapi import Body, Request
from fastapi.responses import StreamingResponse
import asyncio, json
from typing import List, AsyncIterable

from langchain.chat_models import ChatOpenAI
from langchain import LLMChain
from langchain.callbacks import AsyncIteratorCallbackHandler
from langchain.prompts.chat import ChatPromptTemplate


from dev_opsgpt.chat.utils import History, wrap_done
from configs.model_config import (llm_model_dict, LLM_MODEL, VECTOR_SEARCH_TOP_K, SCORE_THRESHOLD)
from dev_opsgpt.utils import BaseResponse
from loguru import logger






def getChatModel(callBack: AsyncIteratorCallbackHandler = None):
    if callBack is None:
        model = ChatOpenAI(
            streaming=True,
            verbose=True,
            openai_api_key=llm_model_dict[LLM_MODEL]["api_key"],
            openai_api_base=llm_model_dict[LLM_MODEL]["api_base_url"],
            model_name=LLM_MODEL
        )
    else:
        model = ChatOpenAI(
            streaming=True,
            verbose=True,
            callBack=[callBack],
            openai_api_key=llm_model_dict[LLM_MODEL]["api_key"],
            openai_api_base=llm_model_dict[LLM_MODEL]["api_base_url"],
            model_name=LLM_MODEL
        )
    return model


class Chat:
    def __init__(
            self,
            engine_name: str = "",
            top_k: int = 1,
            stream: bool = False,
            ) -> None:
        self.engine_name = engine_name
        self.top_k = top_k
        self.stream = stream

    def check_service_status(self, ) -> BaseResponse:
        return BaseResponse(code=200, msg=f"okok")

    def chat(
            self, 
            query: str = Body(..., description="用户输入", examples=["hello"]),
            history: List[History] = Body(
                [], description="历史对话",
                examples=[[{"role": "user", "content": "我们来玩成语接龙，我先来，生龙活虎"}]]
                ),
            engine_name: str = Body(..., description="知识库名称", examples=["samples"]),
            top_k: int = Body(VECTOR_SEARCH_TOP_K, description="匹配向量数"),
            score_threshold: float = Body(SCORE_THRESHOLD, description="知识库匹配相关度阈值，取值范围在0-1之间，SCORE越小，相关度越高，取到1相当于不筛选，建议设置在0.5左右", ge=0, le=1),
            stream: bool = Body(False, description="流式输出"),
            local_doc_url: bool = Body(False, description="知识文件返回本地路径(true)或URL(false)"),
            request: Request = None,
            ):
        self.engine_name = engine_name if isinstance(engine_name, str) else engine_name.default
        self.top_k = top_k if isinstance(top_k, int) else top_k.default
        self.score_threshold = score_threshold if isinstance(score_threshold, float) else score_threshold.default
        self.stream = stream if isinstance(stream, bool) else stream.default
        self.local_doc_url = local_doc_url if isinstance(local_doc_url, bool) else local_doc_url.default
        self.request = request
        return self._chat(query, history)
    
    def _chat(self, query: str, history: List[History]):
        history = [History(**h) if isinstance(h, dict) else h for h in history]
        ## check service dependcy is ok
        service_status = self.check_service_status()
        if service_status.code!=200: return service_status

        def chat_iterator(query: str, history: List[History]):
            model = getChatModel()

            result ,content = self.create_task(query, history, model)
            if self.stream:
                for token in content["text"]:
                    result["answer"] = token
                    yield json.dumps(result, ensure_ascii=False)
            else:
                for token in content["text"]:
                    result["answer"] += token
                yield json.dumps(result, ensure_ascii=False)
        
        return StreamingResponse(chat_iterator(query, history),
                                     media_type="text/event-stream")
        
    def achat(
            self, 
            query: str = Body(..., description="用户输入", examples=["hello"]),
            history: List[History] = Body(
                [], description="历史对话",
                examples=[[{"role": "user", "content": "我们来玩成语接龙，我先来，生龙活虎"}]]
                ),
            engine_name: str = Body(..., description="知识库名称", examples=["samples"]),
            top_k: int = Body(VECTOR_SEARCH_TOP_K, description="匹配向量数"),
            score_threshold: float = Body(SCORE_THRESHOLD, description="知识库匹配相关度阈值，取值范围在0-1之间，SCORE越小，相关度越高，取到1相当于不筛选，建议设置在0.5左右", ge=0, le=1),
            stream: bool = Body(False, description="流式输出"),
            local_doc_url: bool = Body(False, description="知识文件返回本地路径(true)或URL(false)"),
            request: Request = None,
            ):
        self.engine_name = engine_name if isinstance(engine_name, str) else engine_name.default
        self.top_k = top_k if isinstance(top_k, int) else top_k.default
        self.score_threshold = score_threshold if isinstance(score_threshold, float) else score_threshold.default
        self.stream = stream if isinstance(stream, bool) else stream.default
        self.local_doc_url = local_doc_url if isinstance(local_doc_url, bool) else local_doc_url.default
        self.request = request
        return self._achat(query, history)
    
    def _achat(self, query: str, history: List[History]):
        history = [History(**h) if isinstance(h, dict) else h for h in history]
        ## check service dependcy is ok
        service_status = self.check_service_status()
        if service_status.code!=200: return service_status

        async def chat_iterator(query, history):
            callback = AsyncIteratorCallbackHandler()
            model = getChatModel()

            task, result = self.create_atask(query, history, model, callback)
            if self.stream:
                for token in callback["text"]:
                    result["answer"] = token
                    yield json.dumps(result, ensure_ascii=False)
            else:
                for token in callback["text"]:
                    result["answer"] += token
                yield json.dumps(result, ensure_ascii=False)
            await task
        
        return StreamingResponse(chat_iterator(query, history),
                                     media_type="text/event-stream")
        
    def create_task(self, query: str, history: List[History], model):
        '''构建 llm 生成任务'''
        chat_prompt = ChatPromptTemplate.from_messages(
            [i.to_msg_tuple() for i in history] + [("human", "{input}")]
        )
        chain = LLMChain(prompt=chat_prompt, llm=model)
        content = chain({"input": query})
        return {"answer": "", "docs": ""}, content

    def create_atask(self, query, history, model, callback: AsyncIteratorCallbackHandler):
        chat_prompt = ChatPromptTemplate.from_messages(
            [i.to_msg_tuple() for i in history] + [("human", "{input}")]
        )
        chain = LLMChain(prompt=chat_prompt, llm=model)
        task = asyncio.create_task(wrap_done(
            chain.acall({"input": query}), callback.done
        ))
        return task, {"answer": "", "docs": ""}
dev_opsgpt/chat/base_chat.py
codefuse-ai/codefuse-chatbot
from fastapi import Request
import os, asyncio
from urllib.parse import urlencode
from typing import List

from langchain import LLMChain
from langchain.callbacks import AsyncIteratorCallbackHandler
from langchain.prompts.chat import ChatPromptTemplate


from configs.model_config import (
    llm_model_dict, LLM_MODEL, PROMPT_TEMPLATE, 
    VECTOR_SEARCH_TOP_K, SCORE_THRESHOLD)
from dev_opsgpt.chat.utils import History, wrap_done
from dev_opsgpt.utils import BaseResponse
from .base_chat import Chat
from dev_opsgpt.service.kb_api import search_docs, KBServiceFactory
from loguru import logger


class KnowledgeChat(Chat):

    def __init__(
            self,
            engine_name: str = "",
            top_k: int = VECTOR_SEARCH_TOP_K,
            stream: bool = False,
            score_thresold: float = SCORE_THRESHOLD,
            local_doc_url: bool = False,
            request: Request = None,
            ) -> None:
        super().__init__(engine_name, top_k, stream)
        self.score_thresold = score_thresold
        self.local_doc_url = local_doc_url
        self.request = request

    def check_service_status(self) -> BaseResponse:
        kb = KBServiceFactory.get_service_by_name(self.engine_name)
        if kb is None:
            return BaseResponse(code=404, msg=f"未找到知识库 {self.engine_name}")
        return BaseResponse(code=200, msg=f"找到知识库 {self.engine_name}")
    
    def _process(self, query: str, history: List[History], model):
        '''process'''
        docs = search_docs(query, self.engine_name, self.top_k, self.score_threshold)
        context = "\n".join([doc.page_content for doc in docs])
        source_documents = []
        for inum, doc in enumerate(docs):
            filename = os.path.split(doc.metadata["source"])[-1]
            if self.local_doc_url:
                url = "file://" + doc.metadata["source"]
            else:
                parameters = urlencode({"knowledge_base_name": self.engine_name, "file_name":filename})
                url = f"{self.request.base_url}knowledge_base/download_doc?" + parameters
            text = f"""出处 [{inum + 1}] [{filename}]({url}) \n\n{doc.page_content}\n\n"""
            source_documents.append(text)
        chat_prompt = ChatPromptTemplate.from_messages(
            [i.to_msg_tuple() for i in history] + [("human", PROMPT_TEMPLATE)]
        )
        chain = LLMChain(prompt=chat_prompt, llm=model)
        result = {"answer": "", "docs": source_documents}
        return chain, context, result

    def create_task(self, query: str, history: List[History], model):
        '''构建 llm 生成任务'''
        logger.debug(f"query: {query}, history: {history}")
        chain, context, result = self._process(query, history, model)
        try:
            content = chain({"context": context, "question": query})
        except Exception as e:
            content = {"text": str(e)}
        return result, content

    def create_atask(self, query, history, model, callback: AsyncIteratorCallbackHandler):
        chain, context, result = self._process(query, history, model)
        task = asyncio.create_task(wrap_done(
            chain.acall({"context": context, "question": query}), callback.done
        ))
        return task, result

dev_opsgpt/chat/knowledge_chat.py
codefuse-ai/codefuse-chatbot
import asyncio
from typing import List

from langchain import LLMChain
from langchain.callbacks import AsyncIteratorCallbackHandler
from langchain.prompts.chat import ChatPromptTemplate


from dev_opsgpt.chat.utils import History, wrap_done
from .base_chat import Chat
from loguru import logger


class LLMChat(Chat):

    def __init__(
            self,
            engine_name: str = "",
            top_k: int = 1,
            stream: bool = False,
            ) -> None:
        super().__init__(engine_name, top_k, stream)

    def create_task(self, query: str, history: List[History], model):
        '''构建 llm 生成任务'''
        chat_prompt = ChatPromptTemplate.from_messages(
            [i.to_msg_tuple() for i in history] + [("human", "{input}")]
        )
        chain = LLMChain(prompt=chat_prompt, llm=model)
        content = chain({"input": query})
        return {"answer": "", "docs": ""}, content

    def create_atask(self, query, history, model, callback: AsyncIteratorCallbackHandler):
        chat_prompt = ChatPromptTemplate.from_messages(
            [i.to_msg_tuple() for i in history] + [("human", "{input}")]
        )
        chain = LLMChain(prompt=chat_prompt, llm=model)
        task = asyncio.create_task(wrap_done(
            chain.acall({"input": query}), callback.done
        ))
        return task, {"answer": "", "docs": ""}
dev_opsgpt/chat/llm_chat.py
codefuse-ai/codefuse-chatbot
from fastapi import Request
import os, asyncio
from urllib.parse import urlencode
from typing import List, Optional, Dict

from langchain import LLMChain
from langchain.callbacks import AsyncIteratorCallbackHandler
from langchain.utilities import BingSearchAPIWrapper, DuckDuckGoSearchAPIWrapper
from langchain.prompts.chat import ChatPromptTemplate
from langchain.docstore.document import Document

from configs.model_config import (
    PROMPT_TEMPLATE, SEARCH_ENGINE_TOP_K, BING_SUBSCRIPTION_KEY, BING_SEARCH_URL,
    VECTOR_SEARCH_TOP_K, SCORE_THRESHOLD)
from dev_opsgpt.chat.utils import History, wrap_done
from dev_opsgpt.utils import BaseResponse
from .base_chat import Chat

from loguru import logger

from duckduckgo_search import DDGS


def bing_search(text, result_len=SEARCH_ENGINE_TOP_K):
    if not (BING_SEARCH_URL and BING_SUBSCRIPTION_KEY):
        return [{"snippet": "please set BING_SUBSCRIPTION_KEY and BING_SEARCH_URL in os ENV",
                 "title": "env info is not found",
                 "link": "https://python.langchain.com/en/latest/modules/agents/tools/examples/bing_search.html"}]
    search = BingSearchAPIWrapper(bing_subscription_key=BING_SUBSCRIPTION_KEY,
                                  bing_search_url=BING_SEARCH_URL)
    return search.results(text, result_len)


def duckduckgo_search(
        query: str,
        result_len: int = SEARCH_ENGINE_TOP_K,
        region: Optional[str] = "wt-wt", 
        safesearch: str = "moderate", 
        time: Optional[str] = "y",
        backend: str = "api",
        ):
    with DDGS(proxies=os.environ.get("DUCKDUCKGO_PROXY")) as ddgs:
        results = ddgs.text(
            query,
            region=region,
            safesearch=safesearch,
            timelimit=time,
            backend=backend,
        )
        if results is None:
            return [{"Result": "No good DuckDuckGo Search Result was found"}]

        def to_metadata(result: Dict) -> Dict[str, str]:
            if backend == "news":
                return {
                    "date": result["date"],
                    "title": result["title"],
                    "snippet": result["body"],
                    "source": result["source"],
                    "link": result["url"],
                }
            return {
                "snippet": result["body"],
                "title": result["title"],
                "link": result["href"],
            }

        formatted_results = []
        for i, res in enumerate(results, 1):
            if res is not None:
                formatted_results.append(to_metadata(res))
            if len(formatted_results) == result_len:
                break
    return formatted_results


# def duckduckgo_search(text, result_len=SEARCH_ENGINE_TOP_K):
#     search = DuckDuckGoSearchAPIWrapper()
#     return search.results(text, result_len)


SEARCH_ENGINES = {"duckduckgo": duckduckgo_search,
                  "bing": bing_search,
                  }


def search_result2docs(search_results):
    docs = []
    for result in search_results:
        doc = Document(page_content=result["snippet"] if "snippet" in result.keys() else "",
                       metadata={"source": result["link"] if "link" in result.keys() else "",
                                 "filename": result["title"] if "title" in result.keys() else ""})
        docs.append(doc)
    return docs


def lookup_search_engine(
        query: str,
        search_engine_name: str,
        top_k: int = SEARCH_ENGINE_TOP_K,
):
    results = SEARCH_ENGINES[search_engine_name](query, result_len=top_k)
    docs = search_result2docs(results)
    return docs



class SearchChat(Chat):

    def __init__(
            self,
            engine_name: str = "",
            top_k: int = VECTOR_SEARCH_TOP_K,
            stream: bool = False,
            ) -> None:
        super().__init__(engine_name, top_k, stream)

    def check_service_status(self) -> BaseResponse:
        if self.engine_name not in SEARCH_ENGINES.keys():
            return BaseResponse(code=404, msg=f"未支持搜索引擎 {self.engine_name}")
        return BaseResponse(code=200, msg=f"支持搜索引擎 {self.engine_name}")
    
    def _process(self, query: str, history: List[History], model):
        '''process'''
        docs = lookup_search_engine(query, self.engine_name, self.top_k)
        context = "\n".join([doc.page_content for doc in docs])

        source_documents = [
            f"""出处 [{inum + 1}] [{doc.metadata["source"]}]({doc.metadata["source"]}) \n\n{doc.page_content}\n\n"""
            for inum, doc in enumerate(docs)
        ]
        
        chat_prompt = ChatPromptTemplate.from_messages(
            [i.to_msg_tuple() for i in history] + [("human", PROMPT_TEMPLATE)]
        )
        chain = LLMChain(prompt=chat_prompt, llm=model)
        result = {"answer": "", "docs": source_documents}
        return chain, context, result

    def create_task(self, query: str, history: List[History], model):
        '''构建 llm 生成任务'''
        chain, context, result = self._process(query, history, model)
        content = chain({"context": context, "question": query})
        return result, content

    def create_atask(self, query, history, model, callback: AsyncIteratorCallbackHandler):
        chain, context, result = self._process(query, history, model)
        task = asyncio.create_task(wrap_done(
            chain.acall({"context": context, "question": query}), callback.done
        ))
        return task, result

dev_opsgpt/chat/search_chat.py
codefuse-ai/codefuse-chatbot
import asyncio
from typing import Awaitable
from pydantic import BaseModel, Field


async def wrap_done(fn: Awaitable, event: asyncio.Event):
    """Wrap an awaitable with a event to signal when it's done or an exception is raised."""
    try:
        await fn
    except Exception as e:
        # TODO: handle exception
        print(f"Caught exception: {e}")
    finally:
        # Signal the aiter to stop.
        event.set()


class History(BaseModel):
    """
    对话历史
    可从dict生成，如
    h = History(**{"role":"user","content":"你好"})
    也可转换为tuple，如
    h.to_msy_tuple = ("human", "你好")
    """
    role: str = Field(...)
    content: str = Field(...)

    def to_msg_tuple(self):
        return "ai" if self.role=="assistant" else "human", self.content

dev_opsgpt/chat/utils.py
codefuse-ai/codefuse-chatbot
from .json_loader import JSONLoader
from .jsonl_loader import JSONLLoader

__all__ = [
    "JSONLoader", "JSONLLoader"
]
dev_opsgpt/document_loaders/__init__.py
codefuse-ai/codefuse-chatbot
import json
from pathlib import Path
from typing import AnyStr, Callable, Dict, List, Optional, Union

from langchain.docstore.document import Document
from langchain.document_loaders.base import BaseLoader

from dev_opsgpt.utils.common_utils import read_json_file


class JSONLoader(BaseLoader):

    def __init__(
            self,
            file_path: Union[str, Path],
            schema_key: str = "all_text",
            content_key: Optional[str] = None,
            metadata_func: Optional[Callable[[Dict, Dict], Dict]] = None,
            text_content: bool = True,
    ):
        self.file_path = Path(file_path).resolve()
        self.schema_key = schema_key
        self._content_key = content_key
        self._metadata_func = metadata_func
        self._text_content = text_content

    def load(self, ) -> List[Document]:
        """Load and return documents from the JSON file."""
        docs: List[Document] = []
        datas = read_json_file(self.file_path)
        self._parse(datas, docs)
        return docs
    
    def _parse(self, datas: List, docs: List[Document]) -> None:
        for idx, sample in enumerate(datas):
            metadata = dict(
                source=str(self.file_path),
                seq_num=idx,
            )
            text = sample.get(self.schema_key, "")
            docs.append(Document(page_content=text, metadata=metadata))

dev_opsgpt/document_loaders/json_loader.py
codefuse-ai/codefuse-chatbot
import json
from pathlib import Path
from typing import AnyStr, Callable, Dict, List, Optional, Union

from langchain.docstore.document import Document
from langchain.document_loaders.base import BaseLoader

from dev_opsgpt.utils.common_utils import read_jsonl_file


class JSONLLoader(BaseLoader):

    def __init__(
            self,
            file_path: Union[str, Path],
            schema_key: str = "all_text",
            content_key: Optional[str] = None,
            metadata_func: Optional[Callable[[Dict, Dict], Dict]] = None,
            text_content: bool = True,
    ):
        self.file_path = Path(file_path).resolve()
        self.schema_key = schema_key
        self._content_key = content_key
        self._metadata_func = metadata_func
        self._text_content = text_content

    def load(self, ) -> List[Document]:
        """Load and return documents from the JSON file."""
        docs: List[Document] = []
        datas = read_jsonl_file(self.file_path)
        self._parse(datas, docs)
        return docs
    
    def _parse(self, datas: List, docs: List[Document]) -> None:
        for idx, sample in enumerate(datas):
            metadata = dict(
                source=str(self.file_path),
                seq_num=idx,
            )
            text = sample.get(self.schema_key, "")
            docs.append(Document(page_content=text, metadata=metadata))

dev_opsgpt/document_loaders/jsonl_loader.py
codefuse-ai/codefuse-chatbot
from typing import List

from langchain.embeddings.base import Embeddings
from langchain.schema import Document



class BaseVSCService:
    def do_create_kb(self):
        pass

    def do_drop_kb(self):
        pass

    def do_add_doc(self, docs: List[Document], embeddings: Embeddings):
        pass

    def do_clear_vs(self):
        pass

    def vs_type(self) -> str:
        return "default"

    def do_init(self):
        pass

    def do_search(self):
        pass

    def do_insert_multi_knowledge(self):
        pass

    def do_insert_one_knowledge(self):
        pass

    def do_delete_doc(self):
        pass

dev_opsgpt/embeddings/commands/default_vs_cds.py
codefuse-ai/codefuse-chatbot
from functools import lru_cache
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from configs.model_config import embedding_model_dict
from loguru import logger


@lru_cache(1)
def load_embeddings(model: str, device: str):
    logger.info("load embedding model: {}, {}".format(model, embedding_model_dict[model]))
    embeddings = HuggingFaceEmbeddings(model_name=embedding_model_dict[model],
                                       model_kwargs={'device': device})
    return embeddings
dev_opsgpt/embeddings/utils.py
codefuse-ai/codefuse-chatbot
from .db import _engine, Base
from loguru import logger

__all__ = [

]

def create_tables():
    Base.metadata.create_all(bind=_engine)

def reset_tables():
    Base.metadata.drop_all(bind=_engine)
    create_tables()


def check_tables_exist(table_name) -> bool:
    table_exist = _engine.dialect.has_table(_engine.connect(), table_name, schema=None)
    return table_exist

def table_init():
    if (not check_tables_exist("knowledge_base")) or (not check_tables_exist ("knowledge_file")):
        create_tables()

dev_opsgpt/orm/__init__.py
codefuse-ai/codefuse-chatbot
from .document_file_cds import *
from .document_base_cds import *

__all__ = [
    "add_kb_to_db", "list_kbs_from_db", "kb_exists", 
    "load_kb_from_db", "delete_kb_from_db", "get_kb_detail",
    
    "list_docs_from_db", "add_doc_to_db", "delete_file_from_db", 
    "delete_files_from_db", "doc_exists", "get_file_detail",
]
dev_opsgpt/orm/commands/__init__.py
codefuse-ai/codefuse-chatbot
from dev_opsgpt.orm.db import with_session, _engine
from dev_opsgpt.orm.schemas.base_schema import KnowledgeBaseSchema


# @with_session
# def _query_by_condition(session, schema, query_kargs, query_type="first"):
#     if len(query_kargs) >0:
#         if query_type == "first":
#             return session.query(schema).filter_by(query_kargs).first()
#         elif query_type == "all":
#             return session.query(schema).filter_by(query_kargs).first()

# @with_session    
# def _add_to_db(session, schema, query_kargs):
#     kb = schema(**query_kargs)
#     session.add(kb)
#     return True

# @with_session
# def add_to_db(session, schema, query_kargs):
#     kb = _query_by_condition(session, schema, query_kargs, query_type="first")
#     if not kb:
#         _add_to_db(session, schema, query_kargs)
#     else: # update kb with new vs_type and embed_model
#         for k, v in query_kargs.items():
#             if k in kb:
#                 kb[k] = v
#     return True



@with_session
def add_kb_to_db(session, kb_name, vs_type, embed_model):
    # 创建知识库实例
    kb = session.query(KnowledgeBaseSchema).filter_by(kb_name=kb_name).first()
    if not kb:
        kb = KnowledgeBaseSchema(kb_name=kb_name, vs_type=vs_type, embed_model=embed_model)
        session.add(kb)
    else: # update kb with new vs_type and embed_model
        kb.vs_type = vs_type
        kb.embed_model = embed_model
    return True


@with_session
def list_kbs_from_db(session, min_file_count: int = -1):
    kbs = session.query(KnowledgeBaseSchema.kb_name).filter(KnowledgeBaseSchema.file_count > min_file_count).all()
    kbs = [kb[0] for kb in kbs]
    return kbs


@with_session
def kb_exists(session, kb_name):
    kb = session.query(KnowledgeBaseSchema).filter_by(kb_name=kb_name).first()
    status = True if kb else False
    return status


@with_session
def load_kb_from_db(session, kb_name):
    kb = session.query(KnowledgeBaseSchema).filter_by(kb_name=kb_name).first()
    if kb:
        kb_name, vs_type, embed_model = kb.kb_name, kb.vs_type, kb.embed_model
    else:
        kb_name, vs_type, embed_model = None, None, None
    return kb_name, vs_type, embed_model


@with_session
def delete_kb_from_db(session, kb_name):
    kb = session.query(KnowledgeBaseSchema).filter_by(kb_name=kb_name).first()
    if kb:
        session.delete(kb)
    return True


@with_session
def get_kb_detail(session, kb_name: str) -> dict:
    kb: KnowledgeBaseSchema = session.query(KnowledgeBaseSchema).filter_by(kb_name=kb_name).first()
    if kb:
        return {
            "kb_name": kb.kb_name,
            "vs_type": kb.vs_type,
            "embed_model": kb.embed_model,
            "file_count": kb.file_count,
            "create_time": kb.create_time,
        }
    else:
        return {}

dev_opsgpt/orm/commands/document_base_cds.py
codefuse-ai/codefuse-chatbot
from dev_opsgpt.orm.db import with_session, _engine
from dev_opsgpt.orm.schemas.base_schema import KnowledgeFileSchema, KnowledgeBaseSchema
from dev_opsgpt.orm.utils import DocumentFile



@with_session
def list_docs_from_db(session, kb_name):
    files = session.query(KnowledgeFileSchema).filter_by(kb_name=kb_name).all()
    docs = [f.file_name for f in files]
    return docs


@with_session
def add_doc_to_db(session, kb_file: DocumentFile):
    kb = session.query(KnowledgeBaseSchema).filter_by(kb_name=kb_file.kb_name).first()
    if kb:
        # 如果已经存在该文件，则更新文件版本号
        existing_file = session.query(KnowledgeFileSchema).filter_by(file_name=kb_file.filename,
                                                                    kb_name=kb_file.kb_name).first()
        if existing_file:
            existing_file.file_version += 1
        # 否则，添加新文件
        else:
            new_file = KnowledgeFileSchema(
                file_name=kb_file.filename,
                file_ext=kb_file.ext,
                kb_name=kb_file.kb_name,
                document_loader_name=kb_file.document_loader_name,
                text_splitter_name=kb_file.text_splitter_name or "SpacyTextSplitter",
            )
            kb.file_count += 1
            session.add(new_file)
    return True


@with_session
def delete_file_from_db(session, kb_file: DocumentFile):
    existing_file = session.query(KnowledgeFileSchema).filter_by(file_name=kb_file.filename,
                                                                kb_name=kb_file.kb_name).first()
    if existing_file:
        session.delete(existing_file)
        session.commit()

        kb = session.query(KnowledgeBaseSchema).filter_by(kb_name=kb_file.kb_name).first()
        if kb:
            kb.file_count -= 1
            session.commit()
    return True


@with_session
def delete_files_from_db(session, knowledge_base_name: str):
    session.query(KnowledgeFileSchema).filter_by(kb_name=knowledge_base_name).delete()

    kb = session.query(KnowledgeBaseSchema).filter_by(kb_name=knowledge_base_name).first()
    if kb:
        kb.file_count = 0

    session.commit()
    return True


@with_session
def doc_exists(session, kb_file: DocumentFile):
    existing_file = session.query(KnowledgeFileSchema).filter_by(file_name=kb_file.filename,
                                                                kb_name=kb_file.kb_name).first()
    return True if existing_file else False


@with_session
def get_file_detail(session, kb_name: str, filename: str) -> dict:
    file: KnowledgeFileSchema = (session.query(KnowledgeFileSchema)
                                .filter_by(file_name=filename,
                                            kb_name=kb_name).first())
    if file:
        return {
            "kb_name": file.kb_name,
            "file_name": file.file_name,
            "file_ext": file.file_ext,
            "file_version": file.file_version,
            "document_loader": file.document_loader_name,
            "text_splitter": file.text_splitter_name,
            "create_time": file.create_time,
        }
    else:
        return {}

dev_opsgpt/orm/commands/document_file_cds.py
codefuse-ai/codefuse-chatbot
from contextlib import contextmanager
from sqlalchemy.engine import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

from configs.model_config import SQLALCHEMY_DATABASE_URI


_engine = create_engine(SQLALCHEMY_DATABASE_URI)

session_factory = sessionmaker(bind=_engine)

Base = declarative_base()



def init_session():
    session = session_factory()

    try:
        yield session
    finally:
        try:
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()



def with_session(func):
    def wrapper(*args, **kwargs):
        session = session_factory()
        try:
            return func(session, *args, **kwargs)
        finally:
            try:
                session.commit()
            except Exception as e:
                session.rollback()
                raise e
            finally:
                session.close()
    return wrapper


@contextmanager
def session_scope():
    """上下文管理器用于自动获取 Session, 避免错误"""
    session = session_factory(autoflush=True)
    try:
        session.commit()
    except Exception as e:
        session.rollback()
        raise e
    finally:
        session.close()

dev_opsgpt/orm/db.py
codefuse-ai/codefuse-chatbot
from sqlalchemy import Column, Integer, String, DateTime, func

from dev_opsgpt.orm.db import Base


class KnowledgeBaseSchema(Base):
    """
    知识库模型
    """
    __tablename__ = 'knowledge_base'
    id = Column(Integer, primary_key=True, autoincrement=True, comment='知识库ID')
    kb_name = Column(String, comment='知识库名称')
    vs_type = Column(String, comment='嵌入模型类型')
    embed_model = Column(String, comment='嵌入模型名称')
    file_count = Column(Integer, default=0, comment='文件数量')
    create_time = Column(DateTime, default=func.now(), comment='创建时间')

    def __repr__(self):
        return f"""<KnowledgeBase(id='{self.id}', 
        kb_name='{self.kb_name}', 
        vs_type='{self.vs_type}', 
        embed_model='{self.embed_model}', 
        file_count='{self.file_count}', 
        create_time='{self.create_time}')>"""

class KnowledgeFileSchema(Base):
    """
    知识文件模型
    """
    __tablename__ = 'knowledge_file'
    id = Column(Integer, primary_key=True, autoincrement=True, comment='知识文件ID')
    file_name = Column(String, comment='文件名')
    file_ext = Column(String, comment='文件扩展名')
    kb_name = Column(String, comment='所属知识库名称')
    document_loader_name = Column(String, comment='文档加载器名称')
    text_splitter_name = Column(String, comment='文本分割器名称')
    file_version = Column(Integer, default=1, comment='文件版本')
    create_time = Column(DateTime, default=func.now(), comment='创建时间')

    def __repr__(self):
        return f"""<KnowledgeFile(id='{self.id}', 
        file_name='{self.file_name}', 
        file_ext='{self.file_ext}', 
        kb_name='{self.kb_name}', 
        document_loader_name='{self.document_loader_name}', 
        text_splitter_name='{self.text_splitter_name}', 
        file_version='{self.file_version}', 
        create_time='{self.create_time}')>"""

dev_opsgpt/orm/schemas/base_schema.py
codefuse-ai/codefuse-chatbot
import os
from dev_opsgpt.utils.path_utils import get_file_path, get_LoaderClass, SUPPORTED_EXTS


class DocumentFile:
    def __init__(
            self, filename: str, knowledge_base_name: str) -> None:
        self.kb_name = knowledge_base_name
        self.filename = filename
        self.ext = os.path.splitext(filename)[-1].lower()
        if self.ext not in SUPPORTED_EXTS:
            raise ValueError(f"暂未支持的文件格式 {self.ext}")
        self.filepath = get_file_path(knowledge_base_name, filename)
        self.docs = None
        self.document_loader_name = get_LoaderClass(self.ext)

        # TODO: 增加依据文件格式匹配text_splitter
        self.text_splitter_name = None
dev_opsgpt/orm/utils.py
codefuse-ai/codefuse-chatbot
from .basebox import CodeBoxResponse
from .pycodebox import PyCodeBox

__all__ = [
    "CodeBoxResponse", "PyCodeBox"
]
dev_opsgpt/sandbox/__init__.py
codefuse-ai/codefuse-chatbot
from pydantic import BaseModel
from typing import Optional
from pathlib import Path
import sys
from abc import ABC, abstractclassmethod

from configs.server_config import SANDBOX_SERVER


class CodeBoxResponse(BaseModel):
    code_text: str = ""
    code_exe_response: str = ""
    code_exe_type: str = ""
    code_exe_status: int
    do_code_exe: bool

    def __str__(self,):
        return f"""status: {self.code_exe_status}, type: {self.code_exe_type}, response: {self.code_exe_response}"""


class CodeBoxStatus(BaseModel):
    status: str


class CodeBoxFile(BaseModel):
    """
    Represents a file returned from a CodeBox instance.
    """

    name: str
    content: Optional[bytes] = None

    def __str__(self):
        return self.name

    def __repr__(self):
        return f"File({self.name})"
    

class BaseBox(ABC):

    enter_status = False

    def __init__(
            self, 
            remote_url: str = "",
            remote_ip: str = SANDBOX_SERVER["host"],
            remote_port: str = SANDBOX_SERVER["port"],
            token: str = "mytoken",
            do_code_exe: bool = False,
            do_remote: bool = False
    ):
        self.token = token
        self.remote_url = remote_url or remote_ip + ":" + str(remote_port)
        self.remote_ip = remote_ip
        self.remote_port = remote_port
        self.do_code_exe = do_code_exe
        self.do_remote = do_remote
        self.local_pyenv = Path(sys.executable).absolute()
        self.ws = None
        self.aiohttp_session = None
        self.kernel_url = ""

    def chat(self, text: str, file_path: str = None, do_code_exe: bool = None) -> CodeBoxResponse:
        '''执行流'''
        do_code_exe = self.do_code_exe if do_code_exe is None else do_code_exe
        if not do_code_exe:
            return CodeBoxResponse(
                code_exe_response=text, code_text=text, code_exe_type="text", code_exe_status=200,
                do_code_exe=do_code_exe
            )
        
        try:
            code_text = self.decode_code_from_text(text)
            return self.run(code_text, file_path)
        except Exception as e:
            return CodeBoxResponse(
                code_exe_response=str(e), code_text=text, code_exe_type="error", code_exe_status=500,
                do_code_exe=do_code_exe
            )
        
    async def achat(self, text: str, file_path: str = None, do_code_exe: bool = None) -> CodeBoxResponse:
        do_code_exe = self.do_code_exe if do_code_exe is None else do_code_exe
        if not do_code_exe:
            return CodeBoxResponse(
                code_exe_response=text, code_text=text, code_exe_type="text", code_exe_status=200,
                do_code_exe=do_code_exe
            )
        
        try:
            code_text = self.decode_code_from_text(text)
            return await self.arun(code_text, file_path)
        except Exception as e:
            return CodeBoxResponse(
                code_exe_response=str(e), code_text=text, code_exe_type="error", code_exe_status=500,
                do_code_exe=do_code_exe
            )
        
    def run(
            self,
            code_text: str = None,
            file_path: str = None,
            retry=3,
    ) -> CodeBoxResponse:
        '''执行代码'''
        pass

    async def arun(
            self,
            code_text: str = None,
            file_path: str = None,
            retry=3, 
    ) -> CodeBoxResponse:
        '''执行代码'''
        pass

    def decode_code_from_text(self, text):
        pass

    def start(self,):
        pass

    async def astart(self, ):
        pass

    @abstractclassmethod
    def stop(self) -> CodeBoxStatus:
        """Terminate the CodeBox instance"""

    def __enter__(self, ) -> "BaseBox":
        if not self.enter_status:
            self.start()
        return self
    
    async def __aenter__(self, ) -> "BaseBox":
        if not self.enter_status:
            await self.astart()
        return self
    
    def __exit__(self, exc_type, exc_value, traceback) -> None:
        self.stop()


dev_opsgpt/sandbox/basebox.py
codefuse-ai/codefuse-chatbot
import time, os, docker, requests, json, uuid, subprocess, time, asyncio, aiohttp, re, traceback
import psutil
from typing import List, Optional, Union
from pathlib import Path
from loguru import logger

from websockets.sync.client import connect as ws_connect_sync
from websockets.client import connect as ws_connect
from websocket import create_connection
from websockets.client import WebSocketClientProtocol, ClientConnection
from websockets.exceptions import ConnectionClosedError

from configs.server_config import SANDBOX_SERVER
from .basebox import BaseBox, CodeBoxResponse, CodeBoxStatus, CodeBoxFile


class PyCodeBox(BaseBox):

    enter_status: bool = False

    def __init__(
            self, 
            remote_url: str = "",
            remote_ip: str = SANDBOX_SERVER["host"],
            remote_port: str = SANDBOX_SERVER["port"],
            token: str = "mytoken",
            do_code_exe: bool = False,
            do_remote: bool = False
    ):
        super().__init__(remote_url, remote_ip, remote_port, token, do_code_exe, do_remote)
        self.enter_status = True
        asyncio.run(self.astart())

    def decode_code_from_text(self, text: str) -> str:
        pattern = r'```.*?```'
        code_blocks = re.findall(pattern, text, re.DOTALL)
        code_text: str = "\n".join([block.strip('`') for block in code_blocks])
        code_text = code_text[6:] if code_text.startswith("python") else code_text
        code_text = code_text.replace("python\n", "").replace("code", "")
        return code_text
    
    def run(
            self, code_text: Optional[str] = None,
            file_path: Optional[os.PathLike] = None,
            retry = 3,
    ) -> CodeBoxResponse:
        if not code_text and not file_path:
            return CodeBoxResponse(
                code_exe_response="Code or file_path must be specifieds!",
                code_text=code_text,
                code_exe_type="text",
                code_exe_status=502,
                do_code_exe=self.do_code_exe,
            )
        
        if code_text and file_path:
            return CodeBoxResponse(
                code_exe_response="Can only specify code or the file to read_from!",
                code_text=code_text,
                code_exe_type="text",
                code_exe_status=502,
                do_code_exe=self.do_code_exe,
            )
        
        if file_path:
            with open(file_path, "r", encoding="utf-8") as f:
                code_text = f.read()

        # run code in jupyter kernel
        if retry <= 0:
            raise RuntimeError("Could not connect to kernel")
        
        if not self.ws:
            raise RuntimeError("Jupyter not running. Make sure to start it first")
        
        logger.debug(f"code_text: {json.dumps(code_text, ensure_ascii=False)}")
        self.ws.send(
            json.dumps(
                {
                    "header": {
                        "msg_id": (msg_id := uuid.uuid4().hex),
                        "msg_type": "execute_request",
                    },
                    "parent_header": {},
                    "metadata": {},
                    "content": {
                        "code": code_text,
                        "silent": True,
                        "store_history": True,
                        "user_expressions": {},
                        "allow_stdin": False,
                        "stop_on_error": True,
                    },
                    "channel": "shell",
                    "buffers": [],
                }
            )
        )
        result = ""
        while True:
            try:
                if isinstance(self.ws, WebSocketClientProtocol):
                    raise RuntimeError("Mixing asyncio and sync code is not supported")
                received_msg = json.loads(self.ws.recv())
            except ConnectionClosedError:
                logger.debug("box start, ConnectionClosedError!!!")
                self.start()
                return self.run(code_text, file_path, retry - 1)

            if (
                received_msg["header"]["msg_type"] == "stream"
                and received_msg["parent_header"]["msg_id"] == msg_id
            ):
                msg = received_msg["content"]["text"].strip()
                if "Requirement already satisfied:" in msg:
                    continue
                result += msg + "\n"

            elif (
                received_msg["header"]["msg_type"] == "execute_result"
                and received_msg["parent_header"]["msg_id"] == msg_id
            ):
                result += received_msg["content"]["data"]["text/plain"].strip() + "\n"

            elif received_msg["header"]["msg_type"] == "display_data":
                if "image/png" in received_msg["content"]["data"]:
                    return CodeBoxResponse(
                        code_exe_type="image/png",
                        code_text=code_text,
                        code_exe_response=received_msg["content"]["data"]["image/png"],
                        code_exe_status=200,
                        do_code_exe=self.do_code_exe
                    )
                if "text/plain" in received_msg["content"]["data"]:
                    return CodeBoxResponse(
                        code_exe_type="text",
                        code_text=code_text,
                        code_exe_response=received_msg["content"]["data"]["text/plain"],
                        code_exe_status=200,
                        do_code_exe=self.do_code_exe
                    )
                return CodeBoxResponse(
                        code_exe_type="error",
                        code_text=code_text,
                        code_exe_response=received_msg["content"]["data"]["text/plain"],
                        code_exe_status=420,
                        do_code_exe=self.do_code_exe
                    )
            elif (
                received_msg["header"]["msg_type"] == "status"
                and received_msg["parent_header"]["msg_id"] == msg_id
                and received_msg["content"]["execution_state"] == "idle"
            ):
                if len(result) > 500:
                    result = "[...]\n" + result[-500:]
                return CodeBoxResponse(
                        code_exe_type="text",
                        code_text=code_text,
                        code_exe_response=result or "Code run successfully (no output)",
                        code_exe_status=200,
                        do_code_exe=self.do_code_exe
                    )

            elif (
                received_msg["header"]["msg_type"] == "error"
                and received_msg["parent_header"]["msg_id"] == msg_id
            ):
                error = (
                    f"{received_msg['content']['ename']}: "
                    f"{received_msg['content']['evalue']}"
                )
                return CodeBoxResponse(
                        code_exe_type="error",
                        code_text=code_text,
                        code_exe_response=error,
                        code_exe_status=500,
                        do_code_exe=self.do_code_exe
                    )

    def _get_kernelid(self, ) -> None:
        headers = {"Authorization": f'Token {self.token}', 'token': self.token}
        response = requests.get(f"{self.kernel_url}?token={self.token}", headers=headers)
        if len(response.json()) > 0:
            self.kernel_id = response.json()[0]["id"]
        else:
            response = requests.post(f"{self.kernel_url}?token={self.token}", headers=headers)
            self.kernel_id = response.json()["id"]
        if self.kernel_id is None:
            raise Exception("Could not start kernel")
        
    async def _aget_kernelid(self, ) -> None:
        headers = {"Authorization": f'Token {self.token}', 'token': self.token}
        response = requests.get(f"{self.kernel_url}?token={self.token}", headers=headers)
        async with aiohttp.ClientSession() as session:
            async with session.get(f"{self.kernel_url}?token={self.token}", headers=headers) as resp:
                if len(await resp.json()) > 0:
                    self.kernel_id = (await resp.json())[0]["id"]
                else:
                    async with session.post(f"{self.kernel_url}?token={self.token}", headers=headers) as response:
                        self.kernel_id = (await response.json())["id"]

        # if len(response.json()) > 0:
        #     self.kernel_id = response.json()[0]["id"]
        # else:
        #     response = requests.post(f"{self.kernel_url}?token={self.token}", headers=headers)
        #     self.kernel_id = response.json()["id"]
        # if self.kernel_id is None:
        #     raise Exception("Could not start kernel")
    def _check_connect(self, ) -> bool:
        if self.kernel_url == "":
            return False

        try:
             response = requests.get(f"{self.kernel_url}?token={self.token}", timeout=270)
             return response.status_code == 200
        except requests.exceptions.ConnectionError:
            return False

    async def _acheck_connect(self, ) -> bool:
        if self.kernel_url == "":
            return False

        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(f"{self.kernel_url}?token={self.token}", timeout=270) as resp:
                    return resp.status == 200
        except aiohttp.ClientConnectorError:
            pass
        except aiohttp.ServerDisconnectedError:
            pass

    def  _check_port(self, ) -> bool:
        try:
             response = requests.get(f"http://localhost:{self.remote_port}", timeout=270)
             logger.warning(f"Port is conflict, please check your codebox's port {self.remote_port}")
             return response.status_code == 200
        except requests.exceptions.ConnectionError:
            return False
        
    async def _acheck_port(self, ) -> bool:
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(f"http://localhost:{self.remote_port}", timeout=270) as resp:
                    logger.warning(f"Port is conflict, please check your codebox's port {self.remote_port}")
                    return resp.status == 200
        except aiohttp.ClientConnectorError:
            pass
        except aiohttp.ServerDisconnectedError:
            pass

    def _check_connect_success(self, retry_nums: int = 5) -> bool:
        while retry_nums > 0:
            try:
                connect_status = self._check_connect()
                if connect_status:
                    logger.info(f"{self.remote_url} connection success")
                    return True
            except requests.exceptions.ConnectionError:
                logger.info(f"{self.remote_url} connection fail")
            retry_nums -= 1
            time.sleep(5)
        raise BaseException(f"can't connect to {self.remote_url}")
    
    async def _acheck_connect_success(self, retry_nums: int = 5) -> bool:
        while retry_nums > 0:
            try:
                connect_status = await self._acheck_connect()
                if connect_status:
                    logger.info(f"{self.remote_url} connection success")
                    return True
            except requests.exceptions.ConnectionError:
                logger.info(f"{self.remote_url} connection fail")
            retry_nums -= 1
            time.sleep(5)
        raise BaseException(f"can't connect to {self.remote_url}")
    
    def start(self, ):
        '''判断是从外部service执行还是内部启动notebook执行'''
        self.jupyter = None
        if self.do_remote:
            # TODO自动检测日期,并重启容器
            self.kernel_url = self.remote_url + "/api/kernels"
            self._check_connect_success()

            self._get_kernelid()
            logger.debug(self.kernel_url.replace("http", "ws") + f"/{self.kernel_id}/channels?token={self.token}")
            self.wc_url = self.kernel_url.replace("http", "ws") + f"/{self.kernel_id}/channels?token={self.token}"
            headers = {"Authorization": f'Token {self.token}', 'token': self.token}
            self.ws = create_connection(self.wc_url, headers=headers)
        else:
            # TODO 自动检测本地接口
            port_status = self._check_port()
            connect_status = self._check_connect()
            logger.debug(f"port_status: {port_status}, connect_status: {connect_status}")
            if port_status and not connect_status:
                raise BaseException(f"Port is conflict, please check your codebox's port {self.remote_port}")
            
            if not connect_status:
                self.jupyter = subprocess.Popen(
                    [
                        "jupyer", "notebnook",
                        f"--NotebookApp.token={self.token}",
                        f"--port={self.remote_port}",
                        "--no-browser",
                        "--ServerApp.disable_check_xsrf=True",
                    ],
                    stderr=subprocess.PIPE,
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                )
            self.kernel_url = self.remote_url + "/api/kernels"
            self._check_connect_success()
            self._get_kernelid()
            logger.debug(self.kernel_url.replace("http", "ws") + f"/{self.kernel_id}/channels?token={self.token}")
            self.wc_url = self.kernel_url.replace("http", "ws") + f"/{self.kernel_id}/channels?token={self.token}"
            headers = {"Authorization": f'Token {self.token}', 'token': self.token}
            self.ws = create_connection(self.wc_url, headers=headers)

    async def astart(self, ):
        '''判断是从外部service执行还是内部启动notebook执行'''
        self.jupyter = None
        if self.do_remote:
            # TODO自动检测日期,并重启容器
            self.kernel_url = self.remote_url + "/api/kernels"

            await self._acheck_connect_success()
            await self._aget_kernelid()
            self.wc_url = self.kernel_url.replace("http", "ws") + f"/{self.kernel_id}/channels?token={self.token}"
            headers = {"Authorization": f'Token {self.token}', 'token': self.token}
            self.ws = create_connection(self.wc_url, headers=headers)
        else:
            # TODO 自动检测本地接口
            port_status = await self._acheck_port()
            self.kernel_url = self.remote_url + "/api/kernels"
            connect_status = await self._acheck_connect()
            logger.debug(f"port_status: {port_status}, connect_status: {connect_status}")
            if port_status and not connect_status:
                raise BaseException(f"Port is conflict, please check your codebox's port {self.remote_port}")
            
            if not connect_status:
                self.jupyter = subprocess.Popen(
                    [
                        "jupyter", "notebook",
                        f"--NotebookApp.token={self.token}",
                        f"--port={self.remote_port}",
                        "--no-browser",
                        "--ServerApp.disable_check_xsrf=True",
                    ],
                    stderr=subprocess.PIPE,
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                )
            self.kernel_url = self.remote_url + "/api/kernels"
            await self._acheck_connect_success()
            await self._aget_kernelid()
            self.wc_url = self.kernel_url.replace("http", "ws") + f"/{self.kernel_id}/channels?token={self.token}"
            headers = {"Authorization": f'Token {self.token}', 'token': self.token}
            self.ws = create_connection(self.wc_url, headers=headers)

    def status(self,) -> CodeBoxStatus:
        if not self.kernel_id:
            self._get_kernelid()
        
        return CodeBoxStatus(
            status="running" if self.kernel_id
            and requests.get(self.kernel_url, timeout=270).status_code == 200
            else "stopped"
        )
    
    async def astatus(self,) -> CodeBoxStatus:
        if not self.kernel_id:
            await self._aget_kernelid()
        
        return CodeBoxStatus(
            status="running" if self.kernel_id
            and requests.get(self.kernel_url, timeout=270).status_code == 200
            else "stopped"
        )
    
    def restart(self, ) -> CodeBoxStatus:
        return CodeBoxStatus(status="restared")
    
    def stop(self, ) -> CodeBoxStatus:
        try:
            if self.jupyter is not None:
                for process in psutil.process_iter(["pid", "name", "cmdline"]):
                    # 检查进程名是否包含"jupyter"
                    if f'port={self.remote_port}' in str(process.info["cmdline"]).lower() and \
                        "jupyter" in process.info['name'].lower():
                        logger.warning(f'port={self.remote_port}, {process.info}')
                        # 关闭进程
                        process.terminate()

                self.jupyter = None
        except Exception as e:
            logger.error(traceback.format_exc())

        if self.ws is not None:
            try:
                if self.ws is not None:
                    self.ws.close()
                else:
                    loop = asyncio.new_event_loop()
                    loop.run_until_complete(self.ws.close())
            except Exception as e:
                logger.error(traceback.format_exc())
            self.ws = None
        return CodeBoxStatus(status="stopped")
    
    def __del__(self):
        self.stop()

dev_opsgpt/sandbox/pycodebox.py
codefuse-ai/codefuse-chatbot
import nltk
import argparse
import uvicorn, os, sys
from fastapi.middleware.cors import CORSMiddleware
from starlette.responses import RedirectResponse
from typing import List

src_dir = os.path.join(
    os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
)
sys.path.append(src_dir)

sys.path.append(os.path.dirname(os.path.dirname(__file__)))

from configs import VERSION
from configs.model_config import NLTK_DATA_PATH
from configs.server_config import OPEN_CROSS_DOMAIN

from dev_opsgpt.chat import LLMChat, SearchChat, KnowledgeChat
from dev_opsgpt.service.kb_api import *
from dev_opsgpt.utils.server_utils import BaseResponse, ListResponse, FastAPI, MakeFastAPIOffline

nltk.data.path = [NLTK_DATA_PATH] + nltk.data.path

from dev_opsgpt.chat import LLMChat, SearchChat, KnowledgeChat

llmChat = LLMChat()
searchChat = SearchChat()
knowledgeChat = KnowledgeChat()


async def document():
    return RedirectResponse(url="/docs")


def create_app():
    app = FastAPI(
        title="DevOps-ChatBot API Server",
        version=VERSION
    )
    MakeFastAPIOffline(app)
    # Add CORS middleware to allow all origins
    # 在config.py中设置OPEN_DOMAIN=True，允许跨域
    # set OPEN_DOMAIN=True in config.py to allow cross-domain
    if OPEN_CROSS_DOMAIN:
        app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )

    app.get("/",
            response_model=BaseResponse,
            summary="swagger 文档")(document)

    # Tag: Chat
    # app.post("/chat/fastchat",
    #          tags=["Chat"],
    #          summary="与llm模型对话(直接与fastchat api对话)")(openai_chat)

    app.post("/chat/chat",
             tags=["Chat"],
             summary="与llm模型对话(通过LLMChain)")(llmChat.chat)

    app.post("/chat/knowledge_base_chat",
             tags=["Chat"],
             summary="与知识库对话")(knowledgeChat.chat)

    app.post("/chat/search_engine_chat",
             tags=["Chat"],
             summary="与搜索引擎对话")(searchChat.chat)

    # Tag: Knowledge Base Management
    app.get("/knowledge_base/list_knowledge_bases",
            tags=["Knowledge Base Management"],
            response_model=ListResponse,
            summary="获取知识库列表")(list_kbs)

    app.post("/knowledge_base/create_knowledge_base",
             tags=["Knowledge Base Management"],
             response_model=BaseResponse,
             summary="创建知识库"
             )(create_kb)

    app.post("/knowledge_base/delete_knowledge_base",
             tags=["Knowledge Base Management"],
             response_model=BaseResponse,
             summary="删除知识库"
             )(delete_kb)

    app.get("/knowledge_base/list_files",
            tags=["Knowledge Base Management"],
            response_model=ListResponse,
            summary="获取知识库内的文件列表"
            )(list_docs)

    app.post("/knowledge_base/search_docs",
             tags=["Knowledge Base Management"],
             response_model=List[DocumentWithScore],
             summary="搜索知识库"
             )(search_docs)

    app.post("/knowledge_base/upload_docs",
             tags=["Knowledge Base Management"],
             response_model=BaseResponse,
             summary="上传文件到知识库，并/或进行向量化"
             )(upload_doc)

    app.post("/knowledge_base/delete_docs",
             tags=["Knowledge Base Management"],
             response_model=BaseResponse,
             summary="删除知识库内指定文件"
             )(delete_doc)

    app.post("/knowledge_base/update_docs",
             tags=["Knowledge Base Management"],
             response_model=BaseResponse,
             summary="更新现有文件到知识库"
             )(update_doc)

    app.get("/knowledge_base/download_doc",
            tags=["Knowledge Base Management"],
            summary="下载对应的知识文件")(download_doc)

    app.post("/knowledge_base/recreate_vector_store",
             tags=["Knowledge Base Management"],
             summary="根据content中文档重建向量库，流式输出处理进度。"
             )(recreate_vector_store)

    # # LLM模型相关接口
    # app.post("/llm_model/list_models",
    #         tags=["LLM Model Management"],
    #         summary="列出当前已加载的模型",
    #         )(list_llm_models)

    # app.post("/llm_model/stop",
    #         tags=["LLM Model Management"],
    #         summary="停止指定的LLM模型（Model Worker)",
    #         )(stop_llm_model)

    # app.post("/llm_model/change",
    #         tags=["LLM Model Management"],
    #         summary="切换指定的LLM模型（Model Worker)",
    #         )(change_llm_model)

    return app


app = create_app()


def run_api(host, port, **kwargs):
    if kwargs.get("ssl_keyfile") and kwargs.get("ssl_certfile"):
        uvicorn.run(app,
                    host=host,
                    port=port,
                    ssl_keyfile=kwargs.get("ssl_keyfile"),
                    ssl_certfile=kwargs.get("ssl_certfile"),
                    )
    else:
        uvicorn.run(app, host=host, port=port)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(prog='DevOps-ChatBot',
                                     description='About DevOps-ChatBot, local knowledge based LLM with langchain'
                                                 ' ｜ 基于本地知识库的 LLM 问答')
    parser.add_argument("--host", type=str, default="0.0.0.0")
    parser.add_argument("--port", type=int, default=7861)
    parser.add_argument("--ssl_keyfile", type=str)
    parser.add_argument("--ssl_certfile", type=str)
    # 初始化消息
    args = parser.parse_args()
    args_dict = vars(args)
    run_api(host=args.host,
            port=args.port,
            ssl_keyfile=args.ssl_keyfile,
            ssl_certfile=args.ssl_certfile,
            )

dev_opsgpt/service/api.py
codefuse-ai/codefuse-chatbot
from abc import ABC, abstractmethod
from typing import List
import os

from langchain.embeddings.base import Embeddings
from langchain.docstore.document import Document

from configs.model_config import (
    kbs_config, VECTOR_SEARCH_TOP_K, SCORE_THRESHOLD,
    EMBEDDING_MODEL, EMBEDDING_DEVICE
)
from dev_opsgpt.orm.commands import *
from dev_opsgpt.utils.path_utils import *
from dev_opsgpt.orm.utils import DocumentFile
from dev_opsgpt.embeddings.utils import load_embeddings
from dev_opsgpt.text_splitter import LCTextSplitter



class SupportedVSType:
    FAISS = 'faiss'
    # MILVUS = 'milvus'
    # DEFAULT = 'default'
    # PG = 'pg'


class KBService(ABC):

    def __init__(self,
                 knowledge_base_name: str,
                 embed_model: str = EMBEDDING_MODEL,
                 ):
        self.kb_name = knowledge_base_name
        self.embed_model = embed_model
        self.kb_path = get_kb_path(self.kb_name)
        self.doc_path = get_doc_path(self.kb_name)
        self.do_init()

    def _load_embeddings(self, embed_device: str = EMBEDDING_DEVICE) -> Embeddings:
        return load_embeddings(self.embed_model, embed_device)

    def create_kb(self):
        """
        创建知识库
        """
        if not os.path.exists(self.doc_path):
            os.makedirs(self.doc_path)
        self.do_create_kb()
        status = add_kb_to_db(self.kb_name, self.vs_type(), self.embed_model)
        return status

    def clear_vs(self):
        """
        删除向量库中所有内容
        """
        self.do_clear_vs()
        status = delete_files_from_db(self.kb_name)
        return status

    def drop_kb(self):
        """
        删除知识库
        """
        self.do_drop_kb()
        status = delete_kb_from_db(self.kb_name)
        return status

    def add_doc(self, kb_file: DocumentFile, **kwargs):
        """
        向知识库添加文件
        """
        lctTextSplitter = LCTextSplitter(kb_file.filepath)
        docs = lctTextSplitter.file2text()
        if docs:
            self.delete_doc(kb_file)
            embeddings = self._load_embeddings()
            self.do_add_doc(docs, embeddings, **kwargs)
            status = add_doc_to_db(kb_file)
        else:
            status = False
        return status

    def delete_doc(self, kb_file: DocumentFile, delete_content: bool = False, **kwargs):
        """
        从知识库删除文件
        """
        self.do_delete_doc(kb_file, **kwargs)
        status = delete_file_from_db(kb_file)
        if delete_content and os.path.exists(kb_file.filepath):
            os.remove(kb_file.filepath)
        return status

    def update_doc(self, kb_file: DocumentFile, **kwargs):
        """
        使用content中的文件更新向量库
        """
        if os.path.exists(kb_file.filepath):
            self.delete_doc(kb_file, **kwargs)
            return self.add_doc(kb_file, **kwargs)
        
    def exist_doc(self, file_name: str):
        return doc_exists(DocumentFile(knowledge_base_name=self.kb_name,
                                        filename=file_name))

    def list_docs(self):
        return list_docs_from_db(self.kb_name)

    def search_docs(self,
                    query: str,
                    top_k: int = VECTOR_SEARCH_TOP_K,
                    score_threshold: float = SCORE_THRESHOLD,
                    ):
        embeddings = self._load_embeddings()
        docs = self.do_search(query, top_k, score_threshold, embeddings)
        return docs

    @abstractmethod
    def do_create_kb(self):
        """
        创建知识库子类实自己逻辑
        """
        pass

    @staticmethod
    def list_kbs_type():
        return list(kbs_config.keys())

    @classmethod
    def list_kbs(cls):
        return list_kbs_from_db()

    def exists(self, kb_name: str = None):
        kb_name = kb_name or self.kb_name
        return kb_exists(kb_name)

    @abstractmethod
    def vs_type(self) -> str:
        pass

    @abstractmethod
    def do_init(self):
        pass

    @abstractmethod
    def do_drop_kb(self):
        """
        删除知识库子类实自己逻辑
        """
        pass

    @abstractmethod
    def do_search(self,
                  query: str,
                  top_k: int,
                  embeddings: Embeddings,
                  ) -> List[Document]:
        """
        搜索知识库子类实自己逻辑
        """
        pass

    @abstractmethod
    def do_add_doc(self,
                   docs: List[Document],
                   embeddings: Embeddings,
                   ):
        """
        向知识库添加文档子类实自己逻辑
        """
        pass

    @abstractmethod
    def do_delete_doc(self,
                      kb_file: DocumentFile):
        """
        从知识库删除文档子类实自己逻辑
        """
        pass

    @abstractmethod
    def do_clear_vs(self):
        """
        从知识库删除全部向量子类实自己逻辑
        """
        pass

dev_opsgpt/service/base_service.py
codefuse-ai/codefuse-chatbot
import os
import shutil
from typing import List
from functools import lru_cache
from loguru import logger

from langchain.vectorstores import FAISS
from langchain.embeddings.base import Embeddings
from langchain.docstore.document import Document
from langchain.embeddings.huggingface import HuggingFaceEmbeddings

from configs.model_config import (
    KB_ROOT_PATH,
    CACHED_VS_NUM,
    EMBEDDING_MODEL,
    EMBEDDING_DEVICE,
    SCORE_THRESHOLD,
    FAISS_NORMALIZE_L2
)
from .base_service import KBService, SupportedVSType
from dev_opsgpt.utils.path_utils import *
from dev_opsgpt.orm.utils import DocumentFile
from dev_opsgpt.utils.server_utils import torch_gc
from dev_opsgpt.embeddings.utils import load_embeddings


# make HuggingFaceEmbeddings hashable
def _embeddings_hash(self):
    return hash(self.model_name)


HuggingFaceEmbeddings.__hash__ = _embeddings_hash

_VECTOR_STORE_TICKS = {}


@lru_cache(CACHED_VS_NUM)
def load_vector_store(
        knowledge_base_name: str,
        embed_model: str = EMBEDDING_MODEL,
        embed_device: str = EMBEDDING_DEVICE,
        embeddings: Embeddings = None,
        tick: int = 0,  # tick will be changed by upload_doc etc. and make cache refreshed.
):
    print(f"loading vector store in '{knowledge_base_name}'.")
    vs_path = get_vs_path(knowledge_base_name)
    if embeddings is None:
        logger.info("load embedmodel: {}".format(embed_model))
        embeddings = load_embeddings(embed_model, embed_device)

    if not os.path.exists(vs_path):
        os.makedirs(vs_path)

    if "index.faiss" in os.listdir(vs_path):
        search_index = FAISS.load_local(vs_path, embeddings, normalize_L2=FAISS_NORMALIZE_L2)
    else:
        # create an empty vector store
        doc = Document(page_content="init", metadata={})
        search_index = FAISS.from_documents([doc], embeddings, normalize_L2=FAISS_NORMALIZE_L2)
        ids = [k for k, v in search_index.docstore._dict.items()]
        search_index.delete(ids)
        search_index.save_local(vs_path)
    
    if tick == 0: # vector store is loaded first time
        _VECTOR_STORE_TICKS[knowledge_base_name] = 0

    # search_index.embedding_function = embeddings.embed_documents
    return search_index


def refresh_vs_cache(kb_name: str):
    """
    make vector store cache refreshed when next loading
    """
    _VECTOR_STORE_TICKS[kb_name] = _VECTOR_STORE_TICKS.get(kb_name, 0) + 1
    print(f"知识库 {kb_name} 缓存刷新：{_VECTOR_STORE_TICKS[kb_name]}")


class FaissKBService(KBService):
    vs_path: str
    kb_path: str

    def vs_type(self) -> str:
        return SupportedVSType.FAISS

    @staticmethod
    def get_vs_path(knowledge_base_name: str):
        return os.path.join(FaissKBService.get_kb_path(knowledge_base_name), "vector_store")

    @staticmethod
    def get_kb_path(knowledge_base_name: str):
        return os.path.join(KB_ROOT_PATH, knowledge_base_name)

    def do_init(self):
        self.kb_path = FaissKBService.get_kb_path(self.kb_name)
        self.vs_path = FaissKBService.get_vs_path(self.kb_name)

    def do_create_kb(self):
        if not os.path.exists(self.vs_path):
            os.makedirs(self.vs_path)
        load_vector_store(self.kb_name, self.embed_model)

    def do_drop_kb(self):
        self.clear_vs()
        shutil.rmtree(self.kb_path)

    def do_search(self,
                  query: str,
                  top_k: int,
                  score_threshold: float = SCORE_THRESHOLD,
                  embeddings: Embeddings = None,
                  ) -> List[Document]:
        search_index = load_vector_store(self.kb_name,
                                         embeddings=embeddings,
                                         tick=_VECTOR_STORE_TICKS.get(self.kb_name))
        docs = search_index.similarity_search_with_score(query, k=top_k, score_threshold=score_threshold)
        return docs

    def do_add_doc(self,
                   docs: List[Document],
                   embeddings: Embeddings,
                   **kwargs,
                   ):
        vector_store = load_vector_store(self.kb_name,
                                         embeddings=embeddings,
                                         tick=_VECTOR_STORE_TICKS.get(self.kb_name, 0))
        logger.info("docs.lens: {}".format(len(docs)))
        vector_store.add_documents(docs)
        torch_gc()
        if not kwargs.get("not_refresh_vs_cache"):
            vector_store.save_local(self.vs_path)
            refresh_vs_cache(self.kb_name)

    def do_delete_doc(self,
                      kb_file: DocumentFile,
                      **kwargs):
        embeddings = self._load_embeddings()
        vector_store = load_vector_store(self.kb_name,
                                         embeddings=embeddings,
                                         tick=_VECTOR_STORE_TICKS.get(self.kb_name, 0))

        ids = [k for k, v in vector_store.docstore._dict.items() if v.metadata["source"] == kb_file.filepath]
        if len(ids) == 0:
            return None

        vector_store.delete(ids)
        if not kwargs.get("not_refresh_vs_cache"):
            vector_store.save_local(self.vs_path)
            refresh_vs_cache(self.kb_name)

        return True

    def do_clear_vs(self):
        shutil.rmtree(self.vs_path)
        os.makedirs(self.vs_path)
        refresh_vs_cache(self.kb_name)

    def exist_doc(self, file_name: str):
        if super().exist_doc(file_name):
            return "in_db"

        content_path = os.path.join(self.kb_path, "content")
        if os.path.isfile(os.path.join(content_path, file_name)):
            return "in_folder"
        else:
            return False

dev_opsgpt/service/faiss_db_service.py
codefuse-ai/codefuse-chatbot
import urllib, os, json, traceback
from typing import List, Dict
from loguru import logger

from fastapi.responses import StreamingResponse, FileResponse
from fastapi import Body, File, Form, Body, Query, UploadFile
from langchain.docstore.document import Document

from .service_factory import KBServiceFactory
from dev_opsgpt.utils.server_utils import BaseResponse, ListResponse
from dev_opsgpt.utils.path_utils import *
from dev_opsgpt.orm.commands import *
from dev_opsgpt.orm.utils import DocumentFile
from configs.model_config import (
    DEFAULT_VS_TYPE, EMBEDDING_MODEL, VECTOR_SEARCH_TOP_K, SCORE_THRESHOLD
)



async def list_kbs():
    # Get List of Knowledge Base
    return ListResponse(data=list_kbs_from_db())


async def create_kb(knowledge_base_name: str = Body(..., examples=["samples"]),
                    vector_store_type: str = Body("faiss"),
                    embed_model: str = Body(EMBEDDING_MODEL),
                    ) -> BaseResponse:
    # Create selected knowledge base
    if not validate_kb_name(knowledge_base_name):
        return BaseResponse(code=403, msg="Don't attack me")
    if knowledge_base_name is None or knowledge_base_name.strip() == "":
        return BaseResponse(code=404, msg="知识库名称不能为空，请重新填写知识库名称")

    kb = KBServiceFactory.get_service_by_name(knowledge_base_name)
    if kb is not None:
        return BaseResponse(code=404, msg=f"已存在同名知识库 {knowledge_base_name}")

    kb = KBServiceFactory.get_service(knowledge_base_name, vector_store_type, embed_model)
    try:
        kb.create_kb()
    except Exception as e:
        print(e)
        return BaseResponse(code=500, msg=f"创建知识库出错： {e}")

    return BaseResponse(code=200, msg=f"已新增知识库 {knowledge_base_name}")


async def delete_kb(
        knowledge_base_name: str = Body(..., examples=["samples"])
    ) -> BaseResponse:
    # Delete selected knowledge base
    if not validate_kb_name(knowledge_base_name):
        return BaseResponse(code=403, msg="Don't attack me")
    knowledge_base_name = urllib.parse.unquote(knowledge_base_name)

    kb = KBServiceFactory.get_service_by_name(knowledge_base_name)

    if kb is None:
        return BaseResponse(code=404, msg=f"未找到知识库 {knowledge_base_name}")

    try:
        status = kb.clear_vs()
        status = kb.drop_kb()
        if status:
            return BaseResponse(code=200, msg=f"成功删除知识库 {knowledge_base_name}")
    except Exception as e:
        print(e)
        return BaseResponse(code=500, msg=f"删除知识库时出现意外： {e}")

    return BaseResponse(code=500, msg=f"删除知识库失败 {knowledge_base_name}")



class DocumentWithScore(Document):
    score: float = None


def search_docs(query: str = Body(..., description="用户输入", examples=["你好"]),
                knowledge_base_name: str = Body(..., description="知识库名称", examples=["samples"]),
                top_k: int = Body(VECTOR_SEARCH_TOP_K, description="匹配向量数"),
                score_threshold: float = Body(SCORE_THRESHOLD, description="知识库匹配相关度阈值，取值范围在0-1之间，SCORE越小，相关度越高，取到1相当于不筛选，建议设置在0.5左右", ge=0, le=1),
                ) -> List[DocumentWithScore]:
    kb = KBServiceFactory.get_service_by_name(knowledge_base_name)
    if kb is None:
        return []
    docs = kb.search_docs(query, top_k, score_threshold)
    data = [DocumentWithScore(**x[0].dict(), score=x[1]) for x in docs]

    return data


async def list_docs(
    knowledge_base_name: str
) -> ListResponse:
    if not validate_kb_name(knowledge_base_name):
        return ListResponse(code=403, msg="Don't attack me", data=[])

    knowledge_base_name = urllib.parse.unquote(knowledge_base_name)
    kb = KBServiceFactory.get_service_by_name(knowledge_base_name)
    if kb is None:
        return ListResponse(code=404, msg=f"未找到知识库 {knowledge_base_name}", data=[])
    else:
        all_doc_names = kb.list_docs()
        return ListResponse(data=all_doc_names)


async def upload_doc(file: UploadFile = File(..., description="上传文件"),
                     knowledge_base_name: str = Form(..., description="知识库名称", examples=["kb1"]),
                     override: bool = Form(False, description="覆盖已有文件"),
                     not_refresh_vs_cache: bool = Form(False, description="暂不保存向量库（用于FAISS）"),
                     ) -> BaseResponse:
    if not validate_kb_name(knowledge_base_name):
        return BaseResponse(code=403, msg="Don't attack me")

    kb = KBServiceFactory.get_service_by_name(knowledge_base_name)
    if kb is None:
        return BaseResponse(code=404, msg=f"未找到知识库 {knowledge_base_name}")

    file_content = await file.read()  # 读取上传文件的内容

    try:
        kb_file = DocumentFile(filename=file.filename,
                                knowledge_base_name=knowledge_base_name)

        if (os.path.exists(kb_file.filepath)
                and not override
                and os.path.getsize(kb_file.filepath) == len(file_content)
        ):
            # TODO: filesize 不同后的处理
            file_status = f"文件 {kb_file.filename} 已存在。"
            return BaseResponse(code=404, msg=file_status)

        with open(kb_file.filepath, "wb") as f:
            f.write(file_content)
    except Exception as e:
        logger.error(traceback.format_exc())
        return BaseResponse(code=500, msg=f"{kb_file.filename} 文件上传失败，报错信息为: {e}")

    try:
        kb.add_doc(kb_file, not_refresh_vs_cache=not_refresh_vs_cache)
    except Exception as e:
        logger.error(traceback.format_exc())
        return BaseResponse(code=500, msg=f"{kb_file.filename} 文件向量化失败，报错信息为: {e}")

    return BaseResponse(code=200, msg=f"成功上传文件 {kb_file.filename}")


async def delete_doc(knowledge_base_name: str = Body(..., examples=["samples"]),
                     doc_name: str = Body(..., examples=["file_name.md"]),
                     delete_content: bool = Body(False),
                     not_refresh_vs_cache: bool = Body(False, description="暂不保存向量库（用于FAISS）"),
                    ) -> BaseResponse:
    if not validate_kb_name(knowledge_base_name):
        return BaseResponse(code=403, msg="Don't attack me")

    knowledge_base_name = urllib.parse.unquote(knowledge_base_name)
    kb = KBServiceFactory.get_service_by_name(knowledge_base_name)
    if kb is None:
        return BaseResponse(code=404, msg=f"未找到知识库 {knowledge_base_name}")

    if not kb.exist_doc(doc_name):
        return BaseResponse(code=404, msg=f"未找到文件 {doc_name}")

    try:
        kb_file = DocumentFile(filename=doc_name,
                                knowledge_base_name=knowledge_base_name)
        kb.delete_doc(kb_file, delete_content, not_refresh_vs_cache=not_refresh_vs_cache)
    except Exception as e:
        print(e)
        return BaseResponse(code=500, msg=f"{kb_file.filename} 文件删除失败，错误信息：{e}")

    return BaseResponse(code=200, msg=f"{kb_file.filename} 文件删除成功")


async def update_doc(
        knowledge_base_name: str = Body(..., examples=["samples"]),
        file_name: str = Body(..., examples=["file_name"]),
        not_refresh_vs_cache: bool = Body(False, description="暂不保存向量库（用于FAISS）"),
    ) -> BaseResponse:
    '''
    更新知识库文档
    '''
    if not validate_kb_name(knowledge_base_name):
        return BaseResponse(code=403, msg="Don't attack me")

    kb = KBServiceFactory.get_service_by_name(knowledge_base_name)
    if kb is None:
        return BaseResponse(code=404, msg=f"未找到知识库 {knowledge_base_name}")

    try:
        kb_file = DocumentFile(filename=file_name,
                                knowledge_base_name=knowledge_base_name)
        if os.path.exists(kb_file.filepath):
            kb.update_doc(kb_file, not_refresh_vs_cache=not_refresh_vs_cache)
            return BaseResponse(code=200, msg=f"成功更新文件 {kb_file.filename}")
    except Exception as e:
        logger.error(traceback.format_exc())
        return BaseResponse(code=500, msg=f"{kb_file.filename} 文件更新失败，错误信息是：{e}")

    return BaseResponse(code=500, msg=f"{kb_file.filename} 文件更新失败")


async def download_doc(
        knowledge_base_name: str = Query(..., examples=["samples"]),
        file_name: str = Query(..., examples=["test.txt"]),
    ):
    '''
    下载知识库文档
    '''
    if not validate_kb_name(knowledge_base_name):
        return BaseResponse(code=403, msg="Don't attack me")

    kb = KBServiceFactory.get_service_by_name(knowledge_base_name)
    if kb is None:
        return BaseResponse(code=404, msg=f"未找到知识库 {knowledge_base_name}")

    try:
        kb_file = DocumentFile(filename=file_name,
                                knowledge_base_name=knowledge_base_name)

        if os.path.exists(kb_file.filepath):
            return FileResponse(
                path=kb_file.filepath,
                filename=kb_file.filename,
                media_type="multipart/form-data")
    except Exception as e:
        print(e)
        return BaseResponse(code=500, msg=f"{kb_file.filename} 读取文件失败，错误信息是：{e}")

    return BaseResponse(code=500, msg=f"{kb_file.filename} 读取文件失败")


async def recreate_vector_store(
        knowledge_base_name: str = Body(..., examples=["samples"]),
        allow_empty_kb: bool = Body(True),
        vs_type: str = Body(DEFAULT_VS_TYPE),
        embed_model: str = Body(EMBEDDING_MODEL),
    ):
    '''
    recreate vector store from the content.
    this is usefull when user can copy files to content folder directly instead of upload through network.
    by default, get_service_by_name only return knowledge base in the info.db and having document files in it.
    set allow_empty_kb to True make it applied on empty knowledge base which it not in the info.db or having no documents.
    '''

    async def output():
        kb = KBServiceFactory.get_service(knowledge_base_name, vs_type, embed_model)
        if not kb.exists() and not allow_empty_kb:
            yield {"code": 404, "msg": f"未找到知识库 ‘{knowledge_base_name}’"}
        else:
            kb.create_kb()
            kb.clear_vs()
            docs = list_docs_from_folder(knowledge_base_name)
            for i, doc in enumerate(docs):
                try:
                    kb_file = DocumentFile(doc, knowledge_base_name)
                    yield json.dumps({
                        "code": 200,
                        "msg": f"({i + 1} / {len(docs)}): {doc}",
                        "total": len(docs),
                        "finished": i,
                        "doc": doc,
                    }, ensure_ascii=False)
                    if i == len(docs) - 1:
                        not_refresh_vs_cache = False
                    else:
                        not_refresh_vs_cache = True
                    kb.add_doc(kb_file, not_refresh_vs_cache=not_refresh_vs_cache)
                except Exception as e:
                    print(e)
                    yield json.dumps({
                        "code": 500,
                        "msg": f"添加文件‘{doc}’到知识库‘{knowledge_base_name}’时出错：{e}。已跳过。",
                    })

    return StreamingResponse(output(), media_type="text/event-stream")

dev_opsgpt/service/kb_api.py
codefuse-ai/codefuse-chatbot
from multiprocessing import Process, Queue
import multiprocessing as mp
import sys
import os

src_dir = os.path.join(
    os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
)
print(src_dir)
sys.path.append(src_dir)

sys.path.append(os.path.dirname(os.path.dirname(__file__)))
from configs.model_config import llm_model_dict, LLM_MODEL, LLM_DEVICE, LOG_PATH, logger
from dev_opsgpt.utils.server_utils import MakeFastAPIOffline


host_ip = "0.0.0.0"
controller_port = 20001
model_worker_port = 20002
openai_api_port = 8888
base_url = "http://127.0.0.1:{}"

os.environ['PATH'] = os.environ.get("PATH", "") + os.pathsep + r'/d/env_utils/miniconda3/envs/devopsgpt/Lib/site-packages/torch/lib'

def set_httpx_timeout(timeout=60.0):
    import httpx
    httpx._config.DEFAULT_TIMEOUT_CONFIG.connect = timeout
    httpx._config.DEFAULT_TIMEOUT_CONFIG.read = timeout
    httpx._config.DEFAULT_TIMEOUT_CONFIG.write = timeout


def create_controller_app(
        dispatch_method="shortest_queue",
):
    import fastchat.constants
    fastchat.constants.LOGDIR = LOG_PATH
    from fastchat.serve.controller import app, Controller

    controller = Controller(dispatch_method)
    sys.modules["fastchat.serve.controller"].controller = controller

    MakeFastAPIOffline(app)
    app.title = "FastChat Controller"
    return app


def create_model_worker_app(
        worker_address=base_url.format(model_worker_port),
        controller_address=base_url.format(controller_port),
        model_path=llm_model_dict[LLM_MODEL].get("local_model_path"),
        device=LLM_DEVICE,
        gpus=None,
        max_gpu_memory="8GiB",
        load_8bit=False,
        cpu_offloading=None,
        gptq_ckpt=None,
        gptq_wbits=16,
        gptq_groupsize=-1,
        gptq_act_order=False,
        awq_ckpt=None,
        awq_wbits=16,
        awq_groupsize=-1,
        model_names=[LLM_MODEL],
        num_gpus=1, # not in fastchat
        conv_template=None,
        limit_worker_concurrency=5,
        stream_interval=2,
        no_register=False,
):
    import fastchat.constants
    fastchat.constants.LOGDIR = LOG_PATH
    from fastchat.serve.model_worker import app, GptqConfig, AWQConfig, ModelWorker, worker_id
    import argparse
    import threading
    import fastchat.serve.model_worker

    # workaround to make program exit with Ctrl+c
    # it should be deleted after pr is merged by fastchat
    def _new_init_heart_beat(self):
        self.register_to_controller()
        self.heart_beat_thread = threading.Thread(
            target=fastchat.serve.model_worker.heart_beat_worker, args=(self,), daemon=True,
        )
        self.heart_beat_thread.start()
    ModelWorker.init_heart_beat = _new_init_heart_beat

    parser = argparse.ArgumentParser()
    args = parser.parse_args()
    args.model_path = model_path
    args.model_names = model_names
    args.device = device
    args.load_8bit = load_8bit
    args.gptq_ckpt = gptq_ckpt
    args.gptq_wbits = gptq_wbits
    args.gptq_groupsize = gptq_groupsize
    args.gptq_act_order = gptq_act_order
    args.awq_ckpt = awq_ckpt
    args.awq_wbits = awq_wbits
    args.awq_groupsize = awq_groupsize
    args.gpus = gpus
    args.num_gpus = num_gpus
    args.max_gpu_memory = max_gpu_memory
    args.cpu_offloading = cpu_offloading
    args.worker_address = worker_address
    args.controller_address = controller_address
    args.conv_template = conv_template
    args.limit_worker_concurrency = limit_worker_concurrency
    args.stream_interval = stream_interval
    args.no_register = no_register

    if args.gpus:
        if len(args.gpus.split(",")) < args.num_gpus:
            raise ValueError(
                f"Larger --num-gpus ({args.num_gpus}) than --gpus {args.gpus}!"
            )
        os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus

    if gpus and num_gpus is None:
        num_gpus = len(gpus.split(','))
    args.num_gpus = num_gpus

    gptq_config = GptqConfig(
        ckpt=gptq_ckpt or model_path,
        wbits=args.gptq_wbits,
        groupsize=args.gptq_groupsize,
        act_order=args.gptq_act_order,
    )
    awq_config = AWQConfig(
        ckpt=args.awq_ckpt or args.model_path,
        wbits=args.awq_wbits,
        groupsize=args.awq_groupsize,
    )

    # torch.multiprocessing.set_start_method('spawn')
    worker = ModelWorker(
        controller_addr=args.controller_address,
        worker_addr=args.worker_address,
        worker_id=worker_id,
        model_path=args.model_path,
        model_names=args.model_names,
        limit_worker_concurrency=args.limit_worker_concurrency,
        no_register=args.no_register,
        device=args.device,
        num_gpus=args.num_gpus,
        max_gpu_memory=args.max_gpu_memory,
        load_8bit=args.load_8bit,
        cpu_offloading=args.cpu_offloading,
        gptq_config=gptq_config,
        awq_config=awq_config,
        stream_interval=args.stream_interval,
        conv_template=args.conv_template,
    )

    sys.modules["fastchat.serve.model_worker"].worker = worker
    sys.modules["fastchat.serve.model_worker"].args = args
    sys.modules["fastchat.serve.model_worker"].gptq_config = gptq_config
    
    MakeFastAPIOffline(app)
    app.title = f"FastChat LLM Server ({LLM_MODEL})"
    return app


def create_openai_api_app(
        controller_address=base_url.format(controller_port),
        api_keys=[],
):
    import fastchat.constants
    fastchat.constants.LOGDIR = LOG_PATH
    from fastchat.serve.openai_api_server import app, CORSMiddleware, app_settings

    app.add_middleware(
        CORSMiddleware,
        allow_credentials=True,
        allow_origins=["*"],
        allow_methods=["*"],
        allow_headers=["*"],
    )

    app_settings.controller_address = controller_address
    app_settings.api_keys = api_keys

    MakeFastAPIOffline(app)
    app.title = "FastChat OpeanAI API Server"
    return app


def run_controller(q):
    import uvicorn
    app = create_controller_app()

    @app.on_event("startup")
    async def on_startup():
        set_httpx_timeout()
        q.put(1)

    uvicorn.run(app, host=host_ip, port=controller_port)


def run_model_worker(q, *args, **kwargs):
    import uvicorn
    app = create_model_worker_app(*args, **kwargs)

    @app.on_event("startup")
    async def on_startup():
        set_httpx_timeout()
        while True:
            no = q.get()
            if no != 1:
                q.put(no)
            else:
                break
        q.put(2)

    uvicorn.run(app, host=host_ip, port=model_worker_port)


def run_openai_api(q):
    import uvicorn
    app = create_openai_api_app()

    @app.on_event("startup")
    async def on_startup():
        set_httpx_timeout()
        while True:
            no = q.get()
            if no != 2:
                q.put(no)
            else:
                break
        q.put(3)

    uvicorn.run(app, host=host_ip, port=openai_api_port)


if __name__ == "__main__":
    mp.set_start_method("spawn")
    queue = Queue()
    logger.info(llm_model_dict[LLM_MODEL])
    model_path = llm_model_dict[LLM_MODEL]["local_model_path"]

    logger.info(f"如需查看 llm_api 日志，请前往 {LOG_PATH}")

    if not model_path:
        logger.error("local_model_path 不能为空")
    else:
        controller_process = Process(
            target=run_controller,
            name=f"controller({os.getpid()})",
            args=(queue,),
            daemon=True,
        )
        controller_process.start()

        model_worker_process = Process(
            target=run_model_worker,
            name=f"model_worker({os.getpid()})",
            args=(queue,),
            # kwargs={"load_8bit": True},
            daemon=True,
        )
        model_worker_process.start()

        openai_api_process = Process(
            target=run_openai_api,
            name=f"openai_api({os.getpid()})",
            args=(queue,),
            daemon=True,
        )
        openai_api_process.start()

        try:
            model_worker_process.join()
            controller_process.join()
            openai_api_process.join()
        except KeyboardInterrupt:
            model_worker_process.terminate()
            controller_process.terminate()
            openai_api_process.terminate()

# 服务启动后接口调用示例：
# import openai
# openai.api_key = "EMPTY" # Not support yet
# openai.api_base = "http://localhost:8888/v1"

# model = "chatglm2-6b"

# # create a chat completion
# completion = openai.ChatCompletion.create(
#   model=model,
#   messages=[{"role": "user", "content": "Hello! What is your name?"}]
# )
# # print the completion
# print(completion.choices[0].message.content)

dev_opsgpt/service/llm_api.py
codefuse-ai/codefuse-chatbot
import os
from typing import Literal, Callable, Any

from configs.model_config import EMBEDDING_MODEL, DEFAULT_VS_TYPE

from dev_opsgpt.orm.utils import DocumentFile
from dev_opsgpt.orm.commands import add_doc_to_db
from dev_opsgpt.utils.path_utils import *
from .service_factory import KBServiceFactory


def folder2db(
    kb_name: str,
    mode: Literal["recreate_vs", "fill_info_only", "update_in_db", "increament"],
    vs_type: Literal["faiss", "milvus", "pg", "chromadb"] = DEFAULT_VS_TYPE,
    embed_model: str = EMBEDDING_MODEL,
    callback_before: Callable = None,
    callback_after: Callable = None,
):
    '''
    use existed files in local folder to populate database and/or vector store.
    set parameter `mode` to:
        recreate_vs: recreate all vector store and fill info to database using existed files in local folder
        fill_info_only: do not create vector store, fill info to db using existed files only
        update_in_db: update vector store and database info using local files that existed in database only
        increament: create vector store and database info for local files that not existed in database only
    '''
    kb = KBServiceFactory.get_service(kb_name, vs_type, embed_model)
    kb.create_kb()

    if mode == "recreate_vs":
        kb.clear_vs()
        docs = list_docs_from_folder(kb_name)
        for i, doc in enumerate(docs):
            try:
                kb_file = DocumentFile(doc, kb_name)
                if callable(callback_before):
                    callback_before(kb_file, i, docs)
                if i == len(docs) - 1:
                    not_refresh_vs_cache = False
                else:
                    not_refresh_vs_cache = True
                kb.add_doc(kb_file, not_refresh_vs_cache=not_refresh_vs_cache)
                if callable(callback_after):
                    callback_after(kb_file, i, docs)
            except Exception as e:
                print(e)
    elif mode == "fill_info_only":
        docs = list_docs_from_folder(kb_name)
        for i, doc in enumerate(docs):
            try:
                kb_file = DocumentFile(doc, kb_name)
                if callable(callback_before):
                    callback_before(kb_file, i, docs)
                add_doc_to_db(kb_file)
                if callable(callback_after):
                    callback_after(kb_file, i, docs)
            except Exception as e:
                print(e)
    elif mode == "update_in_db":
        docs = kb.list_docs()
        for i, doc in enumerate(docs):
            try:
                kb_file = DocumentFile(doc, kb_name)
                if callable(callback_before):
                    callback_before(kb_file, i, docs)
                if i == len(docs) - 1:
                    not_refresh_vs_cache = False
                else:
                    not_refresh_vs_cache = True
                kb.update_doc(kb_file, not_refresh_vs_cache=not_refresh_vs_cache)
                if callable(callback_after):
                    callback_after(kb_file, i, docs)
            except Exception as e:
                print(e)
    elif mode == "increament":
        db_docs = kb.list_docs()
        folder_docs = list_docs_from_folder(kb_name)
        docs = list(set(folder_docs) - set(db_docs))
        for i, doc in enumerate(docs):
            try:
                kb_file = DocumentFile(doc, kb_name)
                if callable(callback_before):
                    callback_before(kb_file, i, docs)
                if i == len(docs) - 1:
                    not_refresh_vs_cache = False
                else:
                    not_refresh_vs_cache = True
                kb.add_doc(kb_file, not_refresh_vs_cache=not_refresh_vs_cache)
                if callable(callback_after):
                    callback_after(kb_file, i, docs)
            except Exception as e:
                print(e)
    else:
        raise ValueError(f"unspported migrate mode: {mode}")


def recreate_all_vs(
    vs_type: Literal["faiss", "milvus", "pg", "chromadb"] = DEFAULT_VS_TYPE,
    embed_mode: str = EMBEDDING_MODEL,
    **kwargs: Any,
):
    '''
    used to recreate a vector store or change current vector store to another type or embed_model
    '''
    for kb_name in list_kbs_from_folder():
        folder2db(kb_name, "recreate_vs", vs_type, embed_mode, **kwargs)


def prune_db_docs(kb_name: str):
    '''
    delete docs in database that not existed in local folder.
    it is used to delete database docs after user deleted some doc files in file browser
    '''
    kb = KBServiceFactory.get_service_by_name(kb_name)
    if kb.exists():
        docs_in_db = kb.list_docs()
        docs_in_folder = list_docs_from_folder(kb_name)
        docs = list(set(docs_in_db) - set(docs_in_folder))
        for doc in docs:
            kb.delete_doc(DocumentFile(doc, kb_name))
        return docs

def prune_folder_docs(kb_name: str):
    '''
    delete doc files in local folder that not existed in database.
    is is used to free local disk space by delete unused doc files.
    '''
    kb = KBServiceFactory.get_service_by_name(kb_name)
    if kb.exists():
        docs_in_db = kb.list_docs()
        docs_in_folder = list_docs_from_folder(kb_name)
        docs = list(set(docs_in_folder) - set(docs_in_db))
        for doc in docs:
            os.remove(get_file_path(kb_name, doc))
        return docs

dev_opsgpt/service/migrate.py
codefuse-ai/codefuse-chatbot
from typing import List, Union, Dict
import os

from configs.model_config import EMBEDDING_MODEL

from .faiss_db_service import FaissKBService
from .base_service import KBService, SupportedVSType
from dev_opsgpt.orm.commands import *
from dev_opsgpt.utils.path_utils import *



class KBServiceFactory:

    @staticmethod
    def get_service(kb_name: str,
                    vector_store_type: Union[str, SupportedVSType],
                    embed_model: str = EMBEDDING_MODEL,
                    ) -> KBService:
        if isinstance(vector_store_type, str):
            vector_store_type = getattr(SupportedVSType, vector_store_type.upper())
        if SupportedVSType.FAISS == vector_store_type:
            return FaissKBService(kb_name, embed_model=embed_model)
        # if SupportedVSType.PG == vector_store_type:
        #     from server.knowledge_base.kb_service.pg_kb_service import PGKBService
        #     return PGKBService(kb_name, embed_model=embed_model)
        # elif SupportedVSType.MILVUS == vector_store_type:
        #     from server.knowledge_base.kb_service.milvus_kb_service import MilvusKBService
        #     return MilvusKBService(kb_name, embed_model=embed_model) # other milvus parameters are set in model_config.kbs_config
        # elif SupportedVSType.DEFAULT == vector_store_type: # kb_exists of default kbservice is False, to make validation easier.
        #     from server.knowledge_base.kb_service.default_kb_service import DefaultKBService
        #     return DefaultKBService(kb_name)

    @staticmethod
    def get_service_by_name(kb_name: str
                            ) -> KBService:
        _, vs_type, embed_model = load_kb_from_db(kb_name)
        if vs_type is None and os.path.isdir(get_kb_path(kb_name)): # faiss knowledge base not in db
            vs_type = "faiss"
        return KBServiceFactory.get_service(kb_name, vs_type, embed_model)

    @staticmethod
    def get_default():
        return KBServiceFactory.get_service("default", SupportedVSType.DEFAULT)


def get_kb_details() -> List[Dict]:
    kbs_in_folder = list_kbs_from_folder()
    kbs_in_db = KBService.list_kbs()
    result = {}

    for kb in kbs_in_folder:
        result[kb] = {
            "kb_name": kb,
            "vs_type": "",
            "embed_model": "",
            "file_count": 0,
            "create_time": None,
            "in_folder": True,
            "in_db": False,
        }

    for kb in kbs_in_db:
        kb_detail = get_kb_detail(kb)
        if kb_detail:
            kb_detail["in_db"] = True
            if kb in result:
                result[kb].update(kb_detail)
            else:
                kb_detail["in_folder"] = False
                result[kb] = kb_detail

    data = []
    for i, v in enumerate(result.values()):
        v['No'] = i + 1
        data.append(v)
   
    return data


def get_kb_doc_details(kb_name: str) -> List[Dict]:
    kb = KBServiceFactory.get_service_by_name(kb_name)
    docs_in_folder = list_docs_from_folder(kb_name)
    docs_in_db = kb.list_docs()
    result = {}

    for doc in docs_in_folder:
        result[doc] = {
            "kb_name": kb_name,
            "file_name": doc,
            "file_ext": os.path.splitext(doc)[-1],
            "file_version": 0,
            "document_loader": "",
            "text_splitter": "",
            "create_time": None,
            "in_folder": True,
            "in_db": False,
        }
    for doc in docs_in_db:
        doc_detail = get_file_detail(kb_name, doc)
        if doc_detail:
            doc_detail["in_db"] = True
            if doc in result:
                result[doc].update(doc_detail)
            else:
                doc_detail["in_folder"] = False
                result[doc] = doc_detail

    data = []
    for i, v in enumerate(result.values()):
        v['No'] = i + 1
        data.append(v)
   
    return data

dev_opsgpt/service/service_factory.py
codefuse-ai/codefuse-chatbot
from .langchain_splitter import LCTextSplitter

__all__ = ["LCTextSplitter"]
dev_opsgpt/text_splitter/__init__.py
codefuse-ai/codefuse-chatbot
import os
import importlib
from loguru import logger

from langchain.document_loaders.base import BaseLoader
from langchain.text_splitter import (
    SpacyTextSplitter, RecursiveCharacterTextSplitter
)

from configs.model_config import (
    CHUNK_SIZE,
    OVERLAP_SIZE,
    ZH_TITLE_ENHANCE
)
from dev_opsgpt.utils.path_utils import *



class LCTextSplitter:
    '''langchain textsplitter 执行file2text'''
    def __init__(
            self, filepath: str, text_splitter_name: str = None
    ):
        self.filepath = filepath
        self.ext = os.path.splitext(filepath)[-1].lower()
        self.text_splitter_name = text_splitter_name
        if self.ext not in SUPPORTED_EXTS:
            raise ValueError(f"暂未支持的文件格式 {self.ext}")
        self.document_loader_name = get_LoaderClass(self.ext)

    def file2text(self, ):
        loader = self._load_document()
        text_splitter = self._load_text_splitter()
        if self.document_loader_name in ["JSONLoader", "JSONLLoader"]:
            docs = loader.load()
        else:
            docs = loader.load_and_split(text_splitter)
        logger.info(docs[0])
        return docs

    def _load_document(self, ) -> BaseLoader:
        DocumentLoader = EXT2LOADER_DICT[self.ext]
        if self.document_loader_name == "UnstructuredFileLoader":
            loader = DocumentLoader(self.filepath, autodetect_encoding=True)
        else:
            loader = DocumentLoader(self.filepath)
        return loader
    
    def _load_text_splitter(self, ):
        try:
            if self.text_splitter_name is None:
                text_splitter = SpacyTextSplitter(
                    pipeline="zh_core_web_sm",
                    chunk_size=CHUNK_SIZE,
                    chunk_overlap=OVERLAP_SIZE,
                )
                self.text_splitter_name = "SpacyTextSplitter"
            elif self.document_loader_name in ["JSONLoader", "JSONLLoader"]:
                text_splitter = None
            else:
                text_splitter_module = importlib.import_module('langchain.text_splitter')
                TextSplitter = getattr(text_splitter_module, self.text_splitter_name)
                text_splitter = TextSplitter(
                    chunk_size=CHUNK_SIZE,
                    chunk_overlap=OVERLAP_SIZE)
        except Exception as e:
            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=CHUNK_SIZE,
                chunk_overlap=OVERLAP_SIZE,
            )
        return text_splitter

dev_opsgpt/text_splitter/langchain_splitter.py
codefuse-ai/codefuse-chatbot
from .server_utils import BaseResponse, ListResponse
from .common_utils import func_timer

__all__ = [
    "BaseResponse", "ListResponse", "func_timer"
]
dev_opsgpt/utils/__init__.py
codefuse-ai/codefuse-chatbot
import textwrap, time, copy, random, hashlib, json, os
from datetime import datetime, timedelta
from functools import wraps
from loguru import logger



DATE_FORMAT = "%Y-%m-%d %H:%M:%S"


def timestampToDateformat(ts, interval=1000, dateformat=DATE_FORMAT):
    '''将标准时间戳转换标准指定时间格式'''
    return datetime.fromtimestamp(ts//interval).strftime(dateformat)


def datefromatToTimestamp(dt, interval=1000, dateformat=DATE_FORMAT):
    '''将标准时间格式转换未标准时间戳'''
    return datetime.strptime(dt, dateformat).timestamp()*interval


def func_timer():
    '''
    用装饰器实现函数计时
    :param function: 需要计时的函数
    :return: None
    '''
    @wraps(function)
    def function_timer(*args, **kwargs):
        t0 = time.time()
        result = function(*args, **kwargs)
        t1 = time.time()
        logger.info('[Function: {name} finished, spent time: {time:.3f}s]'.format(
            name=function.__name__,
            time=t1 - t0
        ))
        return result
    return function_timer


def read_jsonl_file(filename):
    data = []
    with open(filename, "r", encoding="utf-8") as f:
        for line in f:
            data.append(json.loads(line))
    return data


def save_to_jsonl_file(data, filename):
    dir_name = os.path.dirname(filename)
    if not os.path.exists(dir_name): os.makedirs(dir_name)

    with open(filename, "w", encoding="utf-8") as f:
        for item in data:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")


def read_json_file(filename):
    with open(filename, "r", encoding="utf-8") as f:
        return json.load(f)
    
    
def save_to_json_file(data, filename):
    dir_name = os.path.dirname(filename)
    if not os.path.exists(dir_name): os.makedirs(dir_name)

    with open(filename, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)

dev_opsgpt/utils/common_utils.py
codefuse-ai/codefuse-chatbot
import os
from langchain.document_loaders import CSVLoader, PyPDFLoader, UnstructuredFileLoader, TextLoader, PythonLoader
from langchain.embeddings.huggingface import HuggingFaceEmbeddings

from dev_opsgpt.document_loaders import JSONLLoader, JSONLoader
from configs.model_config import (
    embedding_model_dict,
    KB_ROOT_PATH,
)
from loguru import logger


LOADERNAME2LOADER_DICT = {
    "UnstructuredFileLoader": UnstructuredFileLoader,
    "CSVLoader": CSVLoader,
    "PyPDFLoader": PyPDFLoader,
    "TextLoader": TextLoader,
    "PythonLoader": PythonLoader,
    "JSONLoader": JSONLoader,
    "JSONLLoader": JSONLLoader
}

LOADER2EXT_DICT = {"UnstructuredFileLoader": ['.eml', '.html', '.md', '.msg', '.rst',
                                          '.rtf', '.xml',
                                          '.doc', '.docx', '.epub', '.odt', 
                                          '.ppt', '.pptx', '.tsv'],
               "CSVLoader": [".csv"],
               "PyPDFLoader": [".pdf"],
                "TextLoader": ['.txt'],
                "PythonLoader": ['.py'],
                "JSONLoader": ['.json'],
                "JSONLLoader": ['.jsonl'],
               }

EXT2LOADER_DICT = {ext: LOADERNAME2LOADER_DICT[k] for k, exts in LOADER2EXT_DICT.items() for ext in exts}

SUPPORTED_EXTS = [ext for sublist in LOADER2EXT_DICT.values() for ext in sublist]


def validate_kb_name(knowledge_base_id: str) -> bool:
    # 检查是否包含预期外的字符或路径攻击关键字
    if "../" in knowledge_base_id:
        return False
    return True

def get_kb_path(knowledge_base_name: str):
    return os.path.join(KB_ROOT_PATH, knowledge_base_name)

def get_doc_path(knowledge_base_name: str):
    return os.path.join(get_kb_path(knowledge_base_name), "content")

def get_vs_path(knowledge_base_name: str):
    return os.path.join(get_kb_path(knowledge_base_name), "vector_store")

def get_file_path(knowledge_base_name: str, doc_name: str):
    return os.path.join(get_doc_path(knowledge_base_name), doc_name)

def list_kbs_from_folder():
    return [f for f in os.listdir(KB_ROOT_PATH)
            if os.path.isdir(os.path.join(KB_ROOT_PATH, f))]

def list_docs_from_folder(kb_name: str):
    doc_path = get_doc_path(kb_name)
    return [file for file in os.listdir(doc_path)
            if os.path.isfile(os.path.join(doc_path, file))]

def get_LoaderClass(file_extension):
    for LoaderClass, extensions in LOADER2EXT_DICT.items():
        if file_extension in extensions:
            return LoaderClass
dev_opsgpt/utils/path_utils.py
codefuse-ai/codefuse-chatbot
import pydantic
from pydantic import BaseModel
from typing import List
import torch
from fastapi import FastAPI
from pathlib import Path
import asyncio
from typing import Any, Optional
from loguru import logger


class BaseResponse(BaseModel):
    code: int = pydantic.Field(200, description="API status code")
    msg: str = pydantic.Field("success", description="API status message")

    class Config:
        schema_extra = {
            "example": {
                "code": 200,
                "msg": "success",
            }
        }

class ListResponse(BaseResponse):
    data: List[str] = pydantic.Field(..., description="List of names")

    class Config:
        schema_extra = {
            "example": {
                "code": 200,
                "msg": "success",
                "data": ["doc1.docx", "doc2.pdf", "doc3.txt"],
            }
        }


class ChatMessage(BaseModel):
    question: str = pydantic.Field(..., description="Question text")
    response: str = pydantic.Field(..., description="Response text")
    history: List[List[str]] = pydantic.Field(..., description="History text")
    source_documents: List[str] = pydantic.Field(
        ..., description="List of source documents and their scores"
    )

    class Config:
        schema_extra = {
            "example": {
                "question": "工伤保险如何办理？",
                "response": "根据已知信息，可以总结如下：\n\n1. 参保单位为员工缴纳工伤保险费，以保障员工在发生工伤时能够获得相应的待遇。\n"
                            "2. 不同地区的工伤保险缴费规定可能有所不同，需要向当地社保部门咨询以了解具体的缴费标准和规定。\n"
                            "3. 工伤从业人员及其近亲属需要申请工伤认定，确认享受的待遇资格，并按时缴纳工伤保险费。\n"
                            "4. 工伤保险待遇包括工伤医疗、康复、辅助器具配置费用、伤残待遇、工亡待遇、一次性工亡补助金等。\n"
                            "5. 工伤保险待遇领取资格认证包括长期待遇领取人员认证和一次性待遇领取人员认证。\n"
                            "6. 工伤保险基金支付的待遇项目包括工伤医疗待遇、康复待遇、辅助器具配置费用、一次性工亡补助金、丧葬补助金等。",
                "history": [
                    [
                        "工伤保险是什么？",
                        "工伤保险是指用人单位按照国家规定，为本单位的职工和用人单位的其他人员，缴纳工伤保险费，"
                        "由保险机构按照国家规定的标准，给予工伤保险待遇的社会保险制度。",
                    ]
                ],
                "source_documents": [
                    "出处 [1] 广州市单位从业的特定人员参加工伤保险办事指引.docx：\n\n\t"
                    "( 一)  从业单位  (组织)  按“自愿参保”原则，  为未建 立劳动关系的特定从业人员单项参加工伤保险 、缴纳工伤保 险费。",
                    "出处 [2] ...",
                    "出处 [3] ...",
                ],
            }
        }

def torch_gc():
    if torch.cuda.is_available():
        # with torch.cuda.device(DEVICE):
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
    elif torch.backends.mps.is_available():
        try:
            from torch.mps import empty_cache
            empty_cache()
        except Exception as e:
            print(e)
            print("如果您使用的是 macOS 建议将 pytorch 版本升级至 2.0.0 或更高版本，以支持及时清理 torch 产生的内存占用。")


def run_async(cor):
    '''
    在同步环境中运行异步代码.
    '''
    try:
        loop = asyncio.get_event_loop()
    except:
        loop = asyncio.new_event_loop()
    return loop.run_until_complete(cor)


def iter_over_async(ait, loop):
    '''
    将异步生成器封装成同步生成器.
    '''
    ait = ait.__aiter__()
    async def get_next():
        try:
            obj = await ait.__anext__()
            return False, obj
        except StopAsyncIteration:
            return True, None
    while True:
        done, obj = loop.run_until_complete(get_next())
        if done:
            break
        yield obj




def MakeFastAPIOffline(
    app: FastAPI,
    static_dir = Path(__file__).parent / "static",
    static_url = "/static-offline-docs",
    docs_url: Optional[str] = "/docs",
    redoc_url: Optional[str] = "/redoc",
) -> None:
    """patch the FastAPI obj that doesn't rely on CDN for the documentation page"""
    from fastapi import Request
    from fastapi.openapi.docs import (
        get_redoc_html,
        get_swagger_ui_html,
        get_swagger_ui_oauth2_redirect_html,
    )
    from fastapi.staticfiles import StaticFiles
    from starlette.responses import HTMLResponse

    openapi_url = app.openapi_url
    swagger_ui_oauth2_redirect_url = app.swagger_ui_oauth2_redirect_url

    def remove_route(url: str) -> None:
        '''
        remove original route from app
        '''
        index = None
        for i, r in enumerate(app.routes):
            if r.path.lower() == url.lower():
                index = i
                break
        if isinstance(index, int):
            app.routes.pop(i)

    # Set up static file mount
    app.mount(
        static_url,
        StaticFiles(directory=Path(static_dir).as_posix()),
        name="static-offline-docs",
    )

    if docs_url is not None:
        remove_route(docs_url)
        remove_route(swagger_ui_oauth2_redirect_url)

        # Define the doc and redoc pages, pointing at the right files
        @app.get(docs_url, include_in_schema=False)
        async def custom_swagger_ui_html(request: Request) -> HTMLResponse:
            root = request.scope.get("root_path")
            favicon = f"{root}{static_url}/favicon.png"
            return get_swagger_ui_html(
                openapi_url=f"{root}{openapi_url}",
                title=app.title + " - Swagger UI",
                oauth2_redirect_url=swagger_ui_oauth2_redirect_url,
                swagger_js_url=f"{root}{static_url}/swagger-ui-bundle.js",
                swagger_css_url=f"{root}{static_url}/swagger-ui.css",
                swagger_favicon_url=favicon,
            )

        @app.get(swagger_ui_oauth2_redirect_url, include_in_schema=False)
        async def swagger_ui_redirect() -> HTMLResponse:
            return get_swagger_ui_oauth2_redirect_html()

    if redoc_url is not None:
        remove_route(redoc_url)

        @app.get(redoc_url, include_in_schema=False)
        async def redoc_html(request: Request) -> HTMLResponse:
            root = request.scope.get("root_path")
            favicon = f"{root}{static_url}/favicon.png"

            return get_redoc_html(
                openapi_url=f"{root}{openapi_url}",
                title=app.title + " - ReDoc",
                redoc_js_url=f"{root}{static_url}/redoc.standalone.js",
                with_google_fonts=False,
                redoc_favicon_url=favicon,
            )

dev_opsgpt/utils/server_utils.py
codefuse-ai/codefuse-chatbot
from .dialogue import dialogue_page, chat_box
from .document import knowledge_page
from .prompt import prompt_page
from .utils import ApiRequest

__all__ = [
    "dialogue_page", "chat_box", "prompt_page", "knowledge_page",
    "ApiRequest"
]
dev_opsgpt/webui/__init__.py
codefuse-ai/codefuse-chatbot
import streamlit as st
from streamlit_chatbox import *
from typing import List, Dict
from datetime import datetime

from .utils import *
from dev_opsgpt.utils import *
from dev_opsgpt.chat.search_chat import SEARCH_ENGINES

chat_box = ChatBox(
    assistant_avatar="../sources/imgs/devops-chatbot2.png"
)

GLOBAL_EXE_CODE_TEXT = ""


def get_messages_history(history_len: int) -> List[Dict]:
    def filter(msg):
        '''
        针对当前简单文本对话，只返回每条消息的第一个element的内容
        '''
        content = [x._content for x in msg["elements"] if x._output_method in ["markdown", "text"]]
        return {
            "role": msg["role"],
            "content": content[0] if content else "",
        }

    history = chat_box.filter_history(100000, filter)  # workaround before upgrading streamlit-chatbox.
    user_count = 0
    i = 1
    for i in range(1, len(history) + 1):
        if history[-i]["role"] == "user":
            user_count += 1
            if user_count >= history_len:
                break
    return history[-i:]


def dialogue_page(api: ApiRequest):
    global GLOBAL_EXE_CODE_TEXT
    chat_box.init_session()

    with st.sidebar:
        # TODO: 对话模型与会话绑定
        def on_mode_change():
            mode = st.session_state.dialogue_mode
            text = f"已切换到 {mode} 模式。"
            if mode == "知识库问答":
                cur_kb = st.session_state.get("selected_kb")
                if cur_kb:
                    text = f"{text} 当前知识库： `{cur_kb}`。"
            st.toast(text)
            # sac.alert(text, description="descp", type="success", closable=True, banner=True)

        dialogue_mode = st.selectbox("请选择对话模式",
                                     ["LLM 对话",
                                      "知识库问答",
                                      "搜索引擎问答",
                                      ],
                                     on_change=on_mode_change,
                                     key="dialogue_mode",
                                     )
        history_len = st.number_input("历史对话轮数：", 0, 10, 3)

        # todo: support history len

        def on_kb_change():
            st.toast(f"已加载知识库： {st.session_state.selected_kb}")

        if dialogue_mode == "知识库问答":
            with st.expander("知识库配置", True):
                kb_list = api.list_knowledge_bases(no_remote_api=True)
                selected_kb = st.selectbox(
                    "请选择知识库：",
                    kb_list,
                    on_change=on_kb_change,
                    key="selected_kb",
                )
                kb_top_k = st.number_input("匹配知识条数：", 1, 20, 3)
                score_threshold = st.number_input("知识匹配分数阈值：", 0.0, float(SCORE_THRESHOLD), float(SCORE_THRESHOLD), float(SCORE_THRESHOLD//100))
                # chunk_content = st.checkbox("关联上下文", False, disabled=True)
                # chunk_size = st.slider("关联长度：", 0, 500, 250, disabled=True)
        elif dialogue_mode == "搜索引擎问答":
            with st.expander("搜索引擎配置", True):
                search_engine = st.selectbox("请选择搜索引擎", SEARCH_ENGINES.keys(), 0)
                se_top_k = st.number_input("匹配搜索结果条数：", 1, 20, 3)

        code_interpreter_on = st.toggle("开启代码解释器")
        code_exec_on = st.toggle("自动执行代码")

    # Display chat messages from history on app rerun

    chat_box.output_messages()

    chat_input_placeholder = "请输入对话内容，换行请使用Ctrl+Enter "
    code_text = "" or GLOBAL_EXE_CODE_TEXT
    codebox_res = None

    if prompt := st.chat_input(chat_input_placeholder, key="prompt"):
        history = get_messages_history(history_len)
        chat_box.user_say(prompt)
        if dialogue_mode == "LLM 对话":
            chat_box.ai_say("正在思考...")
            text = ""
            r = api.chat_chat(prompt, history)
            for t in r:
                if error_msg := check_error_msg(t): # check whether error occured
                    st.error(error_msg)
                    break
                text += t["answer"]
                chat_box.update_msg(text)
            logger.debug(f"text: {text}")
            chat_box.update_msg(text, streaming=False)  # 更新最终的字符串，去除光标
            # 判断是否存在代码, 并提高编辑功能，执行功能
            code_text = api.codebox.decode_code_from_text(text)
            GLOBAL_EXE_CODE_TEXT = code_text
            if code_text and code_exec_on:
                codebox_res = api.codebox_chat("```"+code_text+"```", do_code_exe=True)

        elif dialogue_mode == "知识库问答":
            history = get_messages_history(history_len)
            chat_box.ai_say([
                f"正在查询知识库 `{selected_kb}` ...",
                Markdown("...", in_expander=True, title="知识库匹配结果"),
            ])
            text = ""
            for idx_count, d in enumerate(api.knowledge_base_chat(prompt, selected_kb, kb_top_k, score_threshold, history)):
                if error_msg := check_error_msg(d): # check whether error occured
                    st.error(error_msg)
                text += d["answer"]
                if idx_count%10 == 0:
                    chat_box.update_msg(text, element_index=0)
                # chat_box.update_msg("知识库匹配结果: \n\n".join(d["docs"]), element_index=1, streaming=False, state="complete")
            chat_box.update_msg(text, element_index=0, streaming=False)  # 更新最终的字符串，去除光标
            chat_box.update_msg("知识库匹配结果: \n\n".join(d["docs"]), element_index=1, streaming=False, state="complete")
            # 判断是否存在代码, 并提高编辑功能，执行功能
            code_text = api.codebox.decode_code_from_text(text)
            GLOBAL_EXE_CODE_TEXT = code_text
            if code_text and code_exec_on:
                codebox_res = api.codebox_chat("```"+code_text+"```", do_code_exe=True)
        elif dialogue_mode == "搜索引擎问答":
            chat_box.ai_say([
                f"正在执行 `{search_engine}` 搜索...",
                Markdown("...", in_expander=True, title="网络搜索结果"),
            ])
            text = ""
            d = {"docs": []}
            for d in api.search_engine_chat(prompt, search_engine, se_top_k):
                if error_msg := check_error_msg(d): # check whether error occured
                    st.error(error_msg)
                text += d["answer"]
                if idx_count%10 == 0:
                    chat_box.update_msg(text, element_index=0)
                # chat_box.update_msg("搜索匹配结果: \n\n".join(d["docs"]), element_index=1, streaming=False)
            chat_box.update_msg(text, element_index=0, streaming=False)  # 更新最终的字符串，去除光标
            chat_box.update_msg("搜索匹配结果: \n\n".join(d["docs"]), element_index=1, streaming=False, state="complete")
            # 判断是否存在代码, 并提高编辑功能，执行功能
            code_text = api.codebox.decode_code_from_text(text)
            GLOBAL_EXE_CODE_TEXT = code_text
            if code_text and code_exec_on:
                codebox_res = api.codebox_chat("```"+code_text+"```", do_code_exe=True)

    if code_interpreter_on:
        with st.expander("代码编辑执行器", False):
            code_part = st.text_area("代码片段", code_text, key="code_text")
            cols = st.columns(2)
            if cols[0].button(
                "修改对话",
                use_container_width=True,
            ):
                code_text = code_part
                GLOBAL_EXE_CODE_TEXT = code_text
                st.toast("修改对话成功")

            if cols[1].button(
                "执行代码",
                use_container_width=True
            ):
                if code_text:
                    codebox_res = api.codebox_chat("```"+code_text+"```", do_code_exe=True)
                    st.toast("正在执行代码")
                else:
                    st.toast("code 不能为空")

    #TODO 这段信息会被记录到history里
    if codebox_res is not None and codebox_res.code_exe_status != 200:
        st.toast(f"{codebox_res.code_exe_response}")

    if codebox_res is not None and codebox_res.code_exe_status == 200:
        st.toast(f"codebox_chajt {codebox_res}")
        chat_box.ai_say(Markdown(code_text, in_expander=True, title="code interpreter", unsafe_allow_html=True), )
        if codebox_res.code_exe_type == "image/png":
            base_text = f"```\n{code_text}\n```\n\n"
            img_html = "<img src='data:image/png;base64,{}' class='img-fluid'>".format(
                codebox_res.code_exe_response
            )
            chat_box.update_msg(base_text + img_html, streaming=False, state="complete")
        else:
            chat_box.update_msg('```\n'+code_text+'\n```'+"\n\n"+'```\n'+codebox_res.code_exe_response+'\n```', 
                                streaming=False, state="complete")

    now = datetime.now()
    with st.sidebar:

        cols = st.columns(2)
        export_btn = cols[0]
        if cols[1].button(
                "清空对话",
                use_container_width=True,
        ):
            chat_box.reset_history()
            GLOBAL_EXE_CODE_TEXT = ""
            st.experimental_rerun()

    export_btn.download_button(
        "导出记录",
        "".join(chat_box.export2md()),
        file_name=f"{now:%Y-%m-%d %H.%M}_对话记录.md",
        mime="text/markdown",
        use_container_width=True,
    )

dev_opsgpt/webui/dialogue.py
codefuse-ai/codefuse-chatbot
import streamlit as st
import os
import time
import traceback
from typing import Literal, Dict, Tuple
from st_aggrid import AgGrid, JsCode
from st_aggrid.grid_options_builder import GridOptionsBuilder
import pandas as pd

from configs.model_config import embedding_model_dict, kbs_config, EMBEDDING_MODEL, DEFAULT_VS_TYPE, WEB_CRAWL_PATH
from .utils import *
from dev_opsgpt.utils.path_utils import *
from dev_opsgpt.service.service_factory import get_kb_details, get_kb_doc_details
from dev_opsgpt.orm import table_init


# SENTENCE_SIZE = 100

cell_renderer = JsCode("""function(params) {if(params.value==true){return '✓'}else{return '×'}}""")


def config_aggrid(
        df: pd.DataFrame,
        columns: Dict[Tuple[str, str], Dict] = {},
        selection_mode: Literal["single", "multiple", "disabled"] = "single",
        use_checkbox: bool = False,
) -> GridOptionsBuilder:
    gb = GridOptionsBuilder.from_dataframe(df)
    gb.configure_column("No", width=40)
    for (col, header), kw in columns.items():
        gb.configure_column(col, header, wrapHeaderText=True, **kw)
    gb.configure_selection(
        selection_mode=selection_mode,
        use_checkbox=use_checkbox,
        # pre_selected_rows=st.session_state.get("selected_rows", [0]),
    )
    return gb


def file_exists(kb: str, selected_rows: List) -> Tuple[str, str]:
    '''
    check whether a doc file exists in local knowledge base folder.
    return the file's name and path if it exists.
    '''
    if selected_rows:
        file_name = selected_rows[0]["file_name"]
        file_path = get_file_path(kb, file_name)
        if os.path.isfile(file_path):
            return file_name, file_path
    return "", ""


def knowledge_page(api: ApiRequest):
    # 判断表是否存在并进行初始化
    table_init()

    try:
        kb_list = {x["kb_name"]: x for x in get_kb_details()}
    except Exception as e:
        st.error("获取知识库信息错误，请检查是否已按照 `README.md` 中 `4 知识库初始化与迁移` 步骤完成初始化或迁移，或是否为数据库连接错误。")
        st.stop()
    kb_names = list(kb_list.keys())

    if "selected_kb_name" in st.session_state and st.session_state["selected_kb_name"] in kb_names:
        selected_kb_index = kb_names.index(st.session_state["selected_kb_name"])
    else:
        selected_kb_index = 0

    def format_selected_kb(kb_name: str) -> str:
        if kb := kb_list.get(kb_name):
            return f"{kb_name} ({kb['vs_type']} @ {kb['embed_model']})"
        else:
            return kb_name

    selected_kb = st.selectbox(
        "请选择或新建知识库：",
        kb_names + ["新建知识库"],
        format_func=format_selected_kb,
        index=selected_kb_index
    )

    if selected_kb == "新建知识库":
        with st.form("新建知识库"):

            kb_name = st.text_input(
                "新建知识库名称",
                placeholder="新知识库名称，不支持中文命名",
                key="kb_name",
            )

            cols = st.columns(2)

            vs_types = list(kbs_config.keys())
            vs_type = cols[0].selectbox(
                "向量库类型",
                vs_types,
                index=vs_types.index(DEFAULT_VS_TYPE),
                key="vs_type",
            )

            embed_models = list(embedding_model_dict.keys())

            embed_model = cols[1].selectbox(
                "Embedding 模型",
                embed_models,
                index=embed_models.index(EMBEDDING_MODEL),
                key="embed_model",
            )

            submit_create_kb = st.form_submit_button(
                "新建",
                # disabled=not bool(kb_name),
                use_container_width=True,
            )

        if submit_create_kb:
            if not kb_name or not kb_name.strip():
                st.error(f"知识库名称不能为空！")
            elif kb_name in kb_list:
                st.error(f"名为 {kb_name} 的知识库已经存在！")
            else:
                ret = api.create_knowledge_base(
                    knowledge_base_name=kb_name,
                    vector_store_type=vs_type,
                    embed_model=embed_model,
                )
                st.toast(ret.get("msg", " "))
                st.session_state["selected_kb_name"] = kb_name
                st.experimental_rerun()

    elif selected_kb:
        kb = selected_kb

        # 上传文件
        # sentence_size = st.slider("文本入库分句长度限制", 1, 1000, SENTENCE_SIZE, disabled=True)
        files = st.file_uploader("上传知识文件",
                                 [i for ls in LOADER2EXT_DICT.values() for i in ls],
                                 accept_multiple_files=True,
                                 )
        
        base_url = st.text_input(
            "待获取内容的URL地址",
            placeholder="请填写正确可打开的URL地址",
            key="base_url",
        )


        if st.button(
            "添加URL内容到知识库",
            disabled= base_url is None or base_url=="",
            ):
            filename = base_url.replace("https://", " ").\
                replace("http://", " ").replace("/", " ").\
                replace("?", " ").replace("=", " ").replace(".", " ").strip()
            html_name = "_".join(filename.split(" ",) + ["html.jsonl"])
            text_name = "_".join(filename.split(" ",) + ["text.jsonl"])
            html_path = os.path.join(WEB_CRAWL_PATH, html_name,)
            text_path = os.path.join(WEB_CRAWL_PATH, text_name,)
            # if not os.path.exists(text_dir) or :
            st.toast(base_url)
            st.toast(html_path)
            st.toast(text_path)
            res = api.web_crawl(
                base_url=base_url,
                html_dir=html_path,
                text_dir=text_path,
                do_dfs = False,
                reptile_lib="requests",
                method="get",
                time_sleep=2,
                )
            
            if res["status"] == 200:
                st.toast(res["response"], icon="✔")
                data = [{"file": text_path, "filename": text_name, "knowledge_base_name": kb, "not_refresh_vs_cache": False}]
                for k in data:
                    ret = api.upload_kb_doc(**k)
                    logger.info(ret)
                    if msg := check_success_msg(ret):
                        st.toast(msg, icon="✔")
                    elif msg := check_error_msg(ret):
                        st.toast(msg, icon="✖")
                st.session_state.files = []
            else:
                st.toast(res["response"], icon="✖")

            if os.path.exists(html_path):
                os.remove(html_path)

        if st.button(
                "添加文件到知识库",
                # help="请先上传文件，再点击添加",
                # use_container_width=True,
                disabled=len(files) == 0,
        ):
            data = [{"file": f, "knowledge_base_name": kb, "not_refresh_vs_cache": True} for f in files]
            data[-1]["not_refresh_vs_cache"]=False
            for k in data:
                ret = api.upload_kb_doc(**k)
                if msg := check_success_msg(ret):
                    st.toast(msg, icon="✔")
                elif msg := check_error_msg(ret):
                    st.toast(msg, icon="✖")
            st.session_state.files = []

        st.divider()

        # 知识库详情
        # st.info("请选择文件，点击按钮进行操作。")
        doc_details = pd.DataFrame(get_kb_doc_details(kb))
        if not len(doc_details):
            st.info(f"知识库 `{kb}` 中暂无文件")
        else:
            st.write(f"知识库 `{kb}` 中已有文件:")
            st.info("知识库中包含源文件与向量库，请从下表中选择文件后操作")
            doc_details.drop(columns=["kb_name"], inplace=True)
            doc_details = doc_details[[
                "No", "file_name", "document_loader", "text_splitter", "in_folder", "in_db",
            ]]
            # doc_details["in_folder"] = doc_details["in_folder"].replace(True, "✓").replace(False, "×")
            # doc_details["in_db"] = doc_details["in_db"].replace(True, "✓").replace(False, "×")
            gb = config_aggrid(
                doc_details,
                {
                    ("No", "序号"): {},
                    ("file_name", "文档名称"): {},
                    # ("file_ext", "文档类型"): {},
                    # ("file_version", "文档版本"): {},
                    ("document_loader", "文档加载器"): {},
                    ("text_splitter", "分词器"): {},
                    # ("create_time", "创建时间"): {},
                    ("in_folder", "源文件"): {"cellRenderer": cell_renderer},
                    ("in_db", "向量库"): {"cellRenderer": cell_renderer},
                },
                "multiple",
            )

            doc_grid = AgGrid(
                doc_details,
                gb.build(),
                columns_auto_size_mode="FIT_CONTENTS",
                theme="alpine",
                custom_css={
                    "#gridToolBar": {"display": "none"},
                },
                allow_unsafe_jscode=True
            )

            selected_rows = doc_grid.get("selected_rows", [])

            cols = st.columns(4)
            file_name, file_path = file_exists(kb, selected_rows)
            if file_path:
                with open(file_path, "rb") as fp:
                    cols[0].download_button(
                        "下载选中文档",
                        fp,
                        file_name=file_name,
                        use_container_width=True, )
            else:
                cols[0].download_button(
                    "下载选中文档",
                    "",
                    disabled=True,
                    use_container_width=True, )

            st.write()
            # 将文件分词并加载到向量库中
            if cols[1].button(
                    "重新添加至向量库" if selected_rows and (pd.DataFrame(selected_rows)["in_db"]).any() else "添加至向量库",
                    disabled=not file_exists(kb, selected_rows)[0],
                    use_container_width=True,
            ):
                for row in selected_rows:
                    api.update_kb_doc(kb, row["file_name"])
                st.experimental_rerun()

            # 将文件从向量库中删除，但不删除文件本身。
            if cols[2].button(
                    "从向量库删除",
                    disabled=not (selected_rows and selected_rows[0]["in_db"]),
                    use_container_width=True,
            ):
                for row in selected_rows:
                    api.delete_kb_doc(kb, row["file_name"])
                st.experimental_rerun()

            if cols[3].button(
                    "从知识库中删除",
                    type="primary",
                    use_container_width=True,
            ):
                for row in selected_rows:
                    ret = api.delete_kb_doc(kb, row["file_name"], True)
                    st.toast(ret.get("msg", " "))
                st.experimental_rerun()

        st.divider()

        cols = st.columns(3)

        # todo: freezed
        if cols[0].button(
                "依据源文件重建向量库",
                # help="无需上传文件，通过其它方式将文档拷贝到对应知识库content目录下，点击本按钮即可重建知识库。",
                use_container_width=True,
                type="primary",
        ):
            with st.spinner("向量库重构中，请耐心等待，勿刷新或关闭页面。"):
                empty = st.empty()
                empty.progress(0.0, "")
                for d in api.recreate_vector_store(kb):
                    if msg := check_error_msg(d):
                        st.toast(msg)
                    else:
                        empty.progress(d["finished"] / d["total"], f"正在处理： {d['doc']}")
                st.experimental_rerun()

        if cols[2].button(
                "删除知识库",
                use_container_width=True,
        ):
            ret = api.delete_knowledge_base(kb)
            st.toast(ret.get("msg", " "))
            time.sleep(1)
            st.experimental_rerun()

dev_opsgpt/webui/document.py
codefuse-ai/codefuse-chatbot
import streamlit as st
import os
import time
from datetime import datetime
import traceback
from typing import Literal, Dict, Tuple
from st_aggrid import AgGrid, JsCode
from st_aggrid.grid_options_builder import GridOptionsBuilder
import pandas as pd

from configs.model_config import embedding_model_dict, kbs_config, EMBEDDING_MODEL, DEFAULT_VS_TYPE
from .utils import *
from dev_opsgpt.utils.path_utils import *
from dev_opsgpt.service.service_factory import get_kb_details, get_kb_doc_details
from dev_opsgpt.orm import table_init



def prompt_page(api: ApiRequest):
    # 判断表是否存在并进行初始化
    table_init()

    now = datetime.now()
    with st.sidebar:

        cols = st.columns(2)
        export_btn = cols[0]
        if cols[1].button(
                "清空prompt",
                use_container_width=True,
        ):
            st.experimental_rerun()

    export_btn.download_button(
        "导出记录",
        "测试prompt",
        file_name=f"{now:%Y-%m-%d %H.%M}_对话记录.md",
        mime="text/markdown",
        use_container_width=True,
    )

dev_opsgpt/webui/prompt.py
codefuse-ai/codefuse-chatbot
# 该文件包含webui通用工具，可以被不同的webui使用
from typing import *
from pathlib import Path
from io import BytesIO
import httpx
import asyncio
from fastapi.responses import StreamingResponse
import contextlib
import json
import nltk
import traceback
from loguru import logger

from configs.model_config import (
    EMBEDDING_MODEL,
    DEFAULT_VS_TYPE,
    KB_ROOT_PATH,
    LLM_MODEL,
    SCORE_THRESHOLD,
    VECTOR_SEARCH_TOP_K,
    SEARCH_ENGINE_TOP_K,
    NLTK_DATA_PATH,
    logger,
)

from configs.server_config import SANDBOX_SERVER

from dev_opsgpt.utils.server_utils import run_async, iter_over_async
from dev_opsgpt.service.kb_api import *
from dev_opsgpt.chat import LLMChat, SearchChat, KnowledgeChat
from dev_opsgpt.sandbox import PyCodeBox, CodeBoxResponse

from web_crawler.utils.WebCrawler import WebCrawler

nltk.data.path = [NLTK_DATA_PATH] + nltk.data.path


def set_httpx_timeout(timeout=60.0):
    '''
    设置httpx默认timeout到60秒。
    httpx默认timeout是5秒，在请求LLM回答时不够用。
    '''
    httpx._config.DEFAULT_TIMEOUT_CONFIG.connect = timeout
    httpx._config.DEFAULT_TIMEOUT_CONFIG.read = timeout
    httpx._config.DEFAULT_TIMEOUT_CONFIG.write = timeout


KB_ROOT_PATH = Path(KB_ROOT_PATH)
set_httpx_timeout()


class ApiRequest:
    '''
    api.py调用的封装,主要实现:
    1. 简化api调用方式
    2. 实现无api调用(直接运行server.chat.*中的视图函数获取结果),无需启动api.py
    '''
    def __init__(
        self,
        base_url: str = "http://127.0.0.1:7861",
        timeout: float = 60.0,
        no_remote_api: bool = False,   # call api view function directly
    ):
        self.base_url = base_url
        self.timeout = timeout
        self.no_remote_api = no_remote_api
        self.llmChat = LLMChat()
        self.searchChat = SearchChat()
        self.knowledgeChat = KnowledgeChat()
        self.codebox = PyCodeBox(
            remote_url=SANDBOX_SERVER["url"],
            remote_ip=SANDBOX_SERVER["host"], # "http://localhost",
            remote_port=SANDBOX_SERVER["port"],
            token="mytoken",
            do_code_exe=True,
            do_remote=SANDBOX_SERVER["do_remote"]
            )

    def codebox_chat(self, text: str, file_path: str = None, do_code_exe: bool = None) -> CodeBoxResponse:
        return self.codebox.chat(text, file_path, do_code_exe=do_code_exe)

    def _parse_url(self, url: str) -> str:
        if (not url.startswith("http")
                    and self.base_url
                ):
            part1 = self.base_url.strip(" /")
            part2 = url.strip(" /")
            return f"{part1}/{part2}"
        else:
            return url

    def get(
        self,
        url: str,
        params: Union[Dict, List[Tuple], bytes] = None,
        retry: int = 3,
        stream: bool = False,
        **kwargs: Any,
    ) -> Union[httpx.Response, None]:
        url = self._parse_url(url)
        kwargs.setdefault("timeout", self.timeout)
        while retry > 0:
            try:
                if stream:
                    return httpx.stream("GET", url, params=params, **kwargs)
                else:
                    return httpx.get(url, params=params, **kwargs)
            except Exception as e:
                logger.error(e)
                retry -= 1

    async def aget(
        self,
        url: str,
        params: Union[Dict, List[Tuple], bytes] = None,
        retry: int = 3,
        stream: bool = False,
        **kwargs: Any,
    ) -> Union[httpx.Response, None]:
        url = self._parse_url(url)
        kwargs.setdefault("timeout", self.timeout)
        async with httpx.AsyncClient() as client:
            while retry > 0:
                try:
                    if stream:
                        return await client.stream("GET", url, params=params, **kwargs)
                    else:
                        return await client.get(url, params=params, **kwargs)
                except Exception as e:
                    logger.error(e)
                    retry -= 1

    def post(
        self,
        url: str,
        data: Dict = None,
        json: Dict = None,
        retry: int = 3,
        stream: bool = False,
        **kwargs: Any
    ) -> Union[httpx.Response, None]:
        url = self._parse_url(url)
        kwargs.setdefault("timeout", self.timeout)
        while retry > 0:
            try:
                # return requests.post(url, data=data, json=json, stream=stream, **kwargs)
                if stream:
                    return httpx.stream("POST", url, data=data, json=json, **kwargs)
                else:
                    return httpx.post(url, data=data, json=json, **kwargs)
            except Exception as e:
                logger.error(e)
                retry -= 1

    async def apost(
        self,
        url: str,
        data: Dict = None,
        json: Dict = None,
        retry: int = 3,
        stream: bool = False,
        **kwargs: Any
    ) -> Union[httpx.Response, None]:
        url = self._parse_url(url)
        kwargs.setdefault("timeout", self.timeout)
        async with httpx.AsyncClient() as client:
            while retry > 0:
                try:
                    if stream:
                        return await client.stream("POST", url, data=data, json=json, **kwargs)
                    else:
                        return await client.post(url, data=data, json=json, **kwargs)
                except Exception as e:
                    logger.error(e)
                    retry -= 1

    def delete(
        self,
        url: str,
        data: Dict = None,
        json: Dict = None,
        retry: int = 3,
        stream: bool = False,
        **kwargs: Any
    ) -> Union[httpx.Response, None]:
        url = self._parse_url(url)
        kwargs.setdefault("timeout", self.timeout)
        while retry > 0:
            try:
                if stream:
                    return httpx.stream("DELETE", url, data=data, json=json, **kwargs)
                else:
                    return httpx.delete(url, data=data, json=json, **kwargs)
            except Exception as e:
                logger.error(e)
                retry -= 1

    async def adelete(
        self,
        url: str,
        data: Dict = None,
        json: Dict = None,
        retry: int = 3,
        stream: bool = False,
        **kwargs: Any
    ) -> Union[httpx.Response, None]:
        url = self._parse_url(url)
        kwargs.setdefault("timeout", self.timeout)
        async with httpx.AsyncClient() as client:
            while retry > 0:
                try:
                    if stream:
                        return await client.stream("DELETE", url, data=data, json=json, **kwargs)
                    else:
                        return await client.delete(url, data=data, json=json, **kwargs)
                except Exception as e:
                    logger.error(e)
                    retry -= 1

    def _fastapi_stream2generator(self, response: StreamingResponse, as_json: bool =False):
        '''
        将api.py中视图函数返回的StreamingResponse转化为同步生成器
        '''
        try:
            loop = asyncio.get_event_loop()
        except:
            loop = asyncio.new_event_loop()
        
        try:
            for chunk in  iter_over_async(response.body_iterator, loop):
                if as_json and chunk:
                    yield json.loads(chunk)
                elif chunk.strip():
                    yield chunk
        except Exception as e:
            logger.error(traceback.format_exc())

    def _httpx_stream2generator(
        self,
        response: contextlib._GeneratorContextManager,
        as_json: bool = False,
    ):
        '''
        将httpx.stream返回的GeneratorContextManager转化为普通生成器
        '''
        try:
            with response as r:
                for chunk in r.iter_text(None):
                    if as_json and chunk:
                        yield json.loads(chunk)
                    elif chunk.strip():
                        yield chunk
        except httpx.ConnectError as e:
            msg = f"无法连接API服务器，请确认 ‘api.py’ 已正常启动。"
            logger.error(msg)
            logger.error(e)
            yield {"code": 500, "msg": msg}
        except httpx.ReadTimeout as e:
            msg = f"API通信超时，请确认已启动FastChat与API服务（详见RADME '5. 启动 API 服务或 Web UI'）"
            logger.error(msg)
            logger.error(e)
            yield {"code": 500, "msg": msg}
        except Exception as e:
            logger.error(e)
            yield {"code": 500, "msg": str(e)}

    def chat_chat(
        self,
        query: str,
        history: List[Dict] = [],
        stream: bool = True,
        no_remote_api: bool = None,
    ):
        '''
        对应api.py/chat/chat接口
        '''
        if no_remote_api is None:
            no_remote_api = self.no_remote_api

        data = {
            "query": query,
            "history": history,
            "stream": stream,
        }

        if no_remote_api:
            response = self.llmChat.chat(**data)
            return self._fastapi_stream2generator(response, as_json=True)
        else:
            response = self.post("/chat/chat", json=data, stream=True)
            return self._httpx_stream2generator(response)

    def knowledge_base_chat(
        self,
        query: str,
        knowledge_base_name: str,
        top_k: int = VECTOR_SEARCH_TOP_K,
        score_threshold: float = SCORE_THRESHOLD,
        history: List[Dict] = [],
        stream: bool = True,
        no_remote_api: bool = None,
    ):
        '''
        对应api.py/chat/knowledge_base_chat接口
        '''
        if no_remote_api is None:
            no_remote_api = self.no_remote_api

        data = {
            "query": query,
            "engine_name": knowledge_base_name,
            "top_k": top_k,
            "score_threshold": score_threshold,
            "history": history,
            "stream": stream,
            "local_doc_url": no_remote_api,
        }

        if no_remote_api:
            response = self.knowledgeChat.chat(**data)
            return self._fastapi_stream2generator(response, as_json=True)
        else:
            response = self.post(
                "/chat/knowledge_base_chat",
                json=data,
                stream=True,
            )
            return self._httpx_stream2generator(response, as_json=True)

    def search_engine_chat(
        self,
        query: str,
        search_engine_name: str,
        top_k: int = SEARCH_ENGINE_TOP_K,
        stream: bool = True,
        no_remote_api: bool = None,
    ):
        '''
        对应api.py/chat/search_engine_chat接口
        '''
        if no_remote_api is None:
            no_remote_api = self.no_remote_api

        data = {
            "query": query,
            "engine_name": search_engine_name,
            "top_k": top_k,
            "history": [],
            "stream": stream,
        }

        if no_remote_api:
            response = self.searchChat.chat(**data)
            return self._fastapi_stream2generator(response, as_json=True)
        else:
            response = self.post(
                "/chat/search_engine_chat",
                json=data,
                stream=True,
            )
            return self._httpx_stream2generator(response, as_json=True)

    # 知识库相关操作

    def _check_httpx_json_response(
            self,
            response: httpx.Response,
            errorMsg: str = f"无法连接API服务器，请确认已执行python server\\api.py",
        ) -> Dict:
        '''
        check whether httpx returns correct data with normal Response.
        error in api with streaming support was checked in _httpx_stream2enerator
        '''
        try:
            return response.json()
        except Exception as e:
            logger.error(e)
            return {"code": 500, "msg": errorMsg or str(e)}

    def list_knowledge_bases(
        self,
        no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/list_knowledge_bases接口
        '''
        if no_remote_api is None:
            no_remote_api = self.no_remote_api

        if no_remote_api:
            response = run_async(list_kbs())
            return response.data
        else:
            response = self.get("/knowledge_base/list_knowledge_bases")
            data = self._check_httpx_json_response(response)
            return data.get("data", [])

    def create_knowledge_base(
        self,
        knowledge_base_name: str,
        vector_store_type: str = "faiss",
        embed_model: str = EMBEDDING_MODEL,
        no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/create_knowledge_base接口
        '''
        if no_remote_api is None:
            no_remote_api = self.no_remote_api

        data = {
            "knowledge_base_name": knowledge_base_name,
            "vector_store_type": vector_store_type,
            "embed_model": embed_model,
        }

        if no_remote_api:
            response = run_async(create_kb(**data))
            return response.dict()
        else:
            response = self.post(
                "/knowledge_base/create_knowledge_base",
                json=data,
            )
            return self._check_httpx_json_response(response)

    def delete_knowledge_base(
        self,
        knowledge_base_name: str,
        no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/delete_knowledge_base接口
        '''
        if no_remote_api is None:
            no_remote_api = self.no_remote_api

        if no_remote_api:
            response = run_async(delete_kb(knowledge_base_name))
            return response.dict()
        else:
            response = self.post(
                "/knowledge_base/delete_knowledge_base",
                json=f"{knowledge_base_name}",
            )
            return self._check_httpx_json_response(response)

    def list_kb_docs(
        self,
        knowledge_base_name: str,
        no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/list_docs接口
        '''
        if no_remote_api is None:
            no_remote_api = self.no_remote_api

        if no_remote_api:
            response = run_async(list_docs(knowledge_base_name))
            return response.data
        else:
            response = self.get(
                "/knowledge_base/list_docs",
                params={"knowledge_base_name": knowledge_base_name}
            )
            data = self._check_httpx_json_response(response)
            return data.get("data", [])

    def upload_kb_doc(
        self,
        file: Union[str, Path, bytes],
        knowledge_base_name: str,
        filename: str = None,
        override: bool = False,
        not_refresh_vs_cache: bool = False,
        no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/upload_docs接口
        '''
        if no_remote_api is None:
            no_remote_api = self.no_remote_api

        if isinstance(file, bytes): # raw bytes
            file = BytesIO(file)
        elif hasattr(file, "read"): # a file io like object
            filename = filename or file.name
        else: # a local path
            file = Path(file).absolute().open("rb")
            filename = filename or file.name

        if no_remote_api:
            from fastapi import UploadFile
            from tempfile import SpooledTemporaryFile

            temp_file = SpooledTemporaryFile(max_size=10 * 1024 * 1024)
            temp_file.write(file.read())
            temp_file.seek(0)
            response = run_async(upload_doc(
                UploadFile(file=temp_file, filename=filename),
                knowledge_base_name,
                override,
                not_refresh_vs_cache
            ))
            return response.dict()
        else:
            response = self.post(
                "/knowledge_base/upload_doc",
                data={
                    "knowledge_base_name": knowledge_base_name,
                    "override": override,
                    "not_refresh_vs_cache": not_refresh_vs_cache,
                },
                files={"file": (filename, file)},
            )
            return self._check_httpx_json_response(response)

    def delete_kb_doc(
        self,
        knowledge_base_name: str,
        doc_name: str,
        delete_content: bool = False,
        not_refresh_vs_cache: bool = False,
        no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/delete_doc接口
        '''
        if no_remote_api is None:
            no_remote_api = self.no_remote_api

        data = {
            "knowledge_base_name": knowledge_base_name,
            "doc_name": doc_name,
            "delete_content": delete_content,
            "not_refresh_vs_cache": not_refresh_vs_cache,
        }

        if no_remote_api:
            response = run_async(delete_doc(**data))
            return response.dict()
        else:
            response = self.post(
                "/knowledge_base/delete_doc",
                json=data,
            )
            return self._check_httpx_json_response(response)

    def update_kb_doc(
        self,
        knowledge_base_name: str,
        file_name: str,
        not_refresh_vs_cache: bool = False,
        no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/update_doc接口
        '''
        if no_remote_api is None:
            no_remote_api = self.no_remote_api

        if no_remote_api:
            response = run_async(update_doc(knowledge_base_name, file_name, not_refresh_vs_cache))
            return response.dict()
        else:
            response = self.post(
                "/knowledge_base/update_doc",
                json={
                    "knowledge_base_name": knowledge_base_name,
                    "file_name": file_name,
                    "not_refresh_vs_cache": not_refresh_vs_cache,
                },
            )
            return self._check_httpx_json_response(response)

    def recreate_vector_store(
        self,
        knowledge_base_name: str,
        allow_empty_kb: bool = True,
        vs_type: str = DEFAULT_VS_TYPE,
        embed_model: str = EMBEDDING_MODEL,
        no_remote_api: bool = None,
    ):
        '''
        对应api.py/knowledge_base/recreate_vector_store接口
        '''
        if no_remote_api is None:
            no_remote_api = self.no_remote_api

        data = {
            "knowledge_base_name": knowledge_base_name,
            "allow_empty_kb": allow_empty_kb,
            "vs_type": vs_type,
            "embed_model": embed_model,
        }

        if no_remote_api:
            response = run_async(recreate_vector_store(**data))
            return self._fastapi_stream2generator(response, as_json=True)
        else:
            response = self.post(
                "/knowledge_base/recreate_vector_store",
                json=data,
                stream=True,
                timeout=None,
            )
            return self._httpx_stream2generator(response, as_json=True)

    def web_crawl(
            self, 
            base_url: str,
            html_dir: str,
            text_dir: str,
            do_dfs: bool = False,
            reptile_lib: str = "requests",
            method: str = "get",
            time_sleep: float = 2,
            no_remote_api: bool = None
    ):
        '''
        根据url来检索
        '''
        async def _web_crawl(html_dir, text_dir, base_url, reptile_lib, method, time_sleep, do_dfs):
            wc = WebCrawler()
            try:
                if not do_dfs:
                    wc.webcrawler_single(html_dir=html_dir,
                                    text_dir=text_dir,
                                    base_url=base_url,
                                    reptile_lib=reptile_lib,
                                    method=method,
                                    time_sleep=time_sleep
                                    )
                else:
                    wc.webcrawler_1_degree(html_dir=html_dir,
                            text_dir=text_dir,
                            base_url=base_url,
                            reptile_lib=reptile_lib,
                            method=method,
                            time_sleep=time_sleep
                            )
                return {"status": 200, "response": "success"}
            except Exception as e:
                return {"status": 500, "response": str(e)}
                
        if no_remote_api is None:
            no_remote_api = self.no_remote_api

        data = {
            "base_url": base_url,
            "html_dir": html_dir,
            "text_dir": text_dir,
            "do_dfs": do_dfs,
            "reptile_lib": reptile_lib,
            "method": method,
            "time_sleep": time_sleep,
        }
        if no_remote_api:
            response = run_async(_web_crawl(**data))
            return response
        else:
            raise Exception("not impletenion")
    

def check_error_msg(data: Union[str, dict, list], key: str = "errorMsg") -> str:
    '''
    return error message if error occured when requests API
    '''
    if isinstance(data, dict):
        if key in data:
            return data[key]
        if "code" in data and data["code"] != 200:
            return data["msg"]
    return ""


def check_success_msg(data: Union[str, dict, list], key: str = "msg") -> str:
    '''
    return error message if error occured when requests API
    '''
    if (isinstance(data, dict)
        and key in data
        and "code" in data
        and data["code"] == 200):
        return data[key]
    return ""


if __name__ == "__main__":
    api = ApiRequest(no_remote_api=True)

    # print(api.chat_fastchat(
    #     messages=[{"role": "user", "content": "hello"}]
    # ))

    # with api.chat_chat("你好") as r:
    #     for t in r.iter_text(None):
    #         print(t)

    # r = api.chat_chat("你好", no_remote_api=True)
    # for t in r:
    #     print(t)

    # r = api.duckduckgo_search_chat("室温超导最新研究进展", no_remote_api=True)
    # for t in r:
    #     print(t)

    # print(api.list_knowledge_bases())

dev_opsgpt/webui/utils.py
codefuse-ai/codefuse-chatbot
from dataclasses import dataclass, field
import os
from os.path import isdir, isfile
from pathlib import Path
import sys

from transformers import AutoTokenizer


@dataclass
class GptqConfig:
    ckpt: str = field(
        default=None,
        metadata={
            "help": "Load quantized model. The path to the local GPTQ checkpoint."
        },
    )
    wbits: int = field(default=16, metadata={"help": "#bits to use for quantization"})
    groupsize: int = field(
        default=-1,
        metadata={"help": "Groupsize to use for quantization; default uses full row."},
    )
    act_order: bool = field(
        default=True,
        metadata={"help": "Whether to apply the activation order GPTQ heuristic"},
    )


def load_quant_by_autogptq(model):

    from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
    model = AutoGPTQForCausalLM.from_quantized(model, 
                                                inject_fused_attention=False,
                                                inject_fused_mlp=False,
                                                use_cuda_fp16=True,
                                                disable_exllama=False,
                                                device_map='auto'
                                                )
    return model
    
def load_gptq_quantized(model_name, gptq_config: GptqConfig):
    print("Loading GPTQ quantized model...")

    if gptq_config.act_order:
        try:
            script_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
            module_path = os.path.join(script_path, "../repositories/GPTQ-for-LLaMa")

            sys.path.insert(0, module_path)
            from llama import load_quant

            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
            # only `fastest-inference-4bit` branch cares about `act_order`
            model = load_quant(
                model_name,
                find_gptq_ckpt(gptq_config),
                gptq_config.wbits,
                gptq_config.groupsize,
                act_order=gptq_config.act_order,
            )
        except ImportError as e:
            print(f"Error: Failed to load GPTQ-for-LLaMa. {e}")
            print("See https://github.com/lm-sys/FastChat/blob/main/docs/gptq.md")
            sys.exit(-1)

    else:
        # other branches
        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
        model = load_quant_by_autogptq(model_name)

    return model, tokenizer


# def load_gptq_quantized(model_name, gptq_config: GptqConfig):
#     print("Loading GPTQ quantized model...")

#     try:
#         script_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
#         module_path = os.path.join(script_path, "repositories/GPTQ-for-LLaMa")

#         sys.path.insert(0, module_path)
#         from llama import load_quant
#     except ImportError as e:
#         print(f"Error: Failed to load GPTQ-for-LLaMa. {e}")
#         print("See https://github.com/lm-sys/FastChat/blob/main/docs/gptq.md")
#         sys.exit(-1)

#     tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
#     # only `fastest-inference-4bit` branch cares about `act_order`
#     if gptq_config.act_order:
#         model = load_quant(
#             model_name,
#             find_gptq_ckpt(gptq_config),
#             gptq_config.wbits,
#             gptq_config.groupsize,
#             act_order=gptq_config.act_order,
#         )
#     else:
#         # other branches
#         model = load_quant(
#             model_name,
#             find_gptq_ckpt(gptq_config),
#             gptq_config.wbits,
#             gptq_config.groupsize,
#         )

#     return model, tokenizer


def find_gptq_ckpt(gptq_config: GptqConfig):
    if Path(gptq_config.ckpt).is_file():
        return gptq_config.ckpt

    # for ext in ["*.pt", "*.safetensors",]:
    for ext in ["*.pt", "*.bin",]:
        matched_result = sorted(Path(gptq_config.ckpt).glob(ext))
        if len(matched_result) > 0:
            return str(matched_result[-1])

    print("Error: gptq checkpoint not found")
    sys.exit(1)

examples/gptq.py
codefuse-ai/codefuse-chatbot
import docker, sys, os, time, requests

from loguru import logger

src_dir = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)
sys.path.append(src_dir)

from configs.server_config import CONTRAINER_NAME, SANDBOX_SERVER, IMAGE_NAME


if SANDBOX_SERVER["do_remote"]:
    client = docker.from_env()
    for i  in client.containers.list(all=True):
        if i.name == CONTRAINER_NAME:
            container = i
            container.stop()
            container.remove()
        break
    # 启动容器
    logger.info("start ot init container & notebook")
    container = client.containers.run(
        image=IMAGE_NAME,
        command="bash",
        name=CONTRAINER_NAME,
        ports={"5050/tcp": SANDBOX_SERVER["port"]},
        stdin_open=True,
        detach=True,
        tty=True,
    )

    # 启动notebook
    exec_command = container.exec_run("bash jupyter_start.sh")

    # 判断notebook是否启动
    retry_nums = 3
    while retry_nums>0:
        response = requests.get(f"http://localhost:{SANDBOX_SERVER['port']}", timeout=270)
        if response.status_code == 200:
            logger.info("container & notebook init success")
            break
        else:
            retry_nums -= 1
            logger.info(client.containers.list())
            logger.info("wait container running ...")
        time.sleep(5)
else:
    logger.info("启动local的notebook环境支持代码执行")

examples/start_sandbox.py
codefuse-ai/codefuse-chatbot
import docker, sys, os, time, requests

from loguru import logger

src_dir = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)
sys.path.append(src_dir)

from configs.server_config import CONTRAINER_NAME, SANDBOX_SERVER


if SANDBOX_SERVER["do_remote"]:
    # stop and remove the container
    client = docker.from_env()
    for i  in client.containers.list(all=True):
        if i.name == CONTRAINER_NAME:
            container = i
            container.stop()
            container.remove()
        break
else:
    # stop local 
    import psutil
    for process in psutil.process_iter(["pid", "name", "cmdline"]):
        # check process name contains "jupyter" and port=xx
        if f"port={SANDBOX_SERVER['port']}" in str(process.info["cmdline"]).lower() and \
            "jupyter" in process.info['name'].lower():

            logger.warning(f"port={SANDBOX_SERVER['port']}, {process.info}")
            # 关闭进程
            process.terminate()

examples/stop_sandbox.py
codefuse-ai/codefuse-chatbot
# 运行方式：
# 1. 安装必要的包：pip install streamlit-option-menu streamlit-chatbox>=1.1.6
# 2. 运行本机fastchat服务：python server\llm_api.py 或者 运行对应的sh文件
# 3. 运行API服务器：python server/api.py。如果使用api = ApiRequest(no_remote_api=True)，该步可以跳过。
# 4. 运行WEB UI：streamlit run webui.py --server.port 7860

import os, sys
import streamlit as st
from streamlit_option_menu import option_menu

src_dir = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)
sys.path.append(src_dir)

from dev_opsgpt.webui import *
from configs import VERSION, LLM_MODEL



api = ApiRequest(base_url="http://127.0.0.1:7861", no_remote_api=True)


if __name__ == "__main__":
    st.set_page_config(
        "DevOpsGPT-Chat WebUI",
        os.path.join("../sources/imgs", "devops-chatbot.png"),
        initial_sidebar_state="expanded",
        menu_items={
            'Get Help': 'https://github.com/lightislost/devopsgpt',
            'Report a bug': "https://github.com/lightislost/devopsgpt/issues",
            'About': f"""欢迎使用 DevOpsGPT-Chat WebUI {VERSION}！"""
        }
    )

    if not chat_box.chat_inited:
        st.toast(
            f"欢迎使用 [`DevOpsGPT-Chat`](https://github.com/lightislost/devopsgpt) ! \n\n"
            f"当前使用模型`{LLM_MODEL}`, 您可以开始提问了."
        )

    pages = {
        "对话": {
            "icon": "chat",
            "func": dialogue_page,
        },
        "知识库管理": {
            "icon": "hdd-stack",
            "func": knowledge_page,
        },
        # "Prompt管理": {
        #     "icon": "hdd-stack",
        #     "func": prompt_page,
        # },
    }

    with st.sidebar:
        st.image(
            os.path.join(
                "../sources/imgs",
                "devops-chatbot.png"
            ),
            use_column_width=True
        )
        st.caption(
            f"""<p align="right">当前版本：{VERSION}</p>""",
            unsafe_allow_html=True,
        )
        options = list(pages)
        icons = [x["icon"] for x in pages.values()]

        default_index = 0
        selected_page = option_menu(
            "",
            options=options,
            icons=icons,
            # menu_icon="chat-quote",
            default_index=default_index,
        )

    if selected_page in pages:
        pages[selected_page]["func"](api)

examples/webui.py
codefuse-ai/codefuse-chatbot
import os, sys


src_dir = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    )
sys.path.append(src_dir)


from dev_opsgpt.orm import create_tables

create_tables()
tests/db_create_test.py
codefuse-ai/codefuse-chatbot
import time
import docker

st = time.time()
client = docker.from_env()
print(time.time()-st)


st = time.time()
client.containers.run("ubuntu:latest", "echo hello world")
print(time.time()-st)
tests/docker_test.py
codefuse-ai/codefuse-chatbot
import os, sys

src_dir = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)
sys.path.append(src_dir)

from configs import llm_model_dict, LLM_MODEL
import openai
# os.environ["OPENAI_PROXY"] = "socks5h://127.0.0.1:7890"
# os.environ["OPENAI_PROXY"] = "http://127.0.0.1:7890"
os.environ["OPENAI_API_KEY"] = ""


if __name__ == "__main__":
    # print("dsadsa", os.environ.get("OPENAI_PROXY"), os.environ.get("OPENAI_API_KEY"))
    from langchain import PromptTemplate, LLMChain
    from langchain.prompts.chat import ChatPromptTemplate
    from langchain.chat_models import ChatOpenAI
    from langchain.schema import HumanMessage
    
    # chat = ChatOpenAI(temperature=0.1, model_name="gpt-3.5-turbo")
    # print(chat.predict("hi!"))

    print(LLM_MODEL, llm_model_dict[LLM_MODEL]["api_key"], llm_model_dict[LLM_MODEL]["api_base_url"])
    model = ChatOpenAI(
        streaming=True,
        verbose=True,
        openai_api_key=llm_model_dict[LLM_MODEL]["api_key"],
        openai_api_base=llm_model_dict[LLM_MODEL]["api_base_url"],
        model_name=LLM_MODEL
    )
    chat_prompt = ChatPromptTemplate.from_messages([("human", "{input}")])
    chain = LLMChain(prompt=chat_prompt, llm=model)
    content = chain({"input": "hello"})
    print(content)
tests/openai_test.py
codefuse-ai/codefuse-chatbot
import json
from uuid import uuid4
import uuid 

import requests
from websockets.sync.client import connect as ws_connect_sync
from websocket import create_connection
from websockets.client import connect as ws_connect

import sys, os
src_dir = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)
sys.path.append(src_dir)


from dev_opsgpt.sandbox.pycodebox import PyCodeBox
from pathlib import Path


# python = Path(sys.executable).absolute()
# print(Path(python).as_posix())
# print(python)
# print(sys.executable)


# import requests

# # 设置Jupyter Notebook服务器的URL
# url = 'http://localhost:5050'  # 或者是你自己的Jupyter服务器的URL

# # 发送GET请求来获取Jupyter Notebook的登录页面
# response = requests.get(url)

# # 检查响应状态码
# if response.status_code == 200:
#     # 打印响应内容
#     print('connect success')
# else:
#     print('connect fail')

# import subprocess 
# jupyter = subprocess.Popen(
#                     [
#                         "juptyer", "notebnook",
#                         "--NotebookApp.token=mytoken",
#                         f"--port=4949",
#                         "--no-browser",
#                         "--ServerApp.disable_check_xsrf=True",
#                     ],
#                     stderr=subprocess.PIPE,
#                     stdin=subprocess.PIPE,
#                     stdout=subprocess.PIPE,
#                 )

# 测试1
import time, psutil
from loguru import logger
pycodebox = PyCodeBox(remote_url="http://localhost:5050", 
               remote_ip="http://localhost", 
            remote_port="5050", 
            token="mytoken",
            do_code_exe=True, 
            do_remote=False)


reuslt = pycodebox.chat("```print('hello world!')```", do_code_exe=True)
print(reuslt)

reuslt = pycodebox.chat("print('hello world!')", do_code_exe=False)
print(reuslt)

for process in psutil.process_iter(["pid", "name", "cmdline"]):
    # 检查进程名是否包含"jupyter"
    if 'port=5050' in str(process.info["cmdline"]).lower() and \
        "jupyter" in process.info['name'].lower():

        logger.warning(f'port=5050, {process.info}')
        # 关闭进程
        process.terminate()
        

# 测试2
# with PyCodeBox(remote_url="http://localhost:5050", 
#                remote_ip="http://localhost", 
#             remote_port="5050", 
#             token="mytoken",
#             do_code_exe=True, 
#             do_remote=True) as codebox:
    
#     result = codebox.run("print('hello world!')")
#     print(result)



# with PyCodeBox(
#     remote_ip="http://localhost", 
#     remote_port="5555", 
#     token="mytoken",
#     do_code_exe=True, 
#     do_remote=False
#     ) as codebox:
    
#     result = codebox.run("print('hello world!')")
    # print(result)



    
# headers = {'Authorization': 'Token mytoken', 'token': 'mytoken'}

# kernel_url = "http://localhost:5050/api/kernels"

# response = requests.get(kernel_url, headers=headers)
# if len(response.json())>0:
#     kernel_id = response.json()[0]["id"]
# else:
#     response = requests.post(kernel_url, headers=headers)
#     kernel_id = response.json()["id"]


# print(f"ws://localhost:5050/api/kernels/{kernel_id}/channels?token=mytoken")
# ws = create_connection(f"ws://localhost:5050/api/kernels/{kernel_id}/channels?token=mytoken", headers=headers)

# code_text = "print('hello world!')"
# # code_text = "import matplotlib.pyplot as plt\n\nplt.figure(figsize=(4,2))\nplt.plot([1,2,3,4,5])\nplt.show()"

# ws.send(
#             json.dumps(
#                 {
#                     "header": {
#                         "msg_id": (msg_id := uuid4().hex),
#                         "msg_type": "execute_request",
#                     },
#                     "parent_header": {},
#                     "metadata": {},
#                     "content": {
#                         "code": code_text,
#                         "silent": True,
#                         "store_history": True,
#                         "user_expressions": {},
#                         "allow_stdin": False,
#                         "stop_on_error": True,
#                     },
#                     "channel": "shell",
#                     "buffers": [],
#                 }
#             )
#         )

# while True:
#     received_msg = json.loads(ws.recv())
#     if received_msg["msg_type"] == "stream":
#         result_msg = received_msg  # 找到结果消息
#         break
#     elif received_msg["header"]["msg_type"] == "execute_result":
#         result_msg = received_msg  # 找到结果消息
#         break
#     elif received_msg["header"]["msg_type"] == "display_data":
#         result_msg = received_msg  # 找到结果消息
#         break
    
# print(received_msg)

tests/sandbox_test.py
codefuse-ai/codefuse-chatbot
from duckduckgo_search import DDGS

with DDGS(proxies="socks5://127.0.0.1:13659", timeout=20) as ddgs:
    for r in ddgs.text("马克龙、冯德莱恩访华"):
        print(r)

tests/search_test.py
codefuse-ai/codefuse-chatbot
import docker, sys, os, time, requests

from loguru import logger

src_dir = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)
sys.path.append(src_dir)

# from configs.server_config import CONTRAINER_NAME, SANDBOX_SERVER, IMAGE_NAME

CONTRAINER_NAME = "devopsgpt-chatbot_default"
SANDBOX_SERVER = {
    "port": "8501"
}
IMAGE_NAME = "devopsgpt-chatbot:pypy"

client = docker.from_env()
for i  in client.containers.list(all=True):
    if i.name == CONTRAINER_NAME:
        container = i
        container.stop()
        container.remove()
    break
# 启动容器
logger.info("start ot init container & notebook")
container = client.containers.run(
    image=IMAGE_NAME,
    command="bash",
    name=CONTRAINER_NAME,
    ports={"8501/tcp": SANDBOX_SERVER["port"]},
    stdin_open=True,
    detach=True,
    tty=True,
)

# 启动notebook
# exec_command = container.exec_run("bash jupyter_start.sh")


tests/start_webui_docker.py
codefuse-ai/codefuse-chatbot
import os, sys

src_dir = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)
sys.path.append(src_dir)

from dev_opsgpt.orm import create_tables

create_tables()
tests/table_test.py
codefuse-ai/codefuse-chatbot

import os, sys


src_dir = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    )
sys.path.append(src_dir)

from dev_opsgpt.text_splitter import LCTextSplitter

filepath = "D:/project/gitlab/llm/DevOpsGpt/knowledge_base/SAMPLES/content/test.txt" 
lcTextSplitter = LCTextSplitter(filepath)
docs = lcTextSplitter.file2text()

print(docs[0])
tests/text_spliter_test.py
codefuse-ai/codefuse-chatbot
import logging
from utils.WebCrawler import WebCrawler

logging.basicConfig(level=logging.INFO)

if __name__ == '__main__':
    # 保存地址，分别保存html源文件、处理后text文件
    html_dir = "data/html/tmp_csdn_122513786_html.jsonl"
    text_dir = "data/text/tmp_csdn_122513786_text.jsonl"
    # 下载网页数据
    # https://www.langchain.asia/
    # https://blog.csdn.net/weixin_43791511/article/details/122513786
    # https://zhuanlan.zhihu.com/p/645400277
    # https://www.aliyun.com/?utm_content=se_1014243503
    # 'https://cloud.tencent.com/developer/article/1004500?from=15425'
    base_url = 'https://www.langchain.asia/'
    # 爬取方式：
    ## requests和selenium两种方式；requests为简单请求静态网址html内容，js动态数据无法获取；
    ## selenium为模拟人行为请求，可获取全部html数据，但请求时间较长10-20s单网页，尽量设置5s以上的time_sleep。
    reptile_lib = "requests"
    method = "get"  # 目前只支持get请求
    time_sleep = 4  # 每两次请求间隔时间s
    wc = WebCrawler()
    # 爬取base_url单网址
    wc.webcrawler_single(html_dir=html_dir,
                         text_dir=text_dir,
                         base_url=base_url,
                         reptile_lib=reptile_lib,
                         method=method,
                         time_sleep=time_sleep
                         )

    # # 爬取base_url页面所有网址，限制target_url_prefix为前缀，默认target_url_prefix=base_url
    # wc.webcrawler_1_degree(html_dir=html_dir,
    #                        text_dir=text_dir,
    #                        base_url=base_url,
    #                        reptile_lib=reptile_lib,
    #                        method=method,
    #                        time_sleep=time_sleep
    #                        )

web_crawler/main_test.py
codefuse-ai/codefuse-chatbot

import re

# 非打印字符
NON_PRINTING_CHARS_RE = re.compile(
    f"[{''.join(map(chr, list(range(0, 32)) + list(range(127, 160))))}]"
)

class DocTokenizer():
    '''
    文档text处理器。
    '''

    def __init__(self):
        pass

    def doc_process(self, text):
        '''
        去除多余换行、去掉每行非打印字符和开头结尾空格
        '''
        # 去除多余换行
        text = self.remove_excess_lines(text)
        # 将文本拆分成行
        lines = text.split("\n")
        # 去掉每一行的开头和结尾的空格
        lines = [self.remove_non_printing_char_line(
            line.strip()) for line in lines]
        # 将行重新组合成文本
        text_new = "\n".join(lines)
        return text_new

    def remove_excess_lines(self, text):
        '''
        将2个以上的换行符替换为2个，html解析text时会产生大量换行\n
        '''
        pattern = r'\n\n+'
        return re.sub(pattern, '\n\n', text)

    def remove_non_printing_char_line(self, text):
        '''
        去除每一行的非打印字符
        '''
        return NON_PRINTING_CHARS_RE.sub("", text)

web_crawler/utils/DocTokenizer.py
codefuse-ai/codefuse-chatbot
import time
from bs4 import BeautifulSoup
import logging
import json
import os
from tqdm import tqdm
import re
from .DocTokenizer import DocTokenizer

logging.basicConfig(level=logging.INFO)


class Html2Text():
    '''从html中提取text文本内容。
    '''

    def __init__(self):
        pass

    def html2text(self,
                  target_content_tag={},
                  target_tag_list=[],
                  html_dir=None,
                  text_dir=None,
                  mode="w",
                  is_get_all_text=False
                  ):
        '''
        从html中提取text文本内容，需要指定提取html中的tag标签。输入为地址，html文件保存在jsonl文件中，输出也需要指定地址。
        :param target_content_tag: html中正文content所在tag，字典格式限制长度为1，key为选中便签类型name/class/id，vaule为标签取值如div/title/article等
        :param target_tag_list: 指定提取html对应的tag文本，列表，每个元素都与target_content_tag格式相同
        :param is_get_all_text: True则将html页面所有text内容保存到all_text字典中；False不保存all_text
        :param html_dir: html数据地址，注意需要时jsonl格式，一行为一个json字典，有text/url/host_url三个字段
        :param text_dir: 将提取的text内容保存的地址，同样是jsonl格式。
        :return: None
        '''
        assert isinstance(target_content_tag,dict), "target_content_tag请输入字典格式！"
        assert len(target_content_tag.keys()) <= 1,"target_content_tag属性字典只能指定唯一元素！"
        for _ in target_tag_list:
            assert isinstance(_, dict), "target_tag_list列表元素需要字典格式！"
            assert len(_.keys()) <= 1, "target_tag_list列表中的属性字典只能指定唯一元素！"
        # 创建保存目录
        os.makedirs(os.path.dirname(text_dir), exist_ok=True)
        # 读取文件
        logging.info("读取文件中……")
        html_dict_list = self.read_html_jsonl(html_dir)
        url_nums = len(html_dict_list)
        logging.info("共{url_nums}个html网址".format(url_nums=url_nums))
        # 循环处理每行html数据：html提取content正文、指定tag内容
        text_dict_list = []
        for html_dict in tqdm(html_dict_list, mininterval=1):
            # 是否获取全部text内容
            text_dict = self.get_text_dict(
                html_dict=html_dict,
                target_content_tag=target_content_tag,
                target_tag_list=target_tag_list,
                is_get_all_text=is_get_all_text
            )
            text_dict_list.append(text_dict)
        logging.info("保存html提取的text内容……")
        self.save_text_jsonl(json_list=text_dict_list,
                             file_path=text_dir,
                             mode=mode)
        logging.info("保存成功！地址：%s" % text_dir)

    def get_text_dict(self,
                      html_dict={},
                      target_content_tag={},
                      target_tag_list=[],
                      is_get_all_text=True
                      ):
        '''{"name":"div"}
        提取html网页字符中的纯文本内容，采用BeautifulSoup.get_text()获取全部text文本，target_tag_list指定要提取文本的标签。
        :param html_dict: 网页返回的全部文本内容response.text和url
        :param target_content_tag: html中正文content所在tag，字典格式限制长度为1，key为选中便签类型name/class/id，vaule为标签取值如div/title/article等
        :param target_tag_list: 指定提取html对应的tag文本，列表，每个元素都与target_content_tag格式相同
        :return: text_content:{} 提取的text文本内容
        '''
        # 格式定义
        assert isinstance(target_content_tag,dict), "target_content_tag请输入字典格式！"
        assert len(target_content_tag.keys()) <= 1,"target_content_tag属性字典只能指定唯一元素！"
        for _ in target_tag_list:
            assert isinstance(_, dict), "target_tag_list列表元素需要字典格式！"
            assert len(_.keys()) <= 1, "target_tag_list列表中的属性字典只能指定唯一元素！"
        # 提取html的内容
        html_content = html_dict['text']
        url = html_dict['url']
        host_url = html_dict['host_url']
        # 创建BeautifulSoup对象
        soup = BeautifulSoup(html_content, 'html.parser')
        # 处理pre引用代码块，添```引用
        pre_tags = soup.find_all('code')
        for pre_tag in pre_tags:
            pre_tag.string = '\n```code\n' + pre_tag.get_text() + '\n```\n'
        # 提取HTML中的文本内容
        doc_tokenizer = DocTokenizer()
        text_dict = {}
        text_dict['url'] = url
        text_dict['host_url'] = host_url
        # 提取网页的title，不存在则置空
        try:
            text_dict['title'] = soup.title.text
        except:
            text_dict['title'] = None
        # 是否提取全部text，不区分标签
        if is_get_all_text:
            all_text = soup.get_text(separator="", strip=False)
            text_dict['all_text'] = doc_tokenizer.doc_process(all_text)
        # 提取正文tag，可以按照标签的class提取，或按照tag名提取
        if target_content_tag:
            text_dict["content"] = self.soup_find_all_text(soup=soup,doc_tokenizer=doc_tokenizer,attrs=target_content_tag)
        # 提取html中tag内容，每个tag独立作为字段保存
        for target_tag in target_tag_list:
            if target_tag:
                # 提取目标tag名
                tag_ = list(target_tag.values())[0]
                # 提取目标tag内容
                text_dict[tag_] = self.soup_find_all_text(soup,doc_tokenizer,attrs=target_tag)
        return text_dict

    def soup_find_all_text(self,soup,doc_tokenizer,attrs):
        assert isinstance(attrs,dict), "attrs请输入字典格式！"
        assert len(attrs.keys()) == 1,"attrs属性字典只能指定唯一元素！"
        if list(attrs.keys())[0]=="name":
            _tags = soup.find_all(name=attrs["name"])
        else:
            _tags = soup.find_all(attrs=attrs)
        tags_text = ""
        for _tag in _tags:
            tag_text = _tag.get_text(separator="", strip=False)
            tag_text = doc_tokenizer.doc_process(tag_text)
            tags_text += tag_text.strip() + "\n\n"
        return tags_text

    def read_html_jsonl(self, file_name=None):
        '''
        读取html的josnl文件
        '''
        html_dict_list = []
        with open(file_name, "r", encoding="utf-8") as f:
            for k, line in enumerate(f):
                line = json.loads(line)
                html_dict_list.append(line)
        return html_dict_list

    def save_text_jsonl(self, json_list=[], file_path=None, mode="w"):
        '''
        将json_list保存成jsonl格式文件
        '''
        with open(file_path, mode, encoding="utf-8") as f:
            for line in json_list:
                f.write(json.dumps(line, ensure_ascii=False) + "\n")

web_crawler/utils/Html2Text.py
codefuse-ai/codefuse-chatbot
from .WebHtmlExtractor import WebHtmlExtractor
import logging
from .Html2Text import Html2Text


class WebCrawler():
    '''爬取url内容，分为requests和selenium两种方式；selenium需提前下载chrome浏览器与chromedriver，并配置路径。
    安装selenium模拟访问网站，需安装并调试chromedriver，版本与电脑chrome需一致，且正确配置路径。mac电脑路径：打开finder,再按command+shift+G进入/usr/local/bin；windows可配置路径。
    '''

    def __init__(self):
        pass

    def webcrawler_single(self,
                          html_dir=None,
                          text_dir=None,
                          base_url=None,
                          reptile_lib="requests",
                          method="get",
                          mode="w",
                          time_sleep=4,
                          time_out=10,
                          target_content_tag={},
                          target_tag_list=[]
                          ):
        '''
        爬取base_url页网址，分别保存html与解析处理的text
        :param html_dir: 保存html地址，jsonl文件
        :param text_dir: 将提取的text内容保存的地址，同样是jsonl格式。
        :param base_url: 目标网址
        :param reptile_lib: requests和selenium两种方式；requests为简单请求静态网址html内容，js动态数据无法获取；selenium为模拟人行为请求，可获取全部html数据，但请求时间较长，尽量设置5s以上的time_sleep，selenium需提前下载chrome浏览器与chromedriver，并配置路径。。
        :param method: requests请求有get/post两种，selenium只支持get
        :param time_sleep: 等待时间s
        :param time_out: 超时时长s
        :param target_content_tag: html中正文content所在tag，字典格式限制长度为1，key为选中便签类型name/class/id，vaule为标签取值如div/title/article等
        :param target_tag_list: 指定提取html对应的tag文本，列表，每个元素都与target_content_tag格式相同
        :return: None
        '''
        assert method=="get","只支持get请求！"
        # 发送请求获取base_url结果：包含相关页面全部网址
        whe = WebHtmlExtractor(time_sleep=time_sleep, time_out=time_out)
        whe.save_url_html(base_url=base_url, reptile_lib=reptile_lib, method=method, html_dir=html_dir, mode=mode)
        # 读取文件
        h2t = Html2Text()
        # 读取并处理，只按照指定tag获取text，不获取全部text内容
        h2t.html2text(target_content_tag=target_content_tag,
                      target_tag_list=target_tag_list,
                      html_dir=html_dir,
                      text_dir=text_dir,
                      mode="w",
                      is_get_all_text=True)

    def webcrawler_batch(self,
                          html_dir=None,
                          text_dir=None,
                          target_url_list=[],
                          reptile_lib="requests",
                          method="get",
                          mode="w",
                          time_sleep=4,
                          time_out=10,
                          target_content_tag={},
                          target_tag_list=[]
                          ):
        '''
        爬取base_url页网址，分别保存html与解析处理的text
        :param html_dir: 保存html地址，jsonl文件
        :param text_dir: 将提取的text内容保存的地址，同样是jsonl格式。
        :param base_url: 目标网址
        :param reptile_lib: requests和selenium两种方式；requests为简单请求静态网址html内容，js动态数据无法获取；selenium为模拟人行为请求，可获取全部html数据，但请求时间较长，尽量设置5s以上的time_sleep，selenium需提前下载chrome浏览器与chromedriver，并配置路径。。
        :param method: requests请求有get/post两种，selenium只支持get
        :param time_sleep: 等待时间s
        :param time_out: 超时时长s
        :param target_content_tag: html中正文content所在tag，字典格式限制长度为1，key为选中便签类型name/class/id，vaule为标签取值如div/title/article等
        :param target_tag_list: 指定提取html对应的tag文本，列表，每个元素都与target_content_tag格式相同
        :return: None
        '''
        assert method=="get","只支持get请求！"
        # 发送请求获取base_url结果：包含相关页面全部网址
        whe = WebHtmlExtractor(time_sleep=time_sleep, time_out=time_out)
        # 循环调用
        try:
            for k,url in enumerate(target_url_list):
                mode_batch = mode if k==0 else "a"
                whe.save_url_html(base_url=url, reptile_lib=reptile_lib, method=method, html_dir=html_dir, mode=mode_batch)
        except:
            logging.warning("爬取停止！")
        # html中提取text信息，并对doc做基础处理
        h2t = Html2Text()
        h2t.html2text(target_content_tag=target_content_tag,
                      target_tag_list=target_tag_list,
                      html_dir=html_dir,
                      text_dir=text_dir,
                      mode="w",
                      is_get_all_text=True)

    def webcrawler_1_degree(self,
                            html_dir=None,
                            text_dir=None,
                            base_url=None,
                            reptile_lib="requests",
                            method="get",
                            mode="w",
                            time_sleep=4,
                            time_out=10,
                            target_content_tag={},
                            target_tag_list=[],
                            target_url_prefix=None
                            ):
        '''
        爬取base_url页面所有<a href=>网址，限制target_url_prefix为前缀，默认target_url_prefix=base_url，分别保存html与解析处理的text。
        :param html_dir: 保存html地址，jsonl文件
        :param text_dir: 将提取的text内容保存的地址，同样是jsonl格式。
        :param base_url: 目标站点
        :param target_url_prefix: 基于base_url网址，1度跳转链接 且 以target_url_prefix开头。默认为target_url_prefix=base_url（请求返回的当前网址url，中文会自动转为编码）。
        :param reptile_lib: requests和selenium两种方式；requests为简单请求静态网址html内容，js动态数据无法获取；selenium为模拟人行为请求，可获取全部html数据，但请求时间较长，尽量设置5s以上的time_sleep。
        :param method: requests请求有get/post两种，selenium只支持get
        :param time_sleep: 等待时间s
        :param time_out: 超时时长s
        :param target_content_tag: html中正文content所在tag，字典格式限制长度为1，key为选中便签类型name/class/id，vaule为标签取值如div/title/article等
        :param target_tag_list: 指定提取html对应的tag文本，列表，每个元素都与target_content_tag格式相同
        :return: None
        '''
        assert method == "get", "只支持get请求！"
        # 发送请求获取base_url结果：包含相关页面全部网址
        whe = WebHtmlExtractor(time_sleep=time_sleep, time_out=time_out)
        try:
            whe.save_1_jump_url_in_base(base_url=base_url, target_url_prefix=target_url_prefix, reptile_lib=reptile_lib,
                                        method=method, html_dir=html_dir, mode=mode)
        except:
            logging.warning("爬取停止！")
        # 读取文件
        h2t = Html2Text()
        # 读取并处理，只按照指定tag获取text，不获取全部text内容
        h2t.html2text(target_content_tag=target_content_tag,
                      target_tag_list=target_tag_list,
                      html_dir=html_dir,
                      text_dir=text_dir,
                      mode="w",
                      is_get_all_text=True)

web_crawler/utils/WebCrawler.py
codefuse-ai/codefuse-chatbot
import requests
from fake_useragent import UserAgent
import time
from bs4 import BeautifulSoup
from urllib.parse import urljoin
import logging
import json
import os
from requests.exceptions import Timeout
from urllib.parse import urlsplit
import re
import math
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.support.ui import WebDriverWait
from selenium.common.exceptions import WebDriverException, TimeoutException
from selenium.webdriver.common.by import By


class WebHtmlExtractor():
    '''爬取web网站html格式数据，分为requests和selenium两种方式；selenium需提前下载chrome浏览器与chromedriver，并配置路径。
    '''

    def __init__(self, header=None, data={}, time_sleep=1, time_out=20, max_retry_times=3):
        # 最大重试次数
        self.max_retry_times = max_retry_times
        # requests请求的header
        if header:
            self.header = header
        else:
            self.header = {
                'User-Agent': UserAgent().random, 'Cookie': None
            }
        # 发送post请求时的data参数
        self.data = data
        # 请求时间间隔s
        self.time_sleep = time_sleep
        # 请求时间时间s
        self.time_out = time_out

    def save_url_html(self, base_url=None, reptile_lib="requests", method="get", time_sleep=None, time_out=None, html_dir=None,mode="w"):
        '''
        对base_url网址发送请求，爬取base_url网址中1度跳转的网址，限制目标网址以target_url_prefix为前缀
        :param base_url: 目标站点
        :param target_url_prefix: 基于base_url网址，1度跳转链接 且 以target_url_prefix开头。默认为target_url_prefix=base_url（请求返回的当前网址url，中文会自动转为编码）。
        :param reptile_lib: requests和selenium两种方式；requests为简单请求静态网址html内容，js动态数据无法获取；selenium为模拟人行为请求，可获取全部html数据，但请求时间较长，尽量设置5s以上的time_sleep。
        :param method: requests请求有get/post两种，selenium只支持get
        :param time_sleep: 等待时间s
        :param time_out: 超时时长s
        :param html_dir: 保存html地址，jsonl文件
        :return: None
        '''
        if time_out is None:
            time_out = self.time_out
        if time_sleep is None:
            time_sleep = self.time_sleep
        # 创建保存目录
        os.makedirs(os.path.dirname(html_dir), exist_ok=True)
        # 发送请求获取base_url结果：包含当前页面全部网址，默认限制在base_url路径下
        html_dict = self.get_html_dict(
            url=base_url, reptile_lib=reptile_lib, method=method, time_sleep=time_sleep, time_out=time_out)
        # 保存文件
        self.save_jsonl(file_path=html_dir, json_list=[
            html_dict], mode=mode)

    def save_1_jump_url_in_base(self, base_url=None, target_url_prefix=None, reptile_lib="requests", method="get", time_sleep=None, time_out=None, html_dir=None,mode = "w"):
        '''
        对base_url网址发送请求，爬取base_url网址中1度跳转的网址，限制目标网址以target_url_prefix为前缀
        :param base_url: 目标站点
        :param target_url_prefix: 基于base_url网址，1度跳转链接 且 以target_url_prefix开头。默认为target_url_prefix=base_url（请求返回的当前网址url，中文会自动转为编码）。
        :param reptile_lib: requests和selenium两种方式；requests为简单请求静态网址html内容，js动态数据无法获取；selenium为模拟人行为请求，可获取全部html数据，但请求时间较长，尽量设置5s以上的time_sleep。
        :param method: requests请求有get/post两种，selenium只支持get
        :param time_sleep: 等待时间s
        :param time_out: 超时时长s
        :param html_dir: 保存html地址，jsonl文件
        :return: None
        '''
        if time_out is None:
            time_out = self.time_out
        if time_sleep is None:
            time_sleep = self.time_sleep
        # 创建保存目录
        os.makedirs(os.path.dirname(html_dir), exist_ok=True)
        # 发送请求获取base_url结果：包含当前页面全部网址，默认限制在base_url路径下
        html_dict = self.get_html_dict(
            url=base_url, reptile_lib=reptile_lib, method=method, time_sleep=time_sleep, time_out=time_out)
        # 如果target_url_prefix为空，用base_url作为前缀
        if target_url_prefix is None:
            target_url_prefix = html_dict['url']
        sub_url_list = self.get_link_sub_url_list(
            html_dict=html_dict, target_url_prefix=target_url_prefix)
        sub_url_nums = len(sub_url_list)
        logging.info("站点根目录{base_url}包含{url_nums}个网址".format(
            base_url=base_url, url_nums=sub_url_nums))
        # 循环请求1度跳转网址列表，并保存数据
        k = 0
        while k < sub_url_nums:
            # 提取第k个网址
            sub_url = sub_url_list[k]
            # 发送请求并返回response
            logging.info(
                "第{k}个网址，进度{rate}%……".format(k=k + 1, rate=round((k + 1) / sub_url_nums * 100, 1), url=sub_url))
            if k == 0:
                mode = mode
            else:
                mode = "a"
            # 爬取url网页html全部内容，保存为字典
            html_dict = self.get_html_dict(
                url=sub_url, reptile_lib=reptile_lib, method=method, time_sleep=time_sleep, time_out=time_out)
            # 保存文件
            self.save_jsonl(file_path=html_dir, json_list=[
                            html_dict], mode=mode)
            k += 1

    def save_2_jump_url_in_base(self, base_url=None, target_url_prefix=None, reptile_lib="requests", method="get", time_sleep=None, time_out=None, html_dir=None):
        '''
        对base_url网址发送请求，爬取base_url网址中2度跳转的网址，限制目标网址以target_url_prefix为前缀
        :param base_url: 当前根网页站点
        :param target_url_prefix: 基于base_url网址，2度跳转目标页面网址前缀，目标网址需要以target_url_prefix开头
        :param reptile_lib: requests和selenium两种方式；requests为简单请求静态网址html内容，js动态数据无法获取；selenium为模拟人行为请求，可获取全部html数据，但请求时间较长，尽量设置5s以上的time_sleep。
        :param method: requests请求有get/post两种，selenium只支持get
        :param time_sleep: 等待时间s
        :param time_out: 超时时长s
        :param html_dir: 保存html地址，jsonl文件
        :return: None
        '''
        if time_out is None:
            time_out = self.time_out
        if time_sleep is None:
            time_sleep = self.time_sleep
        # 创建保存目录
        os.makedirs(os.path.dirname(html_dir), exist_ok=True)
        # 爬取html网页数据
        html_dict = self.get_html_dict(
            url=base_url, reptile_lib=reptile_lib, method=method, time_sleep=time_sleep, time_out=time_out)
        # 提取html中的链接网址
        sub_url_list = self.get_link_sub_url_list(html_dict=html_dict)
        logging.info("站点根目录{base_url}包含{url_nums}个网址".format(
            base_url=base_url, url_nums=len(sub_url_list)))
        n_url = 0
        # 2度跳转网址，爬取html网页并保存
        for k, sub_url in enumerate(sub_url_list):
            sub_url_html_dict = self.get_html_dict(
                url=sub_url, reptile_lib=reptile_lib, method=method, time_sleep=time_sleep, time_out=time_out)
            sub_url_list_2 = self.get_link_sub_url_list(
                html_dict=sub_url_html_dict, target_url_prefix=target_url_prefix)
            logging.info("第{k}个网址，进度{rate}%,目录{sub_url}包含{url_nums}个网址".format(k=k + 1, rate=round(
                (k + 1) / len(sub_url_list) * 100, 1), sub_url=sub_url, url_nums=len(sub_url_list_2)))
            for i, sub2_url in enumerate(sub_url_list_2):
                if k == 0 and i == 0:
                    mode = "w"
                else:
                    mode = "a"
                # 爬取url网页html全部内容，保存为字典
                html_dict = self.get_html_dict(
                    url=sub2_url, reptile_lib=reptile_lib, method=method, time_sleep=time_sleep, time_out=time_out)
                # 保存文件
                self.save_jsonl(file_path=html_dir, json_list=[
                                html_dict], mode=mode)
                n_url += 1
        logging.info("站点根目录{base_url}两度跳跃包含{url_nums}个网址".format(
            base_url=base_url, url_nums=n_url))

    def get_html_dict(self, url=None, reptile_lib="requests", method="get", selenium_headless=True, time_sleep=None, time_out=None):
        '''
        对url网址发送请求，返回html_dict，requests和selenium两种方式
        :param url: 目标网址
        :param reptile_lib: requests和selenium两种方式；requests为简单请求静态网址html内容，js动态数据无法获取；selenium为模拟人行为请求，可获取全部html数据，但请求时间较长，尽量设置5s以上的time_sleep。
        :param method: requests请求有get/post两种，selenium只支持get
        :param selenium_headless: False则会在电脑自动弹出窗口打开页面；True则不弹出窗口
        :param time_sleep: 等待时间s
        :param time_out: 超时时长s
        :return: html_dict:{}
        '''
        # 爬取url网页html全部内容，保存为字典
        assert reptile_lib in (
            "requests", "selenium"), "reptile_lib请选择requests/selenium方法"
        if time_out is None:
            time_out = self.time_out
        if time_sleep is None:
            time_sleep = self.time_sleep
        if reptile_lib == "requests":
            html_dict = self.get_request_html(
                url=url, method=method, time_sleep=time_sleep, time_out=time_out)
        elif reptile_lib == "selenium":
            html_dict = self.get_selenium_html(
                url=url, method=method, time_sleep=time_sleep, time_out=time_out, headless=selenium_headless)
        return html_dict

    def get_request_html(self, url=None, method="get", time_sleep=1, retry_times=1, header=None, data=None, time_out=None):
        '''
        对url网址发送requests请求，保存html字典
        :param url: 目标网址
        :param method: get/post
        :return: html_dict:{}
        '''
        assert method in ("get", "post"), "method请选择get/post方法"
        # 传入header
        if header:
            self.header = header
        if data:
            self.data = data
        if time_out is None:
            time_out = self.time_out
        if time_sleep is None:
            time_sleep = self.time_sleep
        # 每次发送请求前等待1s
        time.sleep(time_sleep)
        # 发送请求，超时10s重试
        try:
            if method == "get":
                response = requests.get(
                    url, headers=self.header, timeout=time_out)
            elif method == "post":
                response = requests.post(
                    url, headers=self.header, data=self.data, timeout=time_out)
        except Timeout:
            response = None
        # 查看返回结果是否为空，最大重试3次，response保存为html字典
        html_dict = {}
        if response is None:
            html_dict['url'] = url
            html_dict['host_url'] = self.split_host_url(url)
            html_dict['text'] = None
            html_dict['status'] = False
            if retry_times <= self.max_retry_times:
                logging.warning("第{retry_times}次重试……".format(
                    retry_times=retry_times))
                html_dict = self.get_request_html(
                    url=url, method=method, time_sleep=time_sleep, retry_times=retry_times+1)
            else:
                logging.warning("重试{max_retry_times}次，达到最大重试次数！".format(
                    max_retry_times=self.max_retry_times))
        else:
            html_dict['url'] = response.url
            html_dict['host_url'] = self.split_host_url(response.url)
            html_dict['text'] = response.text
            html_dict['status'] = True
            logging.info("请求{url}网址返回成功！".format(url=url))
        return html_dict

    def get_selenium_html(self, url=None, method="get", time_sleep=None, retry_times=1, headless=True, time_out=None):
        '''
        selenium模拟访问网站，保存html字典，需安装并调试chromedriver，版本与电脑chrome需一致，且正确配置路径。mac电脑路径：打开finder,再按command+shift+G进入/usr/local/bin；windows可配置路径。
        :param url: 目标网址
        :param method: 只能取值get
        :param headless: False则会在电脑自动弹出窗口打开页面；True则不弹出窗口
        :return: html_dict : {}
        '''
        assert method == "get", "selenium爬取method只支持get方法"
        if time_out is None:
            time_out = self.time_out
        if time_sleep is None:
            time_sleep = self.time_sleep
        # 创建Chrome浏览器实例，并不在电脑实际展示页面
        options = Options()
        options.headless = headless
        driver = webdriver.Chrome(options=options)
        # 再等待5秒
        time.sleep(time_sleep)
        try:
            # 打开目标网址，等待页面全部元素加载完成
            # url = 'https://market.cloud.tencent.com/products/3027'
            driver.get(url)
            wait = WebDriverWait(driver, time_out)
            wait.until(EC.visibility_of_all_elements_located)
            time.sleep(time_sleep)
            # # 获取整个页面的HTML数据
            page_source = driver.page_source
            current_url = driver.current_url
        except (WebDriverException, TimeoutException) as e:
            page_source = None
            current_url = url
            logging.info("Selenium请求错误:%s" % str(e))
        finally:
            # 关闭浏览器驱动
            driver.quit()

        # 保存html字典
        html_dict = {}
        html_dict['url'] = current_url
        html_dict['host_url'] = self.split_host_url(current_url)
        html_dict['text'] = page_source

        # 请求出错，最大重试3次
        if page_source is None:
            html_dict['status'] = False
            if retry_times <= self.max_retry_times:
                logging.warning("第{retry_times}次重试……".format(
                    retry_times=retry_times))
                html_dict = self.get_selenium_html(
                    url=url, time_sleep=time_sleep, retry_times=retry_times+1, headless=headless)
            else:
                logging.warning("重试{max_retry_times}次，达到最大重试次数！".format(
                    max_retry_times=self.max_retry_times))
        else:
            html_dict['status'] = True
            logging.info("请求{url}网址返回成功！".format(url=url))

        return html_dict

    def get_link_sub_url_list(self, html_dict={}, target_url_prefix=None):
        '''
        提取网页html中出现的相关网址列表，返回结果要以url_prefix为前缀
        :param html_dict: 网页返回的全部文本内容response.text和url
        :param url_prefix: 只返回url_prefix开头的网址，默认为html当前页面的url
        :return: sub_url_list: 网址列表
        '''
        # 从response中提取html_content、base_url
        html_content = html_dict['text']
        url = html_dict['url']
        # 创建BeautifulSoup对象
        soup = BeautifulSoup(html_content, 'html.parser')
        # 提取所有<a>标签的href属性值
        links = soup.find_all('a')
        # 存储提取到的网址
        sub_url_list = [url]
        # 如果没传，以url作为前缀
        if target_url_prefix is None:
            target_url_prefix = url
        # 遍历所有链接
        for link in links:
            href = link.get('href')
            if href:
                # 使用urljoin函数将相对网址转换为绝对网址
                host_url = self.split_host_url(url)
                absolute_url = urljoin(host_url, href)
                # 返回absolute_url以本页url_prefix开头
                if absolute_url.startswith(target_url_prefix):
                    sub_url_list.append(absolute_url)
        # 去重且不改变顺序
        return list(dict.fromkeys(sub_url_list))

    def save_jsonl(self, json_list=[], file_path=None, mode="w"):
        '''
        将json_list保存成jsonl格式文件
        '''
        with open(file_path, mode, encoding="utf-8") as f:
            for line in json_list:
                f.write(json.dumps(line, ensure_ascii=False) + "\n")

    def split_host_url(self, url):
        '''
        从url中提取host域名网址
        '''
        parsed_url = urlsplit(url)
        host_url = parsed_url.scheme + '://' + parsed_url.netloc
        return host_url

web_crawler/utils/WebHtmlExtractor.py
codefuse-ai/codefuse-chatbot