diff --git a/dsRag/StartSSE.py b/dsRag/StartSSE.py new file mode 100644 index 00000000..084f15a6 --- /dev/null +++ b/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) \ No newline at end of file