82 lines
2.4 KiB
Python
82 lines
2.4 KiB
Python
import asyncio
|
||
from pydantic import BaseModel, Field
|
||
from langchain.prompts.chat import ChatMessagePromptTemplate
|
||
from configs import logger, log_verbose
|
||
from server.utils import get_model_worker_config, fschat_openai_api_address
|
||
from langchain.chat_models import ChatOpenAI
|
||
from typing import Awaitable, List, Tuple, Dict, Union, Callable
|
||
|
||
|
||
def get_ChatOpenAI(
|
||
model_name: str,
|
||
temperature: float,
|
||
callbacks: List[Callable] = [],
|
||
) -> ChatOpenAI:
|
||
config = get_model_worker_config(model_name)
|
||
model = ChatOpenAI(
|
||
streaming=True,
|
||
verbose=True,
|
||
callbacks=callbacks,
|
||
openai_api_key=config.get("api_key", "EMPTY"),
|
||
openai_api_base=fschat_openai_api_address(),
|
||
model_name=model_name,
|
||
temperature=temperature,
|
||
openai_proxy=config.get("openai_proxy")
|
||
)
|
||
return model
|
||
|
||
|
||
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
|
||
msg = f"Caught exception: {e}"
|
||
logger.error(f'{e.__class__.__name__}: {msg}',
|
||
exc_info=e if log_verbose else None)
|
||
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
|
||
|
||
def to_msg_template(self, is_raw=True) -> ChatMessagePromptTemplate:
|
||
role_maps = {
|
||
"ai": "assistant",
|
||
"human": "user",
|
||
}
|
||
role = role_maps.get(self.role, self.role)
|
||
if is_raw: # 当前默认历史消息都是没有input_variable的文本。
|
||
content = "{% raw %}" + self.content + "{% endraw %}"
|
||
else:
|
||
content = self.content
|
||
|
||
return ChatMessagePromptTemplate.from_template(
|
||
content,
|
||
"jinja2",
|
||
role=role,
|
||
)
|
||
|
||
@classmethod
|
||
def from_data(cls, h: Union[List, Tuple, Dict]) -> "History":
|
||
if isinstance(h, (list,tuple)) and len(h) >= 2:
|
||
h = cls(role=h[0], content=h[1])
|
||
elif isinstance(h, dict):
|
||
h = cls(**h)
|
||
|
||
return h
|