Project: ကိုယ်ပိုင် Document အတွက် Q&A Assistant

// A private document Q&A pipeline: ingest documents into a vector store,
// then answer questions using only retrieved chunks -- never by retraining
// or fine-tuning anything.
const EMBED_URL = "http://localhost:11434/api/embeddings";
const CHAT_URL = "http://localhost:11434/api/chat";
const EMBED_MODEL = "nomic-embed-text";
const CHAT_MODEL = "llama3.1:8b";

type StoredChunk = {
  vector: number[];
  text: string;
  source: string;
  chunkIndex: number;
};

const vectorStore: StoredChunk[] = []; // in-memory for a single-user tool

function chunkText(text: string, size = 800, overlap = 150): string[] {
  const chunks: string[] = [];
  let start = 0;
  while (start < text.length) {
    const end = Math.min(start + size, text.length);
    chunks.push(text.slice(start, end));
    if (end === text.length) break;
    start = end - overlap; // overlap so an idea at a boundary isn't lost
  }
  return chunks;
}

async function embed(text: string): Promise<number[]> {
  const res = await fetch(EMBED_URL, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({ model: EMBED_MODEL, prompt: text }),
  });
  if (!res.ok) throw new Error(`Embedding request failed: ${res.status}`);
  const data = (await res.json()) as { embedding: number[] };
  return data.embedding;
}

// --- ingest path: run once per uploaded document ---
export async function ingestDocument(source: string, rawText: string) {
  const chunks = chunkText(rawText);
  for (let i = 0; i < chunks.length; i++) {
    const vector = await embed(chunks[i]);
    vectorStore.push({ vector, text: chunks[i], source, chunkIndex: i });
  }
}

function cosineSimilarity(a: number[], b: number[]): number {
  let dot = 0, normA = 0, normB = 0;
  for (let i = 0; i < a.length; i++) {
    dot += a[i] * b[i];
    normA += a[i] * a[i];
    normB += b[i] * b[i];
  }
  return dot / (Math.sqrt(normA) * Math.sqrt(normB));
}

function searchChunks(queryVector: number[], topK = 4): StoredChunk[] {
  return [...vectorStore]
    .map((chunk) => ({ chunk, score: cosineSimilarity(queryVector, chunk.vector) }))
    .sort((a, b) => b.score - a.score)
    .slice(0, topK)
    .map((r) => r.chunk);
}

// --- query path: run on every question ---
export async function answerFromDocuments(question: string) {
  const questionVector = await embed(question);
  const topChunks = searchChunks(questionVector);

  const context = topChunks
    .map((c, i) => `[${i + 1}] (${c.source}, chunk ${c.chunkIndex})\n${c.text}`)
    .join("\n\n");

  const res = await fetch(CHAT_URL, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({
      model: CHAT_MODEL,
      stream: false,
      messages: [
        {
          role: "system",
          content:
            "Answer only using the numbered context below. If the " +
            "context does not contain the answer, say so -- do not guess. " +
            "Cite context numbers like [1] for every claim.\n\n" + context,
        },
        { role: "user", content: question },
      ],
    }),
  });

  if (!res.ok) throw new Error(`Chat request failed: ${res.status}`);
  const data = (await res.json()) as { message: { content: string } };

  return {
    answer: data.message.content,
    citations: topChunks.map((c) => ({ source: c.source, chunk: c.chunkIndex })),
  };
}