Hands-On Project — RAG Chatbot from Scratch

import json
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from fastapi import APIRouter, Request, Depends, HTTPException
from fastapi.responses import StreamingResponse
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session
from app.config import get_settings
from app.database import get_session, SessionLocal
from app.models import LgQuestions, LgRequestChunks, LgAnalyticsChat, LvContentChunks, LvExtractedContent
from app.services.chat import (
    detect_lang,
    expand_query,
    generate_query_embedding,
    retrieve_chunks,
    select_context_chunks,
    build_prompt,
    stream_response,
)
from app.services.sys_params import get_sys_param

BASE_DIR = Path(__file__).resolve().parent.parent
templates = Jinja2Templates(directory=str(BASE_DIR / "templates"))

router = APIRouter()


def _write_analytics_chat(**kwargs):
    s = SessionLocal()
    try:
        s.add(LgAnalyticsChat(**kwargs))
        s.commit()
    except Exception:
        s.rollback()
    finally:
        s.close()


class ChatMessage(BaseModel):
    message: str = Field(min_length=1, max_length=2000)
    search_mode: str = Field(default="semantic")
    rubrics: list[str] | None = Field(default=None)


@router.get("/chat")
def chat_page(request: Request, session: Session = Depends(get_session)):
    is_guest = not request.session.get("authenticated")
    guest_limit = get_sys_param("GUEST_LIMIT", int, 5)

    # Remember the conversation: reload this session's answered Q&A from the DB.
    history = []
    history_sources = {}
    session_id = request.session.get("session_id")
    cleared_at = request.session.get("chat_cleared_at")
    if session_id and not is_guest:
        query = (
            session.query(LgQuestions)
            .filter(LgQuestions.session_id == session_id)
            .filter(LgQuestions.answer_text.isnot(None))
        )
        if cleared_at:
            try:
                query = query.filter(LgQuestions.created_at > datetime.fromisoformat(cleared_at))
            except ValueError:
                pass
        rows = query.order_by(LgQuestions.created_at.asc()).all()

        sources_by_q = {}
        q_ids = [r.id for r in rows]
        if q_ids:
            src_rows = (
                session.query(
                    LgRequestChunks.question_id,
                    LvExtractedContent.title,
                    LvExtractedContent.url,
                    LvContentChunks.chunk_text,
                )
                .join(LvContentChunks, LvContentChunks.id == LgRequestChunks.chunk_id)
                .join(LvExtractedContent, LvExtractedContent.id == LvContentChunks.content_id)
                .filter(LgRequestChunks.question_id.in_(q_ids))
                .filter(LgRequestChunks.was_in_context == "Y")
                .order_by(LgRequestChunks.rank.asc())
                .all()
            )
            for qid, title, url, chunk_text in src_rows:
                sources_by_q.setdefault(qid, []).append({"title": title, "url": url, "text": chunk_text})

        for r in rows:
            history.append({"role": "user", "content": r.question_text})
            history.append({
                "role": "assistant",
                "content": r.answer_text,
                "question_id": str(r.id),
                "rated": r.mark_id is not None,
            })

        history_sources = {str(qid): srcs for qid, srcs in sources_by_q.items()}

    return templates.TemplateResponse("chat.html", {
        "request": request,
        "history": history,
        "history_sources": history_sources,
        "is_guest": is_guest,
        "guest_count": request.session.get("guest_count", 0) if is_guest else 0,
        "guest_limit": guest_limit,
    })


