diff --git a/app.py b/app.py index 52bf80c..b378fd2 100644 --- a/app.py +++ b/app.py @@ -5,6 +5,7 @@ Creates the FastAPI app, registers middleware, mounts all routers. """ import logging import logging.handlers +import os from contextlib import asynccontextmanager from pathlib import Path @@ -13,7 +14,7 @@ from fastapi.responses import HTMLResponse, JSONResponse from fastapi.staticfiles import StaticFiles from fastapi.templating import Jinja2Templates -from config import VERSION, RATE_WINDOW_SECONDS +from config import VERSION, RATE_WINDOW_SECONDS, UPLOAD_DIR from db import init_db from memory import get_memory_count from security import ( @@ -32,6 +33,7 @@ import routers.skills as skills import routers.chat as chat import routers.search_route as search_route import routers.completions as completions +import routers.upload as upload # --- Logging --- log = logging.getLogger("jarvischat") @@ -47,6 +49,7 @@ templates = Jinja2Templates(directory=str(BASE_DIR / "templates")) @asynccontextmanager async def lifespan(app: FastAPI): log.info(f"JarvisChat {VERSION} starting up") + os.makedirs(UPLOAD_DIR, exist_ok=True) init_db() log.info(f"Memory system: {get_memory_count()} memories loaded") yield @@ -138,7 +141,7 @@ async def index(request: Request): for router_module in [ auth_router, conversations.router, memories.router, models.router, presets.router, profile.router, settings.router, skills.router, - chat.router, search_route.router, completions.router, + chat.router, search_route.router, completions.router, upload.router, ]: app.include_router(router_module) diff --git a/config.py b/config.py index b07a605..a552c61 100644 --- a/config.py +++ b/config.py @@ -9,7 +9,7 @@ import logging log = logging.getLogger("jarvischat") -VERSION = "v1.8.9" +VERSION = "v1.9.0" OLLAMA_BASE = os.environ.get("OLLAMA_BASE", "http://localhost:11434") LLAMA_SERVER_BASE = os.environ.get("LLAMA_SERVER_BASE", "http://192.168.50.108:8081") SEARXNG_BASE = "http://localhost:8888" @@ -46,6 +46,13 @@ BODY_LIMIT_DEFAULT_BYTES = 64 * 1024 BODY_LIMIT_CHAT_BYTES = 128 * 1024 BODY_LIMIT_PROFILE_BYTES = 256 * 1024 +# --- Upload --- +UPLOAD_DIR = "/tmp/jarvischat_uploads" +MAX_UPLOAD_BYTES = 20 * 1024 * 1024 +SUPPORTED_UPLOAD_TYPES = {"text/plain", "text/markdown", "application/pdf", "application/json", "text/x-python", "text/html"} +UPLOAD_CONTEXT_EXPIRY_HOURS = 1 +BODY_LIMIT_UPLOAD_BYTES = MAX_UPLOAD_BYTES + MAX_CHAT_MESSAGE_CHARS = 8000 MAX_SEARCH_QUERY_CHARS = 500 MAX_PROFILE_CHARS = 32000 diff --git a/db.py b/db.py index 6bf2098..5fa0750 100644 --- a/db.py +++ b/db.py @@ -68,6 +68,30 @@ def format_active_skills_prompt(skills: list) -> str: return text +def insert_upload_context(db, conversation_id: str, filename: str, content: str, expires_at: str) -> int: + now = datetime.now(timezone.utc).isoformat() + cur = db.execute( + "INSERT INTO upload_context (conversation_id, filename, content, created_at, expires_at) VALUES (?, ?, ?, ?, ?)", + (conversation_id, filename, content, now, expires_at), + ) + return cur.lastrowid + + +def get_upload_context(db, context_id: int): + row = db.execute( + "SELECT id, conversation_id, filename, content, expires_at FROM upload_context WHERE id = ?", + (context_id,), + ).fetchone() + if not row: + return None + expires = datetime.fromisoformat(row["expires_at"]) + if expires < datetime.now(timezone.utc): + db.execute("DELETE FROM upload_context WHERE id = ?", (context_id,)) + db.commit() + return None + return dict(row) + + def init_db(): from security import hash_pin conn = sqlite3.connect(DB_PATH) @@ -108,6 +132,16 @@ def init_db(): fact, topic, source, created_at UNINDEXED ) """) + conn.execute(""" + CREATE TABLE IF NOT EXISTS upload_context ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + conversation_id TEXT, + filename TEXT NOT NULL, + content TEXT NOT NULL, + created_at TEXT NOT NULL, + expires_at TEXT NOT NULL + ) + """) if not conn.execute("SELECT id FROM profile WHERE id = 1").fetchone(): now = datetime.now(timezone.utc).isoformat() diff --git a/rag.py b/rag.py index e0554c9..8c023bc 100644 --- a/rag.py +++ b/rag.py @@ -18,6 +18,24 @@ RAG_COLLECTION = "jarvis_rag" RAG_SCORE_THRESHOLD = 0.25 +def chunk_text(text: str, chunk_size: int = 512, overlap: int = 128) -> list: + words = text.split() + # Approximate token count as len(words) * 1.3 + target_words = int(chunk_size / 1.3) + overlap_words = int(overlap / 1.3) + if not words: + return [] + chunks = [] + start = 0 + while start < len(words): + end = min(start + target_words, len(words)) + chunks.append(" ".join(words[start:end])) + if end == len(words): + break + start += target_words - overlap_words + return chunks + + async def query_rag(query: str, limit: int = 3) -> list: try: async with httpx.AsyncClient() as client: diff --git a/routers/upload.py b/routers/upload.py new file mode 100644 index 0000000..7d06013 --- /dev/null +++ b/routers/upload.py @@ -0,0 +1,98 @@ +"""JarvisChat routers - /api/upload file/document attachment endpoint.""" +import json +import logging +import os +from datetime import datetime, timezone, timedelta + +import httpx +from fastapi import APIRouter, HTTPException, Request, UploadFile, File, Form +from fastapi.responses import JSONResponse + +from config import UPLOAD_DIR, MAX_UPLOAD_BYTES, SUPPORTED_UPLOAD_TYPES, UPLOAD_CONTEXT_EXPIRY_HOURS +from db import get_db, insert_upload_context +from rag import chunk_text, QDRANT_URL, EMBED_URL, EMBED_MODEL, RAG_COLLECTION + +log = logging.getLogger("jarvischat") +router = APIRouter() + + +@router.post("/api/upload") +async def upload_file( + request: Request, + file: UploadFile = File(...), + mode: str = Form("both"), + conversation_id: str = Form(""), +): + if mode not in ("context", "ingest", "both"): + raise HTTPException(status_code=422, detail="mode must be context, ingest, or both") + + if file.size and file.size > MAX_UPLOAD_BYTES: + return JSONResponse(status_code=413, content={"detail": f"File exceeds {MAX_UPLOAD_BYTES} byte limit"}) + + content_type = file.content_type or "application/octet-stream" + if content_type not in SUPPORTED_UPLOAD_TYPES: + return JSONResponse(status_code=415, content={"detail": f"Unsupported file type: {content_type}"}) + + raw_bytes = await file.read() + if not raw_bytes: + raise HTTPException(status_code=422, detail="Empty file") + + if content_type == "application/pdf": + try: + from pypdf import PdfReader + import io + reader = PdfReader(io.BytesIO(raw_bytes)) + extracted = "\n".join(page.extract_text() or "" for page in reader.pages) + except Exception as e: + log.warning(f"PDF extraction error: {e}") + raise HTTPException(status_code=422, detail="Failed to extract text from PDF") + else: + extracted = raw_bytes.decode("utf-8", errors="replace") + + result = {"filename": file.filename, "size_bytes": len(raw_bytes), "mode": mode} + + if mode in ("ingest", "both"): + os.makedirs(UPLOAD_DIR, exist_ok=True) + chunks = chunk_text(extracted) + ingested = 0 + async with httpx.AsyncClient() as client: + for i, chunk in enumerate(chunks): + embed_resp = await client.post( + f"{EMBED_URL}/api/embeddings", + json={"model": EMBED_MODEL, "prompt": chunk}, + timeout=30.0, + ) + if embed_resp.status_code != 200: + log.warning(f"Embedding failed for chunk {i}: {embed_resp.status_code}") + continue + vector = embed_resp.json()["embedding"] + point_id = f"{file.filename}-{i}" + upsert_resp = await client.put( + f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points?wait=true", + json={ + "points": [{ + "id": point_id, + "vector": vector, + "payload": {"text": chunk, "source": file.filename, "upload_date": datetime.now(timezone.utc).isoformat(), "type": "upload"}, + }] + }, + timeout=30.0, + ) + if upsert_resp.status_code in (200, 201): + ingested += 1 + else: + log.warning(f"Qdrant upsert failed for chunk {i}: {upsert_resp.status_code}") + result["chunks_ingested"] = ingested + + if mode in ("context", "both"): + expires = (datetime.now(timezone.utc) + timedelta(hours=UPLOAD_CONTEXT_EXPIRY_HOURS)).isoformat() + db = get_db() + try: + cid = insert_upload_context(db, conversation_id or "", file.filename or "unnamed", extracted, expires) + db.commit() + result["context_id"] = cid + finally: + db.close() + + result["message"] = f"Uploaded {file.filename}" + return result diff --git a/security.py b/security.py index 135397f..2c120eb 100644 --- a/security.py +++ b/security.py @@ -22,7 +22,7 @@ from fastapi import HTTPException, Request from config import ( ALLOWED_NETWORKS, TRUST_X_FORWARDED_FOR, TRUSTED_ORIGINS, - BODY_LIMIT_DEFAULT_BYTES, BODY_LIMIT_CHAT_BYTES, BODY_LIMIT_PROFILE_BYTES, + BODY_LIMIT_DEFAULT_BYTES, BODY_LIMIT_CHAT_BYTES, BODY_LIMIT_PROFILE_BYTES, BODY_LIMIT_UPLOAD_BYTES, RATE_WINDOW_SECONDS, RL_LOGIN_PER_WINDOW, RL_CHAT_PER_WINDOW, RL_SEARCH_PER_WINDOW, RL_STATS_PER_WINDOW, RL_WRITE_PER_WINDOW, RL_DEFAULT_PER_WINDOW, VERSION, @@ -114,6 +114,8 @@ def request_body_limit(path: str) -> int: return BODY_LIMIT_CHAT_BYTES if path == "/api/profile": return BODY_LIMIT_PROFILE_BYTES + if path == "/api/upload": + return BODY_LIMIT_UPLOAD_BYTES return BODY_LIMIT_DEFAULT_BYTES diff --git a/tests/test_upload.py b/tests/test_upload.py new file mode 100644 index 0000000..9257ee3 --- /dev/null +++ b/tests/test_upload.py @@ -0,0 +1,154 @@ +import os +from pathlib import Path + +import httpx +import pytest +from fastapi.testclient import TestClient + +import app +import db +import routers.upload as upload_route +from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS + + +def make_client(tmp_path: Path) -> TestClient: + os.environ["JARVISCHAT_ADMIN_PIN"] = "1234" + db.DB_PATH = tmp_path / "jarvischat-upload.db" + SESSIONS.clear() + PIN_ATTEMPTS.clear() + RATE_EVENTS.clear() + db.init_db() + return TestClient(app.app, raise_server_exceptions=False) + + +def _admin_headers(client: TestClient) -> dict: + login = client.post("/api/auth/login", json={"pin": "1234"}, headers={"Origin": "http://testserver"}) + sid = login.json()["session_id"] + return {"X-Session-ID": sid, "Origin": "http://testserver"} + + +def _guest_headers(client: TestClient) -> dict: + sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()["session_id"] + return {"X-Session-ID": sid, "Origin": "http://testserver"} + + +def test_upload_requires_admin(tmp_path: Path): + with make_client(tmp_path) as client: + resp = client.post("/api/upload", headers=_guest_headers(client), files={"file": ("test.txt", b"hello")}) + assert resp.status_code == 403 + + +def test_upload_unsupported_mime(tmp_path: Path): + with make_client(tmp_path) as client: + resp = client.post( + "/api/upload", headers=_admin_headers(client), + files={"file": ("test.exe", b"fake", "application/x-msdownload")}, + ) + assert resp.status_code == 415 + + +def test_upload_context_mode(tmp_path: Path): + with make_client(tmp_path) as client: + resp = client.post( + "/api/upload", headers=_admin_headers(client), + data={"mode": "context", "conversation_id": "conv-1"}, + files={"file": ("notes.txt", b"Hello world notes")}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["filename"] == "notes.txt" + assert data["mode"] == "context" + assert "context_id" in data + assert "chunks_ingested" not in data + + row = db.get_db().execute("SELECT content FROM upload_context WHERE id = ?", (data["context_id"],)).fetchone() + assert row["content"] == "Hello world notes" + + +def test_upload_ingest_mode(tmp_path: Path, monkeypatch): + embed_count = 0 + + class FakeAsyncClient: + class FakeResponse: + def __init__(self, status, json_data=None): + self.status_code = status + self._json = json_data or {} + + def json(self): + return self._json + + def __init__(self, *a, **kw): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + pass + + async def post(self, url, **kw): + nonlocal embed_count + if "/api/embeddings" in url: + embed_count += 1 + return self.FakeResponse(200, {"embedding": [0.1] * 768}) + return self.FakeResponse(200) + + async def put(self, url, **kw): + return self.FakeResponse(200) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: FakeAsyncClient()) + + with make_client(tmp_path) as client: + resp = client.post( + "/api/upload", headers=_admin_headers(client), + data={"mode": "ingest"}, + files={"file": ("data.txt", b"word " * 1000)}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["mode"] == "ingest" + assert data["chunks_ingested"] > 0 + assert "context_id" not in data + assert embed_count == data["chunks_ingested"] + + +def test_upload_both_mode(tmp_path: Path, monkeypatch): + class FakeAsyncClient: + class FakeResponse: + def __init__(self, status, json_data=None): + self.status_code = status + self._json = json_data or {} + + def json(self): + return self._json + + def __init__(self, *a, **kw): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + pass + + async def post(self, url, **kw): + if "/api/embeddings" in url: + return self.FakeResponse(200, {"embedding": [0.1] * 768}) + return self.FakeResponse(200) + + async def put(self, url, **kw): + return self.FakeResponse(200) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: FakeAsyncClient()) + + with make_client(tmp_path) as client: + resp = client.post( + "/api/upload", headers=_admin_headers(client), + data={"mode": "both", "conversation_id": "conv-2"}, + files={"file": ("both.txt", b"test " * 500)}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["mode"] == "both" + assert "context_id" in data + assert data["chunks_ingested"] > 0