把线上chat模式并入

chatMSA/__init__.py

@@ -0,0 +1,8 @@
+"""
+chatMSA — Multi-turn conversation system built on MSA (Memory Sparse Attention).
+
+Usage:
+    python -m chatMSA.app --model_path ckpt/MSA-4B --port 7860
+"""
+
+__version__ = "0.1.0"

chatMSA/__main__.py

@@ -0,0 +1,5 @@
+"""Allow running chatMSA as a module: python -m chatMSA"""
+
+from chatMSA.app import main
+
+main()

chatMSA/api/__init__.py

@@ -0,0 +1,3 @@
+from chatMSA.api.router import create_fastapi_app
+
+__all__ = ["create_fastapi_app"]

chatMSA/api/chat_routes.py

@@ -0,0 +1,139 @@
+"""
+FastAPI routes for conversation management and chat.
+
+Endpoints:
+    POST   /api/conversations                -> Create conversation
+    GET    /api/conversations                -> List all conversations
+    GET    /api/conversations/{id}           -> Get conversation with messages
+    DELETE /api/conversations/{id}           -> Delete conversation
+    PATCH  /api/conversations/{id}           -> Rename conversation
+    POST   /api/conversations/{id}/messages  -> Send message, get response
+    GET    /api/health                       -> Engine status
+"""
+
+from fastapi import APIRouter, HTTPException
+
+from chatMSA.models.schemas import (
+    ConversationCreateRequest,
+    ConversationDetail,
+    ConversationRenameRequest,
+    ConversationSummary,
+    HealthResponse,
+    MessageResponse,
+    MessageSendRequest,
+)
+from chatMSA.services.chat_service import ChatService
+from chatMSA.services.msa_engine_service import MSAEngineService
+
+
+def create_chat_routes(chat_service: ChatService, engine: MSAEngineService) -> APIRouter:
+    """Create and return the chat API router."""
+    router = APIRouter(prefix="/api")
+
+    # ── Health ──────────────────────────────────────────────────
+
+    @router.get("/health", response_model=HealthResponse)
+    def health():
+        status = "ready" if engine.is_ready else ("loading" if engine.is_loading else "error")
+        import torch
+        gpu_count = len(engine.config.devices) if engine.config.devices else torch.cuda.device_count()
+        return HealthResponse(
+            status=status,
+            model_path=engine.config.model_path,
+            gpu_count=gpu_count,
+            uptime_seconds=engine.uptime,
+            error=engine.error,
+        )
+
+    # ── Conversations ───────────────────────────────────────────
+
+    @router.post("/conversations", response_model=ConversationDetail, status_code=201)
+    def create_conversation(req: ConversationCreateRequest = None):
+        if req is None:
+            req = ConversationCreateRequest()
+        conv = chat_service.create_conversation(title=req.title)
+        return _to_detail(conv)
+
+    @router.get("/conversations", response_model=list)
+    def list_conversations():
+        convs = chat_service.list_conversations()
+        return [_to_summary(c) for c in convs]
+
+    @router.get("/conversations/{conv_id}", response_model=ConversationDetail)
+    def get_conversation(conv_id: str):
+        conv = chat_service.get_conversation(conv_id)
+        if conv is None:
+            raise HTTPException(status_code=404, detail="Conversation not found")
+        return _to_detail(conv)
+
+    @router.delete("/conversations/{conv_id}", status_code=204)
+    def delete_conversation(conv_id: str):
+        deleted = chat_service.delete_conversation(conv_id)
+        if not deleted:
+            raise HTTPException(status_code=404, detail="Conversation not found")
+
+    @router.patch("/conversations/{conv_id}", response_model=ConversationDetail)
+    def rename_conversation(conv_id: str, req: ConversationRenameRequest):
+        conv = chat_service.rename_conversation(conv_id, req.title)
+        if conv is None:
+            raise HTTPException(status_code=404, detail="Conversation not found")
+        return _to_detail(conv)
+
+    # ── Messages ────────────────────────────────────────────────
+
+    @router.post("/conversations/{conv_id}/messages", response_model=MessageResponse)
+    def send_message(conv_id: str, req: MessageSendRequest):
+        try:
+            msg = chat_service.send_message(conv_id, req.content)
+        except ValueError as e:
+            raise HTTPException(status_code=404, detail=str(e))
+        except RuntimeError as e:
+            raise HTTPException(status_code=503, detail=str(e))
+        return MessageResponse(
+            id=msg.id,
+            role=msg.role,
+            content=msg.content,
+            timestamp=msg.timestamp,
+            recall_topk=msg.recall_topk,
+        )
+
+    return router
+
+
+# ── Helpers ─────────────────────────────────────────────────────
+
+def _to_summary(conv) -> ConversationSummary:
+    msg_count = getattr(conv, "_message_count", conv.message_count)
+    turn_count = getattr(conv, "_turn_count", conv.turn_count)
+    last_preview = None
+    if conv.messages:
+        last_preview = conv.messages[-1].content[:100]
+    return ConversationSummary(
+        id=conv.id,
+        title=conv.title,
+        created_at=conv.created_at,
+        updated_at=conv.updated_at,
+        message_count=msg_count,
+        turn_count=turn_count,
+        last_message_preview=last_preview,
+    )
+
+
+def _to_detail(conv) -> ConversationDetail:
+    return ConversationDetail(
+        id=conv.id,
+        title=conv.title,
+        created_at=conv.created_at,
+        updated_at=conv.updated_at,
+        messages=[
+            MessageResponse(
+                id=m.id,
+                role=m.role,
+                content=m.content,
+                timestamp=m.timestamp,
+                recall_topk=m.recall_topk,
+            )
+            for m in conv.messages
+        ],
+        metadata=conv.metadata,
+    )

