Initial commit — kb-app RAG server
This commit is contained in:
Julian Carlson
@@ -0,0 +1,6 @@
|
|||||||
|
venv/
|
||||||
|
__pycache__/
|
||||||
|
*.pyc
|
||||||
|
*.pyo
|
||||||
|
.env
|
||||||
|
*.log
|
||||||
+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())
|
||||||
@@ -0,0 +1,39 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Parallel ingestion of a directory."""
|
||||||
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||||
|
from pathlib import Path
|
||||||
|
import sys
|
||||||
|
sys.path.insert(0, '/home/julianocarlson/kb')
|
||||||
|
|
||||||
|
from app.ingest import ingest_file
|
||||||
|
from app.parsers import PARSERS
|
||||||
|
|
||||||
|
def main(root_path, max_workers=4):
|
||||||
|
root = Path(root_path)
|
||||||
|
files = [p for p in root.rglob("*")
|
||||||
|
if p.is_file() and p.suffix.lower() in PARSERS]
|
||||||
|
|
||||||
|
print(f"Found {len(files)} files to process")
|
||||||
|
|
||||||
|
counts = {"indexed": 0, "unchanged": 0, "error": 0, "skipped": 0, "empty": 0}
|
||||||
|
|
||||||
|
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||||
|
futures = {executor.submit(ingest_file, f): f for f in files}
|
||||||
|
for i, future in enumerate(as_completed(futures), 1):
|
||||||
|
try:
|
||||||
|
result = future.result()
|
||||||
|
status = result.get("status", "unknown")
|
||||||
|
counts[status] = counts.get(status, 0) + 1
|
||||||
|
print(f"[{i}/{len(files)}] {status}: {Path(result['path']).name}")
|
||||||
|
except Exception as e:
|
||||||
|
counts["error"] += 1
|
||||||
|
print(f"[{i}/{len(files)}] EXCEPTION: {e}")
|
||||||
|
|
||||||
|
print("\nSummary:")
|
||||||
|
for k, v in counts.items():
|
||||||
|
print(f" {k}: {v}")
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
path = sys.argv[1] if len(sys.argv) > 1 else "/mnt/nas/transcripts"
|
||||||
|
workers = int(sys.argv[2]) if len(sys.argv) > 2 else 4
|
||||||
|
main(path, workers)
|
||||||
Reference in New Issue
Block a user