File size: 1,861 Bytes
4416e3b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | from flask import Flask, render_template, request, jsonify
from rag_model import get_qa_chain, build_rag_system, PERSIST_DIRECTORY, HuggingFaceEmbeddings, Chroma
from dotenv import load_dotenv
import os
load_dotenv(override=True)
print(f"App: API Key loaded (starting with {os.getenv('GOOGLE_API_KEY', 'None')[:5]}...)")
app = Flask(__name__)
# Load or Build Vector Store
if not os.path.exists(PERSIST_DIRECTORY):
print("Vector store not found. Building...")
vectorstore = build_rag_system()
else:
print("Loading existing vector store...")
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
vectorstore = Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=embeddings)
qa_chain = get_qa_chain(vectorstore)
@app.route('/')
def index():
return render_template('index.html')
@app.route('/ask', methods=['POST'])
def ask():
data = request.json
query = data.get('query')
print(f"Processing query: {query}")
if not query:
return jsonify({"error": "No query provided"}), 400
if not qa_chain:
return jsonify({"error": "QA chain not initialized. Check GOOGLE_API_KEY."}), 500
try:
result = qa_chain({"query": query})
answer = result['result']
sources = []
seen = set()
for doc in result['source_documents']:
source_name = doc.metadata.get('name', 'Unknown')
if source_name not in seen:
sources.append(source_name)
seen.add(source_name)
return jsonify({
"answer": answer,
"sources": sources[:3] # Return top 3 unique names
})
except Exception as e:
return jsonify({"error": str(e)}), 500
if __name__ == '__main__':
app.run(debug=True, port=5000)
|