+35

Cargo.lock

reviewed

··· 267 267 checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" 268 268 269 269 [[package]] 270 270 + name = "futures-macro" 271 271 + version = "0.3.31" 272 272 + source = "registry+https://github.com/rust-lang/crates.io-index" 273 273 + checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" 274 274 + dependencies = [ 275 275 + "proc-macro2", 276 276 + "quote", 277 277 + "syn", 278 278 + ] 279 279 + 280 280 + [[package]] 270 281 name = "futures-sink" 271 282 version = "0.3.31" 272 283 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 286 297 dependencies = [ 287 298 "futures-core", 288 299 "futures-io", 300 300 + "futures-macro", 289 301 "futures-sink", 290 302 "futures-task", 291 303 "memchr", ··· 751 763 "pkg-config", 752 764 "vcpkg", 753 765 ] 766 766 + 767 767 + [[package]] 768 768 + name = "owo-colors" 769 769 + version = "4.2.3" 770 770 + source = "registry+https://github.com/rust-lang/crates.io-index" 771 771 + checksum = "9c6901729fa79e91a0913333229e9ca5dc725089d1c363b2f4b4760709dc4a52" 754 772 755 773 [[package]] 756 774 name = "percent-encoding" ··· 842 860 "sync_wrapper", 843 861 "tokio", 844 862 "tokio-native-tls", 863 863 + "tokio-util", 845 864 "tower", 846 865 "tower-http", 847 866 "tower-service", 848 867 "url", 849 868 "wasm-bindgen", 850 869 "wasm-bindgen-futures", 870 870 + "wasm-streams", 851 871 "web-sys", 852 872 ] 853 873 ··· 1127 1147 dependencies = [ 1128 1148 "anyhow", 1129 1149 "clap", 1150 1150 + "futures-util", 1130 1151 "nom", 1152 1152 + "owo-colors", 1131 1153 "reqwest", 1132 1154 "serde", 1133 1155 "serde_json", ··· 1409 1431 checksum = "bad67dc8b2a1a6e5448428adec4c3e84c43e561d8c9ee8a9e5aabeb193ec41d1" 1410 1432 dependencies = [ 1411 1433 "unicode-ident", 1434 1434 + ] 1435 1435 + 1436 1436 + [[package]] 1437 1437 + name = "wasm-streams" 1438 1438 + version = "0.4.2" 1439 1439 + source = "registry+https://github.com/rust-lang/crates.io-index" 1440 1440 + checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" 1441 1441 + dependencies = [ 1442 1442 + "futures-util", 1443 1443 + "js-sys", 1444 1444 + "wasm-bindgen", 1445 1445 + "wasm-bindgen-futures", 1446 1446 + "web-sys", 1412 1447 ] 1413 1448 1414 1449 [[package]]

+3 -1

Cargo.toml

reviewed

··· 6 6 [dependencies] 7 7 clap = { version = "4.5.48", features = ["derive"] } 8 8 nom = "8" 9 9 - reqwest = { version = "0.12", features = ["json", "blocking"] } 9 9 + reqwest = { version = "0.12", features = ["json", "blocking", "stream"] } 10 10 serde = { version = "1.0", features = ["derive"] } 11 11 serde_json = "1.0" 12 12 anyhow = "1.0" 13 13 tokio = { version = "1" , features = ["macros", "rt-multi-thread"]} 14 14 + owo-colors = "4" 15 15 + futures-util = "0.3"

+222 -80

server/api.py

reviewed

