ocrAPP / rag.py
hkai20000's picture
Update rag.py
76ce598 verified
import os
from typing import List, Dict, Tuple
import numpy as np
from openai import OpenAI
from faq_store import FAQ_ENTRIES, FAQ_VECS
RAG_CONFIDENCE_THRESHOLD = 0.6
MAX_FAQ_MATCHES = 3
_EMBED_MODEL = "text-embedding-3-small"
_CHAT_MODEL = "gpt-4o-mini"
SYSTEM_PROMPT = (
"You are a helpful assistant for ScanAssured, a medical document OCR and NER app. "
"Answer only based on the provided FAQ context. "
"You do NOT have access to any user scan results or personal medical data. "
"For personal medical advice, always direct users to a qualified healthcare professional. "
"Keep answers concise and clear."
)
FALLBACK_MESSAGE = (
"I'm not certain about that. Please consult a qualified healthcare professional "
"for personal medical advice, or refer to the app documentation for usage questions."
)
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
_query_cache: dict[str, np.ndarray] = {}
def cosine(a: np.ndarray, b: np.ndarray) -> float:
return float(a.dot(b) / (np.linalg.norm(a) * np.linalg.norm(b)))
async def get_answer(question: str, history: List[Dict]) -> Tuple[str, List[Dict]]:
if question in _query_cache:
vec = _query_cache[question]
else:
resp = client.embeddings.create(model=_EMBED_MODEL, input=question)
vec = np.array(resp.data[0].embedding, dtype=np.float32)
_query_cache[question] = vec
scores = [(fid, cosine(vec, fvec)) for fid, fvec in FAQ_VECS]
scores.sort(key=lambda x: x[1], reverse=True)
if not scores or scores[0][1] < RAG_CONFIDENCE_THRESHOLD:
return FALLBACK_MESSAGE, []
matches = []
for fid, score in scores[:MAX_FAQ_MATCHES]:
faq = FAQ_ENTRIES[fid]
matches.append({"id": fid, "answer": faq["answer"], "source": faq["source"], "score": score})
messages: List[Dict] = [{"role": "system", "content": SYSTEM_PROMPT}]
for msg in history:
messages.append({"role": msg["role"], "content": msg["content"]})
for faq in matches:
messages.append({"role": "system", "content": faq["answer"]})
messages.append({"role": "user", "content": question})
chat_resp = client.chat.completions.create(
model=_CHAT_MODEL,
messages=messages,
stream=False,
)
answer = chat_resp.choices[0].message.content
citations = [{"id": faq["id"], "source": faq["source"]} for faq in matches]
return answer, citations