83 lines
2.4 KiB
Python
83 lines
2.4 KiB
Python
import sqlite3
|
|
import json
|
|
import os
|
|
from fastapi import FastAPI, Body
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from sentence_transformers import SentenceTransformer
|
|
import uvicorn
|
|
|
|
app = FastAPI()
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
DB_FILE = os.path.join(BASE_DIR, "assets.db")
|
|
MODEL_DIR = os.path.join(BASE_DIR, "local_model_miniLM")
|
|
|
|
if not os.path.exists(MODEL_DIR):
|
|
model = SentenceTransformer('all-MiniLM-L6-v2')
|
|
model.save(MODEL_DIR)
|
|
else:
|
|
model = SentenceTransformer(MODEL_DIR)
|
|
|
|
def init_db():
|
|
conn = sqlite3.connect(DB_FILE)
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS documents
|
|
(
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
title TEXT UNIQUE,
|
|
content TEXT,
|
|
embedding TEXT
|
|
)
|
|
""")
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
init_db()
|
|
|
|
@app.post("/save-document")
|
|
async def save_document(data: dict = Body(...)):
|
|
title = data.get("title")
|
|
content = data.get("content")
|
|
|
|
text_to_vector = f"{title} {str(content)}"
|
|
vector = model.encode(text_to_vector).tolist()
|
|
|
|
conn = sqlite3.connect(DB_FILE)
|
|
try:
|
|
conn.execute("""
|
|
INSERT INTO documents (title, content, embedding)
|
|
VALUES (?, ?, ?) ON CONFLICT(title) DO
|
|
UPDATE SET
|
|
content=excluded.content,
|
|
embedding=excluded.embedding
|
|
""", (title, json.dumps(content), json.dumps(vector)))
|
|
conn.commit()
|
|
return {"status": "success"}
|
|
except Exception as e:
|
|
return {"status": "error", "message": str(e)}
|
|
finally:
|
|
conn.close()
|
|
|
|
@app.get("/load-document")
|
|
async def load_document(title: str = None):
|
|
conn = sqlite3.connect(DB_FILE)
|
|
if title:
|
|
row = conn.execute("SELECT title, content FROM documents WHERE title = ?", (title,)).fetchone()
|
|
else:
|
|
row = conn.execute("SELECT title, content FROM documents ORDER BY id DESC LIMIT 1").fetchone()
|
|
conn.close()
|
|
|
|
if row:
|
|
return {"title": row[0], "content": json.loads(row[1])}
|
|
return {"error": "Nie znaleziono dokumentu"}
|
|
|
|
if __name__ == "__main__":
|
|
uvicorn.run(app, host="127.0.0.1", port=8000)
|