feat: 测试prompt
This commit is contained in:
parent
eb73101ddb
commit
2e2082808a
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Reference in New Issue