··· 22 22 23 23 from fastapi import FastAPI, HTTPException 24 24 from .config import SYSTEM_PROMPT 25 25 - 25 25 + import logging 26 26 import json 27 27 import time 28 28 import uuid ··· 40 40 from server.mem_agent.utils import extract_python_code, extract_reply, extract_thoughts, create_memory_if_not_exists, format_results 41 41 from server.mem_agent.engine import execute_sandboxed_code 42 42 # Global model cache and configuration 43 43 + 44 44 + logger = logging.getLogger("app") 43 45 _model_cache: Dict[str, MLXRunner] = {} 44 46 _current_model_path: Optional[str] = None 45 47 _default_max_tokens: Optional[int] = None # Use dynamic model-aware limits by default ··· 67 69 class ChatCompletionRequest(BaseModel): 68 70 model: str 69 71 messages: List[ChatMessage] 72 72 + chat_start: bool 73 73 + python_code: str 70 74 max_tokens: Optional[int] = None 71 75 temperature: Optional[float] = 0.7 72 76 top_p: Optional[float] = 0.9 ··· 135 139 """Get model from cache or load it if not cached.""" 136 140 global _model_cache, _current_model_path 137 141 138 138 - print(model_spec) 139 142 # Use the existing model path resolution from cache_utils 140 143 141 144 try: 142 145 model_path, model_name, commit_hash = get_model_path(model_spec) 143 146 if not model_path.exists(): 147 147 + logger.info(f"Model {model_spec} not found in cache") 144 148 raise HTTPException(status_code=404, detail=f"Model {model_spec} not found in cache") 145 149 except Exception as e: 150 150 + logger.info(f"Model {model_spec} not found in: {str(e)}") 146 151 raise HTTPException(status_code=404, detail=f"Model {model_spec} not found: {str(e)}") 147 152 148 153 # Check if it's an MLX model 149 154 150 155 model_path_str = str(model_path) 151 156 152 152 - print(_current_model_path) 153 153 - print(model_path_str) 154 157 # Check if we need to load a different model 155 158 if _current_model_path != model_path_str: 156 159 # Proactively clean up any previously loaded runner to release memory ··· 168 171 if verbose: 169 172 print(f"Loading model: {model_name}") 170 173 174 174 + logger.info(f"Loading model: {model_name}") 171 175 runner = MLXRunner(model_path_str, verbose=verbose) 172 176 runner.load_model() 173 177 174 178 _model_cache[model_path_str] = runner 175 179 _current_model_path = model_path_str 180 180 + else: 181 181 + logger.info(f"Model {model_name} already in memory") 176 182 177 183 return _model_cache[model_path_str] 178 184 ··· 196 202 async def start_model(request: StartRequest): 197 203 """Load the model and start the agent""" 198 204 global _messages, _runner,_memory_path 199 199 - print(str(request)) 205 205 + 200 206 _messages = [ChatMessage(role="system", content=SYSTEM_PROMPT)] 201 207 _memory_path = request.memory_path 208 208 + 202 209 try: 203 210 _runner = get_or_load_model(request.model) 204 211 return {"message": "Model loaded"} ··· 212 219 try: 213 220 runner = get_or_load_model(request.model) 214 221 215 215 - # if request.stream: 216 216 - # # Streaming response 217 217 - # return StreamingResponse( 218 218 - # generate_chat_stream(runner, request.messages, request), 219 219 - # media_type="text/plain", 220 220 - # headers={"Cache-Control": "no-cache"} 221 221 - # ) 222 222 - # else: 223 223 - # Non-streaming response 224 224 - completion_id = f"chatcmpl-{uuid.uuid4()}" 225 225 - created = int(time.time()) 226 226 - 227 227 - # Convert messages to dict format for runner 228 228 - # _messages.append(system_message) 229 229 - _messages.extend(request.messages) 230 230 - message_dicts = format_chat_messages_for_runner(_messages) 231 231 - # Let the runner format with chat templates 232 232 - prompt = runner._format_conversation(message_dicts, use_chat_template=True) 233 233 - 234 234 - generated_text = runner.generate_batch( 235 235 - prompt=prompt, 236 236 - max_tokens=runner.get_effective_max_tokens(request.max_tokens or _default_max_tokens, interactive=False), 237 237 - temperature=request.temperature, 238 238 - top_p=request.top_p, 239 239 - repetition_penalty=request.repetition_penalty, 240 240 - use_chat_template=False # Already applied in _format_conversation 241 241 - ) 242 242 - 243 243 - # Token counting 244 244 - # total_prompt = "\n\n".join([msg.content for msg in request.messages]) 245 245 - # prompt_tokens = count_tokens(total_prompt) 246 246 - # completion_tokens = count_tokens(generated_text) 222 222 + if request.stream: 223 223 + result = ({}, "") 224 224 + if request.python_code: 225 225 + create_memory_if_not_exists() 226 226 + result = execute_sandboxed_code( 227 227 + code=request.python_code, 228 228 + allowed_path=_memory_path, 229 229 + import_module="server.mem_agent.tools", 230 230 + ) 247 231 248 248 - thoughts = extract_thoughts(generated_text) 249 249 - reply = extract_reply(generated_text) 250 250 - python_code = extract_python_code(generated_text) 251 251 - print(generated_text) 252 252 - result = ({}, "") 253 253 - if python_code: 254 254 - create_memory_if_not_exists() 255 255 - result = execute_sandboxed_code( 256 256 - code=python_code, 257 257 - allowed_path=_memory_path, 258 258 - import_module="server.mem_agent.tools", 232 232 + _messages.append(ChatMessage(role="user", content=format_results(result[0], result[1]))) 233 233 + 234 234 + # Streaming response 235 235 + return StreamingResponse( 236 236 + generate_chat_stream(runner, request.messages, request), 237 237 + media_type="text/plain", 238 238 + headers={"Cache-Control": "no-cache"} 259 239 ) 240 240 + else: 241 241 + # Non-streaming response 242 242 + completion_id = f"chatcmpl-{uuid.uuid4()}" 243 243 + created = int(time.time()) 260 244 261 261 - print(reply) 262 262 - print(str(result)) 263 263 - 264 264 - remaining_tool_turns = _max_tool_turns 265 265 - while remaining_tool_turns > 0 and not reply: 266 266 - _messages.append(ChatMessage(role="user", content=format_results(result[0], result[1]))) 245 245 + # Convert messages to dict format for runner 246 246 + # _messages.append(system_message) 247 247 + if request.chat_start: 248 248 + _messages.extend(request.messages) 267 249 message_dicts = format_chat_messages_for_runner(_messages) 268 250 # Let the runner format with chat templates 269 251 prompt = runner._format_conversation(message_dicts, use_chat_template=True) 252 252 + 270 253 generated_text = runner.generate_batch( 271 271 - prompt=prompt 254 254 + prompt=prompt, 255 255 + max_tokens=runner.get_effective_max_tokens(request.max_tokens or _default_max_tokens, interactive=False), 256 256 + temperature=request.temperature, 257 257 + top_p=request.top_p, 258 258 + repetition_penalty=request.repetition_penalty, 259 259 + use_chat_template=False # Already applied in _format_conversation 272 260 ) 273 273 - print(generated_text) 274 274 - # Extract the thoughts, reply and python code from the response 261 261 + 262 262 + # Token counting 263 263 + total_prompt = "\n\n".join([msg.content for msg in request.messages]) 264 264 + prompt_tokens = count_tokens(total_prompt) 265 265 + completion_tokens = count_tokens(generated_text) 266 266 + 267 267 + logger.info(f"prompt_token\n{prompt_tokens}") 268 268 + logger.info(f"completion_tokens\n{completion_tokens}") 269 269 + 275 270 thoughts = extract_thoughts(generated_text) 276 271 reply = extract_reply(generated_text) 277 272 python_code = extract_python_code(generated_text) 278 273 279 279 - _messages.append(ChatMessage(role="assistant", content=generated_text)) 274 274 + result = ({}, "") 280 275 if python_code: 281 276 create_memory_if_not_exists() 282 277 result = execute_sandboxed_code( ··· 284 279 allowed_path=_memory_path, 285 280 import_module="server.mem_agent.tools", 286 281 ) 287 287 - else: 288 288 - # Reset result when no Python code is executed 289 289 - result = ({}, "") 290 290 - remaining_tool_turns -= 1 282 282 + 283 283 + logger.info(f"Model thoughts\n{thoughts}") 284 284 + logger.info(f"Model reply\n{reply}") 285 285 + logger.info(f"Model python\n{python_code}") 286 286 + logger.info(f"executed python result\n{str(result)}") 287 287 + 288 288 + # while remaining_tool_turns > 0 and not reply: 289 289 + # logger.info(f"Turn count\n{remaining_tool_turns}") 290 290 + _messages.append(ChatMessage(role="user", content=format_results(result[0], result[1]))) 291 291 + message_dicts = format_chat_messages_for_runner(_messages) 292 292 + # # Let the runner format with chat templates 293 293 + # prompt = runner._format_conversation(message_dicts, use_chat_template=True) 294 294 + # generated_text = runner.generate_batch( 295 295 + # prompt=prompt 296 296 + # ) 297 297 + 298 298 + # total_prompt = "\n\n".join([msg.content for msg in _messages]) 299 299 + # prompt_tokens = count_tokens(total_prompt) 300 300 + # completion_tokens = count_tokens(generated_text) 301 301 + 302 302 + # logger.info(f"prompt_token\n{prompt_tokens}") 303 303 + # logger.info(f"completion_tokens\n{completion_tokens}") 304 304 + 305 305 + # # print(generated_text) 306 306 + # # Extract the thoughts, reply and python code from the response 307 307 + # thoughts = extract_thoughts(generated_text) 308 308 + # reply = extract_reply(generated_text) 309 309 + # python_code = extract_python_code(generated_text) 310 310 + 311 311 + # logger.info(f"Model thoughts\n{thoughts}") 312 312 + # logger.info(f"Model reply\n{reply}") 313 313 + # logger.info(f"Model python\n{python_code}") 314 314 + 315 315 + # _messages.append(ChatMessage(role="assistant", content=generated_text)) 316 316 + # if python_code: 317 317 + # create_memory_if_not_exists() 318 318 + # result = execute_sandboxed_code( 319 319 + # code=python_code, 320 320 + # allowed_path=_memory_path, 321 321 + # import_module="server.mem_agent.tools", 322 322 + # ) 323 323 + # logger.info(f"executed python result\n{str(result)}") 324 324 + # else: 325 325 + # # Reset result when no Python code is executed 326 326 + # result = ({}, "") 327 327 + # logger.info(f"executed python result\n{str(result)}") 328 328 + # remaining_tool_turns -= 1 291 329 292 292 - return ChatCompletionResponse( 293 293 - id=completion_id, 294 294 - created=created, 295 295 - model=request.model, 296 296 - choices=[ 330 330 + return ChatCompletionResponse( 331 331 + id=completion_id, 332 332 + created=created, 333 333 + model=request.model, 334 334 + choices=[ 335 335 + { 336 336 + "index": 0, 337 337 + "message": { 338 338 + "role": "assistant", 339 339 + "content": generated_text 340 340 + }, 341 341 + "finish_reason": "stop" 342 342 + } 343 343 + ], 344 344 + # usage={ 345 345 + # "prompt_tokens": prompt_tokens, 346 346 + # "completion_tokens": completion_tokens, 347 347 + # "total_tokens": prompt_tokens + completion_tokens 348 348 + # } 349 349 + ) 350 350 + except Exception as e: 351 351 + raise HTTPException(status_code=500, detail=str(e)) 352 352 + 353 353 + async def generate_chat_stream( 354 354 + runner: MLXRunner, 355 355 + messages: List[ChatMessage], 356 356 + request: ChatCompletionRequest 357 357 + ) -> AsyncGenerator[str, None]: 358 358 + """Generate streaming chat completion response.""" 359 359 + 360 360 + global _messages 361 361 + completion_id = f"chatcmpl-{uuid.uuid4()}" 362 362 + created = int(time.time()) 363 363 + 364 364 + if request.chat_start: 365 365 + _messages.extend(request.messages) 366 366 + # Convert messages to dict format for runner 367 367 + message_dicts = format_chat_messages_for_runner(_messages) 368 368 + 369 369 + # Let the runner format with chat templates 370 370 + prompt = runner._format_conversation(message_dicts, use_chat_template=True) 371 371 + 372 372 + # Yield initial response 373 373 + initial_response = { 374 374 + "id": completion_id, 375 375 + "object": "chat.completion.chunk", 376 376 + "created": created, 377 377 + "model": request.model, 378 378 + "choices": [ 379 379 + { 380 380 + "index": 0, 381 381 + "delta": {"role": "assistant", "content": ""}, 382 382 + "finish_reason": None 383 383 + } 384 384 + ] 385 385 + } 386 386 + 387 387 + yield f"data: {json.dumps(initial_response)}\n\n" 388 388 + 389 389 + # Stream tokens 390 390 + try: 391 391 + for token in runner.generate_streaming( 392 392 + prompt=prompt, 393 393 + max_tokens=runner.get_effective_max_tokens(request.max_tokens or _default_max_tokens, interactive=False), 394 394 + temperature=request.temperature, 395 395 + top_p=request.top_p, 396 396 + repetition_penalty=request.repetition_penalty, 397 397 + use_chat_template=False, # Already applied in _format_conversation 398 398 + use_chat_stop_tokens=False # Server mode shouldn't stop on chat markers 399 399 + ): 400 400 + chunk_response = { 401 401 + "id": completion_id, 402 402 + "object": "chat.completion.chunk", 403 403 + "created": created, 404 404 + "model": request.model, 405 405 + "choices": [ 406 406 + { 407 407 + "index": 0, 408 408 + "delta": {"content": token}, 409 409 + "finish_reason": None 410 410 + } 411 411 + ] 412 412 + } 413 413 + 414 414 + yield f"data: {json.dumps(chunk_response)}\n\n" 415 415 + 416 416 + # Check for stop sequences 417 417 + if request.stop: 418 418 + stop_sequences = request.stop if isinstance(request.stop, list) else [request.stop] 419 419 + if any(stop in token for stop in stop_sequences): 420 420 + break 421 421 + 422 422 + except Exception as e: 423 423 + error_response = { 424 424 + "id": completion_id, 425 425 + "object": "chat.completion.chunk", 426 426 + "created": created, 427 427 + "model": request.model, 428 428 + "choices": [ 297 429 { 298 430 "index": 0, 299 299 - "message": { 300 300 - "role": "assistant", 301 301 - "content": reply 302 302 - }, 303 303 - "finish_reason": "stop" 431 431 + "delta": {}, 432 432 + "finish_reason": "error" 304 433 } 305 434 ], 306 306 - # usage={ 307 307 - # "prompt_tokens": prompt_tokens, 308 308 - # "completion_tokens": completion_tokens, 309 309 - # "total_tokens": prompt_tokens + completion_tokens 310 310 - # } 311 311 - ) 312 312 - except Exception as e: 313 313 - raise HTTPException(status_code=500, detail=str(e)) 435 435 + "error": str(e) 436 436 + } 437 437 + yield f"data: {json.dumps(error_response)}\n\n" 438 438 + 439 439 + # Final response 440 440 + final_response = { 441 441 + "id": completion_id, 442 442 + "object": "chat.completion.chunk", 443 443 + "created": created, 444 444 + "model": request.model, 445 445 + "choices": [ 446 446 + { 447 447 + "index": 0, 448 448 + "delta": {}, 449 449 + "finish_reason": "stop" 450 450 + } 451 451 + ] 452 452 + } 453 453 + 454 454 + yield f"data: {json.dumps(final_response)}\n\n" 455 455 + yield "data: [DONE]\n\n"

