Coverage for apps/streamchat/streamchat.py: 81%

1""" 

2StreamChat: Generate a video chat response. 

3Starts an HTTP server to accept job submissions and monitor job status. 

4""" 

5 

6import json 

7import re 

8import sys 

9import os 

10import logging 

11import aiofiles 

12import aiofiles.os 

13 

14from typing import override 

15from typing import Dict 

16from typing import Any 

17from typing import List 

18 

19from http import HTTPStatus 

20 

21from quart import request 

22from quart import jsonify 

23 

24from streamchat_job import StreamChatJob 

25 

26# Local relative imports 

27sys.path.append("..") # noqa: E402 

28sys.path.append("../..") # noqa: E402 

29 

30from streamwise_job import StreamWiseJob 

31from streamwise_app import StreamWiseApp 

32from streamwise_app import run_app 

33 

34from quart_utils import QuartReturn 

35 

36 

37# Find latest chatXXX_prompt.jsonl file 

38# matches chat001_prompt.jsonl, chat23_prompt.jsonl, etc. 

39HISTORY_JSONL_PATTERN = re.compile(r"^chat\d+_prompt\.jsonl$") 

40 

41 

42class StreamChatApp(StreamWiseApp): 

43 """Quart app for StreamChat video chat generation.""" 

44 

45 def __init__(self) -> None: 

46 super().__init__("streamchat") 

47 

48 # Register chat route 

49 route = self.app.route 

50 

51 @route("/chat/<job_id>", methods=["POST"]) 

52 async def chat(job_id: str) -> QuartReturn: 

53 logging.info(f"Received chat request for job {job_id}.") 

54 job = self.jobs.get(job_id, None) 

55 if not job: 

56 logging.error(f"Job {job_id} not found.") 

57 return jsonify({ 

58 "status": "error", 

59 "error": f"Job {job_id} not found." 

60 }), HTTPStatus.NOT_FOUND 

61 

62 form = await request.form 

63 files = await request.files 

64 user_message = form.get("message", "").strip() 

65 audio_message = files.get("audio") 

66 if not user_message and not audio_message: 

67 job.logger.error("No text or audio message provided.") 

68 return jsonify({ 

69 "status": "error", 

70 "error": "No user message or audio provided." 

71 }), HTTPStatus.BAD_REQUEST 

72 

73 if audio_message: 

74 self.logger.info("Transcribing audio message.") 

75 audio_path = f"/tmp/{job_id}_audio_input.webm" 

76 audio_message.save(audio_path) # TODO this may not work 

77 user_message = await job.transcribe_audio(audio_path) 

78 self.logger.info(f"Transcribed audio message: {user_message}") 

79 

80 try: 

81 reply = await job.gen_chat(user_message) 

82 response = { 

83 "status": "ok" 

84 } 

85 response.update(reply) 

86 return jsonify(response), HTTPStatus.OK 

87 except Exception as ex: 

88 job.logger.error(f"Error generating chat response: {ex}") 

89 return jsonify({ 

90 "status": "error", 

91 "error": str(ex) 

92 }), HTTPStatus.INTERNAL_SERVER_ERROR 

93 

94 @route("/chat/<job_id>/history", methods=["GET"]) 

95 async def chat_history( 

96 job_id: str 

97 ) -> QuartReturn: 

98 job = self.jobs.get(job_id, None) 

99 if job: 

100 history = await job.get_chat_history() 

101 else: 

102 job_path = f"{self.tmp_dir}/{job_id}" 

103 history = await get_chat_history_from_file(job_path, job_id) 

104 return jsonify({ 

105 "status": "ok", 

106 "history": history 

107 }), HTTPStatus.OK 

108 

109 @override 

110 def create_job( 

111 self, 

112 job_id: str, 

113 job_config: Dict[str, Any] 

114 ) -> StreamWiseJob: 

115 return StreamChatJob( 

116 job_id=job_id, 

117 config=job_config, 

118 service_manager=self.service_manager 

119 ) 

120 

121 

122async def get_chat_history_from_file( 

123 job_path: str, 

124 job_id: str 

125) -> List[Dict[str, str]]: 

126 """ 

127 Load chat history from local JSONL file. 

128 """ 

129 if not await aiofiles.os.path.exists(job_path): 

130 raise FileNotFoundError(f"Job path '{job_path}' does not exist.") 

131 

132 jsonl_files = [ 

133 file 

134 for file in await aiofiles.os.listdir(job_path) 

135 if HISTORY_JSONL_PATTERN.match(file)] 

136 if not jsonl_files: 

137 raise FileNotFoundError(f"History for '{job_id}' does not exist.") 

138 

139 jsonl_files.sort( 

140 key=lambda f: os.path.getmtime(os.path.join(job_path, f)), 

141 reverse=True) 

142 latest_file = os.path.join(job_path, jsonl_files[0]) 

143 history = await parse_chat_history(latest_file) 

144 return history 

145 

146 

147async def parse_chat_history( 

148 file_path: str 

149) -> List[Dict[str, str]]: 

150 """Parse chat history from a JSONL file.""" 

151 history: List[Dict[str, str]] = [] 

152 if not await aiofiles.os.path.exists(file_path): 

153 return history 

154 async with aiofiles.open(file_path, "r", encoding="utf-8") as jsonl_file: 

155 async for line in jsonl_file: 

156 msg = json.loads(line.strip()) 

157 history.append(msg) 

158 return history 

159 

160 

161if __name__ == "__main__": 

162 run_app( 

163 StreamChatApp, 

164 tmp_dir="/tmp/streamchat", 

165 log_files=[ 

166 "streamwise.log", 

167 "streamchat.log" 

168 ], 

169 app_name="StreamChat", 

170 )