mirror of
https://github.com/SirBlobby/Hoya26.git
synced 2026-02-03 19:24:34 -05:00
VectorDB Search limit
This commit is contained in:
@@ -5,4 +5,5 @@ def get_mongo_client ():
|
|||||||
uri = os.environ.get("MONGO_URI")
|
uri = os.environ.get("MONGO_URI")
|
||||||
if not uri:
|
if not uri:
|
||||||
raise ValueError("MONGO_URI environment variable not set")
|
raise ValueError("MONGO_URI environment variable not set")
|
||||||
return MongoClient (uri )
|
print("Connecting to MongoDB Atlas...")
|
||||||
|
return MongoClient(uri, serverSelectionTimeoutMS=5000, connectTimeoutMS=10000, socketTimeoutMS=10000)
|
||||||
|
|||||||
@@ -2,16 +2,12 @@ from flask import Blueprint, request, jsonify
|
|||||||
from src.rag.gemeni import GeminiClient
|
from src.rag.gemeni import GeminiClient
|
||||||
from src.gemini import ask_gemini_with_rag
|
from src.gemini import ask_gemini_with_rag
|
||||||
from src.chroma.vector_store import search_documents
|
from src.chroma.vector_store import search_documents
|
||||||
|
from src.chroma.chroma_store import get_collection
|
||||||
from src.rag.embeddings import get_embedding
|
from src.rag.embeddings import get_embedding
|
||||||
|
|
||||||
gemini_bp = Blueprint('gemini', __name__)
|
# ... (imports)
|
||||||
brain = None
|
|
||||||
|
|
||||||
def get_brain():
|
gemini_bp = Blueprint('gemini', __name__)
|
||||||
global brain
|
|
||||||
if brain is None:
|
|
||||||
brain = GeminiClient()
|
|
||||||
return brain
|
|
||||||
|
|
||||||
@gemini_bp.route('/ask', methods=['POST'])
|
@gemini_bp.route('/ask', methods=['POST'])
|
||||||
def ask():
|
def ask():
|
||||||
@@ -27,33 +23,34 @@ def ask():
|
|||||||
print(f"Generating embedding for prompt: {prompt}")
|
print(f"Generating embedding for prompt: {prompt}")
|
||||||
query_embedding = get_embedding(prompt)
|
query_embedding = get_embedding(prompt)
|
||||||
|
|
||||||
print("Searching Vector Database for context...")
|
# Maximize context retrieval (20 chunks)
|
||||||
search_results = search_documents(query_embedding, num_results=50)
|
print("Searching Vector Database for context (limit=20)...")
|
||||||
|
search_results = search_documents(query_embedding, num_results=20)
|
||||||
|
|
||||||
# Special handling for Georgetown University queries to ensure those docs are included
|
# KEYWORD BOOST: Specifically check for Georgetown documents if mentioned
|
||||||
# even if generic corporate reports outrank them in vector search.
|
# This ensures we get specific entity hits even if semantic similarity is weak for chunks
|
||||||
if "georgetown" in prompt.lower():
|
if "georgetown" in prompt.lower():
|
||||||
try:
|
try:
|
||||||
from src.mongo.connection import get_mongo_client
|
print("Applied keyword boost: 'Georgetown'")
|
||||||
client = get_mongo_client()
|
coll = get_collection()
|
||||||
# Use hardcoded DB name to match vector_store.py
|
# ChromaDB text search filter
|
||||||
db = client["ethix_vectors"]
|
gt_results = coll.get(
|
||||||
collection = db["rag_documents"]
|
where_document={"$contains": "Georgetown"},
|
||||||
|
limit=20
|
||||||
|
)
|
||||||
|
|
||||||
# Fetch docs with Georgetown or MAP_INFO in the filename/source
|
if gt_results and gt_results['documents']:
|
||||||
gt_docs = list(collection.find({"source": {"$regex": "Georgetown|MAP_INFO", "$options": "i"}}).limit(30))
|
print(f"Found {len(gt_results['documents'])} keyword-matched documents.")
|
||||||
|
for i, text in enumerate(gt_results['documents']):
|
||||||
if gt_docs:
|
# Add to results if not already present (check by text content hash or just duplication)
|
||||||
print(f"Direct Match: Found {len(gt_docs)} Georgetown specific documents.")
|
# We'll just append with high score, duplicates handled by set later or LLM tolerance
|
||||||
for doc in gt_docs:
|
|
||||||
# Normalize to match search_results format
|
|
||||||
search_results.append({
|
search_results.append({
|
||||||
"text": doc.get("text", ""),
|
"text": text,
|
||||||
"metadata": doc.get("metadata", {"source": doc.get("source", "Georgetown File")}),
|
"metadata": gt_results['metadatas'][i] if gt_results['metadatas'] else {},
|
||||||
"score": 1.0 # High priority
|
"score": 1.0
|
||||||
})
|
})
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"Error checking Georgetown docs: {e}")
|
print(f"Keyword boost failed: {e}")
|
||||||
|
|
||||||
retrieved_context = ""
|
retrieved_context = ""
|
||||||
if search_results:
|
if search_results:
|
||||||
@@ -87,7 +84,7 @@ def ask():
|
|||||||
full_context = retrieved_context
|
full_context = retrieved_context
|
||||||
|
|
||||||
# Step 3: Ask Gemini
|
# Step 3: Ask Gemini
|
||||||
client = get_brain()
|
client = GeminiClient()
|
||||||
response = client.ask(prompt, full_context)
|
response = client.ask(prompt, full_context)
|
||||||
|
|
||||||
return jsonify({
|
return jsonify({
|
||||||
|
|||||||
Reference in New Issue
Block a user