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.
100 lines
2.9 KiB
100 lines
2.9 KiB
'''
|
|
pip install openai
|
|
'''
|
|
from elasticsearch import Elasticsearch
|
|
from Util.EmbeddingUtil import text_to_embedding
|
|
import Config.Config as config
|
|
from openai import OpenAI
|
|
import json
|
|
|
|
# 初始化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)
|
|
|
|
def search_related_data(query):
|
|
"""搜索向量数据和原始相关数据"""
|
|
# 向量搜索
|
|
vector = text_to_embedding(query)
|
|
vector_results = es.search(
|
|
index='knowledge_base',
|
|
body={
|
|
"size": 5,
|
|
"query": {
|
|
"script_score": {
|
|
"query": {"match_all": {}},
|
|
"script": {
|
|
"source": "cosineSimilarity(params.query_vector, 'vector') + 1.0",
|
|
"params": {"query_vector": vector}
|
|
}
|
|
}
|
|
},
|
|
"_source": ["text"]
|
|
}
|
|
)
|
|
|
|
# 文本精确搜索
|
|
text_results = es.search(
|
|
index='raw_texts',
|
|
body={
|
|
"query": {
|
|
"multi_match": {
|
|
"query": query,
|
|
"fields": ["text"],
|
|
"type": "best_fields"
|
|
}
|
|
},
|
|
"size": 5
|
|
}
|
|
)
|
|
|
|
# 合并结果
|
|
all_results = [hit['_source']['text'] for hit in vector_results['hits']['hits']]
|
|
all_results.extend([hit['_source']['text'] for hit in text_results['hits']['hits']])
|
|
|
|
return "\n\n".join(all_results)
|
|
|
|
def generate_report(query, context):
|
|
"""使用DeepSeek生成报告"""
|
|
prompt = f"""根据以下关于'{query}'的相关信息，整理一份结构化的报告：
|
|
要求：
|
|
1. 分章节组织内容
|
|
2. 包含关键数据和事实
|
|
3. 语言简洁专业
|
|
|
|
相关信息：
|
|
{context}"""
|
|
|
|
response = client.chat.completions.create(
|
|
model="deepseek-chat",
|
|
messages=[
|
|
{"role": "system", "content": "你是一个专业的文档整理助手"},
|
|
{"role": "user", "content": prompt}
|
|
],
|
|
temperature=0.3
|
|
)
|
|
|
|
return response.choices[0].message.content
|
|
|
|
def process_query(query):
|
|
"""处理用户查询并生成报告"""
|
|
print(f"正在搜索与'{query}'相关的数据...")
|
|
context = search_related_data(query)
|
|
print(f"找到{len(context.split(chr(10)+chr(10)))}条相关数据")
|
|
|
|
print("正在生成报告...")
|
|
report = generate_report(query, context)
|
|
|
|
return report
|
|
|
|
if __name__ == "__main__":
|
|
#user_query = input("请输入您的查询要求：")
|
|
user_query = "整理云南省初中在校生情况文档"
|
|
report = process_query(user_query)
|
|
print("\n=== 生成的报告 ===\n")
|
|
print(report) |