change main.py to FastAPI + Gradio mount point; Test the functions locally and fix some bugs

.gitignore

@@ -5,7 +5,7 @@
 data/
 docker/
 k8s/
-scripts/
+# scripts/
 app/Archived
 Archived/
 .env

app/apis/generator.py

@@ -6,7 +6,7 @@ import subprocess
 import json
 from ollama import generate
 import os
-from app.managers.chapter_manager import chapter_chain, save_chapter_to_file
+from app.managers.chapter_manager import chapter_chain, update_chapter
 from app.managers.summary_manager import summary_chain
 
 
@@ -32,9 +32,9 @@ class ChapterOutput(BaseModel):
 
 @router.post("/change")
 def change_chapter(input: ChapterOutput):
-    save_chapter_to_file(input.chapter, chapter_num=input.chapter_num)
+    update_chapter(input.chapter, chapter_num=input.chapter_num)
     summary_chain(input.chapter, chapter_num=input.chapter_num)
-    return True
+    return {"message": "Success"}
     
 @router.post("/get_all")
 def get_all_chapters():
@@ -45,11 +45,14 @@ def get_all_chapters():
                 content = f.read()
                 all_chapters.append({"filename": filename, "content": content})
     all_chapters.sort(key=lambda x: int(x["filename"].split("_")[1].split(".")[0]))
-    return all_chapters
+    return {"all_chapters": all_chapters}
 
+class ChapterNumber(BaseModel):
+    chapter_num: int
+    
 @router.post("/get_one")
-def get_one_chapter(chapter_num: int):
-    filename = f"chapter_{chapter_num:03}.txt"
+def get_one_chapter(input: ChapterNumber):
+    filename = f"chapter_{input.chapter_num:03}.txt"
     filepath = os.path.join("data/samples/raws", filename)
     if os.path.exists(filepath):
         with open(filepath, "r", encoding="utf-8") as f:

app/app.py

@@ -0,0 +1 @@
+from app.main import app

app/front_end/gradio_ui.py

