app.py

from flask import Flask, request, jsonify, Response, stream_with_context,send_file
import requests
import json
import os
from datetime import datetime
import json
from typing import List
from pydantic import BaseModel, Field
from fpdf import FPDF
from urllib.parse import quote

app = Flask(__name__, static_folder='templates',static_url_path="/")

# Configuration
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY', '')
OPENAI_API_BASE = os.getenv('OPENAI_API_BASE', '')
OPENAI_MODEL = os.getenv('OPENAI_MODEL', '')
OPENAI_MODEL_NO_SEARCH = os.getenv('OPENAI_MODEL_NO_SEARCH', '')
GEMINI_OPENAI_API_BASE = os.getenv('GEMINI_OPENAI_API_BASE', '')
GEMINI_OPENAI_MODEL = os.getenv('GEMINI_OPENAI_MODEL', '')
GEMINI_OPENAI_MODEL_NO_SEARCH = os.getenv('GEMINI_OPENAI_MODEL_NO_SEARCH', '')

def get_system_prompt() -> str:
    now = datetime.utcnow().isoformat() + "Z"  # 当前 UTC 时间的 ISO 格式字符串
    return f"""You are an expert researcher. Today is {now}. Follow these instructions when responding:
- You may be asked to research subjects that is after your knowledge cutoff, assume the user is right when presented with news.
- The user is a highly experienced analyst, no need to simplify it, be as detailed as possible and make sure your response is correct.
- Be highly organized.
- Suggest solutions that I didn't think about.
- Be proactive and anticipate my needs.
- Treat me as an expert in all subject matter.
- Mistakes erode my trust, so be accurate and thorough.
- Provide detailed explanations, I'm comfortable with lots of detail.
- Value good arguments over authorities, the source is irrelevant.
- Consider new technologies and contrarian ideas, not just the conventional wisdom.
- You may use high levels of speculation or prediction, just flag it for me."""


class SERPQuery(BaseModel):
    query: str = Field(..., description="The SERP query.")
    researchGoal: str = Field(
        ...,
        description=(
            "First talk about the goal of the research that this query is meant to accomplish, then go deeper into how to advance "
            "the research once the results are found, mention additional research directions. Be as specific as possible, especially "
            "for additional research directions. JSON reserved words should be escaped."
        )
    )


def get_serp_query_schema() -> str:
    """
    返回表示 SERP 查询列表的 JSON Schema 字符串。
    """
    schema = {
        "type": "array",
        "items": SERPQuery.schema()
    }
    return json.dumps(schema, indent=4)


def generate_questions_prompt(query: str) -> str:
    parts = [
        f"Given the following query from the user, ask at least 5 follow-up questions to clarify the research direction: <query>{query}</query>",
        "Questions need to be brief and concise. No need to output content that is irrelevant to the question."
    ]
    return "\n\n".join(parts)


def generate_serp_queries_prompt(query: str) -> str:
    output_schema = get_serp_query_schema()
    parts = [
        f"Given the following query from the user:\n<query>{query}</query>",
        "Based on previous user query, generate a list of SERP queries to further research the topic. Make sure each query is unique and not similar to each other.",
        f"You MUST respond in `JSON` matching this `JSON schema`:\n```json\n{output_schema}\n```",
        "Expected output:\n```json\n[{query: \"This is a sample query. \", researchGoal: \"This is the reason for the query. \"}]\n```"
    ]
    return "\n\n".join(parts)


def process_search_result_prompt(query: str, research_goal: str) -> str:
    parts = [
        f"Please use the following query to get the latest information via google search tool:\n<query>{query}</query>",
        f"You need to organize the searched information according to the following requirements:\n<researchGoal>\n{research_goal}\n</researchGoal>",
        "You need to think like a human researcher. Generate a list of learnings from the search results. Make sure each learning is unique and not similar to each other. The learnings should be to the point, as detailed and information dense as possible. Make sure to include any entities like people, places, companies, products, things, etc in the learnings, as well as any specific entities, metrics, numbers, and dates when available. The learnings will be used to research the topic further."
    ]
    return "\n\n".join(parts)


