Spaces:
Running
Running
| import os | |
| from dotenv import load_dotenv | |
| from langchain.prompts import PromptTemplate | |
| from langchain_groq import ChatGroq | |
| from typing import Literal | |
| # Load environment variables | |
| load_dotenv() | |
| # Initialize LLMs | |
| def initialize_llms(): | |
| """Initialize and return the LLM instances""" | |
| groq_api_key = os.getenv("GROQ_API_KEY") | |
| return { | |
| "llm": ChatGroq( | |
| temperature=0.1, model="llama-3.3-70b-versatile", api_key=groq_api_key | |
| ), | |
| "step_back_llm": ChatGroq( | |
| temperature=0, model="Gemma2-9B-IT", api_key=groq_api_key | |
| ), | |
| } | |
| # Query refinement | |
| def refine_query(query: str, llm: ChatGroq) -> str: | |
| """Enhance pediatric medicine queries for better retrieval while preserving clinical intent""" | |
| template = """ | |
| You are a medical language expert. Your task is to improve the following user question by: | |
| - Correcting any grammatical or spelling errors | |
| - Clarifying vague or ambiguous wording | |
| - Improving sentence structure for readability and precision | |
| - Maintaining the original meaning and clinical focus | |
| Do not add new information. Do not expand abbreviations unless they are unclear. Do not include any commentary or explanation. | |
| Original query: {original_query} | |
| Improved medical question: | |
| """ | |
| prompt = PromptTemplate(input_variables=["original_query"], template=template) | |
| chain = prompt | llm | |
| return chain.invoke({"original_query": query}).content | |
| def query_to_retrieve(query, llm): | |
| """Convert a query to a format suitable for retrieval""" | |
| template = """ | |
| You are an expert in pediatric medical information retrieval. | |
| Your task is to rewrite the following question into a single, concise sentence containing only the most relevant medical and pediatric concepts. This sentence will be used for semantic search in a vector database. | |
| Instructions: | |
| - Include only the core clinical focus (conditions, symptoms, treatments, procedures). | |
| - Mention pediatric-specific entities if relevant (e.g., age group, child-specific medication). | |
| - Remove all conversational language and filler. | |
| - Preserve the original intent. | |
| - Output only one clean, search-optimized sentence. | |
| Original query: {original_query} | |
| Search-ready query: | |
| """ | |
| prompt = PromptTemplate(input_variables=["original_query"], template=template) | |
| chain = prompt | llm | |
| return chain.invoke({"original_query": query}).content | |
| def answer_query_with_chunks( | |
| query: str, | |
| retrieved_docs, | |
| llm: ChatGroq, | |
| ) -> str: | |
| try: | |
| # Embed query using the same embedding function | |
| query_improved = refine_query(query, llm) | |
| if not retrieved_docs: | |
| return "Sorry, no relevant medical information was found." | |
| # Construct context for the LLM | |
| context = "\n\n".join(retrieved_docs) | |
| system_prompt = """ | |
| You are a pediatric medical assistant. | |
| Based only on the provided context, answer the user's question concisely and accurately with explanation. | |
| If the answer is not present in the context, say: "The answer is not available in the current documents." | |
| Context: | |
| {context} | |
| User question: | |
| {query} | |
| Answer: | |
| """ | |
| prompt = PromptTemplate( | |
| input_variables=["context", "query"], | |
| template=system_prompt, | |
| ) | |
| chain = prompt | llm | |
| return chain.invoke({"context": context, "query": query_improved}).content | |
| except Exception as e: | |
| return f"An error occurred while answering the query: {str(e)}" | |