feat(asset): 清理被截断的结尾图片引用，修复整单 502
仓库首次入库（此前仅存在于本地工作区与测试服构建产物中，未上传）。 改动（src/knowledge_adapter/asset_content_processor.py）： - 新增 _TRUNCATED_TRAILING_MARKDOWN_RESOURCE_IMAGE 规则，清除 chunk 末尾被截断的 
姚章浩
@@ -0,0 +1 @@
|
||||
"""Knowledge Adapter Service."""
|
||||
@@ -0,0 +1,201 @@
|
||||
from typing import Annotated, Any
|
||||
from uuid import UUID
|
||||
|
||||
from fastapi import Body, Depends, FastAPI, HTTPException, Request, status
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from fastapi.responses import JSONResponse, RedirectResponse, Response
|
||||
from pydantic import ValidationError
|
||||
from starlette.exceptions import HTTPException as StarletteHTTPException
|
||||
from starlette.types import Lifespan
|
||||
|
||||
from knowledge_adapter.asset_content_processor import ContentProcessingError
|
||||
from knowledge_adapter.asset_redirect import (
|
||||
AssetNotFoundError,
|
||||
AssetRedirector,
|
||||
AssetRedirectUnavailableError,
|
||||
)
|
||||
from knowledge_adapter.auth import DifyKeyAuthenticator
|
||||
from knowledge_adapter.models import RetrievalRequest
|
||||
from knowledge_adapter.observability import (
|
||||
bind_request_id,
|
||||
log_http_request,
|
||||
normalize_request_id,
|
||||
request_started_at,
|
||||
reset_request_id,
|
||||
)
|
||||
from knowledge_adapter.service import InvalidRequestError, RetrievalService
|
||||
from knowledge_adapter.weknora_client import WeKnoraBadResponseError, WeKnoraUnavailableError
|
||||
|
||||
_SAFE_HTTP_DETAILS = frozenset(
|
||||
{
|
||||
"content processing failed",
|
||||
"invalid API key",
|
||||
"invalid request",
|
||||
"invalid upstream response",
|
||||
"upstream unavailable",
|
||||
}
|
||||
)
|
||||
_DEFAULT_ERROR_DETAILS = {
|
||||
400: "invalid request",
|
||||
401: "invalid API key",
|
||||
403: "forbidden",
|
||||
405: "method not allowed",
|
||||
413: "invalid request",
|
||||
422: "invalid request",
|
||||
500: "internal server error",
|
||||
502: "bad gateway",
|
||||
503: "service unavailable",
|
||||
}
|
||||
|
||||
|
||||
def create_app(
|
||||
*,
|
||||
authenticator: DifyKeyAuthenticator | None = None,
|
||||
retrieval_service: RetrievalService | None = None,
|
||||
asset_redirect_service: AssetRedirector | None = None,
|
||||
max_query_chars: int = 2000,
|
||||
lifespan: Lifespan[FastAPI] | None = None,
|
||||
) -> FastAPI:
|
||||
app = FastAPI(
|
||||
docs_url=None,
|
||||
redoc_url=None,
|
||||
openapi_url=None,
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
@app.middleware("http")
|
||||
async def request_id_middleware(request: Request, call_next: Any) -> Response:
|
||||
request_id = normalize_request_id(request.headers.get("X-Request-ID"))
|
||||
request.state.request_id = request_id
|
||||
request_id_token = bind_request_id(request_id)
|
||||
started_at_ns = request_started_at()
|
||||
try:
|
||||
try:
|
||||
response = await call_next(request)
|
||||
except Exception:
|
||||
response = JSONResponse(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
content={"detail": "internal server error"},
|
||||
)
|
||||
response.headers["X-Request-ID"] = request_id
|
||||
route = getattr(request.scope.get("route"), "path", "<unmatched>")
|
||||
log_http_request(
|
||||
request_id=request_id,
|
||||
method=request.method,
|
||||
route=route,
|
||||
status_code=response.status_code,
|
||||
started_at_ns=started_at_ns,
|
||||
)
|
||||
return response
|
||||
finally:
|
||||
reset_request_id(request_id_token)
|
||||
|
||||
@app.exception_handler(RequestValidationError)
|
||||
async def invalid_request_handler(
|
||||
_request: Request,
|
||||
_exc: RequestValidationError,
|
||||
) -> JSONResponse:
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
content={"detail": "invalid request"},
|
||||
)
|
||||
|
||||
@app.exception_handler(StarletteHTTPException)
|
||||
async def safe_http_error_handler(
|
||||
request: Request,
|
||||
exc: StarletteHTTPException,
|
||||
) -> Response:
|
||||
if exc.status_code == status.HTTP_404_NOT_FOUND:
|
||||
if request.url.path.startswith("/i/"):
|
||||
return Response(status_code=status.HTTP_404_NOT_FOUND)
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
content={"detail": "not found"},
|
||||
)
|
||||
detail = (
|
||||
exc.detail
|
||||
if isinstance(exc.detail, str) and exc.detail in _SAFE_HTTP_DETAILS
|
||||
else _DEFAULT_ERROR_DETAILS.get(exc.status_code, "request failed")
|
||||
)
|
||||
headers = (
|
||||
{"WWW-Authenticate": "Bearer"}
|
||||
if exc.status_code == status.HTTP_401_UNAUTHORIZED
|
||||
else None
|
||||
)
|
||||
return JSONResponse(
|
||||
status_code=exc.status_code,
|
||||
content={"detail": detail},
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
@app.get("/healthz")
|
||||
async def healthz() -> dict[str, str]:
|
||||
return {"status": "ok"}
|
||||
|
||||
if asset_redirect_service is not None:
|
||||
|
||||
@app.get("/i/{asset_id:uuid}", response_class=RedirectResponse)
|
||||
async def asset_redirect(asset_id: UUID) -> Response:
|
||||
try:
|
||||
location = await asset_redirect_service.get_redirect_url(asset_id)
|
||||
except AssetNotFoundError:
|
||||
return Response(status_code=status.HTTP_404_NOT_FOUND)
|
||||
except AssetRedirectUnavailableError:
|
||||
return Response(status_code=status.HTTP_503_SERVICE_UNAVAILABLE)
|
||||
return RedirectResponse(
|
||||
url=location,
|
||||
status_code=status.HTTP_302_FOUND,
|
||||
headers={
|
||||
"Cache-Control": "no-store",
|
||||
"Referrer-Policy": "no-referrer",
|
||||
},
|
||||
)
|
||||
|
||||
@app.get("/i/{_invalid_asset_id}")
|
||||
async def invalid_asset_id(_invalid_asset_id: str) -> Response:
|
||||
return Response(status_code=status.HTTP_404_NOT_FOUND)
|
||||
|
||||
if authenticator is not None:
|
||||
|
||||
@app.post("/retrieval", dependencies=[Depends(authenticator)])
|
||||
async def retrieval(
|
||||
payload: Annotated[Any, Body()],
|
||||
) -> dict[str, list[dict[str, object]]]:
|
||||
try:
|
||||
request_model = RetrievalRequest.model_validate(payload)
|
||||
except ValidationError as exc:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="invalid request",
|
||||
) from exc
|
||||
if len(request_model.query) > max_query_chars:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="invalid request",
|
||||
)
|
||||
if retrieval_service is not None:
|
||||
try:
|
||||
return await retrieval_service.retrieve(request_model)
|
||||
except InvalidRequestError as exc:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="invalid request",
|
||||
) from exc
|
||||
except WeKnoraUnavailableError as exc:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
detail="upstream unavailable",
|
||||
) from exc
|
||||
except WeKnoraBadResponseError as exc:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_502_BAD_GATEWAY,
|
||||
detail="invalid upstream response",
|
||||
) from exc
|
||||
except ContentProcessingError as exc:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_502_BAD_GATEWAY,
|
||||
detail="content processing failed",
|
||||
) from exc
|
||||
return {"records": []}
|
||||
|
||||
return app
|
||||
@@ -0,0 +1,166 @@
|
||||
import json
|
||||
import re
|
||||
from typing import Protocol
|
||||
|
||||
from knowledge_adapter.asset_repository import (
|
||||
AssetMapping,
|
||||
AssetMappingConflictError,
|
||||
AssetRepositoryError,
|
||||
)
|
||||
from knowledge_adapter.models import WeKnoraChunk
|
||||
from knowledge_adapter.resource_locator import InvalidResourceReference, parse_resource_reference
|
||||
from knowledge_adapter.weknora_resource_resolver import (
|
||||
ResolvedAsset,
|
||||
ResourceResolutionError,
|
||||
)
|
||||
|
||||
_MARKDOWN_RESOURCE_IMAGE = re.compile(
|
||||
r"(?P<prefix>!\[[^\]\n]*\]\()"
|
||||
r"(?P<source_ref>resource://[A-Za-z0-9_-]{22})"
|
||||
r"(?P<suffix>\))"
|
||||
)
|
||||
_HTML_RESOURCE_IMAGE = re.compile(
|
||||
r"(?P<prefix><img\b[^>]*?\bsrc\s*=\s*(?P<quote>['\"]))"
|
||||
r"(?P<source_ref>resource://[A-Za-z0-9_-]{22})"
|
||||
r"(?P<suffix>(?P=quote)[^>]*>)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_TRUNCATED_TRAILING_MARKDOWN_RESOURCE_IMAGE = re.compile(
|
||||
r"!\[[^\]\n]*\]\(resource://[A-Za-z0-9_-]{0,22}\Z"
|
||||
)
|
||||
_FORBIDDEN_CONTENT = re.compile(
|
||||
r"(?:"
|
||||
r"local://"
|
||||
r"|(?<![A-Za-z0-9])/(?:home|srv|var|etc|root|tmp)/"
|
||||
r"|https?://(?:localhost|127(?:\.\d{1,3}){3}|10(?:\.\d{1,3}){3}|"
|
||||
r"192\.168(?:\.\d{1,3}){2}|172\.(?:1[6-9]|2\d|3[01])(?:\.\d{1,3}){2})"
|
||||
r"(?=[:/])"
|
||||
r"|https?://[^\s)\]]*\.aliyuncs\.com(?=[:/]|$)"
|
||||
r")",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
class ContentProcessingError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class MappingRepository(Protocol):
|
||||
async def get_active_by_source_ref(self, source_ref: str) -> AssetMapping | None: ...
|
||||
|
||||
async def get_or_create(
|
||||
self,
|
||||
*,
|
||||
source_ref: str,
|
||||
object_key: str,
|
||||
content_type: str | None,
|
||||
) -> AssetMapping: ...
|
||||
|
||||
|
||||
class ResourceResolver(Protocol):
|
||||
async def resolve(self, source_ref: str) -> ResolvedAsset: ...
|
||||
|
||||
|
||||
class AssetContentProcessor:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
repository: MappingRepository,
|
||||
resolver: ResourceResolver,
|
||||
public_base_url: str,
|
||||
) -> None:
|
||||
self._repository = repository
|
||||
self._resolver = resolver
|
||||
self._public_base_url = public_base_url.rstrip("/")
|
||||
|
||||
async def _public_url(
|
||||
self,
|
||||
source_ref: str,
|
||||
cache: dict[str, str],
|
||||
) -> str:
|
||||
cached = cache.get(source_ref)
|
||||
if cached is not None:
|
||||
return cached
|
||||
try:
|
||||
mapping = await self._repository.get_active_by_source_ref(source_ref)
|
||||
if mapping is None:
|
||||
resolved = await self._resolver.resolve(source_ref)
|
||||
mapping = await self._repository.get_or_create(
|
||||
source_ref=source_ref,
|
||||
object_key=resolved.object_key,
|
||||
content_type=resolved.content_type,
|
||||
)
|
||||
except (
|
||||
AssetMappingConflictError,
|
||||
AssetRepositoryError,
|
||||
ResourceResolutionError,
|
||||
) as error:
|
||||
raise ContentProcessingError("content processing failed") from error
|
||||
public_url = f"{self._public_base_url}/i/{mapping.asset_id}"
|
||||
cache[source_ref] = public_url
|
||||
return public_url
|
||||
|
||||
async def process(self, chunk: WeKnoraChunk) -> str:
|
||||
parts: list[str] = []
|
||||
public_urls: dict[str, str] = {}
|
||||
cursor = 0
|
||||
for match in _MARKDOWN_RESOURCE_IMAGE.finditer(chunk.content):
|
||||
source_ref = match.group("source_ref")
|
||||
public_url = await self._public_url(source_ref, public_urls)
|
||||
parts.extend(
|
||||
(
|
||||
chunk.content[cursor : match.start()],
|
||||
match.group("prefix"),
|
||||
public_url,
|
||||
match.group("suffix"),
|
||||
)
|
||||
)
|
||||
cursor = match.end()
|
||||
parts.append(chunk.content[cursor:])
|
||||
processed = "".join(parts)
|
||||
|
||||
parts = []
|
||||
cursor = 0
|
||||
for match in _HTML_RESOURCE_IMAGE.finditer(processed):
|
||||
source_ref = match.group("source_ref")
|
||||
public_url = await self._public_url(source_ref, public_urls)
|
||||
parts.extend(
|
||||
(
|
||||
processed[cursor : match.start()],
|
||||
match.group("prefix"),
|
||||
public_url,
|
||||
match.group("suffix"),
|
||||
)
|
||||
)
|
||||
cursor = match.end()
|
||||
parts.append(processed[cursor:])
|
||||
processed = "".join(parts)
|
||||
processed = _TRUNCATED_TRAILING_MARKDOWN_RESOURCE_IMAGE.sub("", processed)
|
||||
|
||||
seen_refs = set(public_urls)
|
||||
try:
|
||||
image_entries = json.loads(chunk.image_info) if chunk.image_info else []
|
||||
except json.JSONDecodeError as error:
|
||||
raise ContentProcessingError("content processing failed") from error
|
||||
if not isinstance(image_entries, list):
|
||||
raise ContentProcessingError("content processing failed")
|
||||
for entry in image_entries:
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
for field in ("url", "original_url"):
|
||||
candidate = entry.get(field)
|
||||
if not isinstance(candidate, str):
|
||||
continue
|
||||
try:
|
||||
source_ref = f"resource://{parse_resource_reference(candidate)}"
|
||||
except InvalidResourceReference:
|
||||
continue
|
||||
if source_ref in seen_refs:
|
||||
continue
|
||||
seen_refs.add(source_ref)
|
||||
public_url = await self._public_url(source_ref, public_urls)
|
||||
separator = "\n\n" if processed else ""
|
||||
processed += f"{separator}"
|
||||
if "resource://" in processed or _FORBIDDEN_CONTENT.search(processed):
|
||||
raise ContentProcessingError("content processing failed")
|
||||
return processed
|
||||
@@ -0,0 +1,56 @@
|
||||
from typing import Protocol
|
||||
from uuid import UUID
|
||||
|
||||
from knowledge_adapter.asset_repository import AssetMapping, AssetRepositoryError
|
||||
|
||||
PRESIGNED_GET_EXPIRES_SECONDS = 600
|
||||
|
||||
|
||||
class AssetNotFoundError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class AssetSignerError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class AssetRedirectUnavailableError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class AssetLookupRepository(Protocol):
|
||||
async def get_active_by_asset_id(self, asset_id: UUID) -> AssetMapping | None: ...
|
||||
|
||||
|
||||
class AssetGetSigner(Protocol):
|
||||
async def sign_get(self, *, object_key: str, expires_seconds: int) -> str: ...
|
||||
|
||||
|
||||
class AssetRedirector(Protocol):
|
||||
async def get_redirect_url(self, asset_id: UUID) -> str: ...
|
||||
|
||||
|
||||
class AssetRedirectService:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
repository: AssetLookupRepository,
|
||||
signer: AssetGetSigner,
|
||||
) -> None:
|
||||
self._repository = repository
|
||||
self._signer = signer
|
||||
|
||||
async def get_redirect_url(self, asset_id: UUID) -> str:
|
||||
try:
|
||||
mapping = await self._repository.get_active_by_asset_id(asset_id)
|
||||
except AssetRepositoryError as error:
|
||||
raise AssetRedirectUnavailableError("asset unavailable") from error
|
||||
if mapping is None:
|
||||
raise AssetNotFoundError("asset not found")
|
||||
try:
|
||||
return await self._signer.sign_get(
|
||||
object_key=mapping.object_key,
|
||||
expires_seconds=PRESIGNED_GET_EXPIRES_SECONDS,
|
||||
)
|
||||
except AssetSignerError as error:
|
||||
raise AssetRedirectUnavailableError("asset unavailable") from error
|
||||
@@ -0,0 +1,113 @@
|
||||
from collections.abc import Mapping
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Any, Protocol, cast
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
_UPSERT_MAPPING = """
|
||||
INSERT INTO asset_mapping (asset_id, source_ref, object_key, content_type, status)
|
||||
VALUES ($1, $2, $3, $4, 'active')
|
||||
ON CONFLICT (source_ref) DO UPDATE
|
||||
SET source_ref = asset_mapping.source_ref
|
||||
WHERE asset_mapping.object_key = EXCLUDED.object_key
|
||||
RETURNING asset_id, source_ref, object_key, content_type, created_at, status
|
||||
"""
|
||||
|
||||
_SELECT_ACTIVE_BY_ASSET_ID = """
|
||||
SELECT asset_id, source_ref, object_key, content_type, created_at, status
|
||||
FROM asset_mapping
|
||||
WHERE asset_id = $1
|
||||
AND status = 'active'
|
||||
"""
|
||||
|
||||
_SELECT_ACTIVE_BY_SOURCE_REF = """
|
||||
SELECT asset_id, source_ref, object_key, content_type, created_at, status
|
||||
FROM asset_mapping
|
||||
WHERE source_ref = $1
|
||||
AND status = 'active'
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AssetMapping:
|
||||
asset_id: UUID
|
||||
source_ref: str
|
||||
object_key: str
|
||||
content_type: str | None
|
||||
created_at: datetime
|
||||
status: str
|
||||
|
||||
|
||||
class AssetMappingConflictError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class AssetRepositoryError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class AssetConnection(Protocol):
|
||||
async def fetchrow(self, query: str, *args: object) -> Mapping[str, Any] | None: ...
|
||||
|
||||
|
||||
class AssetPool(Protocol):
|
||||
def acquire(self) -> AbstractAsyncContextManager[AssetConnection]: ...
|
||||
|
||||
|
||||
def _mapping_from_row(row: Mapping[str, Any]) -> AssetMapping:
|
||||
return AssetMapping(
|
||||
asset_id=UUID(str(row["asset_id"])),
|
||||
source_ref=str(row["source_ref"]),
|
||||
object_key=str(row["object_key"]),
|
||||
content_type=str(row["content_type"]) if row["content_type"] else None,
|
||||
created_at=cast(datetime, row["created_at"]),
|
||||
status=str(row["status"]),
|
||||
)
|
||||
|
||||
|
||||
class AssetRepository:
|
||||
def __init__(self, *, pool: AssetPool) -> None:
|
||||
self._pool = pool
|
||||
|
||||
async def get_active_by_asset_id(self, asset_id: UUID) -> AssetMapping | None:
|
||||
try:
|
||||
async with self._pool.acquire() as connection:
|
||||
row = await connection.fetchrow(_SELECT_ACTIVE_BY_ASSET_ID, asset_id)
|
||||
except Exception as error:
|
||||
raise AssetRepositoryError("asset repository unavailable") from error
|
||||
if row is None:
|
||||
return None
|
||||
return _mapping_from_row(row)
|
||||
|
||||
async def get_active_by_source_ref(self, source_ref: str) -> AssetMapping | None:
|
||||
try:
|
||||
async with self._pool.acquire() as connection:
|
||||
row = await connection.fetchrow(_SELECT_ACTIVE_BY_SOURCE_REF, source_ref)
|
||||
except Exception as error:
|
||||
raise AssetRepositoryError("asset repository unavailable") from error
|
||||
if row is None:
|
||||
return None
|
||||
return _mapping_from_row(row)
|
||||
|
||||
async def get_or_create(
|
||||
self,
|
||||
*,
|
||||
source_ref: str,
|
||||
object_key: str,
|
||||
content_type: str | None,
|
||||
) -> AssetMapping:
|
||||
try:
|
||||
async with self._pool.acquire() as connection:
|
||||
row = await connection.fetchrow(
|
||||
_UPSERT_MAPPING,
|
||||
uuid4(),
|
||||
source_ref,
|
||||
object_key,
|
||||
content_type,
|
||||
)
|
||||
except Exception as error:
|
||||
raise AssetRepositoryError("asset repository unavailable") from error
|
||||
if row is None:
|
||||
raise AssetMappingConflictError("asset mapping conflict")
|
||||
return _mapping_from_row(row)
|
||||
@@ -0,0 +1,38 @@
|
||||
import hashlib
|
||||
import hmac
|
||||
from typing import Annotated, NoReturn
|
||||
|
||||
from fastapi import Header, HTTPException, status
|
||||
|
||||
|
||||
class DifyKeyAuthenticator:
|
||||
def __init__(self, *, current_hash: str, next_hash: str | None = None) -> None:
|
||||
self.current_hash = current_hash
|
||||
self.next_hash = next_hash
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
authorization: Annotated[str | None, Header()] = None,
|
||||
) -> None:
|
||||
if authorization is None:
|
||||
self._reject()
|
||||
|
||||
scheme, separator, token = authorization.partition(" ")
|
||||
if scheme.lower() != "bearer" or separator != " " or not token or " " in token:
|
||||
self._reject()
|
||||
|
||||
digest = hashlib.sha256(token.encode()).hexdigest()
|
||||
current_matches = hmac.compare_digest(digest, self.current_hash)
|
||||
next_matches = False
|
||||
if self.next_hash is not None:
|
||||
next_matches = hmac.compare_digest(digest, self.next_hash)
|
||||
if not (current_matches or next_matches):
|
||||
self._reject()
|
||||
|
||||
@staticmethod
|
||||
def _reject() -> NoReturn:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="invalid API key",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
@@ -0,0 +1,25 @@
|
||||
import re
|
||||
|
||||
TRUNCATION_MARKER = "[内容已由 Knowledge Adapter 截断]"
|
||||
_MARKDOWN_LINK_PATTERN = re.compile(r"!?\[[^\]\n]*\]\([^\)\n]*\)")
|
||||
_FENCED_CODE_PATTERN = re.compile(r"```[\s\S]*?(?:```|\Z)")
|
||||
|
||||
|
||||
def _safe_cutoff(content: str, cutoff: int) -> int:
|
||||
protected_starts = [
|
||||
match.start()
|
||||
for pattern in (_MARKDOWN_LINK_PATTERN, _FENCED_CODE_PATTERN)
|
||||
for match in pattern.finditer(content)
|
||||
if match.start() < cutoff < match.end()
|
||||
]
|
||||
return min(protected_starts, default=cutoff)
|
||||
|
||||
|
||||
def truncate_content(content: str, *, max_chars: int) -> str:
|
||||
if len(content) <= max_chars:
|
||||
return content
|
||||
if max_chars < len(TRUNCATION_MARKER):
|
||||
raise ValueError("max_chars is too small for the truncation marker")
|
||||
prefix_chars = max_chars - len(TRUNCATION_MARKER)
|
||||
prefix_chars = _safe_cutoff(content, prefix_chars)
|
||||
return f"{content[:prefix_chars]}{TRUNCATION_MARKER}"
|
||||
@@ -0,0 +1,9 @@
|
||||
from fastapi import FastAPI
|
||||
|
||||
from knowledge_adapter.factory import build_production_app
|
||||
from knowledge_adapter.settings import Settings
|
||||
|
||||
|
||||
def create_production_app() -> FastAPI:
|
||||
settings = Settings() # type: ignore[call-arg]
|
||||
return build_production_app(settings)
|
||||
@@ -0,0 +1,209 @@
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
import asyncpg
|
||||
import httpx
|
||||
from fastapi import FastAPI
|
||||
|
||||
from knowledge_adapter.app import create_app
|
||||
from knowledge_adapter.asset_content_processor import (
|
||||
AssetContentProcessor,
|
||||
ContentProcessingError,
|
||||
)
|
||||
from knowledge_adapter.asset_redirect import (
|
||||
AssetRedirectService,
|
||||
AssetRedirectUnavailableError,
|
||||
)
|
||||
from knowledge_adapter.asset_repository import AssetRepository
|
||||
from knowledge_adapter.auth import DifyKeyAuthenticator
|
||||
from knowledge_adapter.models import WeKnoraChunk
|
||||
from knowledge_adapter.oss_signer import (
|
||||
AlibabaOssV4Signer,
|
||||
TemporaryCredentialClient,
|
||||
create_ecs_ram_role_credential_client,
|
||||
create_static_ak_credential_client,
|
||||
)
|
||||
from knowledge_adapter.service import RetrievalService
|
||||
from knowledge_adapter.settings import Settings
|
||||
from knowledge_adapter.weknora_client import WeKnoraClient
|
||||
from knowledge_adapter.weknora_resource_resolver import WeKnoraResourceResolver
|
||||
|
||||
|
||||
def build_app(
|
||||
settings: Settings,
|
||||
*,
|
||||
weknora_transport: httpx.AsyncBaseTransport | None = None,
|
||||
) -> FastAPI:
|
||||
authenticator = DifyKeyAuthenticator(
|
||||
current_hash=settings.dify_key_hash_current,
|
||||
next_hash=settings.dify_key_hash_next,
|
||||
)
|
||||
weknora_client = WeKnoraClient(
|
||||
base_url=str(settings.weknora_base_url),
|
||||
api_key=settings.read_weknora_api_key(),
|
||||
timeout_seconds=settings.weknora_timeout_seconds,
|
||||
transport=weknora_transport,
|
||||
)
|
||||
retrieval_service = RetrievalService(
|
||||
search_client=weknora_client,
|
||||
allowed_kb_ids=settings.weknora_allowed_kb_ids,
|
||||
max_top_k=settings.max_top_k,
|
||||
max_chunk_chars=settings.max_chunk_chars,
|
||||
max_total_chars=settings.max_total_chars,
|
||||
)
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
await weknora_client.aclose()
|
||||
|
||||
app = create_app(
|
||||
authenticator=authenticator,
|
||||
retrieval_service=retrieval_service,
|
||||
max_query_chars=settings.max_query_chars,
|
||||
lifespan=lifespan,
|
||||
)
|
||||
app.state.weknora_client = weknora_client
|
||||
return app
|
||||
|
||||
|
||||
PoolFactory = Callable[..., Awaitable[Any]]
|
||||
CredentialClientFactory = Callable[[str], TemporaryCredentialClient]
|
||||
|
||||
|
||||
class _RuntimeAssetRedirector:
|
||||
def __init__(self) -> None:
|
||||
self._delegate: AssetRedirectService | None = None
|
||||
|
||||
def bind(self, delegate: AssetRedirectService) -> None:
|
||||
self._delegate = delegate
|
||||
|
||||
async def get_redirect_url(self, asset_id: UUID) -> str:
|
||||
if self._delegate is None:
|
||||
raise AssetRedirectUnavailableError("asset unavailable")
|
||||
return await self._delegate.get_redirect_url(asset_id)
|
||||
|
||||
|
||||
class _RuntimeContentProcessor:
|
||||
def __init__(self) -> None:
|
||||
self._delegate: AssetContentProcessor | None = None
|
||||
|
||||
def bind(self, delegate: AssetContentProcessor) -> None:
|
||||
self._delegate = delegate
|
||||
|
||||
async def process(self, chunk: WeKnoraChunk) -> str:
|
||||
if self._delegate is None:
|
||||
raise ContentProcessingError("content processing failed")
|
||||
return await self._delegate.process(chunk)
|
||||
|
||||
|
||||
def build_production_app(
|
||||
settings: Settings,
|
||||
*,
|
||||
pool_factory: PoolFactory | None = None,
|
||||
credential_client_factory: CredentialClientFactory | None = None,
|
||||
weknora_transport: httpx.AsyncBaseTransport | None = None,
|
||||
) -> FastAPI:
|
||||
effective_pool_factory = pool_factory or asyncpg.create_pool
|
||||
effective_credential_client_factory = (
|
||||
credential_client_factory or create_ecs_ram_role_credential_client
|
||||
)
|
||||
authenticator = DifyKeyAuthenticator(
|
||||
current_hash=settings.dify_key_hash_current,
|
||||
next_hash=settings.dify_key_hash_next,
|
||||
)
|
||||
weknora_client = WeKnoraClient(
|
||||
base_url=str(settings.weknora_base_url),
|
||||
api_key=settings.read_weknora_api_key(),
|
||||
timeout_seconds=settings.weknora_timeout_seconds,
|
||||
transport=weknora_transport,
|
||||
)
|
||||
runtime_processor = _RuntimeContentProcessor()
|
||||
runtime_redirector = _RuntimeAssetRedirector()
|
||||
retrieval_service = RetrievalService(
|
||||
search_client=weknora_client,
|
||||
allowed_kb_ids=settings.weknora_allowed_kb_ids,
|
||||
max_top_k=settings.max_top_k,
|
||||
max_chunk_chars=settings.max_chunk_chars,
|
||||
max_total_chars=settings.max_total_chars,
|
||||
content_processor=runtime_processor,
|
||||
)
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
|
||||
pools: list[Any] = []
|
||||
try:
|
||||
asset_dsn = settings.asset_database_url_file.read_text(encoding="utf-8").strip()
|
||||
weknora_dsn = settings.weknora_readonly_database_url_file.read_text(
|
||||
encoding="utf-8"
|
||||
).strip()
|
||||
asset_pool = await effective_pool_factory(
|
||||
dsn=asset_dsn,
|
||||
min_size=1,
|
||||
max_size=5,
|
||||
command_timeout=5,
|
||||
server_settings={"statement_timeout": "5000"},
|
||||
)
|
||||
pools.append(asset_pool)
|
||||
weknora_pool = await effective_pool_factory(
|
||||
dsn=weknora_dsn,
|
||||
min_size=1,
|
||||
max_size=2,
|
||||
command_timeout=1,
|
||||
server_settings={
|
||||
"default_transaction_read_only": "on",
|
||||
"statement_timeout": "1000",
|
||||
},
|
||||
)
|
||||
pools.append(weknora_pool)
|
||||
asset_repository = AssetRepository(pool=asset_pool)
|
||||
runtime_processor.bind(
|
||||
AssetContentProcessor(
|
||||
repository=asset_repository,
|
||||
resolver=WeKnoraResourceResolver(
|
||||
pool=weknora_pool,
|
||||
tenant_id=settings.weknora_tenant_id,
|
||||
storage_backend_id=settings.weknora_oss_storage_backend_id,
|
||||
bucket=settings.oss_bucket,
|
||||
),
|
||||
public_base_url=str(settings.asset_public_base_url),
|
||||
)
|
||||
)
|
||||
if settings.oss_credential_mode == "static_ak":
|
||||
credential_client = create_static_ak_credential_client(
|
||||
settings.read_oss_access_key_id(),
|
||||
settings.read_oss_access_key_secret(),
|
||||
)
|
||||
else:
|
||||
if settings.oss_ram_role_name is None:
|
||||
raise RuntimeError("OSS credential unavailable")
|
||||
credential_client = effective_credential_client_factory(settings.oss_ram_role_name)
|
||||
runtime_redirector.bind(
|
||||
AssetRedirectService(
|
||||
repository=asset_repository,
|
||||
signer=AlibabaOssV4Signer(
|
||||
credential_client=credential_client,
|
||||
region=settings.oss_region,
|
||||
bucket=settings.oss_bucket,
|
||||
),
|
||||
)
|
||||
)
|
||||
yield
|
||||
finally:
|
||||
for pool in reversed(pools):
|
||||
await pool.close()
|
||||
await weknora_client.aclose()
|
||||
|
||||
app = create_app(
|
||||
authenticator=authenticator,
|
||||
retrieval_service=retrieval_service,
|
||||
asset_redirect_service=runtime_redirector,
|
||||
max_query_chars=settings.max_query_chars,
|
||||
lifespan=lifespan,
|
||||
)
|
||||
app.state.weknora_client = weknora_client
|
||||
return app
|
||||
@@ -0,0 +1,56 @@
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
||||
|
||||
|
||||
class RetrievalSetting(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
top_k: int = Field(ge=1)
|
||||
score_threshold: float = Field(ge=0, le=1)
|
||||
|
||||
|
||||
class MetadataCondition(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
logical_operator: str | None = None
|
||||
conditions: list[dict[str, Any]] = Field(default_factory=list)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def reject_nonempty_conditions(self) -> "MetadataCondition":
|
||||
if self.conditions:
|
||||
raise ValueError("metadata conditions are not supported")
|
||||
return self
|
||||
|
||||
|
||||
class RetrievalRequest(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
knowledge_id: UUID
|
||||
query: str
|
||||
retrieval_setting: RetrievalSetting
|
||||
metadata_condition: MetadataCondition | None = None
|
||||
|
||||
@field_validator("query")
|
||||
@classmethod
|
||||
def strip_nonempty_query(cls, value: str) -> str:
|
||||
normalized = value.strip()
|
||||
if not normalized:
|
||||
raise ValueError("query must not be empty")
|
||||
return normalized
|
||||
|
||||
|
||||
class WeKnoraChunk(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", allow_inf_nan=False)
|
||||
|
||||
id: UUID
|
||||
content: str
|
||||
score: float
|
||||
knowledge_base_id: UUID
|
||||
knowledge_id: UUID
|
||||
knowledge_title: str = ""
|
||||
knowledge_filename: str = ""
|
||||
chunk_index: int
|
||||
chunk_type: str
|
||||
image_info: str = ""
|
||||
@@ -0,0 +1,74 @@
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from contextvars import ContextVar, Token
|
||||
from time import monotonic_ns
|
||||
from uuid import uuid4
|
||||
|
||||
_REQUEST_ID_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$")
|
||||
_HTTP_LOGGER = logging.getLogger("knowledge_adapter.http")
|
||||
_CURRENT_REQUEST_ID: ContextVar[str | None] = ContextVar(
|
||||
"knowledge_adapter_request_id",
|
||||
default=None,
|
||||
)
|
||||
_ERROR_CODE_BY_STATUS = {
|
||||
400: "invalid_request",
|
||||
401: "unauthorized",
|
||||
403: "forbidden",
|
||||
404: "not_found",
|
||||
405: "method_not_allowed",
|
||||
413: "invalid_request",
|
||||
418: "request_error",
|
||||
422: "invalid_request",
|
||||
500: "internal_error",
|
||||
502: "upstream_error",
|
||||
503: "unavailable",
|
||||
}
|
||||
|
||||
|
||||
def normalize_request_id(value: str | None) -> str:
|
||||
if value is not None and _REQUEST_ID_PATTERN.fullmatch(value):
|
||||
return value
|
||||
return str(uuid4())
|
||||
|
||||
|
||||
def bind_request_id(request_id: str) -> Token[str | None]:
|
||||
return _CURRENT_REQUEST_ID.set(request_id)
|
||||
|
||||
|
||||
def reset_request_id(token: Token[str | None]) -> None:
|
||||
_CURRENT_REQUEST_ID.reset(token)
|
||||
|
||||
|
||||
def current_request_id() -> str | None:
|
||||
return _CURRENT_REQUEST_ID.get()
|
||||
|
||||
|
||||
def request_started_at() -> int:
|
||||
return monotonic_ns()
|
||||
|
||||
|
||||
def log_http_request(
|
||||
*,
|
||||
request_id: str,
|
||||
method: str,
|
||||
route: str,
|
||||
status_code: int,
|
||||
started_at_ns: int,
|
||||
) -> None:
|
||||
duration_ms = max(0, (monotonic_ns() - started_at_ns) // 1_000_000)
|
||||
event = {
|
||||
"duration_ms": duration_ms,
|
||||
"event": "http_request",
|
||||
"method": method,
|
||||
"request_id": request_id,
|
||||
"route": route,
|
||||
"status_code": status_code,
|
||||
}
|
||||
if status_code >= 400:
|
||||
event["error_code"] = _ERROR_CODE_BY_STATUS.get(status_code, "request_error")
|
||||
level = logging.ERROR if status_code >= 500 else logging.INFO
|
||||
_HTTP_LOGGER.log(
|
||||
level,
|
||||
json.dumps(event, separators=(",", ":"), sort_keys=True),
|
||||
)
|
||||
@@ -0,0 +1,124 @@
|
||||
import asyncio
|
||||
from typing import Protocol
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import oss2
|
||||
from alibabacloud_credentials.client import Client as CredentialsClient
|
||||
from alibabacloud_credentials.models import Config, CredentialModel
|
||||
from oss2.credentials import Credentials
|
||||
from pydantic import SecretStr
|
||||
|
||||
from knowledge_adapter.asset_redirect import AssetSignerError
|
||||
from knowledge_adapter.resource_locator import InvalidResourceLocator, validate_oss_object_key
|
||||
|
||||
|
||||
class TemporaryCredentialModel(Protocol):
|
||||
access_key_id: str | None
|
||||
access_key_secret: str | None
|
||||
security_token: str | None
|
||||
|
||||
|
||||
class TemporaryCredentialClient(Protocol):
|
||||
def get_credential(self) -> TemporaryCredentialModel: ...
|
||||
|
||||
|
||||
class OssBucket(Protocol):
|
||||
def sign_url(
|
||||
self,
|
||||
method: str,
|
||||
key: str,
|
||||
expires: int,
|
||||
) -> str: ...
|
||||
|
||||
|
||||
def create_ecs_ram_role_credential_client(role_name: str) -> TemporaryCredentialClient:
|
||||
config = Config(
|
||||
type="ecs_ram_role",
|
||||
role_name=role_name,
|
||||
enable_imds_v2=True,
|
||||
disable_imds_v1=True,
|
||||
connect_timeout=1000,
|
||||
timeout=1000,
|
||||
)
|
||||
return CredentialsClient(config)
|
||||
|
||||
|
||||
class _StaticAkCredentialClient:
|
||||
def __init__(self, access_key_id: SecretStr, access_key_secret: SecretStr) -> None:
|
||||
self._access_key_id = access_key_id
|
||||
self._access_key_secret = access_key_secret
|
||||
|
||||
def get_credential(self) -> CredentialModel:
|
||||
return CredentialModel(
|
||||
access_key_id=self._access_key_id.get_secret_value(),
|
||||
access_key_secret=self._access_key_secret.get_secret_value(),
|
||||
security_token="",
|
||||
type="access_key",
|
||||
)
|
||||
|
||||
|
||||
def create_static_ak_credential_client(
|
||||
access_key_id: SecretStr,
|
||||
access_key_secret: SecretStr,
|
||||
) -> TemporaryCredentialClient:
|
||||
return _StaticAkCredentialClient(access_key_id, access_key_secret)
|
||||
|
||||
|
||||
class _RefreshingOssCredentialsProvider:
|
||||
def __init__(self, client: TemporaryCredentialClient) -> None:
|
||||
self._client = client
|
||||
|
||||
def get_credentials(self) -> Credentials:
|
||||
credential = self._client.get_credential()
|
||||
return Credentials(
|
||||
credential.access_key_id or "",
|
||||
credential.access_key_secret or "",
|
||||
credential.security_token or "",
|
||||
)
|
||||
|
||||
|
||||
class AlibabaOssV4Signer:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
credential_client: TemporaryCredentialClient,
|
||||
region: str,
|
||||
bucket: str,
|
||||
) -> None:
|
||||
endpoint = f"https://oss-{region}.aliyuncs.com"
|
||||
auth = oss2.ProviderAuthV4(_RefreshingOssCredentialsProvider(credential_client))
|
||||
self._expected_hostname = f"{bucket}.oss-{region}.aliyuncs.com"
|
||||
self._bucket: OssBucket = oss2.Bucket(
|
||||
auth,
|
||||
endpoint,
|
||||
bucket,
|
||||
region=region,
|
||||
)
|
||||
|
||||
async def sign_get(self, *, object_key: str, expires_seconds: int) -> str:
|
||||
try:
|
||||
validated_key = validate_oss_object_key(object_key)
|
||||
except InvalidResourceLocator as error:
|
||||
raise AssetSignerError("asset signing failed") from error
|
||||
if expires_seconds != 600:
|
||||
raise AssetSignerError("asset signing failed")
|
||||
try:
|
||||
signed_url = await asyncio.to_thread(
|
||||
self._bucket.sign_url,
|
||||
"GET",
|
||||
validated_key,
|
||||
expires_seconds,
|
||||
)
|
||||
except Exception as error:
|
||||
raise AssetSignerError("asset signing failed") from error
|
||||
parsed = urlparse(signed_url)
|
||||
query = parse_qs(parsed.query)
|
||||
if (
|
||||
parsed.scheme != "https"
|
||||
or parsed.hostname != self._expected_hostname
|
||||
or query.get("x-oss-signature-version") != ["OSS4-HMAC-SHA256"]
|
||||
or query.get("x-oss-expires") != ["600"]
|
||||
or not query.get("x-oss-signature")
|
||||
):
|
||||
raise AssetSignerError("asset signing failed")
|
||||
return signed_url
|
||||
@@ -0,0 +1,57 @@
|
||||
import re
|
||||
|
||||
_RESOURCE_REFERENCE_PATTERN = re.compile(r"resource://([A-Za-z0-9_-]{22})\Z")
|
||||
|
||||
|
||||
class InvalidResourceReference(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
class InvalidResourceLocator(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
def parse_resource_reference(value: str) -> str:
|
||||
match = _RESOURCE_REFERENCE_PATTERN.fullmatch(value.strip())
|
||||
if match is None:
|
||||
raise InvalidResourceReference("invalid resource reference")
|
||||
return match.group(1)
|
||||
|
||||
|
||||
def parse_oss_object_key(
|
||||
*,
|
||||
provider: str,
|
||||
physical_path: str,
|
||||
storage_backend_id: str,
|
||||
expected_storage_backend_id: str,
|
||||
expected_bucket: str,
|
||||
) -> str:
|
||||
if provider != "oss" or storage_backend_id != expected_storage_backend_id:
|
||||
raise InvalidResourceLocator("invalid resource locator")
|
||||
|
||||
storage_prefix = f"storage://{expected_storage_backend_id}/"
|
||||
if not physical_path.startswith(storage_prefix):
|
||||
raise InvalidResourceLocator("invalid resource locator")
|
||||
provider_path = physical_path.removeprefix(storage_prefix)
|
||||
|
||||
oss_prefix = f"oss://{expected_bucket}/"
|
||||
if not provider_path.startswith(oss_prefix):
|
||||
raise InvalidResourceLocator("invalid resource locator")
|
||||
object_key = provider_path.removeprefix(oss_prefix)
|
||||
return validate_oss_object_key(object_key)
|
||||
|
||||
|
||||
def validate_oss_object_key(object_key: str) -> str:
|
||||
segments = object_key.split("/")
|
||||
invalid_character = any(
|
||||
ord(char) < 32 or ord(char) == 127 or char in "\\?#%"
|
||||
for char in object_key
|
||||
)
|
||||
if (
|
||||
not object_key
|
||||
or len(object_key.encode("utf-8")) > 1023
|
||||
or invalid_character
|
||||
or any(segment in {"", ".", ".."} for segment in segments)
|
||||
):
|
||||
raise InvalidResourceLocator("invalid resource locator")
|
||||
return object_key
|
||||
@@ -0,0 +1,96 @@
|
||||
from typing import Any, Protocol
|
||||
from uuid import UUID
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
from knowledge_adapter.content_budget import truncate_content
|
||||
from knowledge_adapter.models import RetrievalRequest, WeKnoraChunk
|
||||
|
||||
|
||||
class InvalidRequestError(Exception):
|
||||
"""A client request violates an adapter-owned policy."""
|
||||
|
||||
|
||||
class SearchClient(Protocol):
|
||||
async def search(
|
||||
self,
|
||||
*,
|
||||
query: str,
|
||||
knowledge_base_id: UUID,
|
||||
) -> list[dict[str, Any]]: ...
|
||||
|
||||
|
||||
class ContentProcessor(Protocol):
|
||||
async def process(self, chunk: WeKnoraChunk) -> str: ...
|
||||
|
||||
|
||||
class IdentityContentProcessor:
|
||||
async def process(self, chunk: WeKnoraChunk) -> str:
|
||||
return chunk.content
|
||||
|
||||
|
||||
class RetrievalService:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
search_client: SearchClient,
|
||||
allowed_kb_ids: frozenset[UUID],
|
||||
max_top_k: int = 5,
|
||||
max_chunk_chars: int = 4000,
|
||||
max_total_chars: int = 12000,
|
||||
content_processor: ContentProcessor | None = None,
|
||||
) -> None:
|
||||
self._search_client = search_client
|
||||
self._allowed_kb_ids = allowed_kb_ids
|
||||
self._max_top_k = max_top_k
|
||||
self._max_chunk_chars = max_chunk_chars
|
||||
self._max_total_chars = max_total_chars
|
||||
self._content_processor = content_processor or IdentityContentProcessor()
|
||||
|
||||
async def retrieve(
|
||||
self,
|
||||
request: RetrievalRequest,
|
||||
) -> dict[str, list[dict[str, object]]]:
|
||||
if request.knowledge_id not in self._allowed_kb_ids:
|
||||
raise InvalidRequestError
|
||||
raw_results = await self._search_client.search(
|
||||
query=request.query,
|
||||
knowledge_base_id=request.knowledge_id,
|
||||
)
|
||||
records: list[dict[str, object]] = []
|
||||
returned_chars = 0
|
||||
effective_top_k = min(request.retrieval_setting.top_k, self._max_top_k)
|
||||
for raw_chunk in raw_results:
|
||||
try:
|
||||
chunk = WeKnoraChunk.model_validate(raw_chunk)
|
||||
except ValidationError:
|
||||
continue
|
||||
if chunk.knowledge_base_id != request.knowledge_id:
|
||||
continue
|
||||
if chunk.score < request.retrieval_setting.score_threshold:
|
||||
continue
|
||||
processed_content = await self._content_processor.process(chunk)
|
||||
content = truncate_content(
|
||||
processed_content,
|
||||
max_chars=self._max_chunk_chars,
|
||||
)
|
||||
if returned_chars + len(content) > self._max_total_chars:
|
||||
break
|
||||
records.append(
|
||||
{
|
||||
"content": content,
|
||||
"score": chunk.score,
|
||||
"title": chunk.knowledge_title or chunk.knowledge_filename,
|
||||
"metadata": {
|
||||
"knowledge_base_id": str(chunk.knowledge_base_id),
|
||||
"knowledge_id": str(chunk.knowledge_id),
|
||||
"chunk_id": str(chunk.id),
|
||||
"chunk_index": chunk.chunk_index,
|
||||
"chunk_type": chunk.chunk_type,
|
||||
},
|
||||
}
|
||||
)
|
||||
returned_chars += len(content)
|
||||
if len(records) >= effective_top_k:
|
||||
break
|
||||
return {"records": records}
|
||||
@@ -0,0 +1,138 @@
|
||||
from pathlib import Path
|
||||
from typing import Annotated, Any, Literal, Self
|
||||
from uuid import UUID
|
||||
|
||||
from pydantic import AnyHttpUrl, Field, SecretStr, field_validator, model_validator
|
||||
from pydantic_settings import BaseSettings, NoDecode, SettingsConfigDict
|
||||
|
||||
from knowledge_adapter.content_budget import TRUNCATION_MARKER
|
||||
|
||||
|
||||
class SecretUnavailableError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
model_config = SettingsConfigDict(
|
||||
env_file=None,
|
||||
case_sensitive=False,
|
||||
extra="ignore",
|
||||
populate_by_name=True,
|
||||
)
|
||||
|
||||
dify_key_hash_current: str = Field(alias="DIFY_KEY_HASH_CURRENT")
|
||||
dify_key_hash_next: str | None = Field(default=None, alias="DIFY_KEY_HASH_NEXT")
|
||||
weknora_base_url: AnyHttpUrl = Field(alias="WEKNORA_BASE_URL")
|
||||
weknora_api_key_file: Path = Field(
|
||||
default=Path("/run/secrets/weknora_api_key"),
|
||||
alias="WEKNORA_API_KEY_FILE",
|
||||
)
|
||||
weknora_allowed_kb_ids: Annotated[frozenset[UUID], NoDecode] = Field(
|
||||
alias="WEKNORA_ALLOWED_KB_IDS"
|
||||
)
|
||||
weknora_tenant_id: int = Field(gt=0, alias="WEKNORA_TENANT_ID")
|
||||
weknora_oss_storage_backend_id: str = Field(
|
||||
min_length=1,
|
||||
max_length=36,
|
||||
alias="WEKNORA_OSS_STORAGE_BACKEND_ID",
|
||||
)
|
||||
oss_region: str = Field(
|
||||
min_length=3,
|
||||
max_length=32,
|
||||
pattern=r"^[a-z0-9]+(?:-[a-z0-9]+)+$",
|
||||
alias="OSS_REGION",
|
||||
)
|
||||
oss_bucket: str = Field(
|
||||
min_length=3,
|
||||
max_length=63,
|
||||
pattern=r"^[a-z0-9][a-z0-9-]{1,61}[a-z0-9]$",
|
||||
alias="OSS_BUCKET",
|
||||
)
|
||||
oss_credential_mode: Literal["ecs_ram_role", "static_ak"] = Field(
|
||||
default="ecs_ram_role",
|
||||
alias="OSS_CREDENTIAL_MODE",
|
||||
)
|
||||
oss_ram_role_name: str | None = Field(
|
||||
default=None,
|
||||
min_length=1,
|
||||
max_length=64,
|
||||
pattern=r"^[A-Za-z0-9.@_-]+$",
|
||||
alias="OSS_RAM_ROLE_NAME",
|
||||
)
|
||||
oss_access_key_id_file: Path = Field(
|
||||
default=Path("/run/secrets/oss_access_key_id"),
|
||||
alias="OSS_ACCESS_KEY_ID_FILE",
|
||||
)
|
||||
oss_access_key_secret_file: Path = Field(
|
||||
default=Path("/run/secrets/oss_access_key_secret"),
|
||||
alias="OSS_ACCESS_KEY_SECRET_FILE",
|
||||
)
|
||||
asset_public_base_url: AnyHttpUrl = Field(alias="ASSET_PUBLIC_BASE_URL")
|
||||
|
||||
max_query_chars: int = Field(default=2000, ge=1, alias="MAX_QUERY_CHARS")
|
||||
max_top_k: int = Field(default=5, ge=1, alias="MAX_TOP_K")
|
||||
max_chunk_chars: int = Field(
|
||||
default=4000,
|
||||
ge=len(TRUNCATION_MARKER),
|
||||
alias="MAX_CHUNK_CHARS",
|
||||
)
|
||||
max_total_chars: int = Field(default=12000, ge=1, alias="MAX_TOTAL_CHARS")
|
||||
weknora_timeout_seconds: float = Field(
|
||||
default=15,
|
||||
gt=0,
|
||||
alias="WEKNORA_TIMEOUT_SECONDS",
|
||||
)
|
||||
|
||||
asset_database_url_file: Path = Path("/run/secrets/asset_database_url")
|
||||
weknora_readonly_database_url_file: Path = Path(
|
||||
"/run/secrets/weknora_readonly_database_url"
|
||||
)
|
||||
|
||||
@field_validator("dify_key_hash_next", mode="before")
|
||||
@classmethod
|
||||
def empty_next_hash_is_none(cls, value: Any) -> Any:
|
||||
if isinstance(value, str) and not value.strip():
|
||||
return None
|
||||
return value
|
||||
|
||||
@field_validator("dify_key_hash_current", "dify_key_hash_next")
|
||||
@classmethod
|
||||
def validate_key_hash(cls, value: str | None) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
normalized = value.lower()
|
||||
if len(normalized) != 64 or any(char not in "0123456789abcdef" for char in normalized):
|
||||
raise ValueError("Dify key hashes must be 64 hexadecimal characters")
|
||||
return normalized
|
||||
|
||||
@field_validator("weknora_allowed_kb_ids", mode="before")
|
||||
@classmethod
|
||||
def parse_allowed_kb_ids(cls, value: Any) -> Any:
|
||||
if isinstance(value, str):
|
||||
return frozenset(UUID(item.strip()) for item in value.split(",") if item.strip())
|
||||
return value
|
||||
|
||||
@model_validator(mode="after")
|
||||
def require_ram_role_for_ram_role_mode(self) -> Self:
|
||||
if self.oss_credential_mode == "ecs_ram_role" and self.oss_ram_role_name is None:
|
||||
raise ValueError("OSS RAM role name is required")
|
||||
return self
|
||||
|
||||
@staticmethod
|
||||
def _read_secret(path: Path) -> SecretStr:
|
||||
try:
|
||||
value = path.read_text(encoding="utf-8").strip()
|
||||
except OSError:
|
||||
raise SecretUnavailableError("required secret unavailable") from None
|
||||
if not value:
|
||||
raise SecretUnavailableError("required secret unavailable")
|
||||
return SecretStr(value)
|
||||
|
||||
def read_weknora_api_key(self) -> SecretStr:
|
||||
return self._read_secret(self.weknora_api_key_file)
|
||||
|
||||
def read_oss_access_key_id(self) -> SecretStr:
|
||||
return self._read_secret(self.oss_access_key_id_file)
|
||||
|
||||
def read_oss_access_key_secret(self) -> SecretStr:
|
||||
return self._read_secret(self.oss_access_key_secret_file)
|
||||
@@ -0,0 +1,89 @@
|
||||
from types import TracebackType
|
||||
from uuid import UUID
|
||||
|
||||
import httpx
|
||||
from pydantic import SecretStr
|
||||
|
||||
from knowledge_adapter.observability import current_request_id
|
||||
|
||||
|
||||
class WeKnoraUnavailableError(Exception):
|
||||
"""WeKnora could not be reached within the configured budget."""
|
||||
|
||||
|
||||
class WeKnoraBadResponseError(Exception):
|
||||
"""WeKnora returned an unusable response."""
|
||||
|
||||
|
||||
class WeKnoraClient:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str,
|
||||
api_key: SecretStr,
|
||||
timeout_seconds: float,
|
||||
transport: httpx.AsyncBaseTransport | None = None,
|
||||
) -> None:
|
||||
short_timeout = min(5.0, timeout_seconds)
|
||||
timeout = httpx.Timeout(
|
||||
timeout_seconds,
|
||||
connect=short_timeout,
|
||||
read=timeout_seconds,
|
||||
write=short_timeout,
|
||||
pool=short_timeout,
|
||||
)
|
||||
limits = httpx.Limits(max_connections=100, max_keepalive_connections=20)
|
||||
self._client = httpx.AsyncClient(
|
||||
base_url=base_url.rstrip("/"),
|
||||
headers={"X-API-Key": api_key.get_secret_value()},
|
||||
timeout=timeout,
|
||||
limits=limits,
|
||||
transport=transport,
|
||||
trust_env=False,
|
||||
)
|
||||
|
||||
async def __aenter__(self) -> "WeKnoraClient":
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> None:
|
||||
await self.aclose()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self._client.aclose()
|
||||
|
||||
async def search(
|
||||
self,
|
||||
*,
|
||||
query: str,
|
||||
knowledge_base_id: UUID,
|
||||
) -> list[dict[str, object]]:
|
||||
request_id = current_request_id()
|
||||
headers = {"X-Request-ID": request_id} if request_id is not None else None
|
||||
try:
|
||||
response = await self._client.post(
|
||||
"/api/v1/knowledge-search",
|
||||
headers=headers,
|
||||
json={
|
||||
"query": query,
|
||||
"knowledge_base_ids": [str(knowledge_base_id)],
|
||||
},
|
||||
)
|
||||
except (httpx.TimeoutException, httpx.NetworkError) as exc:
|
||||
raise WeKnoraUnavailableError("WeKnora unavailable") from exc
|
||||
if not response.is_success:
|
||||
raise WeKnoraBadResponseError("invalid WeKnora response")
|
||||
try:
|
||||
payload = response.json()
|
||||
except ValueError as exc:
|
||||
raise WeKnoraBadResponseError("invalid WeKnora response") from exc
|
||||
if not isinstance(payload, dict) or payload.get("success") is not True:
|
||||
raise WeKnoraBadResponseError("invalid WeKnora response")
|
||||
data = payload.get("data")
|
||||
if not isinstance(data, list) or any(not isinstance(item, dict) for item in data):
|
||||
raise WeKnoraBadResponseError("invalid WeKnora response")
|
||||
return data
|
||||
@@ -0,0 +1,100 @@
|
||||
from collections.abc import Mapping
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Protocol
|
||||
|
||||
from knowledge_adapter.resource_locator import (
|
||||
InvalidResourceLocator,
|
||||
InvalidResourceReference,
|
||||
parse_oss_object_key,
|
||||
parse_resource_reference,
|
||||
)
|
||||
|
||||
_RESOURCE_QUERY = """
|
||||
SELECT tenant_id, storage_backend_id, provider, physical_path, mime_type
|
||||
FROM resources
|
||||
WHERE handle = $1
|
||||
AND tenant_id = $2
|
||||
AND storage_backend_id = $3
|
||||
AND provider = 'oss'
|
||||
AND state = 'active'
|
||||
AND deleted_at IS NULL
|
||||
LIMIT 2
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ResolvedAsset:
|
||||
object_key: str
|
||||
content_type: str | None
|
||||
|
||||
|
||||
class ResourceResolutionError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class ReadonlyConnection(Protocol):
|
||||
def transaction(self, *, readonly: bool) -> AbstractAsyncContextManager[None]: ...
|
||||
|
||||
async def fetch(self, query: str, *args: object) -> list[Mapping[str, Any]]: ...
|
||||
|
||||
|
||||
class ReadonlyPool(Protocol):
|
||||
def acquire(self) -> AbstractAsyncContextManager[ReadonlyConnection]: ...
|
||||
|
||||
|
||||
class WeKnoraResourceResolver:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
pool: ReadonlyPool,
|
||||
tenant_id: int,
|
||||
storage_backend_id: str,
|
||||
bucket: str,
|
||||
) -> None:
|
||||
self._pool = pool
|
||||
self._tenant_id = tenant_id
|
||||
self._storage_backend_id = storage_backend_id
|
||||
self._bucket = bucket
|
||||
|
||||
async def resolve(self, source_ref: str) -> ResolvedAsset:
|
||||
try:
|
||||
handle = parse_resource_reference(source_ref)
|
||||
except InvalidResourceReference as error:
|
||||
raise ResourceResolutionError("resource unavailable") from error
|
||||
try:
|
||||
async with self._pool.acquire() as connection:
|
||||
async with connection.transaction(readonly=True):
|
||||
rows = await connection.fetch(
|
||||
_RESOURCE_QUERY,
|
||||
handle,
|
||||
self._tenant_id,
|
||||
self._storage_backend_id,
|
||||
)
|
||||
except Exception as error:
|
||||
raise ResourceResolutionError("resource unavailable") from error
|
||||
if len(rows) != 1:
|
||||
raise ResourceResolutionError("resource unavailable")
|
||||
row = rows[0]
|
||||
if (
|
||||
row["tenant_id"] != self._tenant_id
|
||||
or str(row["storage_backend_id"]) != self._storage_backend_id
|
||||
or row["provider"] != "oss"
|
||||
):
|
||||
raise ResourceResolutionError("resource unavailable")
|
||||
try:
|
||||
object_key = parse_oss_object_key(
|
||||
provider=str(row["provider"]),
|
||||
physical_path=str(row["physical_path"]),
|
||||
storage_backend_id=str(row["storage_backend_id"]),
|
||||
expected_storage_backend_id=self._storage_backend_id,
|
||||
expected_bucket=self._bucket,
|
||||
)
|
||||
except InvalidResourceLocator as error:
|
||||
raise ResourceResolutionError("resource unavailable") from error
|
||||
|
||||
mime_type = row["mime_type"]
|
||||
return ResolvedAsset(
|
||||
object_key=object_key,
|
||||
content_type=str(mime_type) if mime_type else None,
|
||||
)
|
||||
Reference in New Issue
Block a user