@router.post("/api/chat")
def chat_send(
    body: ChatMessage,
    request: Request,
    session: Session = Depends(get_session),
):
    is_guest = not request.session.get("authenticated")
    user_id = request.session.get("user_id")
    guest_count = request.session.get("guest_count", 0)
    guest_limit = get_sys_param("GUEST_LIMIT", int, 5)
    if is_guest and guest_count >= guest_limit:
        raise HTTPException(status_code=429, detail={
            "code": "guest_limit",
            "used": guest_count,
            "limit": guest_limit,
        })

    if "session_id" not in request.session:
        request.session["session_id"] = uuid.uuid4().hex
    session_id = request.session["session_id"]

    if request.session.get("role") in ("user", "analytics"):
        pending_q = (
            session.query(LgQuestions)
            .filter(LgQuestions.session_id == session_id)
            .filter(LgQuestions.answer_text.isnot(None))
            .filter(LgQuestions.mark_id.is_(None))
        )
        cleared_at = request.session.get("chat_cleared_at")
        if cleared_at:
            try:
                pending_q = pending_q.filter(LgQuestions.created_at > datetime.fromisoformat(cleared_at))
            except ValueError:
                pass
        pending = pending_q.order_by(LgQuestions.created_at.desc()).first()
        if pending is not None:
            raise HTTPException(
                status_code=403,
                detail={
                    "code": "feedback_required",
                    "pending_question_id": str(pending.id),
                },
            )

    question = body.message.strip()
    lang_code = detect_lang(question)

    effective_rubrics = body.rubrics or None
    if not effective_rubrics:
        _default = get_sys_param("DEFAULT_RUBRIC", str, None)
        if _default:
            effective_rubrics = [_default]

    expanded_query = None
    search_text = question
    if get_sys_param("QUERY_EXPANSION", str, "N") == "Y":
        expanded_query = expand_query(question)
        if expanded_query:
            search_text = expanded_query

    question_id = None
    if not is_guest:
        q_row = LgQuestions(
            session_id=session_id,
            user_id=user_id,
            question_text=question,
            expanded_query=expanded_query,
            detected_lang=lang_code,
            rubrics_used=effective_rubrics,
        )
        session.add(q_row)
        session.commit()
        session.refresh(q_row)
        question_id = q_row.id

    settings = get_settings()
    embed_model = settings.EMBED_MODEL
    chat_model = settings.CHAT_MODEL

    start_time = time.perf_counter()
    retrieval_start = start_time

    db_search_ms = 0
    from app.routes.admin import _load_rag_from_db
    rag = _load_rag_from_db(session)
    embed_start = time.perf_counter()
    embedding, embed_error = generate_query_embedding(search_text)
    embedding_ms = int((time.perf_counter() - embed_start) * 1000)
    if embedding is None:
        retrieval_ms = int((time.perf_counter() - retrieval_start) * 1000)
        _write_analytics_chat(
            question_id=question_id, user_id=user_id,
            retrieval_ms=retrieval_ms, embedding_ms=embedding_ms, db_search_ms=db_search_ms,
            generation_ms=0, total_ms=retrieval_ms,
            top_similarity=None, num_chunks_retrieved=0,
            embed_model=embed_model, chat_model=chat_model,
            status="embed_error", error_message=embed_error,
        )
        raise HTTPException(status_code=503, detail="Embedding service unavailable")
    db_start = time.perf_counter()
    chunks = retrieve_chunks(
        session,
        embedding,
        top_k=rag["top_k"],
        threshold=rag["similarity_threshold"],
        lang_code=lang_code,
        rubrics=effective_rubrics,
    )
    db_search_ms = int((time.perf_counter() - db_start) * 1000)
    retrieval_ms = int((time.perf_counter() - retrieval_start) * 1000)

    if not chunks:
        _write_analytics_chat(
            question_id=question_id, user_id=user_id,
            retrieval_ms=retrieval_ms, embedding_ms=embedding_ms, db_search_ms=db_search_ms,
            generation_ms=0, total_ms=retrieval_ms,
            top_similarity=None, num_chunks_retrieved=0,
            embed_model=embed_model, chat_model=chat_model,
            status="no_chunks", error_message=None,
        )
        raise HTTPException(status_code=503, detail="No relevant data found")

    used = select_context_chunks(chunks, question)
    used_ids = {c["chunk_id"] for c in used}

    if question_id is not None:
        threshold_used = rag["similarity_threshold"]
        for rank_idx, c in enumerate(chunks, start=1):
            session.add(LgRequestChunks(
                question_id=question_id,
                chunk_id=c["chunk_id"],
                similarity_score=c["similarity"],
                rank=rank_idx,
                threshold_used=threshold_used,
                role="semantic",
                was_in_context="Y" if c["chunk_id"] in used_ids else "N",
            ))
        try:
            session.commit()
        except Exception:
            session.rollback()

    messages = build_prompt(chunks, question, None)

    sources = [
        {"title": c["title"], "url": c["url"], "text": c["chunk_text"]}
        for c in used
    ]

    if is_guest:
        request.session["guest_count"] = guest_count + 1
        new_guest_count = guest_count + 1
    else:
        new_guest_count = 0

    def event_stream():
        generation_start = time.perf_counter()
        full_response = []
        gen_status = "ok"
        gen_error = None
        try:
            for token in stream_response(messages, chunks=chunks, question=question):
                full_response.append(token)
                yield f"data: {json.dumps({'token': token})}\n\n"
        except Exception as e:
            gen_status = "gen_error"
            gen_error = str(e)
        generation_ms = int((time.perf_counter() - generation_start) * 1000)
        total_ms = int((time.perf_counter() - start_time) * 1000)

        if question_id is not None:
            s2 = SessionLocal()
            try:
                s2.query(LgQuestions).filter(LgQuestions.id == question_id).update({
                    "answer_text": "".join(full_response),
                    "answered_at": datetime.now(timezone.utc),
                })
                s2.commit()
            except Exception:
                s2.rollback()
            finally:
                s2.close()

        _write_analytics_chat(
            question_id=question_id, user_id=user_id,
            retrieval_ms=retrieval_ms, embedding_ms=embedding_ms, db_search_ms=db_search_ms,
            generation_ms=generation_ms, total_ms=total_ms,
            top_similarity=chunks[0]["similarity"] if chunks else None,
            num_chunks_retrieved=len(chunks),
            embed_model=embed_model, chat_model=chat_model,
            status=gen_status, error_message=gen_error,
        )

        done_payload = {"done": True, "sources": sources}
        if question_id is not None:
            done_payload["question_id"] = str(question_id)
        if is_guest:
            done_payload["guest_count"] = new_guest_count
            done_payload["guest_limit"] = guest_limit
        yield f"data: {json.dumps(done_payload)}\n\n"

    return StreamingResponse(event_stream(), media_type="text/event-stream")


