讯飞星火大模型接入库 (opensource)

"""websocket-client bassd Socket Mode client """ import copy import enum import logging from concurrent.futures.thread import ThreadPoolExecutor from logging import Logger from queue import Queue from threading import Lock from typing import Union, Optional, List, Callable, Tuple import websocket from websocket import WebSocketApp, WebSocketException from sparkai.socket_mode.client import BaseSocketModeClient from sparkai.socket_mode.interval_runner import IntervalRunner from sparkai.socket_mode.listeners import ( WebSocketMessageListener, SocketModeRequestListener, ) from sparkai.socket_mode.request import SocketModeRequest from sparkai.memory.chat_memory import BaseChatMessageHistory from sparkai.schema import ChatMessage from sparkai.models.chat import ChatBody, ChatResponse import threading import json # from pynput import keyboard class SparkMessageStatus(enum.IntEnum): DataBegin = 0 # 首数据 DataContinue = 1 # 中间数据 DataEnd = 2 # 尾数据 DataOnce = 3 # 非会话单次输入输出 class ResponseMessage(object): def __init__(self): self.init() pass def init(self): self.content = '' self.role = '' self.question_tokens = 0 self.completion_tokens = 0 self.total_tokens = 0 self.prompt_tokens = 0 self.status = SparkMessageStatus.DataBegin def set_content(self, role, content): self.content: str = content self.role: str = role def set_status(self, status: SparkMessageStatus): self.status = status def set_usage(self, question_tokens, prompt_tokens, completion_tokens, total_tokens): self.question_tokens: int = question_tokens self.prompt_tokens: int = prompt_tokens self.completion_tokens: int = completion_tokens self.total_tokens: int = total_tokens class InteractiveInputer(threading.Thread): def __init__(self, ws, lock, app_id, memory: BaseChatMessageHistory = None, max_token=2048): super(InteractiveInputer, self).__init__() self.is_stopping = False self.ws = ws self.should_input = True self.lock = lock self.max_token = max_token self.app_id = app_id # self.listener = keyboard.Listener(on_press=self.on_press_ctrl_enter) # conversation Memory self.memory = memory # The currently active modifiers self.current = set() self.one_line_complete = False # self.hot_key_trigger = {keyboard.Key.ctrl, keyboard.Key.enter} def stop(self): self.is_stopping = True def clear_key_session(self): try: for key in self.current: self.current.remove(key) except KeyError: pass finally: self.one_line_complete = False def set_ws(self, ws): self.ws = ws def on_press_ctrl_enter(self, key): if key in self.hot_key_trigger: self.current.add(key) if all(k in self.current for k in self.hot_key_trigger): self.one_line_complete = True def run(self) -> None: # self.listener.start() self.get_input() def release_lock(self): self.lock.release() def get_input(self): while not self.is_stopping: if self.should_input: self.lock.acquire() lines = [] words = input("Question: ") # while not self.one_line_complete: while words != 'EOF': lines.append(words) words = input() # 发送数据 rdata = ChatBody(self.app_id, "\n".join(lines), memory=self.memory, max_tokens=self.max_token).json() self.ws.send(rdata) else: self.clear_key_session() class SparkAISocketModeClient(BaseSocketModeClient): logger: Logger app_id: str api_key: str api_secret: str wss_uri: Optional[str] message_queue: Queue message_listeners: List[ Union[ WebSocketMessageListener, Callable[["BaseSocketModeClient", dict, Optional[str]], None], ] ] socket_mode_request_listeners: List[ Union[ SocketModeRequestListener, Callable[["BaseSocketModeClient", SocketModeRequest], None], ] ] current_app_monitor: IntervalRunner current_app_monitor_started: bool message_processor: IntervalRunner message_workers: ThreadPoolExecutor current_session: Optional[WebSocketApp] current_session_runner: IntervalRunner auto_reconnect_enabled: bool default_auto_reconnect_enabled: bool close: bool connect_operation_lock: Lock on_open_listeners: List[Callable[[WebSocketApp], None]] on_message_listeners: List[Callable[[WebSocketApp, str], None]] on_error_listeners: List[Callable[[WebSocketApp, Exception], None]] on_close_listeners: List[Callable[[WebSocketApp], None]] def __init__( self, app_id: str, api_key: str, api_secret: str, conversation_memory: Optional[BaseChatMessageHistory] = None, logger: Optional[Logger] = None, auto_reconnect_enabled: bool = True, auto_reconnect_5m: bool = False, chat_interactive: bool = False, max_token: int = 2048, ping_interval: float = 10, concurrency: int = 10, trace_enabled: bool = False, http_proxy_host: Optional[str] = None, http_proxy_port: Optional[int] = None, http_proxy_auth: Optional[Tuple[str, str]] = None, proxy_type: Optional[str] = None, on_open_listeners: Optional[List[Callable[[WebSocketApp], None]]] = None, on_message_listeners: Optional[List[Callable[[WebSocketApp, str], None]]] = None, on_error_listeners: Optional[List[Callable[[WebSocketApp, Exception], None]]] = None, on_close_listeners: Optional[List[Callable[[WebSocketApp], None]]] = None, ): """ Args: app_id: iflytek xfyun appid api_key: Spark API api key api_secret: Spark API api secret logger: Custom logger chat_interactive: enable chat interactive auto_reconnect_enabled: True if automatic reconnection is enabled (default: True) ping_interval: interval for ping-pong with Slack servers (seconds) concurrency: the size of thread pool (default: 10) http_proxy_host: the HTTP proxy host http_proxy_port: the HTTP proxy port http_proxy_auth: the HTTP proxy username & password proxy_type: the HTTP proxy type on_open_listeners: listener functions for on_open on_message_listeners: listener functions for on_message on_error_listeners: listener functions for on_error on_close_listeners: listener functions for on_close """ self.app_id = app_id self.api_key = api_key self.api_secret = api_secret self.max_token = max_token self.logger = logger or logging.getLogger(__name__) self.default_auto_reconnect_enabled = auto_reconnect_enabled self.auto_reconnect_enabled = self.default_auto_reconnect_enabled self.chat_interactive = chat_interactive self.ping_interval = ping_interval self.wss_uri = None self.message_queue = Queue() self.message_listeners = [] self.socket_mode_request_listeners = [] self.current_session = None self.current_session_runner = IntervalRunner(self._run_current_session, 0.5).start() self.current_app_monitor_started = False self.current_app_monitor = IntervalRunner(self._monitor_current_session, self.ping_interval) self.cl