def review_serp_queries_prompt(query: str, learnings: List[str], suggestion: str) -> str:
    output_schema = get_serp_query_schema()
    learnings_string = "\n".join([f"<learning>\n{learning}\n</learning>" for learning in learnings])
    parts = [
        f"Given the following query from the user:\n<query>{query}</query>",
        f"Here are all the learnings from previous research:\n<learnings>\n{learnings_string}\n</learnings>",
        f"This is the user's suggestion for research direction:\n<suggestion>\n{suggestion}\n</suggestion>",
        "Based on previous research and user research suggestions, determine whether further research is needed. If further research is needed, list of follow-up SERP queries to research the topic further. Make sure each query is unique and not similar to each other. If you believe no further research is needed, you can output an empty queries.",
        f"You MUST respond in `JSON` matching this `JSON schema`: \n```json\n{output_schema}\n```",
        "Expected output:\n```json\n[{query: \"This is a sample query. \", researchGoal: \"This is the reason for the query. \"}]\n```"
    ]
    return "\n\n".join(parts)


def write_final_report_prompt(query: str, learnings: List[str]) -> str:
    learnings_string = "\n".join([f"<learning>\n{learning}\n</learning>" for learning in learnings])
    today = datetime.now().strftime("%Y-%m-%d")
    parts = [
        f"Given the following query from the user, write a final report on the topic using the learnings from research. Make it as as detailed as possible, aim for 3 or more pages, include ALL the learnings from research:\n<query>{query}</query>",
        f"Here are all the learnings from previous research:\n<learnings>\n{learnings_string}\n</learnings>",
        "You need to write this report like a human researcher. Humans don not wrap their writing in markdown blocks. Contains diverse data information such as pictures, katex formulas, mermaid diagrams, etc. in the form of markdown syntax.",
        f"The report must include a generation date below the title, formatted as follows: **报告生成日期:{today}**"
    ]
    return "\n\n".join(parts)

# Helper functions
def get_headers():
    return {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {OPENAI_API_KEY}'
    }


def generate_stream_response(prompt, system_message=None, model=OPENAI_MODEL):
    """Generate a streaming response from the OpenAI-compatible API"""
    messages = []
    content_cache = ""
    
    print("prompt:"+prompt)
    print("system_message"+str(system_message))

    if system_message:
        messages.append({"role": "system", "content": system_message})
    
    messages.append({"role": "user", "content": prompt})
    
    payload = {
        "model": model,
        "messages": messages,
        "output_think": False,
        "show_search": False,
        "stream": True,
        "temperature": 0.1
    }
    
    try:
        response = requests.post(
            f"{OPENAI_API_BASE}/chat/completions",
            headers=get_headers(),
            json=payload,
            stream=True
        )
        
        if response.status_code != 200:
            yield f"data: {json.dumps({'error': f'API Error: {response.status_code}'})}\n\n"
            return
        
        for line in response.iter_lines():
            if line:
                line = line.decode('utf-8')
                if line.startswith('data: '):
                    data = line[6:]
                    if data == '[DONE]':
                        break
                    try:
                        json_data = json.loads(data)
                        if 'choices' in json_data and len(json_data['choices']) > 0:
                            delta = json_data['choices'][0].get('delta', {})
                            if 'content' in delta:
                                yield f"data: {json.dumps({'content': delta['content']}, ensure_ascii=False)}\n\n"
                    except json.JSONDecodeError:
                        continue
    except Exception as e:
        yield f"data: {json.dumps({'error': str(e)})}\n\n"


