Add model selection and local RAG

server.py

@@ -1,11 +1,16 @@
 from __future__ import annotations
 
 import json
+import math
+import re
 import socket
+import time
 import uuid
+from collections import Counter
 from datetime import datetime, timezone
 from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
 from pathlib import Path
+from urllib.parse import urlparse
 
 import ollama
 
@@ -13,9 +18,15 @@ import ollama
 ROOT = Path(__file__).parent
 DATA_DIR = ROOT / "data"
 CHATS_FILE = DATA_DIR / "chats.json"
-MODEL = "qwen2.5:7b"
+DEFAULT_MODEL = "qwen2.5:7b"
+AVAILABLE_MODELS = ["qwen2.5:7b", "llama3.1:8b", "qwen2.5:14b", "qwen2.5:3b"]
 HOST = "0.0.0.0"
 PORT = 8787
+RAG_EXTENSIONS = {".md", ".txt", ".py", ".js", ".html", ".css", ".ps1"}
+RAG_EXCLUDED_DIRS = {".git", ".venv", "__pycache__", "data"}
+
+rag_chunks: list[dict] = []
+rag_indexed_at = ""
 
 
 def now_iso() -> str:
@@ -41,8 +52,114 @@ def chat_title(messages: list[dict]) -> str:
     for message in messages:
         if message.get("role") == "user" and message.get("content"):
             title = " ".join(str(message["content"]).split())
-            return title[:40] or "새 대화"
-    return "새 대화"
+            return title[:40] or "New chat"
+    return "New chat"
+
+
+def installed_models() -> list[str]:
+    try:
+        response = ollama.list()
+    except Exception:
+        return AVAILABLE_MODELS
+
+    names: list[str] = []
+    for model in response.get("models", []):
+        if isinstance(model, dict):
+            name = model.get("name") or model.get("model")
+        else:
+            name = getattr(model, "name", None) or getattr(model, "model", None)
+        if name:
+            names.append(name)
+
+    ordered = [name for name in AVAILABLE_MODELS if name in names]
+    ordered.extend(name for name in names if name not in ordered)
+    return ordered or AVAILABLE_MODELS
+
+
+def tokenize(text: str) -> list[str]:
+    return [token for token in re.findall(r"\w+", text.lower(), flags=re.UNICODE) if len(token) > 1]
+
+
+def iter_rag_files() -> list[Path]:
+    files: list[Path] = []
+    for path in ROOT.rglob("*"):
+        if any(part in RAG_EXCLUDED_DIRS for part in path.relative_to(ROOT).parts):
+            continue
+        if path.is_file() and path.suffix.lower() in RAG_EXTENSIONS:
+            files.append(path)
+    return sorted(files)
+
+
+def chunk_text(path: Path, text: str, size: int = 1800, overlap: int = 250) -> list[dict]:
+    chunks: list[dict] = []
+    start = 0
+    rel = path.relative_to(ROOT).as_posix()
+    while start < len(text):
+        end = min(start + size, len(text))
+        chunk = text[start:end].strip()
+        if chunk:
+            chunks.append({"source": rel, "text": chunk, "tokens": Counter(tokenize(chunk))})
+        if end == len(text):
+            break
+        start = max(0, end - overlap)
+    return chunks
+
+
+def build_rag_index() -> dict:
+    global rag_chunks, rag_indexed_at
+
+    chunks: list[dict] = []
+    for path in iter_rag_files():
+        try:
+            text = path.read_text(encoding="utf-8", errors="replace")
+        except OSError:
+            continue
+        chunks.extend(chunk_text(path, text))
+
+    rag_chunks = chunks
+    rag_indexed_at = now_iso()
+    return rag_status()
+
+
+def rag_status() -> dict:
+    return {
+        "chunks": len(rag_chunks),
+        "indexedAt": rag_indexed_at,
+        "files": len({chunk["source"] for chunk in rag_chunks}),
+    }
+
+
+def search_rag(query: str, limit: int = 5) -> list[dict]:
+    if not rag_chunks:
+        build_rag_index()
+
+    query_tokens = Counter(tokenize(query))
+    if not query_tokens:
+        return []
+
+    scored: list[tuple[float, dict]] = []
+    total = max(1, len(rag_chunks))
+    doc_freq = Counter()
+    for chunk in rag_chunks:
+        for token in chunk["tokens"]:
+            doc_freq[token] += 1
+
+    for chunk in rag_chunks:
+        score = 0.0
+        chunk_tokens = chunk["tokens"]
+        for token, query_count in query_tokens.items():
+            if token not in chunk_tokens:
+                continue
+            idf = math.log((total + 1) / (doc_freq[token] + 1)) + 1
+            score += min(query_count, chunk_tokens[token]) * idf
+        if score:
+            scored.append((score, chunk))
+
+    scored.sort(key=lambda item: item[0], reverse=True)
+    return [
+        {"source": chunk["source"], "text": chunk["text"], "score": round(score, 3)}
+        for score, chunk in scored[:limit]
+    ]
 
 
 class Handler(SimpleHTTPRequestHandler):
@@ -50,21 +167,35 @@ class Handler(SimpleHTTPRequestHandler):
         super().__init__(*args, directory=str(ROOT), **kwargs)
 
     def do_GET(self) -> None:
