Merge pull request #17744 from Classic298/fix-rag-full-context

Fix: Prevent RAG queries when all files are in full context
This commit is contained in:
Tim Jaeryang Baek 2025-09-25 11:55:41 -05:00 committed by GitHub
commit cd417ca0ba
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -641,48 +641,53 @@ async def chat_completion_files_handler(
sources = [] sources = []
if files := body.get("metadata", {}).get("files", None): if files := body.get("metadata", {}).get("files", None):
# Check if all files are in full context mode
all_full_context = all(item.get("context") == "full" for item in files)
queries = [] queries = []
try: if not all_full_context:
queries_response = await generate_queries(
request,
{
"model": body["model"],
"messages": body["messages"],
"type": "retrieval",
},
user,
)
queries_response = queries_response["choices"][0]["message"]["content"]
try: try:
bracket_start = queries_response.find("{") queries_response = await generate_queries(
bracket_end = queries_response.rfind("}") + 1 request,
{
"model": body["model"],
"messages": body["messages"],
"type": "retrieval",
},
user,
)
queries_response = queries_response["choices"][0]["message"]["content"]
if bracket_start == -1 or bracket_end == -1: try:
raise Exception("No JSON object found in the response") bracket_start = queries_response.find("{")
bracket_end = queries_response.rfind("}") + 1
queries_response = queries_response[bracket_start:bracket_end] if bracket_start == -1 or bracket_end == -1:
queries_response = json.loads(queries_response) raise Exception("No JSON object found in the response")
except Exception as e:
queries_response = {"queries": [queries_response]}
queries = queries_response.get("queries", []) queries_response = queries_response[bracket_start:bracket_end]
except: queries_response = json.loads(queries_response)
pass except Exception as e:
queries_response = {"queries": [queries_response]}
queries = queries_response.get("queries", [])
except:
pass
if len(queries) == 0: if len(queries) == 0:
queries = [get_last_user_message(body["messages"])] queries = [get_last_user_message(body["messages"])]
await __event_emitter__( if not all_full_context:
{ await __event_emitter__(
"type": "status", {
"data": { "type": "status",
"action": "queries_generated", "data": {
"queries": queries, "action": "queries_generated",
"done": False, "queries": queries,
}, "done": False,
} },
) }
)
try: try:
# Offload get_sources_from_items to a separate thread # Offload get_sources_from_items to a separate thread
@ -711,7 +716,8 @@ async def chat_completion_files_handler(
r=request.app.state.config.RELEVANCE_THRESHOLD, r=request.app.state.config.RELEVANCE_THRESHOLD,
hybrid_bm25_weight=request.app.state.config.HYBRID_BM25_WEIGHT, hybrid_bm25_weight=request.app.state.config.HYBRID_BM25_WEIGHT,
hybrid_search=request.app.state.config.ENABLE_RAG_HYBRID_SEARCH, hybrid_search=request.app.state.config.ENABLE_RAG_HYBRID_SEARCH,
full_context=request.app.state.config.RAG_FULL_CONTEXT, full_context=all_full_context
or request.app.state.config.RAG_FULL_CONTEXT,
user=user, user=user,
), ),
) )