server.py

import logging
import os
import shutil
import sys
import tempfile
from typing import TypedDict, List

# Compatibility fix for ChromaDB on systems with older SQLite versions (e.g., some Linux distros/Colab)
try:
    __import__('pysqlite3')
    sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
except ImportError:
    pass

from flask import Flask, request, jsonify
from werkzeug.utils import secure_filename

from langchain_chroma import Chroma
from langchain_community.document_loaders import PyPDFLoader
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_ollama import OllamaEmbeddings, ChatOllama
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langgraph.graph import END, StateGraph

# Configuration Params
DB_PATH = "./chroma_db"
LLM_NAME = "phi3.5"
EMBED_NAME = "mxbai-embed-large"
SERVER_PORT = 8000
CHUNK_SIZE = 800
CHUNK_OVERLAP = 100

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

app = Flask(__name__)

llm = ChatOllama(model=LLM_NAME, temperature=0)
embeddings = OllamaEmbeddings(model=EMBED_NAME)

vector_db = None

class AgentState(TypedDict):
    question: str
    generation: str
    documents: List[str]
    steps: List[str]

# --- Workflow Nodes ---

def retrieve_docs(state):
    logger.info("Retrieving documents for query: %s", state["question"])
    if not vector_db:
        return {"documents": [], "steps": state.get("steps", []) + ["retrieve_failed"]}
    
    docs = vector_db.as_retriever(search_kwargs={"k": 3}).invoke(state["question"])
    return {"documents": docs, "steps": state.get("steps", []) + ["retrieve"]}

def filter_docs(state):
    logger.info("Filtering documents for relevance.")
    if not state["documents"]:
        return {"documents": [], "steps": state["steps"] + ["filter_skipped"]}

    doc_text = "\n\n".join([f"[{i}] {d.page_content}" for i, d in enumerate(state["documents"])])

    check_prompt = ChatPromptTemplate.from_template(
        "Evaluate which documents help answer: {question}\n\n{context}\n\n"
        "Return only the indices (e.g. 0,2) or 'NONE'."
    )

    chain = check_prompt | llm | StrOutputParser()
    raw_output = chain.invoke({"question": state["question"], "context": doc_text})

    try:
        relevant_ids = [int(x.strip()) for x in raw_output.split(",") if x.strip().isdigit()]
        filtered = [state["documents"][i] for i in relevant_ids if i < len(state["documents"])]
    except Exception as e:
        logger.error("Error during document filtering: %s", e)
        filtered = state["documents"]

    return {"documents": filtered, "steps": state["steps"] + ["filter"]}

def generate_answer(state):
    logger.info("Generating answer.")
    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are an internal assistant. Use the context below to answer accurately. "
                   "If the answer isn't there, say you don't know. Don't hallucinate."),
        ("user", "Context: {context}\n\nQuestion: {question}")
    ])

    chain = prompt | llm | StrOutputParser()
    context = "\n\n".join([d.page_content for d in state.get("documents", [])]) if state.get("documents") else "No documents found."

    out = chain.invoke({"context": context, "question": state["question"]})
    return {"generation": out, "steps": state["steps"] + ["generate"]}

# --- Graph Construction ---
builder = StateGraph(AgentState)
builder.add_node("retrieve", retrieve_docs)
builder.add_node("filter", filter_docs)
builder.add_node("generate", generate_answer)

builder.set_entry_point("retrieve")
builder.add_edge("retrieve", "filter")
builder.add_edge("filter", "generate")
builder.add_edge("generate", END)

rag_app = builder.compile()

@app.route("/ingest", methods=["POST"])
def ingest_file():
    global vector_db
    if 'file' not in request.files:
        return jsonify({"error": "No file uploaded"}), 400

    file_obj = request.files['file']
    
    # Use tempfile to handle file creation and cleanup automatically
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
        file_obj.save(tmp_file.name)
        tmp_path = tmp_file.name

    try:
        loader = PyPDFLoader(tmp_path)
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP)
        chunks = text_splitter.split_documents(loader.load())

        if os.path.exists(DB_PATH):
            shutil.rmtree(DB_PATH)

        vector_db = Chroma.from_documents(
            documents=chunks,
            embedding=embeddings,
            persist_directory=DB_PATH
        )
        logger.info("Ingestion complete. Processed %d chunks.", len(chunks))
        return jsonify({"status": "success", "chunks": len(chunks)})
    except Exception as e:
        logger.error("Ingestion failed: %s", e)
        return jsonify({"status": "error", "message": str(e)}), 500
    finally:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)

@app.route("/chat", methods=["POST"])
def handle_chat():
    data = request.json
    if not data or "question" not in data:
        return jsonify({"error": "Missing question in payload"}), 400
        
    out = rag_app.invoke({"question": data.get("question"), "steps": []})
    return jsonify({
        "answer": out["generation"],
        "metadata": {"steps": out["steps"]}
    })

if __name__ == "__main__":
    if os.path.exists(DB_PATH):
        vector_db = Chroma(persist_directory=DB_PATH, embedding_function=embeddings)
        logger.info("Loaded existing Vector DB from %s", DB_PATH)
    
    app.run(host="0.0.0.0", port=SERVER_PORT)