70 lines
2.6 KiB
Python
70 lines
2.6 KiB
Python
import json
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import List, Dict
|
||
|
||
from langchain_core.output_parsers import JsonOutputParser
|
||
from langchain_core.prompts import PromptTemplate
|
||
from langchain_huggingface import HuggingFaceEndpointEmbeddings
|
||
from langgraph.runtime import Runtime
|
||
from sentry_sdk.integrations import aiohttp
|
||
|
||
from app.chat_qa.context import QAContext
|
||
from app.chat_qa.qa_state import QAState
|
||
from app.client.embedding_client_manager import embedding_client
|
||
from app.llm import llm
|
||
from app.repository.milvus.message_qa_repository import QARepository
|
||
from app.core.log import logger
|
||
|
||
"""
|
||
知识库入库
|
||
"""
|
||
async def finalize(state: QAState, runtime: Runtime[QAContext]):
|
||
|
||
final_qas: List[Dict] = state.get("final_qas", [])
|
||
chat_history: str = state["chat_history"]
|
||
|
||
logger.info(f"原对胡数量：{chat_history}最终生成的QA数据：{len(final_qas)}")
|
||
|
||
if len(final_qas)==0:
|
||
logger.info(f"生成QA失败，原对胡数量：{chat_history}最终生成的QA数据：{len(final_qas)}")
|
||
return
|
||
|
||
output_dir = Path("E:/qa_output")
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
output_file = output_dir / "qa_result.jsonl" # 改用 .jsonl 后缀
|
||
|
||
try:
|
||
with open(output_file, 'a', encoding='utf-8') as f:
|
||
if isinstance(final_qas, list):
|
||
for qa in final_qas:
|
||
f.write(json.dumps(qa, ensure_ascii=False) + '\n')
|
||
else:
|
||
f.write(json.dumps(final_qas, ensure_ascii=False) + '\n')
|
||
|
||
logger.info(f"成功追加 {len(final_qas) if isinstance(final_qas, list) else 1} 条QA数据到文件: {output_file}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"写入QA文件失败: {e}")
|
||
|
||
# question_list: List[str] = [qa["question"] for qa in final_qas]
|
||
#
|
||
# embedding_batch_size = 10
|
||
# embedding_ed_all: list[list[float]] = []
|
||
# for i in range(0, len(question_list), embedding_batch_size):
|
||
# embedding_part_texts = question_list[i:i + embedding_batch_size]
|
||
# embedding_eds = await embedding_client.aembed_documents(embedding_part_texts)
|
||
# embedding_ed_all.extend(embedding_eds)
|
||
#
|
||
# # 构建对象列表
|
||
# qa_list = []
|
||
# for i, qa in enumerate(final_qas):
|
||
# qa_list.append({
|
||
# "question": qa["question"],
|
||
# "answer": qa.get("answer", ""),
|
||
# "question_dense_vector": embedding_ed_all[i] if i < len(embedding_ed_all) else [],
|
||
# "parent_id": qa.get("parent_id", 0)
|
||
# })
|
||
# await qa_milvus_repository.save_qa_to_milvus(qa_list)
|
||
|