添加一个统一的脚本管理服务

rag2_0/api/query_dinge_qingdan_api.py

@@ -3,6 +3,8 @@ from fastapi import FastAPI, HTTPException, Query
 from pydantic import BaseModel
 from typing import List, Optional, Dict, Any
 import uvicorn
+import asyncio
+import aiosqlite
 import sqlite3
 import sys
 import os
@@ -73,11 +75,13 @@ class QingDanDingEQueryService:
             db_file=self.db_path
         )
     
-    def get_similar_names_by_vector(self, query_text:str, vector_db:SQLiteVSS, field_map:dict, top_k:int=3, scope:str=None):
-        """使用向量检索获取相似名称"""
+    async def get_similar_names_by_vector(self, query_text:str, vector_db:SQLiteVSS, field_map:dict, top_k:int=3, scope:str=None):
+        """使用向量检索获取相似名称（异步包装）"""
         try:
-            # 使用向量数据库进行相似性搜索
-            results = vector_db.similarity_search_with_score(query=query_text, k=30)
+            # 使用线程池包装同步向量检索，避免阻塞事件循环
+            results = await asyncio.to_thread(
+                vector_db.similarity_search_with_score, query_text, 30
+            )
             
             # 提取结果中的元数据
             similar_items = []
@@ -90,16 +94,16 @@ class QingDanDingEQueryService:
                 metadata['similarity_score'] = float(score)
                 similar_items.append(metadata)
             
-            # 按相似度分数排序,分数高的排前面
+            # 分数越小越相似（SQLiteVSS 多为距离），已有代码按升序排序
             similar_items.sort(key=lambda x: x['similarity_score'])
             return similar_items[:top_k]
         except Exception as e:
             print(f"向量检索出错: {str(e)}")
             return []
 
-    def get_db_connection(self):
-        """获取数据库连接"""
-        conn = sqlite3.connect(self.db_path)
+    async def get_db_connection(self):
+        """获取数据库连接（aiosqlite 异步）"""
+        conn = await aiosqlite.connect(self.db_path)
         conn.row_factory = sqlite3.Row  # 设置行工厂，使结果可以通过列名访问
         return conn
     
@@ -138,12 +142,9 @@ class QingDanDingEQueryService:
         # 合并结果，完全匹配的排在前面
         return exact_matches + partial_matches
     
-    def query_ding_e_by_name(self, name, scope=None):
+    async def query_ding_e_by_name(self, name, scope=None):
         """根据定额名称查询定额子目表中详情信息，使用向量检索扩大查询范围"""
         try:
-            conn = self.get_db_connection()
-            cursor = conn.cursor()
-            
             # 获取表名和字段映射
             zimu_table = ExcelToSQLiteProcessor.ding_e_table_names["定额子目"]
             mulu_table = ExcelToSQLiteProcessor.ding_e_table_names["定额目录"]
@@ -151,7 +152,7 @@ class QingDanDingEQueryService:
             field_map = ExcelToSQLiteProcessor.ding_e_field_map
             
             # 1. 先使用向量检索获取相似名称
