kumardatascience/Multi-Source-RAG-AI-System-with-Query-Routing

1"""LLM client wrapper. Handles all communication with Gemini."""2 3import os4from dotenv import load_dotenv5from google import genai6from google.genai import types7 8load_dotenv()9 10 11API_KEY = os.getenv("GOOGLE_API_KEY")12 13# Only create the client if the key exists. Tests can import this module14# without a real key — they only need _build_system_prompt etc.15client = genai.Client(api_key=API_KEY) if API_KEY else None16 17 18MODEL_NAME = "gemini-2.5-flash-lite"19 20# Base system prompt — controls bot behavior across the whole conversation21BASE_SYSTEM_PROMPT = (22    "You are a helpful, friendly AI assistant. "23    "Keep replies clear and concise. "24    "When the user shares personal details (like their name), remember them."25)26 27 28def _build_system_prompt(context: str | None) -> str:29    if not context:30        return BASE_SYSTEM_PROMPT31 32    return (33        BASE_SYSTEM_PROMPT34        + "\n\nYou have been given context to help answer the user's question. "35        + "The context may come from the user's documents (marked 'FROM YOUR DOCUMENTS') "36        + "and/or from a live web search (marked 'FROM THE WEB').\n\n"37        + "Instructions:\n"38        + "1. Use the provided context as your primary source.\n"39        + "2. If multiple sources are provided, synthesize them into one clear answer. "40        + "Mention briefly where key facts came from (e.g., 'According to your documents...' or 'Recent web sources indicate...').\n"41        + "3. If the sources conflict, point this out instead of picking one silently.\n"42        + "4. If the context doesn't contain the answer, say so honestly — don't guess.\n"43        + "5. For general/casual questions not related to the context, answer normally from your own knowledge.\n\n"44        + "=== CONTEXT START ===\n"45        + context46        + "\n=== CONTEXT END ==="47    )48 49 50async def stream_response(history: list[dict], context: str | None = None):51    """52    Send the full chat history to Gemini and stream the reply.53 54    history: list of {"role": "user" | "assistant", "content": "..."}55    context: optional retrieved document chunks to ground the answer56    """57 58    if client is None:59        raise ValueError(60            "GOOGLE_API_KEY not found. Make sure your .env file exists and has the key."61        ) 62 63    gemini_contents = [64        types.Content(65            role="user" if msg["role"] == "user" else "model",66            parts=[types.Part.from_text(text=msg["content"])],67        )68        for msg in history69    ]70 71    response_stream = await client.aio.models.generate_content_stream(72        model=MODEL_NAME,73        contents=gemini_contents,74        config=types.GenerateContentConfig(75            system_instruction=_build_system_prompt(context),76        ),77    )78 79    async for chunk in response_stream:80        if chunk.text:81            yield chunk.text