526 lines
20 KiB
Python
526 lines
20 KiB
Python
import httpx
|
||
from openai import OpenAI
|
||
import os
|
||
import requests
|
||
import json
|
||
import asyncio
|
||
import aiohttp
|
||
from typing import List, Dict, Any
|
||
from tqdm import tqdm # 进度条支持（可选）
|
||
import re
|
||
|
||
class OpenaiAPI:
|
||
# def __init__(self):
|
||
# self.openai_embedding_api_base = os.getenv("OPENAI_API_EMBEDDING_BASE")
|
||
# self.openai_embedding_model = os.getenv("OPENAI_EMBEDDING_MODEL")
|
||
# self.openai_chat_api_base = os.getenv("OPENAI_API_BASE")
|
||
# self.openai_chat_model = os.getenv("OPENAI_MODEL")
|
||
# self.openai_api_key = os.getenv("OPENAI_API_KEY")
|
||
# self.openai_reranker_api_base = os.getenv("OPENAI_API_RERANKER_BASE")
|
||
# self.openai_reranker_model = os.getenv("OPENAI_RERANKER_MODEL")
|
||
|
||
# self.embedding_client = OpenAI(
|
||
# api_key=self.openai_api_key,
|
||
# base_url=self.openai_embedding_api_base,
|
||
# )
|
||
# self.chat_client = OpenAI(
|
||
# api_key=self.openai_api_key,
|
||
# base_url=self.openai_chat_api_base,
|
||
# )
|
||
# self.rerank_client = OpenAI(
|
||
# api_key=self.openai_api_key,
|
||
# base_url=self.openai_reranker_api_base,
|
||
# )
|
||
@staticmethod
|
||
def get_embeddings(embedding_texts: List[str]):
|
||
base_url = os.getenv("OPENAI_API_EMBEDDING_BASE")
|
||
embedding_base_url = os.path.join(base_url, "embeddings")
|
||
with httpx.Client(timeout=60) as client:
|
||
response = client.post(
|
||
embedding_base_url,
|
||
json={"model": "bge-m3", "input": embedding_texts},
|
||
headers={"Authorization": os.getenv("OPENAI_API_KEY")},
|
||
)
|
||
embedding_list = response.json()["data"][0]["embedding"]
|
||
# response = OpenAI(
|
||
# api_key=os.getenv("OPENAI_API_KEY","EMPTY"),
|
||
# base_url=os.getenv("OPENAI_API_EMBEDDING_BASE","http://172.18.30.122:40052/v1"),
|
||
# ).embeddings.create(
|
||
# model=os.getenv("OPENAI_EMBEDDING_MODEL","bge-m3"),
|
||
# input=embedding_texts,
|
||
# )
|
||
# embedding_list = response.data[0].embedding
|
||
return embedding_list
|
||
|
||
@staticmethod
|
||
def batch_embeddings(embedding_texts: List[str]):
|
||
base_url = os.getenv("OPENAI_API_EMBEDDING_BASE")
|
||
embedding_base_url = os.path.join(base_url, "embeddings")
|
||
with httpx.Client(timeout=60) as client:
|
||
response = client.post(
|
||
embedding_base_url,
|
||
json={"model": "bge-m3", "input": embedding_texts},
|
||
headers={"Authorization": os.getenv("OPENAI_API_KEY")},
|
||
)
|
||
embedding_list = response.json()["data"]
|
||
|
||
return embedding_list
|
||
|
||
@staticmethod
|
||
def rerank_query(query: str, documents: List[str], top_n: str = 3):
|
||
"""
|
||
发送重排序请求到 GPUStack API。
|
||
|
||
:param query: 查询字符串
|
||
:param documents: 文档列表（字符串列表）
|
||
:param top_n: 返回的文档数量
|
||
:param model: 使用的模型名称
|
||
:param api_key: API 密钥
|
||
:param server_url: API 服务器地址
|
||
:return: API 的响应结果（JSON 格式）
|
||
|
||
return :
|
||
{
|
||
"id": "rerank-1446b27f0d8642288578ee26707253a1",
|
||
"model": "bge-reranker-v2-m3",
|
||
"usage": {
|
||
"total_tokens": 53
|
||
},
|
||
"results": [
|
||
{
|
||
"index": 2,
|
||
"document": {
|
||
"text": "I have a dog named Dollor"
|
||
},
|
||
"relevance_score": 0.297119140625
|
||
},
|
||
{
|
||
"index": 1,
|
||
"document": {
|
||
"text": "hi"
|
||
},
|
||
"relevance_score": 0.0009255409240722656
|
||
},
|
||
{
|
||
"index": 0,
|
||
"document": {
|
||
"text": "what is panda?"
|
||
},
|
||
"relevance_score": 0.0008897781372070312
|
||
}
|
||
]
|
||
}
|
||
"""
|
||
|
||
# 请求头
|
||
headers = {
|
||
'Content-Type': 'application/json',
|
||
'Authorization': f'Bearer {os.getenv("OPENAI_API_KEY")}'
|
||
}
|
||
|
||
# 请求体
|
||
data = {
|
||
"model": os.getenv("OPENAI_RERANKER_MODEL"),
|
||
"query": query,
|
||
"top_n": top_n,
|
||
"documents": documents
|
||
}
|
||
|
||
# 发送 POST 请求
|
||
response = requests.post(os.getenv("OPENAI_API_RERANKER_BASE") + "/rerank", headers=headers,
|
||
data=json.dumps(data))
|
||
results_list = []
|
||
|
||
# 检查响应状态码
|
||
if response.status_code == 200:
|
||
result_json = response.json() # 返回解析后的 JSON 数据
|
||
for result in result_json["results"]:
|
||
results_list.append(
|
||
{"index": result["index"], "text": result["document"]["text"], "scores": result["relevance_score"]})
|
||
# results_list.append(result["document"]["text"])
|
||
# relevance_scores.append(result["relevance_score"])
|
||
return results_list
|
||
|
||
else:
|
||
raise Exception(f"Error: {response.status_code}\n{response.text}") # 抛出异常
|
||
|
||
@staticmethod
|
||
def embedding_cosine_similarity(text_1, text_2):
|
||
# 请求头
|
||
headers = {
|
||
'Content-Type': 'application/json',
|
||
'Authorization': f'Bearer {os.getenv("OPENAI_API_KEY")}'
|
||
}
|
||
|
||
# 请求体
|
||
data = {
|
||
"text_1": text_1,
|
||
"text_2": text_2,
|
||
"model": "bge-rerank"
|
||
}
|
||
|
||
# # 发送 POST 请求
|
||
# response = requests.post(os.getenv("OPENAI_API_EMBEDDING_BASE") + "/score", headers=headers, data=json.dumps(data))
|
||
response = requests.post(os.getenv("OPENAI_API_RERANKER_BASE") + "/score", headers=headers,
|
||
data=json.dumps(data))
|
||
if response.status_code == 200:
|
||
result_json = response.json()
|
||
score = result_json["data"][0]["score"]
|
||
return score
|
||
else:
|
||
raise Exception(f"Error: {response.status_code}\n{response.text}") # 抛出异常
|
||
|
||
@staticmethod
|
||
def open_api_chat(query: str, model: str):
|
||
response = OpenAI(
|
||
api_key=os.getenv("OPENAI_API_KEY"),
|
||
base_url=os.getenv("OPENAI_API_BASE"),
|
||
).chat.completions.create(
|
||
model=model,
|
||
messages=[
|
||
{"role": "system", "content": "You are a helpful assistant."},
|
||
{"role": "user", "content": query},
|
||
],
|
||
temperature=0.6,
|
||
stream=False
|
||
)
|
||
return response.choices[0].message.content
|
||
|
||
@staticmethod
|
||
def open_api_caht_without_thinking(query: str, model: str):
|
||
response = OpenAI(
|
||
api_key=os.getenv("OPENAI_API_KEY"),
|
||
base_url=os.getenv("OPENAI_API_BASE"),
|
||
).chat.completions.create(
|
||
model=model,
|
||
messages=[
|
||
{"role": "system", "content": "You are a helpful assistant."},
|
||
{"role": "user", "content": query}, ],
|
||
|
||
temperature=0.1,
|
||
stream=False,
|
||
extra_body={"chat_template_kwargs": {"enable_thinking": False}})
|
||
return response.choices[0].message.content
|
||
|
||
@staticmethod
|
||
def open_api_chat_stream(query: str, model: str):
|
||
response = OpenAI(
|
||
api_key=os.getenv("OPENAI_API_KEY"),
|
||
base_url=os.getenv("OPENAI_API_BASE"),
|
||
).chat.completions.create(
|
||
model=model,
|
||
messages=[
|
||
{"role": "system", "content": "You are a helpful assistant."},
|
||
{"role": "user", "content": query},
|
||
],
|
||
temperature=0.6,
|
||
stream=True
|
||
)
|
||
for chunk in response:
|
||
content = chunk.choices[0].delta.content
|
||
if content:
|
||
yield content.encode('utf-8')
|
||
|
||
@staticmethod
|
||
def context_chunk_process(markdown_context: str, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||
"""同步调用入口"""
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
try:
|
||
result = loop.run_until_complete(
|
||
batch_chunk_process(markdown_context, chunks)
|
||
)
|
||
finally:
|
||
loop.close()
|
||
return result
|
||
|
||
|
||
async def batch_chunk_process(
|
||
markdown_context: str,
|
||
chunks: List[Dict[str, Any]],
|
||
batch_size: int = 5,
|
||
timeout: int = 30
|
||
) -> List[Dict[str, Any]]:
|
||
"""
|
||
批量处理文档块（异步并发）
|
||
|
||
Args:
|
||
markdown_context: 完整文档内容
|
||
chunks: 文档块列表，每个块需包含"text"键
|
||
batch_size: 并发请求数 (默认5)
|
||
timeout: 单个请求超时时间(秒)
|
||
|
||
Returns:
|
||
处理后的文档块列表
|
||
"""
|
||
# 初始化参数
|
||
api_key = os.getenv("OPENAI_API_KEY")
|
||
base_url = os.getenv("OPENAI_API_BASE")
|
||
model = os.getenv("OPENAI_MODEL", "qwen3-32b")
|
||
|
||
if not api_key or not base_url:
|
||
raise ValueError("Missing OpenAI API configuration")
|
||
|
||
async def process_single_chunk(
|
||
session: aiohttp.ClientSession,
|
||
chunk: Dict[str, Any],
|
||
semaphore: asyncio.Semaphore
|
||
) -> Dict[str, Any]:
|
||
"""处理单个chunk的异步函数"""
|
||
async with semaphore:
|
||
text = chunk.get("text", "").strip()
|
||
if not text:
|
||
return chunk
|
||
|
||
prompt = f"""
|
||
<document>{markdown_context}</document>
|
||
<chunk>{text}</chunk>
|
||
Please provide a succinct context for this chunk within the document.
|
||
Respond ONLY with the context text, no additional commentary.
|
||
"""
|
||
chinese_prompt = f"""
|
||
<文档>
|
||
{markdown_context}
|
||
</文档>
|
||
以下是需要定位到整体文档中的片段
|
||
<片段>
|
||
{text}
|
||
</片段>
|
||
请简要说明该片段在整体文档中的上下文关系，以提升片段搜索效果。只需给出简洁的上下文说明，无需其他内容。
|
||
"""
|
||
|
||
try:
|
||
async with session.post(
|
||
f"{base_url}/chat/completions",
|
||
headers={"Authorization": f"Bearer {api_key}"},
|
||
json={
|
||
"model": model,
|
||
"messages": [
|
||
{"role": "system", "content": "You are a helpful assistant."},
|
||
{"role": "user", "content": chinese_prompt}
|
||
],
|
||
"temperature": 0.1,
|
||
"extra_body": {"chat_template_kwargs": {"enable_thinking": False}}
|
||
},
|
||
timeout=timeout
|
||
) as response:
|
||
if response.status == 200:
|
||
data = await response.json()
|
||
if data.get("choices"):
|
||
chunk["text"] = data["choices"][0]["message"]["content"].strip()
|
||
else:
|
||
print(f"API error for chunk: {await response.text()}")
|
||
except Exception as e:
|
||
print(f"Error processing chunk: {e}")
|
||
|
||
return chunk
|
||
|
||
# 创建信号量控制并发数
|
||
semaphore = asyncio.Semaphore(batch_size)
|
||
|
||
# 异步处理所有chunks
|
||
async with aiohttp.ClientSession() as session:
|
||
tasks = [
|
||
process_single_chunk(session, chunk, semaphore)
|
||
for chunk in chunks
|
||
]
|
||
|
||
# 使用tqdm显示进度（可选）
|
||
processed_chunks = []
|
||
for future in tqdm(asyncio.as_completed(tasks), total=len(tasks)):
|
||
processed_chunks.append(await future)
|
||
|
||
return processed_chunks
|
||
|
||
|
||
def prepare_rag_prompt(query, documents):
|
||
docs = ""
|
||
for i, doc in enumerate(documents):
|
||
docs += f"文本信息{i + 1}：{doc}\n"
|
||
return f"""你作为一个文档检索问答专家，请你从以下文本信息中选取有用的内容作为参考，仔细思考后用中文回答用户的问题，如果文本中没有关于问题的答案，则回答不知道:
|
||
-----
|
||
{docs}
|
||
-----
|
||
用户提问：{query}
|
||
你的回答："""
|
||
|
||
|
||
def query_rewrite(query):
|
||
prompt = f"""
|
||
你是一个专业的搜索查询优化助手。你的任务是根据用户输入的原始查询，生成更清晰、更易被搜索引擎理解的改写版本。请遵循以下规则：
|
||
|
||
1. **保持原意**：改写后的查询必须与用户意图一致，不能曲解原意。
|
||
2. **简洁清晰**：去掉冗余词，优化表达，但不要过度缩写。
|
||
3. **结构化**：如果查询适合拆解，可以适当结构化（如添加隐含的限定条件）。
|
||
4. **自然语言**：确保改写后的查询仍然符合自然语言习惯。
|
||
5. **多语言支持**：如果查询是中文，优先用中文改写；如果是英文，用英文优化。
|
||
|
||
请改写以下查询，直接返回优化后的内容，不要额外解释：
|
||
【用户查询开始】
|
||
{query}
|
||
【用户查询结束】
|
||
"""
|
||
return OpenaiAPI.open_api_chat(prompt)
|
||
|
||
|
||
def RAG_PROMPT(QUERY, CONTEXT):
|
||
PROMPT = f"""
|
||
### Task:
|
||
Respond to the user query using the provided context, incorporating inline citations in the format [source_id] **only when the <source_id> tag is explicitly provided** in the context.
|
||
|
||
### Guidelines:
|
||
- If you don't know the answer, clearly state that.
|
||
- If uncertain, ask the user for clarification.
|
||
- Respond in the same language as the user's query.
|
||
- If the context is unreadable or of poor quality, inform the user and provide the best possible answer.
|
||
- If the answer isn't present in the context but you possess the knowledge, explain this to the user and provide the answer using your own understanding.
|
||
- When you find corresponding content in the provided sources, try to quote it verbatim as much as possible.
|
||
- When quoting from sources, reproduce the original text verbatim—do not omit, paraphrase, or add content.
|
||
- **Only include inline citations using [source_id] (e.g., [1], [2]) when a `<source_id>` tag is explicitly provided in the context.**
|
||
- If the source file has a .glb extension, directly mark the corresponding content with the source_id in the relevant position.
|
||
- Do not cite if the <source_id> or tag is not provided in the context.
|
||
- Do not use XML tags in your response.
|
||
- **If none of the provided sources contain relevant information to answer the question, respond only with: "没有找到相关来源" (No relevant sources found). Do not add any other content.
|
||
- Ensure citations are concise and directly related to the information provided.
|
||
|
||
### Example of Citation:
|
||
If the user asks about a specific topic and the information is found in "whitepaper.pdf" with a provided <source_id> , the response should include the citation like so:
|
||
* "According to the study, the proposed method increases efficiency by 20% [whitepaper.pdf]."
|
||
If no <source_id> is present, the response should omit the citation.
|
||
|
||
### Output:
|
||
Provide a clear and direct response to the user's query, including inline citations in the format [source_id] only when the <source_id> tag is present in the context.
|
||
|
||
<context>
|
||
{CONTEXT}
|
||
</context>
|
||
|
||
<user_query>
|
||
{QUERY}
|
||
</user_query>
|
||
"""
|
||
return PROMPT
|
||
|
||
# for chunk in response:
|
||
# content = chunk.choices[0].delta.content
|
||
# if content:
|
||
# print(content, end='', flush=True)
|
||
# print('\n')
|
||
# return response
|
||
def split_by_period(text):
|
||
sentences = re.split('(?<=。)', text)
|
||
sentences = [sentence.strip() for sentence in sentences if sentence.strip()]
|
||
return sentences
|
||
|
||
|
||
def rerank_text(query,data):
|
||
for datainfo in data:
|
||
splitdata = split_by_period(datainfo['text'])
|
||
current_score = datainfo['score']
|
||
text_1 = query
|
||
url = os.getenv("OPENAI_API_RERANKER_BASE") + "/score"
|
||
headers = {
|
||
'Authorization': 'Bearer 123456',
|
||
'Content-Type': 'application/json'
|
||
}
|
||
|
||
parse_data = {
|
||
"model": "bge-rerank",
|
||
"text_1": text_1,
|
||
"text_2": splitdata
|
||
}
|
||
response = requests.post(url, headers=headers, json=parse_data)
|
||
parseresult = response.json()
|
||
scores = [item['score'] for item in parseresult['data'] if 'score' in item]
|
||
|
||
max_score = max(scores)
|
||
socore = 0.1 * current_score + 0.9 * max_score
|
||
datainfo['score'] = socore
|
||
return data
|
||
|
||
def rerank_3d(query,data):
|
||
for recall_result in data:
|
||
text_2 = recall_result["text"]
|
||
current_score = recall_result["score"]
|
||
text_1 = query
|
||
url = os.getenv("OPENAI_API_RERANKER_BASE") + "/score"
|
||
headers = {
|
||
'Authorization': 'Bearer 123456',
|
||
'Content-Type': 'application/json'
|
||
}
|
||
|
||
parse_data = {
|
||
"model": "bge-rerank",
|
||
"text_1": text_1,
|
||
"text_2": text_2
|
||
}
|
||
response = requests.post(url, headers=headers, json=parse_data)
|
||
parseresult = response.json()
|
||
rerank_score = parseresult['data'][0]['score']
|
||
score = current_score * 0.2 + rerank_score * 0.8
|
||
recall_result["score"] = score
|
||
result_dict = {}
|
||
for item in data:
|
||
resource = item['resource']
|
||
# 如果该 resource 不存在，或当前 score 更高，则替换
|
||
if resource not in result_dict or item['score'] > result_dict[resource]['score']:
|
||
result_dict[resource] = item
|
||
|
||
# 转换为列表
|
||
filtered_result = list(result_dict.values())
|
||
for new_index, item in enumerate(filtered_result, start=1):
|
||
item['index'] = new_index
|
||
|
||
return filtered_result
|
||
|
||
|
||
class OPENAIAPI:
|
||
@staticmethod
|
||
def openai_chat(url,query, model):
|
||
"""调用OpenAI API"""
|
||
try:
|
||
client = OpenAI(
|
||
api_key="none",
|
||
base_url=url,
|
||
)
|
||
|
||
response = client.chat.completions.create(
|
||
model=model,
|
||
messages=[
|
||
{"role": "system",
|
||
"content": "你是一个专业的武器类查询改写助手，将用户问题改写为一个容易理解、且明确的查询。"},
|
||
{"role": "user", "content": query},
|
||
],
|
||
temperature=0.1,
|
||
stream=False,
|
||
extra_body={"chat_template_kwargs": {"enable_thinking": False}}
|
||
)
|
||
return response.choices[0].message.content
|
||
except Exception as e:
|
||
print(f"调用OpenAI API时出错: {e}")
|
||
return None
|
||
prompt = """
|
||
你是一个查询重写工具。请将以下用户查询改写为一个清晰、具体、完整、无歧义的自然语言问题。
|
||
要求：
|
||
- 保留原始意图，不添加未提及的假设；
|
||
- 若查询模糊（如缩写、代词、片段），基于常识明确实体和领域（例如“AR50”应理解为“AR-50狙击步枪”）；
|
||
- 使用完整陈述句，聚焦信息需求（如定义等）；
|
||
- **仅输出改写后的问题本身，不要任何其他文字、标点前缀、问候、反问或说明。**
|
||
- 不要进行任何思考，直接输出改写后查询
|
||
- 请按照下述示例进行改写
|
||
原始查询：ar50
|
||
改写后查询：AR-50狙击步枪是什么？
|
||
|
||
原始查询：变速器？
|
||
改写后查询：变速器是什么？
|
||
|
||
原始查询：{user_query}
|
||
"""
|
||
def resetQuery(query:str):
|
||
url = os.getenv("OPENAI_API_MODEL_BASE"+"/v1","http://192.168.0.46:59800/v1")
|
||
model = os.getenv("MODEL_NAME","Qwen3-32B")
|
||
query_prompt = prompt.replace("{user_query}",query)
|
||
answer = OPENAIAPI.openai_chat(url,query_prompt, model)
|
||
return answer
|
||
|