-1

server/cache_utils.py

reviewed

··· 182 182 def get_model_path(model_spec): 183 183 model_name, commit_hash = parse_model_spec(model_spec) 184 184 base_cache_dir = MODEL_CACHE / hf_to_cache_dir(model_name) 185 185 - print(base_cache_dir) 186 185 if not base_cache_dir.exists(): 187 186 return None, model_name, commit_hash 188 187 if commit_hash:

+31 -10

server/main.py

reviewed

··· 1 1 - # import os 2 1 import uvicorn 3 2 from .api import app 4 3 from .config import PORT 4 4 + import logging 5 5 + import sys 6 6 + from fastapi import Request 7 7 + 8 8 + # --- logging setup --- 9 9 + logging.basicConfig( 10 10 + level=logging.INFO, 11 11 + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", 12 12 + handlers=[logging.StreamHandler(sys.stdout)], 13 13 + ) 14 14 + logger = logging.getLogger("app") 15 15 + 16 16 + 17 17 + # --- middleware for request logging --- 18 18 + @app.middleware("http") 19 19 + async def log_requests(request: Request, call_next): 20 20 + try: 21 21 + body = await request.json() 22 22 + except Exception: 23 23 + body = None 5 24 6 6 - def run(): 7 7 - # Write PID file 8 8 - # PID_FILE.write_text(str(os.getpid())) 25 25 + logger.info({ 26 26 + "method": request.method, 27 27 + "url": str(request.url), 28 28 + "client": request.client.host, 29 29 + "body": body, 30 30 + }) 9 31 10 10 - # try: 32 32 + response = await call_next(request) 33 33 + logger.info(f"<-- {request.method} {request.url.path} {response.status_code}") 34 34 + return response 35 35 + 36 36 + def run(): 11 37 uvicorn.run(app, host="127.0.0.1", port=PORT) 12 12 - # finally: 13 13 - # if PID_FILE.exists(): 14 14 - # PID_FILE.unlink() 15 15 - # 16 16 - # print("hello from main!") 17 38 18 39 if __name__ == "__main__": 19 40 run()

