基于Playwright数据采集，大模型对接Flask+SQLite，平滑升级向量库与分布式微服务26.5

# app.py 存量基础服务
from flask import Flask, request, jsonify
import sqlite3
from playwright.sync_api import sync_playwright
import os

app = Flask(__name__)
DB_PATH = "crawl_data.db"

# 初始化sqlite表
def init_db():
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS crawl_task (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            url TEXT,
            status TEXT DEFAULT 'pending',
            raw_content TEXT,
            create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    conn.commit()
    conn.close()

# playwright采集函数
def crawl_page(url):
    with sync_playwright() as p:
        browser = p.chromium.launch(headless=True)
        page = browser.new_page()
        page.goto(url, timeout=30000)
        content = page.locator("body").text_content()
        browser.close()
    return content

@app.route("/task", methods=["POST"])
def create_task():
    data = request.get_json()
    target_url = data.get("url")
    conn = sqlite3.connect(DB_PATH)
    cur = conn.cursor()
    cur.execute("INSERT INTO crawl_task(url, status) VALUES (?, 'pending')", (target_url,))
    task_id = cur.lastrowid
    conn.commit()
    conn.close()
    return jsonify({"task_id": task_id, "msg": "任务创建成功"})

@app.route("/task/run/<int:task_id>", methods=["POST"])
def run_task(task_id):
    conn = sqlite3.connect(DB_PATH)
    cur = conn.cursor()
    cur.execute("SELECT url FROM crawl_task WHERE id=?", (task_id,))
    row = cur.fetchone()
    if not row:
        return jsonify({"msg": "任务不存在"}),404
    url = row[0]
    raw_text = crawl_page(url)
    cur.execute("UPDATE crawl_task SET raw_content=?, status='finished' WHERE id=?", (raw_text, task_id))
    conn.commit()
    conn.close()
    return jsonify({"task_id":task_id, "content_len":len(raw_text)})

if __name__ == "__main__":
    init_db()
    app.run(debug=True, port=5000)

# 在原有app.py基础上新增代码
import requests

LLM_API_KEY = "your_api_key"
LLM_ENDPOINT = "https://xxx/v1/chat/completions"

def llm_extract_info(raw_text):
    prompt = f"""
请对下面网页文本做两件事：1.生成200字摘要；2.提取5个核心关键词。
文本：{raw_text[:6000]}
输出JSON格式，包含summary、keywords两个字段
"""
    headers = {"Authorization":f"Bearer {LLM_API_KEY}", "Content-Type":"application/json"}
    payload = {
        "model":"xxx",
        "messages":[{"role":"user","content":prompt}],
        "temperature":0.3
    }
    resp = requests.post(LLM_ENDPOINT, headers=headers, json=payload, timeout=40)
    res = resp.json()
    return res["choices"][0]["message"]["content"]

# 修改数据库初始化，新增llm_result字段
def init_db():
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS crawl_task (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            url TEXT,
            status TEXT DEFAULT 'pending',
            raw_content TEXT,
            llm_result TEXT,
            create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    conn.commit()
    conn.close()

# 修改任务执行接口，采集完成后调用大模型
@app.route("/task/run/<int:task_id>", methods=["POST"])
def run_task(task_id):
    conn = sqlite3.connect(DB_PATH)
    cur = conn.cursor()
    cur.execute("SELECT url FROM crawl_task WHERE id=?", (task_id,))
    row = cur.fetchone()
    if not row:
        return jsonify({"msg": "任务不存在"}),404
    url = row[0]
    raw_text = crawl_page(url)
    try:
        llm_output = llm_extract_info(raw_text)
    except Exception as e:
        llm_output = f"LLM处理失败:{str(e)}"
    cur.execute("UPDATE crawl_task SET raw_content=?, llm_result=?, status='finished' WHERE id=?",
                (raw_text, llm_output, task_id))
    conn.commit()
    conn.close()
    return jsonify({"task_id":task_id, "content_len":len(raw_text), "llm_result":llm_output})

import psycopg2

PG_CONFIG = {
    "host":"127.0.0.1",
    "port":5432,
    "dbname":"crawl_prod",
    "user":"postgres",
    "password":"xxx"
}

def write_to_pg(task_id, url, raw_text, llm_result):
    conn = psycopg2.connect(**PG_CONFIG)
    cur = conn.cursor()
    sql = """
    INSERT INTO crawl_task (id, url, raw_content, llm_result, status)
    VALUES (%s,%s,%s,%s,%s) ON CONFLICT(id) DO UPDATE
    SET raw_content=%s, llm_result=%s, status='finished'
    """
    cur.execute(sql, (task_id, url, raw_text, llm_result, "finished", raw_text, llm_result))
    conn.commit()
    cur.close()
    conn.close()

from chromadb import Client
import requests

client = Client()
collection = client.get_or_create_collection("web_crawl_embedding")

def get_embedding(text):
    emb_url = "https://xxx/v1/embeddings"
    resp = requests.post(emb_url, json={"input":text, "model":"embedding-model"})
    return resp.json()["data"][0]["embedding"]

def save_embedding(task_id, chunk_text):
    vec = get_embedding(chunk_text)
    collection.add(
        embeddings=[vec],
        documents=[chunk_text],
        ids=[f"task_{task_id}_{hash(chunk_text)}"]
    )

# 语义检索
def search_by_query(query, top_n=3):
    q_vec = get_embedding(query)
    res = collection.query(query_embeddings=[q_vec], n_results=top_n)
    return res["documents"][0]

# redis队列简单任务分发示例
import redis
import json

r = redis.Redis(host="127.0.0.1", port=6379, decode_responses=True)

# web网关提交任务
def submit_crawl_task(url, task_id):
    msg = json.dumps({"task_id":task_id, "url":url})
    r.lpush("crawl_task_queue", msg)

# 采集服务消费任务
def consume_crawl_task():
    while True:
        msg = r.brpop("crawl_task_queue", timeout=2)
        if not msg:
            continue
        data = json.loads(msg[1])
        task_url = data["url"]
        task_id = data["task_id"]
        raw_txt = crawl_page(task_url)
        # 推送采集完成事件，交给LLM服务消费
        finish_msg = json.dumps({"task_id":task_id, "raw_text":raw_txt})
        r.lpush("llm_process_queue", finish_msg)

import time

def llm_with_retry(prompt, max_retry=3):
    retry_count = 0
    while retry_count < max_retry:
        try:
            return llm_extract_info(prompt)
        except Exception as e:
            retry_count +=1
            time.sleep(1.5 * retry_count)
    raise Exception("大模型调用多次失败")