fastapi (#14)

* fastapi

* remove unused files

* merge concurrent code

* petits correctifs

---------

Co-authored-by: loiclu7 <loic.lu7@gmail.com>
Co-authored-by: Romain Vanhee <romain.vanhee@yrycom.com>

Author: 3 people

Dockerfile.api

@@ -8,7 +8,7 @@ ENV PYTHONUNBUFFERED=1
 COPY requirements.txt /app/requirements.txt
 RUN pip install --no-cache-dir -r /app/requirements.txt
 
-COPY app /app/app 
+COPY app /app/app
 COPY vector-store /app/vector-store
 
 EXPOSE 8000

RAG.py

@@ -0,0 +1,121 @@
+import os
+import time
+import uuid
+from pathlib import Path
+from typing import List, Optional
+from typing_extensions import TypedDict
+
+from dotenv import load_dotenv
+from fastapi import FastAPI
+from pydantic import BaseModel
+from tqdm import tqdm
+
+from langchain_openai import OpenAIEmbeddings, ChatOpenAI
+from langchain_chroma import Chroma
+from langchain_community.document_loaders import PyMuPDFLoader, Docx2txtLoader, TextLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.documents import Document
+from langgraph.graph import START, StateGraph
+
+load_dotenv()
+LLM_MODEL = os.getenv("MODEL_LLM")
+EMB_MODEL = os.getenv("MODEL_EMBEDDING")
+DOCS_DIR = 'exports'
+PERSIST_DIR = 'vector-store'
+
+class State(TypedDict):
+    question: str
+    context: List[Document]
+    answer: str
+
+class RAGCore:
+    def __init__(self):
+        self.embeddings = OpenAIEmbeddings(model=EMB_MODEL)
+        self.llm = ChatOpenAI(model=LLM_MODEL, temperature=0.2)
+        self.vector_store = Chroma(
+            collection_name="collection_test",
+            embedding_function=self.embeddings,
+            persist_directory=PERSIST_DIR,
+        )
+        self.graph = self._build_graph()
+
+    def _build_graph(self):
+        prompt = ChatPromptTemplate.from_template(
+            "You are an assistant for question-answering tasks. Use the following context to answer.\n"
+            "Context: {context}\nQuestion: {question}\nAnswer:"
+        )
+
+        def retrieve(state: State):
+            # Optimized search: fetch top 5 relevant chunks
+            docs = self.vector_store.similarity_search(state["question"], k=5)
+            return {"context": docs}
+
+        def generate(state: State):
+            docs_content = "\n\n".join(doc.page_content for doc in state["context"])
+            sources = list(set([os.path.basename(doc.metadata.get('source', 'Unknown')) for doc in state["context"]]))
+            messages = prompt.invoke({"question": state["question"], "context": docs_content})
+            response = self.llm.invoke(messages)
+            source_text = "\n\n**Sources:**\n- " + "\n- ".join(sources)
+            final_answer = response.content + source_text
+            return {"answer": final_answer}
+
+        workflow = StateGraph(State)
+        workflow.add_node("retrieve", retrieve)
+        workflow.add_node("generate", generate)
+        workflow.add_edge(START, "retrieve")
+        workflow.add_edge("retrieve", "generate")
+        return workflow.compile()
+
+    def query(self, text: str):
+        return self.graph.invoke({"question": text})
+
+
+app = FastAPI(title="AnythingLLM Custom Bridge")
+engine = RAGCore()
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class ChatCompletionRequest(BaseModel):
+    messages: List[ChatMessage]
+    model: Optional[str] = "custom-rag"
+
+@app.post("/v1/chat/completions")
+async def chat_endpoint(request: ChatCompletionRequest):
+    user_msg = request.messages[-1].content
+    result = engine.query(user_msg)
+    
+    return {
+        "id": f"chatcmpl-{uuid.uuid4()}",
+        "object": "chat.completion",
+        "created": int(time.time()),
+        "model": request.model,
+        "choices": [{
+            "message": {"role": "assistant", "content": result["answer"]},
+            "finish_reason": "stop",
+            "index": 0
+        }]
+    }
+
+def run_ingestion():
+    if not os.path.exists(DOCS_DIR): return
+    
+    paths = list(Path(DOCS_DIR).rglob("*"))
+    new_docs = []
+    for p in tqdm(paths, desc="Processing files"):
+        if p.suffix.lower() == ".pdf": new_docs.extend(PyMuPDFLoader(str(p)).load())
+        elif p.suffix.lower() == ".docx": new_docs.extend(Docx2txtLoader(str(p)).load())
+        elif p.suffix.lower() in {".txt", ".md"}: new_docs.extend(TextLoader(str(p)).load())
+    
+    if new_docs:
+        splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
+        splits = splitter.split_documents(new_docs)
+        engine.vector_store.add_documents(splits)
+        print(f"Successfully indexed {len(splits)} chunks.")
+
+if __name__ == "__main__":
+    import uvicorn
+    run_ingestion()
+    uvicorn.run(app, host="0.0.0.0", port=8000)

app/main.py

@@ -9,7 +9,10 @@
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 
+from app.rag import RAGCore
+
 BACKEND_MODEL_ID = os.getenv("BACKEND_MODEL_ID", "chatbot-rag")
+engine = RAGCore()
 
 
 class ChatMessage(BaseModel):
@@ -47,7 +50,8 @@ def chat_completions(request: ChatCompletionRequest):
     if not request.messages:
         raise HTTPException(status_code=400, detail="messages is required")
 
-    answer = f"Reponse bidon API."
+    user_msg = request.messages[-1].content
+    answer = engine.query(user_msg)
     created = int(time.time())
     completion_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
     model_id = request.model or BACKEND_MODEL_ID

app/rag.py

@@ -0,0 +1,63 @@
+import os
+from typing import List
+from typing_extensions import TypedDict
+
+from dotenv import load_dotenv
+from langchain_openai import OpenAIEmbeddings, ChatOpenAI
+from langchain_chroma import Chroma
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.documents import Document
+from langgraph.graph import START, StateGraph
+
+load_dotenv()
+LLM_MODEL = os.getenv("MODEL_LLM")
+EMB_MODEL = os.getenv("MODEL_EMBEDDING")
+PERSIST_DIR = os.getenv("VECTOR_STORE_DIR", "vector-store")
+COLLECTION_NAME = os.getenv("COLLECTION_NAME", "collection_test")
+
+
+class State(TypedDict):
+    question: str
+    context: List[Document]
+    answer: str
+
+
+class RAGCore:
+    def __init__(self):
+        self.embeddings = OpenAIEmbeddings(model=EMB_MODEL)
+        self.llm = ChatOpenAI(model=LLM_MODEL, temperature=0.2)
+        self.vector_store = Chroma(
+            collection_name=COLLECTION_NAME,
+            embedding_function=self.embeddings,
+            persist_directory=PERSIST_DIR,
+        )
+        self.graph = self._build_graph()
+
+    def _build_graph(self):
+        prompt = ChatPromptTemplate.from_template(
+            "You are an assistant for question-answering tasks. Use the following context to answer.\n"
+            "Context: {context}\nQuestion: {question}\nAnswer:"
+        )
+
+        def retrieve(state: State):
+            docs = self.vector_store.similarity_search(state["question"], k=5)
+            return {"context": docs}
+
+        def generate(state: State):
+            docs_content = "\n\n".join(doc.page_content for doc in state["context"])
+            sources = list(set([os.path.basename(doc.metadata.get("source", "Unknown")) for doc in state["context"]]))
+            messages = prompt.invoke({"question": state["question"], "context": docs_content})
+            response = self.llm.invoke(messages)
+            source_text = "\n\n**Sources:**\n- " + "\n- ".join(sources)
+            return {"answer": response.content + source_text}
+
+        workflow = StateGraph(State)
+        workflow.add_node("retrieve", retrieve)
+        workflow.add_node("generate", generate)
+        workflow.add_edge(START, "retrieve")
+        workflow.add_edge("retrieve", "generate")
+        return workflow.compile()
+
+    def query(self, text: str) -> str:
+        result = self.graph.invoke({"question": text})
+        return result["answer"]

reload_job.sh

@@ -19,5 +19,5 @@ docker compose -f "$COMPOSE_FILE" run --rm --no-deps --build \
   -v "$PROJECT_DIR:/work" \
   -w /work \
   chatbot-api \
-  sh -lc "python export_pages.py && python reload_vector_store.py"
+  sh -lc "python export_pages.py && python vector_store.py"
 echo "[reindex] done"

reload_vector_store.py

@@ -1,6 +1,4 @@
 python-dotenv>=1.0.0
-fastapi>=0.115.0
-uvicorn[standard]>=0.30.0
 langchain-openai
 langchain_chroma
 langchain-text-splitters
@@ -9,3 +7,7 @@ langgraph
 pathlib
 tqdm
 pypdf
+docx2txt
+pymupdf
+fastapi
+uvicorn

requirements.txt

@@ -0,0 +1,64 @@
+
+from dotenv import load_dotenv
+import os
+from pathlib import Path
+# LangChain & LangGraph Imports
+from langchain_openai import OpenAIEmbeddings
+from langchain_chroma import Chroma
+from langchain_community.document_loaders import PyMuPDFLoader, Docx2txtLoader, TextLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+# Configuration & Environment
+load_dotenv()
+# LLM_MODEL = os.getenv("MODEL_LLM")
+EMB_MODEL = os.getenv("MODEL_EMBEDDING")
+DOCS_DIR = 'exports'
+PERSIST_DIR = 'vector-store'
+
+
+
+def ingest():
+    if not os.path.exists(DOCS_DIR):
+        print(f"Error: {DOCS_DIR} folder not found.")
+        return
+    
+    # Document Loading Logic
+    print("--- Loading Documents ---")
+    paths = list(Path(DOCS_DIR).rglob("*"))
+    raw_docs = []
+    
+    for p in paths:
+        try:
+            if p.suffix.lower() == ".pdf":
+                loader = PyMuPDFLoader(str(p))
+                raw_docs.extend(loader.load())
+            elif p.suffix.lower() == ".docx":
+                raw_docs.extend(Docx2txtLoader(str(p)).load())
+            elif p.suffix.lower() in {".txt", ".md"}:
+                raw_docs.extend(TextLoader(str(p), encoding="utf-8").load())
+        except Exception as e:
+            print(f"Skipping {p.name} due to error: {e}")
+
+    if not raw_docs:
+        print("No documents found to index.")
+        return
+
+    print(f"Splitting {len(raw_docs)} documents...")
+    splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=120)
+    all_splits = splitter.split_documents(raw_docs)
+
+    print(f"Generating embeddings and saving to {PERSIST_DIR}...")
+    embeddings = OpenAIEmbeddings(model=EMB_MODEL)
+    
+    # Create the vector store on disk
+    Chroma.from_documents(
+        documents=all_splits,
+        embedding=embeddings,
+        persist_directory=PERSIST_DIR,
+        collection_name="collection_test"
+    )
+    print("Indexing complete!")
+
+if __name__ == "__main__":
+    ingest()
+