@@ -0,0 +1,197 @@
+import gradio as gr
+import requests
+
+API_BASE = "http://localhost:8000"
+
+def get_texts(lang):
+    if lang == "中文":
+        return {
+            "title": "小说生成系统",
+            "generate_tab": "生成章节",
+            "change_tab": "修改章节",
+            "get_all_tab": "查看所有章节",
+            "get_one_tab": "查阅单章",
+            "input_text": "请写下你想要写的新章节包含什么样的故事情节？",
+            "change_text": "请输入修改后的内容",
+            "chapter_num": "章节编号",
+            "generate_button": "生成",
+            "change_button": "修改",
+            "get_all_button": "查看全部",
+            "get_one_button": "查阅",
+            "output": "输出结果",
+            "lang_label": "选择界面语言",
+            "view_one_dropdown": "章节编号",
+            "view_one_output": "章节内容",
+            "view_one_submit": "提交修改",
+            "view_one_edit": "编辑",
+            "view_one_status": "修改状态",
+        }
+    else:
+        return {
+            "title": "AI Novel Generator",
+            "generate_tab": "Generate Chapter",
+            "change_tab": "Edit Chapter",
+            "get_all_tab": "View All",
+            "get_one_tab": "View One",
+            "input_text": "Please write down what kind of storyline you would like the new chapter you are writing to contain?",
+            "change_text": "Enter new content to replace",
+            "chapter_num": "Chapter Number",
+            "generate_button": "Generate",
+            "change_button": "Update",
+            "get_all_button": "Get All",
+            "get_one_button": "Get One",
+            "output": "Output",
+            "lang_label": "Select UI language",
+            "view_one_dropdown": "Chapter Number",
+            "view_one_output": "Chapter Content",
+            "view_one_submit": "Submit Change",
+            "view_one_edit": "Edit chapter",
+            "view_one_status": "Change Status",
+        }
+
+
+
+# ==== API call functions ====
+
+def generate(text):
+    response = requests.post(f"{API_BASE}/generate", json={"query": text})
+    return response.json().get("chapter", "")
+
+def change(chapter_num, new_content):
+    response = requests.post(f"{API_BASE}/change", json={
+        "chapter": new_content,
+        "chapter_num": chapter_num       
+    })
+    return response.json().get("message", "Success")
+
+def get_all():
+    response = requests.post(f"{API_BASE}/get_all")
+    all_chapters = response.json().get("all_chapters", [])
+    chapter_count = len(all_chapters)
+    options = [f"{i+1}" for i in range(chapter_count)]
+    return options, all_chapters
+
+def get_one(chapter_num):
+    response = requests.post(f"{API_BASE}/get_one", json={"chapter_num": chapter_num})
+    return response.json().get("content",  response.json().get("error", "Chapter not found."))
+
+def get_one_cached(chapter_num, cached_chapters):
+    try:
+        print("load_from_cache")
+        return cached_chapters[int(chapter_num) - 1]["content"]
+    except (IndexError, ValueError, TypeError):
+        print("load_from_backend")
+        return get_one(chapter_num)
+
+def on_generate(prompt):
+    chapter = generate(prompt)
+    options, chapters = get_all()
+    return chapter, gr.update(choices=options, value=options[-1]), chapters
+
+def on_load():
+    options, chapters = get_all()
+    return gr.update(choices=options, value=options[-1]), chapters
+
+def on_select_chapter(chapter_num, cached):
+    chapter = get_one_cached(chapter_num, cached)
+    return chapter, gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
+
+def on_edit():
+    return gr.update(interactive=True), gr.update(visible=True), gr.update(visible=False)
+
+def on_submit_edit(chapter_num, new_content):
+    msg = change(chapter_num, new_content)
+    return gr.update(interactive=False), gr.update(visible=False), gr.update(visible=True), msg
+
+# ==== Gradio UI ====
+def build_ui():
+    with gr.Blocks() as demo:
+        
+        def update_labels(lang):
+            txt = get_texts(lang)
+            return (
+                gr.update(value=txt["title"]),
+                gr.update(label=txt["input_text"]), gr.update(label=txt["output"]), 
+                gr.update(value=txt["generate_button"]),
+                gr.update(label=txt["lang_label"]),
+                gr.update(label=txt["generate_tab"]),
+                gr.update(label=txt["view_one_dropdown"]),
+                gr.update(label=txt["view_one_output"]),
+                gr.update(value=txt["view_one_submit"]),
+                gr.update(value=txt["view_one_edit"]),
+                gr.update(label=txt["view_one_status"]),
+                gr.update(label=txt["get_one_tab"]), 
+
+            )
+
+        txt = get_texts("English")
+
+        title = gr.Markdown(f"# {txt['title']}")
+        lang_dropdown = gr.Dropdown(choices=["English", "中文"], value="English", label=txt["lang_label"])
+
+        cached_chapters = gr.State([])
+
+        with gr.Tab(txt["generate_tab"]) as generate_tab:
+            input_box = gr.Textbox(label=txt["input_text"])
+            output_box = gr.Textbox(label=txt["output"])
+            generate_btn = gr.Button(txt["generate_button"])
+            # generate_btn.click(fn=generate, inputs=input_box, outputs=output_box)
+            
+
+        # with gr.Tab(txt["change_tab"]) as change_tab:
+        #     chapter_num = gr.Number(label=txt["chapter_num"])
+        #     new_content = gr.Textbox(label=txt["change_text"])
+        #     result = gr.Textbox(label=txt["output"])
+        #     change_btn = gr.Button(txt["change_button"])
+        #     change_btn.click(fn=change, inputs=[chapter_num, new_content], outputs=result)
+
+        # with gr.Tab(txt["get_all_tab"]) as get_all_tab:
+        #     result_all = gr.Textbox(label=txt["output"], lines=20)
+        #     all_btn = gr.Button(txt["get_all_button"])
+        #     all_btn.click(fn=get_all, outputs=result_all)
+
+        # with gr.Tab(txt["get_one_tab"]) as get_one_tab:
+        #     one_chapter_num = gr.Number(label=txt["chapter_num"])
+        #     one_result = gr.Textbox(label=txt["output"], lines=10)
+        #     one_btn = gr.Button(txt["get_one_button"])
+        #     one_btn.click(fn=get_one, inputs=one_chapter_num, outputs=one_result)
+
+        with gr.Tab("View One") as view_one_tab:
+            view_one_dropdown = gr.Dropdown(label=txt["view_one_dropdown"], choices=[], visible=True)
+            view_one_output = gr.Textbox(label=txt["view_one_output"], lines=10, interactive=False)
+            view_one_submit = gr.Button(txt["view_one_submit"], visible=False)
+            view_one_edit = gr.Button(txt["view_one_edit"], visible=True)
+            view_one_status = gr.Textbox(label=txt["view_one_status"], visible=False)
+
+            view_one_dropdown.change(on_select_chapter, inputs=[view_one_dropdown, cached_chapters], outputs=[view_one_output, view_one_submit, view_one_edit, view_one_status])
+            view_one_edit.click(on_edit, outputs=[view_one_output, view_one_submit, view_one_edit])
+            view_one_submit.click(on_submit_edit, inputs=[view_one_dropdown, view_one_output], outputs=[view_one_output, view_one_submit, view_one_edit, view_one_status])
+
+        generate_btn.click(fn=on_generate, inputs=input_box,
+                        outputs=[output_box, view_one_dropdown, cached_chapters])
+        
+        demo.load(lambda: on_load(), outputs=[view_one_dropdown, cached_chapters])
+
+        lang_dropdown.change(
+            fn=update_labels,
+            inputs=lang_dropdown,
+            outputs=[
+                title,
+                input_box, output_box, generate_btn,
+                lang_dropdown,
+                generate_tab,
+                view_one_dropdown,
+                view_one_output,
+                view_one_submit,
+                view_one_edit,
+                view_one_status,
+                view_one_tab,
+            ]
+        )
+
+    # demo.launch()
+    return demo
+
+if __name__ == "__main__":
+    ui = build_ui()
+    ui.launch()

