418 lines
14 KiB
Python
418 lines
14 KiB
Python
"""Knowledge Base Management API routes"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hmac
|
|
import logging
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime, timezone
|
|
from typing import Any
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, Request, Security, UploadFile, File
|
|
from fastapi.security import APIKeyHeader, APIKeyQuery
|
|
from pydantic import BaseModel
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter(tags=["kb-management"])
|
|
|
|
MAX_UPLOAD_SIZE = 50 * 1024 * 1024 # 50MB
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# API Key Authentication
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
|
|
_api_key_query = APIKeyQuery(name="api_key", auto_error=False)
|
|
|
|
|
|
async def _verify_api_key(
|
|
request: Request,
|
|
api_key_header: str | None = Security(_api_key_header),
|
|
api_key_query: str | None = Security(_api_key_query),
|
|
) -> None:
|
|
"""Verify API key for REST endpoints. Raises HTTPException if invalid."""
|
|
configured_api_key: str | None = None
|
|
if hasattr(request.app.state, "server_config") and request.app.state.server_config:
|
|
configured_api_key = request.app.state.server_config.api_key
|
|
if configured_api_key is None and hasattr(request.app.state, "api_key"):
|
|
configured_api_key = request.app.state.api_key
|
|
|
|
# If no API key is configured, allow all requests (backwards compat)
|
|
if configured_api_key is None:
|
|
return
|
|
|
|
provided = api_key_header or api_key_query
|
|
if not hmac.compare_digest((provided or "").encode(), configured_api_key.encode()):
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="Invalid or missing API key. Provide via X-API-Key header or api_key query parameter.",
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# In-memory Knowledge Source Store
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class KnowledgeSource:
|
|
id: str
|
|
name: str
|
|
type: str # "local" | "feishu" | "confluence" | "http"
|
|
config: dict[str, Any] = field(default_factory=dict)
|
|
status: str = "active"
|
|
document_count: int = 0
|
|
last_synced: str | None = None
|
|
|
|
|
|
@dataclass
|
|
class UploadedDocument:
|
|
document_id: str
|
|
filename: str
|
|
source_id: str
|
|
chunks: int
|
|
status: str
|
|
created_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
|
|
|
|
|
|
class KnowledgeSourceStore:
|
|
def __init__(self, max_sources: int = 500):
|
|
self._sources: dict[str, KnowledgeSource] = {}
|
|
self._documents: list[UploadedDocument] = []
|
|
self._max = max_sources
|
|
|
|
def add_source(self, name: str, source_type: str, config: dict[str, Any]) -> KnowledgeSource:
|
|
source_id = str(uuid.uuid4())
|
|
source = KnowledgeSource(
|
|
id=source_id,
|
|
name=name,
|
|
type=source_type,
|
|
config=config,
|
|
)
|
|
self._sources[source_id] = source
|
|
return source
|
|
|
|
def get_source(self, source_id: str) -> KnowledgeSource | None:
|
|
return self._sources.get(source_id)
|
|
|
|
def remove_source(self, source_id: str) -> bool:
|
|
if source_id in self._sources:
|
|
del self._sources[source_id]
|
|
return True
|
|
return False
|
|
|
|
def list_sources(self) -> list[KnowledgeSource]:
|
|
return list(self._sources.values())
|
|
|
|
def add_document(self, doc: UploadedDocument) -> None:
|
|
self._documents.append(doc)
|
|
# Update source document count
|
|
source = self._sources.get(doc.source_id)
|
|
if source:
|
|
source.document_count += 1
|
|
source.last_synced = doc.created_at
|
|
|
|
def list_documents(self, source_id: str | None = None) -> list[UploadedDocument]:
|
|
if source_id:
|
|
return [d for d in self._documents if d.source_id == source_id]
|
|
return list(self._documents)
|
|
|
|
def delete_document(self, document_id: str) -> bool:
|
|
for i, doc in enumerate(self._documents):
|
|
if doc.document_id == document_id:
|
|
self._documents.pop(i)
|
|
# Update source document count
|
|
source = self._sources.get(doc.source_id)
|
|
if source:
|
|
source.document_count = max(0, source.document_count - 1)
|
|
return True
|
|
return False
|
|
|
|
def update_source(self, source_id: str, data: dict[str, Any]) -> KnowledgeSource | None:
|
|
source = self._sources.get(source_id)
|
|
if source is None:
|
|
return None
|
|
if "name" in data:
|
|
source.name = data["name"]
|
|
if "config" in data:
|
|
source.config = data["config"]
|
|
if "status" in data:
|
|
source.status = data["status"]
|
|
return source
|
|
|
|
|
|
# Module-level singleton
|
|
_source_store = KnowledgeSourceStore()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Request / Response models
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class AddSourceRequest(BaseModel):
|
|
name: str
|
|
type: str # "local" | "feishu" | "confluence" | "http"
|
|
config: dict[str, Any] = {}
|
|
|
|
|
|
class SearchRequest(BaseModel):
|
|
query: str
|
|
sources: list[str] | None = None
|
|
source_ids: list[str] | None = None
|
|
top_k: int = 5
|
|
strategy: str = "vector" # "vector" | "hybrid"
|
|
|
|
|
|
class SearchResult(BaseModel):
|
|
content: str
|
|
source: str
|
|
score: float
|
|
metadata: dict[str, Any] = {}
|
|
|
|
|
|
class UpdateSourceRequest(BaseModel):
|
|
name: str | None = None
|
|
config: dict[str, Any] | None = None
|
|
status: str | None = None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Endpoints
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@router.get("/kb-management/sources")
|
|
async def list_sources(req: Request, _auth: None = Depends(_verify_api_key)):
|
|
"""List all knowledge sources."""
|
|
sources = _source_store.list_sources()
|
|
return {
|
|
"sources": [
|
|
{
|
|
"id": s.id,
|
|
"name": s.name,
|
|
"type": s.type,
|
|
"status": s.status,
|
|
"document_count": s.document_count,
|
|
"last_synced": s.last_synced,
|
|
}
|
|
for s in sources
|
|
]
|
|
}
|
|
|
|
|
|
@router.post("/kb-management/sources", status_code=201)
|
|
async def add_source(
|
|
request: AddSourceRequest, req: Request, _auth: None = Depends(_verify_api_key)
|
|
):
|
|
"""Add a knowledge source."""
|
|
valid_types = {"local", "feishu", "confluence", "http"}
|
|
if request.type not in valid_types:
|
|
raise HTTPException(
|
|
status_code=422,
|
|
detail=f"Invalid source type '{request.type}'. Must be one of: {valid_types}",
|
|
)
|
|
|
|
source = _source_store.add_source(
|
|
name=request.name,
|
|
source_type=request.type,
|
|
config=request.config,
|
|
)
|
|
return {
|
|
"id": source.id,
|
|
"name": source.name,
|
|
"type": source.type,
|
|
"status": source.status,
|
|
"document_count": source.document_count,
|
|
"last_synced": source.last_synced,
|
|
}
|
|
|
|
|
|
@router.delete("/kb-management/sources/{source_id}")
|
|
async def remove_source(source_id: str, req: Request, _auth: None = Depends(_verify_api_key)):
|
|
"""Remove a knowledge source."""
|
|
if not _source_store.remove_source(source_id):
|
|
raise HTTPException(status_code=404, detail=f"Source '{source_id}' not found")
|
|
return {"status": "removed"}
|
|
|
|
|
|
@router.post("/kb-management/sources/{source_id}/sync")
|
|
async def sync_source(source_id: str, _auth: None = Depends(_verify_api_key)):
|
|
"""Trigger source sync."""
|
|
source = _source_store.get_source(source_id)
|
|
if not source:
|
|
raise HTTPException(status_code=404, detail="Source not found")
|
|
# For now, return sync status (actual sync logic depends on source type)
|
|
return {"status": "syncing", "source_id": source_id}
|
|
|
|
|
|
@router.put("/kb-management/sources/{source_id}")
|
|
async def update_source(
|
|
source_id: str, data: UpdateSourceRequest, _auth: None = Depends(_verify_api_key)
|
|
):
|
|
"""Update source config."""
|
|
update_data = {k: v for k, v in data.model_dump().items() if v is not None}
|
|
source = _source_store.update_source(source_id, update_data)
|
|
if not source:
|
|
raise HTTPException(status_code=404, detail="Source not found")
|
|
return {
|
|
"id": source.id,
|
|
"name": source.name,
|
|
"type": source.type,
|
|
"status": source.status,
|
|
"document_count": source.document_count,
|
|
"last_synced": source.last_synced,
|
|
}
|
|
|
|
|
|
@router.get("/kb-management/documents")
|
|
async def list_documents(source_id: str | None = None, _auth: None = Depends(_verify_api_key)):
|
|
"""List all documents, optionally filtered by source_id."""
|
|
documents = _source_store.list_documents(source_id)
|
|
return {
|
|
"documents": [
|
|
{
|
|
"document_id": d.document_id,
|
|
"filename": d.filename,
|
|
"source_id": d.source_id,
|
|
"chunks": d.chunks,
|
|
"status": d.status,
|
|
"created_at": d.created_at,
|
|
}
|
|
for d in documents
|
|
]
|
|
}
|
|
|
|
|
|
@router.delete("/kb-management/documents/{document_id}")
|
|
async def delete_document(document_id: str, _auth: None = Depends(_verify_api_key)):
|
|
"""Delete a document by ID."""
|
|
success = _source_store.delete_document(document_id)
|
|
if not success:
|
|
raise HTTPException(status_code=404, detail="Document not found")
|
|
return {"status": "deleted"}
|
|
|
|
|
|
@router.post("/kb-management/documents/upload")
|
|
async def upload_document(
|
|
req: Request,
|
|
file: UploadFile = File(...),
|
|
source_id: str = "",
|
|
_auth: None = Depends(_verify_api_key),
|
|
):
|
|
"""Upload a document to the knowledge base."""
|
|
if not file.filename:
|
|
raise HTTPException(status_code=422, detail="Filename is required")
|
|
|
|
# Try to use DocumentLoader if available
|
|
chunks = 1
|
|
try:
|
|
from agentkit.memory.document_loader import DocumentLoader
|
|
|
|
content = await file.read(MAX_UPLOAD_SIZE + 1)
|
|
if len(content) > MAX_UPLOAD_SIZE:
|
|
raise HTTPException(
|
|
status_code=413,
|
|
detail=f"File too large. Maximum size is {MAX_UPLOAD_SIZE // (1024 * 1024)}MB",
|
|
)
|
|
loader = DocumentLoader()
|
|
doc = loader.load_bytes(content, file.filename)
|
|
# Estimate chunks based on content length (rough approximation)
|
|
chunks = max(1, len(doc.content) // 500)
|
|
except ImportError:
|
|
# DocumentLoader not available, use basic estimation
|
|
content = await file.read(MAX_UPLOAD_SIZE + 1)
|
|
if len(content) > MAX_UPLOAD_SIZE:
|
|
raise HTTPException(
|
|
status_code=413,
|
|
detail=f"File too large. Maximum size is {MAX_UPLOAD_SIZE // (1024 * 1024)}MB",
|
|
)
|
|
chunks = max(1, len(content) // 500)
|
|
except Exception as e:
|
|
logger.warning(f"Document parsing failed: {e}")
|
|
chunks = 1
|
|
|
|
# Determine source_id - use "local" default if not provided
|
|
effective_source_id = source_id or "local"
|
|
|
|
# Ensure a local source exists
|
|
if effective_source_id == "local":
|
|
local_sources = [s for s in _source_store.list_sources() if s.type == "local"]
|
|
if not local_sources:
|
|
_source_store.add_source("本地文档", "local", {})
|
|
|
|
uploaded = UploadedDocument(
|
|
document_id=str(uuid.uuid4()),
|
|
filename=file.filename,
|
|
source_id=effective_source_id,
|
|
chunks=chunks,
|
|
status="indexed",
|
|
)
|
|
_source_store.add_document(uploaded)
|
|
|
|
return {
|
|
"document_id": uploaded.document_id,
|
|
"filename": uploaded.filename,
|
|
"source_id": uploaded.source_id,
|
|
"chunks": uploaded.chunks,
|
|
"status": uploaded.status,
|
|
"created_at": uploaded.created_at if hasattr(uploaded, "created_at") else "",
|
|
}
|
|
|
|
|
|
@router.post("/kb-management/search")
|
|
async def search_knowledge(
|
|
request: SearchRequest, req: Request, _auth: None = Depends(_verify_api_key)
|
|
):
|
|
"""Test search/retrieval against the knowledge base."""
|
|
# Try to use semantic memory if available
|
|
memory_retriever = getattr(req.app.state, "memory_retriever", None)
|
|
if (
|
|
memory_retriever
|
|
and hasattr(memory_retriever, "semantic_memory")
|
|
and memory_retriever.semantic_memory
|
|
):
|
|
try:
|
|
results = await memory_retriever.semantic_memory.retrieve(
|
|
query=request.query,
|
|
top_k=request.top_k,
|
|
)
|
|
return {
|
|
"results": [
|
|
{
|
|
"content": r.content if hasattr(r, "content") else str(r),
|
|
"source": r.source_id if hasattr(r, "source_id") else "",
|
|
"score": r.score if hasattr(r, "score") else 0.0,
|
|
"metadata": r.metadata if hasattr(r, "metadata") else {},
|
|
}
|
|
for r in results
|
|
]
|
|
}
|
|
except Exception as e:
|
|
logger.warning(f"Semantic search failed: {e}")
|
|
|
|
# Fallback: return empty results with a hint
|
|
return {
|
|
"results": [],
|
|
"message": "未配置语义检索服务，请先配置知识库连接",
|
|
}
|
|
|
|
|
|
@router.get("/kb-management/sources/{source_id}/health")
|
|
async def check_source_health(source_id: str, req: Request, _auth: None = Depends(_verify_api_key)):
|
|
"""Check the health of a knowledge source."""
|
|
source = _source_store.get_source(source_id)
|
|
if source is None:
|
|
raise HTTPException(status_code=404, detail=f"Source '{source_id}' not found")
|
|
|
|
# For local sources, always healthy
|
|
if source.type == "local":
|
|
return {"source_id": source_id, "status": "healthy", "message": "本地存储正常"}
|
|
|
|
# For external sources, try to check connectivity
|
|
# This would use the actual KB adapters in production
|
|
return {"source_id": source_id, "status": "unknown", "message": "健康检查未实现"}
|