def generate_stream_response_gemini(prompt, system_message=None, model=GEMINI_OPENAI_MODEL):
    """Generate a streaming response from the OpenAI-compatible API"""
    messages = []
    content_cache = ""
    
    print("prompt:"+prompt)
    print("system_message"+str(system_message))

    if system_message:
        messages.append({"role": "system", "content": system_message})
    
    messages.append({"role": "user", "content": prompt})
    
    payload = {
        "model": model,
        "messages": messages,
        "output_think": False,
        "show_search": False,
        "stream": True,
        "temperature": 0.1
    }
    
    try:
        response = requests.post(
            f"{GEMINI_OPENAI_API_BASE}/chat/completions",
            headers=get_headers(),
            json=payload,
            stream=True
        )
        
        if response.status_code != 200:
            yield f"data: {json.dumps({'error': f'API Error: {response.status_code}'})}\n\n"
            return
        
        for line in response.iter_lines():
            if line:
                line = line.decode('utf-8')
                if line.startswith('data: '):
                    data = line[6:]
                    if data == '[DONE]':
                        break
                    try:
                        json_data = json.loads(data)
                        if 'choices' in json_data and len(json_data['choices']) > 0:
                            delta = json_data['choices'][0].get('delta', {})
                            if 'content' in delta:
                                yield f"data: {json.dumps({'content': delta['content']}, ensure_ascii=False)}\n\n"
                    except json.JSONDecodeError:
                        continue
    except Exception as e:
        yield f"data: {json.dumps({'error': str(e)})}\n\n"


def generate_response(prompt, system_message=None, model=OPENAI_MODEL):
    """Generate a streaming response from the OpenAI-compatible API"""
    messages = []
    
    print("prompt:"+prompt)
    print("system_message"+str(system_message))

    if system_message:
        messages.append({"role": "system", "content": system_message})
    
    messages.append({"role": "user", "content": prompt})
    
    payload = {
        "model": model,
        "messages": messages,
        "output_think": False,
        "show_search": False,
        "stream": False,
        "temperature": 0.1
    }
    
    response = requests.post(
        f"{OPENAI_API_BASE}/chat/completions",
        headers=get_headers(),
        json=payload,
        stream=True
    )

    return jsonify(response.json()), response.status_code

# Routes
@app.route('/')
def main():
    return app.send_static_file('index.html')

def generate_pdf(text, filename):
    pdf = FPDF()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.add_page()
        # 设置字体，使用支持 UTF-8 的 TTF 字体
    font_path = "NotoSansSC-Regular.ttf" 
    pdf.add_font("NotoSansSC", "", font_path, uni=True)
    pdf.set_font("NotoSansSC", size=12)
    
    pdf.multi_cell(0, 10, text)
    pdf.output(filename)


@app.route('/api/think', methods=['POST'])
def think():
    data = request.json
    topic = data.get('topic', '')
    
    if not topic:
        return jsonify({'error': '请提供研究主题'}), 400
    
    system_message = get_system_prompt()
    
    prompt = generate_questions_prompt(topic)
     
    return generate_stream_response(prompt, system_message, OPENAI_MODEL_NO_SEARCH)

@app.route('/api/research', methods=['POST'])
def research():
    data = request.json
    query = data.get('query', '')
    
    if not query:
        return jsonify({'error': '请提供研究主题'}), 400
    
    system_message = get_system_prompt()
    
    prompt = generate_serp_queries_prompt(query)
    
    return generate_response(prompt,system_message)

@app.route('/api/taskExec', methods=['POST'])
def task_exec():
    data = request.json
    query = data.get('query', '')
    researchGoal = data.get('researchGoal', '')
    
    system_message = get_system_prompt()
    
    prompt = process_search_result_prompt(query,researchGoal)
    
    return generate_stream_response(prompt,system_message)


@app.route('/api/report', methods=['POST'])
def generate_report():
    data = request.json
    query = data.get('query', '')
    learnings = data.get('learnings', '')
    
    system_message = get_system_prompt()
    
    prompt = write_final_report_prompt(query,learnings)
    
    return generate_stream_response(prompt, system_message, OPENAI_MODEL_NO_SEARCH)


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=11997,debug=True)