Create app.py

app.py

@@ -0,0 +1,211 @@
+from fastapi import FastAPI, Request, UploadFile, Form
+from fastapi.responses import JSONResponse
+import os, io, json, tempfile, logging
+import boto3
+from typing import List
+
+from langchain_community.document_loaders import PyMuPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import Chroma
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain.chains import create_history_aware_retriever, create_retrieval_chain
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from chromadb.config import Settings
+
+# ------------------ Configuration ------------------
+S3_BUCKET = "imagesdatalibrary"
+S3_PREFIX = "RetailME Magazine/"
+AWS_REGION = "ap-south-1"
+VECTOR_STORE_DIR = "persistent_vector_store"
+COLLECTION_NAME = "my_documents"
+PROCESSED_FILES_JSON = "processed_files.json"
+
+# ------------------ Setup ------------------
+app = FastAPI()
+logging.basicConfig(level=logging.INFO)
+
+s3 = boto3.client(
+    "s3",
+    aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID", "").strip(),
+    aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY", "").strip(),
+    region_name=AWS_REGION
+)
+
+vector_store = None
+processed_files = set()
+chat_history = [AIMessage(content="Hello! I am your AI assistant. Ask me questions based on your documents.")]
+
+# ------------------ Utilities ------------------
+
+def load_processed_files():
+    global processed_files
+    if os.path.exists(PROCESSED_FILES_JSON):
+        with open(PROCESSED_FILES_JSON, "r") as f:
+            processed_files = set(json.load(f))
+
+def save_processed_files():
+    with open(PROCESSED_FILES_JSON, "w") as f:
+        json.dump(list(processed_files), f)
+
+def get_chroma_settings():
+    return Settings(
+        anonymized_telemetry=False,
+        is_persistent=True,
+        persist_directory=VECTOR_STORE_DIR
+    )
+
+def initialize_vector_store():
+    global vector_store
+    if os.path.exists(VECTOR_STORE_DIR):
+        vector_store = Chroma(
+            persist_directory=VECTOR_STORE_DIR,
+            embedding_function=OpenAIEmbeddings(),
+            collection_name=COLLECTION_NAME,
+            client_settings=get_chroma_settings()
+        )
+    else:
+        vector_store = None
+
+# ------------------ Document Processing ------------------
+
+def list_s3_pdfs():
+    paginator = s3.get_paginator("list_objects_v2")
+    pages = paginator.paginate(Bucket=S3_BUCKET, Prefix=S3_PREFIX)
+    for page in pages:
+        for obj in page.get("Contents", []):
+            key = obj["Key"]
+            if key.endswith(".pdf"):
+                yield key.strip()
+
+def process_new_s3_files():
+    global vector_store, processed_files
+    files_found = False
+    try:
+        s3_keys = list(list_s3_pdfs())
+        unprocessed = [key for key in s3_keys if key not in processed_files]
+        for s3_key in unprocessed:
+            files_found = True
+            file_stream = io.BytesIO()
+            s3.download_fileobj(S3_BUCKET, s3_key, file_stream)
+            file_stream.seek(0)
+
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
+                tmp_file.write(file_stream.read())
+                tmp_path = tmp_file.name
+
+            loader = PyMuPDFLoader(tmp_path)
+            docs = loader.load()
+            os.remove(tmp_path)
+
+            for idx, doc in enumerate(docs):
+                doc.metadata["source"] = s3_key.strip()
+                doc.metadata["page"] = idx + 1
+
+            if not docs:
+                logging.warning(f"No text extracted from {s3_key}")
+                continue
+
+            splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+            chunks = splitter.split_documents(docs)
+
+            if vector_store:
+                vector_store.add_documents(chunks)
+            else:
+                vector_store = Chroma.from_documents(
+                    chunks,
+                    OpenAIEmbeddings(),
+                    persist_directory=VECTOR_STORE_DIR,
+                    collection_name=COLLECTION_NAME,
+                    client_settings=get_chroma_settings()
+                )
+
+            processed_files.add(s3_key.strip())
+            save_processed_files()
+
+        return {"status": "success", "processed": len(unprocessed)}
+
+    except Exception as e:
+        logging.error(f"Error: {e}")
+        return {"status": "error", "message": str(e)}
+
+# ------------------ LangChain Chain ------------------
+
+def get_context_retriever_chain(vector_store):
+    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
+    retriever = vector_store.as_retriever()
+    prompt = ChatPromptTemplate.from_messages([
+        MessagesPlaceholder(variable_name="chat_history"),
+        ("user", "{input}"),
+        ("user", "Generate a search query to help retrieve documents.")
+    ])
+    return create_history_aware_retriever(llm, retriever, prompt)
+
+def get_conversational_rag_chain(retriever_chain):
+    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", "You are a helpful assistant. Only answer using the provided context:\n\n{context}"),
+        MessagesPlaceholder(variable_name="chat_history"),
+        ("user", "{input}")
+    ])
+    doc_chain = create_stuff_documents_chain(llm=llm, prompt=prompt)
+    return create_retrieval_chain(retriever_chain, doc_chain)
+
+# ------------------ Q&A ------------------
+
+def ask_question(question):
+    global chat_history
+    if vector_store is None:
+        return "No documents loaded yet.", []
+
+    retriever_chain = get_context_retriever_chain(vector_store)
+    rag_chain = get_conversational_rag_chain(retriever_chain)
+
+    result = rag_chain.invoke({
+        "input": question.strip(),
+        "chat_history": chat_history
+    })
+
+    answer = result.get("answer", "No answer found.")
+    chat_history.append(HumanMessage(content=question.strip()))
+    chat_history.append(AIMessage(content=answer))
+
+    docs = vector_store.similarity_search(question.strip(), k=3)
+    sources = []
+    for doc in docs:
+        meta = doc.metadata
+        sources.append({
+            "source": meta.get("source", "Unknown").strip(),
+            "page": meta.get("page", "N/A"),
+            "snippet": doc.page_content.strip()[:300]
+        })
+
+    return answer, sources
+
+# ------------------ FastAPI Endpoints ------------------
+
+@app.on_event("startup")
+def startup_event():
+    load_processed_files()
+    initialize_vector_store()
+    process_new_s3_files()
+
+@app.get("/status")
+def get_status():
+    return {"processed_files_count": len(processed_files)}
+
+@app.post("/sync")
+def sync_s3():
+    result = process_new_s3_files()
+    initialize_vector_store()
+    return result
+
+@app.post("/ask")
+def ask_from_documents(payload: dict):
+    question = payload.get("question", "").strip()
+    if not question:
+        return JSONResponse(status_code=400, content={"error": "Question cannot be empty"})
+
+    answer, sources = ask_question(question)
+    return {"answer": answer, "sources": sources}