app/main.py

@@ -1,13 +1,35 @@
+# app/main.py
 from fastapi import FastAPI
-# from app.Archived import test
+from fastapi.middleware.cors import CORSMiddleware
+from starlette.responses import RedirectResponse
+
 from app.apis import generator, extractor
+from app.front_end.gradio_ui import build_ui
+
+import gradio as gr
+import uvicorn
 
 app = FastAPI(title="AI Novelist RAG")
 
+# 允许 CORS（可选）
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# 后端 API
 app.include_router(generator.router)
 app.include_router(extractor.router)
 
+# 前端（Gradio UI）
+demo = build_ui()
+app = gr.mount_gradio_app(app, demo, path="/ui")  # 访问 /ui 查看 Gradio 页面
 
 @app.get("/")
-def read_root():
-    return {"message": "AI Novelist RAG service is running."}
+def root():
+    return RedirectResponse(url="/ui")  # 访问根页面自动跳转到 Gradio UI
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)

app/managers/chapter_manager.py

@@ -52,13 +52,15 @@ def add_chapter(chapter: str, chapter_num: int):
         content=chapter,
         metadata={"type": "chapter", "chapter": chapter_num}
     )
+    save_chapter_to_file(chapter, chapter_num)
 
 def update_chapter(chapter: str, chapter_num: int):
     vm.update_document(
         store_type="chapter",
-        chapter_id=chapter_num,
+        chapter_num=chapter_num,
         new_content=chapter
     )
+    save_chapter_to_file(chapter, chapter_num)
 
 
     
@@ -68,7 +70,6 @@ def chapter_chain(query):
     prompt = setup_prompt(query, context_info)
     chapter = generate_chapter(prompt)
     new_chapter_num = context_info[0] + 1
-    save_chapter_to_file(chapter, new_chapter_num)
     add_chapter(chapter, new_chapter_num)
     return chapter
 

app/managers/summary_manager.py

@@ -25,15 +25,17 @@ def add_summary(summary: str, chapter_num: int):
         content=summary,
         metadata={"type": "summary", "chapter": chapter_num}
     )
+    save_summary_to_file(summary, chapter_num)
 
 
 def update_summary(summary: str, chapter_num: int):
     
     vm.update_document(
         store_type="summary",
-        chapter_id=chapter_num,
+        chapter_num=chapter_num,
         new_content=summary
     )
+    save_summary_to_file(summary, chapter_num)
 
 
 def get_relevant_summaries(query: str, top_k: int = 10):
@@ -46,16 +48,11 @@ def get_relevant_summaries(query: str, top_k: int = 10):
 
 def summary_chain(text: str, chapter_num: int = None) -> str:
     new_chapter_num = get_latest_chapter_num() + 1
-    print("I'm here 1")
     if chapter_num is None or chapter_num >= new_chapter_num:
         chapter_num = new_chapter_num
-        print("I'm here 2")
         summary = generate_summary(text)
-        print("I'm here 3")
-        save_summary_to_file(summary, chapter_num)
         add_summary(summary, chapter_num)
     else:
         summary = generate_summary(text)
-        save_summary_to_file(summary, chapter_num) # cover old file
         update_summary(summary, chapter_num)
     return summary

app/managers/vector_manager.py