-        if self.path == "/api/chats":
+        path = urlparse(self.path).path
+
+        if path == "/api/models":
+            self._send_json(200, {"defaultModel": DEFAULT_MODEL, "models": installed_models()})
+            return
+
+        if path == "/api/rag/status":
+            if not rag_chunks:
+                build_rag_index()
+            self._send_json(200, rag_status())
+            return
+
+        if path == "/api/chats":
             store = read_store()
             chats = [
                 {
                     "id": chat["id"],
-                    "title": chat.get("title", "새 대화"),
+                    "title": chat.get("title", "New chat"),
                     "updatedAt": chat.get("updatedAt", ""),
+                    "model": chat.get("model", DEFAULT_MODEL),
+                    "useRag": chat.get("useRag", True),
                 }
                 for chat in store.get("chats", [])
             ]
             self._send_json(200, {"chats": chats})
             return
 
-        if self.path.startswith("/api/chats/"):
-            chat_id = self.path.rsplit("/", 1)[-1]
+        if path.startswith("/api/chats/"):
+            chat_id = path.rsplit("/", 1)[-1]
             store = read_store()
             chat = next((item for item in store.get("chats", []) if item["id"] == chat_id), None)
             if not chat:
@@ -77,7 +208,12 @@ class Handler(SimpleHTTPRequestHandler):
         super().do_GET()
 
     def do_POST(self) -> None:
-        if self.path != "/api/chat":
+        path = urlparse(self.path).path
+        if path == "/api/rag/reindex":
+            self._send_json(200, build_rag_index())
+            return
+
+        if path != "/api/chat":
             self.send_error(404)
             return
 
@@ -85,38 +221,59 @@ class Handler(SimpleHTTPRequestHandler):
         data = json.loads(self.rfile.read(length) or b"{}")
         chat_id = data.get("chatId")
         messages = data.get("messages", [])
+        model = data.get("model") or DEFAULT_MODEL
+        use_rag = bool(data.get("useRag", True))
 
         try:
+            sources = search_rag(messages[-1]["content"]) if use_rag and messages else []
+            system_prompt = (
+                "You are a concise local assistant. Reply in Korean unless the user asks otherwise. "
+                "If reference context is provided, use it and mention source filenames when useful. "
+                "If the context is not enough, say so clearly."
+            )
+            if sources:
+                context = "\n\n".join(
+                    f"[{index + 1}] {item['source']}\n{item['text']}" for index, item in enumerate(sources)
+                )
+                system_prompt += f"\n\nReference context:\n{context}"
+
+            started = time.perf_counter()
             response = ollama.chat(
-                model=MODEL,
+                model=model,
                 messages=[
-                    {
-                        "role": "system",
-                        "content": "You are a concise local assistant. Reply in Korean unless the user asks otherwise.",
-                    },
+                    {"role": "system", "content": system_prompt},
                     *messages,
                 ],
             )
+            elapsed_ms = round((time.perf_counter() - started) * 1000)
             reply = response["message"]["content"]
             saved_messages = [*messages, {"role": "assistant", "content": reply}]
-            chat = self._save_chat(chat_id, saved_messages)
-            payload = {"chatId": chat["id"], "reply": reply, "chat": chat}
+            chat = self._save_chat(chat_id, saved_messages, model, use_rag)
+            payload = {
+                "chatId": chat["id"],
+                "reply": reply,
+                "chat": chat,
+                "model": model,
+                "elapsedMs": elapsed_ms,
+                "sources": [{"source": item["source"], "score": item["score"]} for item in sources],
+            }
             self._send_json(200, payload)
         except Exception as exc:
             self._send_json(500, {"error": str(exc)})
 
     def do_DELETE(self) -> None:
-        if not self.path.startswith("/api/chats/"):
+        path = urlparse(self.path).path
+        if not path.startswith("/api/chats/"):
             self.send_error(404)
             return
 
-        chat_id = self.path.rsplit("/", 1)[-1]
+        chat_id = path.rsplit("/", 1)[-1]
         store = read_store()
         store["chats"] = [item for item in store.get("chats", []) if item["id"] != chat_id]
         write_store(store)
         self._send_json(200, {"ok": True})
 
-    def _save_chat(self, chat_id: str | None, messages: list[dict]) -> dict:
+    def _save_chat(self, chat_id: str | None, messages: list[dict], model: str, use_rag: bool) -> dict:
         store = read_store()
         chats = store.setdefault("chats", [])
         chat = next((item for item in chats if item["id"] == chat_id), None)
@@ -132,12 +289,14 @@ class Handler(SimpleHTTPRequestHandler):
 
         chat["title"] = chat_title(messages)
         chat["messages"] = messages
+        chat["model"] = model
+        chat["useRag"] = use_rag
         chat["updatedAt"] = now_iso()
         chats.sort(key=lambda item: item.get("updatedAt", ""), reverse=True)
         write_store(store)
         return chat
 
-    def _send_json(self, status: int, payload: dict[str, str]) -> None:
+    def _send_json(self, status: int, payload: dict) -> None:
         body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
         self.send_response(status)
         self.send_header("content-type", "application/json; charset=utf-8")
@@ -147,6 +306,7 @@ class Handler(SimpleHTTPRequestHandler):
 
 
 def main() -> None:
+    build_rag_index()
     server = ThreadingHTTPServer((HOST, PORT), Handler)
     local_ip = socket.gethostbyname(socket.gethostname())
     print(f"LLM web is running locally at http://127.0.0.1:{PORT}")