319 lines
10 KiB
Python
319 lines
10 KiB
Python
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
|
|
|
|
|
|
ROOT = Path(__file__).parent
|
|
DATA_DIR = ROOT / "data"
|
|
CHATS_FILE = DATA_DIR / "chats.json"
|
|
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:
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
|
|
def read_store() -> dict:
|
|
if not CHATS_FILE.exists():
|
|
return {"chats": []}
|
|
|
|
try:
|
|
return json.loads(CHATS_FILE.read_text(encoding="utf-8"))
|
|
except json.JSONDecodeError:
|
|
return {"chats": []}
|
|
|
|
|
|
def write_store(store: dict) -> None:
|
|
DATA_DIR.mkdir(exist_ok=True)
|
|
CHATS_FILE.write_text(json.dumps(store, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
|
|
|
|
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 "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):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, directory=str(ROOT), **kwargs)
|
|
|
|
def do_GET(self) -> None:
|
|
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", "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 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:
|
|
self._send_json(404, {"error": "Chat not found"})
|
|
return
|
|
|
|
self._send_json(200, {"chat": chat})
|
|
return
|
|
|
|
super().do_GET()
|
|
|
|
def do_POST(self) -> None:
|
|
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
|
|
|
|
length = int(self.headers.get("content-length", 0))
|
|
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,
|
|
messages=[
|
|
{"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, 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:
|
|
path = urlparse(self.path).path
|
|
if not path.startswith("/api/chats/"):
|
|
self.send_error(404)
|
|
return
|
|
|
|
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], 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)
|
|
|
|
if not chat:
|
|
chat = {
|
|
"id": str(uuid.uuid4()),
|
|
"createdAt": now_iso(),
|
|
"title": chat_title(messages),
|
|
"messages": [],
|
|
}
|
|
chats.insert(0, chat)
|
|
|
|
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) -> 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")
|
|
self.send_header("content-length", str(len(body)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
|
|
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}")
|
|
print(f"LAN access is available at http://{local_ip}:{PORT}")
|
|
server.serve_forever()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|