-            similar_items = self.get_similar_names_by_vector(query_text=name, 
+            similar_items = await self.get_similar_names_by_vector(query_text=name, 
                                             vector_db=self.ding_e_vector_db, 
                                             field_map=field_map,
                                             scope=scope)
@@ -190,18 +191,16 @@ class QingDanDingEQueryService:
                 query += f" AND attr.{field_map['适用范围']} LIKE ?"
                 params.append(f'%{scope}%')
             
-            cursor.execute(query, params)
-            
-            # 获取结果
-            results = cursor.fetchall()
-            data = [dict(row) for row in results]
+            async with await self.get_db_connection() as conn:
+                cursor = await conn.execute(query, params)
+                # 获取结果
+                results = await cursor.fetchall()
+                data = [dict(row) for row in results]
             
             # 对结果进行排序，将全字匹配的排在前面
             data = self.sort_results_by_exact_match(data, name, field_map['名称'])
             data = data[:self.top_k]
             
-            conn.close()
-            
             if not data:
                 return {"success": True, "message": "未找到匹配的定额信息", "data": []}
             
@@ -219,12 +218,10 @@ class QingDanDingEQueryService:
         except Exception as e:
             return {"success": False, "message": f"查询出错: {str(e)}"}
     
-    def query_ding_e_by_code(self, code, scope=None):
+    async def query_ding_e_by_code(self, code, scope=None):
         """根据定额编码查询定额子目表中详情信息"""
         try:
             code = code.upper()
-            conn = self.get_db_connection()
-            cursor = conn.cursor()
             
             # 获取表名和字段映射
             zimu_table = ExcelToSQLiteProcessor.ding_e_table_names["定额子目"]
@@ -255,18 +252,16 @@ class QingDanDingEQueryService:
                 query += f" AND attr.{field_map['适用范围']} LIKE ?"
                 params.append(f'%{scope}%')
                 
-            cursor.execute(query, params)
-            
-            # 获取结果
-            results = cursor.fetchall()
-            data = [dict(row) for row in results]
+            async with await self.get_db_connection() as conn:
+                cursor = await conn.execute(query, params)
+                # 获取结果
+                results = await cursor.fetchall()
+                data = [dict(row) for row in results]
             
             # 对结果进行排序，将全字匹配的排在前面
             data = self.sort_results_by_exact_match(data, code, field_map['编码'])
             data = data[:self.top_k]
             
-            conn.close()
-            
             if not data:
                 return {"success": True, "message": "未找到匹配的定额信息", "data": []}
             
@@ -284,12 +279,9 @@ class QingDanDingEQueryService:
         except Exception as e:
             return {"success": False, "message": f"查询出错: {str(e)}"}
     
-    def query_qing_dan_by_name(self, name, scope=None):
+    async def query_qing_dan_by_name(self, name, scope=None):
         """根据清单名称查询清单子目表中详情信息，使用向量检索扩大查询范围"""
         try:
-            conn = self.get_db_connection()
-            cursor = conn.cursor()
-            
             # 获取表名和字段映射
             zimu_table = ExcelToSQLiteProcessor.qing_dan_table_names["清单子目"]
             mulu_table = ExcelToSQLiteProcessor.qing_dan_table_names["清单目录"]
@@ -297,7 +289,7 @@ class QingDanDingEQueryService:
             field_map = ExcelToSQLiteProcessor.qing_dan_field_map
             
             # 1. 先使用向量检索获取相似名称
-            similar_items = self.get_similar_names_by_vector(query_text=name, vector_db=self.qing_dan_vector_db, field_map=field_map, scope=scope)
+            similar_items = await self.get_similar_names_by_vector(query_text=name, vector_db=self.qing_dan_vector_db, field_map=field_map, scope=scope)
             similar_names = [item['mc'] for item in similar_items]
             
             # 构建查询条件，始终包含原始名称的模糊匹配
@@ -333,18 +325,16 @@ class QingDanDingEQueryService:
                 query += f" AND attr.{field_map['适用范围']} LIKE ?"
                 params.append(f'%{scope}%')
             
-            cursor.execute(query, params)
-            
-            # 获取结果
-            results = cursor.fetchall()
-            data = [dict(row) for row in results]
+            async with await self.get_db_connection() as conn:
+                cursor = await conn.execute(query, params)
+                # 获取结果
+                results = await cursor.fetchall()
+                data = [dict(row) for row in results]
             
             # 对结果进行排序，将全字匹配的排在前面
             data = self.sort_results_by_exact_match(data, name, field_map['名称'])
             data = data[:self.top_k]
             
-            conn.close()
-            
             if not data:
                 return {"success": True, "message": "未找到匹配的清单信息", "data": []}
             
@@ -362,12 +352,10 @@ class QingDanDingEQueryService:
         except Exception as e:
             return {"success": False, "message": f"查询出错: {str(e)}"}
     
-    def query_qing_dan_by_code(self, code, scope=None):
+    async def query_qing_dan_by_code(self, code, scope=None):
         """根据清单编码查询清单子目表中详情信息"""
         try:
             code = code.upper()
-            conn = self.get_db_connection()
-            cursor = conn.cursor()
             
             # 获取表名和字段映射
             zimu_table = ExcelToSQLiteProcessor.qing_dan_table_names["清单子目"]
@@ -398,18 +386,16 @@ class QingDanDingEQueryService:
                 query += f" AND attr.{field_map['适用范围']} LIKE ?"
                 params.append(f'%{scope}%')
                 
-            cursor.execute(query, params)
-            
-            # 获取结果
-            results = cursor.fetchall()
-            data = [dict(row) for row in results]
+            async with await self.get_db_connection() as conn:
+                cursor = await conn.execute(query, params)
+                # 获取结果
+                results = await cursor.fetchall()
+                data = [dict(row) for row in results]
             
             # 对结果进行排序，将全字匹配的排在前面
             data = self.sort_results_by_exact_match(data, code, field_map['编码'])
             data = data[:self.top_k]
             
-            conn.close()
-            
             if not data:
                 return {"success": True, "message": "未找到匹配的清单信息", "data": []}
             
@@ -427,8 +413,8 @@ class QingDanDingEQueryService:
         except Exception as e:
             return {"success": False, "message": f"查询出错: {str(e)}"}
     
-    def batch_query(self, requests:BatchQueryRequest):
-        """批量查询接口，支持向量检索"""
+    async def batch_query(self, requests:BatchQueryRequest):
+        """批量查询接口，支持向量检索（并发执行）"""
         dinge_results = []
         qingdan_results = []
         tracking_dict = {}  # 用于跟踪已查询过的项目，避免重复
@@ -439,41 +425,48 @@ class QingDanDingEQueryService:
             qingdan_info = requests.dinge_qingdan_info.qingdan_info
             scope = requests.scope
             
+            dinge_tasks = []
+            qingdan_tasks = []
+            
             # 处理定额编码查询
             for code in dinge_info.dinge_code_list or []:
                 key = f"dinge_code_{code}_{scope}"
                 if key not in tracking_dict:
-                    result = self.query_ding_e_by_code(code, scope)
-                    if result["success"] and result["data"]:
-                        dinge_results.extend(result["data"])
-                        tracking_dict[key] = True
+                    dinge_tasks.append(self.query_ding_e_by_code(code, scope))
+                    tracking_dict[key] = True
             
             # 处理定额名称查询
             for name in dinge_info.dinge_name_list or []:
                 key = f"dinge_name_{name}_{scope}"
                 if key not in tracking_dict:
-                    result = self.query_ding_e_by_name(name, scope)
-                    if result["success"] and result["data"]:
-                        dinge_results.extend(result["data"])
-                        tracking_dict[key] = True
+                    dinge_tasks.append(self.query_ding_e_by_name(name, scope))
+                    tracking_dict[key] = True
             
             # 处理清单编码查询
             for code in qingdan_info.qingdan_code_list or []:
                 key = f"qingdan_code_{code}_{scope}"
                 if key not in tracking_dict:
-                    result = self.query_qing_dan_by_code(code, scope)
-                    if result["success"] and result["data"]:
-                        qingdan_results.extend(result["data"])
-                        tracking_dict[key] = True
+                    qingdan_tasks.append(self.query_qing_dan_by_code(code, scope))
+                    tracking_dict[key] = True
             
             # 处理清单名称查询
             for name in qingdan_info.qingdan_name_list or []:
                 key = f"qingdan_name_{name}_{scope}"
                 if key not in tracking_dict:
-                    result = self.query_qing_dan_by_name(name, scope)
-                    if result["success"] and result["data"]:
+                    qingdan_tasks.append(self.query_qing_dan_by_name(name, scope))
+                    tracking_dict[key] = True
+            
+            # 并发执行
+            if dinge_tasks:
+                dinge_outs = await asyncio.gather(*dinge_tasks)
+                for result in dinge_outs:
+                    if result and result.get("success") and result.get("data"):
+                        dinge_results.extend(result["data"])
+            if qingdan_tasks:
+                qingdan_outs = await asyncio.gather(*qingdan_tasks)
+                for result in qingdan_outs:
+                    if result and result.get("success") and result.get("data"):
                         qingdan_results.extend(result["data"])
-                        tracking_dict[key] = True
             
             # 限制返回结果数量
             dinge_results = dinge_results[:self.top_k]
@@ -505,7 +498,7 @@ async def query_ding_e_by_name(
     name: str = Query(..., description="定额名称"),
     scope: Optional[str] = Query(None, description="适用范围")
 ):
-    result = query_service.query_ding_e_by_name(name, scope)
+    result = await query_service.query_ding_e_by_name(name, scope)
     if not result["success"]:
         raise HTTPException(status_code=500, detail=result["message"])
     return QueryResponse(**result)
@@ -516,7 +509,7 @@ async def query_ding_e_by_code(
     code: str = Query(..., description="定额编码"),
     scope: Optional[str] = Query(None, description="适用范围")
 ):
-    result = query_service.query_ding_e_by_code(code, scope)
+    result = await query_service.query_ding_e_by_code(code, scope)
     if not result["success"]:
         raise HTTPException(status_code=500, detail=result["message"])
     return QueryResponse(**result)
@@ -527,7 +520,7 @@ async def query_qing_dan_by_name(
     name: str = Query(..., description="清单名称"),
     scope: Optional[str] = Query(None, description="适用范围")
 ):
-    result = query_service.query_qing_dan_by_name(name, scope)
+    result = await query_service.query_qing_dan_by_name(name, scope)
     if not result["success"]:
         raise HTTPException(status_code=500, detail=result["message"])
     return QueryResponse(**result)
@@ -538,7 +531,7 @@ async def query_qing_dan_by_code(
     code: str = Query(..., description="清单编码"),
     scope: Optional[str] = Query(None, description="适用范围")
 ):
-    result = query_service.query_qing_dan_by_code(code, scope)
+    result = await query_service.query_qing_dan_by_code(code, scope)
     if not result["success"]:
         raise HTTPException(status_code=500, detail=result["message"])
     return QueryResponse(**result)
@@ -546,7 +539,7 @@ async def query_qing_dan_by_code(
 # 5. 批量查询定额和清单信息
 @app.post("/api/batch_query", response_model=BatchQueryResponse)
 async def batch_query(request: BatchQueryRequest):
-    result = query_service.batch_query(request)
+    result = await query_service.batch_query(request)
     if not result["success"]:
         raise HTTPException(status_code=500, detail=result["message"])
     return BatchQueryResponse(**result)
@@ -564,4 +557,4 @@ def main():
 
 if __name__ == "__main__":
     main()
-    # uvicorn rag2_0.dify.query_dinge_qingdan_api:app --host 0.0.0.0 --port 8005 --workers 10
+    # uvicorn rag2_0.api.query_dinge_qingdan_api:app --host 0.0.0.0 --port 8005 --workers 10