Initial commit — kb-app RAG server
This commit is contained in:
Julian Carlson
+337
@@ -0,0 +1,337 @@
|
||||
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
|
||||
from fastapi.responses import HTMLResponse
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
from .config import CONFIG
|
||||
from .embed import embed_one
|
||||
from .ingest import ingest_file
|
||||
from . import db
|
||||
|
||||
app = FastAPI(title="KB Search")
|
||||
|
||||
|
||||
class SearchRequest(BaseModel):
|
||||
query: str
|
||||
category: Optional[str] = None
|
||||
k: int = 10
|
||||
expand_neighbors: bool = True
|
||||
neighbor_window: int = 1
|
||||
rerank: bool = True
|
||||
rerank_pool: int = 25
|
||||
|
||||
|
||||
@app.post("/search")
|
||||
def search(req: SearchRequest):
|
||||
query_emb = embed_one(req.query)
|
||||
|
||||
# Fetch larger pool when reranking
|
||||
initial_k = max(req.rerank_pool, req.k) if req.rerank else req.k
|
||||
|
||||
if req.expand_neighbors:
|
||||
results = db.search_with_neighbors(
|
||||
query_emb, k=initial_k, category=req.category,
|
||||
neighbor_window=req.neighbor_window
|
||||
)
|
||||
else:
|
||||
results = db.search(query_emb, k=initial_k, category=req.category)
|
||||
|
||||
rerank_used = False
|
||||
if req.rerank and results:
|
||||
try:
|
||||
from .rerank import rerank as do_rerank
|
||||
results = do_rerank(req.query, results, top_k=req.k)
|
||||
rerank_used = True
|
||||
except Exception as e:
|
||||
print(f"Rerank failed, falling back to vector: {e}", flush=True)
|
||||
results = results[:req.k]
|
||||
else:
|
||||
results = results[:req.k]
|
||||
|
||||
return {
|
||||
"query": req.query,
|
||||
"count": len(results),
|
||||
"rerank_used": rerank_used,
|
||||
"results": results,
|
||||
}
|
||||
|
||||
|
||||
@app.get("/stats")
|
||||
def stats():
|
||||
return db.get_stats()
|
||||
|
||||
|
||||
@app.post("/upload")
|
||||
async def upload(file: UploadFile = File(...), category: str = Form("uploads")):
|
||||
if not file.filename:
|
||||
raise HTTPException(400, "No filename")
|
||||
|
||||
safe_name = "".join(c for c in file.filename if c.isalnum() or c in "._- ")
|
||||
target_dir = Path(CONFIG.corpus_root) / category
|
||||
target_dir.mkdir(parents=True, exist_ok=True)
|
||||
target = target_dir / safe_name
|
||||
|
||||
with open(target, "wb") as f:
|
||||
shutil.copyfileobj(file.file, f)
|
||||
|
||||
result = ingest_file(target)
|
||||
return result
|
||||
|
||||
|
||||
@app.get("/upload", response_class=HTMLResponse)
|
||||
def upload_page():
|
||||
return """
|
||||
<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
<title>KB Upload</title>
|
||||
<style>
|
||||
body { font-family: -apple-system, sans-serif; max-width: 600px; margin: 50px auto; padding: 20px; }
|
||||
.drop-zone { border: 2px dashed #888; padding: 40px; text-align: center; border-radius: 8px; cursor: pointer; }
|
||||
.drop-zone.dragover { background: #f0f8ff; border-color: #4a9; }
|
||||
select, input { padding: 8px; margin: 8px 0; width: 100%; box-sizing: border-box; }
|
||||
button { padding: 10px 20px; background: #4a9; color: white; border: none; border-radius: 4px; cursor: pointer; }
|
||||
#status { margin-top: 20px; padding: 10px; }
|
||||
.ok { background: #efe; }
|
||||
.err { background: #fee; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<h1>KB Upload</h1>
|
||||
<form id="form" enctype="multipart/form-data">
|
||||
<label>Category:</label>
|
||||
<select name="category">
|
||||
<option value="uploads">uploads (default)</option>
|
||||
<option value="transcripts/claude">transcripts/claude</option>
|
||||
<option value="transcripts/openclaw">transcripts/openclaw</option>
|
||||
<option value="research/strategies">research/strategies</option>
|
||||
<option value="research/notes">research/notes</option>
|
||||
<option value="papers">papers</option>
|
||||
<option value="courses">courses</option>
|
||||
<option value="code-context">code-context</option>
|
||||
</select>
|
||||
<div class="drop-zone" id="dropzone">
|
||||
<p>Drop files here or click to select</p>
|
||||
<input type="file" name="file" id="file" multiple style="display:none">
|
||||
</div>
|
||||
<button type="submit">Upload</button>
|
||||
</form>
|
||||
<div id="status"></div>
|
||||
<script>
|
||||
const form = document.getElementById('form');
|
||||
const fileInput = document.getElementById('file');
|
||||
const dropzone = document.getElementById('dropzone');
|
||||
const status = document.getElementById('status');
|
||||
dropzone.onclick = () => fileInput.click();
|
||||
dropzone.ondragover = (e) => { e.preventDefault(); dropzone.classList.add('dragover'); };
|
||||
dropzone.ondragleave = () => dropzone.classList.remove('dragover');
|
||||
dropzone.ondrop = (e) => {
|
||||
e.preventDefault();
|
||||
dropzone.classList.remove('dragover');
|
||||
fileInput.files = e.dataTransfer.files;
|
||||
dropzone.querySelector('p').textContent = `${fileInput.files.length} file(s) selected`;
|
||||
};
|
||||
fileInput.onchange = () => {
|
||||
dropzone.querySelector('p').textContent = `${fileInput.files.length} file(s) selected`;
|
||||
};
|
||||
form.onsubmit = async (e) => {
|
||||
e.preventDefault();
|
||||
status.innerHTML = 'Uploading...';
|
||||
status.className = '';
|
||||
const category = form.querySelector('[name=category]').value;
|
||||
const results = [];
|
||||
for (const f of fileInput.files) {
|
||||
const fd = new FormData();
|
||||
fd.append('file', f);
|
||||
fd.append('category', category);
|
||||
try {
|
||||
const r = await fetch('/upload', { method: 'POST', body: fd });
|
||||
const j = await r.json();
|
||||
results.push(`${f.name}: ${j.status}` + (j.chunks ? ` (${j.chunks} chunks)` : ''));
|
||||
} catch (err) {
|
||||
results.push(`${f.name}: error - ${err}`);
|
||||
}
|
||||
}
|
||||
status.innerHTML = results.join('<br>');
|
||||
status.className = 'ok';
|
||||
};
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
"""
|
||||
|
||||
|
||||
SEARCH_UI_HTML = """\
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>KB Search</title>
|
||||
<style>
|
||||
* { box-sizing: border-box; margin: 0; padding: 0; }
|
||||
body { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; background: #0f1117; color: #e2e8f0; min-height: 100vh; }
|
||||
.header { padding: 24px 32px 0; display: flex; align-items: center; justify-content: space-between; }
|
||||
.header h1 { font-size: 1.4rem; font-weight: 600; letter-spacing: -0.02em; color: #f8fafc; }
|
||||
.header a { color: #64748b; text-decoration: none; font-size: 0.85rem; }
|
||||
.header a:hover { color: #94a3b8; }
|
||||
.search-box { padding: 24px 32px; display: flex; gap: 10px; align-items: flex-start; flex-wrap: wrap; }
|
||||
.search-box input[type=text] { flex: 1; min-width: 260px; padding: 10px 14px; background: #1e2130; border: 1px solid #334155; border-radius: 8px; color: #e2e8f0; font-size: 0.95rem; outline: none; }
|
||||
.search-box input[type=text]:focus { border-color: #6366f1; }
|
||||
.search-box button { padding: 10px 20px; background: #6366f1; color: #fff; border: none; border-radius: 8px; cursor: pointer; font-size: 0.95rem; white-space: nowrap; }
|
||||
.search-box button:hover { background: #4f46e5; }
|
||||
.search-box button:disabled { opacity: 0.5; cursor: default; }
|
||||
.filters { padding: 0 32px 16px; display: flex; gap: 10px; flex-wrap: wrap; align-items: center; }
|
||||
.filters label { font-size: 0.8rem; color: #64748b; }
|
||||
.filters select, .filters input[type=number] { padding: 6px 10px; background: #1e2130; border: 1px solid #334155; border-radius: 6px; color: #e2e8f0; font-size: 0.82rem; }
|
||||
.filters input[type=checkbox] { accent-color: #6366f1; }
|
||||
.status { padding: 0 32px; font-size: 0.85rem; color: #64748b; min-height: 20px; }
|
||||
.results { padding: 16px 32px 40px; display: flex; flex-direction: column; gap: 12px; }
|
||||
.card { background: #1e2130; border: 1px solid #2d3748; border-radius: 10px; overflow: hidden; }
|
||||
.card-header { padding: 12px 16px; display: flex; align-items: center; gap: 10px; cursor: pointer; user-select: none; }
|
||||
.card-header:hover { background: #252a3a; }
|
||||
.card-rank { font-size: 0.75rem; color: #64748b; min-width: 24px; }
|
||||
.card-source { font-size: 0.78rem; color: #6366f1; font-family: monospace; flex: 1; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
|
||||
.card-topic { font-size: 0.82rem; color: #94a3b8; flex: 2; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
|
||||
.card-score { font-size: 0.75rem; color: #475569; white-space: nowrap; }
|
||||
.card-badges { display: flex; gap: 4px; }
|
||||
.badge { font-size: 0.68rem; padding: 2px 6px; border-radius: 4px; font-weight: 500; }
|
||||
.badge-vec { background: #1e3a5f; color: #60a5fa; }
|
||||
.badge-kw { background: #1a3a2a; color: #4ade80; }
|
||||
.card-body { padding: 0 16px 16px; display: none; }
|
||||
.card-body.open { display: block; }
|
||||
.card-body pre { white-space: pre-wrap; word-break: break-word; font-size: 0.84rem; color: #cbd5e1; line-height: 1.6; font-family: inherit; margin-top: 10px; padding: 12px; background: #141720; border-radius: 6px; border: 1px solid #2d3748; max-height: 400px; overflow-y: auto; }
|
||||
.card-meta { margin-top: 8px; font-size: 0.75rem; color: #475569; }
|
||||
.empty { padding: 48px 32px; text-align: center; color: #475569; }
|
||||
.spinner { display: inline-block; width: 16px; height: 16px; border: 2px solid #334155; border-top-color: #6366f1; border-radius: 50%; animation: spin 0.7s linear infinite; vertical-align: middle; margin-right: 6px; }
|
||||
@keyframes spin { to { transform: rotate(360deg); } }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="header">
|
||||
<h1>🔍 KB Search</h1>
|
||||
<div style="display:flex;gap:16px">
|
||||
<a href="/upload">Upload</a>
|
||||
<a href="/stats">Stats</a>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="search-box">
|
||||
<input type="text" id="q" placeholder="Search your knowledge base..." autofocus>
|
||||
<button id="btn" onclick="doSearch()">Search</button>
|
||||
</div>
|
||||
|
||||
<div class="filters">
|
||||
<label>Category:</label>
|
||||
<select id="cat">
|
||||
<option value="">All</option>
|
||||
<option value="transcripts/openclaw">OpenClaw transcripts</option>
|
||||
<option value="transcripts/claude">Claude transcripts</option>
|
||||
<option value="research">Research</option>
|
||||
<option value="papers">Papers</option>
|
||||
<option value="courses">Courses</option>
|
||||
<option value="code-context">Code</option>
|
||||
<option value="uploads">Uploads</option>
|
||||
</select>
|
||||
<label>Results:</label>
|
||||
<input type="number" id="k" value="10" min="1" max="50" style="width:60px">
|
||||
<label><input type="checkbox" id="rerank" checked> Rerank</label>
|
||||
<label><input type="checkbox" id="expand" checked> Expand neighbors</label>
|
||||
</div>
|
||||
|
||||
<div class="status" id="status"></div>
|
||||
<div class="results" id="results"></div>
|
||||
|
||||
<script>
|
||||
const q = document.getElementById('q');
|
||||
const btn = document.getElementById('btn');
|
||||
const status = document.getElementById('status');
|
||||
const res = document.getElementById('results');
|
||||
|
||||
q.addEventListener('keydown', function(e) { if (e.key === 'Enter') doSearch(); });
|
||||
|
||||
async function doSearch() {
|
||||
var query = q.value.trim();
|
||||
if (!query) return;
|
||||
btn.disabled = true;
|
||||
status.innerHTML = '<span class="spinner"></span>Searching...';
|
||||
res.innerHTML = '';
|
||||
|
||||
try {
|
||||
var body = {
|
||||
query: query,
|
||||
k: parseInt(document.getElementById('k').value) || 10,
|
||||
rerank: document.getElementById('rerank').checked,
|
||||
expand_neighbors: document.getElementById('expand').checked,
|
||||
neighbor_window: 1,
|
||||
rerank_pool: 25
|
||||
};
|
||||
var cat = document.getElementById('cat').value;
|
||||
if (cat) body.category = cat;
|
||||
|
||||
var r = await fetch('/search', {
|
||||
method: 'POST',
|
||||
headers: {'Content-Type': 'application/json'},
|
||||
body: JSON.stringify(body)
|
||||
});
|
||||
var data = await r.json();
|
||||
|
||||
if (!r.ok) { status.textContent = 'Error: ' + (data.detail || r.status); return; }
|
||||
|
||||
var count = data.count || 0;
|
||||
var reranked = data.rerank_used ? ' \u00b7 reranked' : '';
|
||||
status.textContent = count === 0 ? 'No results.' : (count + ' result' + (count !== 1 ? 's' : '') + reranked);
|
||||
|
||||
if (count === 0) { res.innerHTML = '<div class="empty">No matching chunks found.</div>'; return; }
|
||||
|
||||
res.innerHTML = data.results.map(function(result, i) {
|
||||
var src = result.source_path || (result.metadata && result.metadata.source_path) || '\u2014';
|
||||
var srcShort = src.split('/').slice(-2).join('/');
|
||||
var topic = result.topic || (result.metadata && result.metadata.topic) || '';
|
||||
var score = result.score != null ? result.score.toFixed(3) : '';
|
||||
var isVec = result.matched_vector;
|
||||
var isKw = result.matched_keyword;
|
||||
var content = (result.content || '').replace(/</g,'<').replace(/>/g,'>');
|
||||
var tokens = result.token_count ? (result.token_count + ' tokens') : '';
|
||||
var chunk = result.chunk_index != null ? ('chunk ' + result.chunk_index) : '';
|
||||
|
||||
var badges = (isVec ? '<span class="badge badge-vec">vec</span>' : '') +
|
||||
(isKw ? '<span class="badge badge-kw">kw</span>' : '');
|
||||
var meta = [srcShort, chunk, tokens].filter(Boolean).join(' \u00b7 ');
|
||||
|
||||
return '<div class="card">' +
|
||||
'<div class="card-header" onclick="toggle(this)">' +
|
||||
'<span class="card-rank">#' + (i+1) + '</span>' +
|
||||
'<span class="card-source" title="' + src + '">' + srcShort + '</span>' +
|
||||
'<span class="card-topic" title="' + topic + '">' + topic + '</span>' +
|
||||
'<div class="card-badges">' + badges + '</div>' +
|
||||
'<span class="card-score">' + score + '</span>' +
|
||||
'</div>' +
|
||||
'<div class="card-body">' +
|
||||
'<pre>' + content + '</pre>' +
|
||||
'<div class="card-meta">' + meta + '</div>' +
|
||||
'</div>' +
|
||||
'</div>';
|
||||
}).join('');
|
||||
|
||||
} catch(e) {
|
||||
status.textContent = 'Error: ' + e.message;
|
||||
} finally {
|
||||
btn.disabled = false;
|
||||
}
|
||||
}
|
||||
|
||||
function toggle(header) {
|
||||
var body = header.nextElementSibling;
|
||||
body.classList.toggle('open');
|
||||
}
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
"""
|
||||
|
||||
|
||||
@app.get("/", response_class=HTMLResponse)
|
||||
def search_ui():
|
||||
return SEARCH_UI_HTML
|
||||
@@ -0,0 +1,56 @@
|
||||
import tiktoken
|
||||
from .config import CONFIG
|
||||
|
||||
_enc = tiktoken.get_encoding("cl100k_base")
|
||||
|
||||
|
||||
def count_tokens(text):
|
||||
return len(_enc.encode(text))
|
||||
|
||||
|
||||
def chunk_text(text, chunk_size=None, overlap=None):
|
||||
chunk_size = chunk_size or CONFIG.chunk_size_tokens
|
||||
overlap = overlap or CONFIG.chunk_overlap_tokens
|
||||
|
||||
paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
|
||||
|
||||
chunks = []
|
||||
current_chunk = []
|
||||
current_tokens = 0
|
||||
|
||||
for para in paragraphs:
|
||||
para_tokens = count_tokens(para)
|
||||
|
||||
if para_tokens > chunk_size:
|
||||
sentences = para.replace("\n", " ").split(". ")
|
||||
for sent in sentences:
|
||||
if not sent.strip():
|
||||
continue
|
||||
sent_tokens = count_tokens(sent)
|
||||
if current_tokens + sent_tokens > chunk_size and current_chunk:
|
||||
chunks.append("\n\n".join(current_chunk))
|
||||
current_chunk = current_chunk[-1:] if current_chunk else []
|
||||
current_tokens = count_tokens("\n\n".join(current_chunk)) if current_chunk else 0
|
||||
current_chunk.append(sent)
|
||||
current_tokens += sent_tokens
|
||||
else:
|
||||
if current_tokens + para_tokens > chunk_size and current_chunk:
|
||||
chunks.append("\n\n".join(current_chunk))
|
||||
overlap_tokens = 0
|
||||
overlap_paras = []
|
||||
for p in reversed(current_chunk):
|
||||
pt = count_tokens(p)
|
||||
if overlap_tokens + pt <= overlap:
|
||||
overlap_paras.insert(0, p)
|
||||
overlap_tokens += pt
|
||||
else:
|
||||
break
|
||||
current_chunk = overlap_paras
|
||||
current_tokens = overlap_tokens
|
||||
current_chunk.append(para)
|
||||
current_tokens += para_tokens
|
||||
|
||||
if current_chunk:
|
||||
chunks.append("\n\n".join(current_chunk))
|
||||
|
||||
return chunks
|
||||
@@ -0,0 +1,21 @@
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
pg_host: str = os.getenv("KB_PG_HOST", "192.168.1.114")
|
||||
pg_port: int = int(os.getenv("KB_PG_PORT", "5432"))
|
||||
pg_db: str = os.getenv("KB_PG_DB", "kb")
|
||||
pg_user: str = os.getenv("KB_PG_USER", "kb_user")
|
||||
pg_password: str = os.getenv("KB_PG_PASSWORD", "")
|
||||
|
||||
ollama_url: str = os.getenv("KB_OLLAMA_URL", "http://192.168.1.116:11434")
|
||||
embed_model: str = os.getenv("KB_EMBED_MODEL", "nomic-embed-text")
|
||||
|
||||
corpus_root: str = os.getenv("KB_CORPUS_ROOT", "/mnt/nas")
|
||||
|
||||
chunk_size_tokens: int = 600
|
||||
chunk_overlap_tokens: int = 80
|
||||
embed_batch_size: int = 16
|
||||
|
||||
CONFIG = Config()
|
||||
@@ -0,0 +1,150 @@
|
||||
import psycopg
|
||||
from psycopg.rows import dict_row
|
||||
from pgvector.psycopg import register_vector
|
||||
from contextlib import contextmanager
|
||||
from .config import CONFIG
|
||||
|
||||
|
||||
@contextmanager
|
||||
def get_conn():
|
||||
conn = psycopg.connect(
|
||||
host=CONFIG.pg_host,
|
||||
port=CONFIG.pg_port,
|
||||
dbname=CONFIG.pg_db,
|
||||
user=CONFIG.pg_user,
|
||||
password=CONFIG.pg_password,
|
||||
row_factory=dict_row,
|
||||
)
|
||||
register_vector(conn)
|
||||
try:
|
||||
yield conn
|
||||
conn.commit()
|
||||
except Exception:
|
||||
conn.rollback()
|
||||
raise
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def upsert_document(source_path, file_hash, file_type, category, title, last_modified, metadata):
|
||||
with get_conn() as conn:
|
||||
cur = conn.cursor()
|
||||
cur.execute("SELECT id, file_hash FROM kb.documents WHERE source_path = %s", (source_path,))
|
||||
existing = cur.fetchone()
|
||||
|
||||
if existing and existing["file_hash"] == file_hash:
|
||||
return existing["id"], False
|
||||
|
||||
if existing:
|
||||
cur.execute("""
|
||||
UPDATE kb.documents
|
||||
SET file_hash = %s, file_type = %s, category = %s, title = %s,
|
||||
last_modified = %s, indexed_at = NOW(), metadata = %s
|
||||
WHERE id = %s RETURNING id
|
||||
""", (file_hash, file_type, category, title, last_modified,
|
||||
psycopg.types.json.Jsonb(metadata), existing["id"]))
|
||||
doc_id = cur.fetchone()["id"]
|
||||
cur.execute("DELETE FROM kb.chunks WHERE document_id = %s", (doc_id,))
|
||||
else:
|
||||
cur.execute("""
|
||||
INSERT INTO kb.documents
|
||||
(source_path, file_hash, file_type, category, title, last_modified, metadata)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s) RETURNING id
|
||||
""", (source_path, file_hash, file_type, category, title, last_modified,
|
||||
psycopg.types.json.Jsonb(metadata)))
|
||||
doc_id = cur.fetchone()["id"]
|
||||
|
||||
return doc_id, True
|
||||
|
||||
|
||||
def insert_chunks(document_id, chunks):
|
||||
with get_conn() as conn:
|
||||
cur = conn.cursor()
|
||||
for c in chunks:
|
||||
cur.execute("""
|
||||
INSERT INTO kb.chunks
|
||||
(document_id, chunk_index, content, embedding, token_count, metadata)
|
||||
VALUES (%s, %s, %s, %s, %s, %s)
|
||||
""", (document_id, c["chunk_index"], c["content"], c["embedding"],
|
||||
c["token_count"], psycopg.types.json.Jsonb(c.get("metadata", {}))))
|
||||
|
||||
|
||||
def search(query_embedding, k=10, category=None):
|
||||
with get_conn() as conn:
|
||||
cur = conn.cursor()
|
||||
if category:
|
||||
cur.execute("""
|
||||
SELECT
|
||||
c.id as chunk_id, c.content, c.chunk_index,
|
||||
c.metadata as chunk_metadata, c.document_id,
|
||||
d.source_path, d.category, d.title, d.metadata as doc_metadata,
|
||||
1 - (c.embedding <=> %s::vector) as similarity
|
||||
FROM kb.chunks c
|
||||
JOIN kb.documents d ON c.document_id = d.id
|
||||
WHERE d.category = %s
|
||||
ORDER BY c.embedding <=> %s::vector
|
||||
LIMIT %s
|
||||
""", (query_embedding, category, query_embedding, k))
|
||||
else:
|
||||
cur.execute("""
|
||||
SELECT
|
||||
c.id as chunk_id, c.content, c.chunk_index,
|
||||
c.metadata as chunk_metadata, c.document_id,
|
||||
d.source_path, d.category, d.title, d.metadata as doc_metadata,
|
||||
1 - (c.embedding <=> %s::vector) as similarity
|
||||
FROM kb.chunks c
|
||||
JOIN kb.documents d ON c.document_id = d.id
|
||||
ORDER BY c.embedding <=> %s::vector
|
||||
LIMIT %s
|
||||
""", (query_embedding, query_embedding, k))
|
||||
return cur.fetchall()
|
||||
|
||||
|
||||
def search_with_neighbors(query_embedding, k=10, category=None, neighbor_window=1):
|
||||
"""Search and include neighboring chunks for more context."""
|
||||
primary = search(query_embedding, k=k, category=category)
|
||||
|
||||
if not primary or neighbor_window <= 0:
|
||||
return primary
|
||||
|
||||
with get_conn() as conn:
|
||||
cur = conn.cursor()
|
||||
for r in primary:
|
||||
doc_id = r["document_id"]
|
||||
chunk_idx = r["chunk_index"]
|
||||
|
||||
cur.execute("""
|
||||
SELECT chunk_index, content, metadata
|
||||
FROM kb.chunks
|
||||
WHERE document_id = %s
|
||||
AND chunk_index BETWEEN %s AND %s
|
||||
ORDER BY chunk_index
|
||||
""", (doc_id, chunk_idx - neighbor_window, chunk_idx + neighbor_window))
|
||||
|
||||
neighbors = cur.fetchall()
|
||||
# Combine into one expanded passage
|
||||
r["expanded_content"] = "\n\n[...]\n\n".join(n["content"] for n in neighbors)
|
||||
r["neighbor_count"] = len(neighbors) - 1
|
||||
|
||||
return primary
|
||||
|
||||
|
||||
def get_stats():
|
||||
with get_conn() as conn:
|
||||
cur = conn.cursor()
|
||||
cur.execute("SELECT COUNT(*) as total_docs FROM kb.documents")
|
||||
docs = cur.fetchone()
|
||||
cur.execute("SELECT COUNT(*) as total_chunks FROM kb.chunks")
|
||||
chunks = cur.fetchone()
|
||||
cur.execute("""
|
||||
SELECT category, COUNT(*) as count
|
||||
FROM kb.documents
|
||||
GROUP BY category
|
||||
ORDER BY count DESC
|
||||
""")
|
||||
by_cat = cur.fetchall()
|
||||
return {
|
||||
"total_documents": docs["total_docs"],
|
||||
"total_chunks": chunks["total_chunks"],
|
||||
"by_category": by_cat,
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
import requests
|
||||
from .config import CONFIG
|
||||
|
||||
|
||||
def embed_one(text):
|
||||
"""Embed a single text. Handles both old and new Ollama embedding APIs."""
|
||||
try:
|
||||
resp = requests.post(
|
||||
f"{CONFIG.ollama_url}/api/embed",
|
||||
json={"model": CONFIG.embed_model, "input": text},
|
||||
timeout=120,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
if "embeddings" in data:
|
||||
return data["embeddings"][0]
|
||||
if "embedding" in data:
|
||||
return data["embedding"]
|
||||
except Exception:
|
||||
resp = requests.post(
|
||||
f"{CONFIG.ollama_url}/api/embeddings",
|
||||
json={"model": CONFIG.embed_model, "prompt": text},
|
||||
timeout=120,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return resp.json()["embedding"]
|
||||
|
||||
raise RuntimeError(f"Unexpected embedding response: {data}")
|
||||
|
||||
|
||||
def embed_batch(texts):
|
||||
"""Embed multiple texts. Use new API native batching when possible."""
|
||||
try:
|
||||
resp = requests.post(
|
||||
f"{CONFIG.ollama_url}/api/embed",
|
||||
json={"model": CONFIG.embed_model, "input": texts},
|
||||
timeout=300,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
if "embeddings" in data:
|
||||
return data["embeddings"]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return [embed_one(t) for t in texts]
|
||||
+139
@@ -0,0 +1,139 @@
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
from datetime import datetime, timezone
|
||||
from .config import CONFIG
|
||||
from .parsers import parse_file, PARSERS
|
||||
from .chunk import chunk_text, count_tokens
|
||||
from .embed import embed_batch
|
||||
from . import db
|
||||
|
||||
|
||||
def file_hash(path):
|
||||
h = hashlib.sha256()
|
||||
with open(path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(8192), b""):
|
||||
h.update(chunk)
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
def category_from_path(path):
|
||||
try:
|
||||
rel = path.relative_to(CONFIG.corpus_root)
|
||||
except ValueError:
|
||||
return "unknown"
|
||||
parts = rel.parts[:-1]
|
||||
return "/".join(parts) if parts else "uncategorized"
|
||||
|
||||
|
||||
def build_chunk_prefix(title, category, last_modified, doc_metadata):
|
||||
"""Construct a metadata breadcrumb to prepend to each chunk."""
|
||||
parts = []
|
||||
if last_modified:
|
||||
parts.append(last_modified.strftime("%Y-%m-%d"))
|
||||
if category:
|
||||
parts.append(f"Category: {category}")
|
||||
if title:
|
||||
parts.append(f"Title: {title}")
|
||||
fmt = doc_metadata.get("format")
|
||||
if fmt:
|
||||
parts.append(f"Format: {fmt}")
|
||||
return f"[{' | '.join(parts)}]\n\n" if parts else ""
|
||||
|
||||
|
||||
def ingest_file(path):
|
||||
path = Path(path)
|
||||
if path.suffix.lower() not in PARSERS:
|
||||
return {"path": str(path), "status": "skipped", "reason": "unsupported type"}
|
||||
|
||||
fhash = file_hash(path)
|
||||
category = category_from_path(path)
|
||||
|
||||
try:
|
||||
title, text, metadata = parse_file(path)
|
||||
except Exception as e:
|
||||
return {"path": str(path), "status": "error", "reason": f"parse failed: {e}"}
|
||||
|
||||
last_modified = datetime.fromtimestamp(path.stat().st_mtime, tz=timezone.utc)
|
||||
doc_id, changed = db.upsert_document(
|
||||
source_path=str(path),
|
||||
file_hash=fhash,
|
||||
file_type=path.suffix.lower(),
|
||||
category=category,
|
||||
title=title,
|
||||
last_modified=last_modified,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
if not changed:
|
||||
return {"path": str(path), "status": "unchanged", "doc_id": doc_id}
|
||||
|
||||
chunks = chunk_text(text)
|
||||
if not chunks:
|
||||
return {"path": str(path), "status": "empty", "doc_id": doc_id}
|
||||
|
||||
# Build per-chunk topic labels using the LLM
|
||||
from .topic_extractor import extract_topic
|
||||
|
||||
base_prefix = build_chunk_prefix(title, category, last_modified, metadata)
|
||||
|
||||
chunk_records = []
|
||||
embeddings_to_compute = []
|
||||
|
||||
for i, chunk in enumerate(chunks):
|
||||
# Generate per-chunk topic label
|
||||
topic = extract_topic(chunk, doc_title=title, doc_category=category)
|
||||
|
||||
if topic:
|
||||
chunk_with_context = f"{base_prefix}Topic: {topic}\n\n{chunk}"
|
||||
else:
|
||||
chunk_with_context = base_prefix + chunk
|
||||
|
||||
embeddings_to_compute.append(chunk_with_context)
|
||||
chunk_records.append({
|
||||
"chunk_index": i,
|
||||
"content_with_context": chunk_with_context,
|
||||
"raw_chunk": chunk,
|
||||
"topic": topic,
|
||||
})
|
||||
|
||||
# Batch embed
|
||||
embeddings = []
|
||||
for i in range(0, len(embeddings_to_compute), CONFIG.embed_batch_size):
|
||||
batch = embeddings_to_compute[i:i + CONFIG.embed_batch_size]
|
||||
embeddings.extend(embed_batch(batch))
|
||||
|
||||
# Build final records for DB
|
||||
final_records = [
|
||||
{
|
||||
"chunk_index": rec["chunk_index"],
|
||||
"content": rec["content_with_context"],
|
||||
"embedding": emb,
|
||||
"token_count": count_tokens(rec["content_with_context"]),
|
||||
"metadata": {
|
||||
"raw_chunk": rec["raw_chunk"],
|
||||
"topic": rec["topic"],
|
||||
},
|
||||
}
|
||||
for rec, emb in zip(chunk_records, embeddings)
|
||||
]
|
||||
db.insert_chunks(doc_id, final_records)
|
||||
|
||||
return {
|
||||
"path": str(path),
|
||||
"status": "indexed",
|
||||
"doc_id": doc_id,
|
||||
"chunks": len(chunks),
|
||||
"category": category,
|
||||
}
|
||||
|
||||
|
||||
def ingest_directory(root=None):
|
||||
root = Path(root) if root else Path(CONFIG.corpus_root)
|
||||
results = []
|
||||
for path in root.rglob("*"):
|
||||
if path.is_file() and path.suffix.lower() in PARSERS:
|
||||
print(f"Processing: {path}", flush=True)
|
||||
result = ingest_file(path)
|
||||
results.append(result)
|
||||
print(f" → {result['status']}", flush=True)
|
||||
return results
|
||||
@@ -0,0 +1,85 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def parse_text(path):
|
||||
text = path.read_text(encoding="utf-8", errors="replace")
|
||||
return path.stem, text, {}
|
||||
|
||||
|
||||
def parse_markdown(path):
|
||||
text = path.read_text(encoding="utf-8", errors="replace")
|
||||
title = path.stem
|
||||
for line in text.split("\n"):
|
||||
if line.startswith("# "):
|
||||
title = line.lstrip("#").strip()
|
||||
break
|
||||
return title, text, {}
|
||||
|
||||
|
||||
def parse_pdf(path):
|
||||
from pypdf import PdfReader
|
||||
reader = PdfReader(str(path))
|
||||
pages = [p.extract_text() or "" for p in reader.pages]
|
||||
text = "\n\n".join(pages)
|
||||
title = path.stem
|
||||
if reader.metadata and reader.metadata.title:
|
||||
title = reader.metadata.title
|
||||
return title, text, {"page_count": len(pages)}
|
||||
|
||||
|
||||
def parse_claude_export(path):
|
||||
data = json.loads(path.read_text(encoding="utf-8", errors="replace"))
|
||||
|
||||
if isinstance(data, dict) and ("name" in data or "title" in data):
|
||||
return _format_conversation(data, path.stem)
|
||||
elif isinstance(data, list):
|
||||
all_text = []
|
||||
for conv in data:
|
||||
_, text, _ = _format_conversation(conv, "")
|
||||
all_text.append(text)
|
||||
return path.stem, "\n\n---\n\n".join(all_text), {"conversation_count": len(data)}
|
||||
else:
|
||||
return path.stem, json.dumps(data, indent=2), {}
|
||||
|
||||
|
||||
def _format_conversation(conv, default_title):
|
||||
title = conv.get("name") or conv.get("title") or default_title
|
||||
messages = conv.get("messages", []) or conv.get("chat_messages", [])
|
||||
|
||||
parts = [f"# {title}\n"]
|
||||
for msg in messages:
|
||||
role_raw = msg.get("role") or msg.get("sender") or "unknown"
|
||||
# Normalize to USER/ASSISTANT for consistency with OpenClaw transcripts
|
||||
role = "USER" if role_raw == "human" else ("ASSISTANT" if role_raw == "assistant" else role_raw.upper())
|
||||
content = msg.get("content") or msg.get("text") or ""
|
||||
if isinstance(content, list):
|
||||
content = "\n".join(
|
||||
block.get("text", "") for block in content
|
||||
if isinstance(block, dict) and block.get("type") == "text"
|
||||
)
|
||||
if not content.strip():
|
||||
continue # skip empty blocks (tool_use, tool_result, etc.)
|
||||
parts.append(f"\n**{role}**:\n\n{content}\n")
|
||||
|
||||
return title, "\n".join(parts), {
|
||||
"message_count": len(messages),
|
||||
"format": "claude_export",
|
||||
}
|
||||
|
||||
|
||||
PARSERS = {
|
||||
".txt": parse_text,
|
||||
".md": parse_markdown,
|
||||
".markdown": parse_markdown,
|
||||
".pdf": parse_pdf,
|
||||
".json": parse_claude_export,
|
||||
}
|
||||
|
||||
|
||||
def parse_file(path):
|
||||
suffix = path.suffix.lower()
|
||||
parser = PARSERS.get(suffix)
|
||||
if not parser:
|
||||
raise ValueError(f"No parser for {suffix}")
|
||||
return parser(path)
|
||||
@@ -0,0 +1,83 @@
|
||||
"""Reranking via Qwen3-Reranker over Ollama."""
|
||||
import requests
|
||||
import os
|
||||
import re
|
||||
from .config import CONFIG
|
||||
|
||||
RERANK_MODEL = os.getenv("KB_RERANK_MODEL", "dengcao/Qwen3-Reranker-4B:Q5_K_M")
|
||||
|
||||
|
||||
def score_relevance(query: str, document: str) -> float:
|
||||
"""Get a relevance score in [0, 1] for (query, document) pair."""
|
||||
if len(document) > 4000:
|
||||
document = document[:4000] + "..."
|
||||
|
||||
# Pre-fill <think></think> as empty so the model jumps straight to answer
|
||||
prompt = (
|
||||
f"<|im_start|>system\n"
|
||||
f"You evaluate whether a document is relevant to a query about quantitative trading. "
|
||||
f"Be strict: the document must specifically address the query. "
|
||||
f"Output only \"yes\" or \"no\".<|im_end|>\n"
|
||||
f"<|im_start|>user\n"
|
||||
f"<Query>: {query}\n"
|
||||
f"<Document>: {document}<|im_end|>\n"
|
||||
f"<|im_start|>assistant\n"
|
||||
f"<think>\n\n</think>\n\n"
|
||||
)
|
||||
|
||||
try:
|
||||
resp = requests.post(
|
||||
f"{CONFIG.ollama_url}/api/generate",
|
||||
json={
|
||||
"model": RERANK_MODEL,
|
||||
"prompt": prompt,
|
||||
"stream": False,
|
||||
"options": {
|
||||
"num_predict": 5,
|
||||
"temperature": 0.0,
|
||||
"top_p": 1.0,
|
||||
},
|
||||
"raw": True,
|
||||
},
|
||||
timeout=60,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
text = resp.json().get("response", "").strip().lower()
|
||||
text = re.sub(r"</?think>", "", text).strip()
|
||||
resp.raise_for_status()
|
||||
text = resp.json().get("response", "").strip()
|
||||
text = re.sub(r"</?think>", "", text).strip()
|
||||
|
||||
# DEBUG: print every response so we can see what's happening
|
||||
print(f"DEBUG rerank response: {text[:80]!r}", flush=True)
|
||||
|
||||
text = text.lower()
|
||||
|
||||
if text.startswith("yes"):
|
||||
return 1.0
|
||||
elif text.startswith("no"):
|
||||
return 0.0
|
||||
else:
|
||||
print(f"Rerank ambiguous response: {text[:100]!r}", flush=True)
|
||||
return 0.5
|
||||
except Exception as e:
|
||||
print(f"Rerank error: {e}", flush=True)
|
||||
return 0.5
|
||||
|
||||
|
||||
def rerank(query: str, candidates: list, top_k: int = 10) -> list:
|
||||
if not candidates:
|
||||
return []
|
||||
|
||||
scored = []
|
||||
for cand in candidates:
|
||||
chunk_meta = cand.get("chunk_metadata") or {}
|
||||
doc_text = chunk_meta.get("raw_chunk") or cand.get("content", "")
|
||||
score = score_relevance(query, doc_text)
|
||||
cand_with_score = dict(cand)
|
||||
cand_with_score["rerank_score"] = score
|
||||
cand_with_score["original_similarity"] = cand.get("similarity")
|
||||
scored.append(cand_with_score)
|
||||
|
||||
scored.sort(key=lambda x: x["rerank_score"], reverse=True)
|
||||
return scored[:top_k]
|
||||
@@ -0,0 +1,66 @@
|
||||
"""Extract topic labels for chunks using a local LLM."""
|
||||
import requests
|
||||
import os
|
||||
from .config import CONFIG
|
||||
|
||||
# Use a fast model — chosen at config time
|
||||
TOPIC_MODEL = os.getenv("KB_TOPIC_MODEL", "qwen2.5-coder:7b")
|
||||
|
||||
|
||||
def extract_topic(chunk_text, doc_title=None, doc_category=None):
|
||||
"""Generate a brief topic label for a chunk.
|
||||
|
||||
Returns a single-line description of what specific entities,
|
||||
concepts, and topics this chunk discusses.
|
||||
"""
|
||||
# Truncate very long chunks for the LLM
|
||||
if len(chunk_text) > 3000:
|
||||
chunk_text = chunk_text[:3000] + "..."
|
||||
|
||||
context_hint = ""
|
||||
if doc_title:
|
||||
context_hint += f"\nThe broader document is titled: {doc_title}"
|
||||
if doc_category:
|
||||
context_hint += f"\nThe document category: {doc_category}"
|
||||
|
||||
prompt = (
|
||||
f"<|im_start|>system\n"
|
||||
f"You extract topic labels for retrieval. Given a passage, identify the specific "
|
||||
f"entities, concepts, names, and topics it discusses. Output a single concise line "
|
||||
f"of 10-30 words listing what this passage is specifically about. Include proper "
|
||||
f"nouns, technical terms, and metric values. Do not summarize the content; just "
|
||||
f"label what it covers.<|im_end|>\n"
|
||||
f"<|im_start|>user\n"
|
||||
f"Passage:{context_hint}\n\n{chunk_text}<|im_end|>\n"
|
||||
f"<|im_start|>assistant\n"
|
||||
)
|
||||
|
||||
try:
|
||||
resp = requests.post(
|
||||
f"{CONFIG.ollama_url}/api/generate",
|
||||
json={
|
||||
"model": TOPIC_MODEL,
|
||||
"prompt": prompt,
|
||||
"stream": False,
|
||||
"options": {
|
||||
"num_predict": 60,
|
||||
"temperature": 0.0,
|
||||
"top_p": 1.0,
|
||||
},
|
||||
"raw": True,
|
||||
},
|
||||
timeout=60,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
text = resp.json().get("response", "").strip()
|
||||
|
||||
# Take only the first line, drop any "Labels:" prefix
|
||||
first_line = text.split("\n")[0].strip()
|
||||
for prefix in ("Topics:", "Topic:", "Labels:", "Label:", "Subject:"):
|
||||
if first_line.lower().startswith(prefix.lower()):
|
||||
first_line = first_line[len(prefix):].strip()
|
||||
|
||||
return first_line[:300] # cap length
|
||||
except Exception as e:
|
||||
print(f"Topic extraction error: {e}", flush=True)
|
||||
return ""
|
||||
@@ -0,0 +1,27 @@
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from watchfiles import awatch, Change
|
||||
from .config import CONFIG
|
||||
from .ingest import ingest_file
|
||||
from .parsers import PARSERS
|
||||
|
||||
|
||||
async def watch_corpus():
|
||||
print(f"Watching {CONFIG.corpus_root} for changes...", flush=True)
|
||||
async for changes in awatch(CONFIG.corpus_root):
|
||||
for change_type, path_str in changes:
|
||||
path = Path(path_str)
|
||||
if path.suffix.lower() not in PARSERS:
|
||||
continue
|
||||
if change_type in (Change.added, Change.modified):
|
||||
if path.is_file():
|
||||
print(f"Change detected: {path}", flush=True)
|
||||
try:
|
||||
result = ingest_file(path)
|
||||
print(f" → {result['status']}", flush=True)
|
||||
except Exception as e:
|
||||
print(f" → error: {e}", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(watch_corpus())
|
||||
Reference in New Issue
Block a user