import httpx
from openai import OpenAI
import os
import requests
import json
import asyncio
import aiohttp
from typing import List, Dict, Any,Optional
from tqdm import tqdm # 进度条支持（可选）
from openai import AsyncOpenAI
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
async def open_api_chat_async(query: str, model: str = None, temperature: float = 0):
"""
async Chat API(graph search)
Args:
query:
model: (OPENAI_MODEL)
temperature:
Returns:
str: LLM return
"""
if model is None:
model = os.getenv("OPENAI_MODEL")
api_key = os.getenv("OPENAI_API_KEY")
base_url = os.getenv("OPENAI_API_BASE")
if not api_key or not base_url:
raise ValueError("Missing OpenAI API configuration")
async with aiohttp.ClientSession() as session:
async with session.post(
f"{base_url}/chat/completions",
headers={
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json"
},
json={
"model": model,
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": query}
],
"temperature": temperature,
"stream": False
},
timeout=aiohttp.ClientTimeout(total=60)
) as response:
if response.status == 200:
data = await response.json()
return data["choices"][0]["message"]["content"]
else:
error_text = await response.text()
raise Exception(f"API")
@staticmethod
async def open_api_chat_async_json(query: str, model: str = None, temperature: float = 0.6):
"""
async Chat API (graph search) - 强制 JSON 输出模式（若后端支持）
Args:
query: 用户提示（应包含明确的 JSON 输出指令）
model: 模型名称
temperature: 温度参数
Returns:
str: LLM 返回的纯 JSON 字符串（理想情况下）
"""
if model is None:
model = os.getenv("OPENAI_MODEL")
api_key = os.getenv("OPENAI_API_KEY")
base_url = os.getenv("OPENAI_API_BASE")
if not api_key or not base_url:
raise ValueError("Missing OpenAI API configuration")
# 构造请求体
payload = {
"model": model,
"messages": [
{"role": "system", "content": "You are a helpful assistant that outputs only valid JSON."},
{"role": "user", "content": query}
],
"temperature": temperature,
"stream": False,
"response_format": {"type": "json_object"} # ←←← 关键：强制 JSON 模式
}
# 某些后端（如旧版）可能不支持 extra_body，可移除
# 如果你确定后端是 DashScope/Qwen，可以保留或删除 extra_body
# payload["extra_body"] = {"chat_template_kwargs": {"enable_thinking": False}}
async with aiohttp.ClientSession() as session:
async with session.post(
f"{base_url}/chat/completions",
headers={
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json"
},
json=payload,
timeout=aiohttp.ClientTimeout(total=60)
) as response:
if response.status == 200:
data = await response.json()
return data["choices"][0]["message"]["content"]
else:
error_text = await response.text()
raise Exception(f"API调用失败: {response.status} - {error_text}")
@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
@staticmethod
def openai_chat(query:str,timeout:int=60):
"""调用OpenAI API"""
try:
base = os.getenv("OPENAI_API_BASE","http://192.168.0.46:59800/v1")
client = OpenAI(
api_key=os.getenv("OPENAI_API_KEY","none"),
base_url=os.getenv("OPENAI_API_BASE","http://192.168.0.46:59800/v1"),
timeout=timeout # 👈 设置整体请求超时（秒）
)
response = client.chat.completions.create(
model="Qwen3.5-35B-A3B",
messages=[
{"role": "user", "content": query},
],
temperature=0.1,
stream=False,
response_format = {"type": "json_object"},
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
@staticmethod
async def openai_chat_aysnc(query: str, timeout: int = 180) -> str | None:
"""异步调用 OpenAI 兼容 API（如本地部署的 Qwen、vLLM、Ollama 等）"""
try:
base_url = os.getenv("OPENAI_API_BASE", "http://192.168.0.46:59800/v1")
api_key = os.getenv("OPENAI_API_KEY", "none")
client = AsyncOpenAI(
api_key=api_key,
base_url=base_url,
timeout=timeout # 整体请求超时（秒）
)
response = await client.chat.completions.create(
model="Qwen3.5-35B-A3B",
messages=[
{"role": "user", "content": query},
],
temperature=0.1,
stream=False,
response_format={"type": "json_object"},
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
finally:
# 可选：显式关闭客户端（通常不需要，但可避免警告）
await client.close()
@staticmethod
def generate_embedding(text: str) -> Optional[List[float]]:
"""调用OpenAI生成文本embedding"""
try:
client = OpenAI(
api_key=os.getenv("OPENAI_API_KEY", "EMPTY"),
base_url=os.getenv("OPENAI_API_EMBEDDING_BASE", "http://192.168.0.46/v1"),
)
response = client.embeddings.create(
model="bge-m3", # 或你部署的embedding模型名
input=text
)
return response.data[0].embedding
except Exception as e:
print(f"生成embedding失败: {e}")
return None
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.5-35b")
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"""
{markdown_context}
{text}
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 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 `` 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 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 , the response should include the citation like so:
* "According to the study, the proposed method increases efficiency by 20% [whitepaper.pdf]."
If no 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 tag is present in the context.
{CONTEXT}
{QUERY}
"""
return PROMPT
# for chunk in response:
# content = chunk.choices[0].delta.content
# if content:
# print(content, end='', flush=True)
# print('\n')
# return response