import uuid from quart import Blueprint, jsonify, request from applications.config import DEFAULT_MODEL, LOCAL_MODEL_CONFIG, ChunkerConfig, WEIGHT_MAP from applications.api import get_basic_embedding from applications.async_task import ChunkEmbeddingTask from applications.utils.milvus import MilvusSearcher 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(): body = await request.get_json() search_type = body.get("search_type") if not search_type: return jsonify({"error": "missing search_type"}), 400 searcher = MilvusSearcher(vector_db) try: # 统一参数 expr = body.get("expr") search_params = body.get("search_params") or {"metric_type": "COSINE", "params": {"ef": 64}} limit = body.get("limit", 5) query = body.get("query") # 定义不同搜索策略 async def by_pk_id(): pk_id = body.get("id") if not pk_id: return {"error": "missing id"} return await searcher.get_by_id(pk_id) async def by_doc_id(): doc_id, chunk_id = body.get("doc_id"), body.get("chunk_id") if not doc_id or chunk_id is None: return {"error": "missing doc_id or chunk_id"} return await searcher.get_by_doc_and_chunk(doc_id, chunk_id) async def by_vector(): if not query: return {"error": "missing query"} field = body.get("field", "vector_text") query_vec = await get_basic_embedding(text=query, model=DEFAULT_MODEL) return await searcher.vector_search( query_vec=query_vec, anns_field=field, expr=expr, search_params=search_params, limit=limit, ) async def by_filter(): filter_map = body.get("filter_map") if not filter_map: return {"error": "missing filter_map"} return await searcher.filter_search(filter_map) async def hybrid(): if not query: return {"error": "missing query"} field = body.get("field", "vector_text") query_vec = await get_basic_embedding(text=query, model=DEFAULT_MODEL) return await searcher.hybrid_search( query_vec=query_vec, anns_field=field, filters=body.get("filter_map"), limit=limit, ) async def strategy(): if not query: return {"error": "missing query"} query_vec = await get_basic_embedding(text=query, model=DEFAULT_MODEL) return await searcher.search_by_strategy( query_vec=query_vec, weight_map=body.get("weight_map", WEIGHT_MAP), expr=expr, limit=limit, ) # dispatch table handlers = { "pk_id": by_pk_id, "by_doc_id": by_doc_id, "by_vector": by_vector, "by_filter": by_filter, "hybrid": hybrid, "strategy": strategy, } if search_type not in handlers: return jsonify({"error": "invalid search_type"}), 400 result = await handlers[search_type]() return jsonify(result) except Exception as e: return jsonify({"error": str(e)}), 500 return server_bp