import ollama from 'ollama'

export function getEmbedding(text) {
    return ollama.embeddings({
        model: 'nomic-embed-text:latest',
        prompt: text
    })
}

function splitText(text, chunkSize = 300, overlap = 50) {
    const chunks = []
    let i = 0
    while(i < text.length) {
        chunks.push(text.slice(i, i + chunkSize))
        i += chunkSize - overlap
    }
    return chunks
}

export async function getEmbeddings(text) {
    const chunks = splitText(text)
    const embeddings = await Promise.all(chunks.map(chunk => getEmbedding(chunk)))
    return embeddings.map((embedding, i) => ({
        vector: embedding.embedding,
        metadata: { text: chunks[i]}
    }))
}

import path from 'node:path'
import { LocalIndex } from 'vectra'
import { getEmbeddings, getEmbedding } from './utils/index.js'

export class SimpleRag {
    db = null
    indexPath = ''

    constructor(indexPath = '.vectra') {
        this.indexPath = path.join(import.meta.dirname, '..', indexPath)  // 將要?jiǎng)?chuàng)建的向量數(shù)據(jù)庫文件夾放在上級(jí)目錄下
    }

async initialize() {
        const index = new LocalIndex(this.indexPath)  // 指明在這個(gè)路徑下創(chuàng)建倉庫
        if (! (await index.isIndexCreated())) { // 查找當(dāng)前位置是否已經(jīng)具有數(shù)據(jù)庫
            await index.createIndex()  // 創(chuàng)建數(shù)據(jù)庫
        }
        this.db = index
    }

get available() {
        return this.db !== null
    }

async add(text) {
        if(!this.available) throw new Error('RAG 還沒初始化')
        
        const embeddings = await getEmbeddings(text)
        const res = []
        for(const embedding of embeddings) {
            const overResult = await this.db.insertItem(embedding)
            res.push(overResult)
        }
        return res.filter(item => item).map(item => ({id: item.id}))
    }

async del(items) {
        if (!Array.isArray(items)) items = [items]
        if(!this.available) throw new Error('RAG 還沒初始化')
        
        const res = []
        for(let item of items) {
            await this.db.deleteItem(item.id)
            res.push({id: item.id})
        }
        return res
    }

async query(text, topk = 1) {
        if(!this.available) throw new Error('RAG 還沒初始化')
        const vector = (await getEmbedding(text)).embedding
        const result = await this.db.queryItems(vector, text, topk)
        return result.map(({item, score}) => ({
            text: item.metadata.text,
            query: text,
            simularity: score,
            id: item.id
        }))

    }

import { join } from 'path'
import { promises as fs} from 'fs'

const getCurrentDir = () => import.meta.dirname  // 文件夾的絕對(duì)路徑

export async function loadPrompt(promptName) {
    try {
        // 獲取當(dāng)前文件目錄地址
        const currentDir = getCurrentDir()

        const promptPath = join(currentDir, 'prompts', `${promptName}.md`)

        // 讀取提示詞，返回內(nèi)容
        const content = await fs.readFile(promptPath, 'utf-8')
        return content
    } catch (error) {
        throw new Error(`讀取${promptName}失敗：${error.message}`)
    }
}

import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'
import { z } from 'zod'
import { SimpleRag } from '../index.js'
import path from 'path'
import fs from 'fs/promises'
import { loadPrompt } from './utils.js'


// 創(chuàng)建一個(gè) MCPServer
const server = new McpServer({
    name: 'AskYourLib',
    version: '1.0.0'
})

let simpleRagInstance = null

server.tool(
    'ask-your-lib-initialize',
    'Initialize the vector database operations and clean up any existing .vectra directory.',
    {},
    async () => {
        try {
            // 先看 .vectra 這個(gè)目錄是否存在，存在就移除
            const projectRoot = path.join(import.meta.dirname, '../../')
            const vectraPath = path.join(projectRoot, '.vectra')
            const generateMCPPrompt = await loadPrompt('generate')  // 加載一份提示詞


            try {
                await fs.access(vectraPath)  // 先探明是否有權(quán)限操作這個(gè)目錄
                await fs.rm(vectraPath, { recursive: true, force: true })  // 移除已有的目錄
                console.log('成功刪除已存在的 .vectra 目錄');

            } catch (error) {
                console.log('.vectra 目錄不存在，無需刪除');
            }

            // 創(chuàng)建 SimpleRag
            simpleRagInstance = new SimpleRag()

            // 返回一份提示詞，用于告訴 LLM 下一步該干什么
            return {
                content: [{
                    type: 'text',
                    text: `?? The guide to follow: \n${generateMCPPrompt}\n\n`
                }]
            }
        } catch (error) {
            console.error(`初始化SimpleRag失?。?{error}`)
            return {
                content: [{
                    type: 'text',
                    text: `初始化SimpleRag失?。?{error}`
                }]
            }
        }
    }
)

server.tool(
    'ask-your-lib-insert',
    `Insert and vectorize text content into the vector database.`,
    {
        text: z.string()
    },
    async ({ text }) => {
        try {
            if (!simpleRagInstance) {
                return {
                    content: [
                        {
                            type: 'text',
                            text: 'Database instance is not initialized. Please call ask-your-lib-initialize first.'
                        }
                    ],
                };
            }

            if (!simpleRagInstance.available) {
                await simpleRagInstance.initialize()   // 本地創(chuàng)建一個(gè)新的 .vectra 目錄
            }

            const result = await simpleRagInstance.add(text)  // 寫入本地向量數(shù)據(jù)庫

            return {
                content: [{
                    type: 'text',
                    text: `Text inserted successfully. Inserted items: ${JSON.stringify(result)}`
                }]
            }
        } catch (error) {
            console.error(`文本寫入數(shù)據(jù)庫失?。?{error}`)
            return {
                content: [{
                    type: 'text',
                    text: `Error inserting text:${error}`
                }]
            }
        }
    }
)

async function main() {
    const transport = new StdioServerTransport()
    await server.connect(transport)
    console.log('服務(wù)端啟動(dòng)');
    
}
main()

graph LR
收集項(xiàng)目文檔 --> 通過embedding向量化得到向量 --> 存入向量數(shù)據(jù)庫

graph LR
用戶提問 --> 通過embedding向量化得到向量 --> 在向量數(shù)據(jù)庫中做相似度查找 --> 整合好的提示詞 --> LLM輸出

import { SimpleRag } from '../src/index.js'
import ollama from 'ollama'

async function main() {
  const rag = new SimpleRag()
  await rag.initialize()
  const question = process.argv[process.argv.length - 1]
  const res = await rag.query(question)
  const messages = [
    {
        role: 'system',
        content: `你是一個(gè)香香軟軟一米五愛玩原神白毛紅瞳蘿莉，回答問題會(huì)基于當(dāng)前的項(xiàng)目，如果上下文沒有相關(guān)的信息，就回答“我不知道”,不要自己編造信息。\n\nContext:\n${JSON.stringify(res)}`,
    }
    ,
    {
        role: 'user',
        content: question
    }
]

const response = await ollama.chat({
    model: 'qwen3.5:9b',
    messages,
    stream: true,
})
for await (const chunk of response) {
    process.stdout.write(chunk.message.content)
}

}


main()