File size: 3,295 Bytes
577adc4 90b36cb 577adc4 90b36cb 577adc4 90b36cb 577adc4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | /**
* Retrieval utilities: HuggingFace embeddings + TigerGraph vector search
*/
export interface TGChunk {
chunk_id: string;
text: string;
score: number;
}
// In-process embedding cache — avoids re-hitting HF API for the same query text.
// Capped at 256 entries to prevent unbounded memory growth in long-running servers.
const embeddingCache = new Map<string, number[]>();
const EMBED_CACHE_MAX = 256;
/** Generate 384-dim embedding via HF Inference API (all-MiniLM-L6-v2) */
export async function getEmbedding(text: string): Promise<number[] | null> {
const normalized = text.trim().toLowerCase();
const cached = embeddingCache.get(normalized);
if (cached) return cached;
const token = process.env.HUGGING_FACE_HUB_TOKEN || process.env.HF_TOKEN;
if (!token) return null;
try {
const res = await fetch(
"https://api-inference.huggingface.co/pipeline/feature-extraction/sentence-transformers/all-MiniLM-L6-v2",
{
method: "POST",
headers: { Authorization: `Bearer ${token}`, "Content-Type": "application/json" },
body: JSON.stringify({ inputs: text, options: { wait_for_model: true } }),
signal: AbortSignal.timeout(15000),
}
);
if (!res.ok) return null;
const data = await res.json();
if (!Array.isArray(data)) return null;
// Handle both [0.1, 0.2, ...] and [[0.1, 0.2, ...]]
const flat: number[] = Array.isArray(data[0]) ? (data[0] as number[]) : (data as number[]);
if (!flat.every((x) => typeof x === "number")) return null;
if (embeddingCache.size >= EMBED_CACHE_MAX) {
embeddingCache.delete(embeddingCache.keys().next().value!);
}
embeddingCache.set(normalized, flat);
return flat;
} catch {
return null;
}
}
/** Call TigerGraph vectorSearchChunks installed query */
export async function searchChunks(embedding: number[], topK = 5): Promise<TGChunk[]> {
const host = (process.env.TG_HOST || "").replace(/\/$/, "");
const token = process.env.TG_TOKEN;
const graph = process.env.TG_GRAPH || "GraphRAG";
if (!host || !token || !embedding.length) return [];
try {
const res = await fetch(`${host}/restpp/query/${graph}/vectorSearchChunks`, {
method: "POST",
headers: { Authorization: `Bearer ${token}`, "Content-Type": "application/json" },
body: JSON.stringify({ queryVec: embedding, topK }),
signal: AbortSignal.timeout(20000),
});
if (!res.ok) return [];
const data = await res.json();
return (data.results?.[0]?.["@@topChunks"] as TGChunk[]) || [];
} catch {
return [];
}
}
/** Extract compact entity descriptions from chunk text (simulates pre-indexed graph data).
* Entity extraction runs at INGEST TIME so the cost is amortized.
* At query time, we only pay for the compact entity context, not full chunk text. */
export function chunkToEntityContext(text: string, maxChars = 220): string {
// Take first sentence — Wikipedia science articles open with the key entity definition
const firstSentence = text.split(/(?<=[.!?])\s+/)[0].trim();
return firstSentence.slice(0, maxChars);
}
/** Rough token count estimate (1 token ≈ 0.75 words) */
export function estimateTokens(text: string): number {
return Math.ceil(text.split(/\s+/).filter(Boolean).length * 1.33);
}
|