""" ByteAstra backend — FastAPI application entry point and router setup. """ from __future__ import annotations import torch for i in range(1, 8): setattr(torch, f"int{i}", torch.int8) import torch.utils._pytree if not hasattr(torch.utils._pytree, "register_constant"): torch.utils._pytree.register_constant = lambda cls: cls import json import asyncio import logging from contextlib import asynccontextmanager from typing import AsyncGenerator from fastapi import FastAPI, Depends, HTTPException, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse from pymongo.database import Database as DBSession from sse_starlette.sse import EventSourceResponse from app.config import get_settings from app.database import get_db, init_db from app.schemas import ( ChatRequest, ChatResponse, SessionCreate, SessionOut, SessionSummary, MessageOut ) from app.services import sessions as session_svc from app.services import engine from app.services.sessions import parse_citations logging.basicConfig(level=logging.INFO, format="%(levelname)s | %(name)s | %(message)s") logger = logging.getLogger(__name__) settings = get_settings() # ── Lifespan ─────────────────────────────────────────────────────────────────── @asynccontextmanager async def lifespan(app: FastAPI): """Startup and shutdown events.""" logger.info("ByteAstra backend starting up...") init_db() logger.info("Database initialised.") yield logger.info("ByteAstra backend shutting down.") # ── App ──────────────────────────────────────────────────────────────────────── app = FastAPI( title="ByteAstra API", description="Syllabus-grounded AI tutoring engine for Indian professional curricula.", version="0.1.0", lifespan=lifespan, ) app.add_middleware( CORSMiddleware, allow_origins=settings.cors_origins_list, allow_origin_regex=r"https://.*\.vercel\.app", allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # ── Health ───────────────────────────────────────────────────────────────────── @app.get("/health", tags=["system"]) async def health(): from app.services.llm import check_connection is_online = await check_connection() return { "status": "ok", "version": "0.1.0", "engine_mode": "hybrid" if is_online else "fallback" } @app.get("/debug/llm", tags=["system"]) async def debug_llm(): import traceback from app.services.llm import get_llm_client, check_connection from app.config import get_settings settings = get_settings() is_online = await check_connection() res = { "is_online_check": is_online, "resolved_llm_base_url": settings.resolved_llm_base_url, "resolved_model_name": settings.resolved_model_name, "openai_base_url_env": settings.openai_base_url, "llm_base_url_env": settings.llm_base_url, "model_name_env": settings.model_name, "llm_model_name_env": settings.llm_model_name, "error_traceback": None, "models_list": None, "completion_success": False, "completion_response": None, "llama_server_log": None } from pathlib import Path log_path = Path("/code/llama_server.log") if log_path.exists(): try: res["llama_server_log"] = log_path.read_text(encoding="utf-8", errors="replace")[-2000:] except Exception as e: res["llama_server_log"] = f"Error reading log: {e}" else: res["llama_server_log"] = "Log file not found" try: client = get_llm_client() models = await client.models.list() res["models_list"] = [m.id for m in models.data] except Exception as e: res["models_list_error"] = str(e) res["error_traceback"] = traceback.format_exc() return res try: chat_completion = await client.chat.completions.create( model=settings.resolved_model_name, messages=[{"role": "user", "content": "hello"}], max_tokens=10, temperature=0.0 ) res["completion_success"] = True res["completion_response"] = chat_completion.choices[0].message.content except Exception as e: res["completion_error"] = str(e) res["error_traceback"] = traceback.format_exc() return res # ── Sessions ─────────────────────────────────────────────────────────────────── @app.post("/sessions", response_model=SessionOut, tags=["sessions"]) def create_session(body: SessionCreate, db: DBSession = Depends(get_db)): """Create a new tutoring session for a given domain.""" sess = session_svc.create_session(db, domain=body.domain, title=body.title or "New Session", user_id=body.user_id) return _session_to_out(sess) @app.get("/sessions", response_model=list[SessionSummary], tags=["sessions"]) def list_sessions(user_id: str | None = None, db: DBSession = Depends(get_db)): """List all sessions, optionally filtered by user_id.""" return session_svc.list_sessions(db, user_id=user_id) @app.get("/sessions/{session_id}", response_model=SessionOut, tags=["sessions"]) def get_session(session_id: str, db: DBSession = Depends(get_db)): sess = session_svc.get_session(db, session_id) if not sess: raise HTTPException(status_code=404, detail="Session not found") return _session_to_out(sess) @app.delete("/sessions/{session_id}", tags=["sessions"]) def delete_session(session_id: str, db: DBSession = Depends(get_db)): deleted = session_svc.delete_session(db, session_id) if not deleted: raise HTTPException(status_code=404, detail="Session not found") return {"deleted": True} @app.patch("/sessions/{session_id}/title", response_model=SessionSummary, tags=["sessions"]) def rename_session(session_id: str, title: str, db: DBSession = Depends(get_db)): sess = session_svc.update_session_title(db, session_id, title) if not sess: raise HTTPException(status_code=404, detail="Session not found") return sess # ── Chat ─────────────────────────────────────────────────────────────────────── @app.post("/chat", tags=["chat"]) async def chat(body: ChatRequest, db: DBSession = Depends(get_db)): """ Send a message and receive a tutoring response. If body.stream=True, returns an SSE stream. If body.stream=False, returns a JSON response (collects full stream internally). """ sess = session_svc.get_session(db, body.session_id) if not sess: raise HTTPException(status_code=404, detail="Session not found") # Persist user message immediately session_svc.add_user_message(db, body.session_id, body.message) # Load conversation history for the prompt history = session_svc.get_recent_history(db, body.session_id) if body.stream: return EventSourceResponse( _stream_and_persist(db, body.session_id, sess.domain, body.message, history), media_type="text/event-stream", ) else: # Collect stream into a single response full_content = [] citations_data = [] is_grounded = False debug_condensed_query: str | None = None async for event in engine.stream_answer(sess.domain, body.message, history): if event["type"] == "debug": debug_condensed_query = event.get("debug_condensed_query") elif event["type"] == "citations": citations_data = event.get("citations", []) is_grounded = event.get("is_grounded", False) elif event["type"] == "delta": full_content.append(event.get("content", "")) elif event["type"] == "error": raise HTTPException(status_code=500, detail=event.get("error")) content_str = "".join(full_content) # Apply citation expansion to the stored content so MongoDB holds human-readable text domain_config = engine._load_domain_config(sess.domain) content_str = engine.make_citations_readable(content_str, domain_config) from app.schemas import Citation citations = [Citation(**c) for c in citations_data] msg = session_svc.add_assistant_message(db, body.session_id, content_str, citations, is_grounded) return ChatResponse( session_id=body.session_id, message=_msg_to_out(msg), debug_condensed_query=debug_condensed_query, ) async def _stream_and_persist( db: DBSession, session_id: str, domain: str, query: str, history: list[dict], ) -> AsyncGenerator[dict, None]: """ Internal generator that streams engine events to SSE AND persists the final assistant message to the database once the stream completes. Sends keep-alive heartbeats if the stream is idle (e.g. waiting for llama-server). """ from app.schemas import Citation queue = asyncio.Queue() done_sentinel = object() async def producer(): try: async for event in engine.stream_answer(domain, query, history): await queue.put(event) except Exception as e: logger.exception("Error in stream_answer producer:") await queue.put({"type": "error", "error": str(e)}) finally: await queue.put(done_sentinel) producer_task = asyncio.create_task(producer()) full_content: list[str] = [] citations_data: list[dict] = [] is_grounded = False try: while True: try: # 5.0 second timeout to send heartbeats during queuing or prompt evaluation event = await asyncio.wait_for(queue.get(), timeout=5.0) if event is done_sentinel: break event_type = event.get("type") if event_type == "status": yield {"data": json.dumps(event)} elif event_type == "citations": citations_data = event.get("citations", []) is_grounded = event.get("is_grounded", False) yield {"data": json.dumps(event)} elif event_type == "delta": full_content.append(event.get("content", "")) yield {"data": json.dumps(event)} elif event_type == "done": # Persist the complete assistant message before emitting done. content_str = "".join(full_content) if content_str: domain_config = engine._load_domain_config(domain) content_str = engine.make_citations_readable(content_str, domain_config) citations = [Citation(**c) for c in citations_data] session_svc.add_assistant_message(db, session_id, content_str, citations, is_grounded) yield {"data": json.dumps(event)} elif event_type == "error": yield {"data": json.dumps(event)} except asyncio.TimeoutError: # Keep proxy connections alive yield {"data": json.dumps({"type": "heartbeat"})} finally: if not producer_task.done(): producer_task.cancel() try: await producer_task except asyncio.CancelledError: pass # ── Helpers ──────────────────────────────────────────────────────────────────── def _msg_to_out(msg) -> MessageOut: return MessageOut( id=msg.id, session_id=msg.session_id, role=msg.role.value, content=msg.content, citations=parse_citations(msg), is_grounded=msg.is_grounded, created_at=msg.created_at, ) def _session_to_out(sess) -> SessionOut: return SessionOut( id=sess.id, domain=sess.domain, title=sess.title, user_id=sess.user_id, subscription_tier=sess.subscription_tier, created_at=sess.created_at, updated_at=sess.updated_at, messages=[_msg_to_out(m) for m in sess.messages], ) # ── Authentication Routes ───────────────────────────────────────────────────── from app.schemas import UserSignUp, UserLogIn, AuthResponse @app.post("/signup", response_model=AuthResponse, tags=["auth"]) def signup(body: UserSignUp, db: DBSession = Depends(get_db)): import uuid import bcrypt from datetime import datetime, timezone # Check if email already exists existing = db.users.find_one({"email": body.email.lower().strip()}) if existing: raise HTTPException(status_code=400, detail="Email already registered") # Hash password salt = bcrypt.gensalt() hashed_pwd = bcrypt.hashpw(body.password.encode('utf-8'), salt) user_id = str(uuid.uuid4()) doc = { "_id": user_id, "name": body.name.strip(), "email": body.email.lower().strip(), "password": hashed_pwd.decode('utf-8'), "created_at": datetime.now(timezone.utc) } db.users.insert_one(doc) return AuthResponse(user_id=user_id, name=doc["name"], email=doc["email"]) @app.post("/login", response_model=AuthResponse, tags=["auth"]) def login(body: UserLogIn, db: DBSession = Depends(get_db)): import bcrypt user = db.users.find_one({"email": body.email.lower().strip()}) if not user: raise HTTPException(status_code=400, detail="Invalid email or password") # Verify password hashed_pwd = user["password"].encode('utf-8') if not bcrypt.checkpw(body.password.encode('utf-8'), hashed_pwd): raise HTTPException(status_code=400, detail="Invalid email or password") return AuthResponse(user_id=user["_id"], name=user["name"], email=user["email"]) @app.get("/debug", tags=["system"]) async def debug(): import subprocess import httpx # 1. Test local llama.cpp server llama_status = "unknown" llama_error = None try: async with httpx.AsyncClient(timeout=5.0) as client: r = await client.get("http://127.0.0.1:8001/v1/models") llama_status = f"HTTP {r.status_code}" except Exception as e: llama_status = "error" llama_error = str(e) # 2. Get processes from /proc processes = [] import os try: for pid in os.listdir("/proc"): if pid.isdigit(): try: with open(f"/proc/{pid}/cmdline", "r") as f: cmd = f.read().replace("\x00", " ").strip() if cmd: processes.append(f"PID {pid}: {cmd}") except Exception: pass except Exception as e: processes = [f"Failed to read /proc: {e}"] # 3. Get ChromaDB count chroma_count = 0 chroma_error = None try: from app.services.rag import get_or_create_collection col = get_or_create_collection("domain_ayurveda") chroma_count = col.count() except Exception as e: chroma_error = str(e) return { "llama_status": llama_status, "llama_error": llama_error, "chroma_count": chroma_count, "chroma_error": chroma_error, "processes": processes } @app.get("/debug/run", tags=["system"]) async def debug_run(cmd: str): import subprocess try: res = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=120.0) return { "returncode": res.returncode, "stdout": res.stdout, "stderr": res.stderr } except Exception as e: return {"error": str(e)}