DrugBot_Retrieval_Based / Scripts /Answer_Generation.py
Niranjan Sathish
Initial commit
6cf6a92
"""
Answer Generation Module for Retrieval-based Medical QA Chatbot
=================================================================
This module handles answer generation using Groq API with proper error handling.
"""
import os
from openai import OpenAI
# Get API key from environment
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
if GROQ_API_KEY is None:
print("[Warning] GROQ_API_KEY not set!")
client = None
else:
client = OpenAI(
api_key=GROQ_API_KEY,
base_url="https://api.groq.com/openai/v1"
)
# -------------------------------
# Function: Query Groq API
# -------------------------------
def query_groq(prompt, model="meta-llama/llama-4-scout-17b-16e-instruct", max_tokens=300):
"""
Sends a prompt to Groq API and returns the generated response.
Parameters:
prompt (str): The text prompt for the model.
model (str): Model name deployed on Groq API.
max_tokens (int): Maximum tokens allowed in the output.
Returns:
str: Model-generated response text.
"""
if client is None:
return "⚠️ Error: API key not configured. Please contact the administrator."
try:
response = client.chat.completions.create(
model=model,
messages=[
{"role": "system", "content": "You are a helpful biomedical assistant providing accurate drug information."},
{"role": "user", "content": prompt}
],
temperature=0.7,
max_tokens=max_tokens
)
return response.choices[0].message.content.strip()
except Exception as e:
print(f"[Answer Generation] Error calling Groq API: {e}")
return f"⚠️ Error generating answer: {str(e)}"
# -------------------------------
# Function: Build Prompt
# -------------------------------
def build_prompt(question, context):
"""
Constructs a prompt for the model combining the user question and retrieved context.
Parameters:
question (str): User's question.
context (str): Retrieved relevant text chunks.
Returns:
str: Complete prompt text.
"""
return f"""Based strictly on the following medical information, answer the question clearly and concisely.
Question: {question}
Context:
{context}
Instructions:
- Provide a direct, accurate answer based only on the context
- Use clear, simple language
- If the context doesn't contain enough information, say so
- Do not add information not present in the context
"""
# -------------------------------
# Function: Answer Generation
# -------------------------------
def answer_generation(question, top_chunks, top_k=3):
"""
Generates an answer based on retrieved top chunks.
Parameters:
question (str): User's question.
top_chunks (DataFrame): Retrieved top chunks with context.
top_k (int): Number of top chunks to use for answer generation.
Returns:
str: Final generated answer.
"""
try:
# Select top-k chunks
top_chunks = top_chunks.head(top_k)
print(f"[Answer Generation] Using top {len(top_chunks)} chunks")
if top_chunks.empty:
return "⚠️ No relevant information found. Please try rephrasing your question."
# Join context
context = "\n\n".join([
f"Drug: {row['drug_name']}\n"
f"Section: {row['section']}\n"
f"Info: {row['chunk_text']}"
for _, row in top_chunks.iterrows()
])
# Build prompt and query Groq
prompt = build_prompt(question, context)
answer = query_groq(prompt)
return answer
except Exception as e:
print(f"[Answer Generation] Error: {e}")
return f"⚠️ Error generating answer: {str(e)}"