55 lines
2.4 KiB
Python
55 lines
2.4 KiB
Python
import logging
|
||
from typing import Any, Dict, Optional
|
||
|
||
import httpx
|
||
|
||
from app.core.config import get_gateway_settings
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class RagService:
|
||
"""RAG 网关：直接调用 LightRAG 独立服务。
|
||
|
||
说明：历史版本依赖 `siyuan-rag-llm-main`（回档后目录缺失），导致服务无法启动。
|
||
当前实现改为通过 HTTP 调用 LightRAG 的官方 API（/query、/query/data 等）。
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
settings = get_gateway_settings()
|
||
self.base_url = settings.lightrag_url.rstrip("/")
|
||
self.api_key = settings.lightrag_api_key
|
||
|
||
def _headers(self) -> dict:
|
||
headers: dict = {}
|
||
if self.api_key:
|
||
headers["X-API-Key"] = self.api_key
|
||
return headers
|
||
|
||
async def query(self, query: str, top_k: int = 8) -> Dict[str, Any]:
|
||
"""调用 LightRAG /query，返回生成回答与引用。"""
|
||
payload = {"query": query, "top_k": top_k}
|
||
async with httpx.AsyncClient(base_url=self.base_url, timeout=60.0, headers=self._headers()) as client:
|
||
resp = await client.post("/query", json=payload)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
if not isinstance(data, dict):
|
||
return {"response": str(data), "references": []}
|
||
return data
|
||
|
||
async def graph(self, query: str, top_k: int = 60) -> Dict[str, Any]:
|
||
"""调用 LightRAG /query/data 获取结构化检索结果（实体/关系/分块/引用）。"""
|
||
payload = {"query": query, "top_k": top_k, "mode": "mix"}
|
||
async with httpx.AsyncClient(base_url=self.base_url, timeout=60.0, headers=self._headers()) as client:
|
||
resp = await client.post("/query/data", json=payload)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
return data if isinstance(data, dict) else {"status": "failure", "message": "invalid response", "data": data}
|
||
|
||
async def path(self, source: str, target: str) -> Dict[str, Any]:
|
||
"""路径查询：当前 LightRAG API 未提供等价 /path 接口，先返回明确错误。"""
|
||
message = "当前 LightRAG API 未提供 /path 等价接口（rag_gateway:path 暂不可用）"
|
||
logger.warning("%s: source=%s target=%s", message, source, target)
|
||
return {"status": "not_supported", "message": message, "source": source, "target": target}
|
||
|