from contextlib import contextmanager from codetiming import Timer @contextmanager def _timer(name: str, timing_raw): with Timer(name=name, logger=None) as timer: yield timing_raw[name] = timer.last from buffer import SurveyManager from buffer import BufferManager_V2 as BufferManager from vllm import LLM, SamplingParams from transformers import AutoTokenizer import re from fastapi import FastAPI, WebSocket, WebSocketDisconnect from fastapi.middleware.cors import CORSMiddleware import asyncio import argparse from pydantic import BaseModel import json import aiohttp app = FastAPI() # 允许跨域（如果前端和后端端口不同需要加上） app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) active_connections = set() @app.websocket("/ws") async def websocket_endpoint(websocket: WebSocket): await websocket.accept() active_connections.add(websocket) try: while True: await websocket.receive_text() except WebSocketDisconnect: active_connections.remove(websocket) async def post_to_frontend(payload): print(f"Sending payload to frontend: {payload}") # Log the payload being sent for ws in list(active_connections): try: await ws.send_text(payload) except Exception as e: print(f"Error sending to WebSocket: {e}") active_connections.remove(ws) def write_to_json(data, path): with open(path, 'w', encoding='utf8') as f: f.write(json.dumps(data, ensure_ascii=False, indent=4)) class OriginalvLLMRollout: def __init__(self, model_name_or_path): # init vLLM self.rollout_model = LLM( model=model_name_or_path, tokenizer=model_name_or_path, gpu_memory_utilization=0.95, trust_remote_code=True, ) self.sampling_params = SamplingParams( temperature=0.7, top_p=0.8, repetition_penalty=1.05, top_k=20, max_tokens=2748, ) def generate(self, input_texts): generated_texts = [] completions = self.rollout_model.generate(input_texts, self.sampling_params, use_tqdm=False) for output in completions: generated_text = output.outputs[0].text generated_texts.append(generated_text) return generated_texts def chat(self, input_messages): generated_texts = [] completions = self.rollout_model.chat(input_messages, self.sampling_params, use_tqdm=False) for output in completions: generated_text = output.outputs[0].text generated_texts.append(generated_text) return generated_texts async def rollout_with_env(querys, batch_size, max_turns, model_path, url, deploy_port=None): """ Args: querys: [string] """ ############################### #### splited by batch size #### ############################### n = len(querys) // batch_size batch_querys = [] for i in range(n+1): temp_data = querys[i*batch_size: (i+1)*batch_size] if len(temp_data) > 0: batch_querys.append(temp_data) print("QUERY NUMBER with BATCH: ", [len(x) for x in batch_querys]) ################### #### init vllm #### ################### vllm_manager = OriginalvLLMRollout(model_path) ############################ #### init Format Reward #### ############################ tokenizer = AutoTokenizer.from_pretrained(model_path) total_rollout_data = [] for querys in batch_querys: ########################################### #### acquire env configs and init envs #### ########################################### buffer_manager = BufferManager(querys) while True: # Break at max-turns if buffer_manager.step >= max_turns: break ############################### #### prepare input prompts #### ############################### messagess_todo = buffer_manager.build_prompt_for_generator() # breakpoint() # Break when no tasks if len(messagess_todo) == 0: break ########################## #### generate by vLLM #### ########################## timing_raw = {} with _timer('vllm sampling', timing_raw): # response_texts = vllm_manager.chat(messagess_todo) response_texts = await asyncio.to_thread(vllm_manager.chat, messagess_todo) ################################## #### preprocess the responses #### ################################## # 对response的详细处理可以集成到环境类中，因环境而异, 先对Response进行预处理 extracted_results = [] for response_text in response_texts: result = BufferManager.parse_generator_response(response_text) extracted_results.append(result) ################################################# #### execute in environment and get feedback #### ################################################# payload = { "tool_calls": [x["tool_call"] for x in extracted_results] } if buffer_manager.step <=2: payload["topk"] = 20 with _timer('get env feedback', timing_raw): # env_response_batched = requests.post(url, json=payload).json() async with aiohttp.ClientSession() as session: async with session.post(url, json=payload) as resp: env_response_batched = await resp.json() ################################### #### postprocess the feedbacks #### ################################### with _timer('postprocessing', timing_raw): buffer_manager.update_trajectory(extracted_results, env_response_batched) buffer_manager.step += 1 print(timing_raw) if deploy_port is not None: now_text = json_to_markdown(buffer_manager.batch_rollout_data[-1]) now_search_keywords= buffer_manager.batch_rollout_data[-1]["trajectory"][-1]["search_keywords"] now_update = buffer_manager.batch_rollout_data[-1]["trajectory"][-1]["answer_thought"] next_update = buffer_manager.batch_rollout_data[-1]["trajectory"][-1]["tool_call_thought"] now_query = buffer_manager.batch_rollout_data[-1]["query"] trajs = buffer_manager.batch_rollout_data[-1]["trajectory"] updated_success = buffer_manager.batch_rollout_data[-1]["trajectory"][-1]["update_success"] if updated_success: for traj in reversed(trajs): if len(traj["summarys"]) > 0: break summary_num = len(traj["summarys"]) if summary_num == 0: summary_text = "No summaries yet." else: summary_text = "\n".join(traj["summarys"]) frontend_payload = { "markdown": now_text, "searchKeywords": now_search_keywords, "nowUpdate": now_update, "nextUpdate": next_update, "query": now_query, "papers": summary_text } frontend_payload = json.dumps(frontend_payload, ensure_ascii=False) try: await post_to_frontend(frontend_payload) except Exception as e: print(f"Error posting to frontend: {e}") for item in buffer_manager.batch_rollout_data: item["survey_text"] = SurveyManager.convert_survey_dict_to_str(item["state"]["current_survey"]) total_rollout_data.extend(buffer_manager.batch_rollout_data) ##################################### #### clear all envs and shutdown #### ##################################### del buffer_manager return total_rollout_data def json_to_markdown(json_data): text = SurveyManager.convert_survey_dict_to_str(json_data["state"]["current_survey"]) all_summarys = {} for traj in json_data["trajectory"]: for item in traj["summarys"]: split_text = item.split("\n") bibkey = split_text[0].split(":")[1].strip() title_begin_index = item.find("Title:") + len("Title:") title_end_index = item.find("Abstract:") title = item[title_begin_index:title_end_index].strip() arxivid = bibkey.split("arxivid")[-1].strip() html = f"arxiv.org/abs/{arxivid}" all_summarys[bibkey] = f"[{title}](https://{html})" reg = r"\\cite\{(.+?)\}" placeholder_reg = re.compile(r"^#\d+$") reg_bibkeys = re.findall(reg, text) bibkeys = [] for bibkey in reg_bibkeys: single_bib = bibkey.split(",") for bib in single_bib: if not placeholder_reg.match(bib): bib = bib.strip() if bib and bib != "*" and bib not in bibkeys: bibkeys.append(bib) bibkeys_index = {bibkey: i+1 for i, bibkey in enumerate(bibkeys)} def replace_bibkey(bibkey): bibkey = bibkey.group(1) single_bib = bibkey.split(",") new_bibs = [] for bib in single_bib: if not placeholder_reg.match(bib): bib = bib.strip() if bib and bib != "*": if bib in bibkeys_index: new_bibs.append(f"{bibkeys_index[bib]}") else: print(f"Warning: {bib} not found in bibkeys") if len(new_bibs) > 0: return "[" + ",".join(new_bibs) + "]" else: return "" text = re.sub(reg, replace_bibkey, text) reference_text = "\n\n".join([f"[{i}] {all_summarys[bibkey]}" for bibkey, i in bibkeys_index.items()]) text += "\n## References\n" + reference_text return text async def test_surveyGen(model_path, out_path,querys, url, deploy_port=None): total_rollout_data = await rollout_with_env(querys, 1, 1000, model_path, url, deploy_port) all_md_texts = [] for json_data in total_rollout_data: md_text = json_to_markdown(json_data) all_md_texts.append(md_text) all_md_texts = "\n\n".join(all_md_texts) with open(out_path, 'w', encoding='utf8') as f: f.write(all_md_texts) # with jsonlines.open(out_path, 'w') as writer: # for item in total_rollout_data: # writer.write(item) class QueryRequest(BaseModel): query: str @app.post("/generate_survey") async def generate_survey(request: QueryRequest): global args # Ensure args is accessible # 这里可以根据需要处理查询 model_path = args.model_path out_path = args.output_file query = request.query querys = [query] # 将查询转换为列表 url = args.retriver_url deploy_port = args.port if args.port is not None else None try: await test_surveyGen(model_path, out_path, querys, url, deploy_port) return {"status": "success", "message": "Survey generated successfully."} except Exception as e: print(f"Error generating survey: {e}") return {"status": "error", "message": str(e)} if __name__ == "__main__": parser = argparse.ArgumentParser(description="Run survey generation with vLLM.") parser.add_argument("--model_path", type=str, required=True, help="Path to the model.") parser.add_argument("--query", type=str, required=True, help="Query to generate survey.") parser.add_argument("--output_file", type=str, required=True, help="Path to the output Markdown file.") parser.add_argument("--retriver_url", type=str, default="http://localhost:8400", help="URL of the retriever service.") parser.add_argument("--port", type=str, default=None, help="Deploy port, default is None, which means not deploy.") args = parser.parse_args() if args.port is not None: import uvicorn uvicorn.run(app, host="localhost", port=int(args.port))# log_level="debug") # Run the survey generation else: asyncio.run( test_surveyGen( model_path=args.model_path, out_path=args.output_file, querys=[args.query], url=args.retriver_url ) )