feat: 测试prompt

This commit is contained in:
李如威 2025-12-19 15:08:16 +08:00
parent eb73101ddb
commit 2e2082808a
4 changed files with 172 additions and 14 deletions

View File

@ -22,9 +22,6 @@ class Config(TypedDict):
es_port: int
es_user: str
es_password: str
qdrant_host: str
qdrant_port: int
qdrant_api_key: str
def _read_config() -> Config:
@ -44,9 +41,6 @@ def _read_config() -> Config:
"es_port": int(os.getenv("ES_PORT")),
"es_user": os.getenv("ES_USER") or "elastic",
"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()

View File

@ -1,4 +1,5 @@
import httpx
import json
from src.pipeline.config import config
from src.pipeline.core.utils import logger
@ -7,13 +8,30 @@ class AsyncLLm:
def __init__(
self,
timeout: float = 30.0,
max_connections: int = 100,
max_keepalive: int = 20,
max_connections: int = 50,
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_model = config["embedding_model"]
self.api_key = config["embedding_api_key"]
logger.debug(self.embedding_api)
self.embedding_api_key = config["embedding_api_key"]
self.embedding_client = httpx.AsyncClient(
http2=False,
trust_env=False,
@ -24,7 +42,7 @@ class AsyncLLm:
),
headers={
"Content-Type": "application/json",
"Authorization": f"Bearer {self.api_key}",
"Authorization": f"Bearer {self.embedding_api_key}",
},
)
@ -47,14 +65,100 @@ class AsyncLLm:
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):
await self.embedding_client.aclose()
await self.chat_client.aclose()
client: AsyncLLm | None = None
async def init_client():
global client
client = AsyncLLm()
async def close_client():
await client.close()

View File

@ -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)])
start += chunk_size
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

View File

@ -2,12 +2,12 @@ import pytest
import json
from src.pipeline.core.pocket_flow import AsyncFlow
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
async def test_embedding():
return
await llm.init_client()
await es.init_client()
@ -42,8 +42,9 @@ async def test_search():
logger.debug("search from es")
shared = {
"text": "那座山盛产金属矿物",
"text": "哪里盛产金属矿物",
"index": "test_kb",
"top_k": 1,
"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))
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 es.close_client()