feat: 测试prompt
This commit is contained in:
parent
eb73101ddb
commit
2e2082808a
|
|
@ -22,9 +22,6 @@ class Config(TypedDict):
|
||||||
es_port: int
|
es_port: int
|
||||||
es_user: str
|
es_user: str
|
||||||
es_password: str
|
es_password: str
|
||||||
qdrant_host: str
|
|
||||||
qdrant_port: int
|
|
||||||
qdrant_api_key: str
|
|
||||||
|
|
||||||
|
|
||||||
def _read_config() -> Config:
|
def _read_config() -> Config:
|
||||||
|
|
@ -44,9 +41,6 @@ def _read_config() -> Config:
|
||||||
"es_port": int(os.getenv("ES_PORT")),
|
"es_port": int(os.getenv("ES_PORT")),
|
||||||
"es_user": os.getenv("ES_USER") or "elastic",
|
"es_user": os.getenv("ES_USER") or "elastic",
|
||||||
"es_password": os.getenv("ES_PASSWORD") or "",
|
"es_password": os.getenv("ES_PASSWORD") or "",
|
||||||
"qdrant_host": os.getenv("QDRANT_HOST"),
|
|
||||||
"qdrant_port": int(os.getenv("QDRANT_PORT")),
|
|
||||||
"qdrant_api_key": os.getenv("QDRANT_API_KEY") or "",
|
|
||||||
}
|
}
|
||||||
|
|
||||||
config = _read_config()
|
config = _read_config()
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,5 @@
|
||||||
import httpx
|
import httpx
|
||||||
|
import json
|
||||||
from src.pipeline.config import config
|
from src.pipeline.config import config
|
||||||
from src.pipeline.core.utils import logger
|
from src.pipeline.core.utils import logger
|
||||||
|
|
||||||
|
|
@ -7,13 +8,30 @@ class AsyncLLm:
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
timeout: float = 30.0,
|
timeout: float = 30.0,
|
||||||
max_connections: int = 100,
|
max_connections: int = 50,
|
||||||
max_keepalive: int = 20,
|
max_keepalive: int = 10,
|
||||||
):
|
):
|
||||||
|
# chat
|
||||||
|
self.chat_api = config["llm_api_host"].rstrip("/") + "/chat/completions"
|
||||||
|
self.chat_model = config["llm_model"]
|
||||||
|
self.chat_api_key = config["llm_api_key"]
|
||||||
|
self.chat_client = httpx.AsyncClient(
|
||||||
|
http2=False,
|
||||||
|
trust_env=False,
|
||||||
|
timeout=httpx.Timeout(timeout),
|
||||||
|
limits=httpx.Limits(
|
||||||
|
max_connections=max_connections,
|
||||||
|
max_keepalive_connections=max_keepalive,
|
||||||
|
),
|
||||||
|
headers={
|
||||||
|
"Content-Type": "application/json",
|
||||||
|
"Authorization": f"Bearer {self.chat_api_key}",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
# embedding
|
||||||
self.embedding_api = config["embedding_api_host"].rstrip("/") + "/embeddings"
|
self.embedding_api = config["embedding_api_host"].rstrip("/") + "/embeddings"
|
||||||
self.embedding_model = config["embedding_model"]
|
self.embedding_model = config["embedding_model"]
|
||||||
self.api_key = config["embedding_api_key"]
|
self.embedding_api_key = config["embedding_api_key"]
|
||||||
logger.debug(self.embedding_api)
|
|
||||||
self.embedding_client = httpx.AsyncClient(
|
self.embedding_client = httpx.AsyncClient(
|
||||||
http2=False,
|
http2=False,
|
||||||
trust_env=False,
|
trust_env=False,
|
||||||
|
|
@ -24,7 +42,7 @@ class AsyncLLm:
|
||||||
),
|
),
|
||||||
headers={
|
headers={
|
||||||
"Content-Type": "application/json",
|
"Content-Type": "application/json",
|
||||||
"Authorization": f"Bearer {self.api_key}",
|
"Authorization": f"Bearer {self.embedding_api_key}",
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -47,14 +65,100 @@ class AsyncLLm:
|
||||||
|
|
||||||
return []
|
return []
|
||||||
|
|
||||||
|
async def chat(
|
||||||
|
self,
|
||||||
|
messages: list[dict],
|
||||||
|
temperature: float = 0.7,
|
||||||
|
max_tokens: int = 1024,
|
||||||
|
**extra,
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
messages=[
|
||||||
|
{"role": "system", "content": "你是一个专业助手"},
|
||||||
|
{"role": "user", "content": "解释一下什么是 RAG"},
|
||||||
|
]
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
resp = await self.chat_client.post(
|
||||||
|
self.chat_api,
|
||||||
|
json={
|
||||||
|
"model": self.chat_model,
|
||||||
|
"messages": messages,
|
||||||
|
"temperature": temperature,
|
||||||
|
"max_tokens": max_tokens,
|
||||||
|
**extra,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
resp.raise_for_status()
|
||||||
|
data = resp.json()
|
||||||
|
return data["choices"][0]["message"]["content"]
|
||||||
|
|
||||||
|
except Exception:
|
||||||
|
logger.exception("Chat request failed")
|
||||||
|
return ""
|
||||||
|
|
||||||
|
async def stream_chat(
|
||||||
|
self,
|
||||||
|
messages: list[dict],
|
||||||
|
temperature: float = 0.7,
|
||||||
|
max_tokens: int = 1024,
|
||||||
|
**extra,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
messages=[
|
||||||
|
{"role": "system", "content": "你是一个专业助手"},
|
||||||
|
{"role": "user", "content": "解释一下什么是 RAG"},
|
||||||
|
]
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
async with self.chat_client.stream(
|
||||||
|
"POST",
|
||||||
|
self.chat_api,
|
||||||
|
json={
|
||||||
|
"model": self.chat_model,
|
||||||
|
"messages": messages,
|
||||||
|
"temperature": temperature,
|
||||||
|
"max_tokens": max_tokens,
|
||||||
|
"stream": True,
|
||||||
|
**extra,
|
||||||
|
},
|
||||||
|
) as resp:
|
||||||
|
|
||||||
|
resp.raise_for_status()
|
||||||
|
|
||||||
|
async for line in resp.aiter_lines():
|
||||||
|
if not line or not line.startswith("data:"):
|
||||||
|
continue
|
||||||
|
|
||||||
|
data = line[len("data:") :].strip()
|
||||||
|
|
||||||
|
if data == "[DONE]":
|
||||||
|
break
|
||||||
|
|
||||||
|
try:
|
||||||
|
payload = json.loads(data)
|
||||||
|
delta = payload["choices"][0]["delta"]
|
||||||
|
content = delta.get("content")
|
||||||
|
if content:
|
||||||
|
yield content
|
||||||
|
except Exception:
|
||||||
|
continue
|
||||||
|
|
||||||
|
except Exception:
|
||||||
|
logger.exception("Stream chat failed")
|
||||||
|
|
||||||
async def close(self):
|
async def close(self):
|
||||||
await self.embedding_client.aclose()
|
await self.embedding_client.aclose()
|
||||||
|
await self.chat_client.aclose()
|
||||||
|
|
||||||
|
|
||||||
client: AsyncLLm | None = None
|
client: AsyncLLm | None = None
|
||||||
|
|
||||||
|
|
||||||
async def init_client():
|
async def init_client():
|
||||||
global client
|
global client
|
||||||
client = AsyncLLm()
|
client = AsyncLLm()
|
||||||
|
|
||||||
|
|
||||||
async def close_client():
|
async def close_client():
|
||||||
await client.close()
|
await client.close()
|
||||||
|
|
|
||||||
|
|
@ -141,3 +141,52 @@ def fixed_size_chunk(text, chunk_size=500, overlap=50):
|
||||||
chunks.append(text[max(0, start - overlap) : min(len(text), end + overlap)])
|
chunks.append(text[max(0, start - overlap) : min(len(text), end + overlap)])
|
||||||
start += chunk_size
|
start += chunk_size
|
||||||
return chunks
|
return chunks
|
||||||
|
|
||||||
|
|
||||||
|
def rag_system_prompt() -> str:
|
||||||
|
return """
|
||||||
|
你是一个基于知识库的 AI 助手,严格基于提供的文档回答用户问题。
|
||||||
|
- 不允许编造事实。
|
||||||
|
- 只使用提供的内容。
|
||||||
|
- 输出清晰、简洁、准确。
|
||||||
|
- 如果文档中没有明确答案,请回复 "文档未提供相关信息"。
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
def rag_user_prompt(query: str, documents: list[dict]) -> str:
|
||||||
|
"""
|
||||||
|
生成基于 RAG 的用户 Prompt。
|
||||||
|
|
||||||
|
:param query: 用户问题
|
||||||
|
:param documents: 检索到的文档列表,每个文档 dict 至少包含 title, content
|
||||||
|
:return: Prompt 字符串
|
||||||
|
"""
|
||||||
|
# 按相关度排序(如果文档里有 score,可以用)
|
||||||
|
documents = sorted(documents, key=lambda x: x.get("score", 0), reverse=True)
|
||||||
|
|
||||||
|
# 拼接文档内容
|
||||||
|
context_lines = []
|
||||||
|
for i, doc in enumerate(documents, start=1):
|
||||||
|
content = doc.get("content", "").replace("\n", " ").strip()
|
||||||
|
title = doc.get("title", f"文档{i}")
|
||||||
|
context_lines.append(f"文档 {i}:\n标题: {title}\n内容: {content}\n")
|
||||||
|
|
||||||
|
context_text = "\n".join(context_lines)
|
||||||
|
|
||||||
|
# 构建最终 Prompt
|
||||||
|
prompt = f"""
|
||||||
|
以下是从知识库中检索到的内容(按相关度排序):
|
||||||
|
<<<
|
||||||
|
{context_text}
|
||||||
|
>>>
|
||||||
|
|
||||||
|
用户问题: {query}
|
||||||
|
|
||||||
|
回答要求:
|
||||||
|
1. 列出完整列表。
|
||||||
|
2. 引用文档原文时,用“文档 i 原文: …”标注。
|
||||||
|
3. 每条用编号列出。
|
||||||
|
4. 如果文档中没有相关信息,请直接回复: "文档未提供相关信息"。
|
||||||
|
"""
|
||||||
|
logger.debug(prompt)
|
||||||
|
return prompt
|
||||||
|
|
|
||||||
|
|
@ -2,12 +2,12 @@ import pytest
|
||||||
import json
|
import json
|
||||||
from src.pipeline.core.pocket_flow import AsyncFlow
|
from src.pipeline.core.pocket_flow import AsyncFlow
|
||||||
from src.pipeline.core.utils import logger
|
from src.pipeline.core.utils import logger
|
||||||
from src.pipeline.core import llm, es, nodes
|
from src.pipeline.core import llm, es, nodes, utils
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_embedding():
|
async def test_embedding():
|
||||||
|
return
|
||||||
await llm.init_client()
|
await llm.init_client()
|
||||||
await es.init_client()
|
await es.init_client()
|
||||||
|
|
||||||
|
|
@ -42,8 +42,9 @@ async def test_search():
|
||||||
logger.debug("search from es")
|
logger.debug("search from es")
|
||||||
|
|
||||||
shared = {
|
shared = {
|
||||||
"text": "那座山盛产金属矿物",
|
"text": "哪里盛产金属矿物",
|
||||||
"index": "test_kb",
|
"index": "test_kb",
|
||||||
|
"top_k": 1,
|
||||||
"results": [], # [{es_id, doc_id, title, type, created_at, score, content}]
|
"results": [], # [{es_id, doc_id, title, type, created_at, score, content}]
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -56,5 +57,15 @@ async def test_search():
|
||||||
|
|
||||||
logger.debug(json.dumps({**shared, "embedding": shared["embedding"][:4]}, indent=4, ensure_ascii=False))
|
logger.debug(json.dumps({**shared, "embedding": shared["embedding"][:4]}, indent=4, ensure_ascii=False))
|
||||||
|
|
||||||
|
res = await llm.client.chat(
|
||||||
|
messages=[
|
||||||
|
{"role": "system", "content": utils.rag_system_prompt()},
|
||||||
|
{"role": "system", "content": utils.rag_user_prompt(shared["text"], shared["results"])},
|
||||||
|
# {"role": "system", "content": "你好"},
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug(res)
|
||||||
|
|
||||||
await llm.close_client()
|
await llm.close_client()
|
||||||
await es.close_client()
|
await es.close_client()
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue