diff --git a/Makefile b/Makefile
index 22fb64b4..e7feab7f 100644
--- a/Makefile
+++ b/Makefile
@@ -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
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index b7767f32..6798d0d6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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 = [
diff --git a/scripts/batch_upload.py b/scripts/batch_upload.py
index ad908006..680dda40 100644
--- a/scripts/batch_upload.py
+++ b/scripts/batch_upload.py
@@ -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))
diff --git a/scripts/preprocessors/split_data_to_subfiles.py b/scripts/preprocessors/split_data_to_subfiles.py
index f74a1072..6e132f3d 100644
--- a/scripts/preprocessors/split_data_to_subfiles.py
+++ b/scripts/preprocessors/split_data_to_subfiles.py
@@ -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()
diff --git a/scripts/rename_milvus_collections.py b/scripts/rename_milvus_collections.py
index 1baa5eaa..8f23f284 100644
--- a/scripts/rename_milvus_collections.py
+++ b/scripts/rename_milvus_collections.py
@@ -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()
diff --git a/scripts/vllm/test_vllm.py b/scripts/vllm/test_vllm.py
index 204a4a76..2b0b92c7 100644
--- a/scripts/vllm/test_vllm.py
+++ b/scripts/vllm/test_vllm.py
@@ -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)
diff --git a/server/db_manager.py b/server/db_manager.py
index 24308c71..b40056a6 100644
--- a/server/db_manager.py
+++ b/server/db_manager.py
@@ -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()
diff --git a/server/main.py b/server/main.py
index d6e76bac..aeb6c31d 100644
--- a/server/main.py
+++ b/server/main.py
@@ -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)
-
diff --git a/server/models/kb_models.py b/server/models/kb_models.py
index c877ec17..65617a21 100644
--- a/server/models/kb_models.py
+++ b/server/models/kb_models.py
@@ -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 {}, # 确保映射正确
}
diff --git a/server/models/thread_model.py b/server/models/thread_model.py
index 10008373..4bdf764b 100644
--- a/server/models/thread_model.py
+++ b/server/models/thread_model.py
@@ -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")
diff --git a/server/models/user_model.py b/server/models/user_model.py
index 32a5ae88..fc937787 100644
--- a/server/models/user_model.py
+++ b/server/models/user_model.py
@@ -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,
}
diff --git a/server/routers/__init__.py b/server/routers/__init__.py
index 48dd6cd6..bfc7b8b2 100644
--- a/server/routers/__init__.py
+++ b/server/routers/__init__.py
@@ -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/*
diff --git a/server/routers/auth_router.py b/server/routers/auth_router.py
index 51689741..5135ea18 100644
--- a/server/routers/auth_router.py
+++ b/server/routers/auth_router.py
@@ -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
)
# 删除用户
diff --git a/server/routers/chat_router.py b/server/routers/chat_router.py
index 45f58e32..9e548476 100644
--- a/server/routers/chat_router.py
+++ b/server/routers/chat_router.py
@@ -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="对话线程不存在")
diff --git a/server/routers/graph_router.py b/server/routers/graph_router.py
index 0b348e38..04df5bb0 100644
--- a/server/routers/graph_router.py
+++ b/server/routers/graph_router.py
@@ -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"}
diff --git a/server/routers/knowledge_router.py b/server/routers/knowledge_router.py
index 68baa99d..168a9158 100644
--- a/server/routers/knowledge_router.py
+++ b/server/routers/knowledge_router.py
@@ -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)):
"""获取知识库统计信息"""
diff --git a/server/routers/system_router.py b/server/routers/system_router.py
index a4925659..f7945346 100644
--- a/server/routers/system_router.py
+++ b/server/routers/system_router.py
@@ -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服务健康检查完成"}
diff --git a/server/utils/auth_middleware.py b/server/utils/auth_middleware.py
index 346d987f..57920f49 100644
--- a/server/utils/auth_middleware.py
+++ b/server/utils/auth_middleware.py
@@ -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
diff --git a/server/utils/auth_utils.py b/server/utils/auth_utils.py
index 4ea6fe98..08da0380 100644
--- a/server/utils/auth_utils.py
+++ b/server/utils/auth_utils.py
@@ -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:
"""认证工具类"""
diff --git a/server/utils/common_utils.py b/server/utils/common_utils.py
index 73cce32c..1833f18b 100644
--- a/server/utils/common_utils.py
+++ b/server/utils/common_utils.py
@@ -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
diff --git a/src/__init__.py b/src/__init__.py
index 49ecfe27..f9fa2cd7 100644
--- a/src/__init__.py
+++ b/src/__init__.py
@@ -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()
diff --git a/src/agents/__init__.py b/src/agents/__init__.py
index 3598fba7..18be435a 100644
--- a/src/agents/__init__.py
+++ b/src/agents/__init__.py
@@ -3,6 +3,7 @@ import asyncio
from .chatbot.graph import ChatbotAgent
from .react.graph import ReActAgent
+
class AgentManager:
def __init__(self):
self._classes = {}
diff --git a/src/agents/chatbot/context.py b/src/agents/chatbot/context.py
index 83abeb44..816fcc82 100644
--- a/src/agents/chatbot/context.py
+++ b/src/agents/chatbot/context.py
@@ -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服务器列表"},
)
diff --git a/src/agents/chatbot/graph.py b/src/agents/chatbot/graph.py
index b3e65187..5b138695 100644
--- a/src/agents/chatbot/graph.py
+++ b/src/agents/chatbot/graph.py
@@ -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)
diff --git a/src/agents/chatbot/state.py b/src/agents/chatbot/state.py
index c37963cf..28338901 100644
--- a/src/agents/chatbot/state.py
+++ b/src/agents/chatbot/state.py
@@ -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)
diff --git a/src/agents/chatbot/tools.py b/src/agents/chatbot/tools.py
index 1d8fb52a..f4b921c2 100644
--- a/src/agents/chatbot/tools.py
+++ b/src/agents/chatbot/tools.py
@@ -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
-
diff --git a/src/agents/common/base.py b/src/agents/common/base.py
index 41a551f4..3d96cfb3 100644
--- a/src/agents/common/base.py
+++ b/src/agents/common/base.py
@@ -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
diff --git a/src/agents/common/context.py b/src/agents/common/context.py
index e4dc1123..eb00f79e 100644
--- a/src/agents/common/context.py
+++ b/src/agents/common/context.py
@@ -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__ 的字典
diff --git a/src/agents/common/mcp.py b/src/agents/common/mcp.py
index 9efe942e..ae0e5a05 100644
--- a/src/agents/common/mcp.py
+++ b/src/agents/common/mcp.py
@@ -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 = []
diff --git a/src/agents/common/models.py b/src/agents/common/models.py
index 7fe7fe02..befe1329 100644
--- a/src/agents/common/models.py
+++ b/src/agents/common/models.py
@@ -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),
diff --git a/src/agents/common/tools.py b/src/agents/common/tools.py
index ef902058..b986fabe 100644
--- a/src/agents/common/tools.py
+++ b/src/agents/common/tools.py
@@ -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:
diff --git a/src/agents/common/utils.py b/src/agents/common/utils.py
index ced1f75b..c6f70593 100644
--- a/src/agents/common/utils.py
+++ b/src/agents/common/utils.py
@@ -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()
-
diff --git a/src/agents/react/graph.py b/src/agents/react/graph.py
index ab7b966f..08e0a209 100644
--- a/src/agents/react/graph.py
+++ b/src/agents/react/graph.py
@@ -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
diff --git a/src/config/__init__.py b/src/config/__init__.py
index b1b82053..98f1eb6d 100644
--- a/src/config/__init__.py
+++ b/src/config/__init__.py
@@ -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")
diff --git a/src/knowledge/chroma_kb.py b/src/knowledge/chroma_kb.py
index 00613188..17b1a806 100644
--- a/src/knowledge/chroma_kb.py
+++ b/src/knowledge/chroma_kb.py
@@ -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)
diff --git a/src/knowledge/graphbase.py b/src/knowledge/graphbase.py
index 823bd15c..b80edbb2 100644
--- a/src/knowledge/graphbase.py
+++ b/src/knowledge/graphbase.py
@@ -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
diff --git a/src/knowledge/indexing.py b/src/knowledge/indexing.py
index b6590a29..c758e74f 100644
--- a/src/knowledge/indexing.py
+++ b/src/knowledge/indexing.py
@@ -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}"
-
diff --git a/src/knowledge/kb_factory.py b/src/knowledge/kb_factory.py
index b1e4dbc1..35549ca3 100644
--- a/src/knowledge/kb_factory.py
+++ b/src/knowledge/kb_factory.py
@@ -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
diff --git a/src/knowledge/kb_manager.py b/src/knowledge/kb_manager.py
index eea9f7e1..de750d8a 100644
--- a/src/knowledge/kb_manager.py
+++ b/src/knowledge/kb_manager.py
@@ -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
diff --git a/src/knowledge/kb_utils.py b/src/knowledge/kb_utils.py
index cce8d899..c94449d2 100644
--- a/src/knowledge/kb_utils.py
+++ b/src/knowledge/kb_utils.py
@@ -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}")
diff --git a/src/knowledge/knowledge_base.py b/src/knowledge/knowledge_base.py
index 5516aa63..7b67cb06 100644
--- a/src/knowledge/knowledge_base.py
+++ b/src/knowledge/knowledge_base.py
@@ -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}")
diff --git a/src/knowledge/lightrag_kb.py b/src/knowledge/lightrag_kb.py
index 31675043..c8c0ec36 100644
--- a/src/knowledge/lightrag_kb.py
+++ b/src/knowledge/lightrag_kb.py
@@ -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 原生功能导出知识库数据。
[注意] 此功能当前已禁用。
diff --git a/src/knowledge/milvus_kb.py b/src/knowledge/milvus_kb.py
index fe1cba60..a758f359 100644
--- a/src/knowledge/milvus_kb.py
+++ b/src/knowledge/milvus_kb.py
@@ -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
diff --git a/src/models/__init__.py b/src/models/__init__.py
index a900c394..27b7ffbf 100644
--- a/src/models/__init__.py
+++ b/src/models/__init__.py
@@ -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}")
diff --git a/src/models/chat_model.py b/src/models/chat_model.py
index 609173da..e70f68b0 100644
--- a/src/models/chat_model.py
+++ b/src/models/chat_model.py
@@ -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
diff --git a/src/models/embedding.py b/src/models/embedding.py
index 9ae0626f..170f0c6c 100644
--- a/src/models/embedding.py
+++ b/src/models/embedding.py
@@ -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}
diff --git a/src/models/rerank_model.py b/src/models/rerank_model.py
index f51944d7..574b0e62 100644
--- a/src/models/rerank_model.py
+++ b/src/models/rerank_model.py
@@ -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)
diff --git a/src/plugins/_ocr.py b/src/plugins/_ocr.py
index 535a4789..1b387693 100644
--- a/src/plugins/_ocr.py
+++ b/src/plugins/_ocr.py
@@ -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()
diff --git a/src/plugins/mineru.py b/src/plugins/mineru.py
index 67d3bbc7..6acd3cfa 100644
--- a/src/plugins/mineru.py
+++ b/src/plugins/mineru.py
@@ -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).
diff --git a/src/plugins/paddlex.py b/src/plugins/paddlex.py
index 3e26a854..5fff01c0 100644
--- a/src/plugins/paddlex.py
+++ b/src/plugins/paddlex.py
@@ -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
diff --git a/src/utils/__init__.py b/src/utils/__init__.py
index b1c0da2d..40ca0df7 100644
--- a/src/utils/__init__.py
+++ b/src/utils/__init__.py
@@ -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:
diff --git a/src/utils/logging_config.py b/src/utils/logging_config.py
index d1563f29..c3e089c3 100644
--- a/src/utils/logging_config.py
+++ b/src/utils/logging_config.py
@@ -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="{time:MM-DD HH:mm:ss} {level} {name}:{line}: {message}",
- 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)
diff --git a/src/utils/prompts.py b/src/utils/prompts.py
index 540d2193..be740692 100644
--- a/src/utils/prompts.py
+++ b/src/utils/prompts.py
@@ -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"
diff --git a/src/utils/web_search.py b/src/utils/web_search.py
index d8b0e3c4..4fd10671 100644
--- a/src/utils/web_search.py
+++ b/src/utils/web_search.py
@@ -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
diff --git a/test/test_concurrency.py b/test/test_concurrency.py
index 0061d3bb..5a78b9a6 100644
--- a/test/test_concurrency.py
+++ b/test/test_concurrency.py
@@ -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 # 设置并发请求数
diff --git a/test/test_neo4j.py b/test/test_neo4j.py
index 53b967a0..f4c6e1c7 100755
--- a/test/test_neo4j.py
+++ b/test/test_neo4j.py
@@ -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}")
-