111 lines
2.9 KiB
Python
111 lines
2.9 KiB
Python
import chromadb
|
|
from sentence_transformers import SentenceTransformer
|
|
from flask import Flask, request, jsonify, send_from_directory
|
|
|
|
model = SentenceTransformer(
|
|
"Snowflake/snowflake-arctic-embed-m-long",
|
|
device="cpu", # <-- the only line that changes
|
|
trust_remote_code=True
|
|
)
|
|
|
|
# Test Test Test
|
|
client = chromadb.HttpClient(host="chromadb", port=8000)
|
|
|
|
collection = client.get_or_create_collection("news")
|
|
|
|
|
|
def query(question):
|
|
|
|
query_embedding = model.encode(question, normalize_embeddings=True)
|
|
|
|
results = collection.query(
|
|
query_embeddings=query_embedding,
|
|
n_results=5,
|
|
)
|
|
|
|
return results
|
|
|
|
app = Flask(__name__)
|
|
|
|
@app.route("/query", methods=["POST"])
|
|
def query_endpoint():
|
|
data = request.get_json()
|
|
question = data.get("question")
|
|
if not question:
|
|
return jsonify({"error": "Missing 'question' in request body"}), 400
|
|
|
|
results = query(question)
|
|
|
|
# Adapt response to MCP Model standards
|
|
# Example MCP format:
|
|
# {
|
|
# "results": [
|
|
# {
|
|
# "document": ...,
|
|
# "score": ...,
|
|
# "metadata": {...}
|
|
# },
|
|
# ...
|
|
# ]
|
|
# }
|
|
mcp_results = []
|
|
for doc, score, meta in zip(
|
|
results.get("documents", [[]])[0],
|
|
results.get("distances", [[]])[0],
|
|
results.get("metadatas", [[]])[0]):
|
|
mcp_results.append({
|
|
"document": doc,
|
|
"score": float(score),
|
|
"metadata": meta
|
|
})
|
|
|
|
return jsonify({"results": mcp_results})
|
|
|
|
@app.route("/health", methods=["GET"])
|
|
def health():
|
|
return jsonify({"status": "ok"})
|
|
|
|
@app.route("/info", methods=["GET"])
|
|
def info():
|
|
return jsonify({
|
|
"provider": "StockDoc",
|
|
"model": "Snowflake/snowflake-arctic-embed-m-long",
|
|
"embedding_dim": 1024, # or your actual dimension
|
|
"collection": "news"
|
|
})
|
|
|
|
@app.route("/tools", methods=["GET"])
|
|
def tools():
|
|
return jsonify({
|
|
"endpoints": [
|
|
{
|
|
"path": "/query",
|
|
"method": "POST",
|
|
"description": "Query the vector database with a question.",
|
|
"request_format": {"question": "string"},
|
|
"response_format": {
|
|
"results": [
|
|
{"document": "string", "score": "float", "metadata": "object"}
|
|
]
|
|
}
|
|
},
|
|
{
|
|
"path": "/health",
|
|
"method": "GET",
|
|
"description": "Health check endpoint."
|
|
},
|
|
{
|
|
"path": "/info",
|
|
"method": "GET",
|
|
"description": "Returns model and service metadata."
|
|
}
|
|
]
|
|
})
|
|
|
|
@app.route("/openapi.json", methods=["GET"])
|
|
def openapi():
|
|
return send_from_directory('.', 'openapi.json')
|
|
|
|
if __name__ == "__main__":
|
|
app.run(host="0.0.0.0", port=5005, threaded=True, debug=True)
|