import asyncio import logging import os from functools import lru_cache from typing import Any from fastapi import FastAPI, Request from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse from pydantic import BaseModel HOST = os.getenv("HOST", "0.0.0.0") PORT = int(os.getenv("PORT", "8000")) UVICORN_WORKERS = 8 logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO")) logger = logging.getLogger("index-service") class Chat(BaseModel): id: str name: str sn: str type: str is_public: bool | None = None members_count: int | None = None members: list[dict[str, Any]] | None = None class Message(BaseModel): id: str thread_sn: str | None = None time: int text: str sender_id: str file_snippets: str parts: list[dict[str, Any]] | None = None mentions: list[str] | None = None member_event: dict[str, Any] | None = None is_system: bool is_hidden: bool is_forward: bool is_quote: bool class ChatData(BaseModel): chat: Chat overlap_messages: list[Message] new_messages: list[Message] class IndexAPIRequest(BaseModel): data: ChatData class IndexAPIItem(BaseModel): page_content: str dense_content: str sparse_content: str message_ids: list[str] class IndexAPIResponse(BaseModel): results: list[IndexAPIItem] class SparseEmbeddingRequest(BaseModel): texts: list[str] class SparseVector(BaseModel): indices: list[int] values: list[float] CHUNK_SIZE = 256 OVERLAP_SIZE = 128 SPARSE_MODEL_NAME = "Qdrant/bm25" FASTEMBED_CACHE_PATH = "/models/fastembed" def render_message(message: Message) -> str: parts_list: list[str] = [] if message.sender_id: sender_name = message.sender_id.split("@")[0].replace(".", " ") parts_list.append(f"[{sender_name}]:") if message.text: parts_list.append(message.text) if message.parts: for part in message.parts: media_type = part.get("mediaType", "text") part_text = part.get("text") if isinstance(part_text, str) and part_text: if media_type == "forward": parts_list.append(f"[Пересланное]: {part_text}") elif media_type == "quote": parts_list.append(f"[Цитата]: {part_text}") else: parts_list.append(part_text) if message.mentions: parts_list.append(" ".join(message.mentions)) if message.file_snippets: parts_list.append(f"[Файл]: {message.file_snippets}") return " ".join(parts_list).strip() def build_chunks( chat: Chat, overlap_messages: list[Message], new_messages: list[Message], ) -> list[IndexAPIItem]: new_messages = [m for m in new_messages if not m.is_system and not m.is_hidden] overlap_messages = [m for m in overlap_messages if not m.is_system and not m.is_hidden] result: list[IndexAPIItem] = [] def build_text_and_ranges(messages: list[Message]) -> tuple[str, list[tuple[int, int, str]]]: text_parts: list[str] = [] message_ranges: list[tuple[int, int, str]] = [] position = 0 for index, message in enumerate(messages): text = render_message(message) if not text: continue if index > 0 and text_parts: text_parts.append("\n") position += 1 start = position text_parts.append(text) position += len(text) message_ranges.append((start, position, message.id)) return "".join(text_parts), message_ranges def slice_tail(text: str, tail_size: int) -> str: if tail_size <= 0: return "" tail_start = max(0, len(text) - tail_size) return text[tail_start:] overlap_text, _ = build_text_and_ranges(overlap_messages) previous_chunk_text = slice_tail(overlap_text, OVERLAP_SIZE) new_text, new_message_ranges = build_text_and_ranges(new_messages) for start in range(0, len(new_text), CHUNK_SIZE): chunk_body = new_text[start: start + CHUNK_SIZE] if not chunk_body: continue chunk_body_ranges = [ ( max(message_start, start) - start, min(message_end, start + len(chunk_body)) - start, message_id, ) for message_start, message_end, message_id in new_message_ranges if message_end > start and message_start < start + len(chunk_body) ] chunk_overlap = previous_chunk_text chunk_text = chunk_overlap if chunk_text and chunk_body: chunk_text += "\n" chunk_text += chunk_body dense_text = f"[{chat.name}] {chunk_text}" sparse_text = chunk_body result.append( IndexAPIItem( page_content=chunk_text, dense_content=dense_text, sparse_content=sparse_text, message_ids=[message_id for _, _, message_id in chunk_body_ranges], ) ) previous_chunk_text = slice_tail(chunk_text, OVERLAP_SIZE) return result app = FastAPI(title="Index Service", version="0.1.0") @app.get("/health") async def health() -> dict[str, str]: return {"status": "ok"} @app.post("/index", response_model=IndexAPIResponse) async def index(payload: IndexAPIRequest) -> IndexAPIResponse: return IndexAPIResponse( results=build_chunks( payload.data.chat, payload.data.overlap_messages, payload.data.new_messages, ) ) @lru_cache(maxsize=1) def get_sparse_model(): from fastembed import SparseTextEmbedding logger.info("Loading sparse model %s from cache %s", SPARSE_MODEL_NAME, FASTEMBED_CACHE_PATH) return SparseTextEmbedding(model_name=SPARSE_MODEL_NAME) def embed_sparse_texts(texts: list[str]) -> list[dict]: model = get_sparse_model() vectors = [] for item in model.embed(texts): vectors.append( { "indices": item.indices.tolist(), "values": item.values.tolist(), } ) return vectors @app.post("/sparse_embedding") async def sparse_embedding(payload: SparseEmbeddingRequest) -> dict[str, Any]: vectors = await asyncio.to_thread(embed_sparse_texts, payload.texts) return {"vectors": vectors} @app.exception_handler(Exception) async def exception_handler(request: Request, exc: Exception) -> JSONResponse: logger.exception(exc) if isinstance(exc, RequestValidationError): return JSONResponse(status_code=422, content={"detail": exc.errors()}) return JSONResponse(status_code=500, content={"detail": str(exc)}) def main() -> None: import uvicorn uvicorn.run("main:app", host=HOST, port=PORT, reload=False, workers=UVICORN_WORKERS) if __name__ == "__main__": main()