import "dotenv/config";
import { ChatOpenAI, OpenAIEmbeddings } from '@langchain/openai'
import { Document } from '@langchain/core/documents'
import { MemoryVectorStore } from '@langchain/classic/vectorstores/memory'

// 创建模型
const model = new ChatOpenAI({
    model: process.env.MODEL_NAME,
    apiKey: process.env.OPENAI_API_KEY,
    temperature: 0,
    configuration: {
        baseURL: process.env.OPENAI_BASE_URL,
    },
})

// 创建嵌入模型
const embeddings = new OpenAIEmbeddings({
    model: process.env.EMBEDDINGS_MODEL_NAME,
    apiKey: process.env.OPENAI_API_KEY,
    configuration: {
        baseURL: process.env.OPENAI_BASE_URL,
    },
})


// 创建文档
const documents = [
    new Document({
        pageContent: "小胡和小赵是同班同学。开学第一天，教室里人声嘈杂，小胡抱着新领的课本找座位，手里还捏着没来得及喝的热豆浆。小赵看见他站在门口发愣，轻轻拍了拍旁边的空位，说：“这儿没人，你坐。”小胡坐下时，豆浆差点洒出来，小赵连忙递过纸巾。就这样，他们认识了。之后的日子里，他们一起早起占自习座、一起在操场绕圈散步、一起在食堂排队等一份热汤面。很多人以为那只是顺路的照应，可小胡知道——每次小赵回头看他一眼，他心里就会安静下来，像有盏灯被点亮。",
        metadata: {
            chapter: 1,
            character: '小胡和小赵',
            type: '爱情情节',
            mood: '悸动',
        },
    }),
    new Document({
        pageContent: "期中考试前的那段时间，班里人人都在拼命。小胡做题做到半夜，总觉得怎么都赶不上别人的脚步，纸上密密麻麻写满了又擦掉的字。他强撑着说“没事”，可眼神里藏不住疲惫。小赵没说大道理，只是把自己抄得最整齐的笔记放到小胡桌上，又在课间把一颗热乎乎的烤红薯塞进他手里，说：“你先吃点东西，别把自己熬坏。”小胡捧着那点温热，忽然觉得鼻子发酸。他发现，小赵不是来替他解决难题的，而是来把他从孤单里拉出来的。",
        metadata: {
            chapter: 2,
            character: '小胡和小赵',
            type: '爱情情节',
            mood: '心疼',
        },
    }),
    new Document({
        pageContent: "考试成绩出来那天，小胡在公告栏前站了很久，手指发冷。看到自己的名字排在合格线上方，他才长长吐出一口气。人群散去后，小赵还站在他身边，像以前一样把书包带子往肩上提了提，笑着说：“我就知道你可以。”小胡低着头，嗓子发紧：“如果没有你，我可能早就撑不住了。”小赵没躲开他的目光，只轻声说：“你别把自己一个人关起来。我愿意陪你走。”那句话太朴素，却像一把钥匙，轻轻打开了小胡心里最久的门。",
        metadata: {
            chapter: 3,
            character: '小胡和小赵',
            type: '爱情情节',
            mood: '告白前的颤抖',
        },
    }),
    new Document({
        pageContent: "学校办运动会那天，小胡被临时推上去参加接力。发令枪响时，他心里空得发慌，脚步差点乱了。跑到弯道，他听见有人在看台上喊他的名字，声音不大，却一声一声地稳。那是小赵。最后一棒，小胡把接力棒递出去时，手还在抖。赛后他坐在台阶上喘气，小赵递来一瓶水，拧开盖子才递给他，像把一份细心藏在最普通的动作里。小胡忽然明白，输赢都不重要，重要的是无论他跑得快不快，总有人在看台上为他亮着嗓子。",
        metadata: {
            chapter: 4,
            character: '小胡和小赵',
            type: '爱情情节',
            mood: '坚定',
        },
    }),
    new Document({
        pageContent: "真正难熬的是毕业那年。小胡要回老家帮家里撑一段时间，小赵留在城里实习。分别那天，车站风很冷，小胡把围巾绕了两圈，还是觉得脖子发凉。他想说点轻松的话，却怎么都开不了口。小赵把一小包药塞进他口袋，说：“路上别硬扛，头疼就吃。”又把一张折得方方正正的纸条递给他：“到了再看。”车开了，小胡在窗边忍了很久，还是没忍住拆开纸条。上面只有一句话：——“你难的时候，就想想还有我。”字写得很平，却把小胡的眼泪一下子写出来了。他第一次清楚地知道：有些人不是“朋友”两个字能装得下的。",
        metadata: {
            chapter: 5,
            character: '小胡和小赵',
            type: '爱情情节',
            mood: '离别',
        },
    }),
    new Document({
        pageContent: "三年后，小胡回城办事，顺路去母校门口看了一眼。校门口那家老面馆还在，门口的风铃叮当作响。小胡推门进去，抬头就看见小赵坐在靠窗的位置，桌上放着两碗热面，一碗已经晾到刚好入口。小赵像早就知道他会来似的，抬眼笑了：“你瘦了。”小胡喉头一紧，半天只说出一句：“我回来了。”吃到一半，小胡从钱包里摸出那张皱了又展、展了又皱的纸条，放在桌上：“这三年，我靠它走过来。”小赵的眼眶一下红了。小胡低声说：“我不想再把你放在‘以后’里了。小赵，我们能不能……把日子过在一起？”小赵没说很多，只点了点头，眼泪掉进碗里，又被热气悄悄蒸走。",
        metadata: {
            chapter: 6,
            character: '小胡和小赵',
            type: '爱情情节',
            mood: '重逢与告白',
        },
    }),
    new Document({
        pageContent: "日子真要过在一起，才知道并不总是风平浪静。小赵的母亲突然生了一场病，住院那段时间，小赵白天上班，晚上守在病房，整个人像被掏空。小胡赶来时，小赵还逞强说“我撑得住”。小胡没说什么，只把热粥递给他，又去走廊尽头打了盆温水，帮他把皱巴巴的衣领理平。夜里病房安静，小赵终于压着声音哭出来：“我怕。”小胡把他抱住，额头贴着额头，说：“怕就靠着我。你从前把我从苦里拉出来，现在换我陪你熬。”那一晚，窗外的雨一直下，小胡的手一直没松开。小赵忽然明白：爱不是轰轰烈烈的誓言，而是你最狼狈的时候，有人还愿意把你当成宝。",
        metadata: {
            chapter: 7,
            character: '小胡和小赵',
            type: '爱情情节',
            mood: '守护',
        },
    }),
    new Document({
        pageContent: "后来有一次，同学聚会里有人喝多了，拿他们开玩笑，说得难听。小赵脸上挂不住，想起身离开。小胡却先站起来，给那人倒了杯水，语气很平却不容人再放肆：“别拿别人的真心当笑话。”屋子里一下安静了。回家的路上，夜风很凉，小赵走得很快，像怕眼泪被看见。小胡追上去，握住他的手，说：“你不用总是自己扛。你难过，我就替你挡一挡。”小赵停下脚步，低头看着两个人紧扣的手指，忽然像松了一口气，声音哑得厉害：“我这辈子最怕的，是你哪天松开我。”小胡把他的手握得更紧：“我不会松。我们把日子慢慢过下去。”",
        metadata: {
            chapter: 8,
            character: '小胡',
            type: '爱情情节',
            mood: '笃定',
        },
    }),
]