+2 -2

server/mem_agent/engine.py

reviewed

··· 300 300 if result.returncode != 0: 301 301 return None, result.stderr.decode().strip() 302 302 303 303 - print("stderr:", result.stderr.decode()) 304 304 - print("stdout:", result.stdout[:200]) 303 303 + # print("stderr:", result.stderr.decode()) 304 304 + # print("stdout:", result.stdout[:200]) 305 305 306 306 try: 307 307 local_vars, error_msg = pickle.loads(result.stdout)

+2

server/mem_agent/utils.py

reviewed

··· 188 188 """ 189 189 if "<think>" in response and "</think>" in response: 190 190 return response.split("<think>")[1].split("</think>")[0] 191 191 + elif "</think>" in response: 192 192 + return response.split("</think>")[0] 191 193 else: 192 194 return "" 193 195

+1

server/system_prompt.txt

reviewed

··· 40 40 ``` 41 41 42 42 **CRITICAL: Always close ALL tags! Missing </think>, </python>, or </reply> will cause errors!** 43 43 + **CRITICAL: Think block MUST have opening <think> and closing </think>. This block should not only have </think> 43 44 44 45 **NEVER:** 45 46 - Skip the `<think>` block

+134 -15

src/runner/mlx.rs

reviewed

··· 1 1 + use crate::core::modelfile::Modelfile; 1 2 use anyhow::{Context, Result}; 3 3 + use futures_util::StreamExt; 4 4 + use owo_colors::OwoColorize; 2 5 use reqwest::Client; 3 6 use serde_json::{Value, json}; 4 7 use std::io::Write; 5 8 use std::path::PathBuf; 6 9 use std::process::Stdio; 10 10 + use std::str::FromStr; 7 11 use std::{env, fs}; 8 12 use std::{io, process::Command}; 9 9 - 10 10 - use crate::core::modelfile::Modelfile; 13 13 + pub struct ChatResponse { 14 14 + think: String, 15 15 + reply: String, 16 16 + code: String, 17 17 + } 11 18 12 19 pub async fn run(modelfile: Modelfile) { 13 20 let model = modelfile.from.as_ref().unwrap(); ··· 129 136 let modelname = modelfile.from.as_ref().unwrap(); 130 137 load_model(modelname, &memory_path).await.unwrap(); 131 138 println!("Running in interactive mode"); 139 139 + // TODO: Handle "enter" key press or any key press when repl is processing an input 132 140 loop { 133 141 print!(">> "); 134 142 stdout.flush().unwrap(); ··· 141 149 break; 142 150 } 143 151 _ => { 144 144 - if let Ok(response) = chat(input, modelname).await { 145 145 - println!(">> {}", response) 146 146 - } else { 147 147 - println!(">> failed to respond") 152 152 + let mut remaining_count = 6; 153 153 + let mut g_reply: String = "".to_owned(); 154 154 + let mut python_code: String = "".to_owned(); 155 155 + loop { 156 156 + if remaining_count > 0 { 157 157 + let chat_start = if remaining_count == 6 { true } else { false }; 158 158 + if let Ok(response) = chat(input, modelname, chat_start, &python_code).await 159 159 + { 160 160 + if response.reply.is_empty() { 161 161 + if !response.code.is_empty() { 162 162 + python_code = response.code; 163 163 + } 164 164 + remaining_count = remaining_count - 1; 165 165 + } else { 166 166 + g_reply = response.reply.clone(); 167 167 + println!("\n>> {}", response.reply.trim()); 168 168 + break; 169 169 + } 170 170 + } else { 171 171 + println!("\n>> failed to respond"); 172 172 + break; 173 173 + } 174 174 + } 175 175 + } 176 176 + if g_reply.is_empty() { 177 177 + println!(">> No reply") 148 178 } 149 179 } 150 180 } ··· 178 208 } 179 209 } 180 210 181 181 - async fn chat(input: &str, model_name: &str) -> Result<String, String> { 211 211 + async fn chat( 212 212 + input: &str, 213 213 + model_name: &str, 214 214 + chat_start: bool, 215 215 + python_code: &str, 216 216 + ) -> Result<ChatResponse, String> { 182 217 let client = Client::new(); 218 218 + 183 219 let body = json!({ 184 220 "model": model_name, 221 221 + "chat_start": chat_start, 222 222 + "stream": true, 223 223 + "python_code": python_code, 185 224 "messages": [{"role": "user", "content": input}] 186 225 }); 187 226 let res = client ··· 190 229 .send() 191 230 .await 192 231 .unwrap(); 232 232 + 233 233 + let mut stream = res.bytes_stream(); 234 234 + let mut accumulated = String::new(); 235 235 + let mut chat_response = ChatResponse { 236 236 + think: String::new(), 237 237 + reply: String::new(), 238 238 + code: String::new(), 239 239 + }; 240 240 + // let mut inside_python = false; 241 241 + // let mut tag_buffer = String::new(); 242 242 + print!("\n"); 243 243 + while let Some(chunk) = stream.next().await { 244 244 + let chunk = chunk.unwrap(); 245 245 + let s = String::from_utf8_lossy(&chunk); 246 246 + for line in s.lines() { 247 247 + if !line.starts_with("data: ") { 248 248 + continue; 249 249 + } 250 250 + 251 251 + let data = line.trim_start_matches("data: "); 252 252 + 253 253 + if data == "[DONE]" { 254 254 + chat_response = convert_to_chat_response(&accumulated); 255 255 + return Ok(chat_response); 256 256 + } 257 257 + // Parse JSON 258 258 + let v: Value = serde_json::from_str(data).unwrap(); 259 259 + if let Some(delta) = v["choices"][0]["delta"]["content"].as_str() { 260 260 + accumulated.push_str(delta); 261 261 + print!("{}", delta.dimmed()); 262 262 + use std::io::Write; 263 263 + std::io::stdout().flush().ok(); 264 264 + } 265 265 + } 266 266 + } 193 267 // println!("{:?}", res); 194 194 - if res.status() == 200 { 195 195 - let text = res.text().await.unwrap(); 196 196 - let v: Value = serde_json::from_str(&text).unwrap(); 197 197 - let content = v["choices"][0]["message"]["content"] 198 198 - .as_str() 199 199 - .unwrap_or("<no content>"); 200 200 - Ok(content.to_owned()) 268 268 + // if res.status() == 200 { 269 269 + // let text = res.text().await.unwrap(); 270 270 + // let v: Value = serde_json::from_str(&text).unwrap(); 271 271 + // let content = v["choices"][0]["message"]["content"] 272 272 + // .as_str() 273 273 + // .unwrap_or("<no content>"); 274 274 + 275 275 + // // Ok(convert_to_chat_response(content)) 276 276 + // } else { 277 277 + // // Err(String::from("request failed")) 278 278 + // } 279 279 + // unimplemented!() 280 280 + Err(String::from("request failed")) 281 281 + } 282 282 + 283 283 + fn convert_to_chat_response(content: &str) -> ChatResponse { 284 284 + // content.split() 285 285 + ChatResponse { 286 286 + think: extract_think(content), 287 287 + reply: extract_reply(content), 288 288 + code: extract_python(content), 289 289 + } 290 290 + } 291 291 + 292 292 + fn extract_reply(content: &str) -> String { 293 293 + if content.contains("<reply>") && content.contains("</reply>") { 294 294 + let list_a = content.split("<reply>").collect::<Vec<&str>>(); 295 295 + let list_b = list_a[1].split("</reply>").collect::<Vec<&str>>(); 296 296 + list_b[0].to_owned() 201 297 } else { 202 202 - Err(String::from("request failed")) 298 298 + "".to_owned() 299 299 + } 300 300 + } 301 301 + 302 302 + fn extract_python(content: &str) -> String { 303 303 + if content.contains("<python>") && content.contains("</python>") { 304 304 + let list_a = content.split("<python>").collect::<Vec<&str>>(); 305 305 + let list_b = list_a[1].split("</python>").collect::<Vec<&str>>(); 306 306 + list_b[0].to_owned() 307 307 + } else { 308 308 + "".to_owned() 309 309 + } 310 310 + } 311 311 + 312 312 + fn extract_think(content: &str) -> String { 313 313 + if content.contains("<think>") && content.contains("</think>") { 314 314 + let list_a = content.split("<think>").collect::<Vec<&str>>(); 315 315 + let list_b = list_a[1].split("</think>").collect::<Vec<&str>>(); 316 316 + list_b[0].to_owned() 317 317 + } else if content.contains("</think") { 318 318 + let list_a = content.split("</think>").collect::<Vec<&str>>(); 319 319 + list_a[0].to_owned() 320 320 + } else { 321 321 + "".to_owned() 203 322 } 204 323 } 205 324