diff --git a/.env.example b/.env.example index f0dd308..d5adf64 100644 --- a/.env.example +++ b/.env.example @@ -2,3 +2,4 @@ OLLAMA_API_BASE=http://localhost:11434 COMPASS_MODEL=ollama/gemma4:e4b COMPASS_DOCS_PATH=./data/docs COMPASS_WORKSPACE=./data/index +COMPASS_ALLOWED_ORIGINS=http://localhost:3000,http://localhost:5173 diff --git a/backend/chat.py b/backend/chat.py index 104450e..31e9cd0 100644 --- a/backend/chat.py +++ b/backend/chat.py @@ -59,7 +59,7 @@ async def process_chat( messages.append({"role": "user", "content": question}) # 5. Single LLM call — answer + optional suggestion embedded - response = await litellm.acompletion(model=model, messages=messages) + response = await litellm.acompletion(model=model, messages=messages, timeout=60) raw = response.choices[0].message.content # 6. Parse answer and suggestion from the same response diff --git a/backend/database.py b/backend/database.py index b323d3b..6f960ae 100644 --- a/backend/database.py +++ b/backend/database.py @@ -18,6 +18,9 @@ def init_db(): indexed BOOLEAN DEFAULT FALSE ) """) + conn.execute(""" + CREATE INDEX IF NOT EXISTS idx_messages_session_id ON messages(session_id) + """) @contextmanager @@ -45,7 +48,7 @@ def get_session_messages(session_id: str, limit: int = 10) -> list: rows = conn.execute( """SELECT role, content FROM messages WHERE session_id = ? - ORDER BY timestamp DESC LIMIT ?""", + ORDER BY id DESC LIMIT ?""", (session_id, limit), ).fetchall() return [{"role": r["role"], "content": r["content"]} for r in reversed(rows)] diff --git a/backend/main.py b/backend/main.py index 6e3ba30..58a956d 100644 --- a/backend/main.py +++ b/backend/main.py @@ -59,11 +59,17 @@ async def lifespan(app: FastAPI): lifespan=lifespan, ) +ALLOWED_ORIGINS = [ + o.strip() + for o in os.getenv("COMPASS_ALLOWED_ORIGINS", "http://localhost:3000,http://localhost:5173").split(",") + if o.strip() +] + app.add_middleware( CORSMiddleware, - allow_origins=["*"], - allow_methods=["*"], - allow_headers=["*"], + allow_origins=ALLOWED_ORIGINS, + allow_methods=["GET", "POST"], + allow_headers=["Content-Type"], ) @@ -98,19 +104,28 @@ async def chat(body: ChatRequest, request: Request): ) +MAX_UPLOAD_BYTES = 50 * 1024 * 1024 # 50 MB + + @app.post("/upload", status_code=201) async def upload(file: UploadFile = File(...), request: Request = None): indexer = get_indexer(request) - ext = Path(file.filename).suffix.lower() + + safe_name = Path(file.filename).name + ext = Path(safe_name).suffix.lower() if ext not in {".pdf", ".md", ".markdown", ".txt"}: raise HTTPException(status_code=400, detail="Formatos soportados: PDF, Markdown, TXT") - dest = DOCS_PATH / file.filename + contents = await file.read() + if len(contents) > MAX_UPLOAD_BYTES: + raise HTTPException(status_code=413, detail="Archivo demasiado grande. Máximo 50 MB.") + + dest = DOCS_PATH / safe_name with open(dest, "wb") as f: - shutil.copyfileobj(file.file, f) + f.write(contents) doc_id = indexer.index_document(str(dest)) - return {"doc_id": doc_id, "filename": file.filename} + return {"doc_id": doc_id, "filename": safe_name} @app.get("/documents") diff --git a/tests/test_api.py b/tests/test_api.py index b3e8b08..0e82d42 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -138,6 +138,30 @@ def test_201_on_txt_upload(self, client_empty): ) assert r.status_code == 201 + def test_path_traversal_filename_is_sanitized(self, client_empty): + r = client_empty.post( + "/upload", + files={"file": ("../../etc/passwd.md", b"# hacked", "text/markdown")}, + ) + assert r.status_code == 201 + assert r.json()["filename"] == "passwd.md" + + def test_413_when_file_too_large(self, client_empty): + big = b"x" * (50 * 1024 * 1024 + 1) + r = client_empty.post( + "/upload", + files={"file": ("big.md", big, "text/markdown")}, + ) + assert r.status_code == 413 + + def test_file_at_size_limit_is_accepted(self, client_empty): + at_limit = b"x" * (50 * 1024 * 1024) + r = client_empty.post( + "/upload", + files={"file": ("limit.md", at_limit, "text/markdown")}, + ) + assert r.status_code == 201 + # --------------------------------------------------------------------------- # GET /documents diff --git a/tests/test_database.py b/tests/test_database.py new file mode 100644 index 0000000..ffe106c --- /dev/null +++ b/tests/test_database.py @@ -0,0 +1,93 @@ +""" +Unit tests for backend/database.py. +Uses a real in-memory SQLite database — no mocking needed. +""" + +import sqlite3 +import pytest +from unittest.mock import patch +from pathlib import Path + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +@pytest.fixture +def db(tmp_path): + """Isolated database in a temp directory for each test.""" + db_path = tmp_path / "test_compass.db" + with patch("backend.database.DB_PATH", db_path): + from backend.database import init_db, save_message, get_session_messages + init_db() + yield {"save": save_message, "get": get_session_messages, "path": db_path} + + +# --------------------------------------------------------------------------- +# Schema +# --------------------------------------------------------------------------- + +class TestSchema: + def test_messages_table_exists(self, db): + conn = sqlite3.connect(str(db["path"])) + tables = conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' AND name='messages'" + ).fetchone() + conn.close() + assert tables is not None + + def test_session_id_index_exists(self, db): + conn = sqlite3.connect(str(db["path"])) + idx = conn.execute( + "SELECT name FROM sqlite_master WHERE type='index' AND name='idx_messages_session_id'" + ).fetchone() + conn.close() + assert idx is not None + + +# --------------------------------------------------------------------------- +# save_message / get_session_messages +# --------------------------------------------------------------------------- + +class TestMessages: + def test_save_returns_id(self, db): + msg_id = db["save"]("s1", "user", "hello") + assert isinstance(msg_id, int) + assert msg_id > 0 + + def test_get_returns_saved_messages(self, db): + db["save"]("s1", "user", "hello") + db["save"]("s1", "assistant", "hi there") + msgs = db["get"]("s1") + assert len(msgs) == 2 + assert msgs[0]["role"] == "user" + assert msgs[1]["role"] == "assistant" + + def test_get_returns_chronological_order(self, db): + db["save"]("s1", "user", "first") + db["save"]("s1", "assistant", "second") + db["save"]("s1", "user", "third") + msgs = db["get"]("s1") + assert [m["content"] for m in msgs] == ["first", "second", "third"] + + def test_get_isolates_by_session(self, db): + db["save"]("session-a", "user", "message A") + db["save"]("session-b", "user", "message B") + assert len(db["get"]("session-a")) == 1 + assert db["get"]("session-a")[0]["content"] == "message A" + + def test_get_respects_limit(self, db): + for i in range(10): + db["save"]("s1", "user", f"msg {i}") + msgs = db["get"]("s1", limit=3) + assert len(msgs) == 3 + + def test_get_returns_most_recent_when_limited(self, db): + for i in range(5): + db["save"]("s1", "user", f"msg {i}") + msgs = db["get"]("s1", limit=2) + assert msgs[0]["content"] == "msg 3" + assert msgs[1]["content"] == "msg 4" + + def test_get_empty_session_returns_empty_list(self, db): + assert db["get"]("nonexistent") == []