'commit'

dsRag/StartEs.py

@@ -17,11 +17,12 @@ from pydantic import BaseModel, Field, ValidationError
 from starlette.responses import StreamingResponse
 
 from Config.Config import MS_MODEL_PATH, MS_MODEL_LIMIT, MS_HOST, MS_PORT, MS_MAX_CONNECTIONS, MS_NPROBE, \
-    MS_COLLECTION_NAME
+    MS_COLLECTION_NAME, ES_CONFIG
 from Milvus.Utils.MilvusCollectionManager import MilvusCollectionManager
 from Milvus.Utils.MilvusConnectionPool import *
 from Milvus.Utils.MilvusConnectionPool import MilvusConnectionPool
 from Util.ALiYunUtil import ALiYunUtil
+from Util.EsSearchUtil import EsSearchUtil
 
 # 初始化日志
 logger = logging.getLogger(__name__)
@@ -36,13 +37,12 @@ def html_to_word_pandoc(html_file, output_file):
     subprocess.run(['pandoc', html_file, '-o', output_file])
 
 
-@asynccontextmanager
 async def lifespan(app: FastAPI):
     # 初始化阿里云大模型工具
     app.state.aliyun_util = ALiYunUtil()
     yield
-    pass
-
+    # 清理资源
+    await app.state.aliyun_util.close()
 app = FastAPI(lifespan=lifespan)
 
 # 挂载静态文件目录
@@ -80,7 +80,7 @@ async def save_to_word(request: Request):
             f.write(html_content)
 
         # 使用pandoc转换
-        output_file = os.path.join(tempfile.gettempdir(), "小学数学问答.docx")
+        output_file = os.path.join(tempfile.gettempdir(), "【理想大模型】问答.docx")
         subprocess.run(['pandoc', temp_html, '-o', output_file], check=True)
 
         # 读取生成的Word文件
@@ -88,7 +88,7 @@ async def save_to_word(request: Request):
             stream = BytesIO(f.read())
 
         # 返回响应
-        encoded_filename = urllib.parse.quote("小学数学问答.docx")
+        encoded_filename = urllib.parse.quote("【理想大模型】问答.docx")
         return StreamingResponse(
             stream,
             media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
@@ -112,8 +112,84 @@ async def save_to_word(request: Request):
 
 @app.post("/api/rag")
 async def rag_stream(request: Request):
-    pass
-    # todo
+    try:
+        data = await request.json()
+        query = data.get('query', '')
+        query_tags = data.get('tags', [])
+        
+        # 获取EsSearchUtil实例
+        es_search_util = EsSearchUtil(ES_CONFIG)
+        
+        # 执行混合搜索
+        es_conn = es_search_util.es_pool.get_connection()
+        try:
+            # 向量搜索
+            query_embedding = es_search_util.text_to_embedding(query)
+            vector_results = es_conn.search(
+                index=ES_CONFIG['index_name'],
+                body={
+                    "query": {
+                        "script_score": {
+                            "query": {
+                                "bool": {
+                                    "should": [
+                                        {
+                                            "terms": {
+                                                "tags.tags": query_tags
+                                            }
+                                        }
+                                    ],
+                                    "minimum_should_match": 1
+                                }
+                            },
+                            "script": {
+                                "source": "double score = cosineSimilarity(params.query_vector, 'embedding'); return score >= 0 ? score : 0",
+                                "params": {"query_vector": query_embedding}
+                            }
+                        }
+                    },
+                    "size": 3
+                }
+            )
+            
+            # 文本精确搜索
+            text_results = es_conn.search(
+                index=ES_CONFIG['index_name'],
+                body={
+                    "query": {
+                        "bool": {
+                            "must": [
+                                {
+                                    "match": {
+                                        "user_input": query
+                                    }
+                                },
+                                {
+                                    "terms": {
+                                        "tags.tags": query_tags
+                                    }
+                                }
+                            ]
+                        }
+                    },
+                    "size": 3
+                }
+            )
+            
+            # 合并结果
+            results = {
+                "vector_results": [hit['_source'] for hit in vector_results['hits']['hits']],
+                "text_results": [hit['_source'] for hit in text_results['hits']['hits']]
+            }
+            
+            return results
+            
+        finally:
+            es_search_util.es_pool.release_connection(es_conn)
+            
+    except Exception as e:
+        logger.error(f"RAG search error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))