'commit'

dsRag/StartSSE.py

@@ -0,0 +1,111 @@
+"""
+conda activate rag
+pip install fastapi uvicorn sse-starlette
+"""
+import asyncio
+
+from elasticsearch import Elasticsearch
+from fastapi import FastAPI
+from openai import OpenAI
+from sse_starlette.sse import EventSourceResponse
+
+from Config import Config
+
+app = FastAPI()
+
+# 初始化ES连接
+es = Elasticsearch(
+    hosts=Config.ES_CONFIG['hosts'],
+    basic_auth=Config.ES_CONFIG['basic_auth'],
+    verify_certs=Config.ES_CONFIG['verify_certs']
+)
+
+# 初始化DeepSeek客户端
+client = OpenAI(
+    api_key=Config.DEEPSEEK_API_KEY,
+    base_url=Config.DEEPSEEK_URL
+)
+
+def search_related_data(query):
+    """搜索与查询相关的数据"""
+    # 向量搜索
+    vector_results = es.search(
+        index=Config.ES_CONFIG['default_index'],
+        body={
+            "query": {
+                "match": {
+                    "content": {
+                        "query": query,
+                        "analyzer": "ik_smart"
+                    }
+                }
+            },
+            "size": 5
+        }
+    )
+    
+    # 文本精确搜索
+    text_results = es.search(
+        index="raw_texts",
+        body={
+            "query": {
+                "match": {
+                    "text.keyword": query
+                }
+            },
+            "size": 5
+        }
+    )
+    
+    # 合并结果
+    context = ""
+    for hit in vector_results['hits']['hits']:
+        context += f"向量相似度结果(score={hit['_score']}):\n{hit['_source']['text']}\n\n"
+    
+    for hit in text_results['hits']['hits']:
+        context += f"文本精确匹配结果(score={hit['_score']}):\n{hit['_source']['text']}\n\n"
+    
+    return context
+
+async def generate_stream(query):
+    """生成SSE流"""
+    context = search_related_data(query)
+    
+    prompt = f"""根据以下关于'{query}'的相关信息，整理一份结构化的报告：
+要求：
+1. 分章节组织内容
+2. 包含关键数据和事实
+3. 语言简洁专业
+
+相关信息：
+{context}"""
+    
+    try:
+        response = client.chat.completions.create(
+            model="deepseek-chat",
+            messages=[
+                {"role": "system", "content": "你是一个专业的文档整理助手"},
+                {"role": "user", "content": prompt}
+            ],
+            temperature=0.3,
+            stream=True
+        )
+        
+        for chunk in response:
+            if chunk.choices[0].delta.content:
+                yield {"data": chunk.choices[0].delta.content}
+            await asyncio.sleep(0.01)
+    except Exception as e:
+        yield {"data": f"生成报告时出错: {str(e)}"}
+
+@app.get("/api/rag")
+async def rag_stream(query: str):
+    """RAG+DeepSeek流式接口"""
+    return EventSourceResponse(generate_stream(query))
+
+"""
+http://localhost:8000/api/rag?query=整理云南省初中在校生情况文档
+"""
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)