chatMSA/api/router.py

@@ -0,0 +1,45 @@
+"""
+API router aggregation.
+
+Combines all route modules into a single FastAPI application.
+"""
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+
+from chatMSA.api.chat_routes import create_chat_routes
+from chatMSA.services.chat_service import ChatService
+from chatMSA.services.msa_engine_service import MSAEngineService
+
+
+def create_fastapi_app(chat_service: ChatService, engine: MSAEngineService) -> FastAPI:
+    """
+    Create and configure the FastAPI application.
+
+    Args:
+        chat_service: The chat business logic service.
+        engine: The MSA engine service (for health checks).
+
+    Returns:
+        Configured FastAPI app with all routes mounted.
+    """
+    app = FastAPI(
+        title="chatMSA API",
+        description="Multi-turn conversation system built on MSA (Memory Sparse Attention)",
+        version="0.1.0",
+    )
+
+    # CORS — allow Gradio frontend and local development
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
+    # Mount API routes
+    chat_router = create_chat_routes(chat_service, engine)
+    app.include_router(chat_router)
+
+    return app

chatMSA/app.py

@@ -0,0 +1,85 @@
+"""
+chatMSA entry point.
+
+Launches FastAPI (REST API) with embedded Gradio (chat UI).
+
+Usage:
+    python -m chatMSA.app --model_path ckpt/MSA-4B --port 7860
+    python -m chatMSA.app --help
+"""
+
+import signal
+import sys
+import pathlib
+
+# Ensure project root is importable
+_project_root = pathlib.Path(__file__).parent.parent
+sys.path.insert(0, str(_project_root))
+
+import gradio as gr
+import uvicorn
+
+from chatMSA.config import ChatConfig
+from chatMSA.services.chat_service import ChatService
+from chatMSA.services.msa_engine_service import MSAEngineService
+from chatMSA.storage.sqlite_store import SQLiteConversationStore
+from chatMSA.api.router import create_fastapi_app
+from chatMSA.frontend.chat_ui import create_gradio_app
+
+
+def main():
+    """Main entry point."""
+    # 1. Parse config
+    config = ChatConfig.from_args()
+    print(f"[chatMSA] Config: model={config.model_path}, db={config.db_path}")
+    print(f"[chatMSA] Server: {config.host}:{config.port}")
+
+    # 2. Initialize storage
+    store = SQLiteConversationStore(config.db_path)
+    store.initialize()
+    print(f"[chatMSA] Storage initialized: {config.db_path}")
+
+    # 3. Initialize engine (lazy — heavy loading happens in engine.start())
+    engine = MSAEngineService(config)
+
+    # 4. Initialize chat service
+    chat_service = ChatService(engine, store, config)
+
+    # 5. Create FastAPI app
+    fastapi_app = create_fastapi_app(chat_service, engine)
+
+    # 6. Create and mount Gradio app
+    gradio_app = create_gradio_app(chat_service)
+    fastapi_app = gr.mount_gradio_app(fastapi_app, gradio_app, path="/")
+
+    # 7. Graceful shutdown
+    def shutdown_handler(sig, frame):
+        print("\n[chatMSA] Shutting down...")
+        engine.stop()
+        store.close()
+        sys.exit(0)
+
+    signal.signal(signal.SIGINT, shutdown_handler)
+    signal.signal(signal.SIGTERM, shutdown_handler)
+
+    # 8. Start engine (heavy: model loading + memory prefill)
+    #    This runs before the HTTP server so the first request doesn't wait.
+    #    TODO: Move engine.start() to a background thread and serve a "loading"
+    #    page from Gradio while the engine initializes. This would improve UX
+    #    for large models with long startup times.
+    print("[chatMSA] Starting MSA engine (this may take a while)...")
+    try:
+        engine.start()
+        print("[chatMSA] Engine ready!")
+    except Exception as e:
+        print(f"[chatMSA] Engine failed to start: {e}")
+        print("[chatMSA] Continuing in degraded mode (chat will return errors)")
+
+    # 9. Launch server
+    print(f"[chatMSA] UI: http://{config.host}:{config.port}")
+    print(f"[chatMSA] API: http://{config.host}:{config.port}/docs")
+    uvicorn.run(fastapi_app, host=config.host, port=config.port, log_level="info")
+
+
+if __name__ == "__main__":
+    main()