You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

184 lines
6.6 KiB

4 weeks ago
import asyncio
4 weeks ago
from contextlib import asynccontextmanager
4 weeks ago
from logging.handlers import RotatingFileHandler
4 weeks ago
4 weeks ago
import jieba # 导入 jieba 分词库
4 weeks ago
import uvicorn
4 weeks ago
from fastapi import FastAPI, Request, HTTPException
from pydantic import BaseModel, Field, ValidationError
from fastapi.staticfiles import StaticFiles
4 weeks ago
from openai import OpenAI
4 weeks ago
from sse_starlette.sse import EventSourceResponse
from gensim.models import KeyedVectors
4 weeks ago
from Config import Config
4 weeks ago
from Config.Config import MS_MODEL_PATH, MS_MODEL_LIMIT, MS_HOST, MS_PORT, MS_MAX_CONNECTIONS, MS_NPROBE, DEEPSEEK_API_KEY, DEEPSEEK_URL, MS_COLLECTION_NAME
4 weeks ago
from Milvus.Utils.MilvusCollectionManager import MilvusCollectionManager
from Milvus.Utils.MilvusConnectionPool import *
from Milvus.Utils.MilvusConnectionPool import MilvusConnectionPool
4 weeks ago
4 weeks ago
# 初始化日志
4 weeks ago
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
4 weeks ago
handler = RotatingFileHandler('Logs/start.log', maxBytes=1024 * 1024, backupCount=5)
4 weeks ago
handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
logger.addHandler(handler)
4 weeks ago
4 weeks ago
# 1. 加载预训练的 Word2Vec 模型
model_path = MS_MODEL_PATH # 替换为你的 Word2Vec 模型路径
model = KeyedVectors.load_word2vec_format(model_path, binary=False, limit=MS_MODEL_LIMIT)
4 weeks ago
logger.info(f"模型加载成功,词向量维度: {model.vector_size}")
4 weeks ago
4 weeks ago
@asynccontextmanager
async def lifespan(app: FastAPI):
4 weeks ago
# 初始化Milvus连接池
app.state.milvus_pool = MilvusConnectionPool(host=MS_HOST, port=MS_PORT, max_connections=MS_MAX_CONNECTIONS)
4 weeks ago
4 weeks ago
# 初始化集合管理器
app.state.collection_manager = MilvusCollectionManager(MS_COLLECTION_NAME)
app.state.collection_manager.load_collection()
4 weeks ago
4 weeks ago
# 初始化DeepSeek客户端
app.state.deepseek_client = OpenAI(
api_key=Config.DEEPSEEK_API_KEY,
base_url=Config.DEEPSEEK_URL
)
4 weeks ago
yield
4 weeks ago
4 weeks ago
# 关闭Milvus连接池
app.state.milvus_pool.close()
4 weeks ago
4 weeks ago
4 weeks ago
app = FastAPI(lifespan=lifespan)
4 weeks ago
4 weeks ago
# 挂载静态文件目录
app.mount("/static", StaticFiles(directory="Static"), name="static")
4 weeks ago
4 weeks ago
# 将文本转换为嵌入向量
def text_to_embedding(text):
words = jieba.lcut(text) # 使用 jieba 分词
print(f"文本: {text}, 分词结果: {words}")
embeddings = [model[word] for word in words if word in model]
4 weeks ago
logger.info(f"有效词向量数量: {len(embeddings)}")
4 weeks ago
if embeddings:
avg_embedding = sum(embeddings) / len(embeddings)
4 weeks ago
logger.info(f"生成的平均向量: {avg_embedding[:5]}...") # 打印前 5 维
4 weeks ago
return avg_embedding
else:
4 weeks ago
logger.warning("未找到有效词,返回零向量")
4 weeks ago
return [0.0] * model.vector_size
4 weeks ago
4 weeks ago
async def generate_stream(client, milvus_pool, collection_manager, query):
4 weeks ago
"""生成SSE流"""
4 weeks ago
# 从连接池获取连接
connection = milvus_pool.get_connection()
4 weeks ago
4 weeks ago
try:
# 1. 将查询文本转换为向量
current_embedding = text_to_embedding(query)
# 2. 搜索相关数据
search_params = {
"metric_type": "L2", # 使用 L2 距离度量方式
"params": {"nprobe": MS_NPROBE} # 设置 IVF_FLAT 的 nprobe 参数
}
# 7. 将文本转换为嵌入向量
4 weeks ago
results = collection_manager.search(current_embedding, search_params, limit=5) # 返回 2 条结果
4 weeks ago
4 weeks ago
# 3. 处理搜索结果
4 weeks ago
logger.info("最相关的历史对话:")
4 weeks ago
context = ""
4 weeks ago
if results:
for hits in results:
for hit in hits:
try:
# 查询非向量字段
record = collection_manager.query_by_id(hit.id)
4 weeks ago
logger.info(f"ID: {hit.id}")
logger.info(f"标签: {record['tags']}")
logger.info(f"用户问题: {record['user_input']}")
4 weeks ago
# 获取完整内容
full_content = record['tags'].get('full_content', record['user_input'])
context = context + full_content
4 weeks ago
logger.info(f"时间: {record['timestamp']}")
logger.info(f"距离: {hit.distance}")
logger.info("-" * 40) # 分隔线
4 weeks ago
except Exception as e:
4 weeks ago
logger.error(f"查询失败: {e}")
4 weeks ago
else:
4 weeks ago
logger.warning("未找到相关历史对话,请检查查询参数或数据。")
4 weeks ago
4 weeks ago
prompt = f"""
信息检索与回答助手
根据以下关于'{query}'的相关信息
基本信息
- 语言: 中文
- 描述: 根据提供的材料检索信息并回答问题
- 特点: 快速准确提取关键信息清晰简洁地回答
相关信息
{context}
回答要求
1. 准确无误来源可靠
2. 客观公正避免主观判断
3. 使用HTML格式返回包含适当的段落列表和标题标签
4. 确保内容结构清晰便于前端展示
"""
4 weeks ago
4 weeks ago
response = client.chat.completions.create(
model="deepseek-chat",
messages=[
{"role": "system", "content": "你是一个专业的文档整理助手"},
{"role": "user", "content": prompt}
],
temperature=0.3,
4 weeks ago
stream=False
4 weeks ago
)
4 weeks ago
4 weeks ago
yield {"data": response.choices[0].message.content}
4 weeks ago
except Exception as e:
yield {"data": f"生成报告时出错: {str(e)}"}
4 weeks ago
finally:
# 释放连接
milvus_pool.release_connection(connection)
4 weeks ago
4 weeks ago
"""
4 weeks ago
http://10.10.21.22:8000/static/ai.html
4 weeks ago
小学数学中有哪些模型
4 weeks ago
"""
4 weeks ago
4 weeks ago
class QueryRequest(BaseModel):
query: str = Field(..., description="用户查询的问题")
4 weeks ago
@app.post("/api/rag")
4 weeks ago
async def rag_stream(request: Request):
try:
data = await request.json()
query_request = QueryRequest(**data)
except ValidationError as e:
logger.error(f"请求体验证失败: {e.errors()}")
raise HTTPException(status_code=422, detail=e.errors())
except Exception as e:
logger.error(f"请求解析失败: {str(e)}")
raise HTTPException(status_code=400, detail="无效的请求格式")
4 weeks ago
"""RAG+DeepSeek接口"""
async for chunk in generate_stream(
request.app.state.deepseek_client,
request.app.state.milvus_pool,
request.app.state.collection_manager,
query_request.query
):
return chunk
4 weeks ago
4 weeks ago
4 weeks ago
if __name__ == "__main__":
4 weeks ago
uvicorn.run(app, host="0.0.0.0", port=8000)