516 lines
20 KiB
Python
516 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,Optional
|
||
from tqdm import tqdm # 进度条支持（可选）
|
||
from openai import AsyncOpenAI
|
||
from neo4j_graphrag.embeddings.base import Embedder
|
||
|
||
#now_model = "Qwen3.5-35B-A3B"
|
||
#now_model = "Qwen3-32B"
|
||
|
||
from config import LLM_CONFIG, EMBEDDING_CONFIG
|
||
|
||
|
||
class OpenaiAPI:
|
||
|
||
@staticmethod
|
||
def get_embeddings(embedding_texts: List[str]):
|
||
# ✅ 2. 使用 EMBEDDING_CONFIG
|
||
base_url = EMBEDDING_CONFIG.get("base_url_v2", EMBEDDING_CONFIG["base_url"].replace("/embeddings", ""))
|
||
embedding_base_url = os.path.join(base_url, "embeddings")
|
||
|
||
with httpx.Client(timeout=60) as client:
|
||
response = client.post(
|
||
embedding_base_url,
|
||
json={"model": EMBEDDING_CONFIG["model"], "input": embedding_texts},
|
||
headers={"Authorization": f"Bearer {EMBEDDING_CONFIG['api_key']}"},
|
||
)
|
||
# 注意：batch_embeddings 返回的是列表，get_embeddings 原代码只取了第一个元素的 embedding
|
||
# 这里保持原逻辑，如果输入是列表，通常期望返回所有 embedding，原代码可能有误，此处仅替换配置源
|
||
data = response.json()["data"]
|
||
if isinstance(data, list) and len(data) > 0:
|
||
return data[0]["embedding"]
|
||
return []
|
||
|
||
@staticmethod
|
||
def batch_embeddings(embedding_texts: List[str]):
|
||
# ✅ 3. 使用 EMBEDDING_CONFIG
|
||
base_url = EMBEDDING_CONFIG.get("base_url_v2", EMBEDDING_CONFIG["base_url"].replace("/embeddings", ""))
|
||
embedding_base_url = os.path.join(base_url, "embeddings")
|
||
|
||
with httpx.Client(timeout=60) as client:
|
||
response = client.post(
|
||
embedding_base_url,
|
||
json={"model": EMBEDDING_CONFIG["model"], "input": embedding_texts},
|
||
headers={"Authorization": f"Bearer {EMBEDDING_CONFIG['api_key']}"},
|
||
)
|
||
embedding_list = response.json()["data"]
|
||
return embedding_list
|
||
|
||
@staticmethod
|
||
def rerank_query(query: str, documents: List[str], top_n: str = 3):
|
||
headers = {
|
||
'Content-Type': 'application/json',
|
||
'Authorization': f'Bearer {EMBEDDING_CONFIG["api_key"]}' # 假设 Rerank 也使用相同的 Key 或在此处扩展配置
|
||
}
|
||
|
||
data = {
|
||
"model": os.getenv("OPENAI_RERANKER_MODEL", "bge-reranker-v2-m3"), # Rerank 模型未在 config.py 定义，暂保留环境变量或可添加到 config
|
||
"query": query,
|
||
"top_n": top_n,
|
||
"documents": documents
|
||
}
|
||
|
||
# 假设 Rerank 地址与 Embedding 地址相近，或者需要单独配置。此处暂用环境变量或默认逻辑
|
||
rerank_base = os.getenv("OPENAI_API_RERANKER_BASE", EMBEDDING_CONFIG["base_url_v2"])
|
||
response = requests.post(f"{rerank_base}/rerank", headers=headers, data=json.dumps(data))
|
||
|
||
results_list = []
|
||
if response.status_code == 200:
|
||
result_json = response.json()
|
||
for result in result_json["results"]:
|
||
results_list.append({
|
||
"index": result["index"],
|
||
"text": result["document"]["text"],
|
||
"scores": 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 {EMBEDDING_CONFIG["api_key"]}'
|
||
}
|
||
data = {
|
||
"text_1": text_1,
|
||
"text_2": text_2,
|
||
"model": "bge-rerank"
|
||
}
|
||
rerank_base = os.getenv("OPENAI_API_RERANKER_BASE", EMBEDDING_CONFIG["base_url_v2"])
|
||
response = requests.post(f"{rerank_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 = None):
|
||
# ✅ 4. 使用 LLM_CONFIG
|
||
if model is None:
|
||
model = LLM_CONFIG["model"]
|
||
|
||
response = OpenAI(
|
||
api_key=LLM_CONFIG["api_key"],
|
||
base_url=LLM_CONFIG["base_url"],
|
||
).chat.completions.create(
|
||
model=model,
|
||
messages=[
|
||
{"role": "system", "content": "You are a helpful assistant."},
|
||
{"role": "user", "content": query},
|
||
],
|
||
temperature=LLM_CONFIG["temperature"],
|
||
stream=False
|
||
)
|
||
return response.choices[0].message.content
|
||
|
||
@staticmethod
|
||
async def open_api_chat_async(query: str, model: str = None, temperature: float = None) -> str:
|
||
if model is None:
|
||
model = LLM_CONFIG["model"]
|
||
if temperature is None:
|
||
temperature = LLM_CONFIG["temperature"]
|
||
|
||
client = AsyncOpenAI(
|
||
api_key=LLM_CONFIG["api_key"],
|
||
base_url=LLM_CONFIG["base_url"]
|
||
)
|
||
|
||
response = await client.chat.completions.create(
|
||
model=model,
|
||
messages=[
|
||
{"role": "system", "content": "You are a helpful assistant."},
|
||
{"role": "user", "content": query}
|
||
],
|
||
max_tokens=LLM_CONFIG["max_tokens"],
|
||
temperature=temperature,
|
||
stream=False,
|
||
extra_body={"chat_template_kwargs": {"enable_thinking": False}}
|
||
)
|
||
return response.choices[0].message.content
|
||
|
||
@staticmethod
|
||
async def open_api_chat_async_json(query: str, model: str = None, temperature: float = None) -> str:
|
||
if model is None:
|
||
model = LLM_CONFIG["model"]
|
||
if temperature is None:
|
||
temperature = LLM_CONFIG["temperature"]
|
||
|
||
client = AsyncOpenAI(
|
||
api_key=LLM_CONFIG["api_key"],
|
||
base_url=LLM_CONFIG["base_url"]
|
||
)
|
||
|
||
response = await client.chat.completions.create(
|
||
model=model,
|
||
messages=[
|
||
{"role": "system", "content": "You are a helpful assistant that outputs only valid JSON."},
|
||
{"role": "user", "content": query}
|
||
],
|
||
temperature=temperature,
|
||
stream=False,
|
||
max_tokens=LLM_CONFIG["max_tokens"],
|
||
response_format={"type": "json_object"},
|
||
extra_body={"chat_template_kwargs": {"enable_thinking": False}}
|
||
)
|
||
return response.choices[0].message.content
|
||
|
||
@staticmethod
|
||
def open_api_caht_without_thinking(query: str, model: str = None):
|
||
if model is None:
|
||
model = LLM_CONFIG["model"]
|
||
|
||
response = OpenAI(
|
||
api_key=LLM_CONFIG["api_key"],
|
||
base_url=LLM_CONFIG["base_url"],
|
||
).chat.completions.create(
|
||
model=model,
|
||
messages=[
|
||
{"role": "system", "content": "You are a helpful assistant."},
|
||
{"role": "user", "content": query},
|
||
],
|
||
temperature=LLM_CONFIG["temperature"],
|
||
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 = None):
|
||
if model is None:
|
||
model = LLM_CONFIG["model"]
|
||
|
||
response = OpenAI(
|
||
api_key=LLM_CONFIG["api_key"],
|
||
base_url=LLM_CONFIG["base_url"],
|
||
).chat.completions.create(
|
||
model=model,
|
||
messages=[
|
||
{"role": "system", "content": "You are a helpful assistant."},
|
||
{"role": "user", "content": query},
|
||
],
|
||
temperature=LLM_CONFIG["temperature"],
|
||
stream=True
|
||
)
|
||
for chunk in response:
|
||
content = chunk.choices[0].delta.content
|
||
if content:
|
||
yield content.encode('utf-8')
|
||
|
||
@staticmethod
|
||
def openai_chat(query: str, timeout: int = None):
|
||
if timeout is None:
|
||
timeout = LLM_CONFIG["timeout"]
|
||
|
||
try:
|
||
client = OpenAI(
|
||
api_key=LLM_CONFIG["api_key"],
|
||
base_url=LLM_CONFIG["base_url"],
|
||
timeout=timeout
|
||
)
|
||
|
||
response = client.chat.completions.create(
|
||
model=LLM_CONFIG["model"],
|
||
messages=[
|
||
{"role": "user", "content": query},
|
||
],
|
||
temperature=LLM_CONFIG["temperature"],
|
||
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 = None) -> str | None:
|
||
if timeout is None:
|
||
timeout = LLM_CONFIG["timeout"]
|
||
|
||
try:
|
||
client = AsyncOpenAI(
|
||
api_key=LLM_CONFIG["api_key"],
|
||
base_url=LLM_CONFIG["base_url"],
|
||
timeout=timeout
|
||
)
|
||
|
||
response = await client.chat.completions.create(
|
||
model=LLM_CONFIG["model"],
|
||
messages=[
|
||
{"role": "user", "content": query},
|
||
],
|
||
temperature=LLM_CONFIG["temperature"],
|
||
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
|
||
async def openai_chat_aysnc_nothink(query: str, timeout: int = None) -> str | None:
|
||
if timeout is None:
|
||
timeout = LLM_CONFIG["timeout"]
|
||
|
||
try:
|
||
client = AsyncOpenAI(
|
||
api_key=LLM_CONFIG["api_key"],
|
||
base_url=LLM_CONFIG["base_url"],
|
||
timeout=timeout
|
||
)
|
||
|
||
# 注意：原代码中此方法使用了不同的模型 "46-qwen3.6-35B"
|
||
# 如果需要统一，请使用 LLM_CONFIG["model"]，否则需单独配置
|
||
model_to_use = LLM_CONFIG["model"]
|
||
|
||
response = await client.chat.completions.create(
|
||
model=model_to_use,
|
||
messages=[
|
||
{"role": "user", "content": query},
|
||
],
|
||
temperature=LLM_CONFIG["temperature"],
|
||
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
|
||
finally:
|
||
await client.close()
|
||
|
||
@staticmethod
|
||
def generate_embedding(text: str) -> Optional[List[float]]:
|
||
try:
|
||
client = OpenAI(
|
||
api_key=EMBEDDING_CONFIG["api_key"],
|
||
base_url=EMBEDDING_CONFIG.get("base_url_v2", EMBEDDING_CONFIG["base_url"].replace("/embeddings", "")),
|
||
)
|
||
response = client.embeddings.create(
|
||
model=EMBEDDING_CONFIG["model"],
|
||
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"""
|
||
<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
|
||
class LocalBgeM3Embeddings(Embedder):
|
||
def __init__(self, base_url: str, api_key: str, model_name: str = "bge-m3"):
|
||
self.base_url = base_url.rstrip("/")
|
||
self.api_key = api_key
|
||
self.model_name = model_name
|
||
|
||
def embed_query(self, text: str) -> List[float]:
|
||
return self.embed_documents([text])[0]
|
||
|
||
def embed_documents(self, texts: List[str]) -> List[List[float]]:
|
||
embedding_base_url = os.path.join(self.base_url, "embeddings")
|
||
with httpx.Client(timeout=60) as client:
|
||
response = client.post(
|
||
embedding_base_url,
|
||
json={"model": self.model_name, "input": texts},
|
||
headers={"Authorization": f"Bearer {self.api_key}"},
|
||
)
|
||
embedding_list = [item["embedding"] for item in response.json()["data"]]
|
||
return embedding_list
|
||
|
||
|
||
|