86 lines
2.3 KiB
Python
86 lines
2.3 KiB
Python
from flask import Flask, request, jsonify
|
|
from sentence_transformers import SentenceTransformer
|
|
import weaviate
|
|
import openai
|
|
import os
|
|
|
|
app = Flask(__name__)
|
|
|
|
# Load local model
|
|
model = SentenceTransformer('all-MiniLM-L6-v2')
|
|
|
|
# Configure Weaviate client (assumes running on localhost)
|
|
weaviate_client = weaviate.Client("http://localhost:8083")
|
|
|
|
# OpenAI setup
|
|
openai.api_key = os.getenv("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")
|
|
|
|
# Create Weaviate schema (if not exists)
|
|
CLASS_NAME = "Document"
|
|
|
|
if not weaviate_client.schema.contains({"class": CLASS_NAME}):
|
|
weaviate_client.schema.create_class({
|
|
"class": CLASS_NAME,
|
|
"vectorizer": "none", # Because we use external embeddings
|
|
"properties": [
|
|
{
|
|
"name": "content",
|
|
"dataType": ["text"]
|
|
}
|
|
]
|
|
})
|
|
|
|
|
|
@app.route("/ingest", methods=["POST"])
|
|
def ingest():
|
|
content = request.json.get("content")
|
|
vector = model.encode([content])[0].tolist()
|
|
|
|
weaviate_client.data_object.create(
|
|
data_object={"content": content},
|
|
class_name=CLASS_NAME,
|
|
vector=vector
|
|
)
|
|
return jsonify({"status": "saved"})
|
|
|
|
|
|
@app.route("/query", methods=["POST"])
|
|
def query():
|
|
question = request.json.get("question")
|
|
q_vector = model.encode([question])[0].tolist()
|
|
|
|
response = weaviate_client.query.get(CLASS_NAME, ["content"])\
|
|
.with_near_vector({"vector": q_vector})\
|
|
.with_limit(3).do()
|
|
|
|
relevant_chunks = [d["content"] for d in response["data"]["Get"][CLASS_NAME]]
|
|
context = "\n".join(relevant_chunks)
|
|
|
|
prompt = f"""
|
|
Answer the following question based on the context below.
|
|
|
|
Context:
|
|
{context}
|
|
|
|
Question: {question}
|
|
"""
|
|
|
|
completion = openai.ChatCompletion.create(
|
|
model="gpt-3.5-turbo",
|
|
messages=[
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{"role": "user", "content": prompt}
|
|
]
|
|
)
|
|
|
|
answer = completion['choices'][0]['message']['content'].strip()
|
|
return jsonify({"answer": answer})
|
|
|
|
|
|
@app.route("/")
|
|
def root():
|
|
return "Embedding/QA API is running."
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app.run(host="0.0.0.0", port=5000)
|