@@ -53,26 +53,32 @@ def add_document(store_type: str, content: str, metadata: dict):
     save_vectorstore(vs, store_type)
 
 
-def update_document(store_type: str, chapter_id: str, new_content: str):
+def update_document(store_type: str, chapter_num: str, new_content: str):
+    print(f"Updating chapter {chapter_num} in store: {store_type}")
+    
     vs = load_vectorstore(store_type)
 
-    doc_ids = [
-        k for k, doc in vs.docstore._dict.items()
-        if doc.metadata.get("chapter") == chapter_id
+    target_doc_ids = [
+        doc_id for doc_id, doc in vs.docstore._dict.items()
+        if doc.metadata.get("chapter") == chapter_num
     ]
-    if not doc_ids:
-        raise ValueError(f"No document found for chapter: {chapter_id}")
+    
+    if not target_doc_ids:
+        raise ValueError(f"No document found for chapter: {chapter_num}")
 
-    docs = [
-        doc for k, doc in vs.docstore._dict.items()
-        if doc.metadata.get("chapter") != chapter_id
-    ]
-    new_doc = Document(page_content=new_content, metadata={"type": store_type, "chapter": chapter_id})
-    docs.append(new_doc)
+    for doc_id in target_doc_ids:
+        vs.delete([doc_id])
+
+    new_doc = Document(
+        page_content=new_content,
+        metadata={"type": store_type, "chapter": chapter_num}
+    )
+    vs.add_documents([new_doc])
+
+    save_vectorstore(vs, store_type)
+
+    print(f"Chapter {chapter_num} successfully updated.")
 
-    new_vs = create_new_vectorstore(Embedding_model)
-    new_vs.add_documents(docs)
-    save_vectorstore(new_vs, store_type)
 
 
 def get_relevant_documents(store_type: str, query: str, top_k: int = 10) -> List[Document]:

app/tests/langchain_test.ipynb

@@ -2,20 +2,515 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "16ea7d3d",
    "metadata": {},
    "outputs": [],
    "source": [
-    "from langchain.chains import RetrievalQA\n",
-    "from langchain.embeddings import OpenAIEmbeddings\n",
-    "from langchain.vectorstores import pinecone"
+    "import os\n",
+    "from langchain.vectorstores import FAISS\n",
+    "from langchain.docstore import InMemoryDocstore\n",
+    "from langchain_huggingface import HuggingFaceEmbeddings\n",
+    "from langchain.schema import Document\n",
+    "import faiss\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "31e17418",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "os.path.exists(\"conda_env_test.ipynb\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "cb6dccc3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "os.chdir('..\\..')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "f15a094d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "c:\\Users\\nanfangwuyu\\.conda\\envs\\novel_rag\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n",
+      "Device set to use cuda:0\n"
+     ]
+    }
+   ],
+   "source": [
+    "from app.models.bart_large_cnn import BartSummaryModel\n",
+    "Summary_Model = BartSummaryModel()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "d5b232b0",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'Jack is a young mage with a spirit as untamed as the tempest itself. He has a rare affinity for all five elements that set him apart from others. Despite his burgeoning power, he felt the weight of an uncertain destiny resting heavily on his shoulders. Little did he know, the world was on the brink of change.'"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "text = \"\"\"\n",
+    "**Chapter 1: Embers of Destiny**\n",
+    "\n",
+    "In the realm of Eldoria, where the five elements danced in a delicate balance, the air crackled with a potential that whispered of ancient magics. Fire, fierce and unyielding, flickered in the heart of the volcanoes; water flowed with grace through the rivers and oceans, nurturing life in its embrace; wind roamed free, a playful spirit that caressed the mountains; electricity surged in the stormy skies, a wild tempest of power; and earth, the steadfast guardian, cradled the very foundations of existence. Here, in this vibrant tapestry woven by elemental forces, the fate of kingdoms hung by a thread.\n",
+    "\n",
+    "Amidst the verdant valleys of Eldoria stood Jack, a young mage with a spirit as untamed as the tempest itself. He was a figure of quiet strength, with tousled chestnut hair that caught the light of the setting sun and eyes as deep as the ocean’s depths—reflecting a soul that yearned for adventure. Jack had long been aware of his gifts, a rare affinity for all five elements that set him apart from others. Yet, despite his burgeoning power, he felt the weight of an uncertain destiny resting heavily on his shoulders.\n",
+    "\n",
+    "As he stood atop the cliffs of Varyn, the wind tousling his hair, Jack gazed out at the swirling mists of the horizon. The clouds, tinged with the colors of twilight, whispered promises of journeys yet to be taken. Today, a flicker of intuition ignited within him, a stirring that foretold an awakening—a call to harness the elemental forces that lay dormant, waiting for the right moment to erupt. With a heart full of determination and a mind sharpened by dreams, Jack took a deep breath, ready to carve his path through the ever-shifting tides of fate. Little did he know, the world was on the brink of change, and he was destined to be its catalyst.\n",
+    "\"\"\"\n",
+    "Summary_Model.summarize(text)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "864c3438",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "c:\\Users\\nanfangwuyu\\.conda\\envs\\novel_rag\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n",
+      "Device set to use cuda:0\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Reload existing faiss\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "0"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "\n",
+    "from app.managers import vector_manager as vm\n",
+    "import os\n",
+    "\n",
+    "SAVE_DIR = \"data/samples/raws\"\n",
+    "os.makedirs(SAVE_DIR, exist_ok=True)\n",
+    "\n",
+    "\n",
+    "\n",
+    "def get_latest_chapter_num(store_type=\"summary\"):\n",
+    "    vectorstore = vm.load_vectorstore(store_type)\n",
+    "    all_docs = vectorstore.docstore._dict.values()\n",
+    "    chapter_numbers = [doc.metadata.get(\"chapter\", 0) for doc in all_docs if isinstance(doc.metadata.get(\"chapter\", 0), int)]\n",
+    "    latest_chapter_num = max(chapter_numbers) if chapter_numbers else 0\n",
+    "    return latest_chapter_num\n",
+    "\n",
+    "get_latest_chapter_num()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "2421634b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def create_new_vectorstore(embedding_model):\n",
+    "    index = faiss.IndexFlatL2(len(embedding_model.embed_query(\"hello world\")))\n",
+    "    vectorstore = vector_store = FAISS(\n",
+    "        embedding_function=embedding_model,\n",
+    "        index=index,\n",
+    "        docstore=InMemoryDocstore(),\n",
+    "        index_to_docstore_id={},\n",
+    "    )\n",
+    "    return vectorstore"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "e42f4e7b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "True\n"
+     ]
+    }
+   ],
+   "source": [
+    "VECTORSTORE_PATH = \"data/vectorstore/summary_index\"\n",
+    "embedding_model = HuggingFaceEmbeddings(model_name=\"sentence-transformers/all-MiniLM-L6-v2\")\n",
+    "\n",
+    "# Step 0: 加载或初始化向量库\n",
+    "if os.path.exists(VECTORSTORE_PATH):\n",
+    "    print(True)\n",
+    "    vectorstore = FAISS.load_local(VECTORSTORE_PATH, embedding_model, allow_dangerous_deserialization=True)\n",
+    "else:\n",
+    "    vectorstore = create_new_vectorstore(embedding_model)\n",
+    "    "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "3ccd8712",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{}"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "vectorstore = create_new_vectorstore(embedding_model)\n",
+    "vectorstore.docstore._dict"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "id": "5f179882",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Step 1: 获取所有文档 & 最新章节编号\n",
+    "all_docs = vectorstore.docstore._dict.values()\n",
+    "chapter_numbers = [doc.metadata.get(\"chapter\", 0) for doc in all_docs if isinstance(doc.metadata.get(\"chapter\", 0), int)]\n",
+    "latest_chapter = max(chapter_numbers) if chapter_numbers else 0\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "id": "c219e6ef",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(dict_values([Document(id='e948cd24-4678-436f-b98f-9514db311f42', metadata={'chapter': 1}, page_content='In the magical realm of Eldoria, where fire, water, wind, electricity, and earth govern existence, Jack, a unique mage cloaked in indigo, possesses the rare ability to master all five elements. Revered and envied by townsfolk, he grapples with the loneliness of his gifts. As he stands at the edge of the Whispering Woods one evening, feeling the elemental energies around him, he senses a change in the air. Unbeknownst to him, the universe is weaving a fate that will challenge the core of his powers and alter his destiny.'), Document(id='27fff7cd-20be-4b44-be42-50af318b6e80', metadata={'chapter': 2}, page_content='In the twilight of Eldoria, Jack stands at the edge of the Whispering Woods, feeling a deep connection to the ancient trees. Drawn by an electric anticipation, he ventures into the forest, where reality blurs and whispers grow louder. Amidst the scents of earth and nightshade, he senses a change marked by a foreign smoke. Following a flickering glow, he discovers a clearing with a stone pedestal radiating energy. Compelled by an inner voice, Jack realizes that his destiny is intertwined with the secrets of this sacred place.')]),\n",
+       " 2)"
+      ]
+     },
+     "execution_count": 29,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "all_docs, latest_chapter"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "id": "3319424c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Step 2: 获取最新章节 summary\n",
+    "latest_summary_docs = [doc for doc in all_docs if doc.metadata.get(\"chapter\") == latest_chapter]\n",
+    "latest_summary_doc = latest_summary_docs[0] if latest_summary_docs else \"\"\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "id": "ea1a62ac",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Step 3: 用户主题输入\n",
+    "if latest_chapter == 0:\n",
+    "    user_query = \"Help me write a fantasy novel. Introduces the worldview and the male protagonist. The world has five elements: fire, water, wind, electricity, and earth. The male protagonist names Jack, is a powerful mage.\" \n",
+    "else:\n",
+    "    user_query = \"Continue writing.\"\n",
+    "\n",
+    "# 检索相关章节概要\n",
+    "retriever = vectorstore.as_retriever(search_type=\"similarity\", search_kwargs={\"k\": 3})\n",
+    "related_docs = retriever.get_relevant_documents(user_query)\n",
+    "# related_chapter_numbers = [doc.metadata.get(\"chapter\", 0) for doc in related_docs if isinstance(doc.metadata.get(\"chapter\", 0), int)]\n",
+    "if latest_summary_doc in related_docs:\n",
+    "    related_docs.remove(latest_summary_doc)\n",
+    "related_summaries = \"\\n\\n\".join([\"Summary of Chapter {}:\\n\".format(doc.metadata.get(\"chapter\")) + doc.page_content for doc in related_docs])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "id": "b39d9d19",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "([Document(id='e948cd24-4678-436f-b98f-9514db311f42', metadata={'chapter': 1}, page_content='In the magical realm of Eldoria, where fire, water, wind, electricity, and earth govern existence, Jack, a unique mage cloaked in indigo, possesses the rare ability to master all five elements. Revered and envied by townsfolk, he grapples with the loneliness of his gifts. As he stands at the edge of the Whispering Woods one evening, feeling the elemental energies around him, he senses a change in the air. Unbeknownst to him, the universe is weaving a fate that will challenge the core of his powers and alter his destiny.')],\n",
+       " 'Summary of Chapter 1:\\nIn the magical realm of Eldoria, where fire, water, wind, electricity, and earth govern existence, Jack, a unique mage cloaked in indigo, possesses the rare ability to master all five elements. Revered and envied by townsfolk, he grapples with the loneliness of his gifts. As he stands at the edge of the Whispering Woods one evening, feeling the elemental energies around him, he senses a change in the air. Unbeknownst to him, the universe is weaving a fate that will challenge the core of his powers and alter his destiny.')"
+      ]
+     },
+     "execution_count": 32,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "related_docs, related_summaries"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "id": "f37f3f8d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "You are a novelist working on a long novel.\n",
+      "\n",
+      "Here is a summary of the previously on few related chapters:\n",
+      "\n",
+      "\"Summary of Chapter 1:\n",
+      "In the magical realm of Eldoria, where fire, water, wind, electricity, and earth govern existence, Jack, a unique mage cloaked in indigo, possesses the rare ability to master all five elements. Revered and envied by townsfolk, he grapples with the loneliness of his gifts. As he stands at the edge of the Whispering Woods one evening, feeling the elemental energies around him, he senses a change in the air. Unbeknownst to him, the universe is weaving a fate that will challenge the core of his powers and alter his destiny.\"\n",
+      "\n",
+      "And here is a summary of the previously on the last chapter (Chapter 2): \n",
+      "\n",
+      "\"In the twilight of Eldoria, Jack stands at the edge of the Whispering Woods, feeling a deep connection to the ancient trees. Drawn by an electric anticipation, he ventures into the forest, where reality blurs and whispers grow louder. Amidst the scents of earth and nightshade, he senses a change marked by a foreign smoke. Following a flickering glow, he discovers a clearing with a stone pedestal radiating energy. Compelled by an inner voice, Jack realizes that his destiny is intertwined with the secrets of this sacred place.\"\n",
+      "\n",
+      "Please continue to WRITE the CONTENT of the NEXT CHAPTER - chapter 3, keeping it logically coherent, within about 300 words. \n",
+      "\n",
+      "Here is what the user expects to see in the next chapter: \"Continue writing.\"\n",
+      "\n",
+      "    \n"
+     ]
+    }
+   ],
+   "source": [
+    "# Step 4: 构建 prompt 生成内容\n",
+    "if latest_chapter > 1:\n",
+    "\n",
+    "    prompt = f\"\"\"You are a novelist working on a long novel.\n",
+    "\n",
+    "Here is a summary of the previously on few related chapters:\n",
+    "\n",
+    "\"{related_summaries}\"\n",
+    "\n",
+    "And here is a summary of the previously on the last chapter (Chapter {latest_summary_doc.metadata.get(\"chapter\")}): \n",
+    "\n",
+    "\"{latest_summary_doc.page_content}\"\n",
+    "\n",
+    "Please continue to WRITE the CONTENT of the NEXT CHAPTER - chapter {latest_chapter + 1}, keeping it logically coherent, within about 300 words. \n",
+    "\n",
+    "Here is what the user expects to see in the next chapter: \"{user_query}\"\n",
+    "\n",
+    "    \"\"\"\n",
+    "elif latest_chapter == 1:\n",
+    "    prompt = f\"\"\"You are a novelist working on a long novel.\n",
+    "\n",
+    "Here is a summary of the previously on the first chapter (Chapter {latest_summary_doc.metadata.get(\"chapter\")}): \n",
+    "\n",
+    "\"{latest_summary_doc.page_content}\"\n",
+    "\n",
+    "Please continue to WRITE the CONTENT of the NEXT CHAPTER - chapter {latest_chapter + 1}, keeping it logically coherent, within about 300 words. \n",
+    "\n",
+    "Here is what the user expects to see in the next chapter: \"{user_query}\"\n",
+    "\"\"\"\n",
+    "else:\n",
+    "    prompt = f\"\"\"You are a novelist working on a long novel.\n",
+    "\n",
+    "Please WRITE the content of the chapter 1 within about 300 words. \n",
+    "\n",
+    "Here is what the user expects to see in the chapter 1: \"{user_query}\"\n",
+    "\n",
+    "    \"\"\"\n",
+    "\n",
+    "#   # 你封装的模型类，含 generate & summarize 方法\n",
+    "\n",
+    "print(prompt)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "93d0b566",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from app.models.model import LLM\n",
+    "generated_chapter = LLM.generate(prompt, max_tokens=512)\n",
+    "generated_summary = LLM.summarize(generated_chapter)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 35,
+   "id": "1b51a228",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Generated Chapter:\n",
+      "\n",
+      "**Chapter 3: The Awakening of Shadows**\n",
+      "\n",
+      "Jack stepped into the clearing, the air crackling with an otherworldly energy that hummed beneath his skin. The stone pedestal, a relic of ages past, loomed before him, its surface etched with runes that glowed faintly in the twilight. The whispers of the forest grew louder, transforming into a symphony of voices, each one tugging at the corners of his mind, urging him to unlock the secrets hidden within the stone.\n",
+      "\n",
+      "As he approached, the flickering glow intensified, casting dancing shadows across the trees. Jack could feel the weight of history pressing upon him, as if the very essence of the forest had gathered to witness this moment. A sense of purpose surged within him, igniting a flicker of hope amidst the loneliness he had often felt. Was this the connection he had long sought—the intertwining of his fate with something greater than himself?\n",
+      "\n",
+      "Tentatively, he reached out, his fingers brushing against the cool stone. Instantly, a jolt of energy coursed through him, igniting every element he commanded. Fire sparked to life in his fingertips, water swirled around his ankles, earth pulsed beneath his feet, wind whipped through his hair, and electricity crackled in the air. In that moment, Jack felt not like an outsider but rather a conduit of the very forces that shaped Eldoria.\n",
+      "\n",
+      "The runes flared brighter, revealing a vision—a shadowy figure cloaked in darkness, standing at the precipice of chaos. Jack’s heart raced as he realized that this figure held the key to his destiny. The whispers coalesced into a single voice, resonating deep within him: \"You must choose, Jack. Embrace the light or succumb to the shadows.\"\n",
+      "\n",
+      "With his heart pounding, Jack stood at the crossroads of fate, poised to confront the challenges that lay ahead. The echoes of the forest faded into silence, leaving him alone with the weight of his choice. \n",
+      "\n",
+      "Generated Summary:\n",
+      "\n",
+      "In Chapter 3, \"The Awakening of Shadows,\" Jack enters a mystical clearing in the forest of Eldoria, where he discovers a glowing stone pedestal inscribed with ancient runes. As he approaches, he feels a surge of energy and connection to the elemental forces around him. The whispers of the forest reveal a shadowy figure representing chaos, urging Jack to choose between light and darkness. Empowered and filled with purpose, Jack stands at a pivotal moment in his destiny, ready to confront the challenges that await him.\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(f\"Generated Chapter:\\n\\n{generated_chapter} \\n\\nGenerated Summary:\\n\\n{generated_summary}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "16893feb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# generated_summary = \"Jack and Jill went up the hill to fetch a pail of water. Jack fell down and broke his crown, and Jill came tumbling after.\"  # 👈 替换成模型生成的内容\n",
+    "# generated_summary = \"Lucy and her brother went to the market to buy some groceries. They met a friendly dog on the way.\"  # 👈 替换成模型生成的内容\n",
+    "# generated_summary = \"The sun was shining brightly as they walked home. They had a great time at the market.\"  # 👈 替换成模型生成的内容"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "id": "90e8904e",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['a855b3b5-75ab-4e3b-88c7-81763ee7fd7f']"
+      ]
+     },
+     "execution_count": 36,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Step 5: 构建并添加文档\n",
+    "latest_chapter += 1\n",
+    "new_doc = Document(\n",
+    "    page_content=generated_summary,\n",
+    "    metadata={\"chapter\": latest_chapter}\n",
+    ")\n",
+    "vectorstore.add_documents([new_doc])\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 37,
+   "id": "7260c9ef",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'e948cd24-4678-436f-b98f-9514db311f42': Document(id='e948cd24-4678-436f-b98f-9514db311f42', metadata={'chapter': 1}, page_content='In the magical realm of Eldoria, where fire, water, wind, electricity, and earth govern existence, Jack, a unique mage cloaked in indigo, possesses the rare ability to master all five elements. Revered and envied by townsfolk, he grapples with the loneliness of his gifts. As he stands at the edge of the Whispering Woods one evening, feeling the elemental energies around him, he senses a change in the air. Unbeknownst to him, the universe is weaving a fate that will challenge the core of his powers and alter his destiny.'),\n",
+       " '27fff7cd-20be-4b44-be42-50af318b6e80': Document(id='27fff7cd-20be-4b44-be42-50af318b6e80', metadata={'chapter': 2}, page_content='In the twilight of Eldoria, Jack stands at the edge of the Whispering Woods, feeling a deep connection to the ancient trees. Drawn by an electric anticipation, he ventures into the forest, where reality blurs and whispers grow louder. Amidst the scents of earth and nightshade, he senses a change marked by a foreign smoke. Following a flickering glow, he discovers a clearing with a stone pedestal radiating energy. Compelled by an inner voice, Jack realizes that his destiny is intertwined with the secrets of this sacred place.'),\n",
+       " 'a855b3b5-75ab-4e3b-88c7-81763ee7fd7f': Document(id='a855b3b5-75ab-4e3b-88c7-81763ee7fd7f', metadata={'chapter': 3}, page_content='In Chapter 3, \"The Awakening of Shadows,\" Jack enters a mystical clearing in the forest of Eldoria, where he discovers a glowing stone pedestal inscribed with ancient runes. As he approaches, he feels a surge of energy and connection to the elemental forces around him. The whispers of the forest reveal a shadowy figure representing chaos, urging Jack to choose between light and darkness. Empowered and filled with purpose, Jack stands at a pivotal moment in his destiny, ready to confront the challenges that await him.')}"
+      ]
+     },
+     "execution_count": 37,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "vectorstore.docstore._dict"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "cd92e6e8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Step 6: 保存向量库\n",
+    "vectorstore.save_local(VECTORSTORE_PATH)\n",
+    "\n",
+    "# Step 7: 保存章节内容和总结到本地（你已有逻辑）\n",
+    "# save_chapter_text(latest_chapter + 1, generated_chapter)\n",
+    "# save_summary_text(latest_chapter + 1, generated_summary)\n"
    ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "novel_rag",
    "language": "python",
    "name": "python3"
   },

app/utils/utils.py

@@ -1,11 +1 @@
 
-
-
-# def summarize_a_chapter(text: str) -> str:
-#     """
-#     Summarize a chapter using the model.
-#     """
-#     # Assuming `model` is a pre-defined object that can summarize text
-#     summary = model.summarize(text)
-#     return summary
-