""" Engimind Python Backend — FastAPI application All API routes matching the original Go/Wails backend functionality. """ from __future__ import annotations import asyncio import base64 import logging import os import shutil import time from contextlib import asynccontextmanager from pathlib import Path from typing import Optional from fastapi import FastAPI, HTTPException, UploadFile, File, Form from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse from pydantic import BaseModel from sse_starlette.sse import EventSourceResponse from database import db from models import ( LLMProvider, VectorDBConfig, EmbeddingModelConfig, Project, SourceFile, ChatMessage, TemplateChapter, TextChunk, DeliveryStandard, LLMProviderSchema, VectorDBConfigSchema, EmbeddingConfigSchema, ProjectSchema, SourceFileSchema, ChatMessageSchema, TemplateChapterSchema, DeliveryStandardSchema, ) from llm_client import llm_client from vector_service import ( embedding_service, vector_store, rag_service, chunk_text, search_text_chunks_keyword, ) from parsers.registry import ( parse_file, detect_file_type, categorize_file, get_file_size, SUPPORTED_EXTENSIONS, ) from parsers.excel_parser import parse_excel_to_chunks, pre_parse_excel logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s") logger = logging.getLogger("engimind") @asynccontextmanager async def lifespan(app: FastAPI): db.init_global() logger.info("Engimind backend started") # Try connecting to Qdrant if configured with db.global_session() as s: vdb = s.query(VectorDBConfig).first() if vdb and vdb.endpoint: try: vector_store.connect(vdb.endpoint) except Exception as e: logger.warning("Qdrant connection failed at startup: %s", e) # Auto-open the most recent project so API calls work immediately with db.global_session() as s: latest = s.query(Project).order_by(Project.updated_at.desc()).first() if latest and latest.path: try: db.open_project(latest.id, latest.path) logger.info("Auto-opened project: %s (%s)", latest.name, latest.id) except Exception as e: logger.warning("Failed to auto-open project: %s", e) yield await llm_client.close() await embedding_service.close() vector_store.disconnect() logger.info("Engimind backend stopped") app = FastAPI(title="Engimind", version="2.0.0", lifespan=lifespan) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # ══════════════════════════════════════════════ # Health # ══════════════════════════════════════════════ @app.get("/health") async def health(): return {"status": "ok"} # ══════════════════════════════════════════════ # Config: LLM Providers # ══════════════════════════════════════════════ @app.get("/api/providers", response_model=list[LLMProviderSchema]) async def get_all_providers(): with db.global_session() as s: rows = s.query(LLMProvider).all() return [LLMProviderSchema.from_orm_obj(r) for r in rows] @app.post("/api/providers") async def save_provider(p: LLMProviderSchema): with db.global_session() as s: obj = s.query(LLMProvider).filter_by(id=p.id).first() if obj: obj.name, obj.provider, obj.base_url = p.name, p.provider, p.url obj.api_key, obj.model_id, obj.enabled = p.key, p.model, p.enabled else: obj = LLMProvider(id=p.id, name=p.name, provider=p.provider, base_url=p.url, api_key=p.key, model_id=p.model, enabled=p.enabled) s.add(obj) s.commit() return {"ok": True} @app.delete("/api/providers/{provider_id}") async def delete_provider(provider_id: str): with db.global_session() as s: s.query(LLMProvider).filter_by(id=provider_id).delete() s.commit() return {"ok": True} # ══════════════════════════════════════════════ # Config: VectorDB # ══════════════════════════════════════════════ @app.get("/api/vector-db", response_model=VectorDBConfigSchema) async def get_vector_db(): with db.global_session() as s: cfg = s.query(VectorDBConfig).first() if not cfg: return VectorDBConfigSchema() return VectorDBConfigSchema(endpoint=cfg.endpoint, apiKey=cfg.api_key, status=cfg.status) @app.post("/api/vector-db") async def save_vector_db(c: VectorDBConfigSchema): with db.global_session() as s: cfg = s.query(VectorDBConfig).first() if cfg: cfg.endpoint, cfg.api_key, cfg.status = c.endpoint, c.apiKey, c.status else: s.add(VectorDBConfig(id=1, endpoint=c.endpoint, api_key=c.apiKey, status=c.status)) s.commit() return {"ok": True} @app.post("/api/vector-db/test") async def test_vector_db(body: dict): endpoint = body.get("endpoint", "") try: from qdrant_client import QdrantClient client = QdrantClient(url=endpoint, timeout=5) client.get_collections() client.close() return {"ok": True} except Exception as e: raise HTTPException(400, f"连接失败: {e}") # ══════════════════════════════════════════════ # Config: Embedding # ══════════════════════════════════════════════ @app.get("/api/embedding", response_model=EmbeddingConfigSchema) async def get_embedding_config(): with db.global_session() as s: cfg = s.query(EmbeddingModelConfig).first() if not cfg: return EmbeddingConfigSchema() return EmbeddingConfigSchema.from_orm_obj(cfg) @app.post("/api/embedding") async def save_embedding_config(c: EmbeddingConfigSchema): with db.global_session() as s: cfg = s.query(EmbeddingModelConfig).first() if cfg: cfg.provider, cfg.base_url, cfg.api_key = c.provider, c.url, c.key cfg.model_id, cfg.enabled = c.model, c.enabled cfg.dimensions = c.dimensions else: s.add(EmbeddingModelConfig(id=1, provider=c.provider, base_url=c.url, api_key=c.key, model_id=c.model, enabled=c.enabled, dimensions=c.dimensions)) s.commit() return {"ok": True} @app.post("/api/embedding/test") async def test_embedding(body: dict): try: vec = await embedding_service.get_embedding( "hello", body["url"], body["model"], body.get("key", ""), body["provider"]) if not vec: raise HTTPException(400, "返回的向量为空") return {"ok": True, "dim": len(vec)} except HTTPException: raise except Exception as e: raise HTTPException(400, f"Embedding 测试失败: {e}") @app.post("/api/llm/test") async def test_llm(body: dict): provider = body.get("provider", "") base_url = body.get("url", "").rstrip("/") api_key = body.get("key", "") import httpx if provider == "Ollama": url = base_url + "/api/tags" elif provider in ("DeepSeek", "OpenAI"): url = base_url + "/models" else: url = base_url + "/v1/models" headers = {} if api_key and provider != "Ollama": headers["Authorization"] = f"Bearer {api_key}" try: async with httpx.AsyncClient(timeout=8) as client: resp = await client.get(url, headers=headers) if 200 <= resp.status_code < 400: return {"ok": True} if resp.status_code == 401: raise HTTPException(400, "API Key 无效") raise HTTPException(400, f"HTTP {resp.status_code}") except httpx.HTTPError as e: raise HTTPException(400, f"网络连通性异常: {e}") # ══════════════════════════════════════════════ # Projects # ══════════════════════════════════════════════ @app.get("/api/projects", response_model=list[ProjectSchema]) async def list_projects(): with db.global_session() as s: return [ProjectSchema.model_validate(p) for p in s.query(Project).order_by(Project.created_at.desc()).all()] @app.post("/api/projects", response_model=ProjectSchema) async def create_project(body: dict): name = body.get("name", "未命名项目") pid = f"p-{int(time.time() * 1000)}" proj_dir = db.projects_base_dir() / pid proj_dir.mkdir(parents=True, exist_ok=True) db_path = str(proj_dir / "project.db") with db.global_session() as s: proj = Project(id=pid, name=name, path=db_path) s.add(proj) s.commit() db.open_project(pid, db_path) return ProjectSchema(id=pid, name=name) @app.post("/api/projects/{project_id}/switch") async def switch_project(project_id: str): with db.global_session() as s: proj = s.query(Project).filter_by(id=project_id).first() if not proj: raise HTTPException(404, "项目不存在") db.open_project(proj.id, proj.path) return {"ok": True} @app.delete("/api/projects/{project_id}") async def delete_project(project_id: str): if db.current_project_id == project_id: db.close_project() with db.global_session() as s: proj = s.query(Project).filter_by(id=project_id).first() if proj: proj_dir = os.path.dirname(proj.path) if os.path.isdir(proj_dir): shutil.rmtree(proj_dir, ignore_errors=True) s.delete(proj) s.commit() return {"ok": True} # ══════════════════════════════════════════════ # Materials # ══════════════════════════════════════════════ @app.get("/api/materials", response_model=list[SourceFileSchema]) async def get_project_files(): session = db.project_session() if not session or not db.current_project_id: return [] with session as s: files = s.query(SourceFile).filter_by(project_id=db.current_project_id)\ .order_by(SourceFile.created_at.desc()).all() return [SourceFileSchema.from_orm_obj(f) for f in files] @app.get("/api/materials/{file_id}/content") async def get_material_content(file_id: str): """Return file content for online preview. - Excel: structured JSON with sheets -> rows (for table rendering) - Others: markdown text """ session = db.project_session() if not session: raise HTTPException(400, "项目数据库未就绪") with session as s: sf = s.query(SourceFile).filter_by(id=file_id).first() if not sf: raise HTTPException(404, "素材未找到") file_path = sf.file_path file_type = sf.type cached_content = sf.parsed_content or "" if not file_path or not os.path.isfile(file_path): raise HTTPException(400, "文件不存在") if file_type == "excel": # Return structured table data for rich rendering from parsers.excel_parser import _prepare_grid, _cell_str from openpyxl import load_workbook wb = load_workbook(file_path, data_only=True) sheets = [] for sheet_name in wb.sheetnames: ws = wb[sheet_name] grid = _prepare_grid(ws) if not grid: continue rows = [] for row in grid: rows.append([_cell_str(c) for c in row]) sheets.append({"name": sheet_name, "rows": rows}) wb.close() return {"type": "excel", "sheets": sheets} else: # For non-excel: return cached markdown or re-parse if cached_content: return {"type": file_type, "content": cached_content} result = parse_file(file_path) return {"type": file_type, "content": result.get("markdown", "")} class UploadRequest(BaseModel): filePaths: list[str] @app.post("/api/materials/upload", response_model=list[SourceFileSchema]) async def upload_materials(req: UploadRequest): """Upload materials by local file paths (Electron sends paths from dialog).""" pid = db.current_project_id if not pid: raise HTTPException(400, "请先创建或选择一个项目") session = db.project_session() if not session: raise HTTPException(400, "项目数据库未就绪") files_dir = db.projects_base_dir() / pid / "files" files_dir.mkdir(parents=True, exist_ok=True) results = [] for i, src_path in enumerate(req.filePaths): if not os.path.isfile(src_path): continue file_name = os.path.basename(src_path) file_id = f"f-{int(time.time() * 1000)}-{i}" file_type = detect_file_type(src_path) file_size = get_file_size(src_path) dst_path = str(files_dir / f"{file_id}_{file_name}") shutil.copy2(src_path, dst_path) with session as s: # Remove duplicates existing = s.query(SourceFile).filter_by(project_id=pid, name=file_name).all() for old in existing: if old.file_path and os.path.isfile(old.file_path): os.remove(old.file_path) s.query(TextChunk).filter_by(source_id=old.id).delete() s.delete(old) if vector_store.connected: vector_store.delete_by_source(pid, old.id) sf = SourceFile(id=file_id, project_id=pid, name=file_name, type=file_type, category=categorize_file(file_type), file_path=dst_path, size=file_size, vector_status="pending") s.add(sf) s.commit() results.append(SourceFileSchema(id=file_id, name=file_name, type=file_type, category=categorize_file(file_type), size=file_size, vectorStatus="pending")) # Parse and index in background asyncio.create_task(_parse_and_index(pid, file_id, file_name, file_type, dst_path)) return results @app.delete("/api/materials/{file_id}") async def delete_material(file_id: str): session = db.project_session() if not session: raise HTTPException(400, "项目数据库未就绪") pid = db.current_project_id with session as s: sf = s.query(SourceFile).filter_by(id=file_id).first() if not sf: raise HTTPException(404, "素材未找到") if sf.file_path and os.path.isfile(sf.file_path): os.remove(sf.file_path) s.query(TextChunk).filter_by(source_id=file_id).delete() s.delete(sf) s.commit() if vector_store.connected and pid: vector_store.delete_by_source(pid, file_id) return {"ok": True} # ══════════════════════════════════════════════ # Chat # ══════════════════════════════════════════════ @app.get("/api/chat/messages", response_model=list[ChatMessageSchema]) async def get_chat_messages(): session = db.project_session() if not session or not db.current_project_id: return [] with session as s: msgs = s.query(ChatMessage).filter_by(project_id=db.current_project_id)\ .order_by(ChatMessage.created_at).all() return [ChatMessageSchema.from_orm_obj(m) for m in msgs] class SaveMessageRequest(BaseModel): role: str content: str sources: str = "" citations: str = "" @app.post("/api/chat/messages") async def save_chat_message(req: SaveMessageRequest): session = db.project_session() if not session or not db.current_project_id: raise HTTPException(400, "无活动项目") with session as s: s.add(ChatMessage(project_id=db.current_project_id, role=req.role, content=req.content, sources=req.sources, citations=req.citations)) s.commit() return {"ok": True} @app.delete("/api/chat/messages") async def clear_chat_messages(): session = db.project_session() if not session or not db.current_project_id: return {"ok": True} with session as s: s.query(ChatMessage).filter_by(project_id=db.current_project_id).delete() s.commit() return {"ok": True} class SendMessageRequest(BaseModel): content: str selectedFileIds: list[str] = [] modelId: str _CHAT_SYSTEM_PROMPT = ( "你是一位专业的工程技术助手。你只能从提供的 Context 素材中提取数值和信息来回答问题。\n" "重要约束:\n" "1. 如果 Context 中没有与问题相关的数据,必须诚实回答『在当前素材中未找到相关数据』,严禁编造数值或结论。\n" "2. 绝对不要联网搜索或自行猜测数据。\n" "3. 素材中的数据路径可能与用户提问的词序不同。例如素材写『湿地 > 国家所有』而用户问『国家所有湿地』," "你必须识别出这是同一个属性并正确引用数值。\n" "4. 如果素材信息不完整,应指出缺失部分,而非自行填补。\n" "5. 引用来源时使用 [N] 标注。" ) @app.post("/api/chat/send") async def send_message(req: SendMessageRequest): """Non-streaming chat with RAG context.""" provider = _get_provider(req.modelId) context_text = _search_material_context(req.content, req.selectedFileIds, 5) messages = [ {"role": "system", "content": _CHAT_SYSTEM_PROMPT}, ] if context_text: messages.append({"role": "user", "content": f"参考以下工程素材:\n\n{context_text}\n\n---\n\n用户提问:{req.content}"}) else: messages.append({"role": "user", "content": req.content}) result = await llm_client.complete(provider["url"], provider["key"], provider["model"], messages) if result.get("choices"): return {"content": result["choices"][0]["message"]["content"]} raise HTTPException(500, "模型无响应") @app.post("/api/chat/stream") async def stream_message(req: SendMessageRequest): """Streaming chat with RAG context via SSE.""" provider = _get_provider(req.modelId) context_text = _search_material_context(req.content, req.selectedFileIds, 5) messages = [ {"role": "system", "content": _CHAT_SYSTEM_PROMPT}, ] if context_text: messages.append({"role": "user", "content": f"参考以下工程素材:\n\n{context_text}\n\n---\n\n用户提问:{req.content}"}) else: messages.append({"role": "user", "content": req.content}) async def event_gen(): import json async for chunk in llm_client.stream_complete(provider["url"], provider["key"], provider["model"], messages): yield {"data": json.dumps(chunk, ensure_ascii=False)} return EventSourceResponse(event_gen()) # ══════════════════════════════════════════════ # Template / Delivery Standard # ══════════════════════════════════════════════ @app.get("/api/template/chapters", response_model=list[TemplateChapterSchema]) async def get_template_chapters(): session = db.project_session() if not session or not db.current_project_id: return [] with session as s: chapters = s.query(TemplateChapter).filter_by(project_id=db.current_project_id)\ .order_by(TemplateChapter.sort_order).all() return [TemplateChapterSchema.from_orm_obj(c) for c in chapters] @app.post("/api/template/chapters") async def save_template_chapters(chapters: list[TemplateChapterSchema]): session = db.project_session() if not session or not db.current_project_id: raise HTTPException(400, "无活动项目") with session as s: s.query(TemplateChapter).filter_by(project_id=db.current_project_id).delete() for i, c in enumerate(chapters): s.add(TemplateChapter( id=c.id, project_id=db.current_project_id, title=c.title, status=c.status, progress=c.progress, content=c.content, sort_order=i, )) s.commit() return {"ok": True} @app.delete("/api/template/chapters/{chapter_id}") async def delete_template_chapter(chapter_id: str): session = db.project_session() if not session: raise HTTPException(400, "无活动项目") with session as s: s.query(TemplateChapter).filter_by(id=chapter_id).delete() s.commit() return {"ok": True} @app.post("/api/template/extract-directory") async def stream_template_directory(body: dict): """Stream LLM directory extraction from delivery standard text.""" content = body.get("content", "") model_id = body.get("modelId", "") provider = _get_provider(model_id) prompt = ( "你是一个工程标准的目录解析助手。请从下面提供的交付标准文本中提取完整的章节目录结构(包括主章节和小节),并以 JSON 数组的格式返回。\n\n" "### 要求:\n" "1. 只返回 JSON 数组,不包含其他废话或者回答前缀。\n" "2. 必须提取所有层级:主章节和其下的小节(子章节)。\n" "3. 输出格式必须严格符合:\n" '[{"id": "chapter-1", "title": "1 原材料进场检验", "content": "简要描述(可选)", ' '"children": [{"id": "chapter-1-1", "title": "1.1 钢材检验", "content": ""}]}]\n' "4. 如果某个章节没有小节,children 可以为空数组 []。\n\n" f"### 交付标准内容:\n{content}" ) messages = [ {"role": "system", "content": "你是一个专业的结构化数据抽取工具。你只输出合法的 JSON,不要使用 Markdown 代码块包裹,也不要给出任何其他解释。"}, {"role": "user", "content": prompt}, ] async def event_gen(): import json async for chunk in llm_client.stream_complete(provider["url"], provider["key"], provider["model"], messages): yield {"data": json.dumps(chunk, ensure_ascii=False)} return EventSourceResponse(event_gen()) @app.get("/api/delivery-standard", response_model=Optional[DeliveryStandardSchema]) async def get_delivery_standard(): session = db.project_session() if not session or not db.current_project_id: return None with session as s: ds = s.query(DeliveryStandard).filter_by(project_id=db.current_project_id).first() if not ds: return None return DeliveryStandardSchema(fileName=ds.file_name, content=ds.content) @app.post("/api/delivery-standard") async def save_delivery_standard(body: DeliveryStandardSchema): session = db.project_session() if not session or not db.current_project_id: raise HTTPException(400, "无活动项目") with session as s: s.query(DeliveryStandard).filter_by(project_id=db.current_project_id).delete() s.add(DeliveryStandard(project_id=db.current_project_id, file_name=body.fileName, content=body.content)) s.commit() return {"ok": True} @app.post("/api/parse-file") async def parse_file_endpoint(body: dict): """Parse a local file and return markdown. Used for delivery standard parsing.""" file_path = body.get("filePath", "") if not file_path or not os.path.isfile(file_path): raise HTTPException(400, "文件不存在") result = parse_file(file_path) return result # ══════════════════════════════════════════════ # Excel Pre-parse & Final Ingest # ══════════════════════════════════════════════ class PreParseRequest(BaseModel): fileId: str startRow: Optional[int] = None # optional override for "refresh preview" @app.post("/api/materials/pre-parse") async def pre_parse_material(req: PreParseRequest): """Interface A: Pre-parse an Excel file and return preview data.""" session = db.project_session() if not session: raise HTTPException(400, "项目数据库未就绪") with session as s: sf = s.query(SourceFile).filter_by(id=req.fileId).first() if not sf: raise HTTPException(404, "素材未找到") file_path = sf.file_path file_type = sf.type if file_type != "excel": raise HTTPException(400, "仅支持 Excel 文件的预解析") if not file_path or not os.path.isfile(file_path): raise HTTPException(400, "文件不存在") result = pre_parse_excel(file_path, start_row=req.startRow) return result class FinalIngestRequest(BaseModel): fileId: str startRow: int @app.post("/api/materials/ingest") async def final_ingest_material(req: FinalIngestRequest): """Interface B: Final ingest — parse with confirmed start_row and vector index.""" pid = db.current_project_id if not pid: raise HTTPException(400, "请先选择一个项目") session = db.project_session() if not session: raise HTTPException(400, "项目数据库未就绪") with session as s: sf = s.query(SourceFile).filter_by(id=req.fileId).first() if not sf: raise HTTPException(404, "素材未找到") file_path = sf.file_path file_name = sf.name file_type = sf.type if file_type != "excel": raise HTTPException(400, "仅支持 Excel 文件") if not file_path or not os.path.isfile(file_path): raise HTTPException(400, "文件不存在") # Launch background ingest task asyncio.create_task(_ingest_excel(pid, req.fileId, file_name, file_path, req.startRow)) return {"ok": True, "message": f"开始入库,数据起始行: {req.startRow}"} async def _ingest_excel(project_id: str, file_id: str, file_name: str, file_path: str, start_row: int): """Background: parse Excel with confirmed start_row, clear old vectors, re-index.""" import json session = db.project_session() if not session: return def emit(status: str, step: str): _broadcast_event(json.dumps({ "type": "material_status_update", "fileId": file_id, "fileName": file_name, "status": status, "step": step, }, ensure_ascii=False)) try: with session as s: sf = s.query(SourceFile).filter_by(id=file_id).first() if sf: sf.vector_status = "processing" s.commit() emit("processing", f"📄 开始重新解析「{file_name}」(起始行: {start_row})") # Step 1: Parse with confirmed start_row raw_chunks = parse_excel_to_chunks(file_path, start_row=start_row) text_chunks = [c['content'] for c in raw_chunks] metadata_list = [c['metadata'] for c in raw_chunks] emit("processing", f"🔪 文本分块: 共 {len(text_chunks)} 个片段") # Step 2: Clear old data with session as s: s.query(TextChunk).filter_by(source_id=file_id).delete() s.commit() if vector_store.connected: vector_store.delete_by_source(project_id, file_id) emit("processing", "🗑️ 已清除旧向量数据") # Step 3: Store chunks in SQLite with session as s: for i, chunk in enumerate(text_chunks): s.add(TextChunk(id=f"{file_id}-chunk-{i}", project_id=project_id, source_id=file_id, content=chunk, chunk_idx=i)) s.commit() # Step 4: Vector index emb_cfg = _get_embedding_config() if emb_cfg and vector_store.connected: emit("processing", f"🔗 开始向量化索引 ({len(text_chunks)} 分块)...") dim = emb_cfg.get("dimensions", 1024) def on_progress(done: int, total: int): emit("processing", f"🔗 向量化进度: {done}/{total}") await rag_service.index_chunks( project_id, file_id, text_chunks, emb_cfg, dim=dim, metadata_list=metadata_list, on_progress=on_progress, ) emit("done", "✓ 素材完全就绪,向量索引成功") else: emit("done", f"✓ 素材就绪 ({len(text_chunks)} 分块)") with session as s: sf = s.query(SourceFile).filter_by(id=file_id).first() if sf: sf.vector_status = "done" s.commit() except Exception as e: logger.exception("Excel ingest failed for %s", file_name) emit("error", f"入库失败: {e}") with session as s: sf = s.query(SourceFile).filter_by(id=file_id).first() if sf: sf.vector_status = "error" s.commit() # ══════════════════════════════════════════════ # Chapter Generation # ══════════════════════════════════════════════ class GenerateChapterRequest(BaseModel): chapterTitle: str selectedFileIds: list[str] = [] modelId: str @app.post("/api/chat/generate-chapter") async def stream_generate_chapter(req: GenerateChapterRequest): provider = _get_provider(req.modelId) context_text = _search_material_context(req.chapterTitle, req.selectedFileIds, 8) prompt = ( f"你是工程报告撰写专家。请根据以下工程素材,按照章节要求撰写报告内容。\n\n" f"## 章节要求\n{req.chapterTitle}\n\n" f"## 参考素材\n{context_text}\n\n" f"## 输出要求\n" f"1. 使用 Markdown 格式\n" f"2. 引用素材时使用 [N] 标注\n" f"3. 内容专业、结构清晰\n" f"4. 包含具体数据和分析结论" ) messages = [ {"role": "system", "content": ( "你是一位资深工程报告撰写专家,擅长根据工程素材生成结构化的技术报告章节。\n" "重要约束:只能使用参考素材中的数据和信息。如果素材不足以支撑某个章节,请明确标注" "『素材中未包含此部分数据,需补充』,绝不可编造数值或结论。" )}, {"role": "user", "content": prompt}, ] async def event_gen(): import json async for chunk in llm_client.stream_complete(provider["url"], provider["key"], provider["model"], messages): yield {"data": json.dumps(chunk, ensure_ascii=False)} return EventSourceResponse(event_gen()) # ══════════════════════════════════════════════ # SSE Event Stream (material status broadcast) # ══════════════════════════════════════════════ # Simple in-memory event bus for material processing status _event_subscribers: list[asyncio.Queue] = [] @app.get("/api/events") async def event_stream(): """SSE endpoint for real-time events (material status, etc.).""" queue: asyncio.Queue = asyncio.Queue() _event_subscribers.append(queue) async def gen(): try: while True: event = await queue.get() yield {"data": event} except asyncio.CancelledError: pass finally: _event_subscribers.remove(queue) return EventSourceResponse(gen()) def _broadcast_event(event_data: str): for q in _event_subscribers: q.put_nowait(event_data) # ══════════════════════════════════════════════ # Internal helpers # ══════════════════════════════════════════════ def _get_provider(model_id: str) -> dict: with db.global_session() as s: p = s.query(LLMProvider).filter_by(id=model_id, enabled=True).first() if not p: raise HTTPException(400, f"模型 {model_id} 未找到或已禁用") return {"url": p.base_url, "key": p.api_key, "model": p.model_id, "provider": p.provider} def _get_embedding_config() -> dict | None: with db.global_session() as s: cfg = s.query(EmbeddingModelConfig).first() if cfg and cfg.enabled and cfg.base_url: return {"base_url": cfg.base_url, "model": cfg.model_id, "api_key": cfg.api_key, "provider": cfg.provider, "dimensions": cfg.dimensions or 1024} return None def _search_material_context(query: str, file_ids: list[str], top_k: int) -> str: pid = db.current_project_id if not pid: return "" chunks = [] # Try vector search emb_cfg = _get_embedding_config() if emb_cfg and vector_store.connected: try: import asyncio loop = asyncio.get_event_loop() chunks_coro = rag_service.search_context(pid, query, top_k, emb_cfg, file_ids or None) # We're in sync context here, but called from async — use create_task workaround # Actually this helper is called from async routes, so just run sync approach import concurrent.futures with concurrent.futures.ThreadPoolExecutor() as pool: future = pool.submit(asyncio.run, chunks_coro) chunks = future.result() except Exception as e: logger.warning("Vector search failed, falling back: %s", e) chunks = [] # Fallback: keyword search if not chunks: session = db.project_session() if session: with session as s: chunks = search_text_chunks_keyword(s, pid, query, file_ids or None, top_k) if not chunks: return "" return "\n\n---\n\n".join(c.get("text", "") for c in chunks) async def _parse_and_index(project_id: str, file_id: str, file_name: str, file_type: str, file_path: str): """Background task: parse file, store chunks, optionally vector index.""" import json session = db.project_session() if not session: return def emit(status: str, step: str): _broadcast_event(json.dumps({ "type": "material_status_update", "fileId": file_id, "fileName": file_name, "status": status, "step": step, }, ensure_ascii=False)) try: # Update status with session as s: sf = s.query(SourceFile).filter_by(id=file_id).first() if sf: sf.vector_status = "processing" s.commit() emit("processing", f"📄 开始处理「{file_name}」") emit("processing", f"📖 正在解析 {file_type.upper()} 文件内容...") result = parse_file(file_path) content = result.get("markdown", "") if not content: emit("error", f"解析失败: {result.get('error', '无内容')}") with session as s: sf = s.query(SourceFile).filter_by(id=file_id).first() if sf: sf.vector_status = "error" s.commit() return emit("processing", f"✓ 文件解析完成 (提取 {len(content)} 字符)") # Save parsed content with session as s: sf = s.query(SourceFile).filter_by(id=file_id).first() if sf: sf.parsed_content = content s.commit() # Chunk text if file_type == "excel": raw_chunks = parse_excel_to_chunks(file_path, start_row=None) text_chunks = [c['content'] for c in raw_chunks] metadata_list = [c['metadata'] for c in raw_chunks] else: text_chunks = chunk_text(content, 800, 80) metadata_list = None emit("processing", f"🔪 文本分块: 共 {len(text_chunks)} 个片段") # Store chunks in SQLite with session as s: for i, chunk in enumerate(text_chunks): s.add(TextChunk(id=f"{file_id}-chunk-{i}", project_id=project_id, source_id=file_id, content=chunk, chunk_idx=i)) s.commit() # Optional: vector indexing (use pre-chunked text instead of re-chunking) emb_cfg = _get_embedding_config() logger.info("Vector indexing check: emb_cfg=%s, connected=%s", {k: v for k, v in emb_cfg.items() if k != 'api_key'} if emb_cfg else None, vector_store.connected) if emb_cfg and vector_store.connected: emit("processing", f"🔗 开始向量化索引 ({len(text_chunks)} 分块)...") try: dim = emb_cfg.get("dimensions", 1024) logger.info("Calling index_chunks: project=%s, source=%s, chunks=%d, dim=%d", project_id, file_id, len(text_chunks), dim) def on_progress(done: int, total: int): emit("processing", f"🔗 向量化进度: {done}/{total}") await rag_service.index_chunks( project_id, file_id, text_chunks, emb_cfg, dim=dim, metadata_list=metadata_list, on_progress=on_progress, ) emit("done", "✓ 素材完全就绪,向量索引成功") except Exception as e: logger.exception("Vector indexing failed for %s", file_id) emit("done", f"⚠️ 素材已就绪,向量索引失败: {e}") else: emit("done", f"✓ 素材就绪 ({len(text_chunks)} 分块)") with session as s: sf = s.query(SourceFile).filter_by(id=file_id).first() if sf: sf.vector_status = "done" s.commit() except Exception as e: logger.exception("Parse and index failed for %s", file_name) emit("error", f"处理失败: {e}") with session as s: sf = s.query(SourceFile).filter_by(id=file_id).first() if sf: sf.vector_status = "error" s.commit() # ══════════════════════════════════════════════ # Entry point # ══════════════════════════════════════════════ if __name__ == "__main__": import uvicorn port = int(os.environ.get("ENGIMIND_PORT", "9231")) uvicorn.run(app, host="127.0.0.1", port=port, log_level="info")