// 创建向量存储
const vectorStore = await MemoryVectorStore.fromDocuments(documents, embeddings)

// 创建检索器，k=3表示返回3个最相关的文档
const retriever = vectorStore.asRetriever({ k: 3 })

const questions = [
    "小胡和小赵是怎样一步步走到一起的？"
    // '小胡和小赵有没有结婚？',
    // '小张和小李的故事是怎样的？',
]

for (const question of questions) {
    console.log('--------------------------------')
    console.log(`问题: ${question}`)
    console.log('--------------------------------')
    

    // 检索相关文档
    const retrieverDocs = await retriever.invoke(question)

    // 计算相似度
    const scoredDocs = await vectorStore.similaritySearchWithScore(question, 3)
    

    console.log('\n 【检索到的文档和相似度评分】');
    for (const [doc, score] of scoredDocs) {
        console.log(`文档: ${doc.pageContent}`)
        console.log(`相似度评分: ${score}`)
        console.log();
    }
    console.log('--------------------------------')


    const context = retrieverDocs.map(doc => doc.pageContent).join('\n')

    const prompt = `你是一个讲故事的老师，基于以下故事片段回答问题，用温暖生动的语言。如果故事中没有提到，就说“这个故事没有提到这个情节”。
    故事片段: ${context}
    问题: ${question}
    回答:
    `

    const response = await model.invoke(prompt)
    console.log('\n--故事回答------------------------------\n')
    console.log(response.content)
    console.log('\n----------------------------------------\n')
}