Updates libraries impporting (especially langchain), add more apis, complete test for generation

.gitignore

@@ -6,3 +6,6 @@ data/
 docker/
 k8s/
 scripts/
+app/Archived
+Archived/
+.env

app/apis/generator.py

@@ -21,8 +21,10 @@ class ChapterInput(BaseModel):
 @router.post("/generate")
 def generate_chapter(input: ChapterInput):
     result = chapter_chain(query=input.query)
-    summary_chain(result) # auto-call summary_chain
-    return result
+    print("chapter generated")
+    summary = summary_chain(result) # auto-call summary_chain
+    print("summary generated")
+    return {"chapter": result, "summary": summary}
 
 class ChapterOutput(BaseModel):
     chapter: str
@@ -34,6 +36,28 @@ def change_chapter(input: ChapterOutput):
     summary_chain(input.chapter, chapter_num=input.chapter_num)
     return True
     
+@router.post("/get_all")
+def get_all_chapters():
+    all_chapters = []
+    for filename in os.listdir("data/samples/raws"):
+        if filename.endswith(".txt"):
+            with open(os.path.join("data/samples/raws", filename), "r", encoding="utf-8") as f:
+                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
+
+@router.post("/get_one")
+def get_one_chapter(chapter_num: int):
+    filename = f"chapter_{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:
+            content = f.read()
+        return {"filename": filename, "content": content}
+    else:
+        return {"error": "Chapter not found."}
+
 
 # @router.post("/generate")
 # def generate_text(request: GenerateRequest):

app/main.py

@@ -5,7 +5,6 @@ from app.apis import generator, extractor
 app = FastAPI(title="AI Novelist RAG")
 
 app.include_router(generator.router)
-app.include_router(test.router)
 app.include_router(extractor.router)
 
 

app/managers/chapter_manager.py

@@ -1,26 +1,15 @@
-from langchain.chains import RetrievalQA
-from langchain.chains.llm import LLMChain
-from langchain.prompts import PromptTemplate
-from langchain.llms import OpenAI
-# from app.utils.memory import faiss_index
-from langchain.vectorstores import FAISS
-import numpy as np
 from app.models.model import LLM
 from app.utils.prompts import get_chapter_generation_prompt
 from app.managers import vector_manager as vm
 import os
-from langchain.docstore import InMemoryDocstore
-from langchain_huggingface import HuggingFaceEmbeddings
-from langchain.schema import Document
-import faiss
 
 SAVE_DIR = "data/samples/raws"
 os.makedirs(SAVE_DIR, exist_ok=True)
 
 
 
-def get_latest_chapter_num():
-    vectorstore = vm.load_vectorstore("summary")
+def get_latest_chapter_num(store_type="summary"):
+    vectorstore = vm.load_vectorstore(store_type)
     all_docs = vectorstore.docstore._dict.values()
     chapter_numbers = [doc.metadata.get("chapter", 0) for doc in all_docs if isinstance(doc.metadata.get("chapter", 0), int)]
     latest_chapter_num = max(chapter_numbers) if chapter_numbers else 0
@@ -48,7 +37,7 @@ def setup_prompt(query: str, context_info):
 
 
 def generate_chapter(prompt):
-    generated_chapter = LLM.generate(prompt, max_tokens=512)
+    generated_chapter = LLM.generate(prompt, max_tokens=2048)
     return generated_chapter
 
 def save_chapter_to_file(chapter: str, chapter_num: int):
@@ -78,8 +67,8 @@ def chapter_chain(query):
     context_info = get_latest_and_relevant_chapter_summaries(query)
     prompt = setup_prompt(query, context_info)
     chapter = generate_chapter(prompt)
-    chapter_num = context_info[0] + 1
-    save_chapter_to_file(chapter, chapter_num)
-    add_chapter(chapter, chapter_num)
-    return True
+    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

@@ -1,5 +1,4 @@
 import os
-from langchain.schema import Document
 from app.models.model import Summary_Model
 from app.managers import vector_manager as vm
 from app.managers.chapter_manager import get_latest_chapter_num, chapter_chain
@@ -47,9 +46,12 @@ 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:

app/managers/vector_manager.py

@@ -1,8 +1,9 @@
 import os
 import faiss
 from typing import List, Optional
-from langchain.vectorstores import FAISS
-from langchain.docstore import InMemoryDocstore
+from langchain_community.vectorstores import FAISS
+# from langchain.docstore import InMemoryDocstore
+from langchain_community.docstore.in_memory import InMemoryDocstore
 from langchain.schema import Document
 from app.models.model import Embedding_model
 
@@ -29,11 +30,13 @@ def load_vectorstore(store_type: str) -> FAISS:
     assert store_type in VECTORSTORE_TYPES, "Invalid vectorstore type."
     path = os.path.join(BASE_PATH, VECTORSTORE_TYPES[store_type])
 
-    if os.path.exists(path):
+    if os.path.exists(os.path.join(path, 'index.faiss')):
+        print("Reload existing faiss")
         return FAISS.load_local(path, Embedding_model, allow_dangerous_deserialization=True)
     else:
+        print("Create new faiss")
         vs = create_new_vectorstore(Embedding_model)
-        vs.save_local(path)
+        save_vectorstore(vs, store_type)
         return vs
 
 

app/models/model.py

@@ -1,4 +1,4 @@
-from app.models.tinyllama import TinyLlamaModel
+# from app.models.tinyllama import TinyLlamaModel
 from app.models.gpt4omini import GPT4OMini
 from app.models.bart_large_cnn import BartSummaryModel
 from langchain_huggingface import HuggingFaceEmbeddings