@router.post("/api/chat/clear")
def chat_clear(request: Request):
    request.session["chat_cleared_at"] = datetime.now(timezone.utc).isoformat()
    return {"status": "ok"}

import json
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path

from fastapi import FastAPI, Request, Depends, HTTPException, Response
from fastapi.responses import StreamingResponse
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel, ConfigDict, Field
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session
from starlette.middleware.sessions import SessionMiddleware

from app.config import get_settings
from app.database import get_session, SessionLocal
from app.models import (
    LgQuestions,
    LgRequestChunks,
    LgAnalyticsChat,
    LvFeedbackMarks,
)
from app.services.chat import (
    detect_lang,
    expand_query,
    generate_query_embedding,
    retrieve_chunks,
    build_prompt,
    stream_response,
)
from app.services.sys_params import get_sys_param
from public.rate_limit import check as rate_check

BASE_DIR = Path(__file__).resolve().parent
templates = Jinja2Templates(directory=str(BASE_DIR / "templates"))


class AskBody(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True)
    message: str = Field(min_length=1, max_length=2000)


class FeedbackBody(BaseModel):
    question_id: uuid.UUID
    mark_id: int


def _client_ip(request: Request) -> str:
    fwd = request.headers.get("x-forwarded-for")
    if fwd:
        return fwd.split(",")[0].strip()
    return request.client.host if request.client else "unknown"


def _write_analytics(**kwargs):
    s = SessionLocal()
    try:
        s.add(LgAnalyticsChat(**kwargs))
        s.commit()
    except Exception:
        s.rollback()
    finally:
        s.close()


