2064 lines
77 KiB
Python
2064 lines
77 KiB
Python
"""
|
||
FastAPI main application
|
||
"""
|
||
import warnings
|
||
warnings.filterwarnings("ignore", message=".*pkg_resources is deprecated.*", category=UserWarning)
|
||
|
||
from contextlib import asynccontextmanager
|
||
from pathlib import Path
|
||
from fastapi import FastAPI, HTTPException, UploadFile, File, Form
|
||
from fastapi.responses import StreamingResponse, Response, FileResponse
|
||
from fastapi.staticfiles import StaticFiles
|
||
from fastapi.openapi.docs import get_swagger_ui_html, get_redoc_html
|
||
from pydantic import BaseModel, Field
|
||
from typing import Optional, List, Dict, Any
|
||
from loguru import logger
|
||
from config import settings
|
||
from db_utils import get_db_connection, init_session_db
|
||
from rag import VectorStoreManager, RAGEngine, FileParser
|
||
from sync_service import SyncServiceManager
|
||
import requests
|
||
from datetime import datetime
|
||
import time
|
||
import uuid
|
||
import threading
|
||
import logging
|
||
from fastapi import Request
|
||
from fastapi.responses import PlainTextResponse
|
||
import sqlite3
|
||
import json
|
||
import os
|
||
import asyncio
|
||
import re
|
||
import markdown2
|
||
from pathlib import Path
|
||
import hashlib
|
||
from sync.base_sync import BaseSync
|
||
|
||
|
||
# Global instances
|
||
vector_store_manager: Optional[VectorStoreManager] = None
|
||
rag_engine: Optional[RAGEngine] = None
|
||
sync_manager: Optional[SyncServiceManager] = None
|
||
file_parser: Optional[FileParser] = None
|
||
auto_sync_task = None # Keep reference to prevent garbage collection
|
||
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(app: FastAPI):
|
||
"""Lifespan context manager for startup and shutdown events"""
|
||
global vector_store_manager, rag_engine, sync_manager, file_parser, auto_sync_task
|
||
|
||
# Startup
|
||
try:
|
||
logger.info("Initializing RAG services...")
|
||
|
||
# Initialize vector store
|
||
try:
|
||
vector_store_manager = VectorStoreManager()
|
||
logger.info("✓ VectorStoreManager initialized")
|
||
except Exception as e:
|
||
logger.error(f"Failed to initialize VectorStoreManager: {e}")
|
||
raise
|
||
|
||
# Initialize RAG engine
|
||
try:
|
||
rag_engine = RAGEngine(vector_store_manager)
|
||
logger.info("✓ RAGEngine initialized")
|
||
except Exception as e:
|
||
logger.error(f"Failed to initialize RAGEngine: {e}")
|
||
raise
|
||
|
||
# Initialize file parser
|
||
try:
|
||
file_parser = FileParser()
|
||
logger.info("✓ FileParser initialized")
|
||
except Exception as e:
|
||
logger.error(f"Failed to initialize FileParser: {e}")
|
||
raise
|
||
|
||
|
||
logger.info("✓ Core RAG services initialized")
|
||
|
||
# Initialize sync service in background (non-blocking)
|
||
# This allows API to start immediately even if MySQL connection fails
|
||
async def init_sync_service():
|
||
"""Initialize sync service and start sync in background"""
|
||
global sync_manager
|
||
try:
|
||
logger.info("Initializing sync services in background...")
|
||
|
||
# Run SyncServiceManager initialization in thread pool to avoid blocking event loop
|
||
# SyncServiceManager.__init__() contains synchronous data source connection checks
|
||
loop = asyncio.get_event_loop()
|
||
sync_manager = await loop.run_in_executor(None, SyncServiceManager)
|
||
|
||
# Start all sync services in background
|
||
# This will start auto sync with recovery for each data source
|
||
if settings.AUTO_SYNC:
|
||
logger.info(f"Starting auto sync services (interval: {settings.SYNC_INTERVAL}s)...")
|
||
# Start all sync services in background
|
||
await sync_manager.start_all_sync_services()
|
||
else:
|
||
logger.info("Auto sync is disabled, skipping auto sync service startup")
|
||
except Exception as e:
|
||
logger.error(f"Failed to initialize sync services: {e}")
|
||
logger.warning("API is still available, but sync services are disabled. Please check data source connections.")
|
||
# Don't raise - allow API to continue running
|
||
|
||
# Start sync service initialization in background (non-blocking)
|
||
# This ensures API can respond immediately while sync runs in background
|
||
sync_init_task = asyncio.create_task(init_sync_service())
|
||
auto_sync_task = sync_init_task # Store reference for shutdown
|
||
|
||
logger.info("✓ RAG API is now available (sync running in background)")
|
||
except Exception as e:
|
||
logger.error(f"Failed to initialize core services: {e}")
|
||
raise
|
||
|
||
yield
|
||
|
||
# Shutdown
|
||
if auto_sync_task and not auto_sync_task.done():
|
||
logger.info("Stopping sync services...")
|
||
# Wait a bit for the task to finish gracefully
|
||
try:
|
||
await asyncio.wait_for(auto_sync_task, timeout=10.0)
|
||
except asyncio.TimeoutError:
|
||
logger.warning("Sync task did not stop in time, cancelling...")
|
||
auto_sync_task.cancel()
|
||
try:
|
||
await auto_sync_task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
if sync_manager:
|
||
sync_manager.stop_all_sync_services()
|
||
sync_manager.close_all()
|
||
logger.info("RAG services shut down")
|
||
|
||
|
||
# Check if static assets are available for offline use
|
||
STATIC_ROOT_DIR = Path(__file__).parent.parent / "static"
|
||
STATIC_DIR = STATIC_ROOT_DIR / "swagger-ui"
|
||
SWAGGER_UI_BUNDLE = STATIC_DIR / "swagger-ui-bundle.js"
|
||
SWAGGER_UI_CSS = STATIC_DIR / "swagger-ui.css"
|
||
REDOC_BUNDLE = STATIC_DIR / "redoc.standalone.js"
|
||
|
||
# Determine if we should use local assets or CDN
|
||
USE_LOCAL_ASSETS = (
|
||
SWAGGER_UI_BUNDLE.exists() and
|
||
SWAGGER_UI_CSS.exists() and
|
||
REDOC_BUNDLE.exists()
|
||
)
|
||
|
||
if USE_LOCAL_ASSETS:
|
||
logger.info("Using local Swagger UI assets for offline mode")
|
||
else:
|
||
logger.warning(
|
||
"Local Swagger UI assets not found. Swagger UI will use CDN resources. "
|
||
"For offline use, run: python download_swagger_assets.py"
|
||
)
|
||
|
||
# Initialize FastAPI app
|
||
app = FastAPI(
|
||
title=settings.API_TITLE,
|
||
version=settings.API_VERSION,
|
||
description="RAG API for local knowledge base retrieval and generation",
|
||
lifespan=lifespan,
|
||
swagger_ui_parameters={
|
||
"persistAuthorization": True,
|
||
},
|
||
)
|
||
|
||
# Mount static files directories if they exist
|
||
if STATIC_ROOT_DIR.exists():
|
||
# Mount the entire static directory to serve all static assets
|
||
app.mount("/static", StaticFiles(directory=str(STATIC_ROOT_DIR)), name="static")
|
||
logger.info(f"Mounted static directory: {STATIC_ROOT_DIR}")
|
||
|
||
# Add route for chat interface
|
||
@app.get("/chat", include_in_schema=False)
|
||
async def chat_interface():
|
||
"""Serve the chat interface"""
|
||
return FileResponse(STATIC_ROOT_DIR / "chat" / "index.html")
|
||
|
||
# Add route for configuration management interface
|
||
@app.get("/config", include_in_schema=False)
|
||
async def config_interface():
|
||
"""Serve the configuration management interface"""
|
||
return FileResponse(STATIC_ROOT_DIR / "config" / "index.html")
|
||
|
||
|
||
# Add favicon route to prevent 404 errors
|
||
@app.get("/favicon.ico", include_in_schema=False)
|
||
async def favicon():
|
||
"""Return empty favicon to prevent 404 errors"""
|
||
# Return a minimal 1x1 transparent PNG
|
||
# This prevents browser from requesting favicon and getting 404
|
||
return Response(
|
||
content=b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\nIDATx\x9cc\x00\x01\x00\x00\x05\x00\x01\r\n-\xdb\x00\x00\x00\x00IEND\xaeB`\x82',
|
||
media_type="image/png"
|
||
)
|
||
|
||
|
||
# Override Swagger UI to use local assets when available
|
||
@app.get("/docs", include_in_schema=False)
|
||
async def custom_swagger_ui_html():
|
||
"""Custom Swagger UI that uses local assets in offline mode"""
|
||
if USE_LOCAL_ASSETS:
|
||
# Use local static files
|
||
return get_swagger_ui_html(
|
||
openapi_url=app.openapi_url,
|
||
title=app.title + " - Swagger UI",
|
||
swagger_js_url="/static/swagger-ui/swagger-ui-bundle.js",
|
||
swagger_css_url="/static/swagger-ui/swagger-ui.css",
|
||
swagger_favicon_url="/favicon.ico",
|
||
oauth2_redirect_url=app.swagger_ui_oauth2_redirect_url,
|
||
init_oauth=app.swagger_ui_init_oauth,
|
||
swagger_ui_parameters=app.swagger_ui_parameters,
|
||
)
|
||
else:
|
||
# Fallback to default (CDN)
|
||
return get_swagger_ui_html(
|
||
openapi_url=app.openapi_url,
|
||
title=app.title + " - Swagger UI",
|
||
swagger_favicon_url="/favicon.ico",
|
||
oauth2_redirect_url=app.swagger_ui_oauth2_redirect_url,
|
||
init_oauth=app.swagger_ui_init_oauth,
|
||
swagger_ui_parameters=app.swagger_ui_parameters,
|
||
)
|
||
|
||
|
||
# Override ReDoc to use local assets when available
|
||
@app.get("/redoc", include_in_schema=False)
|
||
async def custom_redoc_html():
|
||
"""Custom ReDoc that uses local assets in offline mode"""
|
||
if USE_LOCAL_ASSETS:
|
||
# Use local static files
|
||
return get_redoc_html(
|
||
openapi_url=app.openapi_url,
|
||
title=app.title + " - ReDoc",
|
||
redoc_js_url="/static/swagger-ui/redoc.standalone.js",
|
||
redoc_favicon_url="/favicon.ico",
|
||
with_google_fonts=False, # Disable Google Fonts for offline use
|
||
)
|
||
else:
|
||
# Fallback to default (CDN)
|
||
return get_redoc_html(
|
||
openapi_url=app.openapi_url,
|
||
title=app.title + " - ReDoc",
|
||
redoc_favicon_url="/favicon.ico",
|
||
with_google_fonts=True,
|
||
)
|
||
|
||
|
||
# Request/Response models
|
||
class QueryRequest(BaseModel):
|
||
"""Query request model"""
|
||
query: str = Field(..., description="User query string", min_length=1)
|
||
top_k: Optional[int] = Field(None, description="Number of documents to retrieve", ge=1, le=20)
|
||
stream: bool = Field(True, description="Whether to stream the response")
|
||
|
||
|
||
class RetrieveRequest(BaseModel):
|
||
"""Retrieve request model - only retrieves documents from ChromaDB, no LLM generation"""
|
||
query: str = Field(..., description="Query string for retrieval", min_length=1)
|
||
top_k: Optional[int] = Field(None, description="Number of documents to retrieve", ge=1, le=20)
|
||
|
||
|
||
class RetrievedDocument(BaseModel):
|
||
"""Retrieved document model"""
|
||
content: str = Field(..., description="Document content/text")
|
||
score: Optional[float] = Field(None, description="Similarity score")
|
||
metadata: dict = Field(default_factory=dict, description="Document metadata")
|
||
doc_id: Optional[str] = Field(None, description="Document ID")
|
||
|
||
|
||
class RetrieveResponse(BaseModel):
|
||
"""Retrieve response model"""
|
||
query: str = Field(..., description="Original query")
|
||
documents: list[RetrievedDocument] = Field(..., description="Retrieved documents")
|
||
count: int = Field(..., description="Number of documents retrieved")
|
||
|
||
|
||
class SyncRequest(BaseModel):
|
||
"""Manual sync request model"""
|
||
full_sync: bool = Field(False, description="Whether to perform full sync")
|
||
force: bool = Field(False, description="Whether to force re-processing of all documents (even if they exist)")
|
||
source_name: Optional[str] = Field(None, description="Specific data source to sync (optional, syncs all if not provided)")
|
||
|
||
|
||
class HealthResponse(BaseModel):
|
||
"""Health check response"""
|
||
status: str
|
||
message: str
|
||
|
||
|
||
class DocumentChunk(BaseModel):
|
||
"""Document chunk model"""
|
||
id: str = Field(..., description="Chunk ID")
|
||
text: str = Field(..., description="Chunk content")
|
||
metadata: dict = Field(default_factory=dict, description="Chunk metadata")
|
||
chunk_index: Optional[int] = Field(None, description="Chunk index (if document was chunked)")
|
||
|
||
|
||
class DocumentResponse(BaseModel):
|
||
"""Document response model"""
|
||
doc_id: str = Field(..., description="Document ID")
|
||
chunks: List[DocumentChunk] = Field(..., description="Document chunks")
|
||
total_chunks: int = Field(..., description="Total number of chunks")
|
||
full_text: str = Field(..., description="Full document text (all chunks combined)")
|
||
|
||
|
||
class UploadResponse(BaseModel):
|
||
"""File upload response model"""
|
||
doc_id: str = Field(..., description="Document ID")
|
||
filename: str = Field(..., description="Original filename")
|
||
file_type: str = Field(..., description="File type/extension")
|
||
chunks: int = Field(..., description="Number of chunks created")
|
||
message: str = Field(..., description="Upload status message")
|
||
|
||
# User authentication models
|
||
class UserRegisterRequest(BaseModel):
|
||
"""Request model for user registration"""
|
||
username: str
|
||
password: str
|
||
|
||
class UserLoginRequest(BaseModel):
|
||
"""Request model for user login"""
|
||
username: str
|
||
password: str
|
||
|
||
class UserResponse(BaseModel):
|
||
"""Response model for user operations"""
|
||
id: str
|
||
username: str
|
||
message: str
|
||
|
||
|
||
# --- Compatibility layer: endpoints copied from example.py (Flask) ---
|
||
# These endpoints provide a backward-compatible interface so existing clients
|
||
# that used the Flask-based API can keep working.
|
||
|
||
# In-memory conversation store and lock (thread-safe)
|
||
conversations = {}
|
||
conv_lock = threading.Lock()
|
||
|
||
# Configure logging for this compatibility layer
|
||
compat_logger = logging.getLogger("compat_api")
|
||
compat_logger.setLevel(logging.INFO)
|
||
|
||
# Session persistence (SQLite)
|
||
DATA_DIR = Path(__file__).parent.parent / "data"
|
||
DATA_DIR.mkdir(parents=True, exist_ok=True)
|
||
DB_PATH = DATA_DIR / "sessions.db"
|
||
|
||
|
||
def hash_password(password: str) -> str:
|
||
"""Hash a password using SHA-256"""
|
||
return hashlib.sha256(password.encode()).hexdigest()
|
||
|
||
|
||
|
||
|
||
|
||
def load_sessions():
|
||
conn = sqlite3.connect(DB_PATH)
|
||
try:
|
||
cur = conn.execute("SELECT id, user_login, title, data, update_time FROM sessions")
|
||
rows = cur.fetchall()
|
||
with conv_lock:
|
||
for rid, user_login, title, data_text, update_time in rows:
|
||
try:
|
||
data = json.loads(data_text) if data_text else {"messages": []}
|
||
except Exception:
|
||
data = {"messages": []}
|
||
conversations[rid] = {
|
||
"title": title or "",
|
||
"user_login": user_login or "",
|
||
"messages": data.get("messages", []),
|
||
"create_time": update_time,
|
||
"update_time": update_time,
|
||
"active_stream": False
|
||
}
|
||
finally:
|
||
conn.close()
|
||
|
||
def serialize_conversation(conv):
|
||
"""安全的序列化函数，处理可能包含协程的情况"""
|
||
messages = conv.get("messages", [])
|
||
|
||
# 清理消息列表，移除不可序列化的对象
|
||
cleaned_messages = []
|
||
for msg in messages:
|
||
if asyncio.iscoroutine(msg):
|
||
# 如果是协程，记录错误或跳过
|
||
compat_logger.warning(f"发现协程对象在消息列表中: {msg}")
|
||
continue
|
||
elif isinstance(msg, dict):
|
||
# 递归清理字典中的值
|
||
cleaned_msg = {}
|
||
for key, value in msg.items():
|
||
if not asyncio.iscoroutine(value):
|
||
cleaned_msg[key] = value
|
||
else:
|
||
cleaned_msg[key] = "<coroutine>"
|
||
cleaned_messages.append(cleaned_msg)
|
||
else:
|
||
cleaned_messages.append(msg)
|
||
|
||
return json.dumps(
|
||
{"messages": cleaned_messages},
|
||
ensure_ascii=False,
|
||
default=str # 处理其他不可序列化的类型
|
||
)
|
||
|
||
|
||
def save_session(conv_id: str):
|
||
# Upsert session into SQLite
|
||
with conv_lock:
|
||
conv = conversations.get(conv_id)
|
||
if not conv:
|
||
return
|
||
title = conv.get("title", "")
|
||
user_login = conv.get("user_login", "")
|
||
data_text = serialize_conversation(conv)
|
||
update_time = conv.get("update_time", datetime.now().isoformat())
|
||
|
||
conn = sqlite3.connect(DB_PATH)
|
||
try:
|
||
conn.execute(
|
||
"REPLACE INTO sessions (id, user_login, title, data, update_time) VALUES (?, ?, ?, ?, ?)",
|
||
(conv_id, user_login, title, data_text, update_time)
|
||
)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
# Initialize DB and load sessions at import
|
||
init_session_db()
|
||
load_sessions()
|
||
|
||
|
||
def build_prompt_from_history(conv_id: str, max_messages: int = 10) -> str:
|
||
"""Build a single prompt string that includes recent conversation history."""
|
||
|
||
def filter_coroutines_and_join(parts):
|
||
"""
|
||
过滤掉协程对象和协程字符串，只保留真正的内容
|
||
"""
|
||
filtered_parts = []
|
||
|
||
for item in parts:
|
||
# 如果是协程对象，直接跳过
|
||
if asyncio.iscoroutine(item):
|
||
continue
|
||
|
||
# 如果是字符串，检查是否包含协程描述
|
||
if isinstance(item, str):
|
||
# 检查是否包含协程标记
|
||
if contains_coroutine_marker(item):
|
||
continue
|
||
|
||
# 保留非协程相关的字符串
|
||
filtered_parts.append(item)
|
||
# 其他非字符串类型（如果不是协程）
|
||
else:
|
||
filtered_parts.append(str(item))
|
||
|
||
return "\n".join(filtered_parts)
|
||
|
||
def contains_coroutine_marker(text):
|
||
"""检查文本是否包含协程标记"""
|
||
coroutine_patterns = [
|
||
r'<coroutine object',
|
||
r'<coroutine>',
|
||
r'coroutine object at 0x',
|
||
r'coroutine at 0x'
|
||
]
|
||
|
||
for pattern in coroutine_patterns:
|
||
if re.search(pattern, text, re.IGNORECASE):
|
||
return True
|
||
|
||
return False
|
||
|
||
with conv_lock:
|
||
conv = conversations.get(conv_id)
|
||
if not conv:
|
||
return None
|
||
msgs = conv.get("messages", [])[-max_messages:]
|
||
|
||
if len(msgs) <= 1: # only new query content
|
||
return None
|
||
|
||
parts = []
|
||
for m in msgs[:-1]: # Exclude the last message (current user query)
|
||
role = m.get("role", "user")
|
||
content = m.get("content", "")
|
||
if role == "user":
|
||
parts.append(f"User: {content}")
|
||
else:
|
||
parts.append(f"Assistant: {content}")
|
||
if len(parts) == 0:
|
||
return None
|
||
else:
|
||
return filter_coroutines_and_join(parts) # 防止coroutine对象被放入messages中
|
||
|
||
|
||
# API endpoints
|
||
@app.get("/", response_model=HealthResponse)
|
||
async def root():
|
||
"""Root endpoint"""
|
||
return HealthResponse(
|
||
status="ok",
|
||
message="RAG API is running"
|
||
)
|
||
|
||
|
||
@app.get("/health", response_model=HealthResponse)
|
||
async def health_check():
|
||
"""Health check endpoint"""
|
||
try:
|
||
if vector_store_manager is None or rag_engine is None:
|
||
raise HTTPException(status_code=503, detail="Services not initialized")
|
||
|
||
return HealthResponse(
|
||
status="healthy",
|
||
message="All services are running"
|
||
)
|
||
except Exception as e:
|
||
raise HTTPException(status_code=503, detail=f"Service unhealthy: {str(e)}")
|
||
|
||
|
||
@app.post("/register", response_model=UserResponse)
|
||
async def register_user(request: UserRegisterRequest):
|
||
"""Register a new user"""
|
||
try:
|
||
# Check if username already exists
|
||
conn = sqlite3.connect(DB_PATH)
|
||
cur = conn.execute("SELECT * FROM users WHERE username = ?", (request.username,))
|
||
if cur.fetchone():
|
||
conn.close()
|
||
raise HTTPException(status_code=400, detail="Username already exists")
|
||
|
||
# Create new user
|
||
user_id = str(uuid.uuid4())
|
||
hashed_pw = hash_password(request.password)
|
||
create_time = datetime.now().isoformat()
|
||
|
||
conn.execute(
|
||
"INSERT INTO users (id, username, password, create_time) VALUES (?, ?, ?, ?)",
|
||
(user_id, request.username, hashed_pw, create_time)
|
||
)
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
return UserResponse(
|
||
id=user_id,
|
||
username=request.username,
|
||
message="User registered successfully"
|
||
)
|
||
except HTTPException:
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"Registration failed: {e}")
|
||
raise HTTPException(status_code=500, detail="Registration failed")
|
||
|
||
|
||
@app.post("/login", response_model=UserResponse)
|
||
async def login_user(request: UserLoginRequest):
|
||
"""Login a user"""
|
||
try:
|
||
conn = sqlite3.connect(DB_PATH)
|
||
cur = conn.execute("SELECT id, password FROM users WHERE username = ?", (request.username,))
|
||
user = cur.fetchone()
|
||
conn.close()
|
||
|
||
if not user or user[1] != hash_password(request.password):
|
||
raise HTTPException(status_code=401, detail="Invalid username or password")
|
||
|
||
return UserResponse(
|
||
id=user[0],
|
||
username=request.username,
|
||
message="Login successful"
|
||
)
|
||
except HTTPException:
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"Login failed: {e}")
|
||
raise HTTPException(status_code=500, detail="Login failed")
|
||
|
||
|
||
@app.post("/query")
|
||
async def query(request: QueryRequest):
|
||
"""
|
||
Query the RAG system (retrieves documents and generates LLM response)
|
||
|
||
Args:
|
||
request: Query request with query string and options
|
||
|
||
Returns:
|
||
Streaming response or complete response
|
||
"""
|
||
if rag_engine is None:
|
||
raise HTTPException(status_code=503, detail="RAG engine not initialized")
|
||
|
||
try:
|
||
if request.stream:
|
||
# Stream response
|
||
|
||
return StreamingResponse(
|
||
rag_engine.query_stream(request.query, None, request.top_k),
|
||
media_type="text/event-stream",
|
||
headers={
|
||
"X-Accel-Buffering": "no",
|
||
"Cache-Control": "no-cache",
|
||
"Connection": "keep-alive"
|
||
}
|
||
)
|
||
else:
|
||
# Return complete response (run in thread pool for better concurrency)
|
||
response = await rag_engine.query(request.query, None, request.top_k)
|
||
return response
|
||
except Exception as e:
|
||
logger.error(f"Error processing query: {e}")
|
||
raise HTTPException(status_code=500, detail=f"Error processing query: {str(e)}")
|
||
|
||
|
||
# ----------------------
|
||
# Backward-compatible Flask-style endpoints (from example.py)
|
||
# ----------------------
|
||
|
||
|
||
@app.get('/conversations')
|
||
def list_conversations():
|
||
with conv_lock:
|
||
conv_list = [
|
||
{
|
||
"id": conv_id,
|
||
"title": conv["title"],
|
||
"update_time": conv["update_time"],
|
||
"active_stream": conv.get("active_stream", False)
|
||
}
|
||
for conv_id, conv in conversations.items()
|
||
]
|
||
|
||
conv_list.sort(key=lambda x: x["update_time"], reverse=True)
|
||
return conv_list
|
||
|
||
@app.get('/conversations/{user_login}')
|
||
def list_conversations_by_user(user_login: str):
|
||
with conv_lock:
|
||
conv_list = [
|
||
{
|
||
"id": conv_id,
|
||
"user_login": conv.get("user_login", ""),
|
||
"title": conv["title"],
|
||
"update_time": conv["update_time"],
|
||
"active_stream": conv.get("active_stream", False)
|
||
}
|
||
for conv_id, conv in conversations.items()
|
||
if conv.get("user_login") == user_login
|
||
]
|
||
|
||
conv_list.sort(key=lambda x: x["update_time"], reverse=True)
|
||
return conv_list
|
||
|
||
|
||
@app.post('/conversations')
|
||
def create_conversation(payload: dict):
|
||
title = payload.get("title", "新对话")
|
||
user_login = payload.get("user_login")
|
||
if not user_login:
|
||
raise HTTPException(status_code=400, detail="user_login is required")
|
||
conv_id = str(uuid.uuid4())
|
||
with conv_lock:
|
||
conversations[conv_id] = {
|
||
"title": title,
|
||
"user_login": user_login,
|
||
"messages": [],
|
||
"create_time": datetime.now().isoformat(),
|
||
"update_time": datetime.now().isoformat(),
|
||
"active_stream": False
|
||
}
|
||
# Persist
|
||
save_session(conv_id)
|
||
return {"id": conv_id, "title": title}
|
||
|
||
|
||
# @app.get('/conversations/{conv_id}')
|
||
# def get_conversation(conv_id: str):
|
||
# with conv_lock:
|
||
# if conv_id not in conversations:
|
||
# raise HTTPException(status_code=404, detail="对话不存在")
|
||
# conv = conversations[conv_id]
|
||
# return {
|
||
# "id": conv_id,
|
||
# "title": conv["title"],
|
||
# "messages": conv["messages"],
|
||
# "create_time": conv["create_time"],
|
||
# "update_time": conv["update_time"],
|
||
# "active_stream": conv.get("active_stream", False)
|
||
# }
|
||
|
||
@app.get('/conversations/{user_login}/{conv_id}')
|
||
def get_conversation_by_user(user_login: str, conv_id: str):
|
||
with conv_lock:
|
||
if conv_id not in conversations:
|
||
raise HTTPException(status_code=404, detail="对话不存在")
|
||
conv = conversations[conv_id]
|
||
if conv.get("user_login") != user_login:
|
||
raise HTTPException(status_code=403, detail="无权访问该对话")
|
||
return {
|
||
"id": conv_id,
|
||
"title": conv["title"],
|
||
"messages": conv["messages"],
|
||
"create_time": conv["create_time"],
|
||
"update_time": conv["update_time"],
|
||
"active_stream": conv.get("active_stream", False)
|
||
}
|
||
|
||
|
||
@app.delete('/conversations/{conv_id}')
|
||
def delete_conversation(conv_id: str):
|
||
with conv_lock:
|
||
if conv_id not in conversations:
|
||
raise HTTPException(status_code=404, detail="对话不存在")
|
||
del conversations[conv_id]
|
||
return {"message": "删除成功"}
|
||
|
||
@app.delete('/conversations/{user_login}/{conv_id}')
|
||
def delete_conversation_by_user(user_login: str, conv_id: str):
|
||
with conv_lock:
|
||
if conv_id not in conversations:
|
||
raise HTTPException(status_code=404, detail="对话不存在")
|
||
conv = conversations[conv_id]
|
||
if conv.get("user_login") != user_login:
|
||
raise HTTPException(status_code=403, detail="无权删除该对话")
|
||
del conversations[conv_id]
|
||
return {"message": "删除成功"}
|
||
|
||
|
||
@app.post('/conversations/{conv_id}/messages')
|
||
async def post_message(conv_id: str, payload: dict):
|
||
if conv_id not in conversations:
|
||
raise HTTPException(status_code=404, detail="对话不存在")
|
||
|
||
if not isinstance(payload, dict) or "content" not in payload:
|
||
raise HTTPException(status_code=400, detail="Content required in JSON body")
|
||
|
||
user_msg = {"role": "user", "content": payload["content"], "time": datetime.now().isoformat()}
|
||
with conv_lock:
|
||
conv = conversations[conv_id]
|
||
conv["messages"].append(user_msg)
|
||
conv["update_time"] = datetime.now().isoformat()
|
||
if len(conv["messages"]) == 1:
|
||
conv["title"] = user_msg["content"][:20] + ("..." if len(user_msg["content"]) > 20 else "")
|
||
|
||
# persist after adding user message
|
||
save_session(conv_id)
|
||
|
||
# Use local RAG engine to generate response instead of external API
|
||
if rag_engine is None:
|
||
ai_content = "RAG engine not initialized"
|
||
else:
|
||
try:
|
||
history = build_prompt_from_history(conv_id, max_messages=10)
|
||
ai_content = await rag_engine.query(payload["content"], history)
|
||
except Exception as e:
|
||
compat_logger.error(f"RAG query failed: {e}")
|
||
ai_content = f"RAG错误: {e}"
|
||
|
||
ai_msg = {"role": "assistant", "content": ai_content, "time": datetime.now().isoformat()}
|
||
with conv_lock:
|
||
conversations[conv_id]["messages"].append(ai_msg)
|
||
conversations[conv_id]["update_time"] = datetime.now().isoformat()
|
||
|
||
# persist after adding assistant message
|
||
save_session(conv_id)
|
||
|
||
return ai_msg
|
||
|
||
|
||
@app.post('/conversations/{conv_id}/stream')
|
||
async def stream_message(conv_id: str, payload: dict):
|
||
if conv_id not in conversations:
|
||
raise HTTPException(status_code=404, detail="对话不存在")
|
||
|
||
with conv_lock:
|
||
if conversations[conv_id].get("active_stream", False):
|
||
raise HTTPException(status_code=429, detail="该对话已有活动流，请先中止当前流")
|
||
conversations[conv_id]["active_stream"] = True
|
||
|
||
if not isinstance(payload, dict) or "content" not in payload:
|
||
with conv_lock:
|
||
conversations[conv_id]["active_stream"] = False
|
||
raise HTTPException(status_code=400, detail="Content required in JSON body")
|
||
|
||
if not isinstance(payload, dict) or "user_login" not in payload:
|
||
with conv_lock:
|
||
conversations[conv_id]["active_stream"] = False
|
||
raise HTTPException(status_code=400, detail="user_login required in JSON body")
|
||
|
||
if conversations[conv_id].get("user_login") != payload["user_login"]:
|
||
with conv_lock:
|
||
conversations[conv_id]["active_stream"] = False
|
||
raise HTTPException(status_code=403, detail="用户无权操作该对话")
|
||
|
||
user_msg = {"role": "user", "content": payload["content"], "time": datetime.now().isoformat()}
|
||
with conv_lock:
|
||
conversations[conv_id]["messages"].append(user_msg)
|
||
conversations[conv_id]["update_time"] = datetime.now().isoformat()
|
||
if len(conversations[conv_id]["messages"]) == 1:
|
||
conversations[conv_id]["title"] = user_msg["content"][:20] + ("..." if len(user_msg["content"]) > 20 else "")
|
||
|
||
# persist after adding user message
|
||
save_session(conv_id)
|
||
|
||
# DEEPSEEK_API = getattr(settings, 'DEEPSEEK_API', os.getenv('DEEPSEEK_API', None))
|
||
# MODEL_ID = getattr(settings, 'DEEPSEEK_MODEL_ID', os.getenv('MODEL_ID', None))
|
||
|
||
# Use local RAG engine streaming API
|
||
async def generate_async():
|
||
full_response = ""
|
||
try:
|
||
if rag_engine is None:
|
||
yield "RAG engine not initialized"
|
||
return
|
||
|
||
chunk_counter = 0
|
||
check_interval = 5 # 每生成5个chunk检查一次
|
||
|
||
# Use rag_engine.query_stream which yields chunks asynchronously
|
||
try:
|
||
history = build_prompt_from_history(conv_id, max_messages=10)
|
||
async for chunk in rag_engine.query_stream(payload["content"], history):
|
||
# 定期检查active_stream状态
|
||
chunk_counter += 1
|
||
if chunk_counter % check_interval == 0:
|
||
with conv_lock:
|
||
if conv_id not in conversations or not conversations[conv_id].get("active_stream", True):
|
||
compat_logger.info(f"流被中止: {conv_id}")
|
||
return # 直接返回，中止生成
|
||
if chunk:
|
||
# normalize chunk to str
|
||
text_chunk = str(chunk.decode('utf-8')) if isinstance(chunk, bytes) else str(chunk)
|
||
# markdown_text = markdown2.markdown(text_chunk)
|
||
full_response += text_chunk
|
||
yield text_chunk.encode('utf-8') if isinstance(chunk, str) else text_chunk
|
||
except Exception as e:
|
||
compat_logger.error(f"RAG stream error: {e}", exc_info=True)
|
||
yield f"流式响应错误: {e}"
|
||
return
|
||
|
||
# 最终检查是否被中止
|
||
with conv_lock:
|
||
if conv_id not in conversations or not conversations[conv_id].get("active_stream", True):
|
||
compat_logger.info(f"流在完成前被中止: {conv_id}")
|
||
return
|
||
|
||
# Append final assistant message to conversation
|
||
ai_msg = {"role": "assistant", "content": full_response, "time": datetime.now().isoformat()}
|
||
with conv_lock:
|
||
if conv_id in conversations:
|
||
conversations[conv_id]["messages"].append(ai_msg)
|
||
conversations[conv_id]["active_stream"] = False
|
||
|
||
# persist after adding assistant message
|
||
save_session(conv_id)
|
||
|
||
finally:
|
||
with conv_lock:
|
||
if conv_id in conversations:
|
||
# 只在正常完成时才设为False，如果被中止，abort接口已经设为False了
|
||
if conversations[conv_id].get("active_stream", True):
|
||
conversations[conv_id]["active_stream"] = False
|
||
|
||
# return StreamingResponse(generate_async(), media_type='text/plain; charset=utf-8')
|
||
return StreamingResponse(
|
||
generate_async(),
|
||
media_type="text/event-stream",
|
||
headers={
|
||
"X-Accel-Buffering": "no",
|
||
"Cache-Control": "no-cache",
|
||
"Connection": "keep-alive"
|
||
}
|
||
)
|
||
|
||
|
||
@app.post('/conversations/{conv_id}/abort')
|
||
def abort_stream(conv_id: str):
|
||
with conv_lock:
|
||
if conv_id not in conversations:
|
||
raise HTTPException(status_code=404, detail="对话不存在")
|
||
if not conversations[conv_id].get("active_stream", False):
|
||
raise HTTPException(status_code=400, detail="该对话没有活动流")
|
||
conversations[conv_id]["active_stream"] = False
|
||
return {"message": "流式响应已中止"}
|
||
|
||
|
||
@app.post("/retrieve", response_model=RetrieveResponse)
|
||
async def retrieve(request: RetrieveRequest):
|
||
"""
|
||
Retrieve documents from ChromaDB based on query (no LLM generation)
|
||
|
||
This endpoint only performs vector similarity search and returns the retrieved
|
||
documents with their content, metadata, and similarity scores.
|
||
|
||
Args:
|
||
request: Retrieve request with query string and top_k option
|
||
|
||
Returns:
|
||
Retrieved documents with content, metadata, and scores
|
||
"""
|
||
if vector_store_manager is None:
|
||
raise HTTPException(status_code=503, detail="Vector store not initialized")
|
||
|
||
try:
|
||
|
||
# Get retriever (works even if collection is empty, will return empty results)
|
||
retriever = vector_store_manager.get_retriever(top_k=request.top_k or settings.TOP_K)
|
||
|
||
# Retrieve documents (run in thread pool since retriever.retrieve() is synchronous)
|
||
def retrieve_docs():
|
||
nodes = retriever.retrieve(request.query)
|
||
return nodes
|
||
|
||
nodes = await asyncio.to_thread(retrieve_docs)
|
||
|
||
# Convert nodes to response format
|
||
documents = []
|
||
for node in nodes:
|
||
# Extract similarity score if available
|
||
# LlamaIndex retriever returns NodeWithScore objects
|
||
score = None
|
||
if hasattr(node, 'score'):
|
||
score = node.score
|
||
elif hasattr(node, 'node') and hasattr(node.node, 'score'):
|
||
score = node.node.score
|
||
|
||
# Get the actual node object (NodeWithScore.node or the node itself)
|
||
actual_node = node.node if hasattr(node, 'node') else node
|
||
|
||
# Get content
|
||
content = actual_node.text if hasattr(actual_node, 'text') else str(actual_node)
|
||
|
||
# Get metadata
|
||
metadata = {}
|
||
if hasattr(actual_node, 'metadata'):
|
||
metadata = actual_node.metadata.copy() if actual_node.metadata else {}
|
||
|
||
# Get doc_id
|
||
doc_id = metadata.get('doc_id') or metadata.get('original_doc_id')
|
||
if not doc_id and hasattr(actual_node, 'node_id'):
|
||
doc_id = actual_node.node_id
|
||
|
||
documents.append(RetrievedDocument(
|
||
content=content,
|
||
score=score,
|
||
metadata=metadata,
|
||
doc_id=doc_id
|
||
))
|
||
|
||
return RetrieveResponse(
|
||
query=request.query,
|
||
documents=documents,
|
||
count=len(documents)
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"Error retrieving documents: {e}")
|
||
raise HTTPException(status_code=500, detail=f"Error retrieving documents: {str(e)}")
|
||
|
||
|
||
@app.delete("/documents/source/{source_name}")
|
||
async def delete_documents_by_source(source_name: str):
|
||
"""
|
||
Delete all documents from a specific data source
|
||
|
||
Args:
|
||
source_name: Name of the data source to delete documents from
|
||
|
||
Returns:
|
||
Deletion status
|
||
"""
|
||
if vector_store_manager is None:
|
||
raise HTTPException(status_code=503, detail="Vector store manager not initialized")
|
||
|
||
try:
|
||
# Delete documents by source name
|
||
vector_store_manager.delete_documents_by_source(source_name)
|
||
return {"status": "success", "message": f"Deleted all documents from {source_name}"}
|
||
except Exception as e:
|
||
logger.error(f"Error deleting documents by source: {e}")
|
||
raise HTTPException(status_code=500, detail=f"Error deleting documents: {str(e)}")
|
||
|
||
|
||
@app.post("/sync")
|
||
async def manual_sync(request: SyncRequest):
|
||
"""
|
||
Manually trigger synchronization
|
||
|
||
Args:
|
||
request: Sync request with options
|
||
|
||
Returns:
|
||
Sync status
|
||
"""
|
||
if sync_manager is None:
|
||
raise HTTPException(status_code=503, detail="Sync service not initialized")
|
||
|
||
try:
|
||
message = ""
|
||
if request.source_name:
|
||
specific_sync_service = sync_manager.get_sync_service(request.source_name)
|
||
if not specific_sync_service:
|
||
raise HTTPException(status_code=404, detail=f"Sync service not found for source: {request.source_name}")
|
||
|
||
if request.full_sync:
|
||
# Start auto sync service with recovery for this data source
|
||
asyncio.create_task(specific_sync_service.start_auto_sync_with_recovery(skip_initial_sync=False))
|
||
message = f"Started full synchronization for data source: {request.source_name}"
|
||
message += f" (Auto sync service started)"
|
||
else:
|
||
# Start auto sync service with recovery for this data source
|
||
asyncio.create_task(specific_sync_service.start_auto_sync_with_recovery(skip_initial_sync=True))
|
||
message = f"Started synchronization for data source: {request.source_name}"
|
||
message += f" (Auto sync service started)"
|
||
else:
|
||
# Sync all data sources
|
||
# This is a SyncServiceManager, it starts auto sync services which do initial sync
|
||
await sync_manager.start_all_sync_services()
|
||
message = "Started synchronization for all data sources"
|
||
|
||
return {"status": "success", "message": message}
|
||
except Exception as e:
|
||
logger.error(f"Error during sync: {e}")
|
||
raise HTTPException(status_code=500, detail=f"Sync failed: {str(e)}")
|
||
|
||
|
||
@app.get("/stats")
|
||
async def get_stats():
|
||
"""
|
||
Get system statistics
|
||
|
||
Returns:
|
||
System statistics
|
||
"""
|
||
if vector_store_manager is None:
|
||
raise HTTPException(status_code=503, detail="Vector store not initialized")
|
||
|
||
try:
|
||
collection = vector_store_manager.collection
|
||
count = collection.count()
|
||
|
||
return {
|
||
"vector_store": {
|
||
"collection_name": settings.CHROMA_COLLECTION_NAME,
|
||
"document_count": count
|
||
},
|
||
"config": {
|
||
"model": settings.OLLAMA_MODEL,
|
||
"top_k": settings.TOP_K,
|
||
"chunk_size": settings.CHUNK_SIZE
|
||
}
|
||
}
|
||
except Exception as e:
|
||
logger.error(f"Error getting stats: {e}")
|
||
raise HTTPException(status_code=500, detail=f"Error getting stats: {str(e)}")
|
||
|
||
|
||
@app.post("/documents/upload", response_model=UploadResponse)
|
||
async def upload_document(
|
||
file: UploadFile = File(..., description="File to upload"),
|
||
doc_id: Optional[str] = Form(None, description="Custom document ID (optional, will use filename if not provided)"),
|
||
metadata: Optional[str] = Form(None, description="Additional metadata as JSON string (optional)")
|
||
):
|
||
"""
|
||
Upload and parse a document file
|
||
|
||
Supports various file formats:
|
||
- Text files: .txt, .md, .markdown
|
||
- PDF files: .pdf
|
||
- Word documents: .docx, .doc
|
||
- HTML files: .html, .htm
|
||
- CSV files: .csv
|
||
- JSON files: .json
|
||
|
||
File size limit: Configurable via MAX_UPLOAD_SIZE_MB in .env (default: 10MB)
|
||
|
||
Args:
|
||
file: File to upload (max size configured in MAX_UPLOAD_SIZE_MB)
|
||
doc_id: Optional custom document ID
|
||
metadata: Optional JSON string with additional metadata
|
||
|
||
Returns:
|
||
Upload response with document ID and chunk count
|
||
|
||
Raises:
|
||
HTTPException: 400 if file format is unsupported
|
||
HTTPException: 413 if file size exceeds 10MB
|
||
HTTPException: 400 if document parsing fails
|
||
"""
|
||
if file_parser is None or vector_store_manager is None:
|
||
raise HTTPException(status_code=503, detail="Services not initialized")
|
||
|
||
try:
|
||
|
||
# Check if file format is supported
|
||
if not file_parser.is_supported(file.filename):
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Unsupported file format. Supported formats: {', '.join(file_parser.SUPPORTED_EXTENSIONS)}"
|
||
)
|
||
|
||
# Read file content
|
||
content = await file.read()
|
||
|
||
# Check file size limit (from config, default: 10MB)
|
||
MAX_FILE_SIZE = settings.MAX_UPLOAD_SIZE_MB * 1024 * 1024 # Convert MB to bytes
|
||
file_size = len(content)
|
||
if file_size > MAX_FILE_SIZE:
|
||
raise HTTPException(
|
||
status_code=413,
|
||
detail=f"File size ({file_size / (1024 * 1024):.2f}MB) exceeds maximum allowed size ({settings.MAX_UPLOAD_SIZE_MB}MB)"
|
||
)
|
||
|
||
# Parse metadata if provided
|
||
extra_metadata = {}
|
||
if metadata:
|
||
try:
|
||
extra_metadata = json.loads(metadata)
|
||
except json.JSONDecodeError:
|
||
raise HTTPException(status_code=400, detail="Invalid metadata JSON format")
|
||
|
||
# Use custom doc_id or generate from filename
|
||
if not doc_id:
|
||
doc_id = Path(file.filename).stem
|
||
|
||
# Parse file (run in thread pool to avoid blocking event loop for large files)
|
||
logger.info(f"Parsing uploaded file: {file.filename} (doc_id: {doc_id}, size: {file_size / 1024:.2f}KB)")
|
||
try:
|
||
# 获取事件循环，用于在异步环境中运行同步函数
|
||
loop = asyncio.get_event_loop()
|
||
# 使用线程池执行器运行文件解析函数（同步函数），避免阻塞事件循环
|
||
documents = await loop.run_in_executor(
|
||
None, # 使用默认线程池
|
||
file_parser.parse_file_content, # 要执行的同步解析函数
|
||
content, # 文件内容（字节流）
|
||
file.filename, # 文件名
|
||
doc_id, # 文档ID
|
||
extra_metadata# 额外元数据
|
||
)
|
||
except Exception as parse_error:
|
||
logger.error(f"Error parsing file {file.filename}: {parse_error}", exc_info=True)
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Failed to parse document: {str(parse_error)}"
|
||
)
|
||
|
||
if not documents:
|
||
raise HTTPException(status_code=400, detail="Failed to parse document or document is empty")
|
||
|
||
logger.info(f"Successfully parsed {file.filename}: {len(documents)} document(s)")
|
||
|
||
# Chunk documents if needed
|
||
logger.info(f"Chunking {len(documents)} document(s)...")
|
||
chunked_documents = BaseSync.chunk_documents(documents)
|
||
logger.info(f"Chunked into {len(chunked_documents)} chunk(s)")
|
||
|
||
# Add to vector store (this may take time for large documents due to embedding generation)
|
||
logger.info(f"Adding {len(chunked_documents)} chunks to vector store (this may take time for large files)...")
|
||
try:
|
||
# Run in thread pool to avoid blocking event loop during embedding generation
|
||
loop = asyncio.get_event_loop()
|
||
await loop.run_in_executor(
|
||
None,
|
||
vector_store_manager.add_documents,
|
||
chunked_documents,
|
||
False # skip_existing
|
||
)
|
||
logger.info(f"Successfully added {len(chunked_documents)} chunks to vector store")
|
||
except Exception as add_error:
|
||
logger.error(f"Error adding documents to vector store: {add_error}", exc_info=True)
|
||
raise HTTPException(
|
||
status_code=500,
|
||
detail=f"Failed to add documents to vector store: {str(add_error)}"
|
||
)
|
||
|
||
file_type = Path(file.filename).suffix.lstrip('.')
|
||
|
||
return UploadResponse(
|
||
doc_id=doc_id,
|
||
filename=file.filename,
|
||
file_type=file_type,
|
||
chunks=len(chunked_documents),
|
||
message=f"Document uploaded and processed successfully"
|
||
)
|
||
|
||
except HTTPException:
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"Error uploading document: {e}")
|
||
raise HTTPException(status_code=500, detail=f"Error uploading document: {str(e)}")
|
||
|
||
|
||
@app.get("/documents/{doc_id}", response_model=DocumentResponse)
|
||
async def get_document_by_id(doc_id: str):
|
||
"""
|
||
Get all content for a document by its ID
|
||
|
||
This endpoint retrieves all chunks/content that belong to a document,
|
||
even if the document was split into multiple chunks during processing.
|
||
|
||
Args:
|
||
doc_id: Document ID to retrieve
|
||
|
||
Returns:
|
||
Document response with all chunks and full text
|
||
"""
|
||
if vector_store_manager is None:
|
||
raise HTTPException(status_code=503, detail="Vector store not initialized")
|
||
|
||
try:
|
||
# Get all chunks for this document
|
||
chunks = vector_store_manager.get_document_by_id(doc_id)
|
||
|
||
if not chunks:
|
||
raise HTTPException(status_code=404, detail=f"Document with ID '{doc_id}' not found")
|
||
|
||
# Convert to response format
|
||
document_chunks = []
|
||
full_text_parts = []
|
||
|
||
for chunk in chunks:
|
||
document_chunks.append(DocumentChunk(
|
||
id=chunk['id'],
|
||
text=chunk['text'],
|
||
metadata=chunk.get('metadata', {}),
|
||
chunk_index=chunk.get('chunk_index')
|
||
))
|
||
full_text_parts.append(chunk['text'])
|
||
|
||
# Combine all chunks into full text
|
||
full_text = "\n\n".join(full_text_parts)
|
||
|
||
return DocumentResponse(
|
||
doc_id=doc_id,
|
||
chunks=document_chunks,
|
||
total_chunks=len(document_chunks),
|
||
full_text=full_text
|
||
)
|
||
|
||
except HTTPException:
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"Error retrieving document {doc_id}: {e}")
|
||
raise HTTPException(status_code=500, detail=f"Error retrieving document: {str(e)}")
|
||
|
||
|
||
# Configuration management endpoints
|
||
from fastapi import HTTPException
|
||
import os
|
||
import json
|
||
from typing import List, Dict, Any
|
||
|
||
# Configuration management endpoints
|
||
@app.get("/folder-configs")
|
||
async def get_folder_configs():
|
||
"""
|
||
Get all folder configurations
|
||
|
||
Returns:
|
||
List of all configurations with their IDs and details
|
||
"""
|
||
try:
|
||
# Load configs from SQLite database
|
||
DATA_DIR = Path(__file__).parent.parent / "data"
|
||
DATA_DIR.mkdir(parents=True, exist_ok=True)
|
||
DB_PATH = DATA_DIR / "sessions.db"
|
||
|
||
conn = sqlite3.connect(DB_PATH)
|
||
cursor = conn.cursor()
|
||
|
||
# 查询所有数据源配置
|
||
try:
|
||
cursor.execute('SELECT name, config, update_at FROM data_sources')
|
||
rows = cursor.fetchall()
|
||
|
||
configurations = []
|
||
for row in rows:
|
||
name, config_json, update_at = row
|
||
try:
|
||
config_data = json.loads(config_json)
|
||
configurations.append({
|
||
"id": name,
|
||
"type": config_data.get("type", "unknown"),
|
||
"config": config_data,
|
||
"update_at": update_at
|
||
})
|
||
except Exception as e:
|
||
logger.error(f"Error parsing config from database: {name}, error: {e}")
|
||
except sqlite3.OperationalError as e:
|
||
# 表不存在的情况，返回空列表
|
||
logger.warning(f"SQLite table error: {e}. Returning empty config list.")
|
||
configurations = []
|
||
|
||
# 关闭数据库连接
|
||
conn.close()
|
||
|
||
return {"configurations": configurations}
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error getting folder configs: {e}")
|
||
# 返回空列表而不是错误响应
|
||
return {"configurations": []}
|
||
|
||
|
||
@app.post("/folder-configs")
|
||
async def create_config(config: Dict[str, Any]):
|
||
"""
|
||
Create a new configuration
|
||
|
||
Args:
|
||
config: Configuration data for the new config
|
||
|
||
Returns:
|
||
Created configuration with ID
|
||
"""
|
||
conn = None
|
||
try:
|
||
# Ensure config has required fields
|
||
if "type" not in config:
|
||
raise HTTPException(status_code=400, detail="Config type is required")
|
||
|
||
config_type = config["type"]
|
||
|
||
# Check for unique identifier based on config type
|
||
# 必须有数据源唯一标识才能创建配置
|
||
unique_id_parts = [config_type]
|
||
|
||
if config_type == "database":
|
||
# 数据库配置需要：主机、端口、数据库名、表名
|
||
if not config.get("mysql_host") or not config.get("mysql_port") or not config.get("database") or not config.get("table_name"):
|
||
raise HTTPException(status_code=400, detail="数据库配置必须包含主机、端口、数据库名和表名")
|
||
# 添加数据库类型到唯一标识符
|
||
db_type = config.get("db_type", "mysql").lower()
|
||
unique_id_parts.extend([
|
||
db_type,
|
||
config["mysql_host"].lower(),
|
||
str(config["mysql_port"]),
|
||
config["database"].lower(),
|
||
config["table_name"].lower()
|
||
])
|
||
elif config_type == "folder":
|
||
# 文件夹配置需要：主机、文件夹路径
|
||
if not config.get("host") or not config.get("folder_path"):
|
||
raise HTTPException(status_code=400, detail="文件夹配置必须包含主机和文件夹路径")
|
||
# 替换路径中的特殊字符为下划线
|
||
folder_path = config["folder_path"].lower().replace("/", "_").replace(":", "_").replace("\\", "_")
|
||
# 去除开头的下划线
|
||
if folder_path.startswith("_"):
|
||
folder_path = folder_path[1:]
|
||
unique_id_parts.extend([
|
||
config["host"].lower(),
|
||
folder_path
|
||
])
|
||
elif config_type == "git":
|
||
# Git配置需要：仓库URL
|
||
if not config.get("git_url"):
|
||
raise HTTPException(status_code=400, detail="Git配置必须包含仓库URL")
|
||
# 添加Git仓库URL到唯一标识符
|
||
# 替换URL中的特殊字符为下划线
|
||
git_url = config["git_url"].lower().replace("/", "_").replace(":", "_").replace(".", "_")
|
||
# 截取URL的一部分作为唯一标识
|
||
git_url_part = git_url[:100] # 限制长度
|
||
unique_id_parts.extend([
|
||
git_url_part
|
||
])
|
||
else:
|
||
raise HTTPException(status_code=400, detail=f"不支持的配置类型: {config_type}")
|
||
|
||
# Generate config ID based on unique identifier
|
||
config_id = "_".join(unique_id_parts)
|
||
|
||
# Save config to SQLite database
|
||
conn, cursor = get_db_connection()
|
||
|
||
# 检查是否存在同源配置
|
||
try:
|
||
# 先尝试查询现有配置
|
||
cursor.execute('SELECT name, config FROM data_sources')
|
||
existing_configs = cursor.fetchall()
|
||
|
||
for existing_name, existing_config_json in existing_configs:
|
||
existing_config_data = json.loads(existing_config_json)
|
||
|
||
# Check if it's the same type
|
||
if existing_config_data.get('type') == config_type:
|
||
if config_type == 'database':
|
||
# For database configs, same source means same db type, host, port, and db name
|
||
if (existing_config_data.get('db_type', 'mysql') == config.get('db_type', 'mysql') and
|
||
existing_config_data.get('mysql_host') == config.get('mysql_host') and
|
||
existing_config_data.get('mysql_port') == config.get('mysql_port') and
|
||
existing_config_data.get('database') == config.get('database')):
|
||
raise HTTPException(
|
||
status_code=409,
|
||
detail=f"已存在相同源的数据库配置。如需调整，请点击配置列表中的配置并修改配置内容。"
|
||
)
|
||
elif config_type == 'folder':
|
||
# For folder configs, same source means same host and folder path
|
||
if (existing_config_data.get('host') == config.get('host') and
|
||
existing_config_data.get('folder_path') == config.get('folder_path')):
|
||
raise HTTPException(
|
||
status_code=409,
|
||
detail=f"已存在相同服务器和路径的文件夹配置。如需调整，请点击配置列表中的配置并修改配置内容。"
|
||
)
|
||
elif config_type == 'git':
|
||
# For git configs, same source means same git url
|
||
if existing_config_data.get('git_url') == config.get('git_url'):
|
||
raise HTTPException(
|
||
status_code=409,
|
||
detail=f"已存在相同Git仓库的配置。如需调整，请点击配置列表中的配置并修改配置内容。"
|
||
)
|
||
except sqlite3.OperationalError as e:
|
||
# 表不存在的情况，会在后面创建表
|
||
logger.warning(f"SQLite table error: {e}. This is expected if the table doesn't exist yet.")
|
||
|
||
|
||
|
||
# Check if a config with the same name already exists
|
||
cursor.execute('SELECT name FROM data_sources WHERE name = ?', (config_id,))
|
||
existing_config = cursor.fetchone()
|
||
|
||
if existing_config:
|
||
raise HTTPException(
|
||
status_code=409,
|
||
detail=f"配置 '{config['name']}' 已存在。如需调整，请点击配置列表中的配置并修改配置内容。"
|
||
)
|
||
|
||
# Insert new config into data_sources table
|
||
config_json = json.dumps(config, ensure_ascii=False, indent=2)
|
||
|
||
cursor.execute(
|
||
'INSERT INTO data_sources (name, config) VALUES (?, ?)',
|
||
(config_id, config_json)
|
||
)
|
||
|
||
conn.commit()
|
||
|
||
# Create sync service for this new data source
|
||
global sync_manager
|
||
if sync_manager is not None:
|
||
# Create appropriate data source config object
|
||
from config import BaseDataSourceConfig, DatabaseDataSourceConfig, FolderDataSourceConfig, GitDataSourceConfig
|
||
|
||
if config_type == "database":
|
||
source_config = DatabaseDataSourceConfig(
|
||
name=config_id,
|
||
database=config.get("database"),
|
||
table_name=config.get("table_name"),
|
||
id_column=config.get("id_column"),
|
||
content_column=config.get("content_column"),
|
||
file_column=config.get("file_column"),
|
||
title_column=config.get("title_column"),
|
||
metadata_columns=config.get("metadata_columns"),
|
||
content_separator=config.get("content_separator"),
|
||
updated_at_column=config.get("updated_at_column"),
|
||
host=config.get("mysql_host"),
|
||
port=config.get("mysql_port"),
|
||
user=config.get("mysql_user"),
|
||
password=config.get("mysql_password"),
|
||
db_type=config.get("db_type", "mysql"),
|
||
file_source_type=config.get("file_source_type"),
|
||
file_system_base_path=config.get("file_system_base_path"),
|
||
scp_host=config.get("scp_host"),
|
||
scp_port=config.get("scp_port"),
|
||
scp_username=config.get("scp_username"),
|
||
scp_password=config.get("scp_password")
|
||
)
|
||
elif config_type == "folder":
|
||
source_config = FolderDataSourceConfig(
|
||
name=config_id,
|
||
folder_path=config.get("folder_path"),
|
||
host=config.get("host"),
|
||
port=config.get("port"),
|
||
username=config.get("username"),
|
||
password=config.get("password"),
|
||
recursive=config.get("recursive", True),
|
||
ignore_patterns=config.get("ignore_patterns")
|
||
)
|
||
elif config_type == "git":
|
||
source_config = GitDataSourceConfig(
|
||
name=config_id,
|
||
git_url=config.get("git_url"),
|
||
branch=config.get("branch", "main"),
|
||
protocol=config.get("protocol", "https"),
|
||
ssh_key=config.get("ssh_key"),
|
||
https_token=config.get("https_token"),
|
||
local_repo_path=config.get("local_repo_path"),
|
||
poll_interval=config.get("poll_interval", 300),
|
||
support_lang=config.get("support_lang"),
|
||
latest_commit_id=config.get("latest_commit_id"),
|
||
last_sync_time=config.get("last_sync_time")
|
||
)
|
||
else:
|
||
logger.warning(f"Unknown config type: {config_type}")
|
||
# Create a base config as fallback
|
||
source_config = BaseDataSourceConfig(name=config_id, type=config_type)
|
||
|
||
# Create the sync service
|
||
sync_manager.create_or_update_sync_service(source_config)
|
||
|
||
# Return the created config
|
||
return {
|
||
"id": config_id,
|
||
"type": config_type,
|
||
"config": config
|
||
}
|
||
|
||
except HTTPException:
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"Error creating config: {e}")
|
||
raise HTTPException(status_code=500, detail=f"Error creating configuration: {str(e)}")
|
||
finally:
|
||
if conn:
|
||
conn.close()
|
||
|
||
@app.post("/folder/test-connection")
|
||
async def test_folder_connection(connection_data: Dict[str, Any]):
|
||
"""
|
||
Test SSH connection for folder configuration
|
||
|
||
Args:
|
||
connection_data: Connection data including host, port, username, password
|
||
|
||
Returns:
|
||
Success message if connection is successful
|
||
"""
|
||
try:
|
||
host = connection_data.get("host")
|
||
port = connection_data.get("port", 22)
|
||
username = connection_data.get("username")
|
||
password = connection_data.get("password")
|
||
|
||
if not host or not username:
|
||
raise HTTPException(status_code=400, detail="主机地址和用户名是必填项")
|
||
|
||
# Import SSHClient here to avoid circular imports
|
||
from sync.folder_sync import SSHClient
|
||
|
||
# Test SSH connection
|
||
ssh_client = SSHClient(
|
||
host=host,
|
||
port=port,
|
||
username=username,
|
||
password=password
|
||
)
|
||
|
||
# Try to connect
|
||
success = ssh_client.test_connection()
|
||
|
||
if success:
|
||
return {"message": "SSH连接成功！"}
|
||
else:
|
||
raise HTTPException(status_code=500, detail="SSH连接失败")
|
||
|
||
except HTTPException:
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"Error testing SSH connection: {e}")
|
||
raise HTTPException(status_code=500, detail=f"SSH连接失败: {str(e)}")
|
||
|
||
|
||
@app.post("/git/test-connection")
|
||
async def test_git_connection(connection_data: Dict[str, Any]):
|
||
"""
|
||
Test Git connection for Git repository configuration
|
||
|
||
Args:
|
||
connection_data: Connection data including git_url, protocol, branch, https_token, ssh_key
|
||
|
||
Returns:
|
||
Success message if connection is successful
|
||
"""
|
||
try:
|
||
git_url = connection_data.get("git_url")
|
||
protocol = connection_data.get("protocol", "https")
|
||
branch = connection_data.get("branch", "main")
|
||
https_token = connection_data.get("https_token")
|
||
ssh_key = connection_data.get("ssh_key")
|
||
|
||
if not git_url:
|
||
raise HTTPException(status_code=400, detail="Git仓库URL是必填项")
|
||
|
||
# Import GitTool here to avoid circular imports
|
||
from utils.git_tool import GitTool
|
||
|
||
# Create a temporary GitTool instance to test connection
|
||
git_tool = GitTool(
|
||
git_url=git_url,
|
||
branch=branch,
|
||
protocol=protocol,
|
||
https_token=https_token,
|
||
ssh_key=ssh_key,
|
||
local_repo_path=None # 测试连接不需要本地路径
|
||
)
|
||
|
||
# Try to test connection
|
||
success = git_tool.test_connection()
|
||
|
||
if success:
|
||
return {"message": "Git连接成功！"}
|
||
else:
|
||
raise HTTPException(status_code=500, detail="Git连接失败")
|
||
|
||
except HTTPException:
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"Error testing Git connection: {e}")
|
||
raise HTTPException(status_code=500, detail=f"Git连接失败: {str(e)}")
|
||
|
||
|
||
@app.post("/folder-configs/remote")
|
||
async def create_remote_folder_config(config: Dict[str, Any]):
|
||
"""
|
||
Create a new remote folder configuration (deprecated, use /folder-configs instead)
|
||
|
||
Args:
|
||
config: Configuration data for the new remote folder
|
||
|
||
Returns:
|
||
Created configuration with ID
|
||
"""
|
||
# Set type to folder if not provided
|
||
if "type" not in config:
|
||
config["type"] = "folder"
|
||
|
||
# Call the generic create_config function
|
||
return await create_config(config)
|
||
|
||
|
||
@app.put("/folder-configs/{config_id}")
|
||
async def update_folder_config(config_id: str, config: Dict[str, Any]):
|
||
"""
|
||
Update an existing folder configuration
|
||
|
||
Args:
|
||
config_id: ID of the configuration to update
|
||
config: Updated configuration data
|
||
|
||
Returns:
|
||
Updated configuration
|
||
"""
|
||
conn = None
|
||
try:
|
||
conn, cursor = get_db_connection()
|
||
|
||
# Check if config exists
|
||
cursor.execute('SELECT * FROM data_sources WHERE name = ?', (config_id,))
|
||
existing_config = cursor.fetchone()
|
||
|
||
if not existing_config:
|
||
raise HTTPException(status_code=404, detail=f"Configuration with ID '{config_id}' not found")
|
||
|
||
# Ensure config has required fields
|
||
if "type" not in config:
|
||
raise HTTPException(status_code=400, detail="Configuration type is required")
|
||
|
||
# Generate new config ID based on the complete configuration information
|
||
config_type = config["type"]
|
||
new_config_id = None
|
||
|
||
# 根据不同类型的配置生成有意义的ID
|
||
if config_type == "database":
|
||
# 数据库配置：使用数据库类型、数据库名和表名生成ID
|
||
if not config.get("database"):
|
||
raise HTTPException(status_code=400, detail="Database name is required for database configuration")
|
||
if not config.get("table_name"):
|
||
raise HTTPException(status_code=400, detail="Table name is required for database configuration")
|
||
db_type = config.get("db_type", "mysql").lower()
|
||
new_config_id = f"{config_type}_{db_type}_{config['database'].lower()}_{config['table_name'].lower()}"
|
||
elif config_type == "folder":
|
||
# 文件夹配置：根据是否有host字段区分本地和远程
|
||
if not config.get("folder_path"):
|
||
raise HTTPException(status_code=400, detail="Folder path is required for folder configuration")
|
||
# 替换路径中的特殊字符为下划线
|
||
folder_path = config["folder_path"].lower().replace("/", "_").replace(":", "_").replace("\\", "_")
|
||
# 去除开头的下划线
|
||
if folder_path.startswith("_"):
|
||
folder_path = folder_path[1:]
|
||
# 根据是否有host字段生成不同的ID
|
||
if config.get("host") and config["host"] != "localhost":
|
||
# 远程文件夹：使用主机和文件夹路径生成ID
|
||
new_config_id = f"{config_type}_{config['host'].lower()}_{folder_path}"
|
||
else:
|
||
# 本地文件夹：使用文件夹路径生成ID
|
||
new_config_id = f"{config_type}_{folder_path}"
|
||
|
||
# 如果无法生成新的有意义的ID，保留原来的ID
|
||
if not new_config_id:
|
||
new_config_id = config_id
|
||
|
||
# Update config in data_sources table
|
||
# Since name is primary key, we need to use INSERT OR REPLACE
|
||
config_json = json.dumps(config, ensure_ascii=False, indent=2)
|
||
|
||
# First, delete the old config if new ID is different
|
||
if new_config_id != config_id:
|
||
cursor.execute('DELETE FROM data_sources WHERE name = ?', (config_id,))
|
||
|
||
# Then insert the updated config with new ID, update_at value to null
|
||
cursor.execute(
|
||
'INSERT OR REPLACE INTO data_sources (name, config, update_at) VALUES (?, ?, ?)',
|
||
(new_config_id, config_json, None)
|
||
)
|
||
|
||
conn.commit()
|
||
|
||
# Update the sync service for this data source
|
||
global sync_manager
|
||
if sync_manager is not None:
|
||
# Create appropriate data source config object
|
||
from config import BaseDataSourceConfig, DatabaseDataSourceConfig, FolderDataSourceConfig
|
||
|
||
if config_type == "database":
|
||
source_config = DatabaseDataSourceConfig(
|
||
name=new_config_id,
|
||
database=config.get("database"),
|
||
table_name=config.get("table_name"),
|
||
id_column=config.get("id_column"),
|
||
content_column=config.get("content_column"),
|
||
file_column=config.get("file_column"),
|
||
title_column=config.get("title_column"),
|
||
metadata_columns=config.get("metadata_columns"),
|
||
content_separator=config.get("content_separator"),
|
||
updated_at_column=config.get("updated_at_column"),
|
||
host=config.get("mysql_host"),
|
||
port=config.get("mysql_port"),
|
||
user=config.get("mysql_user"),
|
||
password=config.get("mysql_password"),
|
||
db_type=config.get("db_type", "mysql"),
|
||
file_source_type=config.get("file_source_type"),
|
||
file_system_base_path=config.get("file_system_base_path"),
|
||
scp_host=config.get("scp_host"),
|
||
scp_port=config.get("scp_port"),
|
||
scp_username=config.get("scp_username"),
|
||
scp_password=config.get("scp_password")
|
||
)
|
||
elif config_type == "folder":
|
||
source_config = FolderDataSourceConfig(
|
||
name=new_config_id,
|
||
folder_path=config.get("folder_path"),
|
||
host=config.get("host"),
|
||
port=config.get("port"),
|
||
username=config.get("username"),
|
||
password=config.get("password"),
|
||
recursive=config.get("recursive", True),
|
||
ignore_patterns=config.get("ignore_patterns")
|
||
)
|
||
else:
|
||
logger.warning(f"Unknown config type: {config_type}")
|
||
# Create a base config as fallback
|
||
source_config = BaseDataSourceConfig(name=new_config_id, type=config_type)
|
||
|
||
# Update the sync service
|
||
sync_manager.create_or_update_sync_service(source_config)
|
||
|
||
# Return the updated config with new ID
|
||
return {
|
||
"id": new_config_id,
|
||
"type": config_type,
|
||
"config": config
|
||
}
|
||
|
||
except HTTPException:
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"Error updating folder config: {e}")
|
||
raise HTTPException(status_code=500, detail=f"Error updating configuration: {str(e)}")
|
||
finally:
|
||
if conn:
|
||
conn.close()
|
||
|
||
|
||
@app.delete("/folder-configs/{config_id}")
|
||
async def delete_folder_config(config_id: str):
|
||
"""
|
||
Delete a folder configuration
|
||
|
||
Args:
|
||
config_id: ID of the configuration to delete
|
||
|
||
Returns:
|
||
Success message
|
||
"""
|
||
conn = None
|
||
try:
|
||
conn, cursor = get_db_connection()
|
||
|
||
# Check if config exists
|
||
cursor.execute('SELECT * FROM data_sources WHERE name = ?', (config_id,))
|
||
existing_config = cursor.fetchone()
|
||
|
||
if not existing_config:
|
||
raise HTTPException(status_code=404, detail=f"Configuration with ID '{config_id}' not found")
|
||
|
||
# Delete config from data_sources table
|
||
cursor.execute('DELETE FROM data_sources WHERE name = ?', (config_id,))
|
||
|
||
conn.commit()
|
||
|
||
# Remove the sync service for this data source
|
||
global sync_manager
|
||
if sync_manager is not None:
|
||
sync_manager.remove_sync_service(config_id)
|
||
|
||
return {"status": "success", "message": f"Configuration '{config_id}' deleted successfully"}
|
||
|
||
except HTTPException:
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"Error deleting folder config: {e}")
|
||
raise HTTPException(status_code=500, detail=f"Error deleting configuration: {str(e)}")
|
||
finally:
|
||
if conn:
|
||
conn.close()
|
||
|
||
|
||
# Database schema exploration endpoints
|
||
from pydantic import BaseModel
|
||
import mysql.connector
|
||
|
||
|
||
|
||
class DatabaseConnectionParams(BaseModel):
|
||
"""Database connection parameters"""
|
||
host: str = Field(..., description="Database host")
|
||
port: int = Field(default=3306, description="Database port")
|
||
username: str = Field(..., description="Database username")
|
||
password: str = Field(..., description="Database password")
|
||
db_type: str = Field(default="mysql", description="Database type: mysql or dameng")
|
||
|
||
class DatabaseParams(BaseModel):
|
||
"""Database parameters"""
|
||
host: str = Field(..., description="Database host")
|
||
port: int = Field(default=3306, description="Database port")
|
||
username: str = Field(..., description="Database username")
|
||
password: str = Field(..., description="Database password")
|
||
database: str = Field(..., description="Database name")
|
||
db_type: str = Field(default="mysql", description="Database type: mysql or dameng")
|
||
|
||
class TableParams(BaseModel):
|
||
"""Table parameters"""
|
||
host: str = Field(..., description="Database host")
|
||
port: int = Field(default=3306, description="Database port")
|
||
username: str = Field(..., description="Database username")
|
||
password: str = Field(..., description="Database password")
|
||
database: str = Field(..., description="Database name")
|
||
table_name: str = Field(..., description="Table name")
|
||
db_type: str = Field(default="mysql", description="Database type: mysql or dameng")
|
||
|
||
|
||
|
||
|
||
@app.post("/database/databases")
|
||
async def get_databases(params: DatabaseConnectionParams):
|
||
"""
|
||
Get list of databases from database server
|
||
|
||
Args:
|
||
params: Database connection parameters
|
||
|
||
Returns:
|
||
List of available databases
|
||
"""
|
||
try:
|
||
if params.db_type == "dameng":
|
||
# DaMeng database connection
|
||
import dmPython
|
||
connection = dmPython.connect(
|
||
user=params.username,
|
||
password=params.password,
|
||
server=params.host,
|
||
port=params.port
|
||
)
|
||
|
||
cursor = connection.cursor()
|
||
# Get schemas for DaMeng (similar to databases in MySQL)
|
||
cursor.execute("SELECT NAME AS SCHEMA_NAME FROM SYSOBJECTS WHERE TYPE$ = 'SCH';")
|
||
|
||
databases = []
|
||
for (schema_name,) in cursor:
|
||
databases.append(schema_name)
|
||
|
||
cursor.close()
|
||
connection.close()
|
||
|
||
else:
|
||
# MySQL database connection
|
||
connection = mysql.connector.connect(
|
||
host=params.host,
|
||
port=params.port,
|
||
user=params.username,
|
||
password=params.password
|
||
)
|
||
|
||
cursor = connection.cursor()
|
||
cursor.execute("SHOW DATABASES")
|
||
|
||
databases = []
|
||
for (database_name,) in cursor:
|
||
# Skip system databases
|
||
if database_name not in ['information_schema', 'mysql', 'performance_schema', 'sys']:
|
||
databases.append(database_name)
|
||
|
||
cursor.close()
|
||
connection.close()
|
||
|
||
return {"databases": databases}
|
||
|
||
except Exception as e:
|
||
logger.error(f"Database connection error: {e}")
|
||
raise HTTPException(status_code=500, detail=f"Database connection failed: {str(e)}")
|
||
|
||
|
||
@app.post("/database/tables")
|
||
async def get_tables(params: DatabaseParams):
|
||
"""
|
||
Get list of tables from specified database
|
||
|
||
Args:
|
||
params: Database parameters including database name
|
||
|
||
Returns:
|
||
List of tables in the specified database
|
||
"""
|
||
try:
|
||
if params.db_type == "dameng":
|
||
# DaMeng database connection
|
||
import dmPython
|
||
connection = dmPython.connect(
|
||
user=params.username,
|
||
password=params.password,
|
||
server=params.host,
|
||
port=params.port
|
||
)
|
||
|
||
cursor = connection.cursor()
|
||
# Get tables for DaMeng
|
||
cursor.execute(f"SELECT TABLE_NAME FROM ALL_TABLES WHERE OWNER = '{params.database.upper()}'")
|
||
|
||
tables = []
|
||
for (table_name,) in cursor:
|
||
tables.append(table_name)
|
||
|
||
cursor.close()
|
||
connection.close()
|
||
|
||
else:
|
||
# MySQL database connection
|
||
connection = mysql.connector.connect(
|
||
host=params.host,
|
||
port=params.port,
|
||
user=params.username,
|
||
password=params.password,
|
||
database=params.database
|
||
)
|
||
|
||
cursor = connection.cursor()
|
||
cursor.execute("SHOW TABLES")
|
||
|
||
tables = []
|
||
for (table_name,) in cursor:
|
||
tables.append(table_name)
|
||
|
||
cursor.close()
|
||
connection.close()
|
||
|
||
return {"tables": tables}
|
||
|
||
except Exception as e:
|
||
logger.error(f"Database connection error: {e}")
|
||
raise HTTPException(status_code=500, detail=f"Database connection failed: {str(e)}")
|
||
|
||
|
||
@app.post("/database/table-structure")
|
||
async def get_table_structure(params: TableParams):
|
||
"""
|
||
Get structure of specified table
|
||
|
||
Args:
|
||
params: Table parameters including table name
|
||
|
||
Returns:
|
||
Table structure with column details
|
||
"""
|
||
try:
|
||
if params.db_type == "dameng":
|
||
# DaMeng database connection
|
||
import dmPython
|
||
connection = dmPython.connect(
|
||
user=params.username,
|
||
password=params.password,
|
||
server=params.host,
|
||
port=params.port
|
||
)
|
||
|
||
cursor = connection.cursor()
|
||
# Get table structure for DaMeng
|
||
cursor.execute(f"SELECT COLUMN_NAME, DATA_TYPE, NULLABLE, COLUMN_ID FROM ALL_TAB_COLUMNS WHERE OWNER = '{params.database.upper()}' AND TABLE_NAME = '{params.table_name.upper()}' ORDER BY COLUMN_ID")
|
||
|
||
columns = []
|
||
for (column_name, data_type, nullable, column_id) in cursor:
|
||
columns.append({
|
||
"name": column_name,
|
||
"type": data_type,
|
||
"null": "YES" if nullable == 'Y' else "NO",
|
||
"key": "", # DaMeng doesn't return key info in this query
|
||
"default": None,
|
||
"extra": ""
|
||
})
|
||
|
||
cursor.close()
|
||
connection.close()
|
||
|
||
else:
|
||
# MySQL database connection
|
||
connection = mysql.connector.connect(
|
||
host=params.host,
|
||
port=params.port,
|
||
user=params.username,
|
||
password=params.password,
|
||
database=params.database
|
||
)
|
||
|
||
cursor = connection.cursor()
|
||
cursor.execute(f"DESCRIBE {params.table_name}")
|
||
|
||
columns = []
|
||
for (field, type, null, key, default, extra) in cursor:
|
||
columns.append({
|
||
"name": field,
|
||
"type": type,
|
||
"null": null,
|
||
"key": key,
|
||
"default": default,
|
||
"extra": extra
|
||
})
|
||
|
||
cursor.close()
|
||
connection.close()
|
||
|
||
return {"columns": columns}
|
||
|
||
except Exception as e:
|
||
logger.error(f"Database connection error: {e}")
|
||
raise HTTPException(status_code=500, detail=f"Database connection failed: {str(e)}")
|
||
|
||
|
||
|
||
|