285 lines
10 KiB
Python
285 lines
10 KiB
Python
"""سرویس چت‌بات سایت آرامش — FastAPI با استریم SSE و بازیابی ترکیبی روی محتوای دیتابیس."""
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import time
|
||
import uuid
|
||
from collections import defaultdict, deque
|
||
from contextlib import asynccontextmanager
|
||
|
||
from fastapi import FastAPI, HTTPException, Query, Request
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from fastapi.responses import StreamingResponse
|
||
|
||
from langchain_core.messages import AIMessage, HumanMessage
|
||
|
||
from . import content as content_mod
|
||
from . import pipeline
|
||
from . import vectorstore
|
||
from .config import settings
|
||
from .indexing import INDEX, sync_index
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||
)
|
||
log = logging.getLogger("chatbot")
|
||
|
||
|
||
# --- چرخه حیات: ساخت ایندکس در شروع + بازسازی خودکار وقتی محتوا عوض شود ---
|
||
|
||
_refresh_task: asyncio.Task | None = None
|
||
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(_app: FastAPI):
|
||
try:
|
||
await asyncio.to_thread(sync_index)
|
||
except Exception as e:
|
||
log.critical("Initial index sync failed: %s — chat will be unavailable until reindex", e)
|
||
global _refresh_task
|
||
_refresh_task = asyncio.create_task(_refresh_loop())
|
||
yield
|
||
_refresh_task.cancel()
|
||
|
||
|
||
async def _refresh_loop() -> None:
|
||
interval = settings.index_refresh_seconds
|
||
if interval <= 0:
|
||
return
|
||
while True:
|
||
await asyncio.sleep(interval)
|
||
try:
|
||
# همگام‌سازی افزایشی است — فقط سطرهای جدید/تغییریافته امبد می‌شوند
|
||
await asyncio.to_thread(sync_index)
|
||
except Exception as e:
|
||
log.warning("Background index refresh failed: %s", e)
|
||
|
||
|
||
api = FastAPI(
|
||
title="Aramesh Site Chatbot",
|
||
version="1.0.0",
|
||
description="دستیار گفتگوی محتوای سایت — بازیابی ترکیبی (pgvector + BM25) با ری‌رنکر BGE",
|
||
lifespan=lifespan,
|
||
)
|
||
|
||
api.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_credentials=False,
|
||
allow_methods=["GET", "POST"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
|
||
# --- محدودسازی نرخ درخواست سبک در حافظه (پنجره لغزان per-IP) ---
|
||
|
||
_hits: dict[str, deque] = defaultdict(deque)
|
||
_last_sweep = 0.0
|
||
|
||
|
||
def _rate_limited(ip: str) -> bool:
|
||
global _last_sweep
|
||
now = time.time()
|
||
# پاک‌سازی دوره‌ای کلیدهای بی‌مصرف تا دیکشنری بی‌نهایت بزرگ نشود
|
||
if now - _last_sweep > settings.rate_limit_window:
|
||
stale = [k for k, q in _hits.items() if not q or now - q[-1] > settings.rate_limit_window]
|
||
for k in stale:
|
||
_hits.pop(k, None)
|
||
_last_sweep = now
|
||
window = _hits[ip]
|
||
while window and now - window[0] > settings.rate_limit_window:
|
||
window.popleft()
|
||
if len(window) >= settings.rate_limit_messages:
|
||
return True
|
||
window.append(now)
|
||
return False
|
||
|
||
|
||
def _client_ip(request: Request | None) -> str:
|
||
if request is None:
|
||
return "?"
|
||
forwarded = request.headers.get("x-forwarded-for")
|
||
return forwarded.split(",")[0].strip() if forwarded else (request.client.host if request.client else "?")
|
||
|
||
|
||
# --- کمکی‌های SSE ---
|
||
|
||
def _sse(payload: dict) -> str:
|
||
return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
|
||
|
||
|
||
def _chunk_text(chunk) -> str:
|
||
"""استخراج متن از قطعه استریم LLM (رشته یا لیست بخش‌ها)."""
|
||
content = getattr(chunk, "content", "")
|
||
if isinstance(content, str):
|
||
return content
|
||
if isinstance(content, list):
|
||
return "".join(p.get("text", "") for p in content if isinstance(p, dict))
|
||
return ""
|
||
|
||
|
||
async def _stream_answer(message: str, thread_id: str):
|
||
if not INDEX.ready:
|
||
yield _sse({"error": "ایندکس محتوا هنوز آماده نیست؛ چند لحظه بعد دوباره تلاش کنید.", "done": True})
|
||
return
|
||
if pipeline.get_llm() is None:
|
||
yield _sse({"error": pipeline.LLM_UNAVAILABLE, "done": True})
|
||
return
|
||
|
||
config = {"configurable": {"thread_id": thread_id}}
|
||
inputs = {"messages": [HumanMessage(content=message)]}
|
||
try:
|
||
async for event in pipeline.graph.astream_events(inputs, config=config, version="v2"):
|
||
# فقط توکن‌های پاسخ نهایی — نه فراخوانی‌های کمکی (بازنویسی/گسترش پرسش)
|
||
if event["event"] == "on_chat_model_stream" and "generate" in (event.get("tags") or []):
|
||
text = _chunk_text(event["data"]["chunk"])
|
||
if text:
|
||
yield _sse({"content": text, "done": False})
|
||
yield _sse({"content": "", "done": True})
|
||
except Exception as e:
|
||
log.exception("Streaming failed")
|
||
yield _sse({"error": f"خطا در تولید پاسخ: {e}", "done": True})
|
||
|
||
|
||
# --- مسیرها ---
|
||
|
||
@api.get("/chat/healthz")
|
||
async def healthz():
|
||
return {
|
||
"status": "ok",
|
||
"index_ready": INDEX.ready,
|
||
"documents": len(INDEX.docs),
|
||
"vectors": await asyncio.to_thread(vectorstore.vector_count),
|
||
"llm_configured": bool(settings.llm_base_url),
|
||
"reranker_enabled": settings.reranker_enabled,
|
||
}
|
||
|
||
|
||
@api.get("/chat/suggestions")
|
||
async def suggestions():
|
||
return {"questions": _get_suggestions()}
|
||
|
||
|
||
@api.get("/chat/eval")
|
||
async def eval_question(message: str = Query(...)):
|
||
"""مخصوص ارزیابی RAGAS — پاسخ + اسناد بازیابی‌شده، با thread تازه."""
|
||
if not INDEX.ready:
|
||
return {"error": "index not ready", "question": message, "answer": "", "contexts": []}
|
||
if pipeline.get_llm() is None:
|
||
return {"error": "llm not configured", "question": message, "answer": "", "contexts": []}
|
||
|
||
thread_id = f"eval_{uuid.uuid4().hex[:10]}"
|
||
config = {"configurable": {"thread_id": thread_id}}
|
||
inputs = {"messages": [HumanMessage(content=message)]}
|
||
final_state = await pipeline.graph.ainvoke(inputs, config=config)
|
||
|
||
answer = ""
|
||
for msg in reversed(final_state.get("messages", [])):
|
||
if isinstance(msg, AIMessage):
|
||
answer = msg.content if isinstance(msg.content, str) else str(msg.content)
|
||
break
|
||
return {
|
||
"question": message,
|
||
"contextualized_query": final_state.get("contextualized_query") or message,
|
||
"answer": answer,
|
||
"contexts": final_state.get("documents", []),
|
||
"thread_id": thread_id,
|
||
}
|
||
|
||
|
||
@api.post("/chat/reindex")
|
||
async def reindex(request: Request):
|
||
"""همگام‌سازی اجباری (امبد دوباره همه سطرها) — با توکن محافظت می‌شود."""
|
||
if settings.reindex_token and request.headers.get("x-reindex-token") != settings.reindex_token:
|
||
raise HTTPException(status_code=403, detail="invalid reindex token")
|
||
stats = await asyncio.to_thread(sync_index, True)
|
||
return {"status": "ok", "documents": len(INDEX.docs), "stats": stats}
|
||
|
||
|
||
# سازگاری با eval_ragas.py قدیمی که آدرس را /chat/{username}/eval می‌سازد
|
||
|
||
@api.get("/chat/{scope}/stream")
|
||
async def stream_chat_scoped(scope: str, message: str = Query(...), thread_id: str = Query(default=""), request: Request = None):
|
||
return await stream_chat(message=message, thread_id=thread_id, request=request)
|
||
|
||
|
||
@api.get("/chat/{scope}/eval")
|
||
async def eval_scoped(scope: str, message: str = Query(...)):
|
||
return await eval_question(message=message)
|
||
|
||
|
||
@api.get("/chat/{scope}/suggestions")
|
||
async def suggestions_scoped(scope: str):
|
||
return {"questions": _get_suggestions()}
|
||
|
||
|
||
@api.get("/chat/stream")
|
||
async def stream_chat(message: str = Query(...), thread_id: str = Query(default=""), request: Request = None):
|
||
"""استریم توکن‌به‌توکن پاسخ با SSE."""
|
||
if _rate_limited(_client_ip(request)):
|
||
raise HTTPException(status_code=429, detail="تعداد پیام‌های شما زیاد است؛ کمی بعد دوباره تلاش کنید.")
|
||
|
||
thread_id = thread_id or f"t_{uuid.uuid4().hex[:10]}"
|
||
|
||
async def generate():
|
||
async for chunk in _stream_answer(message, thread_id):
|
||
yield chunk
|
||
|
||
return StreamingResponse(
|
||
generate(),
|
||
media_type="text/event-stream",
|
||
headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
|
||
)
|
||
|
||
|
||
# --- پیشنهاد پرسش‌ها (کش‌شده) ---
|
||
|
||
_suggestions_cache: tuple[float, list[str]] = (0.0, [])
|
||
|
||
|
||
def _get_suggestions() -> list[str]:
|
||
global _suggestions_cache
|
||
now = time.time()
|
||
if _suggestions_cache[1] and now - _suggestions_cache[0] < settings.suggestions_ttl:
|
||
return _suggestions_cache[1]
|
||
|
||
questions: list[str] = []
|
||
if INDEX.ready:
|
||
faqs = [d for d in INDEX.docs if d["type"] == "faq"][:8]
|
||
titles = [d["title"] for d in INDEX.docs if d["type"] == "article"][:10]
|
||
model = pipeline.get_llm()
|
||
if model and (faqs or titles):
|
||
material = "\n".join([d["text"][:200] for d in faqs] + titles)
|
||
try:
|
||
prompt = (
|
||
"بر اساس محتوای سایت «گروه ساختمانی آرامش» ۴ پرسش کوتاه و جذاب بنویس که بازدیدکننده "
|
||
"احتمالاً می‌پرسد. هر پرسش حداکثر ۱۰ کلمه و به فارسی.\n"
|
||
'فقط آرایه JSON رشته‌ای برگردان: ["پرسش ۱", "پرسش ۲", ...]\n\n'
|
||
f"محتوا:\n{material[:2000]}"
|
||
)
|
||
resp = model.invoke([HumanMessage(content=prompt)])
|
||
text = resp.content if isinstance(resp.content, str) else ""
|
||
if text.startswith("```"):
|
||
text = text.split("```")[1]
|
||
text = text[4:] if text.startswith("json") else text
|
||
import re
|
||
match = re.search(r"\[.*\]", text, re.DOTALL)
|
||
if match:
|
||
questions = [q.strip() for q in json.loads(match.group()) if isinstance(q, str) and q.strip()][:4]
|
||
except Exception as e:
|
||
log.warning("Suggestion generation failed: %s", e)
|
||
if not questions and faqs:
|
||
# جایگزین: خود پرسش‌های متداول
|
||
questions = [d["text"].split("\n")[0].removeprefix("پرسش: ") for d in faqs[:4]]
|
||
|
||
_suggestions_cache = (now, questions)
|
||
return questions
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import uvicorn
|
||
|
||
uvicorn.run(api, host="0.0.0.0", port=8000)
|