# encoding: utf-8 ''' @author: 温进 @file: code_chat.py @time: 2023/10/24 下午4:04 @desc: ''' from fastapi import Request, Body import os, asyncio from typing import List from fastapi.responses import StreamingResponse 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, CODE_PROMPT_TEMPLATE) from coagent.connector.configs.prompts import CODE_PROMPT_TEMPLATE from coagent.chat.utils import History, wrap_done from coagent.utils import BaseResponse from .base_chat import Chat from coagent.llm_models import getChatModel, getChatModelFromConfig from coagent.llm_models.llm_config import LLMConfig, EmbedConfig from coagent.service.cb_api import search_code, cb_exists_api from loguru import logger import json class CodeChat(Chat): def __init__( self, code_base_name: str = '', code_limit: int = 1, stream: bool = False, request: Request = None, ) -> None: super().__init__(engine_name=code_base_name, stream=stream) self.engine_name = code_base_name self.code_limit = code_limit self.request = request self.history_node_list = [] def check_service_status(self) -> BaseResponse: cb = cb_exists_api(self.engine_name) if not cb: 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, llm_config: LLMConfig, embed_config: EmbedConfig): '''process''' codes_res = search_code(query=query, cb_name=self.engine_name, code_limit=self.code_limit, search_type=self.cb_search_type, history_node_list=self.history_node_list, api_key=llm_config.api_key, api_base_url=llm_config.api_base_url, model_name=llm_config.model_name, temperature=llm_config.temperature, embed_model=embed_config.embed_model, embed_model_path=embed_config.embed_model_path, embed_engine=embed_config.embed_engine, model_device=embed_config.model_device, ) context = codes_res['context'] related_vertices = codes_res['related_vertices'] # update node names # node_names = [node[0] for node in nodes] # self.history_node_list.extend(node_names) # self.history_node_list = list(set(self.history_node_list)) source_nodes = [] for inum, node_name in enumerate(related_vertices[0:5]): source_nodes.append(f'{inum + 1}. 节点名: `{node_name}`') logger.info('history={}'.format(history)) logger.info('message={}'.format([i.to_msg_tuple() for i in history] + [("human", CODE_PROMPT_TEMPLATE)])) chat_prompt = ChatPromptTemplate.from_messages( [i.to_msg_tuple() for i in history] + [("human", CODE_PROMPT_TEMPLATE)] ) logger.info('chat_prompt={}'.format(chat_prompt)) chain = LLMChain(prompt=chat_prompt, llm=model) result = {"answer": "", "codes": source_nodes} return chain, context, result 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"]), code_limit: int = Body(1, examples=['1']), cb_search_type: str = Body('', examples=['1']), stream: bool = Body(False, description="流式输出"), local_doc_url: bool = Body(False, description="知识文件返回本地路径(true)或URL(false)"), request: Request = None, api_key: str = Body(os.environ.get("OPENAI_API_KEY")), api_base_url: str = Body(os.environ.get("API_BASE_URL")), embed_model: str = Body("", ), embed_model_path: str = Body("", ), embed_engine: str = Body("", ), model_name: str = Body("", ), temperature: float = Body(0.5, ), model_device: str = Body("", ), **kargs ): params = locals() params.pop("self") llm_config: LLMConfig = LLMConfig(**params) embed_config: EmbedConfig = EmbedConfig(**params) self.engine_name = engine_name if isinstance(engine_name, str) else engine_name.default self.code_limit = code_limit 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 self.cb_search_type = cb_search_type return self._chat(query, history, llm_config, embed_config, **kargs) def _chat(self, query: str, history: List[History], llm_config: LLMConfig, embed_config: EmbedConfig, **kargs): history = [History(**h) if isinstance(h, dict) else h for h in history] service_status = self.check_service_status() if service_status.code != 200: return service_status def chat_iterator(query: str, history: List[History]): # model = getChatModel() model = getChatModelFromConfig(llm_config) result, content = self.create_task(query, history, model, llm_config, embed_config, **kargs) # logger.info('result={}'.format(result)) # logger.info('content={}'.format(content)) 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 create_task(self, query: str, history: List[History], model, llm_config: LLMConfig, embed_config: EmbedConfig): '''构建 llm 生成任务''' chain, context, result = self._process(query, history, model, llm_config, embed_config) logger.info('chain={}'.format(chain)) 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, llm_config: LLMConfig, embed_config: EmbedConfig, callback: AsyncIteratorCallbackHandler): chain, context, result = self._process(query, history, model, llm_config, embed_config) task = asyncio.create_task(wrap_done( chain.acall({"context": context, "question": query}), callback.done )) return task, result