135 lines
3.9 KiB
Python
135 lines
3.9 KiB
Python
import sqlite3
|
|
import os
|
|
import numpy as np
|
|
from fastapi import FastAPI, Body, HTTPException
|
|
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 get_db_connection():
|
|
conn = sqlite3.connect(DB_FILE)
|
|
|
|
conn.execute("PRAGMA journal_mode=WAL;")
|
|
conn.row_factory = sqlite3.Row
|
|
return conn
|
|
|
|
|
|
def init_db():
|
|
with get_db_connection() as conn:
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS documents
|
|
(
|
|
id
|
|
INTEGER
|
|
PRIMARY
|
|
KEY
|
|
AUTOINCREMENT,
|
|
title
|
|
TEXT
|
|
UNIQUE,
|
|
content
|
|
BLOB,
|
|
content_type
|
|
TEXT,
|
|
embedding
|
|
BLOB
|
|
)
|
|
""")
|
|
conn.commit()
|
|
|
|
|
|
init_db()
|
|
|
|
|
|
@app.post("/save-document")
|
|
async def save_document(
|
|
title: str = Body(...),
|
|
content: str = Body(...),
|
|
content_type: str = Body("text/plain")
|
|
):
|
|
|
|
vector = model.encode(f"{title} {content}").astype(np.float32).tobytes()
|
|
|
|
try:
|
|
with get_db_connection() as conn:
|
|
conn.execute("""
|
|
INSERT INTO documents (title, content, content_type, embedding)
|
|
VALUES (?, ?, ?, ?) ON CONFLICT(title) DO
|
|
UPDATE SET
|
|
content=excluded.content,
|
|
content_type=excluded.content_type,
|
|
embedding=excluded.embedding
|
|
""", (title, content.encode('utf-8'), content_type, vector))
|
|
conn.commit()
|
|
return {"status": "success", "message": f"Dokument '{title}' zapisany."}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@app.post("/search")
|
|
async def search_similar(query: str = Body(..., embed=True), top_k: int = 3):
|
|
"""Wyszukiwanie semantyczne (Vector Search)"""
|
|
query_vector = model.encode(query).astype(np.float32)
|
|
|
|
with get_db_connection() as conn:
|
|
cursor = conn.execute("SELECT title, content, embedding FROM documents")
|
|
rows = cursor.fetchall()
|
|
|
|
results = []
|
|
for row in rows:
|
|
db_vector = np.frombuffer(row['embedding'], dtype=np.float32)
|
|
|
|
|
|
score = np.dot(query_vector, db_vector) / (np.linalg.norm(query_vector) * np.linalg.norm(db_vector))
|
|
|
|
results.append({
|
|
"title": row['title'],
|
|
"content": row['content'].decode('utf-8', errors='ignore'),
|
|
"score": float(score)
|
|
})
|
|
|
|
|
|
results = sorted(results, key=lambda x: x['score'], reverse=True)[:top_k]
|
|
return {"results": results}
|
|
|
|
|
|
@app.get("/load-document")
|
|
async def load_document(title: str = None):
|
|
with get_db_connection() as conn:
|
|
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()
|
|
|
|
if row:
|
|
return {
|
|
"title": row['title'],
|
|
"content": row['content'].decode('utf-8', errors='ignore')
|
|
}
|
|
return {"error": "Nie znaleziono dokumentu"}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
uvicorn.run(app, host="127.0.0.1", port=8000)
|