163 lines
6.3 KiB
Python
163 lines
6.3 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
from collections.abc import AsyncGenerator, AsyncIterator
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Annotated
|
|
|
|
from anyio import CancelScope
|
|
from fastapi import Depends, FastAPI, HTTPException, Path as ApiPath, Request
|
|
from fastapi.concurrency import run_in_threadpool
|
|
from fastapi.responses import Response, StreamingResponse
|
|
from fastapi.staticfiles import StaticFiles
|
|
from pydantic import BaseModel
|
|
|
|
from .groq_client import GenerationError, InsufficientContent, stream_analysis
|
|
from .storage import Note, NoteStore
|
|
|
|
logger = logging.getLogger(__name__)
|
|
NDJSON = "application/x-ndjson"
|
|
|
|
|
|
class NoteInput(BaseModel):
|
|
title: str
|
|
content: str
|
|
|
|
def validated_text(self) -> tuple[str, str]:
|
|
title, content = self.title.strip(), self.content.strip()
|
|
if not 1 <= len(title) <= 120:
|
|
raise HTTPException(422, "Заголовок должен содержать от 1 до 120 символов.")
|
|
if not 1 <= len(content) <= 12_000:
|
|
raise HTTPException(422, "Текст заметки должен содержать от 1 до 12 000 символов.")
|
|
return title, content
|
|
|
|
|
|
NoteId = Annotated[int, ApiPath(ge=1, le=2**63 - 1)]
|
|
|
|
|
|
@dataclass
|
|
class Generation:
|
|
note: Note
|
|
store: NoteStore
|
|
upstream: AsyncGenerator[str | tuple[str, list[str]], None]
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
app.state.store = NoteStore(Path(os.getenv("NOTES_DB_PATH", "./data/notes.sqlite3")))
|
|
await run_in_threadpool(app.state.store.initialize)
|
|
app.state.generating = set()
|
|
app.state.groq_key = os.getenv("GROQ_API_KEY", "").strip()
|
|
app.state.groq_model = os.getenv("GROQ_MODEL", "").strip() or "openai/gpt-oss-20b"
|
|
yield
|
|
|
|
|
|
app = FastAPI(title="AI-заметки", lifespan=lifespan)
|
|
|
|
|
|
@app.get("/api/notes")
|
|
def list_notes(request: Request) -> list[Note]:
|
|
return request.app.state.store.list_notes()
|
|
|
|
|
|
@app.post("/api/notes", status_code=201)
|
|
def create_note(payload: NoteInput, request: Request) -> Note:
|
|
title, content = payload.validated_text()
|
|
return request.app.state.store.create_note(title, content)
|
|
|
|
|
|
@app.put("/api/notes/{note_id}")
|
|
def update_note(note_id: NoteId, payload: NoteInput, request: Request) -> Note:
|
|
title, content = payload.validated_text()
|
|
note = request.app.state.store.update_note(note_id, title, content)
|
|
if note is None:
|
|
raise HTTPException(404, "Заметка не найдена.")
|
|
return note
|
|
|
|
|
|
@app.delete("/api/notes/{note_id}", status_code=204, response_class=Response)
|
|
def delete_note(note_id: NoteId, request: Request) -> None:
|
|
request.app.state.store.delete_note(note_id)
|
|
|
|
|
|
async def reserve_generation(note_id: NoteId, request: Request) -> AsyncIterator[Generation]:
|
|
state = request.app.state
|
|
note = await run_in_threadpool(state.store.get_note, note_id)
|
|
if note is None:
|
|
raise HTTPException(404, "Заметка не найдена.")
|
|
if not state.groq_key:
|
|
raise HTTPException(503, "На сервере не задан ключ Groq. Добавьте GROQ_API_KEY в .env.")
|
|
if note_id in state.generating:
|
|
raise HTTPException(409, "Для этой заметки уже идёт генерация.")
|
|
|
|
upstream = stream_analysis(
|
|
api_key=state.groq_key,
|
|
model=state.groq_model,
|
|
title=note.title,
|
|
content=note.content,
|
|
)
|
|
state.generating.add(note_id)
|
|
try:
|
|
yield Generation(note=note, store=state.store, upstream=upstream)
|
|
finally:
|
|
# This dependency stays active until the response ends, including a disconnect.
|
|
try:
|
|
with CancelScope(shield=True):
|
|
await upstream.aclose()
|
|
finally:
|
|
state.generating.discard(note_id)
|
|
|
|
|
|
def stream_event(event_type: str, **fields: object) -> str:
|
|
return json.dumps({"type": event_type, **fields}, ensure_ascii=False) + "\n"
|
|
|
|
|
|
async def analysis_events(generation: Generation) -> AsyncIterator[str]:
|
|
try:
|
|
async for part in generation.upstream:
|
|
if isinstance(part, str):
|
|
yield stream_event("delta", text=part)
|
|
continue
|
|
summary, tags = part
|
|
note = await run_in_threadpool(
|
|
generation.store.save_analysis, generation.note, summary, tags
|
|
)
|
|
if note is None:
|
|
raise GenerationError("Заметка изменилась или была удалена. Обновите страницу.")
|
|
yield stream_event("done", note=note.model_dump())
|
|
return
|
|
raise GenerationError("Groq завершил ответ без результата. Попробуйте ещё раз.")
|
|
except GenerationError as error:
|
|
yield stream_event("error", message=str(error))
|
|
except InsufficientContent:
|
|
try:
|
|
if generation.note.summary is not None or generation.note.tags:
|
|
await run_in_threadpool(generation.store.save_analysis, generation.note, None, [])
|
|
except Exception as error:
|
|
logger.error("Could not clear previous analysis (%s)", type(error).__name__)
|
|
yield stream_event("error", message="Не удалось обновить резюме. Попробуйте ещё раз.")
|
|
return
|
|
yield stream_event("insufficient", message="Добавьте в заметку тему, факт или задачу, чтобы получить резюме.")
|
|
except Exception as error:
|
|
logger.error("Generation failed (%s)", type(error).__name__)
|
|
yield stream_event("error", message="Не удалось завершить генерацию. Попробуйте ещё раз.")
|
|
|
|
|
|
@app.post(
|
|
"/api/notes/{note_id}/generate",
|
|
response_class=StreamingResponse,
|
|
responses={200: {"content": {NDJSON: {}}}},
|
|
)
|
|
async def generate(generation: Annotated[Generation, Depends(reserve_generation)]) -> StreamingResponse:
|
|
return StreamingResponse(
|
|
analysis_events(generation),
|
|
media_type=NDJSON,
|
|
headers={"Cache-Control": "no-cache, no-transform", "X-Accel-Buffering": "no"},
|
|
)
|
|
|
|
|
|
if web_dist := os.getenv("NOTES_WEB_DIST"):
|
|
app.mount("/", StaticFiles(directory=web_dist, html=True), name="frontend")
|