from fastapi import FastAPI, Form, Request,Response from fastapi.responses import HTMLResponse from jinja2 import Template import markdown import time from datetime import datetime, timedelta from apscheduler.schedulers.background import BackgroundScheduler from agents import DeepResearchAgent, get_llms import threading import logging from queue import Queue import uuid lock = threading.Lock() app = FastAPI() # 每日最大回复次数 MAX_REPLIES_PER_DAY = 100 # 当日回复次数计数器 reply_count = 0 # 启动时设置计数器重置 last_reset_time = datetime.now() # HTML模板 html_template = """ CoI Agent online demo 😊

CoI Agent online demo 😊

Time Taken: {{ time_taken }} seconds

Today's Replies: {{ reply_count }}/100

Idea

Example Input:

{% if error %}

Error

{% endif %}

""" # 重置每日计数器 def reset_counter(): global reply_count reply_count = 0 # 设置定时任务每天0点重置计数器 scheduler = BackgroundScheduler() scheduler.add_job(reset_counter, 'cron', hour=0, minute=0) scheduler.start() queue = Queue() def fix_markdown(text): lines = text.split('\n') # Initialize the result list result = [] # Iterate through the lines for i, line in enumerate(lines): # Check if the current line starts with a numbered list item numbers = ['1.', '2.', '3.', '4.', '5.', '6.', '7.', '8.', '9.', '10.', '11.', '12.', '13.', '14.', '15.', '16.', '17.', '18.', '19.', '20.','21.','22.','23.','24.','25.','26.','27.','28.','29.','30.'] if line.lstrip().startswith(tuple(numbers)): # If it's not the first line and the previous line is not blank, add a blank line if i > 0 and lines[i - 1].strip(): result.append('') # Append the current line to the result result.append(line) # Join the result list into a single string with newline characters return '
'.join(result) @app.middleware("http") async def add_user_id_and_state_cookie(request: Request, call_next): user_id = request.cookies.get("user_id") state = request.cookies.get("state") if not user_id: user_id = str(uuid.uuid4()) state = "generate" response = await call_next(request) response.set_cookie(key="user_id", value=user_id) response.set_cookie(key="state", value=state) elif not state: state = "generate" response = await call_next(request) response.set_cookie(key="state", value=state) else: response = await call_next(request) return response @app.get("/user_id") async def get_user_id(): user_id = str(uuid.uuid4()) return {"user_id": user_id} @app.get("/", response_class=HTMLResponse) def form_get(): return Template(html_template).render(idea= "This is a example of the idea geneartion", error=None, reply_count=reply_count,button_text="Generate",loading_text="Generating content, Usually takes 2-3 minutes, please wait...") @app.post("/", response_class=HTMLResponse) def form_post(request: Request,response: Response,topic: str = Form(...)): global reply_count start_time = time.time() # 获得每个网页的user_id user_id = request.cookies.get("user_id") state = request.cookies.get("state", "generate") print(user_id,state) loading_text = "Generating content, Usually takes 2-3 minutes, please wait..." if state == "generate": if not queue.empty(): queue_len = queue.qsize() if queue_len + reply_count >= MAX_REPLIES_PER_DAY: error_message = "Today's maximum number of replies has been reached. Please try again tomorrow." return Template(html_template).render(idea="", error=error_message, reply_count=reply_count, button_text="Generate",loading_text=loading_text) error_message = "There are currently {} requests being processed. If you want to queue, please click the Continue button and you will enter the queue.".format(queue_len) new_state = "continue" new_button_text = "Continue" response.set_cookie(key="state", value=new_state) return Template(html_template).render(idea="", error=error_message, reply_count=reply_count, button_text=new_button_text,loading_text=f"Generating content, Usually takes {queue_len*2}-{queue_len*3} minutes, please wait...") else: queue.put(user_id) elif state == "continue": queue.put(user_id) new_state = "generate" new_button_text = "Generate" response.set_cookie(key="state", value=new_state) queue_len = queue.qsize() # 判断当前是否轮到该用户,如果没轮到则一直等待到轮到为止 while queue.queue[0] != user_id: continue with lock: logging.info(f"Processing request for topic: {topic}") start_time = time.time() error_message = None idea = "" time_taken = 0 # 检查是否超过每日最大回复次数 if reply_count >= MAX_REPLIES_PER_DAY: error_message = "Today's maximum number of replies has been reached. Please try again tomorrow." logging.info(f"Today's maximum number of replies has been reached. Please try again tomorrow.") try: main_llm, cheap_llm = get_llms() deep_research_agent = DeepResearchAgent(llm=main_llm, cheap_llm=cheap_llm, improve_cnt=1, max_chain_length=5, min_chain_length=3, max_chain_numbers=1) print(f"begin to generate idea of topic {topic}") idea, related_experiments, entities, idea_chain, ideas, trend, future, human, year = deep_research_agent.generate_idea_with_chain(topic) idea = fix_markdown(idea) idea = markdown.markdown(idea) # 更新每日回复次数 reply_count += 1 end_time = time.time() time_taken = round(end_time - start_time, 2) logging.info(f"Successfully generated idea for topic: {topic}") except Exception as e: end_time = time.time() time_taken = round(end_time - start_time, 2) logging.error(f"Failed to generate idea for topic: {topic}, Error: {str(e)}") error_message = str(e) # 从队列中移除当前用户 queue.get() return Template(html_template).render(idea=idea, error=error_message, reply_count=reply_count, time_taken=time_taken,button_text=new_button_text,loading_text=loading_text)