def create_app() -> FastAPI:
    settings = get_settings()
    if not settings.SESSION_SECRET:
        raise RuntimeError("SESSION_SECRET is not set")

    application = FastAPI(title="Nihol — Помощник", docs_url=None, redoc_url=None)
    application.add_middleware(SessionMiddleware, secret_key=settings.SESSION_SECRET)

    @application.get("/")
    def ask_page(request: Request):
        return templates.TemplateResponse("ask.html", {"request": request})

    @application.get("/api/feedback-marks")
    def get_marks(session: Session = Depends(get_session)):
        rows = (
            session.query(LvFeedbackMarks)
            .order_by(LvFeedbackMarks.display_order)
            .all()
        )
        return {"data": [
            {"id": r.id, "code": r.code, "name": r.name, "color_hex": r.color_hex}
            for r in rows
        ]}

    @application.post("/api/feedback")
    def post_feedback(
        body: FeedbackBody,
        session: Session = Depends(get_session),
    ):
        q = session.query(LgQuestions).filter(LgQuestions.id == body.question_id).first()
        if q is None:
            raise HTTPException(status_code=404, detail="Question not found")
        if q.mark_id is not None:
            raise HTTPException(status_code=409, detail="Already rated")
        mark = session.query(LvFeedbackMarks.id).filter(LvFeedbackMarks.id == body.mark_id).first()
        if mark is None:
            raise HTTPException(status_code=400, detail="Invalid mark_id")
        q.mark_id = body.mark_id
        try:
            session.commit()
        except IntegrityError:
            session.rollback()
            raise HTTPException(status_code=409, detail="Already rated")
        return {"status": "ok"}

    @application.get("/api/feedback/pending")
    def get_pending(request: Request, session: Session = Depends(get_session)):
        session_id = request.session.get("session_id")
        if not session_id:
            return Response(status_code=204)
        pending = (
            session.query(LgQuestions)
            .filter(LgQuestions.session_id == session_id)
            .filter(LgQuestions.answer_text.isnot(None))
            .filter(LgQuestions.mark_id.is_(None))
            .order_by(LgQuestions.created_at.desc())
            .first()
        )
        if pending is None:
            return Response(status_code=204)
        return {
            "question_id": str(pending.id),
            "question_text": pending.question_text,
            "answer_text": pending.answer_text,
        }

    @application.post("/api/ask")
    def ask_send(
        body: AskBody,
        request: Request,
        session: Session = Depends(get_session),
    ):
        ip = _client_ip(request)
        rate_limit = get_sys_param("PUBLIC_RATE_LIMIT_PER_MIN", int, settings.PUBLIC_RATE_LIMIT_PER_MIN or 10)
        if not rate_check(ip, rate_limit):
            raise HTTPException(status_code=429, detail="Too many requests")

        if "session_id" not in request.session:
            request.session["session_id"] = uuid.uuid4().hex
        session_id = request.session["session_id"]

        # Grading is optional — no gate. Users may rate the previous answer or just continue.
        question = body.message.strip()
        lang_code = detect_lang(question)

        public_rubrics = get_sys_param("PUBLIC_RUBRICS", str, "")
        rubrics = [r.strip() for r in public_rubrics.split(",") if r.strip()]
        if not rubrics:
            fallback = get_sys_param("DEFAULT_RUBRIC", str, "project")
            rubrics = [fallback] if fallback else ["project"]

        expanded_query = None
        search_text = question
        if get_sys_param("QUERY_EXPANSION", str, "N") == "Y":
            expanded_query = expand_query(question)
            if expanded_query:
                search_text = expanded_query

        settings_local = get_settings()
        embed_model = settings_local.EMBED_MODEL
        chat_model = settings_local.CHAT_MODEL

        q_row = LgQuestions(
            session_id=session_id,
            question_text=question,
            expanded_query=expanded_query,
            detected_lang=lang_code,
            rubrics_used=rubrics,
        )
        session.add(q_row)
        session.commit()
        session.refresh(q_row)
        question_id = q_row.id

        start_time = time.perf_counter()
        retrieval_start = start_time
        db_search_ms = 0

        embed_start = time.perf_counter()
        embedding, embed_error = generate_query_embedding(search_text)
        embedding_ms = int((time.perf_counter() - embed_start) * 1000)
        if embedding is None:
            retrieval_ms = int((time.perf_counter() - retrieval_start) * 1000)
            _write_analytics(
                question_id=question_id, user_id=None,
                retrieval_ms=retrieval_ms, embedding_ms=embedding_ms, db_search_ms=db_search_ms,
                generation_ms=0, total_ms=retrieval_ms,
                top_similarity=None, num_chunks_retrieved=0,
                embed_model=embed_model, chat_model=chat_model,
                status="embed_error", error_message=embed_error,
            )
            raise HTTPException(status_code=503, detail="Сервис недоступен")

        from app.routes.admin import _load_rag_from_db
        rag = _load_rag_from_db(session)
        db_start = time.perf_counter()
        chunks = retrieve_chunks(
            session,
            embedding,
            top_k=rag["top_k"],
            threshold=rag["similarity_threshold"],
            lang_code=lang_code,
            rubrics=rubrics,
        )
        db_search_ms = int((time.perf_counter() - db_start) * 1000)
        retrieval_ms = int((time.perf_counter() - retrieval_start) * 1000)

        if not chunks:
            _write_analytics(
                question_id=question_id, user_id=None,
                retrieval_ms=retrieval_ms, embedding_ms=embedding_ms, db_search_ms=db_search_ms,
                generation_ms=0, total_ms=retrieval_ms,
                top_similarity=None, num_chunks_retrieved=0,
                embed_model=embed_model, chat_model=chat_model,
                status="no_chunks", error_message=None,
            )
            raise HTTPException(status_code=404, detail="По этому вопросу информация не найдена")

        threshold_used = rag["similarity_threshold"]
        for rank_idx, c in enumerate(chunks, start=1):
            session.add(LgRequestChunks(
                question_id=question_id,
                chunk_id=c["chunk_id"],
                similarity_score=c["similarity"],
                rank=rank_idx,
                threshold_used=threshold_used,
                role="semantic",
                was_in_context="Y",
            ))
        try:
            session.commit()
        except Exception:
            session.rollback()

        messages = build_prompt(chunks, question, None)

        def event_stream():
            generation_start = time.perf_counter()
            full_response = []
            gen_status = "ok"
            gen_error = None
            try:
                for token in stream_response(messages, chunks=chunks, question=question):
                    full_response.append(token)
                    yield f"data: {json.dumps({'token': token})}\n\n"
            except Exception as e:
                gen_status = "gen_error"
                gen_error = str(e)

            generation_ms = int((time.perf_counter() - generation_start) * 1000)
            total_ms = int((time.perf_counter() - start_time) * 1000)

            s2 = SessionLocal()
            try:
                s2.query(LgQuestions).filter(LgQuestions.id == question_id).update({
                    "answer_text": "".join(full_response),
                    "answered_at": datetime.now(timezone.utc),
                })
                s2.commit()
            except Exception:
                s2.rollback()
            finally:
                s2.close()

            _write_analytics(
                question_id=question_id, user_id=None,
                retrieval_ms=retrieval_ms, embedding_ms=embedding_ms, db_search_ms=db_search_ms,
                generation_ms=generation_ms, total_ms=total_ms,
                top_similarity=chunks[0]["similarity"] if chunks else None,
                num_chunks_retrieved=len(chunks),
                embed_model=embed_model, chat_model=chat_model,
                status=gen_status, error_message=gen_error,
            )

            yield f"data: {json.dumps({'done': True, 'question_id': str(question_id)})}\n\n"

        return StreamingResponse(event_stream(), media_type="text/event-stream")

    return application


app = create_app()