kris/ForcePilot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

stylle: 使用更加严格的风格检查以及代码格式化方法

Browse Source
This commit is contained in:
Wenjie Zhang 2025-09-01 22:37:03 +08:00
parent f16dc7ae49
commit 7641ab3699
56 changed files with 1211 additions and 1279 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
Makefile
View File

@ -12,3 +12,16 @@ logs:
@echo "\n\nBranch: $$(git branch --show-current)"
@echo "Commit ID: $$(git rev-parse HEAD)"
@echo "System: $$(uname -a)"
######################
# LINTING AND FORMATTING
######################
lint:
uv run python -m ruff check .
uv run python -m ruff format src --diff
uv run python -m ruff check --select I src
format format_diff:
uv run ruff format
uv run ruff check --select I --fix

3
pyproject.toml
View File

@ -54,14 +54,13 @@ dependencies = [
"tabulate>=0.9.0",
]
[tool.ruff]
line-length = 210 # 代码最大行宽
line-length = 120 # 代码最大行宽
lint.select = [ # 选择的规则
"F",
"E",
"W",
"UP",
]
lint.ignore = ["F401", "E501"] # 忽略的规则
[dependency-groups]
dev = [

88
scripts/batch_upload.py
View File

@ -1,4 +1,3 @@
import asyncio
import hashlib
import json
@ -7,7 +6,7 @@ import pathlib
import httpx
import typer
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn, TimeElapsedColumn
from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
app = typer.Typer()
console = Console()
@ -44,12 +43,14 @@ async def upload_file(
f"{base_url}/knowledge/files/upload",
params={"db_id": db_id},
files=files,
timeout=300, # 5 minutes timeout for large files
timeout=300, # 5 minutes timeout for large files
)
response.raise_for_status()
return response.json().get("file_path")
except httpx.HTTPStatusError as e:
console.print(f"[bold red]Failed to upload {file_path.name}: {e.response.status_code} - {e.response.text}[/bold red]")
console.print(
f"[bold red]Failed to upload {file_path.name}: {e.response.status_code} - {e.response.text}[/bold red]"
)
return None
except httpx.RequestError as e:
console.print(f"[bold red]Failed to upload {file_path.name}: {e}[/bold red]")
@ -75,21 +76,23 @@ async def process_document(
"enable_ocr": enable_ocr,
"use_qa_split": use_qa_split,
"qa_separator": qa_separator,
"content_type": "file"
"content_type": "file",
}
try:
response = await client.post(
f"{base_url}/knowledge/databases/{db_id}/documents",
json={"items": [server_file_path], "params": params},
timeout=600, # 10 minutes timeout for processing
timeout=600, # 10 minutes timeout for processing
)
response.raise_for_status()
result = response.json()
# Check if the overall request was successful
if result.get("status") != "success":
console.print(f"[bold yellow]Processing warning for {server_file_path}: {result.get('message')}[/bold yellow]")
console.print(
f"[bold yellow]Processing warning for {server_file_path}: {result.get('message')}[/bold yellow]"
)
return False
# Check the specific file's processing status in the items array
@ -129,7 +132,9 @@ async def process_document(
return False
except httpx.HTTPStatusError as e:
console.print(f"[bold red]Failed to process {server_file_path}: {e.response.status_code} - {e.response.text}[/bold red]")
console.print(
f"[bold red]Failed to process {server_file_path}: {e.response.status_code} - {e.response.text}[/bold red]"
)
return False
except httpx.RequestError as e:
console.print(f"[bold red]Failed to process {server_file_path}: {e}[/bold red]")
@ -159,17 +164,20 @@ async def worker(
progress.update(upload_task_id, advance=1, postfix=f"Uploaded {file_path.name}")
if not server_file_path:
progress.update(process_task_id, advance=1) # Mark as processed to not hang the progress bar
progress.update(process_task_id, advance=1) # Mark as processed to not hang the progress bar
return file_path, file_hash, "upload_failed"
# 2. Process file
success = await process_document(
client, base_url, db_id, server_file_path,
client,
base_url,
db_id,
server_file_path,
enable_ocr=enable_ocr,
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
use_qa_split=use_qa_split,
qa_separator=qa_separator
qa_separator=qa_separator,
)
progress.update(process_task_id, advance=1, postfix=f"Processed {file_path.name}")
@ -193,7 +201,7 @@ def load_processed_files(record_file: pathlib.Path) -> set[str]:
try:
with open(record_file) as f:
data = json.load(f)
return set(data.get('processed_files', []))
return set(data.get("processed_files", []))
except (OSError, json.JSONDecodeError) as e:
console.print(f"[bold yellow]Warning: Could not load processed files record: {e}[/bold yellow]")
return set()
@ -205,8 +213,8 @@ def save_processed_files(record_file: pathlib.Path, processed_files: set[str]):
record_file.parent.mkdir(parents=True, exist_ok=True)
try:
with open(record_file, 'w') as f:
json.dump({'processed_files': list(processed_files)}, f, indent=2)
with open(record_file, "w") as f:
json.dump({"processed_files": list(processed_files)}, f, indent=2)
except OSError as e:
console.print(f"[bold red]Error: Could not save processed files record: {e}[/bold red]")
@ -232,7 +240,9 @@ async def convert_to_markdown(
console.print(f"[bold red]Failed to convert {server_file_path}: {result.get('message')}[/bold red]")
return None
except httpx.HTTPStatusError as e:
console.print(f"[bold red]Failed to convert {server_file_path}: {e.response.status_code} - {e.response.text}[/bold red]")
console.print(
f"[bold red]Failed to convert {server_file_path}: {e.response.status_code} - {e.response.text}[/bold red]"
)
return None
except httpx.RequestError as e:
console.print(f"[bold red]Request failed for {server_file_path}: {e}[/bold red]")
@ -267,7 +277,7 @@ async def trans_worker(
try:
output_path = output_dir / file_path.with_suffix(".md").name
output_path.parent.mkdir(parents=True, exist_ok=True)
with open(output_path, 'w', encoding='utf-8') as f:
with open(output_path, "w", encoding="utf-8") as f:
f.write(markdown_content)
progress.update(task_id, advance=1, postfix=f"[green]Converted {file_path.name}[/green]")
return file_path, "success"
@ -280,14 +290,18 @@ async def trans_worker(
@app.command()
def upload(
db_id: str = typer.Option(..., help="The ID of the knowledge base."),
directory: pathlib.Path = typer.Option(..., help="The directory containing files to upload.", exists=True, file_okay=False),
directory: pathlib.Path = typer.Option(
..., help="The directory containing files to upload.", exists=True, file_okay=False
),
pattern: str = typer.Option("*.md", help="The glob pattern for files to upload (e.g., '*.pdf', '**/*.txt')."),
base_url: str = typer.Option("http://127.0.0.1:5050/api", help="The base URL of the API server."),
username: str = typer.Option(..., help="Admin username for login."),
password: str = typer.Option(..., help="Admin password for login."),
concurrency: int = typer.Option(1, help="The number of concurrent upload/process tasks."),
recursive: bool = typer.Option(False, "--recursive", "-r", help="Search for files recursively in subdirectories."),
record_file: pathlib.Path = typer.Option("scripts/tmp/batch_processed_files.txt", help="File to store processed files record."),
record_file: pathlib.Path = typer.Option(
"scripts/tmp/batch_processed_files.txt", help="File to store processed files record."
),
chunk_size: int = typer.Option(1000, help="Chunk size for document processing."),
chunk_overlap: int = typer.Option(200, help="Chunk overlap for document processing."),
enable_ocr: str = typer.Option("paddlex_ocr", help="OCR engine to use (paddlex_ocr, mineru_ocr, disable)."),
@ -325,7 +339,9 @@ def upload(
files_to_upload.append((file_path, file_hash))
if not files_to_upload:
console.print(f"[bold green]All {len(all_files)} files have already been processed. Nothing to do.[/bold green]")
console.print(
f"[bold green]All {len(all_files)} files have already been processed. Nothing to do.[/bold green]"
)
raise typer.Exit()
console.print(f"Found {len(all_files)} total files:")
@ -361,13 +377,20 @@ def upload(
for file_path, file_hash in files_to_upload:
task = asyncio.create_task(
worker(
semaphore, client, base_url, db_id, file_path, file_hash,
progress, upload_task_id, process_task_id,
semaphore,
client,
base_url,
db_id,
file_path,
file_hash,
progress,
upload_task_id,
process_task_id,
enable_ocr=enable_ocr,
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
use_qa_split=use_qa_split,
qa_separator=qa_separator
qa_separator=qa_separator,
)
)
tasks.append(task)
@ -381,12 +404,12 @@ def upload(
newly_processed_hashes = set()
for file_path, file_hash, status in results:
if status == 'success':
if status == "success":
successful_files.append(file_path)
newly_processed_hashes.add(file_hash)
elif status == 'upload_failed':
elif status == "upload_failed":
upload_failures.append(file_path)
elif status == 'processing_failed':
elif status == "processing_failed":
processing_failures.append(file_path)
# Don't add to processed files if processing failed
@ -394,7 +417,9 @@ def upload(
if newly_processed_hashes:
all_processed_files = processed_files | newly_processed_hashes
save_processed_files(record_file, all_processed_files)
console.print(f"[bold green]Updated processed files record with {len(newly_processed_hashes)} new entries.[/bold green]")
console.print(
f"[bold green]Updated processed files record with {len(newly_processed_hashes)} new entries.[/bold green]"
)
console.print("[bold green]Batch operation complete.[/bold green]")
console.print(f" - [green]Successful:[/green] {len(successful_files)}")
@ -413,7 +438,9 @@ def upload(
@app.command()
def trans(
db_id: str = typer.Option(..., help="The ID of the knowledge base (for temporary file upload)."),
directory: pathlib.Path = typer.Option(..., help="The directory containing files to convert.", exists=True, file_okay=False),
directory: pathlib.Path = typer.Option(
..., help="The directory containing files to convert.", exists=True, file_okay=False
),
output_dir: pathlib.Path = typer.Option("output_markdown", help="The directory to save converted markdown files."),
pattern: str = typer.Option("*.docx", help="The glob pattern for files to convert (e.g., '*.pdf', '*.docx')."),
base_url: str = typer.Option("http://127.0.0.1:5050/api", help="The base URL of the API server."),
@ -467,10 +494,7 @@ def trans(
for file_path in files_to_convert:
task = asyncio.create_task(
trans_worker(
semaphore, client, base_url, db_id, file_path,
output_dir, progress, task_id
)
trans_worker(semaphore, client, base_url, db_id, file_path, output_dir, progress, task_id)
)
tasks.append(task)
@ -481,7 +505,7 @@ def trans(
failed_files = []
for file_path, status in results:
if status == 'success':
if status == "success":
successful_files.append(file_path)
else:
failed_files.append((file_path, status))

21
scripts/preprocessors/split_data_to_subfiles.py
View File

@ -1,26 +1,30 @@
import typer
import pandas as pd
import json
import random
import re
from pathlib import Path
import pandas as pd
import typer
app = typer.Typer()
def sanitize_filename(name: str) -> str:
return re.sub(r'[\\/*?:"<>|]', "_", str(name).strip())
def random_suffix() -> str:
return f"_{random.randint(10000000, 99999999)}"
def read_table(file_path: Path) -> pd.DataFrame:
suffix = file_path.suffix.lower()
if suffix in ['.xlsx', '.xls']:
if suffix in [".xlsx", ".xls"]:
return pd.read_excel(file_path)
elif suffix == '.csv':
elif suffix == ".csv":
return pd.read_csv(file_path)
elif suffix == '.json':
with open(file_path, encoding='utf-8') as f:
elif suffix == ".json":
with open(file_path, encoding="utf-8") as f:
data = json.load(f)
if isinstance(data, list) and len(data) > 1 and isinstance(data[0], dict):
return pd.DataFrame(data)
@ -29,6 +33,7 @@ def read_table(file_path: Path) -> pd.DataFrame:
else:
raise ValueError(f"不支持的文件格式：{suffix}")
def export_txts(df: pd.DataFrame, output_dir: Path, title_field: str = "标题"):
output_dir.mkdir(parents=True, exist_ok=True)
df.columns = [c.strip() for c in df.columns]
@ -54,11 +59,12 @@ def export_txts(df: pd.DataFrame, output_dir: Path, title_field: str = "标题")
typer.echo(f"✅ 成功导出 {len(df)} 个文件到目录：{output_dir}")
@app.command()
def convert(
input_file: Path = typer.Argument(..., help="输入文件（.xlsx/.xls/.csv/.json）"),
out_dir: Path = typer.Option("output", help="输出目录"),
title_field: str = typer.Option("标题", help="标题字段名（用于文件名）")
title_field: str = typer.Option("标题", help="标题字段名（用于文件名）"),
):
"""
将结构化数据文件（Excel/CSV/JSON）转换为多个 .txt 文件。
@ -70,5 +76,6 @@ def convert(
typer.echo(f"❌ 错误：{e}", err=True)
raise typer.Exit(code=1)
if __name__ == "__main__":
app()

14
scripts/rename_milvus_collections.py
View File

@ -1,5 +1,7 @@
import os
from pymilvus import utility, connections, Collection
from pymilvus import Collection, connections, utility
def get_collection_info(collection_name, alias):
"""Safely gets a collection object and its number of entities."""
@ -11,13 +13,14 @@ def get_collection_info(collection_name, alias):
print(f"Error getting info for collection '{collection_name}': {e}")
return None, 0
def rename_and_resolve_duplicates():
"""
Connects to Milvus, renames collections from 'kb_kb_' to 'kb_',
and resolves duplicates by keeping the collection with more rows.
"""
milvus_uri = os.getenv('MILVUS_URI', 'http://localhost:19530')
milvus_token = os.getenv('MILVUS_TOKEN', '')
milvus_uri = os.getenv("MILVUS_URI", "http://localhost:19530")
milvus_token = os.getenv("MILVUS_TOKEN", "")
connection_alias = "rename_script"
try:
@ -26,7 +29,7 @@ def rename_and_resolve_duplicates():
print("Successfully connected to Milvus.")
all_collections = utility.list_collections(using=connection_alias)
collections_to_rename = [c for c in all_collections if c.startswith('kb_kb_')]
collections_to_rename = [c for c in all_collections if c.startswith("kb_kb_")]
if not collections_to_rename:
print("No collections with the prefix 'kb_kb_' found. Nothing to do.")
@ -35,7 +38,7 @@ def rename_and_resolve_duplicates():
print(f"Found {len(collections_to_rename)} collections with 'kb_kb_' prefix to process.")
for old_name in collections_to_rename:
new_name = old_name.replace('kb_kb_', 'kb_', 1)
new_name = old_name.replace("kb_kb_", "kb_", 1)
try:
print(f"Attempting to rename '{old_name}' to '{new_name}'...")
utility.rename_collection(old_name, new_name, using=connection_alias)
@ -74,5 +77,6 @@ def rename_and_resolve_duplicates():
connections.disconnect(connection_alias)
print("Disconnected from Milvus.")
if __name__ == "__main__":
rename_and_resolve_duplicates()

3
scripts/vllm/test_vllm.py
View File

@ -1,4 +1,5 @@
from openai import OpenAI
# Set OpenAI's API key and API base to use vLLM's API server.
openai_api_key = "EMPTY"
openai_api_base = "http://localhost:8080/v1"
@ -13,6 +14,6 @@ chat_response = client.chat.completions.create(
messages=[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Tell me a joke."},
]
],
)
print("Chat response:", chat_response)

13
server/db_manager.py
View File

@ -1,16 +1,18 @@
import os
import pathlib
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from contextlib import contextmanager
from src import config
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from server.models import Base
from server.models.user_model import User
from server.models.thread_model import Thread
from server.models.kb_models import KnowledgeDatabase, KnowledgeFile, KnowledgeNode
from server.models.thread_model import Thread
from server.models.user_model import User
from src import config
from src.utils import logger
class DBManager:
"""数据库管理器 - 只提供基础的数据库连接和会话管理"""
@ -65,5 +67,6 @@ class DBManager:
finally:
session.close()
# 创建全局数据库管理器实例
db_manager = DBManager()

6
server/main.py
View File

@ -1,6 +1,5 @@
import uvicorn
from fastapi import FastAPI, Request, HTTPException, status, Depends
from fastapi import Depends, FastAPI, HTTPException, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
@ -24,6 +23,7 @@ app.add_middleware(
allow_headers=["*"],
)
# 鉴权中间件
class AuthMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: Request, call_next):
@ -56,9 +56,9 @@ class AuthMiddleware(BaseHTTPMiddleware):
# 继续处理请求
return await call_next(request)
# 添加鉴权中间件
app.add_middleware(AuthMiddleware)
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=5050, threads=10, workers=10, reload=True)

27
server/models/kb_models.py
View File

@ -1,13 +1,16 @@
from sqlalchemy import Column, Integer, String, DateTime, JSON, ForeignKey, Text
import time
from sqlalchemy import JSON, Column, DateTime, ForeignKey, Integer, String, Text
from sqlalchemy.orm import relationship
from sqlalchemy.sql import func
import time
from server.models import Base
class KnowledgeDatabase(Base):
"""知识库模型"""
__tablename__ = 'knowledge_databases'
__tablename__ = "knowledge_databases"
id = Column(Integer, primary_key=True, autoincrement=True)
db_id = Column(String, nullable=False, unique=True, index=True) # 数据库ID
@ -31,7 +34,7 @@ class KnowledgeDatabase(Base):
"embed_model": self.embed_model,
"dimension": self.dimension,
"metadata": self.meta_info or {},
"created_at": self.created_at.isoformat() if self.created_at else None
"created_at": self.created_at.isoformat() if self.created_at else None,
}
# 添加文件信息
if self.files:
@ -40,13 +43,15 @@ class KnowledgeDatabase(Base):
result["files"] = {}
return result
class KnowledgeFile(Base):
"""知识库文件模型"""
__tablename__ = 'knowledge_files'
__tablename__ = "knowledge_files"
id = Column(Integer, primary_key=True, autoincrement=True)
file_id = Column(String, nullable=False, index=True) # 文件ID
database_id = Column(String, ForeignKey('knowledge_databases.db_id'), nullable=False) # 所属数据库ID
database_id = Column(String, ForeignKey("knowledge_databases.db_id"), nullable=False) # 所属数据库ID
filename = Column(String, nullable=False) # 文件名
path = Column(String, nullable=False) # 文件路径
file_type = Column(String, nullable=False) # 文件类型
@ -71,18 +76,20 @@ class KnowledgeFile(Base):
"type": self.file_type,
"status": self.status,
"node_count": self.computed_node_count,
"created_at": self.created_at.timestamp() if self.created_at else time.time()
"created_at": self.created_at.timestamp() if self.created_at else time.time(),
}
if with_nodes:
result["nodes"] = [node.to_dict() for node in self.nodes] if self.nodes else []
return result
class KnowledgeNode(Base):
"""知识块模型"""
__tablename__ = 'knowledge_nodes'
__tablename__ = "knowledge_nodes"
id = Column(Integer, primary_key=True, autoincrement=True)
file_id = Column(String, ForeignKey('knowledge_files.file_id'), nullable=False) # 所属文件ID
file_id = Column(String, ForeignKey("knowledge_files.file_id"), nullable=False) # 所属文件ID
text = Column(Text, nullable=False) # 文本内容
hash = Column(String, nullable=True) # 文本哈希值
start_char_idx = Column(Integer, nullable=True) # 开始字符索引
@ -101,5 +108,5 @@ class KnowledgeNode(Base):
"hash": self.hash,
"start_char_idx": self.start_char_idx,
"end_char_idx": self.end_char_idx,
"metadata": self.meta_info or {} # 确保映射正确
"metadata": self.meta_info or {}, # 确保映射正确
}

7
server/models/thread_model.py
View File

@ -1,13 +1,14 @@
from sqlalchemy import Column, String, Integer, DateTime, Text, ForeignKey
from sqlalchemy.sql import func
from sqlalchemy.orm import relationship
from sqlalchemy import Column, DateTime, ForeignKey, Integer, String, Text
from sqlalchemy.dialects.mysql import JSON
from sqlalchemy.orm import relationship
from sqlalchemy.sql import func
from server.models import Base
class Thread(Base):
"""对话线程表"""
__tablename__ = "thread"
id = Column(String(64), primary_key=True, index=True, comment="线程ID")

20
server/models/user_model.py
View File

@ -1,17 +1,19 @@
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Text
from sqlalchemy.sql import func
from sqlalchemy import Column, DateTime, ForeignKey, Integer, String, Text
from sqlalchemy.orm import relationship
from sqlalchemy.sql import func
from server.models import Base
class User(Base):
"""用户模型"""
__tablename__ = 'users'
__tablename__ = "users"
id = Column(Integer, primary_key=True, autoincrement=True)
username = Column(String, nullable=False, unique=True, index=True)
password_hash = Column(String, nullable=False)
role = Column(String, nullable=False, default='user') # 角色: superadmin, admin, user
role = Column(String, nullable=False, default="user") # 角色: superadmin, admin, user
created_at = Column(DateTime, default=func.now())
last_login = Column(DateTime, nullable=True)
@ -24,18 +26,20 @@ class User(Base):
"username": self.username,
"role": self.role,
"created_at": self.created_at.isoformat() if self.created_at else None,
"last_login": self.last_login.isoformat() if self.last_login else None
"last_login": self.last_login.isoformat() if self.last_login else None,
}
if include_password:
result["password_hash"] = self.password_hash
return result
class OperationLog(Base):
"""操作日志模型"""
__tablename__ = 'operation_logs'
__tablename__ = "operation_logs"
id = Column(Integer, primary_key=True, autoincrement=True)
user_id = Column(Integer, ForeignKey('users.id'), nullable=False)
user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
operation = Column(String, nullable=False)
details = Column(Text, nullable=True)
ip_address = Column(String, nullable=True)
@ -51,5 +55,5 @@ class OperationLog(Base):
"operation": self.operation,
"details": self.details,
"ip_address": self.ip_address,
"timestamp": self.timestamp.isoformat() if self.timestamp else None
"timestamp": self.timestamp.isoformat() if self.timestamp else None,
}

15
server/routers/__init__.py
View File

@ -1,15 +1,16 @@
from fastapi import APIRouter
from server.routers.system_router import system
from server.routers.auth_router import auth
from server.routers.chat_router import chat
from server.routers.knowledge_router import knowledge
from server.routers.graph_router import graph
from server.routers.knowledge_router import knowledge
from server.routers.system_router import system
router = APIRouter()
# 注册路由结构
router.include_router(system) # /api/system/*
router.include_router(auth) # /api/auth/*
router.include_router(chat) # /api/chat/*
router.include_router(knowledge) # /api/knowledge/*
router.include_router(graph) # /api/graph/*
router.include_router(system) # /api/system/*
router.include_router(auth) # /api/auth/*
router.include_router(chat) # /api/chat/*
router.include_router(knowledge) # /api/knowledge/*
router.include_router(graph) # /api/graph/*

94
server/routers/auth_router.py
View File

@ -1,18 +1,20 @@
from datetime import datetime
from fastapi import APIRouter, Depends, HTTPException, Request, status
from fastapi.security import OAuth2PasswordRequestForm
from pydantic import BaseModel
from sqlalchemy.orm import Session
from datetime import datetime
from server.db_manager import db_manager
from server.models.user_model import User, OperationLog
from server.models.user_model import OperationLog, User
from server.utils.auth_middleware import get_admin_user, get_current_user, get_db, get_superadmin_user, oauth2_scheme
from server.utils.auth_utils import AuthUtils
from server.utils.auth_middleware import get_db, get_current_user, get_admin_user, get_superadmin_user, oauth2_scheme
from server.utils.common_utils import log_operation
# 创建路由器
auth = APIRouter(prefix="/auth", tags=["authentication"])
# 请求和响应模型
class Token(BaseModel):
access_token: str
@ -21,16 +23,19 @@ class Token(BaseModel):
username: str
role: str
class UserCreate(BaseModel):
username: str
password: str
role: str = "user"
class UserUpdate(BaseModel):
username: str | None = None
password: str | None = None
role: str | None = None
class UserResponse(BaseModel):
id: int
username: str
@ -38,10 +43,12 @@ class UserResponse(BaseModel):
created_at: str
last_login: str | None = None
class InitializeAdmin(BaseModel):
username: str
password: str
# =============================================================================
# === 工具函数 ===
# =============================================================================
@ -52,11 +59,9 @@ class InitializeAdmin(BaseModel):
# === 认证分组 ===
# =============================================================================
@auth.post("/token", response_model=Token)
async def login_for_access_token(
form_data: OAuth2PasswordRequestForm = Depends(),
db: Session = Depends(get_db)
):
async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends(), db: Session = Depends(get_db)):
# 查找用户
user = db.query(User).filter(User.username == form_data.username).first()
@ -84,21 +89,20 @@ async def login_for_access_token(
"token_type": "bearer",
"user_id": user.id,
"username": user.username,
"role": user.role
"role": user.role,
}
# 路由：校验是否需要初始化管理员
@auth.get("/check-first-run")
async def check_first_run():
is_first_run = db_manager.check_first_run()
return {"first_run": is_first_run}
# 路由：初始化管理员账户
@auth.post("/initialize", response_model=Token)
async def initialize_admin(
admin_data: InitializeAdmin,
db: Session = Depends(get_db)
):
async def initialize_admin(admin_data: InitializeAdmin, db: Session = Depends(get_db)):
# 检查是否是首次运行
if not db_manager.check_first_run():
raise HTTPException(
@ -110,10 +114,7 @@ async def initialize_admin(
hashed_password = AuthUtils.hash_password(admin_data.password)
new_admin = User(
username=admin_data.username,
password_hash=hashed_password,
role="superadmin",
last_login=datetime.now()
username=admin_data.username, password_hash=hashed_password, role="superadmin", last_login=datetime.now()
)
db.add(new_admin)
@ -132,29 +133,30 @@ async def initialize_admin(
"token_type": "bearer",
"user_id": new_admin.id,
"username": new_admin.username,
"role": new_admin.role
"role": new_admin.role,
}
# 路由：获取当前用户信息
# =============================================================================
# === 用户信息分组 ===
# =============================================================================
@auth.get("/me", response_model=UserResponse)
async def read_users_me(current_user: User = Depends(get_current_user)):
return current_user.to_dict()
# 路由：创建新用户（管理员权限）
# =============================================================================
# === 用户管理分组 ===
# =============================================================================
@auth.post("/users", response_model=UserResponse)
async def create_user(
user_data: UserCreate,
request: Request,
current_user: User = Depends(get_admin_user),
db: Session = Depends(get_db)
user_data: UserCreate, request: Request, current_user: User = Depends(get_admin_user), db: Session = Depends(get_db)
):
# 检查用户名是否已存在
existing_user = db.query(User).filter(User.username == user_data.username).first()
@ -182,45 +184,30 @@ async def create_user(
detail="管理员只能创建普通用户账户",
)
new_user = User(
username=user_data.username,
password_hash=hashed_password,
role=user_data.role
)
new_user = User(username=user_data.username, password_hash=hashed_password, role=user_data.role)
db.add(new_user)
db.commit()
db.refresh(new_user)
# 记录操作
log_operation(
db,
current_user.id,
"创建用户",
f"创建用户: {user_data.username}, 角色: {user_data.role}",
request
)
log_operation(db, current_user.id, "创建用户", f"创建用户: {user_data.username}, 角色: {user_data.role}", request)
return new_user.to_dict()
# 路由：获取所有用户（管理员权限）
@auth.get("/users", response_model=list[UserResponse])
async def read_users(
skip: int = 0,
limit: int = 100,
current_user: User = Depends(get_admin_user),
db: Session = Depends(get_db)
skip: int = 0, limit: int = 100, current_user: User = Depends(get_admin_user), db: Session = Depends(get_db)
):
users = db.query(User).offset(skip).limit(limit).all()
return [user.to_dict() for user in users]
# 路由：获取特定用户信息（管理员权限）
@auth.get("/users/{user_id}", response_model=UserResponse)
async def read_user(
user_id: int,
current_user: User = Depends(get_admin_user),
db: Session = Depends(get_db)
):
async def read_user(user_id: int, current_user: User = Depends(get_admin_user), db: Session = Depends(get_db)):
user = db.query(User).filter(User.id == user_id).first()
if user is None:
raise HTTPException(
@ -229,6 +216,7 @@ async def read_user(
)
return user.to_dict()
# 路由：更新用户信息（管理员权限）
@auth.put("/users/{user_id}", response_model=UserResponse)
async def update_user(
@ -236,7 +224,7 @@ async def update_user(
user_data: UserUpdate,
request: Request,
current_user: User = Depends(get_admin_user),
db: Session = Depends(get_db)
db: Session = Depends(get_db),
):
user = db.query(User).filter(User.id == user_id).first()
if user is None:
@ -284,23 +272,15 @@ async def update_user(
db.commit()
# 记录操作
log_operation(
db,
current_user.id,
"更新用户",
f"更新用户ID {user_id}: {', '.join(update_details)}",
request
)
log_operation(db, current_user.id, "更新用户", f"更新用户ID {user_id}: {', '.join(update_details)}", request)
return user.to_dict()
# 路由：删除用户（管理员权限）
@auth.delete("/users/{user_id}", response_model=dict)
async def delete_user(
user_id: int,
request: Request,
current_user: User = Depends(get_admin_user),
db: Session = Depends(get_db)
user_id: int, request: Request, current_user: User = Depends(get_admin_user), db: Session = Depends(get_db)
):
user = db.query(User).filter(User.id == user_id).first()
if user is None:
@ -335,11 +315,7 @@ async def delete_user(
# 记录操作
log_operation(
db,
current_user.id,
"删除用户",
f"删除用户: {user.username}, ID: {user.id}, 角色: {user.role}",
request
db, current_user.id, "删除用户", f"删除用户: {user.username}, ID: {user.id}, 角色: {user.role}", request
)
# 删除用户

144
server/routers/chat_router.py
View File

@ -1,24 +1,25 @@
import os
import json
import asyncio
import json
import os
import time
import traceback
import uuid
import time
from fastapi import APIRouter, Body, Depends, HTTPException, Query
from fastapi.responses import StreamingResponse
from langchain_core.messages import AIMessageChunk, HumanMessage
from sqlalchemy.orm import Session
from pydantic import BaseModel
from sqlalchemy.orm import Session
from src import executor, config
from server.models.thread_model import Thread
from server.models.user_model import User
from server.routers.auth_router import get_admin_user
from server.utils.auth_middleware import get_db, get_required_user
from src import config, executor
from src.agents import agent_manager
from src.agents.common.tools import gen_tool_info, get_buildin_tools
from src.models import select_model
from src.utils.logging_config import logger
from src.agents.common.tools import get_buildin_tools, gen_tool_info
from server.routers.auth_router import get_admin_user
from server.utils.auth_middleware import get_required_user, get_db
from server.models.user_model import User
from server.models.thread_model import Thread
chat = APIRouter(prefix="/chat", tags=["chat"])
@ -26,6 +27,7 @@ chat = APIRouter(prefix="/chat", tags=["chat"])
# > === 智能体管理分组 ===
# =============================================================================
@chat.get("/default_agent")
async def get_default_agent(current_user: User = Depends(get_required_user)):
"""获取默认智能体ID（需要登录）"""
@ -42,8 +44,9 @@ async def get_default_agent(current_user: User = Depends(get_required_user)):
logger.error(f"获取默认智能体出错: {e}")
raise HTTPException(status_code=500, detail=f"获取默认智能体出错: {str(e)}")
@chat.post("/set_default_agent")
async def set_default_agent(request_data: dict = Body(...), current_user = Depends(get_admin_user)):
async def set_default_agent(request_data: dict = Body(...), current_user=Depends(get_admin_user)):
"""设置默认智能体ID (仅管理员)"""
try:
agent_id = request_data.get("agent_id")
@ -69,15 +72,18 @@ async def set_default_agent(request_data: dict = Body(...), current_user = Depen
logger.error(f"设置默认智能体出错: {e}")
raise HTTPException(status_code=500, detail=f"设置默认智能体出错: {str(e)}")
# =============================================================================
# > === 对话分组 ===
# =============================================================================
@chat.post("/call")
async def call(query: str = Body(...), meta: dict = Body(None), current_user: User = Depends(get_required_user)):
"""调用模型进行简单问答（需要登录）"""
meta = meta or {}
model = select_model(model_provider=meta.get("model_provider"), model_name=meta.get("model_name"))
async def predict_async(query):
loop = asyncio.get_event_loop()
return await loop.run_in_executor(executor, model.predict, query)
@ -87,6 +93,7 @@ async def call(query: str = Body(...), meta: dict = Body(None), current_user: Us
return {"response": response.content}
@chat.get("/agent")
async def get_agent(current_user: User = Depends(get_required_user)):
"""获取所有可用智能体（需要登录）"""
@ -94,35 +101,39 @@ async def get_agent(current_user: User = Depends(get_required_user)):
# logger.debug(f"agents: {agents}")
return {"agents": agents}
@chat.post("/agent/{agent_id}")
async def chat_agent(agent_id: str,
query: str = Body(...),
config: dict = Body({}),
meta: dict = Body({}),
current_user: User = Depends(get_required_user)):
async def chat_agent(
agent_id: str,
query: str = Body(...),
config: dict = Body({}),
meta: dict = Body({}),
current_user: User = Depends(get_required_user),
):
"""使用特定智能体进行对话（需要登录）"""
logger.info(f"agent_id: {agent_id}, query: {query}, config: {config}, meta: {meta}")
meta.update({
"query": query,
"agent_id": agent_id,
"server_model_name": config.get("model", agent_id),
"thread_id": config.get("thread_id"),
"user_id": current_user.id
})
meta.update(
{
"query": query,
"agent_id": agent_id,
"server_model_name": config.get("model", agent_id),
"thread_id": config.get("thread_id"),
"user_id": current_user.id,
}
)
# 将meta和thread_id整合到config中
def make_chunk(content=None, **kwargs):
return json.dumps({
"request_id": meta.get("request_id"),
"response": content,
**kwargs
}, ensure_ascii=False).encode('utf-8') + b"\n"
return (
json.dumps(
{"request_id": meta.get("request_id"), "response": content, **kwargs}, ensure_ascii=False
).encode("utf-8")
+ b"\n"
)
async def stream_messages():
# 代表服务端已经收到了请求
yield make_chunk(status="init", meta=meta, msg=HumanMessage(content=query).model_dump())
@ -145,39 +156,38 @@ async def chat_agent(agent_id: str,
async for msg, metadata in agent.stream_messages(messages, input_context=input_context):
# logger.debug(f"msg: {msg.model_dump()}, metadata: {metadata}")
if isinstance(msg, AIMessageChunk):
yield make_chunk(content=msg.content,
msg=msg.model_dump(),
metadata=metadata,
status="loading")
yield make_chunk(content=msg.content, msg=msg.model_dump(), metadata=metadata, status="loading")
else:
yield make_chunk(msg=msg.model_dump(),
metadata=metadata,
status="loading")
yield make_chunk(msg=msg.model_dump(), metadata=metadata, status="loading")
yield make_chunk(status="finished", meta=meta)
except Exception as e:
logger.error(f"Error streaming messages: {e}, {traceback.format_exc()}")
yield make_chunk(message=f"Error streaming messages: {e}", status="error")
return StreamingResponse(stream_messages(), media_type='application/json')
return StreamingResponse(stream_messages(), media_type="application/json")
# =============================================================================
# > === 模型管理分组 ===
# =============================================================================
@chat.get("/models")
async def get_chat_models(model_provider: str, current_user: User = Depends(get_admin_user)):
"""获取指定模型提供商的模型列表（需要登录）"""
model = select_model(model_provider=model_provider)
return {"models": model.get_models()}
@chat.post("/models/update")
async def update_chat_models(model_provider: str, model_names: list[str], current_user = Depends(get_admin_user)):
async def update_chat_models(model_provider: str, model_names: list[str], current_user=Depends(get_admin_user)):
"""更新指定模型提供商的模型列表 (仅管理员)"""
config.model_names[model_provider]["models"] = model_names
config._save_models_to_file()
return {"models": config.model_names[model_provider]["models"]}
@chat.get("/tools")
async def get_tools(agent_id: str, current_user: User = Depends(get_required_user)):
"""获取所有可用工具（需要登录）"""
@ -193,12 +203,9 @@ async def get_tools(agent_id: str, current_user: User = Depends(get_required_use
tools_info = gen_tool_info(tools)
return {"tools": {tool["id"]: tool for tool in tools_info}}
@chat.post("/agent/{agent_id}/config")
async def save_agent_config(
agent_id: str,
config: dict = Body(...),
current_user: User = Depends(get_required_user)
):
async def save_agent_config(agent_id: str, config: dict = Body(...), current_user: User = Depends(get_required_user)):
"""保存智能体配置到YAML文件（需要登录）"""
try:
# 获取Agent实例和配置类
@ -217,12 +224,9 @@ async def save_agent_config(
logger.error(f"保存智能体配置出错: {e}, {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=f"保存智能体配置出错: {str(e)}")
@chat.get("/agent/{agent_id}/history")
async def get_agent_history(
agent_id: str,
thread_id: str,
current_user: User = Depends(get_required_user)
):
async def get_agent_history(agent_id: str, thread_id: str, current_user: User = Depends(get_required_user)):
"""获取智能体历史消息（需要登录）"""
try:
# 获取Agent实例和配置类
@ -237,11 +241,9 @@ async def get_agent_history(
logger.error(f"获取智能体历史消息出错: {e}, {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=f"获取智能体历史消息出错: {str(e)}")
@chat.get("/agent/{agent_id}/config")
async def get_agent_config(
agent_id: str,
current_user: User = Depends(get_required_user)
):
async def get_agent_config(agent_id: str, current_user: User = Depends(get_required_user)):
"""从YAML文件加载智能体配置（需要登录）"""
try:
# 检查智能体是否存在
@ -256,8 +258,10 @@ async def get_agent_config(
logger.error(f"加载智能体配置出错: {e}, {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=f"加载智能体配置出错: {str(e)}")
# ==================== 线程管理 API ====================
class ThreadCreate(BaseModel):
title: str | None = None
agent_id: str
@ -279,11 +283,10 @@ class ThreadResponse(BaseModel):
# > === 会话管理分组 ===
# =============================================================================
@chat.post("/thread", response_model=ThreadResponse)
async def create_thread(
thread: ThreadCreate,
db: Session = Depends(get_db),
current_user: User = Depends(get_required_user)
thread: ThreadCreate, db: Session = Depends(get_db), current_user: User = Depends(get_required_user)
):
"""创建新对话线程"""
thread_id = str(uuid.uuid4())
@ -313,11 +316,7 @@ async def create_thread(
@chat.get("/threads", response_model=list[ThreadResponse])
async def list_threads(
agent_id: str,
db: Session = Depends(get_db),
current_user: User = Depends(get_required_user)
):
async def list_threads(agent_id: str, db: Session = Depends(get_db), current_user: User = Depends(get_required_user)):
"""获取用户的所有对话线程"""
assert agent_id, "agent_id 不能为空"
query = db.query(Thread).filter(
@ -345,16 +344,9 @@ async def list_threads(
@chat.delete("/thread/{thread_id}")
async def delete_thread(
thread_id: str,
db: Session = Depends(get_db),
current_user: User = Depends(get_required_user)
):
async def delete_thread(thread_id: str, db: Session = Depends(get_db), current_user: User = Depends(get_required_user)):
"""删除对话线程"""
thread = db.query(Thread).filter(
Thread.id == thread_id,
Thread.user_id == str(current_user.id)
).first()
thread = db.query(Thread).filter(Thread.id == thread_id, Thread.user_id == str(current_user.id)).first()
if not thread:
raise HTTPException(status_code=404, detail="对话线程不存在")
@ -376,14 +368,14 @@ async def update_thread(
thread_id: str,
thread_update: ThreadUpdate,
db: Session = Depends(get_db),
current_user: User = Depends(get_required_user)
current_user: User = Depends(get_required_user),
):
"""更新对话线程信息"""
thread = db.query(Thread).filter(
Thread.id == thread_id,
Thread.user_id == str(current_user.id),
Thread.status == 1
).first()
thread = (
db.query(Thread)
.filter(Thread.id == thread_id, Thread.user_id == str(current_user.id), Thread.status == 1)
.first()
)
if not thread:
raise HTTPException(status_code=404, detail="对话线程不存在")

159
server/routers/graph_router.py
View File

@ -1,9 +1,10 @@
import traceback
from fastapi import APIRouter, Query, HTTPException, Depends, Body
from server.utils.auth_middleware import get_admin_user
from server.models.user_model import User
from src import knowledge_base, graph_base
from fastapi import APIRouter, Body, Depends, HTTPException, Query
from server.models.user_model import User
from server.utils.auth_middleware import get_admin_user
from src import graph_base, knowledge_base
from src.utils.logging_config import logger
graph = APIRouter(prefix="/graph", tags=["graph"])
@ -13,13 +14,14 @@ graph = APIRouter(prefix="/graph", tags=["graph"])
# === 子图查询分组 ===
# =============================================================================
@graph.get("/lightrag/subgraph")
async def get_lightrag_subgraph(
db_id: str = Query(..., description="数据库ID"),
node_label: str = Query(..., description="节点标签或实体名称"),
max_depth: int = Query(2, description="最大深度", ge=1, le=5),
max_nodes: int = Query(100, description="最大节点数", ge=1, le=1000),
current_user: User = Depends(get_admin_user)
current_user: User = Depends(get_admin_user),
):
"""
使用 LightRAG 原生方法获取知识图谱子图
@ -34,13 +36,14 @@ async def get_lightrag_subgraph(
包含节点和边的知识图谱数据
"""
try:
logger.info(f"获取子图数据 - db_id: {db_id}, node_label: {node_label}, max_depth: {max_depth}, max_nodes: {max_nodes}")
logger.info(
f"获取子图数据 - db_id: {db_id}, node_label: {node_label}, max_depth: {max_depth}, max_nodes: {max_nodes}"
)
# 检查是否是 LightRAG 数据库
if not knowledge_base.is_lightrag_database(db_id):
raise HTTPException(
status_code=400,
detail=f"数据库 {db_id} 不是 LightRAG 类型，图谱功能仅支持 LightRAG 知识库"
status_code=400, detail=f"数据库 {db_id} 不是 LightRAG 类型，图谱功能仅支持 LightRAG 知识库"
)
# 获取 LightRAG 实例
@ -50,30 +53,32 @@ async def get_lightrag_subgraph(
# 使用 LightRAG 的原生 get_knowledge_graph 方法
knowledge_graph = await rag_instance.get_knowledge_graph(
node_label=node_label,
max_depth=max_depth,
max_nodes=max_nodes
node_label=node_label, max_depth=max_depth, max_nodes=max_nodes
)
# 将 LightRAG 的 KnowledgeGraph 格式转换为前端需要的格式
nodes = []
for node in knowledge_graph.nodes:
nodes.append({
"id": node.id,
"labels": node.labels,
"entity_type": node.properties.get("entity_type", "unknown"),
"properties": node.properties
})
nodes.append(
{
"id": node.id,
"labels": node.labels,
"entity_type": node.properties.get("entity_type", "unknown"),
"properties": node.properties,
}
)
edges = []
for edge in knowledge_graph.edges:
edges.append({
"id": edge.id,
"source": edge.source,
"target": edge.target,
"type": edge.type,
"properties": edge.properties
})
edges.append(
{
"id": edge.id,
"source": edge.source,
"target": edge.target,
"type": edge.type,
"properties": edge.properties,
}
)
result = {
"success": True,
@ -82,8 +87,8 @@ async def get_lightrag_subgraph(
"edges": edges,
"is_truncated": knowledge_graph.is_truncated,
"total_nodes": len(nodes),
"total_edges": len(edges)
}
"total_edges": len(edges),
},
}
logger.info(f"成功获取子图 - 节点数: {len(nodes)}, 边数: {len(edges)}")
@ -99,9 +104,7 @@ async def get_lightrag_subgraph(
@graph.get("/lightrag/databases")
async def get_lightrag_databases(
current_user: User = Depends(get_admin_user)
):
async def get_lightrag_databases(current_user: User = Depends(get_admin_user)):
"""
获取所有可用的 LightRAG 数据库
@ -110,25 +113,21 @@ async def get_lightrag_databases(
"""
try:
lightrag_databases = knowledge_base.get_lightrag_databases()
return {
"success": True,
"data": {
"databases": lightrag_databases
}
}
return {"success": True, "data": {"databases": lightrag_databases}}
except Exception as e:
logger.error(f"获取 LightRAG 数据库列表失败: {e}")
raise HTTPException(status_code=500, detail=f"获取 LightRAG 数据库列表失败: {str(e)}")
# =============================================================================
# === 节点管理分组 ===
# =============================================================================
@graph.get("/lightrag/labels")
async def get_lightrag_labels(
db_id: str = Query(..., description="数据库ID"),
current_user: User = Depends(get_admin_user)
db_id: str = Query(..., description="数据库ID"), current_user: User = Depends(get_admin_user)
):
"""
获取知识图谱中的所有标签
@ -145,8 +144,7 @@ async def get_lightrag_labels(
# 检查是否是 LightRAG 数据库
if not knowledge_base.is_lightrag_database(db_id):
raise HTTPException(
status_code=400,
detail=f"数据库 {db_id} 不是 LightRAG 类型，图谱功能仅支持 LightRAG 知识库"
status_code=400, detail=f"数据库 {db_id} 不是 LightRAG 类型，图谱功能仅支持 LightRAG 知识库"
)
# 获取 LightRAG 实例
@ -157,12 +155,7 @@ async def get_lightrag_labels(
# 使用 LightRAG 的原生方法获取所有标签
labels = await rag_instance.get_graph_labels()
return {
"success": True,
"data": {
"labels": labels
}
}
return {"success": True, "data": {"labels": labels}}
except HTTPException:
# 重新抛出 HTTP 异常
@ -177,7 +170,7 @@ async def get_lightrag_labels(
async def get_neo4j_nodes(
kgdb_name: str = Query(..., description="知识图谱数据库名称"),
num: int = Query(100, description="节点数量", ge=1, le=1000),
current_user: User = Depends(get_admin_user)
current_user: User = Depends(get_admin_user),
):
"""
获取图谱节点样本数据
@ -190,20 +183,16 @@ async def get_neo4j_nodes(
result = graph_base.get_sample_nodes(kgdb_name, num)
return {
"success": True,
"result": result,
"message": "success"
}
return {"success": True, "result": result, "message": "success"}
except Exception as e:
logger.error(f"获取图节点数据失败: {e}\n{traceback.format_exc()}")
raise HTTPException(status_code=500, detail=f"获取图节点数据失败: {str(e)}")
@graph.get("/neo4j/node")
async def get_neo4j_node(
entity_name: str = Query(..., description="实体名称"),
current_user: User = Depends(get_admin_user)
entity_name: str = Query(..., description="实体名称"), current_user: User = Depends(get_admin_user)
):
"""
根据实体名称查询图节点
@ -214,16 +203,13 @@ async def get_neo4j_node(
result = graph_base.query_node(entity_name=entity_name)
return {
"success": True,
"result": result,
"message": "success"
}
return {"success": True, "result": result, "message": "success"}
except Exception as e:
logger.error(f"查询图节点失败: {e}\n{traceback.format_exc()}")
raise HTTPException(status_code=500, detail=f"查询图节点失败: {str(e)}")
# =============================================================================
# === 边管理分组 ===
# =============================================================================
@ -234,10 +220,10 @@ async def get_neo4j_node(
# === 图谱分析分组 ===
# =============================================================================
@graph.get("/lightrag/stats")
async def get_lightrag_stats(
db_id: str = Query(..., description="数据库ID"),
current_user: User = Depends(get_admin_user)
db_id: str = Query(..., description="数据库ID"), current_user: User = Depends(get_admin_user)
):
"""
获取知识图谱统计信息
@ -248,8 +234,7 @@ async def get_lightrag_stats(
# 检查是否是 LightRAG 数据库
if not knowledge_base.is_lightrag_database(db_id):
raise HTTPException(
status_code=400,
detail=f"数据库 {db_id} 不是 LightRAG 类型，图谱功能仅支持 LightRAG 知识库"
status_code=400, detail=f"数据库 {db_id} 不是 LightRAG 类型，图谱功能仅支持 LightRAG 知识库"
)
# 获取 LightRAG 实例
@ -261,7 +246,7 @@ async def get_lightrag_stats(
knowledge_graph = await rag_instance.get_knowledge_graph(
node_label="*",
max_depth=1,
max_nodes=10000 # 设置较大值以获取完整统计
max_nodes=10000, # 设置较大值以获取完整统计
)
# 统计实体类型分布
@ -271,8 +256,7 @@ async def get_lightrag_stats(
entity_types[entity_type] = entity_types.get(entity_type, 0) + 1
entity_types_list = [
{"type": k, "count": v}
for k, v in sorted(entity_types.items(), key=lambda x: x[1], reverse=True)
{"type": k, "count": v} for k, v in sorted(entity_types.items(), key=lambda x: x[1], reverse=True)
]
return {
@ -281,8 +265,8 @@ async def get_lightrag_stats(
"total_nodes": len(knowledge_graph.nodes),
"total_edges": len(knowledge_graph.edges),
"entity_types": entity_types_list,
"is_truncated": knowledge_graph.is_truncated
}
"is_truncated": knowledge_graph.is_truncated,
},
}
except HTTPException:
@ -293,6 +277,7 @@ async def get_lightrag_stats(
logger.error(f"Traceback: {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=f"获取图谱统计信息失败: {str(e)}")
@graph.get("/neo4j/info")
async def get_neo4j_info(current_user: User = Depends(get_admin_user)):
"""获取Neo4j图数据库信息"""
@ -300,26 +285,21 @@ async def get_neo4j_info(current_user: User = Depends(get_admin_user)):
graph_info = graph_base.get_graph_info()
if graph_info is None:
raise HTTPException(status_code=400, detail="图数据库获取出错")
return {
"success": True,
"data": graph_info
}
return {"success": True, "data": graph_info}
except Exception as e:
logger.error(f"获取图数据库信息失败: {e}")
raise HTTPException(status_code=500, detail=f"获取图数据库信息失败: {str(e)}")
@graph.post("/neo4j/index-entities")
async def index_neo4j_entities(
data: dict = Body(default={}),
current_user: User = Depends(get_admin_user)
):
async def index_neo4j_entities(data: dict = Body(default={}), current_user: User = Depends(get_admin_user)):
"""为Neo4j图谱节点添加嵌入向量索引"""
try:
if not graph_base.is_running():
raise HTTPException(status_code=400, detail="图数据库未启动")
# 获取参数或使用默认值
kgdb_name = data.get('kgdb_name', 'neo4j')
kgdb_name = data.get("kgdb_name", "neo4j")
# 调用GraphDatabase的add_embedding_to_nodes方法
count = graph_base.add_embedding_to_nodes(kgdb_name=kgdb_name)
@ -328,37 +308,24 @@ async def index_neo4j_entities(
"success": True,
"status": "success",
"message": f"已成功为{count}个节点添加嵌入向量",
"indexed_count": count
"indexed_count": count,
}
except Exception as e:
logger.error(f"索引节点失败: {e}")
raise HTTPException(status_code=500, detail=f"索引节点失败: {str(e)}")
@graph.post("/neo4j/add-entities")
async def add_neo4j_entities(
file_path: str = Body(...),
kgdb_name: str | None = Body(None),
current_user: User = Depends(get_admin_user)
file_path: str = Body(...), kgdb_name: str | None = Body(None), current_user: User = Depends(get_admin_user)
):
"""通过JSONL文件添加图谱实体到Neo4j"""
try:
if not file_path.endswith('.jsonl'):
return {
"success": False,
"message": "文件格式错误，请上传jsonl文件",
"status": "failed"
}
if not file_path.endswith(".jsonl"):
return {"success": False, "message": "文件格式错误，请上传jsonl文件", "status": "failed"}
await graph_base.jsonl_file_add_entity(file_path, kgdb_name)
return {
"success": True,
"message": "实体添加成功",
"status": "success"
}
return {"success": True, "message": "实体添加成功", "status": "success"}
except Exception as e:
logger.error(f"添加实体失败: {e}, {traceback.format_exc()}")
return {
"success": False,
"message": f"添加实体失败: {e}",
"status": "failed"
}
return {"success": False, "message": f"添加实体失败: {e}", "status": "failed"}

154
server/routers/knowledge_router.py
View File

@ -1,14 +1,15 @@
import os
import asyncio
import os
import traceback
from fastapi import APIRouter, File, UploadFile, HTTPException, Depends, Body, Form, Query
from fastapi import APIRouter, Body, Depends, File, Form, HTTPException, Query, UploadFile
from fastapi.responses import FileResponse
from src import executor, config, knowledge_base
from src.utils import logger, hashstr
from src.knowledge.indexing import process_file_to_markdown
from server.utils.auth_middleware import get_admin_user
from server.models.user_model import User
from server.utils.auth_middleware import get_admin_user
from src import config, executor, knowledge_base
from src.knowledge.indexing import process_file_to_markdown
from src.utils import hashstr, logger
knowledge = APIRouter(prefix="/knowledge", tags=["knowledge"])
@ -16,6 +17,7 @@ knowledge = APIRouter(prefix="/knowledge", tags=["knowledge"])
# === 数据库管理分组 ===
# =============================================================================
@knowledge.get("/databases")
async def get_databases(current_user: User = Depends(get_admin_user)):
"""获取所有知识库"""
@ -26,6 +28,7 @@ async def get_databases(current_user: User = Depends(get_admin_user)):
logger.error(f"获取数据库列表失败 {e}, {traceback.format_exc()}")
return {"message": f"获取数据库列表失败 {e}", "databases": []}
@knowledge.post("/databases")
async def create_database(
database_name: str = Body(...),
@ -34,23 +37,21 @@ async def create_database(
kb_type: str = Body("lightrag"),
additional_params: dict = Body({}),
llm_info: dict = Body(None),
current_user: User = Depends(get_admin_user)
current_user: User = Depends(get_admin_user),
):
"""创建知识库"""
logger.debug(f"Create database {database_name} with kb_type {kb_type}, additional_params {additional_params}, llm_info {llm_info}")
logger.debug(
f"Create database {database_name} with kb_type {kb_type}, additional_params {additional_params}, llm_info {llm_info}"
)
try:
embed_info = config.embed_model_names[embed_model_name]
database_info = await knowledge_base.create_database(
database_name,
description,
kb_type=kb_type,
embed_info=embed_info,
llm_info=llm_info,
**additional_params
database_name, description, kb_type=kb_type, embed_info=embed_info, llm_info=llm_info, **additional_params
)
# 需要重新加载所有智能体，因为工具刷新了
from src.agents import agent_manager
await agent_manager.reload_all()
return database_info
@ -58,6 +59,7 @@ async def create_database(
logger.error(f"创建数据库失败 {e}, {traceback.format_exc()}")
return {"message": f"创建数据库失败 {e}", "status": "failed"}
@knowledge.get("/databases/{db_id}")
async def get_database_info(db_id: str, current_user: User = Depends(get_admin_user)):
"""获取知识库详细信息"""
@ -66,12 +68,10 @@ async def get_database_info(db_id: str, current_user: User = Depends(get_admin_u
raise HTTPException(status_code=404, detail="Database not found")
return database
@knowledge.put("/databases/{db_id}")
async def update_database_info(
db_id: str,
name: str = Body(...),
description: str = Body(...),
current_user: User = Depends(get_admin_user)
db_id: str, name: str = Body(...), description: str = Body(...), current_user: User = Depends(get_admin_user)
):
"""更新知识库信息"""
logger.debug(f"Update database {db_id} info: {name}, {description}")
@ -82,6 +82,7 @@ async def update_database_info(
logger.error(f"更新数据库失败 {e}, {traceback.format_exc()}")
raise HTTPException(status_code=400, detail=f"更新数据库失败: {e}")
@knowledge.delete("/databases/{db_id}")
async def delete_database(db_id: str, current_user: User = Depends(get_admin_user)):
"""删除知识库"""
@ -91,6 +92,7 @@ async def delete_database(db_id: str, current_user: User = Depends(get_admin_use
# 需要重新加载所有智能体，因为工具刷新了
from src.agents import agent_manager
await agent_manager.reload_all()
return {"message": "删除成功"}
@ -98,19 +100,18 @@ async def delete_database(db_id: str, current_user: User = Depends(get_admin_use
logger.error(f"删除数据库失败 {e}, {traceback.format_exc()}")
raise HTTPException(status_code=400, detail=f"删除数据库失败: {e}")
@knowledge.get("/databases/{db_id}/export")
async def export_database(
db_id: str,
format: str = Query("csv", enum=["csv", "xlsx", "md", "txt"]),
include_vectors: bool = Query(False, description="是否在导出中包含向量数据"),
current_user: User = Depends(get_admin_user)
current_user: User = Depends(get_admin_user),
):
"""导出知识库数据"""
logger.debug(f"Exporting database {db_id} with format {format}")
try:
file_path = await knowledge_base.export_data(
db_id, format=format, include_vectors=include_vectors
)
file_path = await knowledge_base.export_data(db_id, format=format, include_vectors=include_vectors)
if not os.path.exists(file_path):
raise HTTPException(status_code=404, detail="Exported file not found.")
@ -119,15 +120,11 @@ async def export_database(
"csv": "text/csv",
"xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
"md": "text/markdown",
"txt": "text/plain"
"txt": "text/plain",
}
media_type = media_types.get(format, "application/octet-stream")
return FileResponse(
path=file_path,
filename=os.path.basename(file_path),
media_type=media_type
)
return FileResponse(path=file_path, filename=os.path.basename(file_path), media_type=media_type)
except NotImplementedError as e:
logger.warning(f"A disabled feature was accessed: {e}")
raise HTTPException(status_code=501, detail=str(e))
@ -135,38 +132,34 @@ async def export_database(
logger.error(f"导出数据库失败 {e}, {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=f"导出数据库失败: {e}")
# =============================================================================
# === 文档管理分组 ===
# =============================================================================
@knowledge.post("/databases/{db_id}/documents")
async def add_documents(
db_id: str,
items: list[str] = Body(...),
params: dict = Body(...),
current_user: User = Depends(get_admin_user)
db_id: str, items: list[str] = Body(...), params: dict = Body(...), current_user: User = Depends(get_admin_user)
):
"""添加文档到知识库"""
logger.debug(f"Add documents for db_id {db_id}: {items} {params=}")
content_type = params.get('content_type', 'file')
content_type = params.get("content_type", "file")
try:
processed_items = await knowledge_base.add_content(db_id, items, params=params)
item_type = "URLs" if content_type == 'url' else "files"
processed_failed_count = len([_p for _p in processed_items if _p['status'] == 'failed'])
item_type = "URLs" if content_type == "url" else "files"
processed_failed_count = len([_p for _p in processed_items if _p["status"] == "failed"])
processed_info = f"Processed {len(processed_items)} {item_type}, {processed_failed_count} {item_type} failed"
return {"message": processed_info, "items": processed_items, "status": "success"}
except Exception as e:
logger.error(f"Failed to process {content_type}s: {e}, {traceback.format_exc()}")
return {"message": f"Failed to process {content_type}s: {e}", "status": "failed"}
@knowledge.get("/databases/{db_id}/documents/{doc_id}")
async def get_document_info(
db_id: str,
doc_id: str,
current_user: User = Depends(get_admin_user)
):
async def get_document_info(db_id: str, doc_id: str, current_user: User = Depends(get_admin_user)):
"""获取文档详细信息"""
logger.debug(f"GET document {doc_id} info in {db_id}")
@ -177,12 +170,9 @@ async def get_document_info(
logger.error(f"Failed to get file info, {e}, {db_id=}, {doc_id=}, {traceback.format_exc()}")
return {"message": "Failed to get file info", "status": "failed"}
@knowledge.delete("/databases/{db_id}/documents/{doc_id}")
async def delete_document(
db_id: str,
doc_id: str,
current_user: User = Depends(get_admin_user)
):
async def delete_document(db_id: str, doc_id: str, current_user: User = Depends(get_admin_user)):
"""删除文档"""
logger.debug(f"DELETE document {doc_id} info in {db_id}")
try:
@ -192,16 +182,15 @@ async def delete_document(
logger.error(f"删除文档失败 {e}, {traceback.format_exc()}")
raise HTTPException(status_code=400, detail=f"删除文档失败: {e}")
# =============================================================================
# === 查询分组 ===
# =============================================================================
@knowledge.post("/databases/{db_id}/query")
async def query_knowledge_base(
db_id: str,
query: str = Body(...),
meta: dict = Body(...),
current_user: User = Depends(get_admin_user)
db_id: str, query: str = Body(...), meta: dict = Body(...), current_user: User = Depends(get_admin_user)
):
"""查询知识库"""
logger.debug(f"Query knowledge base {db_id}: {query}")
@ -212,12 +201,10 @@ async def query_knowledge_base(
logger.error(f"知识库查询失败 {e}, {traceback.format_exc()}")
return {"message": f"知识库查询失败: {e}", "status": "failed"}
@knowledge.post("/databases/{db_id}/query-test")
async def query_test(
db_id: str,
query: str = Body(...),
meta: dict = Body(...),
current_user: User = Depends(get_admin_user)
db_id: str, query: str = Body(...), meta: dict = Body(...), current_user: User = Depends(get_admin_user)
):
"""测试查询知识库"""
logger.debug(f"Query test in {db_id}: {query}")
@ -228,11 +215,9 @@ async def query_test(
logger.error(f"测试查询失败 {e}, {traceback.format_exc()}")
return {"message": f"测试查询失败: {e}", "status": "failed"}
@knowledge.get("/databases/{db_id}/query-params")
async def get_knowledge_base_query_params(
db_id: str,
current_user: User = Depends(get_admin_user)
):
async def get_knowledge_base_query_params(db_id: str, current_user: User = Depends(get_admin_user)):
"""获取知识库类型特定的查询参数"""
try:
# 获取数据库信息
@ -258,21 +243,21 @@ async def get_knowledge_base_query_params(
{"value": "hybrid", "label": "Hybrid", "description": "本地和全局混合"},
{"value": "naive", "label": "Naive", "description": "基本搜索"},
{"value": "mix", "label": "Mix", "description": "知识图谱和向量检索混合"},
]
],
},
{
"key": "only_need_context",
"label": "只使用上下文",
"type": "boolean",
"default": True,
"description": "只返回上下文，不生成回答"
"description": "只返回上下文，不生成回答",
},
{
"key": "only_need_prompt",
"label": "只使用提示",
"type": "boolean",
"default": False,
"description": "只返回提示，不进行检索"
"description": "只返回提示，不进行检索",
},
{
"key": "top_k",
@ -281,9 +266,9 @@ async def get_knowledge_base_query_params(
"default": 10,
"min": 1,
"max": 100,
"description": "返回的最大结果数量"
}
]
"description": "返回的最大结果数量",
},
],
}
elif kb_type == "chroma":
params = {
@ -296,7 +281,7 @@ async def get_knowledge_base_query_params(
"default": 10,
"min": 1,
"max": 100,
"description": "返回的最大结果数量"
"description": "返回的最大结果数量",
},
{
"key": "similarity_threshold",
@ -306,16 +291,16 @@ async def get_knowledge_base_query_params(
"min": 0.0,
"max": 1.0,
"step": 0.1,
"description": "过滤相似度低于此值的结果"
"description": "过滤相似度低于此值的结果",
},
{
"key": "include_distances",
"label": "显示相似度",
"type": "boolean",
"default": True,
"description": "在结果中显示相似度分数"
}
]
"description": "在结果中显示相似度分数",
},
],
}
elif kb_type == "milvus":
params = {
@ -328,7 +313,7 @@ async def get_knowledge_base_query_params(
"default": 10,
"min": 1,
"max": 100,
"description": "返回的最大结果数量"
"description": "返回的最大结果数量",
},
{
"key": "similarity_threshold",
@ -338,14 +323,14 @@ async def get_knowledge_base_query_params(
"min": 0.0,
"max": 1.0,
"step": 0.1,
"description": "过滤相似度低于此值的结果"
"description": "过滤相似度低于此值的结果",
},
{
"key": "include_distances",
"label": "显示相似度",
"type": "boolean",
"default": True,
"description": "在结果中显示相似度分数"
"description": "在结果中显示相似度分数",
},
{
"key": "metric_type",
@ -355,11 +340,11 @@ async def get_knowledge_base_query_params(
"options": [
{"value": "COSINE", "label": "余弦相似度", "description": "适合文本语义相似度"},
{"value": "L2", "label": "欧几里得距离", "description": "适合数值型数据"},
{"value": "IP", "label": "内积", "description": "适合标准化向量"}
{"value": "IP", "label": "内积", "description": "适合标准化向量"},
],
"description": "向量相似度计算方法"
}
]
"description": "向量相似度计算方法",
},
],
}
else:
# 未知类型，返回基本参数
@ -373,9 +358,9 @@ async def get_knowledge_base_query_params(
"default": 10,
"min": 1,
"max": 100,
"description": "返回的最大结果数量"
"description": "返回的最大结果数量",
}
]
],
}
return {"params": params, "message": "success"}
@ -384,15 +369,15 @@ async def get_knowledge_base_query_params(
logger.error(f"获取知识库查询参数失败 {e}, {traceback.format_exc()}")
return {"message": f"获取知识库查询参数失败 {e}", "params": {}}
# =============================================================================
# === 文件管理分组 ===
# =============================================================================
@knowledge.post("/files/upload")
async def upload_file(
file: UploadFile = File(...),
db_id: str | None = Query(None),
current_user: User = Depends(get_admin_user)
file: UploadFile = File(...), db_id: str | None = Query(None), current_user: User = Depends(get_admin_user)
):
"""上传文件"""
if not file.filename:
@ -414,11 +399,9 @@ async def upload_file(
return {"message": "File successfully uploaded", "file_path": file_path, "db_id": db_id}
@knowledge.post("/files/markdown")
async def mark_it_down(
file: UploadFile = File(...),
current_user: User = Depends(get_admin_user)
):
async def mark_it_down(file: UploadFile = File(...), current_user: User = Depends(get_admin_user)):
"""调用 src.knowledge.indexing 下面的 process_file_to_markdown 解析为 markdown，参数是文件，需要管理员权限"""
try:
content = await file.read()
@ -428,10 +411,12 @@ async def mark_it_down(
logger.error(f"文件解析失败 {e}, {traceback.format_exc()}")
return {"message": f"文件解析失败 {e}", "markdown_content": ""}
# =============================================================================
# === 知识库类型分组 ===
# =============================================================================
@knowledge.get("/types")
async def get_knowledge_base_types(current_user: User = Depends(get_admin_user)):
"""获取支持的知识库类型"""
@ -442,6 +427,7 @@ async def get_knowledge_base_types(current_user: User = Depends(get_admin_user))
logger.error(f"获取知识库类型失败 {e}, {traceback.format_exc()}")
return {"message": f"获取知识库类型失败 {e}", "kb_types": {}}
@knowledge.get("/stats")
async def get_knowledge_base_statistics(current_user: User = Depends(get_admin_user)):
"""获取知识库统计信息"""

99
server/routers/system_router.py
View File

@ -1,64 +1,63 @@
import os
import yaml
import requests
from pathlib import Path
from fastapi import Request, Body, Depends, HTTPException
from fastapi import APIRouter
from collections import deque
from pathlib import Path
import requests
import yaml
from fastapi import APIRouter, Body, Depends, HTTPException, Request
from src import config, knowledge_base, graph_base
from server.utils.auth_middleware import get_admin_user, get_superadmin_user
from server.models.user_model import User
from server.utils.auth_middleware import get_admin_user, get_superadmin_user
from src import config, graph_base, knowledge_base
from src.utils.logging_config import logger
system = APIRouter(prefix="/system", tags=["system"])
# =============================================================================
# === 健康检查分组 ===
# =============================================================================
@system.get("/health")
async def health_check():
"""系统健康检查接口（公开接口）"""
return {"status": "ok", "message": "服务正常运行"}
# =============================================================================
# === 配置管理分组 ===
# =============================================================================
@system.get("/config")
def get_config(current_user: User = Depends(get_admin_user)):
"""获取系统配置"""
return config.dump_config()
@system.post("/config")
async def update_config_single(
key = Body(...),
value = Body(...),
current_user: User = Depends(get_admin_user)
) -> dict:
async def update_config_single(key=Body(...), value=Body(...), current_user: User = Depends(get_admin_user)) -> dict:
"""更新单个配置项"""
config[key] = value
config.save()
return config.dump_config()
@system.post("/config/update")
async def update_config_batch(
items: dict = Body(...),
current_user: User = Depends(get_admin_user)
) -> dict:
async def update_config_batch(items: dict = Body(...), current_user: User = Depends(get_admin_user)) -> dict:
"""批量更新配置项"""
config.update(items)
config.save()
return config.dump_config()
@system.post("/restart")
async def restart_system(current_user: User = Depends(get_superadmin_user)):
"""重启系统（仅超级管理员）"""
graph_base.start()
return {"message": "系统已重启"}
@system.get("/logs")
def get_system_logs(current_user: User = Depends(get_admin_user)):
"""获取系统日志"""
@ -68,16 +67,18 @@ def get_system_logs(current_user: User = Depends(get_admin_user)):
with open(LOG_FILE) as f:
last_lines = deque(f, maxlen=1000)
log = ''.join(last_lines)
log = "".join(last_lines)
return {"log": log, "message": "success", "log_file": LOG_FILE}
except Exception as e:
logger.error(f"获取系统日志失败: {e}")
raise HTTPException(status_code=500, detail=f"获取系统日志失败: {str(e)}")
# =============================================================================
# === 信息管理分组 ===
# =============================================================================
def load_info_config():
"""加载信息配置文件"""
try:
@ -91,7 +92,7 @@ def load_info_config():
config_path = Path("src/static/info.template.yaml")
# 读取配置文件
with open(config_path, encoding='utf-8') as file:
with open(config_path, encoding="utf-8") as file:
config = yaml.safe_load(file)
return config
@ -100,61 +101,49 @@ def load_info_config():
logger.error(f"Failed to load info config: {e}")
return get_default_info_config()
def get_default_info_config():
"""获取默认信息配置"""
return {
"organization": {
"name": "江南语析",
"logo": "/favicon.svg",
"avatar": "/avatar.jpg"
},
"organization": {"name": "江南语析", "logo": "/favicon.svg", "avatar": "/avatar.jpg"},
"branding": {
"name": "Yuxi-Know",
"title": "Yuxi-Know",
"subtitle": "大模型驱动的知识库管理工具",
"description": "结合知识库与知识图谱，提供更准确、更全面的回答"
"description": "结合知识库与知识图谱，提供更准确、更全面的回答",
},
"features": [
"📚 灵活知识库",
"🕸️ 知识图谱集成",
"🤖 多模型支持"
],
"footer": {
"copyright": "© 江南语析 2025 [WIP] v0.2.0"
}
"features": ["📚 灵活知识库", "🕸️ 知识图谱集成", "🤖 多模型支持"],
"footer": {"copyright": "© 江南语析 2025 [WIP] v0.2.0"},
}
@system.get("/info")
async def get_info_config():
"""获取系统信息配置（公开接口，无需认证）"""
try:
config = load_info_config()
return {
"success": True,
"data": config
}
return {"success": True, "data": config}
except Exception as e:
logger.error(f"获取信息配置失败: {e}")
raise HTTPException(status_code=500, detail="获取信息配置失败")
@system.post("/info/reload")
async def reload_info_config(current_user: User = Depends(get_admin_user)):
"""重新加载信息配置"""
try:
config = load_info_config()
return {
"success": True,
"message": "配置重新加载成功",
"data": config
}
return {"success": True, "message": "配置重新加载成功", "data": config}
except Exception as e:
logger.error(f"重新加载信息配置失败: {e}")
raise HTTPException(status_code=500, detail="重新加载信息配置失败")
# =============================================================================
# === OCR服务分组 ===
# =============================================================================
@system.get("/ocr/stats")
async def get_ocr_stats(current_user: User = Depends(get_admin_user)):
"""
@ -163,20 +152,13 @@ async def get_ocr_stats(current_user: User = Depends(get_admin_user)):
"""
try:
from src.plugins._ocr import get_ocr_stats
stats = get_ocr_stats()
return {
"status": "success",
"stats": stats,
"message": "OCR统计信息获取成功"
}
return {"status": "success", "stats": stats, "message": "OCR统计信息获取成功"}
except Exception as e:
logger.error(f"获取OCR统计信息失败: {str(e)}")
return {
"status": "error",
"stats": {},
"message": f"获取OCR统计信息失败: {str(e)}"
}
return {"status": "error", "stats": {}, "message": f"获取OCR统计信息失败: {str(e)}"}
@system.get("/ocr/health")
@ -188,12 +170,14 @@ async def check_ocr_services_health(current_user: User = Depends(get_admin_user)
health_status = {
"rapid_ocr": {"status": "unknown", "message": ""},
"mineru_ocr": {"status": "unknown", "message": ""},
"paddlex_ocr": {"status": "unknown", "message": ""}
"paddlex_ocr": {"status": "unknown", "message": ""},
}
# 检查 RapidOCR (ONNX) 模型
try:
model_dir_root = os.getenv("MODEL_DIR") if not os.getenv("RUNNING_IN_DOCKER") else os.getenv("MODEL_DIR_IN_DOCKER")
model_dir_root = (
os.getenv("MODEL_DIR") if not os.getenv("RUNNING_IN_DOCKER") else os.getenv("MODEL_DIR_IN_DOCKER")
)
model_dir = os.path.join(model_dir_root, "SWHL/RapidOCR")
det_model_path = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_det_infer.onnx")
rec_model_path = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_rec_infer.onnx")
@ -201,6 +185,7 @@ async def check_ocr_services_health(current_user: User = Depends(get_admin_user)
if os.path.exists(model_dir) and os.path.exists(det_model_path) and os.path.exists(rec_model_path):
# 尝试初始化RapidOCR
from rapidocr_onnxruntime import RapidOCR
test_ocr = RapidOCR(det_box_thresh=0.3, det_model_path=det_model_path, rec_model_path=rec_model_path) # noqa: F841
health_status["rapid_ocr"]["status"] = "healthy"
health_status["rapid_ocr"]["message"] = "RapidOCR模型已加载"
@ -258,8 +243,4 @@ async def check_ocr_services_health(current_user: User = Depends(get_admin_user)
# 计算整体健康状态
overall_status = "healthy" if any(svc["status"] == "healthy" for svc in health_status.values()) else "unhealthy"
return {
"overall_status": overall_status,
"services": health_status,
"message": "OCR服务健康检查完成"
}
return {"overall_status": overall_status, "services": health_status, "message": "OCR服务健康检查完成"}

23
server/utils/auth_middleware.py
View File

@ -1,8 +1,9 @@
import re
from fastapi import Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer
from sqlalchemy.orm import Session
from jose import JWTError, jwt
import re
from sqlalchemy.orm import Session
from server.db_manager import db_manager
from server.models.user_model import User
@ -13,14 +14,15 @@ oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/auth/token", auto_error=Fals
# 公开路径列表，无需登录即可访问
PUBLIC_PATHS = [
r"^/api/auth/token$", # 登录
r"^/api/auth/token$", # 登录
r"^/api/auth/check-first-run$", # 检查是否首次运行
r"^/api/auth/initialize$", # 初始化系统
r"^/api$", # Health Check
r"^/api/system/health$", # Health Check
r"^/api/system/info$", # 获取系统信息配置
r"^/api/auth/initialize$", # 初始化系统
r"^/api$", # Health Check
r"^/api/system/health$", # Health Check
r"^/api/system/info$", # 获取系统信息配置
]
# 获取数据库会话
def get_db():
db = db_manager.get_session()
@ -29,6 +31,7 @@ def get_db():
finally:
db.close()
# 获取当前用户
async def get_current_user(token: str | None = Depends(oauth2_scheme), db: Session = Depends(get_db)):
credentials_exception = HTTPException(
@ -65,6 +68,7 @@ async def get_current_user(token: str | None = Depends(oauth2_scheme), db: Sessi
return user
# 获取已登录用户（抛出401如果未登录）
async def get_required_user(user: User | None = Depends(get_current_user)):
if user is None:
@ -75,6 +79,7 @@ async def get_required_user(user: User | None = Depends(get_current_user)):
)
return user
# 获取管理员用户
async def get_admin_user(current_user: User = Depends(get_required_user)):
if current_user.role not in ["admin", "superadmin"]:
@ -84,6 +89,7 @@ async def get_admin_user(current_user: User = Depends(get_required_user)):
)
return current_user
# 获取超级管理员用户
async def get_superadmin_user(current_user: User = Depends(get_required_user)):
if current_user.role != "superadmin":
@ -93,9 +99,10 @@ async def get_superadmin_user(current_user: User = Depends(get_required_user)):
)
return current_user
# 检查路径是否为公开路径
def is_public_path(path: str) -> bool:
path = path.rstrip('/') # 去除尾部斜杠以便于匹配
path = path.rstrip("/") # 去除尾部斜杠以便于匹配
for pattern in PUBLIC_PATHS:
if re.match(pattern, path):
return True

4
server/utils/auth_utils.py
View File

@ -1,14 +1,16 @@
import hashlib
import os
import jwt
from datetime import datetime, timedelta
from typing import Any
import jwt
# JWT配置
JWT_SECRET_KEY = os.environ.get("JWT_SECRET_KEY", "yuxi_know_secure_key")
JWT_ALGORITHM = "HS256"
JWT_EXPIRATION = 24 * 60 * 60 # 24小时过期
class AuthUtils:
"""认证工具类"""

25
server/utils/common_utils.py
View File

@ -1,19 +1,18 @@
"""通用工具函数"""
import logging
from sqlalchemy.orm import Session
from fastapi import Request
from server.models.user_model import User, OperationLog
from sqlalchemy.orm import Session
from server.models.user_model import OperationLog, User
def setup_logging():
"""配置应用程序日志格式"""
# 配置日志格式
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s %(levelname)s: %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
force=True
level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s", datefmt="%Y-%m-%d %H:%M:%S", force=True
)
# 确保uvicorn的日志也使用相同格式
@ -21,10 +20,7 @@ def setup_logging():
uvicorn_access_logger = logging.getLogger("uvicorn.access")
# 创建格式化器
formatter = logging.Formatter(
fmt="%(asctime)s %(levelname)s: %(message)s",
datefmt="%m-%d %H:%M:%S"
)
formatter = logging.Formatter(fmt="%(asctime)s %(levelname)s: %(message)s", datefmt="%m-%d %H:%M:%S")
# 为所有处理器设置格式化器
for handler in uvicorn_logger.handlers:
@ -39,12 +35,7 @@ def log_operation(db: Session, user_id: int, operation: str, details: str = None
if request:
ip_address = request.client.host if request.client else None
log = OperationLog(
user_id=user_id,
operation=operation,
details=details,
ip_address=ip_address
)
log = OperationLog(user_id=user_id, operation=operation, details=details, ip_address=ip_address)
db.add(log)
db.commit()
@ -60,6 +51,6 @@ def convert_serializable(obj):
return [convert_serializable(item) for item in obj]
if isinstance(obj, dict):
return {k: convert_serializable(v) for k, v in obj.items()}
if hasattr(obj, '__dict__'):
if hasattr(obj, "__dict__"):
return convert_serializable(vars(obj))
return obj

30
src/__init__.py
View File

@ -1,38 +1,39 @@
import os
from dotenv import load_dotenv
load_dotenv("src/.env", override=True)
from concurrent.futures import ThreadPoolExecutor # noqa: E402
executor = ThreadPoolExecutor()
from src.config import Config # noqa: E402
config = Config()
# 导入知识库相关模块
from src.knowledge.chroma_kb import ChromaKB # noqa: E402
from src.knowledge.kb_factory import KnowledgeBaseFactory # noqa: E402
from src.knowledge.kb_manager import KnowledgeBaseManager # noqa: E402
from src.knowledge.lightrag_kb import LightRagKB # noqa: E402
from src.knowledge.chroma_kb import ChromaKB # noqa: E402
from src.knowledge.milvus_kb import MilvusKB # noqa: E402
# 注册知识库类型
KnowledgeBaseFactory.register("chroma", ChromaKB, {
"chunk_size": 1000,
"chunk_overlap": 200,
"description": "基于 ChromaDB 的轻量级向量知识库，适合开发和小规模部署"
})
KnowledgeBaseFactory.register(
"chroma",
ChromaKB,
{"chunk_size": 1000, "chunk_overlap": 200, "description": "基于 ChromaDB 的轻量级向量知识库，适合开发和小规模部署"},
)
KnowledgeBaseFactory.register("milvus", MilvusKB, {
"chunk_size": 1000,
"chunk_overlap": 200,
"description": "基于 Milvus 的生产级向量知识库，适合大规模高性能部署"
})
KnowledgeBaseFactory.register(
"milvus",
MilvusKB,
{"chunk_size": 1000, "chunk_overlap": 200, "description": "基于 Milvus 的生产级向量知识库，适合大规模高性能部署"},
)
KnowledgeBaseFactory.register("lightrag", LightRagKB, {
"description": "基于图检索的知识库，支持实体关系构建和复杂查询"
})
KnowledgeBaseFactory.register("lightrag", LightRagKB, {"description": "基于图检索的知识库，支持实体关系构建和复杂查询"})
# 创建知识库管理器
@ -40,4 +41,5 @@ work_dir = os.path.join(config.save_dir, "knowledge_base_data")
knowledge_base = KnowledgeBaseManager(work_dir)
from src.knowledge import GraphDatabase # noqa: E402
graph_base = GraphDatabase()

1
src/agents/__init__.py
View File

@ -3,6 +3,7 @@ import asyncio
from .chatbot.graph import ChatbotAgent
from .react.graph import ReActAgent
class AgentManager:
def __init__(self):
self._classes = {}

20
src/agents/chatbot/context.py
View File

@ -1,22 +1,18 @@
from typing import Annotated
from dataclasses import dataclass, field
from typing import Annotated
from src.agents.common.context import BaseContext
from src.agents.common.tools import gen_tool_info
from src.agents.common.mcp import MCP_SERVERS
from src.agents.common.tools import gen_tool_info
from .tools import get_tools
@dataclass(kw_only=True)
class Context(BaseContext):
model: Annotated[str, {"__template_metadata__": {"kind": "llm"}}] = field(
default="siliconflow/Qwen/Qwen3-235B-A22B-Instruct-2507",
metadata={
"name": "智能体模型",
"options": [],
"description": "智能体的驱动模型"
},
metadata={"name": "智能体模型", "options": [], "description": "智能体的驱动模型"},
)
tools: Annotated[list[dict], {"__template_metadata__": {"kind": "tools"}}] = field(
@ -24,15 +20,11 @@ class Context(BaseContext):
metadata={
"name": "工具",
"options": gen_tool_info(get_tools()), # 这里的选择是所有的工具
"description": "工具列表"
"description": "工具列表",
},
)
mcps: list[str] = field(
default_factory=list,
metadata={
"name": "MCP服务器",
"options": list(MCP_SERVERS.keys()),
"description": "MCP服务器列表"
},
metadata={"name": "MCP服务器", "options": list(MCP_SERVERS.keys()), "description": "MCP服务器列表"},
)

29
src/agents/chatbot/graph.py
View File

@ -1,28 +1,26 @@
import os
import uuid
from typing import Any, cast
from pathlib import Path
from typing import Any, cast
from langchain_core.messages import AIMessage, ToolMessage
from langgraph.graph import StateGraph, START, END
from langgraph.runtime import Runtime
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver, aiosqlite
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver, aiosqlite
from langgraph.graph import END, START, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.runtime import Runtime
from src import config as sys_config
from src.utils import logger
from src.agents.common.base import BaseAgent
from src.agents.common.models import load_chat_model
from src.agents.common.mcp import get_mcp_tools
from src.agents.common.models import load_chat_model
from src.utils import logger
from .state import State
from .context import Context
from .state import State
from .tools import get_tools
class ChatbotAgent(BaseAgent):
name = "智能体助手"
description = "基础的对话机器人，可以回答问题，默认不使用任何工具，可在配置中启用需要的工具。"
@ -69,16 +67,11 @@ class ChatbotAgent(BaseAgent):
# 使用异步调用
response = cast(
AIMessage,
await model.ainvoke(
[{"role": "system", "content": runtime.context.system_prompt}, *state.messages]
),
await model.ainvoke([{"role": "system", "content": runtime.context.system_prompt}, *state.messages]),
)
return {"messages": [response]}
async def dynamic_tools_node(
self, state: State, runtime: Runtime[Context]
) -> dict[str, list[ToolMessage]]:
async def dynamic_tools_node(self, state: State, runtime: Runtime[Context]) -> dict[str, list[ToolMessage]]:
"""Execute tools dynamically based on configuration.
This function gets the available tools based on the current configuration
@ -133,6 +126,7 @@ class ChatbotAgent(BaseAgent):
"""获取异步存储实例"""
return AsyncSqliteSaver(await self.get_async_conn())
def main():
agent = ChatbotAgent(Context)
@ -140,6 +134,7 @@ def main():
config = {"configurable": {"thread_id": thread_id}}
from src.agents.utils import agent_cli
agent_cli(agent, config)

8
src/agents/chatbot/state.py
View File

@ -2,12 +2,12 @@
from __future__ import annotations
from dataclasses import dataclass, field
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import Annotated
from langchain_core.messages import AnyMessage
from langgraph.graph import add_messages
from typing import Annotated
@dataclass
@ -17,6 +17,4 @@ class State:
This class is used to define the initial state and structure of incoming data.
"""
messages: Annotated[Sequence[AnyMessage], add_messages] = field(
default_factory=list
)
messages: Annotated[Sequence[AnyMessage], add_messages] = field(default_factory=list)

4
src/agents/chatbot/tools.py
View File

@ -2,8 +2,8 @@ from typing import Any
from langchain_core.tools import tool
from src.utils import logger
from src.agents.common.tools import get_buildin_tools
from src.utils import logger
@tool
@ -26,9 +26,9 @@ def calculator(a: float, b: float, operation: str) -> float:
logger.error(f"Calculator error: {e}")
raise
def get_tools() -> dict[str, Any]:
"""获取所有可运行的工具（给大模型使用）"""
tools = get_buildin_tools()
tools.append(calculator)
return tools

20
src/agents/common/base.py
View File

@ -1,14 +1,14 @@
from __future__ import annotations
from abc import abstractmethod
from langgraph.graph.state import CompiledStateGraph
from src.utils import logger
from src.agents.common.context import BaseContext
from src.utils import logger
class BaseAgent:
"""
定义一个基础 Agent 供 各类 graph 继承
"""
@ -23,7 +23,7 @@ class BaseAgent:
@property
def module_name(self) -> str:
"""Get the module name of the agent class."""
return self.__class__.__module__.split('.')[-2]
return self.__class__.__module__.split(".")[-2]
@property
def id(self) -> str:
@ -42,18 +42,20 @@ class BaseAgent:
async def get_config(self):
return self.context_schema.from_file(module_name=self.module_name)
async def stream_values(self, messages: list[str], input_context = None, **kwargs):
async def stream_values(self, messages: list[str], input_context=None, **kwargs):
graph = await self.get_graph()
context = self.context_schema.from_file(module_name=self.module_name, input_context=input_context)
for event in graph.astream({"messages": messages}, stream_mode="values", context=context):
yield event["messages"]
async def stream_messages(self, messages: list[str], input_context = None, **kwargs):
async def stream_messages(self, messages: list[str], input_context=None, **kwargs):
graph = await self.get_graph()
context = self.context_schema.from_file(module_name=self.module_name, input_context=input_context)
logger.debug(f"stream_messages: {context}")
# TODO 的 Checkpointer 似乎还没有适配最新的 Context API
async for msg, metadata in graph.astream({"messages": messages}, stream_mode="messages", context=context, config={"configurable": input_context}):
async for msg, metadata in graph.astream(
{"messages": messages}, stream_mode="messages", context=context, config={"configurable": input_context}
):
yield msg, metadata
async def check_checkpointer(self):
@ -76,12 +78,12 @@ class BaseAgent:
result = []
if state:
messages = state.values.get('messages', [])
messages = state.values.get("messages", [])
for msg in messages:
if hasattr(msg, 'model_dump'):
if hasattr(msg, "model_dump"):
msg_dict = msg.model_dump() # 转换成字典
else:
msg_dict = dict(msg) if hasattr(msg, '__dict__') else {"content": str(msg)}
msg_dict = dict(msg) if hasattr(msg, "__dict__") else {"content": str(msg)}
result.append(msg_dict)
return result

42
src/agents/common/context.py
View File

@ -3,11 +3,12 @@
from __future__ import annotations
import os
import yaml
import uuid
from dataclasses import dataclass, field, fields, MISSING
from dataclasses import MISSING, dataclass, field, fields
from pathlib import Path
from typing import get_origin, get_args
from typing import get_args, get_origin
import yaml
from src import config as sys_config
from src.utils import logger
@ -32,28 +33,17 @@ class BaseContext:
thread_id: str = field(
default_factory=lambda: str(uuid.uuid4()),
metadata={
"name": "线程ID",
"configurable": False,
"description": "用来描述智能体的角色和行为"
},
metadata={"name": "线程ID", "configurable": False, "description": "用来描述智能体的角色和行为"},
)
user_id: str = field(
default_factory=lambda: str(uuid.uuid4()),
metadata={
"name": "用户ID",
"configurable": False,
"description": "用来描述智能体的角色和行为"
},
metadata={"name": "用户ID", "configurable": False, "description": "用来描述智能体的角色和行为"},
)
system_prompt: str = field(
default="You are a helpful assistant.",
metadata={
"name": "系统提示词",
"description": "用来描述智能体的角色和行为"
},
metadata={"name": "系统提示词", "description": "用来描述智能体的角色和行为"},
)
@classmethod
@ -66,7 +56,7 @@ class BaseContext:
if module_name is not None and os.path.exists(config_file_path):
file_config = {}
try:
with open(config_file_path, encoding='utf-8') as f:
with open(config_file_path, encoding="utf-8") as f:
file_config = yaml.safe_load(f) or {}
except Exception as e:
logger.error(f"加载智能体配置文件出错: {e}")
@ -92,7 +82,7 @@ class BaseContext:
config_file_path = Path(sys_config.save_dir) / "agents" / module_name / "config.yaml"
# 确保目录存在
os.makedirs(os.path.dirname(config_file_path), exist_ok=True)
with open(config_file_path, 'w', encoding='utf-8') as f:
with open(config_file_path, "w", encoding="utf-8") as f:
yaml.dump(configurable_config, f, indent=2, allow_unicode=True)
return True
@ -118,7 +108,11 @@ class BaseContext:
"type": type_name,
"name": f.metadata.get("name", f.name),
"options": f.metadata.get("options", []),
"default": f.default if f.default is not MISSING else f.default_factory() if f.default_factory is not MISSING else None,
"default": f.default
if f.default is not MISSING
else f.default_factory()
if f.default_factory is not MISSING
else None,
"description": f.metadata.get("description", ""),
"template_metadata": template_metadata, # Annotated 的额外元数据
}
@ -132,8 +126,8 @@ class BaseContext:
if get_origin(field_type) is not None:
# 处理泛型类型如 list[str], Annotated[str, {...}]
origin = get_origin(field_type)
if hasattr(origin, '__name__'):
if origin.__name__ == 'Annotated':
if hasattr(origin, "__name__"):
if origin.__name__ == "Annotated":
# Annotated 类型，获取真实类型
args = get_args(field_type)
if args:
@ -141,7 +135,7 @@ class BaseContext:
return origin.__name__
else:
return str(origin)
elif hasattr(field_type, '__name__'):
elif hasattr(field_type, "__name__"):
return field_type.__name__
else:
return str(field_type)
@ -151,7 +145,7 @@ class BaseContext:
"""从 Annotated 类型中提取模板元数据"""
if get_origin(field_type) is not None:
origin = get_origin(field_type)
if hasattr(origin, '__name__') and origin.__name__ == 'Annotated':
if hasattr(origin, "__name__") and origin.__name__ == "Annotated":
args = get_args(field_type)
if len(args) > 1:
# 查找包含 __template_metadata__ 的字典

5
src/agents/common/mcp.py
View File

@ -1,13 +1,13 @@
"""MCP Client setup and management for LangGraph ReAct Agent."""
import traceback
from typing import Any, cast
from collections.abc import Callable
from typing import Any, cast
from langchain_mcp_adapters.tools import load_mcp_tools
from langchain_mcp_adapters.client import ( # type: ignore[import-untyped]
MultiServerMCPClient,
)
from langchain_mcp_adapters.tools import load_mcp_tools
from src.utils import logger
@ -70,6 +70,7 @@ async def get_mcp_tools(server_name: str) -> list[Callable[..., Any]]:
logger.opt(exception=True).warning(f"Failed to load tools from MCP server '{server_name}'")
return []
async def get_all_mcp_tools() -> list[Callable[..., Any]]:
"""Get all tools from all configured MCP servers."""
all_tools = []

11
src/agents/common/models.py
View File

@ -1,13 +1,12 @@
import os
import traceback
from src import config
from src.utils import get_docker_safe_url
from src.models import get_custom_model
from langchain_core.language_models import BaseChatModel
from pydantic import SecretStr
from src import config
from src.models import get_custom_model
from src.utils import get_docker_safe_url
def load_chat_model(fully_specified_name: str, **kwargs) -> BaseChatModel:
@ -18,6 +17,7 @@ def load_chat_model(fully_specified_name: str, **kwargs) -> BaseChatModel:
if provider == "custom":
from langchain_openai import ChatOpenAI
model_info = get_custom_model(model)
api_key = model_info.get("api_key") or "custom_model"
base_url = get_docker_safe_url(model_info["api_base"])
@ -34,6 +34,7 @@ def load_chat_model(fully_specified_name: str, **kwargs) -> BaseChatModel:
if provider in ["deepseek", "dashscope"]:
from langchain_deepseek import ChatDeepSeek
return ChatDeepSeek(
model=model,
api_key=SecretStr(api_key),
@ -43,6 +44,7 @@ def load_chat_model(fully_specified_name: str, **kwargs) -> BaseChatModel:
elif provider == "together":
from langchain_together import ChatTogether
return ChatTogether(
model=model,
api_key=SecretStr(api_key),
@ -52,6 +54,7 @@ def load_chat_model(fully_specified_name: str, **kwargs) -> BaseChatModel:
else:
try: # 其他模型，默认使用OpenAIBase, like openai, zhipuai
from langchain_openai import ChatOpenAI
return ChatOpenAI(
model=model,
api_key=SecretStr(api_key),

53
src/agents/common/tools.py
View File

@ -1,10 +1,11 @@
import inspect
import asyncio
import inspect
import traceback
from typing import Annotated, Any
from pydantic import BaseModel, Field
from langchain_core.tools import StructuredTool, tool
from langchain_tavily import TavilySearch
from pydantic import BaseModel, Field
from src import config, graph_base, knowledge_base
from src.utils import logger
@ -15,13 +16,16 @@ def query_knowledge_graph(query: Annotated[str, "The keyword to query knowledge
"""Use this to query knowledge graph, which include some food domain knowledge."""
try:
logger.debug(f"Querying knowledge graph with: {query}")
result = graph_base.query_node(query, hops=2, return_format='triples')
logger.debug(f"Knowledge graph query returned {len(result.get('triples', [])) if isinstance(result, dict) else 'N/A'} triples")
result = graph_base.query_node(query, hops=2, return_format="triples")
logger.debug(
f"Knowledge graph query returned {len(result.get('triples', [])) if isinstance(result, dict) else 'N/A'} triples"
)
return result
except Exception as e:
logger.error(f"Knowledge graph query error: {e}, {traceback.format_exc()}")
return f"知识图谱查询失败: {str(e)}"
def get_static_tools() -> list:
"""注册静态工具"""
static_tools = [
@ -38,13 +42,11 @@ def get_static_tools() -> list:
class KnowledgeRetrieverModel(BaseModel):
query_text: str = Field(
description=(
"查询的关键词，查询的时候，应该尽量以可能帮助回答这个问题的关键词进行查询，"
"不要直接使用用户的原始输入去查询。"
"查询的关键词，查询的时候，应该尽量以可能帮助回答这个问题的关键词进行查询，不要直接使用用户的原始输入去查询。"
)
)
def get_kb_based_tools() -> list:
"""获取所有知识库基于的工具"""
# 获取所有知识库
@ -53,6 +55,7 @@ def get_kb_based_tools() -> list:
def _create_retriever_wrapper(db_id: str, retriever_info: dict[str, Any]):
"""创建检索器包装函数的工厂函数，避免闭包变量捕获问题"""
async def async_retriever_wrapper(query_text: str) -> Any:
"""异步检索器包装函数"""
retriever = retriever_info["retriever"]
@ -70,17 +73,13 @@ def get_kb_based_tools() -> list:
return async_retriever_wrapper
for db_id, retrieve_info in retrievers.items():
try:
# 使用改进的工具ID生成策略
tool_id = f"query_{db_id[:8]}"
# 构建工具描述
description = (
f"使用 {retrieve_info['name']} 知识库进行检索。\n"
f"下面是这个知识库的描述：\n{retrieve_info['description'] or '没有描述。'} "
)
description = f"使用 {retrieve_info['name']} 知识库进行检索。\n下面是这个知识库的描述：\n{retrieve_info['description'] or '没有描述。'} "
# 使用工厂函数创建检索器包装函数，避免闭包问题
retriever_wrapper = _create_retriever_wrapper(db_id, retrieve_info)
@ -91,9 +90,7 @@ def get_kb_based_tools() -> list:
name=tool_id,
description=description,
args_schema=KnowledgeRetrieverModel,
metadata=retrieve_info["metadata"] | {
"tag": ["knowledgebase"]
}
metadata=retrieve_info["metadata"] | {"tag": ["knowledgebase"]},
)
kb_tools.append(tool)
@ -129,30 +126,34 @@ def gen_tool_info(tools) -> list[dict[str, Any]]:
# 获取注册的工具信息
for tool_obj in tools:
try:
metadata = getattr(tool_obj, 'metadata', {}) or {}
metadata = getattr(tool_obj, "metadata", {}) or {}
info = {
"id": tool_obj.name,
"name": metadata.get('name', tool_obj.name),
"name": metadata.get("name", tool_obj.name),
"description": tool_obj.description,
'metadata': metadata,
"metadata": metadata,
"args": [],
# "is_async": is_async # Include async information
}
if hasattr(tool_obj, 'args_schema') and tool_obj.args_schema:
if hasattr(tool_obj, "args_schema") and tool_obj.args_schema:
schema = tool_obj.args_schema.schema()
for arg_name, arg_info in schema.get('properties', {}).items():
info["args"].append({
"name": arg_name,
"type": arg_info.get('type', ''),
"description": arg_info.get('description', '')
})
for arg_name, arg_info in schema.get("properties", {}).items():
info["args"].append(
{
"name": arg_name,
"type": arg_info.get("type", ""),
"description": arg_info.get("description", ""),
}
)
tools_info.append(info)
# logger.debug(f"Successfully processed tool info for {tool_obj.name}")
except Exception as e:
logger.error(f"Failed to process tool {getattr(tool_obj, 'name', 'unknown')}: {e}\n{traceback.format_exc()}")
logger.error(
f"Failed to process tool {getattr(tool_obj, 'name', 'unknown')}: {e}\n{traceback.format_exc()}"
)
continue
except Exception as e:

20
src/agents/common/utils.py
View File

@ -1,18 +1,17 @@
from datetime import datetime, timezone, UTC
import asyncio
import os
import traceback
from datetime import UTC, datetime, timezone
from src import config
from src.utils import get_docker_safe_url
from src.models import get_custom_model
from src.agents.common.base import BaseAgent
from langchain_core.language_models import BaseChatModel
from langchain_core.runnables import RunnableConfig
from langchain_core.messages import AIMessageChunk, ToolMessage
from langchain_core.runnables import RunnableConfig
from pydantic import SecretStr
from src import config
from src.agents.common.base import BaseAgent
from src.models import get_custom_model
from src.utils import get_docker_safe_url
def load_chat_model(fully_specified_name: str, **kwargs) -> BaseChatModel:
@ -23,6 +22,7 @@ def load_chat_model(fully_specified_name: str, **kwargs) -> BaseChatModel:
if provider == "custom":
from langchain_openai import ChatOpenAI
model_info = get_custom_model(model)
api_key = model_info.get("api_key") or "custom_model"
base_url = get_docker_safe_url(model_info["api_base"])
@ -39,6 +39,7 @@ def load_chat_model(fully_specified_name: str, **kwargs) -> BaseChatModel:
if provider in ["deepseek", "dashscope"]:
from langchain_deepseek import ChatDeepSeek
return ChatDeepSeek(
model=model,
api_key=SecretStr(api_key),
@ -48,6 +49,7 @@ def load_chat_model(fully_specified_name: str, **kwargs) -> BaseChatModel:
elif provider == "together":
from langchain_together import ChatTogether
return ChatTogether(
model=model,
api_key=SecretStr(api_key),
@ -57,6 +59,7 @@ def load_chat_model(fully_specified_name: str, **kwargs) -> BaseChatModel:
else:
try: # 其他模型，默认使用OpenAIBase, like openai, zhipuai
from langchain_openai import ChatOpenAI
return ChatOpenAI(
model=model,
api_key=SecretStr(api_key),
@ -66,7 +69,6 @@ def load_chat_model(fully_specified_name: str, **kwargs) -> BaseChatModel:
raise ValueError(f"Model provider {provider} load failed, {e} \n {traceback.format_exc()}")
async def agent_cli(agent: BaseAgent, config: RunnableConfig | None = None):
config = config or {}
if "configurable" not in config:
@ -100,6 +102,6 @@ async def agent_cli(agent: BaseAgent, config: RunnableConfig | None = None):
if isinstance(msg, ToolMessage):
print(f"Tool: {msg.content}")
def get_cur_time_with_utc():
return datetime.now(tz=UTC).isoformat()

16
src/agents/react/graph.py
View File

@ -1,26 +1,27 @@
import os
from pathlib import Path
from langchain_core.messages import AnyMessage, SystemMessage
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver, aiosqlite
from langgraph.prebuilt import create_react_agent
from langchain_core.messages import AnyMessage, SystemMessage
from langgraph.runtime import get_runtime
from src import config as sys_config
from src.utils import logger
from src.agents.common.context import BaseContext
from src.agents.common.base import BaseAgent
from src.agents.common.context import BaseContext
from src.agents.common.models import load_chat_model
from src.agents.common.tools import get_buildin_tools
from src.utils import logger
model = load_chat_model("siliconflow/Qwen/Qwen3-235B-A22B-Instruct-2507")
def prompt(state) -> list[AnyMessage]:
runtime = get_runtime(BaseContext)
system_msg = SystemMessage(content=runtime.context.system_prompt)
return [system_msg] + state["messages"]
class ReActAgent(BaseAgent):
name = "ReAct (all tools)"
description = "A react agent that can answer questions and help with tasks."
@ -38,12 +39,7 @@ class ReActAgent(BaseAgent):
available_tools = get_buildin_tools()
sqlite_checkpointer = AsyncSqliteSaver(await aiosqlite.connect(self.workdir / "react_history.db"))
graph = create_react_agent(
model,
tools=available_tools,
checkpointer=sqlite_checkpointer,
prompt=prompt
)
graph = create_react_agent(model, tools=available_tools, checkpointer=sqlite_checkpointer, prompt=prompt)
self.graph = graph
logger.info("ReActAgent使用SQLite checkpointer构建成功")
return graph

66
src/config/__init__.py
View File

@ -1,13 +1,15 @@
import os
import json
import yaml
import os
from pathlib import Path
import yaml
from src.utils.logging_config import logger
DEFAULT_MOCK_API = 'this_is_mock_api_key_in_frontend'
DEFAULT_MOCK_API = "this_is_mock_api_key_in_frontend"
class SimpleConfig(dict):
def __key(self, key):
return "" if key is None else key # 目前忘记了这里为什么要 lower 了，只能说配置项最好不要有大写的
@ -35,11 +37,10 @@ class SimpleConfig(dict):
class Config(SimpleConfig):
def __init__(self):
super().__init__()
self._config_items = {}
self.save_dir = os.getenv('SAVE_DIR', 'saves')
self.save_dir = os.getenv("SAVE_DIR", "saves")
self.filename = str(Path(f"{self.save_dir}/config/base.yaml"))
os.makedirs(os.path.dirname(self.filename), exist_ok=True)
@ -48,7 +49,11 @@ class Config(SimpleConfig):
### >>> 默认配置
# 功能选项
self.add_item("enable_reranker", default=False, des="是否开启重排序")
self.add_item("enable_web_search", default=False, des="是否开启网页搜索（注：现阶段会根据 TAVILY_API_KEY 自动开启，无法手动配置，将会在下个版本移除此配置项）") # noqa: E501
self.add_item(
"enable_web_search",
default=False,
des="是否开启网页搜索（注：现阶段会根据 TAVILY_API_KEY 自动开启，无法手动配置，将会在下个版本移除此配置项）",
) # noqa: E501
# 默认智能体配置
self.add_item("default_agent_id", default="", des="默认智能体ID")
# 模型配置
@ -57,8 +62,18 @@ class Config(SimpleConfig):
self.add_item("model_provider", default="siliconflow", des="模型提供商", choices=list(self.model_names.keys()))
self.add_item("model_name", default="zai-org/GLM-4.5", des="模型名称")
self.add_item("embed_model", default="siliconflow/BAAI/bge-m3", des="Embedding 模型", choices=list(self.embed_model_names.keys()))
self.add_item("reranker", default="siliconflow/BAAI/bge-reranker-v2-m3", des="Re-Ranker 模型", choices=list(self.reranker_names.keys())) # noqa: E501
self.add_item(
"embed_model",
default="siliconflow/BAAI/bge-m3",
des="Embedding 模型",
choices=list(self.embed_model_names.keys()),
)
self.add_item(
"reranker",
default="siliconflow/BAAI/bge-reranker-v2-m3",
des="Re-Ranker 模型",
choices=list(self.reranker_names.keys()),
) # noqa: E501
### <<< 默认配置结束
self.load()
@ -66,11 +81,7 @@ class Config(SimpleConfig):
def add_item(self, key, default, des=None, choices=None):
self.__setattr__(key, default)
self._config_items[key] = {
"default": default,
"des": des,
"choices": choices
}
self._config_items[key] = {"default": default, "des": des, "choices": choices}
def __dict__(self):
blocklist = [
@ -87,12 +98,12 @@ class Config(SimpleConfig):
从 models.yaml 和 models.private.yml 中更新 MODEL_NAMES
"""
with open(Path("src/static/models.yaml"), encoding='utf-8') as f:
with open(Path("src/static/models.yaml"), encoding="utf-8") as f:
_models = yaml.safe_load(f)
# 尝试打开一个 models.private.yml 文件，用来覆盖 models.yaml 中的配置
try:
with open(Path("src/static/models.private.yml"), encoding='utf-8') as f:
with open(Path("src/static/models.private.yml"), encoding="utf-8") as f:
_models_private = yaml.safe_load(f)
except FileNotFoundError:
_models_private = {}
@ -110,7 +121,7 @@ class Config(SimpleConfig):
"EMBED_MODEL_INFO": self.embed_model_names,
"RERANKER_LIST": self.reranker_names,
}
with open(Path("src/static/models.private.yml"), 'w', encoding='utf-8') as f:
with open(Path("src/static/models.private.yml"), "w", encoding="utf-8") as f:
yaml.dump(_models, f, indent=2, allow_unicode=True)
def handle_self(self):
@ -121,10 +132,14 @@ class Config(SimpleConfig):
if self.model_dir:
if os.path.exists(self.model_dir):
logger.debug(f"The model directory （{self.model_dir}） contains the following folders: {os.listdir(self.model_dir)}")
logger.debug(
f"The model directory （{self.model_dir}） contains the following folders: {os.listdir(self.model_dir)}"
)
else:
logger.warning(f"Warning: The model directory （{self.model_dir}） does not exist. If not configured, please ignore it. If configured, please check if the configuration is correct;"
"For example, the mapping in the docker-compose file")
logger.warning(
f"Warning: The model directory （{self.model_dir}） does not exist. If not configured, please ignore it. If configured, please check if the configuration is correct;"
"For example, the mapping in the docker-compose file"
)
# 检查模型提供商的环境变量
conds = {}
@ -138,13 +153,14 @@ class Config(SimpleConfig):
self.enable_web_search = True
self.valuable_model_provider = [k for k, v in self.model_provider_status.items() if v]
assert len(self.valuable_model_provider) > 0, f"No model provider available, please check your `.env` file. API_KEY_LIST: {conds}"
assert len(self.valuable_model_provider) > 0, (
f"No model provider available, please check your `.env` file. API_KEY_LIST: {conds}"
)
def load(self):
"""根据传入的文件覆盖掉默认配置"""
logger.info(f"Loading config from {self.filename}")
if self.filename is not None and os.path.exists(self.filename):
if self.filename.endswith(".json"):
with open(self.filename) as f:
content = f.read()
@ -173,14 +189,14 @@ class Config(SimpleConfig):
os.makedirs(os.path.dirname(self.filename), exist_ok=True)
if self.filename.endswith(".json"):
with open(self.filename, 'w+') as f:
with open(self.filename, "w+") as f:
json.dump(self.__dict__(), f, indent=4, ensure_ascii=False)
elif self.filename.endswith(".yaml"):
with open(self.filename, 'w+') as f:
with open(self.filename, "w+") as f:
yaml.dump(self.__dict__(), f, indent=2, allow_unicode=True)
else:
logger.warning(f"Unknown config file type {self.filename}, save as json")
with open(self.filename, 'w+') as f:
with open(self.filename, "w+") as f:
json.dump(self, f, indent=4)
logger.info(f"Config file {self.filename} saved")

90
src/knowledge/chroma_kb.py
View File

@ -1,21 +1,26 @@
import json
import os
import time
import traceback
import json
from pathlib import Path
from typing import Optional, Any
from datetime import datetime
from pathlib import Path
from typing import Any, Optional
import chromadb
from chromadb.api.types import Documents, EmbeddingFunction, Embeddings
from chromadb.config import Settings
from chromadb.api.types import EmbeddingFunction, Documents, Embeddings
from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction
from src.knowledge.indexing import process_url_to_markdown, process_file_to_markdown
from src.knowledge.knowledge_base import KnowledgeBase
from src.knowledge.kb_utils import split_text_into_chunks, split_text_into_qa_chunks, prepare_item_metadata, get_embedding_config
from src.utils import logger, hashstr
from src import config
from src.knowledge.indexing import process_file_to_markdown, process_url_to_markdown
from src.knowledge.kb_utils import (
get_embedding_config,
prepare_item_metadata,
split_text_into_chunks,
split_text_into_qa_chunks,
)
from src.knowledge.knowledge_base import KnowledgeBase
from src.utils import hashstr, logger
class ChromaKB(KnowledgeBase):
@ -40,8 +45,7 @@ class ChromaKB(KnowledgeBase):
# 初始化 ChromaDB 客户端
self.chroma_client = chromadb.PersistentClient(
path=self.chroma_db_path,
settings=Settings(anonymized_telemetry=False)
path=self.chroma_db_path, settings=Settings(anonymized_telemetry=False)
)
# 存储集合映射 {db_id: collection}
@ -68,10 +72,7 @@ class ChromaKB(KnowledgeBase):
try:
# 尝试获取现有集合
collection = self.chroma_client.get_collection(
name=collection_name,
embedding_function=embedding_function
)
collection = self.chroma_client.get_collection(name=collection_name, embedding_function=embedding_function)
logger.info(f"Retrieved existing collection: {collection_name}")
# 检查现有集合的配置是否匹配当前的 embed_info
@ -82,7 +83,9 @@ class ChromaKB(KnowledgeBase):
logger.debug(f"Collection {collection_name} uses model '{current_model}', but expected '{expected_model}'.")
# 如果模型不匹配，删除现有集合并重新创建
if current_model != expected_model:
logger.warning(f"Collection {collection_name} uses model '{current_model}', but expected '{expected_model}'. Recreating collection.")
logger.warning(
f"Collection {collection_name} uses model '{current_model}', but expected '{expected_model}'. Recreating collection."
)
self.chroma_client.delete_collection(name=collection_name)
raise Exception("Model mismatch, recreating collection")
@ -92,12 +95,10 @@ class ChromaKB(KnowledgeBase):
collection_metadata = {
"db_id": db_id,
"created_at": datetime.now().isoformat(),
"embedding_model": embed_info.get("name") if embed_info else "default"
"embedding_model": embed_info.get("name") if embed_info else "default",
}
collection = self.chroma_client.create_collection(
name=collection_name,
embedding_function=embedding_function,
metadata=collection_metadata
name=collection_name, embedding_function=embedding_function, metadata=collection_metadata
)
logger.info(f"Created new collection: {collection_name}")
@ -114,7 +115,7 @@ class ChromaKB(KnowledgeBase):
return OpenAIEmbeddingFunction(
model_name=config_dict["model"],
api_key=config_dict["api_key"],
api_base=config_dict["base_url"].replace('/embeddings', '')
api_base=config_dict["base_url"].replace("/embeddings", ""),
)
async def _get_chroma_collection(self, db_id: str):
@ -141,11 +142,11 @@ class ChromaKB(KnowledgeBase):
def _split_text_into_chunks(self, text: str, file_id: str, filename: str, params: dict) -> list[dict]:
"""将文本分割成块"""
# 检查是否使用QA分割模式
use_qa_split = params.get('use_qa_split', False)
use_qa_split = params.get("use_qa_split", False)
if use_qa_split:
# 使用QA分割模式
qa_separator = params.get('qa_separator', '\n\n\n')
qa_separator = params.get("qa_separator", "\n\n\n")
chunks = split_text_into_qa_chunks(text, file_id, filename, qa_separator, params)
else:
# 使用传统分割模式
@ -157,13 +158,12 @@ class ChromaKB(KnowledgeBase):
"source": chunk["source"],
"chunk_id": chunk["chunk_id"],
"full_doc_id": file_id,
"chunk_type": chunk.get("chunk_type", "normal") # 添加chunk类型标识
"chunk_type": chunk.get("chunk_type", "normal"), # 添加chunk类型标识
}
return chunks
async def add_content(self, db_id: str, items: list[str],
params:dict | None) -> list[dict]:
async def add_content(self, db_id: str, items: list[str], params: dict | None) -> list[dict]:
"""添加内容（文件/URL）"""
if db_id not in self.databases_meta:
raise ValueError(f"Database {db_id} not found")
@ -172,7 +172,7 @@ class ChromaKB(KnowledgeBase):
if not collection:
raise ValueError(f"Failed to get ChromaDB collection for {db_id}")
content_type = params.get('content_type', 'file') if params else 'file'
content_type = params.get("content_type", "file") if params else "file"
processed_items_info = []
for item in items:
@ -205,24 +205,20 @@ class ChromaKB(KnowledgeBase):
ids = [chunk["id"] for chunk in chunks]
# 插入到 ChromaDB
collection.add(
documents=documents,
metadatas=metadatas,
ids=ids
)
collection.add(documents=documents, metadatas=metadatas, ids=ids)
logger.info(f"Inserted {content_type} {item} into ChromaDB. Done.")
# 更新状态为完成
self.files_meta[file_id]["status"] = "done"
self._save_metadata()
file_record['status'] = "done"
file_record["status"] = "done"
except Exception as e:
logger.error(f"处理{content_type} {item} 失败: {e}, {traceback.format_exc()}")
self.files_meta[file_id]["status"] = "failed"
self._save_metadata()
file_record['status'] = "failed"
file_record["status"] = "failed"
finally:
self._remove_from_processing_queue(file_id)
@ -241,9 +237,7 @@ class ChromaKB(KnowledgeBase):
similarity_threshold = kwargs.get("similarity_threshold", 0.0)
results = collection.query(
query_texts=[query_text],
n_results=top_k,
include=["documents", "metadatas", "distances"]
query_texts=[query_text], n_results=top_k, include=["documents", "metadatas", "distances"]
)
if not results or not results.get("documents") or not results["documents"][0]:
@ -262,14 +256,10 @@ class ChromaKB(KnowledgeBase):
metadata = metadatas[i] if i < len(metadatas) else {}
# 确保 file_id 在元数据中，并使用统一的键名
if 'full_doc_id' in metadata:
metadata['file_id'] = metadata.pop('full_doc_id')
if "full_doc_id" in metadata:
metadata["file_id"] = metadata.pop("full_doc_id")
retrieved_chunks.append({
"content": doc,
"metadata": metadata,
"score": similarity
})
retrieved_chunks.append({"content": doc, "metadata": metadata, "score": similarity})
logger.debug(f"ChromaDB query response: {len(retrieved_chunks)} chunks found (after similarity filtering)")
return retrieved_chunks
@ -284,10 +274,7 @@ class ChromaKB(KnowledgeBase):
if collection:
try:
# 查找所有相关的chunks
results = collection.get(
where={"full_doc_id": file_id},
include=["metadatas"]
)
results = collection.get(where={"full_doc_id": file_id}, include=["metadatas"])
# 删除所有相关chunks
if results and results.get("ids"):
@ -312,10 +299,7 @@ class ChromaKB(KnowledgeBase):
if collection:
try:
# 获取文档的所有chunks
results = collection.get(
where={"full_doc_id": file_id},
include=["documents", "metadatas"]
)
results = collection.get(where={"full_doc_id": file_id}, include=["documents", "metadatas"])
# 构建chunks数据
doc_chunks = []
@ -325,7 +309,9 @@ class ChromaKB(KnowledgeBase):
"id": chunk_id,
"content": results["documents"][i] if i < len(results["documents"]) else "",
"metadata": results["metadatas"][i] if i < len(results["metadatas"]) else {},
"chunk_order_index": results["metadatas"][i].get("chunk_index", i) if i < len(results["metadatas"]) else i
"chunk_order_index": results["metadatas"][i].get("chunk_index", i)
if i < len(results["metadatas"])
else i,
}
doc_chunks.append(chunk_data)

189
src/knowledge/graphbase.py
View File

@ -1,7 +1,7 @@
import os
import json
import warnings
import os
import traceback
import warnings
from neo4j import GraphDatabase as GD
from neo4j import Query
@ -15,6 +15,7 @@ warnings.filterwarnings("ignore", category=UserWarning)
UIE_MODEL = None
class GraphDatabase:
def __init__(self):
self.driver = None
@ -55,10 +56,11 @@ class GraphDatabase:
"""检查图数据库是否正在运行"""
return self.status == "open" or self.status == "processing"
def get_sample_nodes(self, kgdb_name='neo4j', num=50):
def get_sample_nodes(self, kgdb_name="neo4j", num=50):
"""获取指定数据库的 num 个节点信息，优先返回连通的节点子图"""
assert self.driver is not None, "Database is not connected"
self.use_database(kgdb_name)
def query(tx, num):
"""Note: 使用连通性查询获取集中的节点子图"""
# 首先尝试获取一个连通的子图
@ -104,25 +106,25 @@ class GraphDatabase:
try:
results = tx.run(query_str, num=int(num))
formatted_results = {'nodes': [], 'edges': []}
formatted_results = {"nodes": [], "edges": []}
node_ids = set()
for item in results:
h_node = item['h']
t_node = item['t']
h_node = item["h"]
t_node = item["t"]
# 避免重复添加节点
if h_node['id'] not in node_ids:
formatted_results['nodes'].append(h_node)
node_ids.add(h_node['id'])
if t_node['id'] not in node_ids:
formatted_results['nodes'].append(t_node)
node_ids.add(t_node['id'])
if h_node["id"] not in node_ids:
formatted_results["nodes"].append(h_node)
node_ids.add(h_node["id"])
if t_node["id"] not in node_ids:
formatted_results["nodes"].append(t_node)
node_ids.add(t_node["id"])
formatted_results['edges'].append(item['r'])
formatted_results["edges"].append(item["r"])
# 如果连通查询没有返回足够的结果，回退到原始查询
if len(formatted_results['nodes']) < num // 2:
if len(formatted_results["nodes"]) < num // 2:
fallback_query = """
MATCH (n:Entity)-[r]-(m:Entity)
WHERE elementId(n) < elementId(m)
@ -133,11 +135,11 @@ class GraphDatabase:
LIMIT $num
"""
fallback_results = tx.run(fallback_query, num=int(num))
formatted_results = {'nodes': [], 'edges': []}
formatted_results = {"nodes": [], "edges": []}
for item in fallback_results:
formatted_results['nodes'].extend([item['h'], item['t']])
formatted_results['edges'].append(item['r'])
formatted_results["nodes"].extend([item["h"], item["t"]])
formatted_results["edges"].append(item["r"])
return formatted_results
@ -154,11 +156,11 @@ class GraphDatabase:
LIMIT $num
"""
results = tx.run(fallback_query, num=int(num))
formatted_results = {'nodes': [], 'edges': []}
formatted_results = {"nodes": [], "edges": []}
for item in results:
formatted_results['nodes'].extend([item['h'], item['t']])
formatted_results['edges'].append(item['r'])
formatted_results["nodes"].extend([item["h"], item["t"]])
formatted_results["edges"].append(item["r"])
return formatted_results
@ -171,7 +173,7 @@ class GraphDatabase:
assert self.driver is not None, "Database is not connected"
with self.driver.session() as session:
existing_databases = session.run("SHOW DATABASES")
existing_db_names = [db['name'] for db in existing_databases]
existing_db_names = [db["name"] for db in existing_databases]
if existing_db_names:
print(f"已存在数据库: {existing_db_names[0]}")
@ -183,15 +185,17 @@ class GraphDatabase:
def use_database(self, kgdb_name="neo4j"):
"""切换到指定数据库"""
assert kgdb_name == self.kgdb_name, f"传入的数据库名称 '{kgdb_name}' 与当前实例的数据库名称 '{self.kgdb_name}' 不一致"
assert kgdb_name == self.kgdb_name, (
f"传入的数据库名称 '{kgdb_name}' 与当前实例的数据库名称 '{self.kgdb_name}' 不一致"
)
if self.status == "closed":
self.start()
async def txt_add_vector_entity(self, triples, kgdb_name='neo4j'):
async def txt_add_vector_entity(self, triples, kgdb_name="neo4j"):
"""添加实体三元组"""
assert self.driver is not None, "Database is not connected"
self.use_database(kgdb_name)
def _index_exists(tx, index_name):
"""检查索引是否存在"""
result = tx.run("SHOW INDEXES")
@ -203,11 +207,16 @@ class GraphDatabase:
def _create_graph(tx, data):
"""添加一个三元组"""
for entry in data:
tx.run("""
tx.run(
"""
MERGE (h:Entity:Upload {name: $h})
MERGE (t:Entity:Upload {name: $t})
MERGE (h)-[r:RELATION {type: $r}]->(t)
""", h=entry['h'], t=entry['t'], r=entry['r'])
""",
h=entry["h"],
t=entry["t"],
r=entry["r"],
)
def _create_vector_index(tx, dim):
"""创建向量索引"""
@ -232,42 +241,50 @@ class GraphDatabase:
param_placeholders = ", ".join([f"${key}" for key in params.keys()])
# 执行查询
result = tx.run(f"""
result = tx.run(
f"""
MATCH (n:Entity)
WHERE n.name IN [{param_placeholders}] AND n.embedding IS NULL
RETURN n.name AS name
""", params)
""",
params,
)
return [record["name"] for record in result]
def _batch_set_embeddings(tx, entity_embedding_pairs):
"""批量设置实体的嵌入向量"""
for entity_name, embedding in entity_embedding_pairs:
tx.run("""
tx.run(
"""
MATCH (e:Entity {name: $name})
CALL db.create.setNodeVectorProperty(e, 'embedding', $embedding)
""", name=entity_name, embedding=embedding)
""",
name=entity_name,
embedding=embedding,
)
# 判断模型名称是否匹配
self.embed_model_name = self.embed_model_name or config.embed_model
cur_embed_info = config.embed_model_names.get(self.embed_model_name)
logger.warning(f"embed_model_name={self.embed_model_name}, {cur_embed_info=}")
assert self.embed_model_name == config.embed_model or self.embed_model_name is None, \
assert self.embed_model_name == config.embed_model or self.embed_model_name is None, (
f"embed_model_name={self.embed_model_name}, {config.embed_model=}"
)
with self.driver.session() as session:
logger.info(f"Adding entity to {kgdb_name}")
session.execute_write(_create_graph, triples)
logger.info(f"Creating vector index for {kgdb_name} with {config.embed_model}")
session.execute_write(_create_vector_index, cur_embed_info['dimension'])
session.execute_write(_create_vector_index, cur_embed_info["dimension"])
# 收集所有需要处理的实体名称，去重
all_entities = []
for entry in triples:
if entry['h'] not in all_entities:
all_entities.append(entry['h'])
if entry['t'] not in all_entities:
all_entities.append(entry['t'])
if entry["h"] not in all_entities:
all_entities.append(entry["h"])
if entry["t"] not in all_entities:
all_entities.append(entry["t"])
# 筛选出没有embedding的节点
nodes_without_embedding = session.execute_read(_get_nodes_without_embedding, all_entities)
@ -282,11 +299,9 @@ class GraphDatabase:
total_entities = len(nodes_without_embedding)
for i in range(0, total_entities, max_batch_size):
batch_entities = nodes_without_embedding[i:i+max_batch_size]
batch_entities = nodes_without_embedding[i : i + max_batch_size]
logger.debug(
f"Processing entities batch "
f"{i//max_batch_size + 1}/{(total_entities-1)//max_batch_size + 1} "
f"({len(batch_entities)} entities)"
f"Processing entities batch {i // max_batch_size + 1}/{(total_entities - 1) // max_batch_size + 1} ({len(batch_entities)} entities)"
)
# 批量获取嵌入向量
@ -301,15 +316,15 @@ class GraphDatabase:
# 数据添加完成后保存图信息
self.save_graph_info()
async def jsonl_file_add_entity(self, file_path, kgdb_name='neo4j'):
async def jsonl_file_add_entity(self, file_path, kgdb_name="neo4j"):
assert self.driver is not None, "Database is not connected"
self.status = "processing"
kgdb_name = kgdb_name or 'neo4j'
kgdb_name = kgdb_name or "neo4j"
self.use_database(kgdb_name) # 切换到指定数据库
logger.info(f"Start adding entity to {kgdb_name} with {file_path}")
def read_triples(file_path):
with open(file_path, encoding='utf-8') as file:
with open(file_path, encoding="utf-8") as file:
for line in file:
if line.strip():
yield json.loads(line.strip())
@ -347,7 +362,9 @@ class GraphDatabase:
"""
tx.run(query)
def query_node(self, entity_name, threshold=0.9, kgdb_name='neo4j', hops=2, max_entities=5, return_format='graph', **kwargs):
def query_node(
self, entity_name, threshold=0.9, kgdb_name="neo4j", hops=2, max_entities=5, return_format="graph", **kwargs
):
"""知识图谱查询节点的入口:"""
assert self.driver is not None, "Database is not connected"
assert self.is_running(), "图数据库未启动"
@ -363,29 +380,33 @@ class GraphDatabase:
logger.debug(f"Graph Query Entities: {entity_name}, {qualified_entities=}")
# 对每个合格的实体进行查询
all_query_results = {'nodes': [], 'edges': [], 'triples': []}
all_query_results = {"nodes": [], "edges": [], "triples": []}
for entity in qualified_entities:
query_result = self._query_specific_entity(entity_name=entity, kgdb_name=kgdb_name, hops=hops)
if return_format == 'graph':
all_query_results['nodes'].extend(query_result['nodes'])
all_query_results['edges'].extend(query_result['edges'])
elif return_format == 'triples':
all_query_results['triples'].extend(query_result['triples'])
if return_format == "graph":
all_query_results["nodes"].extend(query_result["nodes"])
all_query_results["edges"].extend(query_result["edges"])
elif return_format == "triples":
all_query_results["triples"].extend(query_result["triples"])
else:
raise ValueError(f"Invalid return_format: {return_format}")
return all_query_results
def _query_with_fuzzy_match(self, keyword, kgdb_name='neo4j'):
def _query_with_fuzzy_match(self, keyword, kgdb_name="neo4j"):
"""模糊查询"""
assert self.driver is not None, "Database is not connected"
self.use_database(kgdb_name)
def query_fuzzy_match(tx, keyword):
result = tx.run("""
result = tx.run(
"""
MATCH (n:Entity)
WHERE n.name CONTAINS $keyword
RETURN DISTINCT n.name AS name
""", keyword=keyword)
""",
keyword=keyword,
)
values = result.values()
logger.debug(f"Fuzzy Query Results: {values}")
return values
@ -393,7 +414,7 @@ class GraphDatabase:
with self.driver.session() as session:
return session.execute_read(query_fuzzy_match, keyword)
def _query_with_vector_sim(self, keyword, kgdb_name='neo4j', threshold=0.9):
def _query_with_vector_sim(self, keyword, kgdb_name="neo4j", threshold=0.9):
"""向量查询"""
assert self.driver is not None, "Database is not connected"
self.use_database(kgdb_name)
@ -409,22 +430,26 @@ class GraphDatabase:
def query_by_vector(tx, text, threshold):
# 首先检查索引是否存在
if not _index_exists(tx, "entityEmbeddings"):
raise Exception("向量索引不存在，请先创建索引，或当前图谱中未上传任何三元组（知识库中自动构建的，不会在此处展示和检索）。")
raise Exception(
"向量索引不存在，请先创建索引，或当前图谱中未上传任何三元组（知识库中自动构建的，不会在此处展示和检索）。"
)
embedding = self.get_embedding(text)
result = tx.run("""
result = tx.run(
"""
CALL db.index.vector.queryNodes('entityEmbeddings', 10, $embedding)
YIELD node AS similarEntity, score
RETURN similarEntity.name AS name, score
""", embedding=embedding)
""",
embedding=embedding,
)
return [r for r in result if r["score"] > threshold]
with self.driver.session() as session:
results = session.execute_read(query_by_vector, keyword, threshold=threshold)
return results
def _query_specific_entity(self, entity_name, kgdb_name='neo4j', hops=2, limit=100):
def _query_specific_entity(self, entity_name, kgdb_name="neo4j", hops=2, limit=100):
"""查询指定实体三元组信息（无向关系）"""
assert self.driver is not None, "Database is not connected"
if not entity_name:
@ -467,12 +492,12 @@ class GraphDatabase:
logger.info(f"未找到实体 {entity_name} 的相关信息")
return {}
formatted_results = {'nodes': [], 'edges': [], 'triples': []}
formatted_results = {"nodes": [], "edges": [], "triples": []}
for item in results:
formatted_results['nodes'].extend([item['h'], item['t']])
formatted_results['edges'].append(item['r'])
formatted_results['triples'].append((item['h']['name'], item['r']['type'], item['t']['name']))
formatted_results["nodes"].extend([item["h"], item["t"]])
formatted_results["edges"].append(item["r"])
formatted_results["triples"].append((item["h"]["name"], item["r"]["type"], item["t"]["name"]))
logger.debug(f"Query Results: {results}")
return formatted_results
@ -506,20 +531,27 @@ class GraphDatabase:
return outputs
def set_embedding(self, tx, entity_name, embedding):
tx.run("""
tx.run(
"""
MATCH (e:Entity {name: $name})
CALL db.create.setNodeVectorProperty(e, 'embedding', $embedding)
""", name=entity_name, embedding=embedding)
""",
name=entity_name,
embedding=embedding,
)
def get_graph_info(self, graph_name="neo4j"):
assert self.driver is not None, "Database is not connected"
self.use_database(graph_name)
def query(tx):
# 只统计包含Entity标签的节点
entity_count = tx.run("MATCH (n:Entity) RETURN count(n) AS count").single()["count"]
# 只统计包含RELATION标签的关系
relationship_count = tx.run("MATCH ()-[r:RELATION]->() RETURN count(r) AS count").single()["count"]
triples_count = tx.run("MATCH (n:Entity)-[r:RELATION]->(m:Entity) RETURN count(n) AS count").single()["count"]
triples_count = tx.run("MATCH (n:Entity)-[r:RELATION]->(m:Entity) RETURN count(n) AS count").single()[
"count"
]
# 获取所有标签
labels = tx.run("CALL db.labels() YIELD label RETURN collect(label) AS labels").single()["labels"]
@ -532,7 +564,7 @@ class GraphDatabase:
"labels": labels,
"status": self.status,
"embed_model_name": self.embed_model_name,
"unindexed_node_count": self.query_nodes_without_embedding(graph_name)
"unindexed_node_count": self.query_nodes_without_embedding(graph_name),
}
try:
@ -543,6 +575,7 @@ class GraphDatabase:
# 添加时间戳
from datetime import datetime
graph_info["last_updated"] = datetime.now().isoformat()
return graph_info
else:
@ -565,7 +598,7 @@ class GraphDatabase:
return False
info_file_path = os.path.join(self.work_dir, "graph_info.json")
with open(info_file_path, 'w', encoding='utf-8') as f:
with open(info_file_path, "w", encoding="utf-8") as f:
json.dump(graph_info, f, ensure_ascii=False, indent=2)
# logger.info(f"图数据库信息已保存到：{info_file_path}")
@ -574,7 +607,7 @@ class GraphDatabase:
logger.error(f"保存图数据库信息失败：{e}")
return False
def query_nodes_without_embedding(self, kgdb_name='neo4j'):
def query_nodes_without_embedding(self, kgdb_name="neo4j"):
"""查询没有嵌入向量的节点
Returns:
@ -605,7 +638,7 @@ class GraphDatabase:
logger.debug(f"图数据库信息文件不存在：{info_file_path}")
return False
with open(info_file_path, encoding='utf-8') as f:
with open(info_file_path, encoding="utf-8") as f:
graph_info = json.load(f)
# 更新对象属性
@ -621,7 +654,7 @@ class GraphDatabase:
logger.error(f"加载图数据库信息失败：{e}")
return False
def add_embedding_to_nodes(self, node_names=None, kgdb_name='neo4j'):
def add_embedding_to_nodes(self, node_names=None, kgdb_name="neo4j"):
"""为节点添加嵌入向量
Args:
@ -655,13 +688,12 @@ class GraphDatabase:
edges = []
for item in results:
nodes.extend([item['h'], item['t']])
edges.append(item['r'])
nodes.extend([item["h"], item["t"]])
edges.append(item["r"])
formatted_results = {"nodes": nodes, "edges": edges}
return formatted_results
def _extract_relationship_info(self, relationship, source_name=None, target_name=None, node_dict=None):
"""
提取关系信息并返回格式化的节点和边信息
@ -701,12 +733,13 @@ class GraphDatabase:
return node_info, edge_info
def clean_triples_embedding(triples):
for item in triples:
if hasattr(item[0], '_properties'):
item[0]._properties['embedding'] = None
if hasattr(item[2], '_properties'):
item[2]._properties['embedding'] = None
if hasattr(item[0], "_properties"):
item[0]._properties["embedding"] = None
if hasattr(item[2], "_properties"):
item[2]._properties["embedding"] = None
return triples

90
src/knowledge/indexing.py
View File

@ -1,16 +1,17 @@
import os
import asyncio
import os
from pathlib import Path
from langchain.schema.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import (
TextLoader,
PyPDFLoader,
Docx2txtLoader,
UnstructuredMarkdownLoader,
UnstructuredHTMLLoader,
CSVLoader,
JSONLoader
Docx2txtLoader,
JSONLoader,
PyPDFLoader,
TextLoader,
UnstructuredHTMLLoader,
UnstructuredMarkdownLoader,
)
from src.utils import hashstr, logger
@ -31,22 +32,22 @@ def chunk_with_parser(file_path, params=None):
file_type = Path(file_path).suffix.lower()
# 选择合适的加载器
if file_type in ['.txt']:
if file_type in [".txt"]:
loader = TextLoader(file_path)
elif file_type in ['.md']:
elif file_type in [".md"]:
loader = UnstructuredMarkdownLoader(file_path)
elif file_type in ['.docx', '.doc']:
elif file_type in [".docx", ".doc"]:
loader = Docx2txtLoader(file_path)
elif file_type in ['.html', '.htm']:
elif file_type in [".html", ".htm"]:
loader = UnstructuredHTMLLoader(file_path)
elif file_type in ['.json']:
elif file_type in [".json"]:
loader = JSONLoader(file_path, jq_schema=".")
elif file_type in ['.csv']:
elif file_type in [".csv"]:
loader = CSVLoader(file_path)
else:
@ -73,6 +74,7 @@ def chunk_with_parser(file_path, params=None):
return nodes
def chunk_text(text, params=None):
"""
将文本切分成固定大小的块
@ -83,9 +85,7 @@ def chunk_text(text, params=None):
# 创建文本分割器
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
separators=["\n\n", "\n", ".", " ", ""]
chunk_size=chunk_size, chunk_overlap=chunk_overlap, separators=["\n\n", "\n", ".", " ", ""]
)
# 分割文档
@ -95,9 +95,11 @@ def chunk_text(text, params=None):
nodes = [{"text": node, "metadata": {"chunk_idx": i}} for i, node in enumerate(nodes)]
return nodes
def chunk(text_or_path, params=None):
raise NotImplementedError("chunk is deprecated, use chunk_with_parser or chunk_text instead")
def pdfreader(file_path, params=None):
"""读取PDF文件并返回text文本"""
if isinstance(file_path, str):
@ -114,6 +116,7 @@ def pdfreader(file_path, params=None):
text = "\n\n".join([d.page_content for d in docs])
return text
def plainreader(file_path):
"""读取普通文本文件并返回text文本"""
assert os.path.exists(file_path), "File not found"
@ -124,6 +127,7 @@ def plainreader(file_path):
text = "\n\n".join([d.page_content for d in docs])
return text
def parse_pdf(file, params=None):
"""
解析PDF文件，支持多种OCR方式
@ -149,14 +153,17 @@ def parse_pdf(file, params=None):
try:
if opt_ocr == "onnx_rapid_ocr":
from src.plugins import ocr
return ocr.process_pdf(file, params=params)
elif opt_ocr == "mineru_ocr":
from src.plugins import ocr
return ocr.process_file_mineru(file, params=params)
elif opt_ocr == "paddlex_ocr":
from src.plugins import ocr
return ocr.process_file_paddlex(file, params=params)
else:
@ -167,11 +174,7 @@ def parse_pdf(file, params=None):
raise
except Exception as e:
logger.error(f"PDF parsing failed: {str(e)}")
raise OCRServiceException(
f"PDF解析失败: {str(e)}",
opt_ocr,
"parsing_failed"
)
raise OCRServiceException(f"PDF解析失败: {str(e)}", opt_ocr, "parsing_failed")
def parse_image(file, params=None):
@ -190,14 +193,17 @@ def parse_image(file, params=None):
try:
if opt_ocr == "onnx_rapid_ocr":
from src.plugins import ocr
return ocr.process_image(file, params=params)
elif opt_ocr == "mineru_ocr":
from src.plugins import ocr
return ocr.process_file_mineru(file, params=params)
elif opt_ocr == "paddlex_ocr":
from src.plugins import ocr
return ocr.process_file_paddlex(file, params=params)
else:
@ -208,18 +214,17 @@ def parse_image(file, params=None):
raise
except Exception as e:
logger.error(f"Image parsing failed: {str(e)}")
raise OCRServiceException(
f"Image解析失败: {str(e)}",
opt_ocr,
"parsing_failed"
)
raise OCRServiceException(f"Image解析失败: {str(e)}", opt_ocr, "parsing_failed")
async def parse_pdf_async(file, params=None):
return await asyncio.to_thread(parse_pdf, file, params=params)
async def parse_image_async(file, params=None):
return await asyncio.to_thread(parse_image, file, params=params)
async def process_file_to_markdown(file_path: str, params: dict | None = None) -> str:
"""
将不同类型的文件转换为markdown格式
@ -234,40 +239,43 @@ async def process_file_to_markdown(file_path: str, params: dict | None = None) -
file_path_obj = Path(file_path)
file_ext = file_path_obj.suffix.lower()
if file_ext == '.pdf':
if file_ext == ".pdf":
# 使用 OCR 处理 PDF
text = await parse_pdf_async(str(file_path_obj), params=params)
return f"# {file_path_obj.name}\n\n{text}"
elif file_ext in ['.txt', '.md']:
elif file_ext in [".txt", ".md"]:
# 直接读取文本文件
with open(file_path_obj, encoding='utf-8') as f:
with open(file_path_obj, encoding="utf-8") as f:
content = f.read()
return f"# {file_path_obj.name}\n\n{content}"
elif file_ext in ['.doc', '.docx']:
elif file_ext in [".doc", ".docx"]:
# 处理 Word 文档
from docx import Document # type: ignore
doc = Document(file_path_obj)
text = '\n'.join([para.text for para in doc.paragraphs])
text = "\n".join([para.text for para in doc.paragraphs])
return f"# {file_path_obj.name}\n\n{text}"
elif file_ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']:
elif file_ext in [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"]:
# 使用 OCR 处理图片
text = await parse_image_async(str(file_path_obj), params=params)
return f"# {file_path_obj.name}\n\n{text}"
elif file_ext in ['.html', '.htm']:
elif file_ext in [".html", ".htm"]:
# 使用 BeautifulSoup 处理 HTML 文件
from markdownify import markdownify as md
with open(file_path_obj, encoding='utf-8') as f:
with open(file_path_obj, encoding="utf-8") as f:
content = f.read()
text = md(content, heading_style="ATX")
return f"# {file_path_obj.name}\n\n{text}"
elif file_ext == '.csv':
elif file_ext == ".csv":
# 处理 CSV 文件
import pandas as pd
df = pd.read_csv(file_path_obj)
# 将每一行数据与表头组合成独立的表格
markdown_content = f"# {file_path_obj.name}\n\n"
@ -280,9 +288,10 @@ async def process_file_to_markdown(file_path: str, params: dict | None = None) -
return markdown_content.strip()
elif file_ext in ['.xls', '.xlsx']:
elif file_ext in [".xls", ".xlsx"]:
# 处理 Excel 文件
import pandas as pd
# 读取所有工作表
excel_file = pd.ExcelFile(file_path_obj)
markdown_content = f"# {file_path_obj.name}\n\n"
@ -300,10 +309,11 @@ async def process_file_to_markdown(file_path: str, params: dict | None = None) -
return markdown_content.strip()
elif file_ext == '.json':
elif file_ext == ".json":
# 处理 JSON 文件
import json
with open(file_path_obj, encoding='utf-8') as f:
with open(file_path_obj, encoding="utf-8") as f:
data = json.load(f)
# 将 JSON 数据格式化为 markdown 代码块
json_str = json.dumps(data, ensure_ascii=False, indent=2)
@ -313,6 +323,7 @@ async def process_file_to_markdown(file_path: str, params: dict | None = None) -
# 尝试作为文本文件读取
raise ValueError(f"Unsupported file type: {file_ext}")
async def process_url_to_markdown(url: str, params: dict | None = None) -> str:
"""
将URL转换为markdown格式
@ -329,10 +340,9 @@ async def process_url_to_markdown(url: str, params: dict | None = None) -> str:
try:
response = requests.get(url, timeout=30)
soup = BeautifulSoup(response.content, 'html.parser')
soup = BeautifulSoup(response.content, "html.parser")
text_content = soup.get_text()
return f"# {url}\n\n{text_content}"
except Exception as e:
logger.error(f"Failed to process URL {url}: {e}")
return f"# {url}\n\nFailed to process URL: {e}"

11
src/knowledge/kb_factory.py
View File

@ -1,5 +1,6 @@
from typing import Any
from src.knowledge.knowledge_base import KnowledgeBase, KBNotFoundError
from src.knowledge.knowledge_base import KBNotFoundError, KnowledgeBase
from src.utils import logger
@ -13,8 +14,7 @@ class KnowledgeBaseFactory:
_default_configs: dict[str, dict] = {}
@classmethod
def register(cls, kb_type: str, kb_class: type[KnowledgeBase],
default_config: dict = None):
def register(cls, kb_type: str, kb_class: type[KnowledgeBase], default_config: dict = None):
"""
注册知识库类型
@ -48,8 +48,7 @@ class KnowledgeBaseFactory:
"""
if kb_type not in cls._kb_types:
available_types = list(cls._kb_types.keys())
raise KBNotFoundError(f"Unknown knowledge base type: {kb_type}. "
f"Available types: {available_types}")
raise KBNotFoundError(f"Unknown knowledge base type: {kb_type}. Available types: {available_types}")
kb_class = cls._kb_types[kb_type]
@ -79,7 +78,7 @@ class KnowledgeBaseFactory:
result[kb_type] = {
"class_name": kb_class.__name__,
"description": kb_class.__doc__ or "",
"default_config": cls._default_configs[kb_type]
"default_config": cls._default_configs[kb_type],
}
return result

41
src/knowledge/kb_manager.py
View File

@ -1,12 +1,12 @@
import os
import json
import time
import asyncio
from typing import Any
import json
import os
import time
from datetime import datetime
from typing import Any
from src.knowledge.knowledge_base import KnowledgeBase, KBNotFoundError, KBOperationError
from src.knowledge.kb_factory import KnowledgeBaseFactory
from src.knowledge.knowledge_base import KBNotFoundError, KBOperationError, KnowledgeBase
from src.utils import logger
@ -49,7 +49,7 @@ class KnowledgeBaseManager:
meta_file = os.path.join(self.work_dir, "global_metadata.json")
if os.path.exists(meta_file):
try:
with open(meta_file, encoding='utf-8') as f:
with open(meta_file, encoding="utf-8") as f:
data = json.load(f)
self.global_databases_meta = data.get("databases", {})
logger.info(f"Loaded global metadata for {len(self.global_databases_meta)} databases")
@ -59,12 +59,8 @@ class KnowledgeBaseManager:
def _save_global_metadata(self):
"""保存全局元数据"""
meta_file = os.path.join(self.work_dir, "global_metadata.json")
data = {
"databases": self.global_databases_meta,
"updated_at": datetime.now().isoformat(),
"version": "2.0"
}
with open(meta_file, 'w', encoding='utf-8') as f:
data = {"databases": self.global_databases_meta, "updated_at": datetime.now().isoformat(), "version": "2.0"}
with open(meta_file, "w", encoding="utf-8") as f:
json.dump(data, f, ensure_ascii=False, indent=2)
def _initialize_existing_kbs(self):
@ -140,7 +136,9 @@ class KnowledgeBaseManager:
return {"databases": all_databases}
async def create_database(self, database_name: str, description: str, kb_type: str, embed_info: dict | None = None, **kwargs) -> dict:
async def create_database(
self, database_name: str, description: str, kb_type: str, embed_info: dict | None = None, **kwargs
) -> dict:
"""
创建数据库
@ -160,8 +158,7 @@ class KnowledgeBaseManager:
kb_instance = self._get_or_create_kb_instance(kb_type)
db_info = kb_instance.create_database(database_name, description,
embed_info, **kwargs)
db_info = kb_instance.create_database(database_name, description, embed_info, **kwargs)
db_id = db_info["db_id"]
async with self._metadata_lock:
@ -170,7 +167,7 @@ class KnowledgeBaseManager:
"description": description,
"kb_type": kb_type,
"created_at": datetime.now().isoformat(),
"additional_params": kwargs.copy()
"additional_params": kwargs.copy(),
}
self._save_global_metadata()
@ -203,7 +200,7 @@ class KnowledgeBaseManager:
kb_instance = self._get_kb_for_database(db_id)
return await kb_instance.aquery(query_text, db_id, **kwargs)
async def export_data(self, db_id: str, format: str = 'zip', **kwargs) -> str:
async def export_data(self, db_id: str, format: str = "zip", **kwargs) -> str:
"""导出知识库数据"""
kb_instance = self._get_kb_for_database(db_id)
return await kb_instance.export_data(db_id, format=format, **kwargs)
@ -294,17 +291,13 @@ class KnowledgeBaseManager:
info[kb_type] = {
"work_dir": kb_instance.work_dir,
"database_count": len(kb_instance.databases_meta),
"file_count": len(kb_instance.files_meta)
"file_count": len(kb_instance.files_meta),
}
return info
def get_statistics(self) -> dict:
"""获取统计信息"""
stats = {
"total_databases": len(self.global_databases_meta),
"kb_types": {},
"total_files": 0
}
stats = {"total_databases": len(self.global_databases_meta), "kb_types": {}, "total_files": 0}
# 按知识库类型统计
for db_meta in self.global_databases_meta.values():
@ -352,7 +345,7 @@ class KnowledgeBaseManager:
kb_instance = self._get_kb_for_database(db_id)
# 如果不是 LightRagKB 实例，返回错误
if not hasattr(kb_instance, '_get_lightrag_instance'):
if not hasattr(kb_instance, "_get_lightrag_instance"):
logger.error(f"Knowledge base instance for {db_id} is not LightRagKB")
return None

80
src/knowledge/kb_utils.py
View File

@ -2,9 +2,11 @@ import os
import time
from pathlib import Path
from typing import Any
from langchain_text_splitters import MarkdownTextSplitter
from src.utils import hashstr, get_docker_safe_url, logger
from src import config
from src.utils import get_docker_safe_url, hashstr, logger
def split_text_into_chunks(text: str, file_id: str, filename: str, params: dict = {}) -> list[dict]:
@ -12,8 +14,8 @@ def split_text_into_chunks(text: str, file_id: str, filename: str, params: dict
将文本分割成块，使用 LangChain 的 MarkdownTextSplitter 进行智能分割
"""
chunks = []
chunk_size = params.get('chunk_size', 1000)
chunk_overlap = params.get('chunk_overlap', 200)
chunk_size = params.get("chunk_size", 1000)
chunk_overlap = params.get("chunk_overlap", 200)
# 使用 MarkdownTextSplitter 进行智能分割
# MarkdownTextSplitter 会尝试沿着 Markdown 格式的标题进行分割
@ -27,15 +29,17 @@ def split_text_into_chunks(text: str, file_id: str, filename: str, params: dict
# 转换为标准格式
for chunk_index, chunk_content in enumerate(text_chunks):
if chunk_content.strip(): # 跳过空块
chunks.append({
"id": f"{file_id}_chunk_{chunk_index}",
"content": chunk_content.strip(),
"file_id": file_id,
"filename": filename,
"chunk_index": chunk_index,
"source": filename,
"chunk_id": f"{file_id}_chunk_{chunk_index}"
})
chunks.append(
{
"id": f"{file_id}_chunk_{chunk_index}",
"content": chunk_content.strip(),
"file_id": file_id,
"filename": filename,
"chunk_index": chunk_index,
"source": filename,
"chunk_id": f"{file_id}_chunk_{chunk_index}",
}
)
logger.debug(f"Successfully split text into {len(chunks)} chunks using MarkdownTextSplitter")
return chunks
@ -64,15 +68,16 @@ def prepare_item_metadata(item: str, content_type: str, db_id: str) -> dict:
"file_type": file_type,
"status": "processing",
"created_at": time.time(),
"file_id": file_id
"file_id": file_id,
}
def split_text_into_qa_chunks(text: str, file_id: str, filename: str,
qa_separator: None | str = None, params: dict = {}) -> list[dict]:
def split_text_into_qa_chunks(
text: str, file_id: str, filename: str, qa_separator: None | str = None, params: dict = {}
) -> list[dict]:
"""
将文本按QA对分割成块，使用 LangChain 的 CharacterTextSplitter 进行分割"""
qa_separator = qa_separator or '\n\n'
qa_separator = qa_separator or "\n\n"
text_chunks = text.split(qa_separator)
# 转换为标准格式
@ -80,19 +85,23 @@ def split_text_into_qa_chunks(text: str, file_id: str, filename: str,
for chunk_index, chunk_content in enumerate(text_chunks):
if chunk_content.strip(): # 跳过空块
chunk_content = chunk_content.strip()[:4096]
chunks.append({
"id": f"{file_id}_qa_chunk_{chunk_index}",
"content": chunk_content.strip(),
"file_id": file_id,
"filename": filename,
"chunk_index": chunk_index,
"source": filename,
"chunk_id": f"{file_id}_qa_chunk_{chunk_index}",
"chunk_type": "qa" # 标识为QA类型的chunk
})
chunks.append(
{
"id": f"{file_id}_qa_chunk_{chunk_index}",
"content": chunk_content.strip(),
"file_id": file_id,
"filename": filename,
"chunk_index": chunk_index,
"source": filename,
"chunk_id": f"{file_id}_qa_chunk_{chunk_index}",
"chunk_type": "qa", # 标识为QA类型的chunk
}
)
logger.debug(f"QA chunks: {chunks[0]}")
logger.debug(f"Successfully split QA text into {len(chunks)} chunks using CharacterTextSplitter with `{qa_separator=}`")
logger.debug(
f"Successfully split QA text into {len(chunks)} chunks using CharacterTextSplitter with `{qa_separator=}`"
)
return chunks
@ -110,17 +119,18 @@ def get_embedding_config(embed_info: dict) -> dict:
try:
if embed_info:
config_dict['model'] = embed_info["name"]
config_dict['api_key'] = os.getenv(embed_info["api_key"], embed_info["api_key"])
config_dict['base_url'] = embed_info["base_url"]
config_dict['dimension'] = embed_info.get("dimension", 1024)
config_dict["model"] = embed_info["name"]
config_dict["api_key"] = os.getenv(embed_info["api_key"], embed_info["api_key"])
config_dict["base_url"] = embed_info["base_url"]
config_dict["dimension"] = embed_info.get("dimension", 1024)
else:
from src.models import select_embedding_model
default_model = select_embedding_model(config.embed_model)
config_dict['model'] = default_model.model
config_dict['api_key'] = default_model.api_key
config_dict['base_url'] = default_model.base_url
config_dict['dimension'] = getattr(default_model, 'dimension', 1024)
config_dict["model"] = default_model.model
config_dict["api_key"] = default_model.api_key
config_dict["base_url"] = default_model.base_url
config_dict["dimension"] = getattr(default_model, "dimension", 1024)
except Exception as e:
logger.error(f"Error in get_embedding_config: {e}, {embed_info}")

59
src/knowledge/knowledge_base.py
View File

@ -1,27 +1,30 @@
import os
import json
import os
import time
from abc import ABC, abstractmethod
from typing import Any
from collections.abc import AsyncGenerator
from pathlib import Path
from datetime import datetime
from pathlib import Path
from typing import Any
from src.utils import logger
class KnowledgeBaseException(Exception):
"""知识库统一异常基类"""
pass
class KBNotFoundError(KnowledgeBaseException):
"""知识库不存在错误"""
pass
class KBOperationError(KnowledgeBaseException):
"""知识库操作错误"""
pass
@ -84,8 +87,14 @@ class KnowledgeBase(ABC):
"""
pass
def create_database(self, database_name: str, description: str,
embed_info: dict | None = None, llm_info: dict | None = None, **kwargs) -> dict:
def create_database(
self,
database_name: str,
description: str,
embed_info: dict | None = None,
llm_info: dict | None = None,
**kwargs,
) -> dict:
"""
创建数据库
@ -110,7 +119,7 @@ class KnowledgeBase(ABC):
"embed_info": embed_info,
"llm_info": llm_info,
"metadata": kwargs,
"created_at": datetime.now().isoformat()
"created_at": datetime.now().isoformat(),
}
self._save_metadata()
@ -137,8 +146,7 @@ class KnowledgeBase(ABC):
"""
if db_id in self.databases_meta:
# 删除相关文件记录
files_to_delete = [fid for fid, finfo in self.files_meta.items()
if finfo.get("database_id") == db_id]
files_to_delete = [fid for fid, finfo in self.files_meta.items() if finfo.get("database_id") == db_id]
for file_id in files_to_delete:
del self.files_meta[file_id]
@ -150,6 +158,7 @@ class KnowledgeBase(ABC):
working_dir = os.path.join(self.work_dir, db_id)
if os.path.exists(working_dir):
import shutil
try:
shutil.rmtree(working_dir)
except Exception as e:
@ -158,8 +167,7 @@ class KnowledgeBase(ABC):
return {"message": "删除成功"}
@abstractmethod
async def add_content(self, db_id: str, items: list[str],
params: dict | None = None) -> list[dict]:
async def add_content(self, db_id: str, items: list[str], params: dict | None = None) -> list[dict]:
"""
添加内容（文件/URL）
@ -188,7 +196,7 @@ class KnowledgeBase(ABC):
"""
pass
async def export_data(self, db_id: str, format: str = 'zip', **kwargs) -> str:
async def export_data(self, db_id: str, format: str = "zip", **kwargs) -> str:
pass
def query(self, query_text: str, db_id: str, **kwargs) -> list[dict]:
@ -204,6 +212,7 @@ class KnowledgeBase(ABC):
一个包含字典的列表，每个字典代表一个检索到的文档块。
"""
import asyncio
logger.warning("query is deprecated, use aquery instead")
return asyncio.run(self.aquery(query_text, db_id, **kwargs))
@ -236,7 +245,7 @@ class KnowledgeBase(ABC):
"path": file_info.get("path", ""),
"type": file_info.get("file_type", ""),
"status": file_info.get("status", "done"),
"created_at": file_info.get("created_at", time.time())
"created_at": file_info.get("created_at", time.time()),
}
# 按创建时间倒序排序文件列表
@ -272,7 +281,7 @@ class KnowledgeBase(ABC):
"path": file_info.get("path", ""),
"type": file_info.get("file_type", ""),
"status": file_info.get("status", "done"),
"created_at": file_info.get("created_at", time.time())
"created_at": file_info.get("created_at", time.time()),
}
# 按创建时间倒序排序文件列表
@ -323,8 +332,6 @@ class KnowledgeBase(ABC):
with cls._processing_lock:
return file_id in cls._processing_files
def _check_and_fix_processing_status(self, db_id: str) -> None:
"""
检查并修复异常的processing状态
@ -338,14 +345,16 @@ class KnowledgeBase(ABC):
# 检查该数据库下所有processing状态的文件
for file_id, file_info in self.files_meta.items():
if (file_info.get("database_id") == db_id and
file_info.get("status") == "processing"):
if file_info.get("database_id") == db_id and file_info.get("status") == "processing":
# 检查文件是否真的在处理队列中
if not self._is_file_in_processing_queue(file_id):
logger.warning(f"File {file_id} has processing status but is not in processing queue, marking as error")
logger.warning(
f"File {file_id} has processing status but is not in processing queue, marking as error"
)
self.files_meta[file_id]["status"] = "error"
self.files_meta[file_id]["error"] = "Processing interrupted - file not found in processing queue"
self.files_meta[file_id]["error"] = (
"Processing interrupted - file not found in processing queue"
)
status_changed = True
# 如果有状态变更，保存元数据
@ -430,16 +439,18 @@ class KnowledgeBase(ABC):
"""
retrievers = {}
for db_id, meta in self.databases_meta.items():
def make_retriever(db_id):
async def retriever(query_text):
return await self.aquery(query_text, db_id)
return retriever
retrievers[db_id] = {
"name": meta["name"],
"description": meta["description"],
"retriever": make_retriever(db_id),
"metadata": meta
"metadata": meta,
}
return retrievers
@ -448,7 +459,7 @@ class KnowledgeBase(ABC):
meta_file = os.path.join(self.work_dir, f"metadata_{self.kb_type}.json")
if os.path.exists(meta_file):
try:
with open(meta_file, encoding='utf-8') as f:
with open(meta_file, encoding="utf-8") as f:
data = json.load(f)
self.databases_meta = data.get("databases", {})
self.files_meta = data.get("files", {})
@ -464,9 +475,9 @@ class KnowledgeBase(ABC):
"databases": self.databases_meta,
"files": self.files_meta,
"kb_type": self.kb_type,
"updated_at": datetime.now().isoformat()
"updated_at": datetime.now().isoformat(),
}
with open(meta_file, 'w', encoding='utf-8') as f:
with open(meta_file, "w", encoding="utf-8") as f:
json.dump(data, f, ensure_ascii=False, indent=2)
except Exception as e:
logger.error(f"Failed to save {self.kb_type} metadata: {e}")

68
src/knowledge/lightrag_kb.py
View File

@ -3,17 +3,16 @@ import traceback
from datetime import datetime
from lightrag import LightRAG, QueryParam
from lightrag.kg.shared_storage import initialize_pipeline_status
from lightrag.llm.openai import openai_complete_if_cache, openai_embed
from lightrag.utils import EmbeddingFunc, setup_logger
from lightrag.kg.shared_storage import initialize_pipeline_status
from pymilvus import connections, utility
from neo4j import GraphDatabase
from pymilvus import connections, utility
from src.knowledge.indexing import process_file_to_markdown, process_url_to_markdown
from src.knowledge.kb_utils import get_embedding_config, prepare_item_metadata
from src.knowledge.knowledge_base import KnowledgeBase
from src.knowledge.indexing import process_url_to_markdown, process_file_to_markdown
from src.knowledge.kb_utils import prepare_item_metadata, get_embedding_config
from src.utils import logger, hashstr
from src.utils import hashstr, logger
LIGHTRAG_LLM_PROVIDER = os.getenv("LIGHTRAG_LLM_PROVIDER", "siliconflow")
LIGHTRAG_LLM_NAME = os.getenv("LIGHTRAG_LLM_NAME", "zai-org/GLM-4.5-Air")
@ -38,8 +37,9 @@ class LightRagKB(KnowledgeBase):
# 设置 LightRAG 日志
log_dir = os.path.join(work_dir, "logs", "lightrag")
os.makedirs(log_dir, exist_ok=True)
setup_logger("lightrag", log_file_path=os.path.join(
log_dir, f"lightrag_{datetime.now().strftime('%Y-%m-%d')}.log"))
setup_logger(
"lightrag", log_file_path=os.path.join(log_dir, f"lightrag_{datetime.now().strftime('%Y-%m-%d')}.log")
)
logger.info("LightRagKB initialized")
@ -52,15 +52,11 @@ class LightRagKB(KnowledgeBase):
"""删除数据库，同时清除Milvus和Neo4j中的数据"""
# Drop Milvus collection
try:
milvus_uri = os.getenv('MILVUS_URI', 'http://localhost:19530')
milvus_token = os.getenv('MILVUS_TOKEN', '')
milvus_uri = os.getenv("MILVUS_URI", "http://localhost:19530")
milvus_token = os.getenv("MILVUS_TOKEN", "")
connection_alias = f"lightrag_{hashstr(db_id, 6)}"
connections.connect(
alias=connection_alias,
uri=milvus_uri,
token=milvus_token
)
connections.connect(alias=connection_alias, uri=milvus_uri, token=milvus_token)
# 删除 LightRAG 创建的三个集合
collection_names = [f"{db_id}_chunks", f"{db_id}_relationships", f"{db_id}_entities"]
@ -76,24 +72,28 @@ class LightRagKB(KnowledgeBase):
logger.error(f"Failed to drop Milvus collection {db_id}: {e}")
# Delete Neo4j data
neo4j_uri = os.getenv('NEO4J_URI', 'bolt://localhost:7687')
neo4j_username = os.getenv('NEO4J_USERNAME', 'neo4j')
neo4j_password = os.getenv('NEO4J_PASSWORD', '0123456789')
neo4j_uri = os.getenv("NEO4J_URI", "bolt://localhost:7687")
neo4j_username = os.getenv("NEO4J_USERNAME", "neo4j")
neo4j_password = os.getenv("NEO4J_PASSWORD", "0123456789")
try:
driver = GraphDatabase.driver(neo4j_uri, auth=(neo4j_username, neo4j_password))
with driver.session() as session:
# 删除带有特定 db_id 标签的节点和关系
session.run("""
MATCH (n:`""" + db_id + """`)
session.run(
"""
MATCH (n:`"""
+ db_id
+ """`)
DETACH DELETE n
""")
"""
)
logger.info(f"Deleted Neo4j nodes and relationships for workspace {db_id}")
except Exception as e:
logger.error(f"Failed to delete Neo4j data for {db_id}: {e}")
finally:
if 'driver' in locals():
if "driver" in locals():
driver.close()
# Delete local files and metadata
@ -194,6 +194,7 @@ class LightRagKB(KnowledgeBase):
base_url=model.base_url,
**kwargs,
)
return llm_model_func
def _get_embedding_func(self, embed_info: dict):
@ -211,8 +212,7 @@ class LightRagKB(KnowledgeBase):
),
)
async def add_content(self, db_id: str, items: list[str],
params: dict | None = None) -> list[dict]:
async def add_content(self, db_id: str, items: list[str], params: dict | None = None) -> list[dict]:
"""添加内容（文件/URL）"""
if db_id not in self.databases_meta:
raise ValueError(f"Database {db_id} not found")
@ -221,7 +221,7 @@ class LightRagKB(KnowledgeBase):
if not rag:
raise ValueError(f"Failed to get LightRAG instance for {db_id}")
content_type = params.get('content_type', 'file') if params else 'file'
content_type = params.get("content_type", "file") if params else "file"
processed_items_info = []
for item in items:
@ -240,24 +240,20 @@ class LightRagKB(KnowledgeBase):
# 根据内容类型处理内容
if content_type == "file":
markdown_content = await process_file_to_markdown(item, params=params)
markdown_content_lines = markdown_content[:100].replace('\n', ' ')
markdown_content_lines = markdown_content[:100].replace("\n", " ")
logger.info(f"Markdown content: {markdown_content_lines}...")
else: # URL
markdown_content = await process_url_to_markdown(item, params=params)
# 使用 LightRAG 插入内容
await rag.ainsert(
input=markdown_content,
ids=file_id,
file_paths=item_path
)
await rag.ainsert(input=markdown_content, ids=file_id, file_paths=item_path)
logger.info(f"Inserted {content_type} {item} into LightRAG. Done.")
# 更新状态为完成
self.files_meta[file_id]["status"] = "done"
self._save_metadata()
file_record['status'] = "done"
file_record["status"] = "done"
except Exception as e:
error_msg = str(e)
@ -265,8 +261,8 @@ class LightRagKB(KnowledgeBase):
self.files_meta[file_id]["status"] = "failed"
self.files_meta[file_id]["error"] = error_msg
self._save_metadata()
file_record['status'] = "failed"
file_record['error'] = error_msg
file_record["status"] = "failed"
file_record["error"] = error_msg
finally:
self._remove_from_processing_queue(file_id)
@ -324,7 +320,7 @@ class LightRagKB(KnowledgeBase):
if rag:
try:
# 获取文档的所有 chunks
assert hasattr(rag.text_chunks, 'get_all'), "text_chunks does not have get_all method"
assert hasattr(rag.text_chunks, "get_all"), "text_chunks does not have get_all method"
all_chunks = await rag.text_chunks.get_all() # type: ignore
# 筛选属于该文档的 chunks
@ -344,7 +340,7 @@ class LightRagKB(KnowledgeBase):
return {"lines": []}
async def export_data(self, db_id: str, format: str = 'csv', **kwargs) -> str:
async def export_data(self, db_id: str, format: str = "csv", **kwargs) -> str:
"""
使用 LightRAG 原生功能导出知识库数据。
[注意] 此功能当前已禁用。

103
src/knowledge/milvus_kb.py
View File

@ -1,24 +1,26 @@
import asyncio
import json
import os
import time
import traceback
import json
import asyncio
from pathlib import Path
from typing import Any
from datetime import datetime
from functools import partial
from pathlib import Path
from typing import Any
from pymilvus import (
connections, utility, Collection, CollectionSchema,
FieldSchema, DataType, db
)
from pymilvus import Collection, CollectionSchema, DataType, FieldSchema, connections, db, utility
from src import config
from src.models.embedding import OtherEmbedding
from src.knowledge.indexing import process_file_to_markdown, process_url_to_markdown
from src.knowledge.kb_utils import (
get_embedding_config,
prepare_item_metadata,
split_text_into_chunks,
split_text_into_qa_chunks,
)
from src.knowledge.knowledge_base import KnowledgeBase
from src.knowledge.indexing import process_url_to_markdown, process_file_to_markdown
from src.knowledge.kb_utils import split_text_into_chunks, split_text_into_qa_chunks, prepare_item_metadata, get_embedding_config
from src.utils import logger, hashstr
from src.models.embedding import OtherEmbedding
from src.utils import hashstr, logger
MILVUS_AVAILABLE = True
@ -42,9 +44,9 @@ class MilvusKB(KnowledgeBase):
# Milvus 配置
# self.milvus_host = kwargs.get('milvus_host', os.getenv('MILVUS_HOST', 'localhost'))
# self.milvus_port = kwargs.get('milvus_port', int(os.getenv('MILVUS_PORT', '19530')))
self.milvus_token = kwargs.get('milvus_token', os.getenv('MILVUS_TOKEN', ''))
self.milvus_uri = kwargs.get('milvus_uri', os.getenv('MILVUS_URI', 'http://localhost:19530'))
self.milvus_db = kwargs.get('milvus_db', 'yuxi_know')
self.milvus_token = kwargs.get("milvus_token", os.getenv("MILVUS_TOKEN", ""))
self.milvus_uri = kwargs.get("milvus_uri", os.getenv("MILVUS_URI", "http://localhost:19530"))
self.milvus_db = kwargs.get("milvus_db", "yuxi_know")
# 连接名称
self.connection_alias = f"milvus_{hashstr(work_dir, 6)}"
@ -53,8 +55,8 @@ class MilvusKB(KnowledgeBase):
self.collections: dict[str, Any] = {}
# 分块配置
self.chunk_size = kwargs.get('chunk_size', 1000)
self.chunk_overlap = kwargs.get('chunk_overlap', 200)
self.chunk_size = kwargs.get("chunk_size", 1000)
self.chunk_overlap = kwargs.get("chunk_overlap", 200)
# 元数据锁
self._metadata_lock = asyncio.Lock()
@ -73,11 +75,7 @@ class MilvusKB(KnowledgeBase):
"""初始化 Milvus 连接"""
try:
# 连接到 Milvus
connections.connect(
alias=self.connection_alias,
uri=self.milvus_uri,
token=self.milvus_token
)
connections.connect(alias=self.connection_alias, uri=self.milvus_uri, token=self.milvus_token)
# 创建数据库（如果不存在）
try:
@ -106,10 +104,7 @@ class MilvusKB(KnowledgeBase):
try:
# 检查集合是否存在
if utility.has_collection(collection_name, using=self.connection_alias):
collection = Collection(
name=collection_name,
using=self.connection_alias
)
collection = Collection(name=collection_name, using=self.connection_alias)
# 检查嵌入模型是否匹配
description = collection.description
@ -137,27 +132,18 @@ class MilvusKB(KnowledgeBase):
FieldSchema(name="chunk_id", dtype=DataType.VARCHAR, max_length=100),
FieldSchema(name="file_id", dtype=DataType.VARCHAR, max_length=100),
FieldSchema(name="chunk_index", dtype=DataType.INT64),
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=embedding_dim)
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=embedding_dim),
]
schema = CollectionSchema(
fields=fields,
description=f"Knowledge base collection for {db_id} using {model_name}"
fields=fields, description=f"Knowledge base collection for {db_id} using {model_name}"
)
# 创建集合
collection = Collection(
name=collection_name,
schema=schema,
using=self.connection_alias
)
collection = Collection(name=collection_name, schema=schema, using=self.connection_alias)
# 创建索引
index_params = {
"metric_type": "COSINE",
"index_type": "IVF_FLAT",
"params": {"nlist": 1024}
}
index_params = {"metric_type": "COSINE", "index_type": "IVF_FLAT", "params": {"nlist": 1024}}
collection.create_index("embedding", index_params)
logger.info(f"Created new Milvus collection: {collection_name}")
@ -218,18 +204,17 @@ class MilvusKB(KnowledgeBase):
def _split_text_into_chunks(self, text: str, file_id: str, filename: str, params: dict) -> list[dict]:
"""将文本分割成块"""
# 检查是否使用QA分割模式
use_qa_split = params.get('use_qa_split', False)
use_qa_split = params.get("use_qa_split", False)
if use_qa_split:
# 使用QA分割模式
qa_separator = params.get('qa_separator', '\n\n\n')
qa_separator = params.get("qa_separator", "\n\n\n")
return split_text_into_qa_chunks(text, file_id, filename, qa_separator, params)
else:
# 使用传统分割模式
return split_text_into_chunks(text, file_id, filename, params)
async def add_content(self, db_id: str, items: list[str],
params: dict | None = {}) -> list[dict]:
async def add_content(self, db_id: str, items: list[str], params: dict | None = {}) -> list[dict]:
"""添加内容（文件/URL）"""
if db_id not in self.databases_meta:
raise ValueError(f"Database {db_id} not found")
@ -241,7 +226,7 @@ class MilvusKB(KnowledgeBase):
embed_info = self.databases_meta[db_id].get("embed_info", {})
embedding_function = self._get_async_embedding_function(embed_info)
content_type = params.get('content_type', 'file') if params else 'file'
content_type = params.get("content_type", "file") if params else "file"
processed_items_info = []
for item in items:
@ -280,7 +265,7 @@ class MilvusKB(KnowledgeBase):
[chunk["chunk_id"] for chunk in chunks],
[chunk["file_id"] for chunk in chunks],
[chunk["chunk_index"] for chunk in chunks],
embeddings
embeddings,
]
def _insert_and_flush():
@ -294,7 +279,7 @@ class MilvusKB(KnowledgeBase):
async with self._metadata_lock:
self.files_meta[file_id]["status"] = "done"
self._save_metadata()
file_record['status'] = "done"
file_record["status"] = "done"
# 从处理队列中移除
self._remove_from_processing_queue(file_id)
@ -303,7 +288,7 @@ class MilvusKB(KnowledgeBase):
async with self._metadata_lock:
self.files_meta[file_id]["status"] = "failed"
self._save_metadata()
file_record['status'] = "failed"
file_record["status"] = "failed"
# 从处理队列中移除
self._remove_from_processing_queue(file_id)
finally:
@ -334,7 +319,7 @@ class MilvusKB(KnowledgeBase):
anns_field="embedding",
param=search_params,
limit=top_k,
output_fields=["content", "source", "chunk_id", "file_id", "chunk_index"]
output_fields=["content", "source", "chunk_id", "file_id", "chunk_index"],
)
if not results or len(results) == 0 or len(results[0]) == 0:
@ -352,14 +337,12 @@ class MilvusKB(KnowledgeBase):
"source": entity.get("source", "未知来源"),
"chunk_id": entity.get("chunk_id"),
"file_id": entity.get("file_id"),
"chunk_index": entity.get("chunk_index")
"chunk_index": entity.get("chunk_index"),
}
retrieved_chunks.append({
"content": entity.get("content", ""),
"metadata": metadata,
"score": similarity
})
retrieved_chunks.append(
{"content": entity.get("content", ""), "metadata": metadata, "score": similarity}
)
logger.debug(f"Milvus query response: {len(retrieved_chunks)} chunks found (after similarity filtering)")
return retrieved_chunks
@ -376,11 +359,7 @@ class MilvusKB(KnowledgeBase):
# 先查询文件是否存在，避免不必要的删除操作
try:
expr = f'file_id == "{file_id}"'
results = collection.query(
expr=expr,
output_fields=["id"],
limit=1
)
results = collection.query(expr=expr, output_fields=["id"], limit=1)
if not results:
logger.info(f"File {file_id} not found in Milvus, skipping delete operation")
@ -417,7 +396,7 @@ class MilvusKB(KnowledgeBase):
results = collection.query(
expr=expr,
output_fields=["content", "chunk_id", "chunk_index"],
limit=10000 # 假设单个文件不会超过10000个chunks
limit=10000, # 假设单个文件不会超过10000个chunks
)
# 构建chunks数据
@ -426,7 +405,7 @@ class MilvusKB(KnowledgeBase):
chunk_data = {
"id": result.get("chunk_id", ""),
"content": result.get("content", ""),
"chunk_order_index": result.get("chunk_index", 0)
"chunk_order_index": result.get("chunk_index", 0),
}
doc_chunks.append(chunk_data)
@ -457,7 +436,7 @@ class MilvusKB(KnowledgeBase):
def __del__(self):
"""清理连接"""
try:
if hasattr(self, 'connection_alias'):
if hasattr(self, "connection_alias"):
connections.disconnect(self.connection_alias)
except Exception: # noqa: S110
pass

7
src/models/__init__.py
View File

@ -2,9 +2,10 @@ import os
import traceback
from src import config
from src.utils.logging_config import logger
from src.models.chat_model import OpenAIBase
from src.models.embedding import OllamaEmbedding, OtherEmbedding
from src.utils.logging_config import logger
def select_model(model_provider, model_name=None):
"""根据模型提供者选择模型"""
@ -16,12 +17,14 @@ def select_model(model_provider, model_name=None):
if model_provider == "openai":
from src.models.chat_model import OpenModel
return OpenModel(model_name)
if model_provider == "custom":
model_info = get_custom_model(model_name)
from src.models.chat_model import CustomModel
return CustomModel(model_info)
# 其他模型，默认使用OpenAIBase
@ -37,7 +40,7 @@ def select_model(model_provider, model_name=None):
def select_embedding_model(model_id):
provider, model_name = model_id.split('/', 1) if model_id else ("", "")
provider, model_name = model_id.split("/", 1) if model_id else ("", "")
support_embed_models = config.embed_model_names.keys()
assert model_id in support_embed_models, f"Unsupported embed model: {model_id}, only support {support_embed_models}"
logger.debug(f"Loading embedding model {model_id}")

21
src/models/chat_model.py
View File

@ -1,8 +1,11 @@
import os
import requests
from openai import OpenAI
from src.utils import logger, get_docker_safe_url
from langchain_openai import ChatOpenAI
from openai import OpenAI
from src.utils import get_docker_safe_url, logger
class OpenAIBase:
def __init__(self, api_key, base_url, model_name, **kwargs):
@ -14,7 +17,7 @@ class OpenAIBase:
def predict(self, message, stream=False):
if isinstance(message, str):
messages=[{"role": "user", "content": message}]
messages = [{"role": "user", "content": message}]
else:
messages = message
@ -31,8 +34,8 @@ class OpenAIBase:
stream=True,
)
for chunk in response:
if len(chunk.choices) > 0:
yield chunk.choices[0].delta
if len(chunk.choices) > 0:
yield chunk.choices[0].delta
except Exception as e:
err = f"Error streaming response: {e}, URL: {self.base_url}, API Key: {self.api_key[:5]}***, Model: {self.model_name}"
@ -49,11 +52,7 @@ class OpenAIBase:
def get_models(self):
try:
return self.client.models.list(
extra_query={
"type": "text"
}
)
return self.client.models.list(extra_query={"type": "text"})
except Exception as e:
logger.error(f"Error getting models: {e}")
return []
@ -67,7 +66,6 @@ class OpenModel(OpenAIBase):
super().__init__(api_key=api_key, base_url=base_url, model_name=model_name)
class CustomModel(OpenAIBase):
def __init__(self, model_info):
model_name = model_info["name"]
@ -83,5 +81,6 @@ class GeneralResponse:
self.content = content
self.is_full = False
if __name__ == "__main__":
pass

41
src/models/embedding.py
View File

@ -1,16 +1,16 @@
import os
import asyncio
import json
import os
from abc import ABC, abstractmethod
import httpx
import requests
import asyncio
from abc import abstractmethod, ABC
from src import config
from src.utils import hashstr, logger, get_docker_safe_url
from src.utils import get_docker_safe_url, hashstr, logger
class BaseEmbeddingModel(ABC):
def __init__(self, model=None, name=None, dimension=None, url=None, base_url=None, api_key=None):
"""
Args:
@ -60,22 +60,18 @@ class BaseEmbeddingModel(ABC):
task_id = None
if len(messages) > batch_size:
task_id = hashstr(messages)
self.embed_state[task_id] = {
'status': 'in-progress',
'total': len(messages),
'progress': 0
}
self.embed_state[task_id] = {"status": "in-progress", "total": len(messages), "progress": 0}
for i in range(0, len(messages), batch_size):
group_msg = messages[i:i+batch_size]
group_msg = messages[i : i + batch_size]
logger.info(f"Encoding [{i}/{len(messages)}] messages (bsz={batch_size})")
response = self.encode(group_msg)
data.extend(response)
if task_id:
self.embed_state[task_id]['progress'] = i + len(group_msg)
self.embed_state[task_id]["progress"] = i + len(group_msg)
if task_id:
self.embed_state[task_id]['status'] = 'completed'
self.embed_state[task_id]["status"] = "completed"
return data
@ -84,15 +80,11 @@ class BaseEmbeddingModel(ABC):
task_id = None
if len(messages) > batch_size:
task_id = hashstr(messages)
self.embed_state[task_id] = {
'status': 'in-progress',
'total': len(messages),
'progress': 0
}
self.embed_state[task_id] = {"status": "in-progress", "total": len(messages), "progress": 0}
tasks = []
for i in range(0, len(messages), batch_size):
group_msg = messages[i:i+batch_size]
group_msg = messages[i : i + batch_size]
tasks.append(self.aencode(group_msg))
results = await asyncio.gather(*tasks)
@ -100,11 +92,12 @@ class BaseEmbeddingModel(ABC):
data.extend(res)
if task_id:
self.embed_state[task_id]['progress'] = len(messages)
self.embed_state[task_id]['status'] = 'completed'
self.embed_state[task_id]["progress"] = len(messages)
self.embed_state[task_id]["status"] = "completed"
return data
class OllamaEmbedding(BaseEmbeddingModel):
"""
Ollama Embedding Model
@ -149,13 +142,9 @@ class OllamaEmbedding(BaseEmbeddingModel):
class OtherEmbedding(BaseEmbeddingModel):
def __init__(self, **kwargs) -> None:
super().__init__(**kwargs)
self.headers = {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json"
}
self.headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
def build_payload(self, message: list[str] | str) -> dict:
return {"model": self.model, "input": message}

26
src/models/rerank_model.py
View File

@ -1,25 +1,25 @@
import os
import json
import requests
import os
import numpy as np
import requests
from src import config
from src.utils import logger, get_docker_safe_url
from src.utils import get_docker_safe_url, logger
def sigmoid(x):
return 1 / (1 + np.exp(-x))
class OnlineReranker:
def __init__(self, model_name, api_key, base_url, **kwargs):
self.url = get_docker_safe_url(base_url)
self.model = model_name
self.api_key = api_key
self.headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json"
}
self.headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
def compute_score(self, sentence_pairs, batch_size = 256, max_length = 512, normalize = False):
def compute_score(self, sentence_pairs, batch_size=256, max_length=512, normalize=False):
# TODO 还没实现 batch_size
query, sentences = sentence_pairs[0], sentence_pairs[1]
payload = self.build_payload(query, sentences, max_length)
@ -35,7 +35,7 @@ class OnlineReranker:
return all_scores
def build_payload(self, query, sentences, max_length = 512):
def build_payload(self, query, sentences, max_length=512):
return {
"model": self.model,
"query": query,
@ -43,6 +43,7 @@ class OnlineReranker:
"max_chunks_per_doc": max_length,
}
def get_reranker(model_id, **kwargs):
support_rerankers = config.reranker_names.keys()
assert model_id in support_rerankers, f"Unsupported Reranker: {model_id}, only support {support_rerankers}"
@ -51,9 +52,4 @@ def get_reranker(model_id, **kwargs):
base_url = model_info["base_url"]
api_key = os.getenv(model_info["api_key"], model_info["api_key"])
assert api_key, f"{model_info['name']} api_key is required"
return OnlineReranker(
model_name=model_info["name"],
api_key=api_key,
base_url=base_url,
**kwargs
)
return OnlineReranker(model_name=model_info["name"], api_key=api_key, base_url=base_url, **kwargs)

105
src/plugins/_ocr.py
View File

@ -1,27 +1,22 @@
import os
import uuid
import time
from pathlib import Path
import uuid
from argparse import ArgumentParser
from collections import defaultdict
from pathlib import Path
import fitz # fitz就是pip install PyMuPDF
import numpy as np # Added import for numpy
from PIL import Image
from tqdm import tqdm
from rapidocr_onnxruntime import RapidOCR
from tqdm import tqdm
from src.utils import logger, is_text_pdf
from src.utils import is_text_pdf, logger
GOLBAL_STATE = {}
# OCR服务监控统计
OCR_STATS = {
"requests": defaultdict(int),
"failures": defaultdict(int),
"service_status": defaultdict(str)
}
OCR_STATS = {"requests": defaultdict(int), "failures": defaultdict(int), "service_status": defaultdict(str)}
def log_ocr_request(service_name: str, file_path: str, success: bool, processing_time: float, error_msg: str = None):
@ -50,7 +45,7 @@ def get_ocr_stats():
"success_count": success_count,
"failure_count": OCR_STATS["failures"][service],
"success_rate": f"{success_rate:.2%}",
"status": OCR_STATS["service_status"][service]
"status": OCR_STATS["service_status"][service],
}
return stats
@ -58,6 +53,7 @@ def get_ocr_stats():
class OCRServiceException(Exception):
"""OCR服务异常"""
def __init__(self, message, service_name=None, status_code=None):
super().__init__(message)
self.service_name = service_name
@ -69,8 +65,10 @@ class OCRPlugin:
def __init__(self, **kwargs):
self.ocr = None
self.det_box_thresh = kwargs.get('det_box_thresh', 0.3)
self.model_dir_root = os.getenv("MODEL_DIR") if not os.getenv("RUNNING_IN_DOCKER") else os.getenv("MODEL_DIR_IN_DOCKER")
self.det_box_thresh = kwargs.get("det_box_thresh", 0.3)
self.model_dir_root = (
os.getenv("MODEL_DIR") if not os.getenv("RUNNING_IN_DOCKER") else os.getenv("MODEL_DIR_IN_DOCKER")
)
def _check_rapid_ocr_availability(self):
"""检查RapidOCR模型是否可用"""
@ -81,16 +79,14 @@ class OCRPlugin:
if not os.path.exists(model_dir):
raise OCRServiceException(
f"模型目录不存在: {model_dir}。请下载 SWHL/RapidOCR 模型",
"rapid_ocr",
"model_not_found"
f"模型目录不存在: {model_dir}。请下载 SWHL/RapidOCR 模型", "rapid_ocr", "model_not_found"
)
if not os.path.exists(det_model_dir) or not os.path.exists(rec_model_dir):
raise OCRServiceException(
f"模型文件缺失。请确认模型文件完整: {det_model_dir}, {rec_model_dir}",
"rapid_ocr",
"model_incomplete"
"model_incomplete",
)
return True
@ -99,11 +95,7 @@ class OCRPlugin:
if isinstance(e, OCRServiceException):
raise
else:
raise OCRServiceException(
f"RapidOCR模型检查失败: {str(e)}",
"rapid_ocr",
"check_failed"
)
raise OCRServiceException(f"RapidOCR模型检查失败: {str(e)}", "rapid_ocr", "check_failed")
def load_model(self):
"""加载 OCR 模型"""
@ -120,11 +112,7 @@ class OCRPlugin:
self.ocr = RapidOCR(det_box_thresh=0.3, det_model_path=det_model_dir, rec_model_path=rec_model_dir)
logger.info(f"OCR Plugin for det_box_thresh = {self.det_box_thresh} loaded.")
except Exception as e:
raise OCRServiceException(
f"RapidOCR模型加载失败: {str(e)}",
"rapid_ocr",
"load_failed"
)
raise OCRServiceException(f"RapidOCR模型加载失败: {str(e)}", "rapid_ocr", "load_failed")
def process_image(self, image, params=None):
"""
@ -165,7 +153,7 @@ class OCRPlugin:
# 提取文本
if result:
text = '\n'.join([line[1] for line in result])
text = "\n".join([line[1] for line in result])
log_ocr_request("rapid_ocr", image_path, True, processing_time)
return text
else:
@ -189,11 +177,11 @@ class OCRPlugin:
str: 临时文件路径
"""
# 为临时文件创建目录（如果不存在）
tmp_dir = os.path.join(os.getcwd(), 'tmp')
tmp_dir = os.path.join(os.getcwd(), "tmp")
os.makedirs(tmp_dir, exist_ok=True)
# 生成临时文件路径
temp_filename = f'ocr_temp_{uuid.uuid4().hex[:8]}.png'
temp_filename = f"ocr_temp_{uuid.uuid4().hex[:8]}.png"
image_path = os.path.join(tmp_dir, temp_filename)
# 根据图像类型保存文件
@ -224,7 +212,7 @@ class OCRPlugin:
pdfDoc = fitz.open(pdf_path)
totalPage = pdfDoc.page_count
for pg in tqdm(range(totalPage), desc='to images', ncols=100):
for pg in tqdm(range(totalPage), desc="to images", ncols=100):
page = pdfDoc[pg]
rotate, zoom_x, zoom_y = 0, 2, 2
mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
@ -234,12 +222,12 @@ class OCRPlugin:
# 处理每个图像并合并文本
all_text = []
for img_path in tqdm(images, desc='to txt', ncols=100):
for img_path in tqdm(images, desc="to txt", ncols=100):
text = self.process_image(img_path)
all_text.append(text)
logger.debug(f"PDF OCR result: {all_text[:50]}(...) total {len(all_text)} pages.")
return '\n\n'.join(all_text)
return "\n\n".join(all_text)
except Exception as e:
logger.error(f"PDF processing error: {str(e)}")
@ -253,6 +241,7 @@ class OCRPlugin:
:return: 提取的文本
"""
import requests
from .mineru import parse_doc
mineru_ocr_uri = os.getenv("MINERU_OCR_URI", "http://localhost:30000")
@ -269,28 +258,20 @@ class OCRPlugin:
error_detail = health_check_response.text
raise OCRServiceException(
f"MinerU OCR服务健康检查失败: {error_detail}",
"mineru_ocr",
"health_check_failed"
f"MinerU OCR服务健康检查失败: {error_detail}", "mineru_ocr", "health_check_failed"
)
except Exception as e:
if isinstance(e, OCRServiceException):
raise
raise OCRServiceException(
f"MinerU OCR服务检查失败: {str(e)}",
"mineru_ocr",
"service_error"
)
raise OCRServiceException(f"MinerU OCR服务检查失败: {str(e)}", "mineru_ocr", "service_error")
try:
start_time = time.time()
file_path_list = [file_path]
output_dir = os.path.join(os.getcwd(), "tmp", "mineru_ocr")
text = parse_doc(file_path_list, output_dir,
backend="vlm-sglang-client",
server_url=mineru_ocr_uri)[0]
text = parse_doc(file_path_list, output_dir, backend="vlm-sglang-client", server_url=mineru_ocr_uri)[0]
processing_time = time.time() - start_time
log_ocr_request("mineru_ocr", file_path, True, processing_time)
@ -303,11 +284,7 @@ class OCRPlugin:
error_msg = f"MinerU OCR处理失败: {str(e)}"
log_ocr_request("mineru_ocr", file_path, False, processing_time, error_msg)
raise OCRServiceException(
error_msg,
"mineru_ocr",
"processing_failed"
)
raise OCRServiceException(error_msg, "mineru_ocr", "processing_failed")
def process_file_paddlex(self, pdf_path, params=None):
"""
@ -331,18 +308,12 @@ class OCRPlugin:
error_detail = health_check_response.text
raise OCRServiceException(
f"PaddleX OCR服务健康检查失败: {error_detail}",
"paddlex_ocr",
"health_check_failed"
f"PaddleX OCR服务健康检查失败: {error_detail}", "paddlex_ocr", "health_check_failed"
)
except Exception as e:
if isinstance(e, OCRServiceException):
raise
raise OCRServiceException(
f"PaddleX OCR服务检查失败: {str(e)}",
"paddlex_ocr",
"service_error"
)
raise OCRServiceException(f"PaddleX OCR服务检查失败: {str(e)}", "paddlex_ocr", "service_error")
try:
start_time = time.time()
@ -353,11 +324,7 @@ class OCRPlugin:
error_msg = f"PaddleX OCR处理失败: {result['error']}"
log_ocr_request("paddlex_ocr", pdf_path, False, processing_time, error_msg)
raise OCRServiceException(
error_msg,
"paddlex_ocr",
"processing_failed"
)
raise OCRServiceException(error_msg, "paddlex_ocr", "processing_failed")
log_ocr_request("paddlex_ocr", pdf_path, True, processing_time)
return result["full_text"]
@ -365,19 +332,17 @@ class OCRPlugin:
except Exception as e:
if isinstance(e, OCRServiceException):
raise
processing_time = time.time() - start_time if 'start_time' in locals() else 0
processing_time = time.time() - start_time if "start_time" in locals() else 0
error_msg = f"PaddleX OCR处理失败: {str(e)}"
log_ocr_request("paddlex_ocr", pdf_path, False, processing_time, error_msg)
raise OCRServiceException(
error_msg,
"paddlex_ocr",
"processing_failed"
)
raise OCRServiceException(error_msg, "paddlex_ocr", "processing_failed")
def get_state(task_id):
return GOLBAL_STATE.get(task_id, {})
def plainreader(file_path):
"""读取普通文本文件并返回text文本"""
assert os.path.exists(file_path), "File not found"
@ -389,8 +354,8 @@ def plainreader(file_path):
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument('--pdf-path', type=str, required=True, help='Path to the PDF file')
parser.add_argument('--return-text', action='store_true', help='Return the extracted text')
parser.add_argument("--pdf-path", type=str, required=True, help="Path to the PDF file")
parser.add_argument("--return-text", action="store_true", help="Return the extracted text")
args = parser.parse_args()
ocr = OCRPlugin()

94
src/plugins/mineru.py
View File

@ -3,17 +3,17 @@ import copy
import json
import os
from pathlib import Path
from tqdm import tqdm
from mineru.backend.pipeline.model_json_to_middle_json import result_to_middle_json as pipeline_result_to_middle_json
from mineru.backend.pipeline.pipeline_analyze import doc_analyze as pipeline_doc_analyze
from mineru.backend.pipeline.pipeline_middle_json_mkcontent import union_make as pipeline_union_make
from mineru.backend.vlm.vlm_analyze import doc_analyze as vlm_doc_analyze
from mineru.backend.vlm.vlm_middle_json_mkcontent import union_make as vlm_union_make
from mineru.cli.common import convert_pdf_bytes_to_bytes_by_pypdfium2, prepare_env, read_fn
from mineru.data.data_reader_writer import FileBasedDataWriter
from mineru.utils.draw_bbox import draw_layout_bbox, draw_span_bbox
from mineru.utils.enum_class import MakeMode
from mineru.backend.vlm.vlm_analyze import doc_analyze as vlm_doc_analyze
from mineru.backend.pipeline.pipeline_analyze import doc_analyze as pipeline_doc_analyze
from mineru.backend.pipeline.pipeline_middle_json_mkcontent import union_make as pipeline_union_make
from mineru.backend.pipeline.model_json_to_middle_json import result_to_middle_json as pipeline_result_to_middle_json
from mineru.backend.vlm.vlm_middle_json_mkcontent import union_make as vlm_union_make
from tqdm import tqdm
from src.utils.logging_config import logger
@ -39,17 +39,20 @@ def do_parse(
start_page_id=0, # Start page ID for parsing, default is 0
end_page_id=None, # End page ID for parsing, default is None (parse all pages until the end of the document)
) -> list[str]:
if backend == "pipeline":
for idx, pdf_bytes in enumerate(pdf_bytes_list):
new_pdf_bytes = convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes, start_page_id, end_page_id)
pdf_bytes_list[idx] = new_pdf_bytes
result = pipeline_doc_analyze(pdf_bytes_list, p_lang_list, parse_method=parse_method, formula_enable=p_formula_enable,table_enable=p_table_enable)
result = pipeline_doc_analyze(
pdf_bytes_list,
p_lang_list,
parse_method=parse_method,
formula_enable=p_formula_enable,
table_enable=p_table_enable,
)
infer_results, all_image_lists, all_pdf_docs, lang_list, ocr_enabled_list = result
md_results = []
for idx, model_list in enumerate(infer_results):
model_json = copy.deepcopy(model_list)
@ -61,7 +64,9 @@ def do_parse(
pdf_doc = all_pdf_docs[idx]
_lang = lang_list[idx]
_ocr_enable = ocr_enabled_list[idx]
middle_json = pipeline_result_to_middle_json(model_list, images_list, pdf_doc, image_writer, _lang, _ocr_enable, p_formula_enable)
middle_json = pipeline_result_to_middle_json(
model_list, images_list, pdf_doc, image_writer, _lang, _ocr_enable, p_formula_enable
)
pdf_info = middle_json["pdf_info"]
@ -124,7 +129,9 @@ def do_parse(
pdf_bytes = convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes, start_page_id, end_page_id)
local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name, parse_method)
image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(local_md_dir)
middle_json, infer_result = vlm_doc_analyze(pdf_bytes, image_writer=image_writer, backend=backend, server_url=server_url)
middle_json, infer_result = vlm_doc_analyze(
pdf_bytes, image_writer=image_writer, backend=backend, server_url=server_url
)
pdf_info = middle_json["pdf_info"]
@ -176,35 +183,35 @@ def do_parse(
def parse_doc(
path_list: list[Path],
output_dir,
lang="ch",
backend="pipeline",
method="auto",
server_url=None,
start_page_id=0, # Start page ID for parsing, default is 0
end_page_id=None # End page ID for parsing, default is None (parse all pages until the end of the document)
path_list: list[Path],
output_dir,
lang="ch",
backend="pipeline",
method="auto",
server_url=None,
start_page_id=0, # Start page ID for parsing, default is 0
end_page_id=None, # End page ID for parsing, default is None (parse all pages until the end of the document)
) -> list[str]:
"""
Parameter description:
path_list: List of document paths to be parsed, can be PDF or image files.
output_dir: Output directory for storing parsing results.
lang: Language option, default is 'ch', optional values include['ch', 'ch_server', 'ch_lite', 'en', 'korean', 'japan', 'chinese_cht', 'ta', 'te', 'ka']。
Input the languages in the pdf (if known) to improve OCR accuracy. Optional.
Adapted only for the case where the backend is set to "pipeline"
backend: the backend for parsing pdf:
pipeline: More general.
vlm-transformers: More general.
vlm-sglang-engine: Faster(engine).
vlm-sglang-client: Faster(client).
without method specified, pipeline will be used by default.
method: the method for parsing pdf:
auto: Automatically determine the method based on the file type.
txt: Use text extraction method.
ocr: Use OCR method for image-based PDFs.
Without method specified, 'auto' will be used by default.
Adapted only for the case where the backend is set to "pipeline".
server_url: When the backend is `sglang-client`, you need to specify the server_url, for example:`http://127.0.0.1:30000`
Parameter description:
path_list: List of document paths to be parsed, can be PDF or image files.
output_dir: Output directory for storing parsing results.
lang: Language option, default is 'ch', optional values include['ch', 'ch_server', 'ch_lite', 'en', 'korean', 'japan', 'chinese_cht', 'ta', 'te', 'ka']。
Input the languages in the pdf (if known) to improve OCR accuracy. Optional.
Adapted only for the case where the backend is set to "pipeline"
backend: the backend for parsing pdf:
pipeline: More general.
vlm-transformers: More general.
vlm-sglang-engine: Faster(engine).
vlm-sglang-client: Faster(client).
without method specified, pipeline will be used by default.
method: the method for parsing pdf:
auto: Automatically determine the method based on the file type.
txt: Use text extraction method.
ocr: Use OCR method for image-based PDFs.
Without method specified, 'auto' will be used by default.
Adapted only for the case where the backend is set to "pipeline".
server_url: When the backend is `sglang-client`, you need to specify the server_url, for example:`http://127.0.0.1:30000`
"""
file_name_list = []
pdf_bytes_list = []
@ -225,21 +232,22 @@ def parse_doc(
parse_method=method,
server_url=server_url,
start_page_id=start_page_id,
end_page_id=end_page_id
end_page_id=end_page_id,
)
return result if result else [""]
if __name__ == "__main__":
pdf_files_dir = "/home/zwj/workspace/projects/Yuxi-Know/test/struct_pdf"
output_dir = "/home/zwj/workspace/projects/Yuxi-Know/test/struct_pdf_output"
pdf_suffixes = [".pdf"]
image_suffixes = [".png", ".jpeg", ".jpg"]
doc_path_list = []
for doc_path in Path(pdf_files_dir).glob('*'):
for doc_path in Path(pdf_files_dir).glob("*"):
if doc_path.suffix in pdf_suffixes + image_suffixes:
doc_path_list.append(doc_path)
parse_doc(doc_path_list, output_dir, backend="vlm-sglang-client", server_url="http://172.19.13.5:30000") # faster(client).
parse_doc(
doc_path_list, output_dir, backend="vlm-sglang-client", server_url="http://172.19.13.5:30000"
) # faster(client).

106
src/plugins/paddlex.py
View File

@ -1,30 +1,29 @@
import requests
import json
import base64
import json
import os
import time
from typing import Optional, Any
from pathlib import Path
from typing import Any, Optional
import requests
if __name__ == "__main__":
from loguru import logger
import typer
from loguru import logger
else:
from src.utils import logger
class PaddleXLayoutParser:
"""PaddleX 版面解析服务客户端"""
def __init__(self, base_url: str = "http://localhost:8080"):
self.base_url = base_url.rstrip('/')
self.base_url = base_url.rstrip("/")
self.endpoint = f"{self.base_url}/layout-parsing"
def encode_file_to_base64(self, file_path: str) -> str:
with open(file_path, 'rb') as file:
encoded = base64.b64encode(file.read()).decode('utf-8')
with open(file_path, "rb") as file:
encoded = base64.b64encode(file.read()).decode("utf-8")
return encoded
def _process_file_input(self, file_input: str) -> str:
@ -43,7 +42,7 @@ class PaddleXLayoutParser:
raise
# 检查是否为URL
elif file_input.startswith(('http://', 'https://')):
elif file_input.startswith(("http://", "https://")):
logger.info(f"🌐 检测到URL: {file_input}")
return file_input
@ -52,21 +51,23 @@ class PaddleXLayoutParser:
logger.info(f"📝 假设为Base64编码内容，长度: {len(file_input)} 字符")
return file_input
def layout_parsing(self,
file_input: str,
file_type: int | None = None,
use_textline_orientation: bool | None = None,
use_seal_recognition: bool | None = None,
use_table_recognition: bool | None = None,
use_formula_recognition: bool | None = None,
use_chart_recognition: bool | None = None,
use_region_detection: bool | None = None,
layout_threshold: float | None = None,
layout_nms: bool | None = None,
use_doc_orientation_classify: bool = True,
use_doc_unwarping: bool | None = False,
use_wired_table_cells_trans_to_html: bool = True, # 启用则直接基于有线表单元格检测结果的几何关系构建HTML。
**kwargs) -> dict[str, Any]:
def layout_parsing(
self,
file_input: str,
file_type: int | None = None,
use_textline_orientation: bool | None = None,
use_seal_recognition: bool | None = None,
use_table_recognition: bool | None = None,
use_formula_recognition: bool | None = None,
use_chart_recognition: bool | None = None,
use_region_detection: bool | None = None,
layout_threshold: float | None = None,
layout_nms: bool | None = None,
use_doc_orientation_classify: bool = True,
use_doc_unwarping: bool | None = False,
use_wired_table_cells_trans_to_html: bool = True, # 启用则直接基于有线表单元格检测结果的几何关系构建HTML。
**kwargs,
) -> dict[str, Any]:
"""
调用版面解析API：https://paddlepaddle.github.io/PaddleX/latest/pipeline_usage/tutorials/ocr_pipelines/PP-StructureV3.html#22-python
"""
@ -102,10 +103,7 @@ class PaddleXLayoutParser:
try:
response = requests.post(
self.endpoint,
json=payload,
headers={"Content-Type": "application/json"},
timeout=300
self.endpoint, json=payload, headers={"Content-Type": "application/json"}, timeout=300
)
if response.status_code == 200:
@ -132,7 +130,6 @@ class PaddleXLayoutParser:
return {"error": str(e)}
def _parse_recognition_result(api_result: dict[str, Any], file_path: str) -> dict[str, Any]:
# 基本信息
parsed_result = {
@ -143,7 +140,7 @@ def _parse_recognition_result(api_result: dict[str, Any], file_path: str) -> dic
"total_pages": 0,
"pages": [],
"full_text": "",
"summary": {}
"summary": {},
}
result_data = api_result.get("result", {})
@ -155,7 +152,7 @@ def _parse_recognition_result(api_result: dict[str, Any], file_path: str) -> dic
parsed_result["document_info"] = {
"type": data_info.get("type", "unknown"),
"total_pages": data_info.get("numPages", len(layout_results)),
"page_dimensions": data_info.get("pages", [])
"page_dimensions": data_info.get("pages", []),
}
# 统计信息
@ -168,11 +165,7 @@ def _parse_recognition_result(api_result: dict[str, Any], file_path: str) -> dic
# 解析每页结果
for page_index, page_result in enumerate(layout_results):
page_info = {
"page_number": page_index + 1,
"content": {},
"statistics": {}
}
page_info = {"page_number": page_index + 1, "content": {}, "statistics": {}}
# Markdown内容
if "markdown" in page_result:
@ -240,37 +233,28 @@ def _parse_recognition_result(api_result: dict[str, Any], file_path: str) -> dic
"total_charts": total_charts,
"total_seals": total_seals,
"total_characters": len(parsed_result["full_text"]),
"average_elements_per_page": round(total_elements / max(1, len(layout_results)), 2)
"average_elements_per_page": round(total_elements / max(1, len(layout_results)), 2),
}
return parsed_result
def analyze_document(file_path: str, base_url: str = "http://localhost:8080") -> dict[str, Any]:
# 检查文件是否存在
if not os.path.exists(file_path):
return {
"success": False,
"error": f"文件不存在: {file_path}",
"file_path": file_path
}
return {"success": False, "error": f"文件不存在: {file_path}", "file_path": file_path}
# 初始化客户端
client = PaddleXLayoutParser(base_url=base_url)
# 判断文件类型
file_ext = os.path.splitext(file_path)[1].lower()
if file_ext == '.pdf':
if file_ext == ".pdf":
file_type = 0
elif file_ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']:
elif file_ext in [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"]:
file_type = 1
else:
return {
"success": False,
"error": f"不支持的文件类型: {file_ext}",
"file_path": file_path
}
return {"success": False, "error": f"不支持的文件类型: {file_ext}", "file_path": file_path}
logger.info(f"📄 开始分析文档: {os.path.basename(file_path)}")
logger.info(f"📏 文件大小: {os.path.getsize(file_path) / 1024 / 1024:.2f} MB")
@ -286,7 +270,7 @@ def analyze_document(file_path: str, base_url: str = "http://localhost:8080") ->
"success": False,
"error": result.get("errorMsg", "API调用失败"),
"file_path": file_path,
"raw_result": result
"raw_result": result,
}
# 解析结果
@ -294,11 +278,7 @@ def analyze_document(file_path: str, base_url: str = "http://localhost:8080") ->
return analysis_result
except Exception as e:
return {
"success": False,
"error": f"处理异常: {str(e)}",
"file_path": file_path
}
return {"success": False, "error": f"处理异常: {str(e)}", "file_path": file_path}
def check_paddlex_health(base_url: str = "http://localhost:8080") -> bool:
@ -318,7 +298,7 @@ def analyze_folder(input_dir: str, output_dir: str, base_url: str = "http://loca
output_path.mkdir(parents=True, exist_ok=True)
# 获取所有支持的文件
supported_extensions = {'.pdf', '.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
supported_extensions = {".pdf", ".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"}
files = []
for root, dirs, filenames in os.walk(input_dir):
for filename in filenames:
@ -343,12 +323,16 @@ def analyze_folder(input_dir: str, output_dir: str, base_url: str = "http://loca
if result.get("success"):
# 保持目录结构
relative_path = file_path.relative_to(input_path)
output_file = output_path / relative_path.with_suffix('.txt')
output_file = output_path / relative_path.with_suffix(".txt")
output_file.parent.mkdir(parents=True, exist_ok=True)
# 写入文本内容
text_content = result.get("full_text", "未提取到内容") if result.get("success") else f"分析失败：{result.get('error', '未知错误')}"
with open(output_file, 'w', encoding='utf-8') as f:
text_content = (
result.get("full_text", "未提取到内容")
if result.get("success")
else f"分析失败：{result.get('error', '未知错误')}"
)
with open(output_file, "w", encoding="utf-8") as f:
f.write(text_content)
success_count += 1

12
src/utils/__init__.py
View File

@ -1,10 +1,13 @@
import time
import hashlib
import os
import time
from src.utils.logging_config import logger
def is_text_pdf(pdf_path):
import fitz
doc = fitz.open(pdf_path)
total_pages = len(doc)
if total_pages == 0:
@ -22,6 +25,7 @@ def is_text_pdf(pdf_path):
# 如果超过50%的页面有文本内容，则认为是文本PDF
return text_ratio > 0.5
def hashstr(input_string, length=None, with_salt=False):
"""生成字符串的哈希值
Args:
@ -31,14 +35,14 @@ def hashstr(input_string, length=None, with_salt=False):
"""
try:
# 尝试直接编码
encoded_string = str(input_string).encode('utf-8')
encoded_string = str(input_string).encode("utf-8")
except UnicodeEncodeError:
# 如果编码失败，替换无效字符
encoded_string = str(input_string).encode('utf-8', errors='replace')
encoded_string = str(input_string).encode("utf-8", errors="replace")
if with_salt:
salt = str(time.time())
encoded_string = (encoded_string.decode('utf-8') + salt).encode('utf-8')
encoded_string = (encoded_string.decode("utf-8") + salt).encode("utf-8")
hash = hashlib.md5(encoded_string).hexdigest()
if length:

16
src/utils/logging_config.py
View File

@ -2,12 +2,12 @@ import os
from datetime import datetime
import pytz
from loguru import logger as loguru_logger
SAVE_DIR = os.getenv('SAVE_DIR', 'saves')
DATETIME = datetime.now(pytz.timezone('Asia/Shanghai')).strftime('%Y-%m-%d')
LOG_FILE = f'{SAVE_DIR}/logs/yuxi-{DATETIME}.log'
SAVE_DIR = os.getenv("SAVE_DIR", "saves")
DATETIME = datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y-%m-%d")
LOG_FILE = f"{SAVE_DIR}/logs/yuxi-{DATETIME}.log"
def setup_logger(name, level="DEBUG", console=True):
"""使用 loguru 设置日志记录器"""
@ -24,7 +24,7 @@ def setup_logger(name, level="DEBUG", console=True):
encoding="utf-8",
rotation="10 MB", # 文件大小达到 10MB 时轮转
retention="30 days", # 保留30天的日志
compression="zip" # 压缩旧日志文件
compression="zip", # 压缩旧日志文件
)
# 添加控制台日志（有颜色）
@ -33,16 +33,16 @@ def setup_logger(name, level="DEBUG", console=True):
lambda msg: print(msg, end=""),
level=level,
format="<green>{time:MM-DD HH:mm:ss}</green> <level>{level}</level> <cyan>{name}:{line}</cyan>: <level>{message}</level>",
colorize=True
colorize=True,
)
return loguru_logger
# 设置根日志记录器
logger = setup_logger('Yuxi')
logger = setup_logger("Yuxi")
__all__ = ['logger']
__all__ = ["logger"]
# If you want to disable logging from external libraries
# logging.getLogger('some_external_library').setLevel(logging.CRITICAL)

15
src/utils/prompts.py
View File

@ -1,7 +1,8 @@
from datetime import datetime
def get_system_prompt():
return (f"当前时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
return f"当前时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
knowbase_qa_template = """
@ -65,14 +66,4 @@ keywords_prompt_template = """
<文本>{text}</文本>
"""
HYDE_PROMPT_TEMPLATE = (
"Please write a passage to answer the question\n"
"Try to include as many key details as possible.\n"
"\n"
"\n"
"{context_str}\n"
"\n"
"{query}\n"
"\n"
'Passage:\n'
)
HYDE_PROMPT_TEMPLATE = "Please write a passage to answer the question\nTry to include as many key details as possible.\n\n\n{context_str}\n\n{query}\n\nPassage:\n"

25
src/utils/web_search.py
View File

@ -1,7 +1,10 @@
import os
from tavily import TavilyClient
from src.utils.logging_config import logger
class WebSearcher:
def __init__(self):
api_key = os.getenv("TAVILY_API_KEY")
@ -22,21 +25,19 @@ class WebSearcher:
搜索结果列表
"""
try:
search_results = self.client.search(
query=query,
search_depth="basic",
max_results=max_results
)
search_results = self.client.search(query=query, search_depth="basic", max_results=max_results)
# 提取需要的信息
formatted_results = []
for result in search_results['results'][:max_results]:
formatted_results.append({
'title': result.get('title', ''),
'content': result.get('content', ''),
'url': result.get('url', ''),
'score': result.get('score', 0)
})
for result in search_results["results"][:max_results]:
formatted_results.append(
{
"title": result.get("title", ""),
"content": result.get("content", ""),
"url": result.get("url", ""),
"score": result.get("score", 0),
}
)
return formatted_results

29
test/test_concurrency.py
View File

@ -1,15 +1,13 @@
import asyncio
import aiohttp
import time
import aiohttp
async def make_request(session: aiohttp.ClientSession, request_id: int) -> dict:
"""发送单个请求到API"""
url = "http://localhost:5000/chat/call"
payload = {
"query": "写一个冒泡排序",
"meta": {}
}
payload = {"query": "写一个冒泡排序", "meta": {}}
start_time = time.time()
print(f"请求 {request_id} 开始时间: {time.strftime('%H:%M:%S', time.localtime(start_time))}")
@ -19,19 +17,23 @@ async def make_request(session: aiohttp.ClientSession, request_id: int) -> dict:
print(f"请求 {request_id} 结果: {result}")
end_time = time.time()
duration = end_time - start_time
print(f"请求 {request_id} 完成时间: {time.strftime('%H:%M:%S', time.localtime(end_time))} (耗时: {duration:.2f}秒)")
print(
f"请求 {request_id} 完成时间: {time.strftime('%H:%M:%S', time.localtime(end_time))} (耗时: {duration:.2f}秒)"
)
return {
"request_id": request_id,
"status": response.status,
"time": duration,
"start_time": start_time,
"end_time": end_time,
"success": True
"success": True,
}
except Exception as e:
end_time = time.time()
duration = end_time - start_time
print(f"请求 {request_id} 失败时间: {time.strftime('%H:%M:%S', time.localtime(end_time))} (耗时: {duration:.2f}秒)")
print(
f"请求 {request_id} 失败时间: {time.strftime('%H:%M:%S', time.localtime(end_time))} (耗时: {duration:.2f}秒)"
)
return {
"request_id": request_id,
"status": None,
@ -39,15 +41,17 @@ async def make_request(session: aiohttp.ClientSession, request_id: int) -> dict:
"start_time": start_time,
"end_time": end_time,
"success": False,
"error": str(e)
"error": str(e),
}
async def run_concurrent_test(num_requests: int = 10) -> list[dict]:
"""运行并发测试"""
async with aiohttp.ClientSession() as session:
tasks = [make_request(session, i) for i in range(num_requests)]
return await asyncio.gather(*tasks)
def analyze_results(results: list[dict]) -> None:
"""分析并打印测试结果"""
total_requests = len(results)
@ -97,9 +101,9 @@ def analyze_results(results: list[dict]) -> None:
active_requests = [t for t in active_requests if t > start_time]
active_requests.append(end_time)
print(f"{result['request_id']:^7} {time.strftime('%H:%M:%S', time.localtime(start_time))} "
f"{time.strftime('%H:%M:%S', time.localtime(end_time))} "
f"{result['time']:^8.2f} {len(active_requests):^6}")
print(
f"{result['request_id']:^7} {time.strftime('%H:%M:%S', time.localtime(start_time))} {time.strftime('%H:%M:%S', time.localtime(end_time))} {result['time']:^8.2f} {len(active_requests):^6}"
)
# 计算最大并发数
max_concurrent = 0
@ -116,6 +120,7 @@ def analyze_results(results: list[dict]) -> None:
print(f"\n最大并发请求数: {max_concurrent}")
if __name__ == "__main__":
NUM_REQUESTS = 100 # 设置并发请求数

5
test/test_neo4j.py
View File

@ -1,8 +1,9 @@
from neo4j import GraphDatabase
from neo4j.exceptions import ServiceUnavailable, AuthError
from neo4j.exceptions import AuthError, ServiceUnavailable
# sudo ln -s /snap/core22/1586/usr/sbin/iptables /usr/sbin/iptables
def check_neo4j_status(uri="bolt://localhost:7687", username="neo4j", password="0123456789"):
"""
检查 Neo4j 数据库是否可以连接并正常工作。
@ -30,7 +31,7 @@ def check_neo4j_status(uri="bolt://localhost:7687", username="neo4j", password="
# 确保关闭驱动
driver.close()
# 测试函数
status = check_neo4j_status()
print(f"Neo4j status: {status}")