diff --git a/src/pipeline/config.py b/src/pipeline/config.py index 7dd43b0..e322641 100644 --- a/src/pipeline/config.py +++ b/src/pipeline/config.py @@ -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() diff --git a/src/pipeline/core/llm.py b/src/pipeline/core/llm.py index 5d1082a..b97396d 100644 --- a/src/pipeline/core/llm.py +++ b/src/pipeline/core/llm.py @@ -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() diff --git a/src/pipeline/core/utils.py b/src/pipeline/core/utils.py index cb8c4ca..b2024cf 100644 --- a/src/pipeline/core/utils.py +++ b/src/pipeline/core/utils.py @@ -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 diff --git a/src/tests/test_nodes.py b/src/tests/test_nodes.py index 49137bb..0c308e7 100644 --- a/src/tests/test_nodes.py +++ b/src/tests/test_nodes.py @@ -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()