mirror of
https://github.com/SirBlobby/Hoya26.git
synced 2026-02-03 19:24:34 -05:00
95 lines
2.8 KiB
Python
95 lines
2.8 KiB
Python
from flask import Blueprint, request, jsonify
|
|
from src.rag.gemeni import GeminiClient
|
|
from src.gemini import ask_gemini_with_rag
|
|
from src.chroma.vector_store import search_documents
|
|
from src.rag.embeddings import get_embedding
|
|
|
|
gemini_bp = Blueprint('gemini', __name__)
|
|
brain = None
|
|
|
|
def get_brain():
|
|
global brain
|
|
if brain is None:
|
|
brain = GeminiClient()
|
|
return brain
|
|
|
|
@gemini_bp.route('/ask', methods=['POST'])
|
|
def ask():
|
|
data = request.json
|
|
prompt = data.get("prompt")
|
|
context = data.get("context", "")
|
|
|
|
if not prompt:
|
|
return jsonify({"error": "No prompt provided"}), 400
|
|
|
|
try:
|
|
# Step 1: Retrieve relevant context from ChromaDB (RAG)
|
|
print(f"Generating embedding for prompt: {prompt}")
|
|
query_embedding = get_embedding(prompt)
|
|
|
|
print("Searching ChromaDB for context...")
|
|
search_results = search_documents(query_embedding, num_results=15)
|
|
|
|
retrieved_context = ""
|
|
if search_results:
|
|
print(f"Found {len(search_results)} documents.")
|
|
retrieved_context = "RELEVANT INFORMATION FROM DATABASE:\n"
|
|
for res in search_results:
|
|
# Include metadata if useful, e.g. brand name or date
|
|
meta = res.get('metadata', {})
|
|
source_info = f"[Source: {meta.get('type', 'doc')} - {meta.get('product_name', 'Unknown')}]"
|
|
retrieved_context += f"{source_info}\n{res['text']}\n\n"
|
|
else:
|
|
print("No relevant documents found.")
|
|
|
|
# Step 2: Combine manual context (if any) with retrieved context
|
|
full_context = context
|
|
if retrieved_context:
|
|
if full_context:
|
|
full_context += "\n\n" + retrieved_context
|
|
else:
|
|
full_context = retrieved_context
|
|
|
|
# Step 3: Ask Gemini
|
|
client = get_brain()
|
|
response = client.ask(prompt, full_context)
|
|
|
|
return jsonify({
|
|
"status": "success",
|
|
"reply": response
|
|
})
|
|
except Exception as e:
|
|
print(f"Error in RAG flow: {e}")
|
|
return jsonify({
|
|
"status": "error",
|
|
"message": str(e)
|
|
}), 500
|
|
|
|
@gemini_bp .route ('/rag',methods =['POST'])
|
|
def rag ():
|
|
data =request .json
|
|
prompt =data .get ("prompt")
|
|
category =data .get ("category")
|
|
|
|
if not prompt :
|
|
return jsonify ({"error":"No prompt provided"}),400
|
|
|
|
try :
|
|
response =ask_gemini_with_rag (prompt ,category =category )
|
|
return jsonify ({
|
|
"status":"success",
|
|
"reply":response
|
|
})
|
|
except Exception as e :
|
|
return jsonify ({
|
|
"status":"error",
|
|
"message":str (e )
|
|
}),500
|
|
|
|
@gemini_bp .route ('/vision',methods =['POST'])
|
|
def vision ():
|
|
return jsonify ({
|
|
"status":"error",
|
|
"message":"Vision endpoint not yet implemented"
|
|
}),501
|