rag_server/routes/buleprint.py

import uuid

from quart import Blueprint, jsonify, request

from applications.config import DEFAULT_MODEL, LOCAL_MODEL_CONFIG, ChunkerConfig
from applications.api import get_basic_embedding
from applications.async_task import ChunkEmbeddingTask

server_bp = Blueprint("api", __name__, url_prefix="/api")


def server_routes(mysql_db, vector_db):

    @server_bp.route("/embed", methods=["POST"])
    async def embed():
        body = await request.get_json()
        text = body.get("text")
        model_name = body.get("model", DEFAULT_MODEL)
        if not LOCAL_MODEL_CONFIG.get(model_name):
            return jsonify({"error": "error  model"})

        embedding = await get_basic_embedding(text, model_name)
        return jsonify({"embedding": embedding})

    @server_bp.route("/chunk", methods=["POST"])
    async def chunk():
        body = await request.get_json()
        text = body.get("text", "")
        text = text.strip()
        if not text:
            return jsonify({"error": "error  text"})
        doc_id = f"doc-{uuid.uuid4()}"
        chunk_task = ChunkEmbeddingTask(mysql_db, vector_db, cfg=ChunkerConfig(), doc_id=doc_id)
        doc_id = await chunk_task.deal(body)
        return jsonify({"doc_id": doc_id})

    @server_bp.route("/search", methods=["POST"])
    async def search():
        pass

    return server_bp