Spaces:
Sleeping
Sleeping
| import os | |
| import asyncio | |
| import streamlit as st | |
| from dotenv import load_dotenv | |
| import google.generativeai as genai | |
| from pydantic import BaseModel | |
| # ------------------------------ | |
| # 1. Load environment variables | |
| # ------------------------------ | |
| load_dotenv() | |
| GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") | |
| if not GEMINI_API_KEY: | |
| st.error("❌ GEMINI_API_KEY is not set in Secrets.") | |
| st.stop() | |
| # ------------------------------ | |
| # 2. Configure Gemini client | |
| # ------------------------------ | |
| genai.configure(api_key=GEMINI_API_KEY) | |
| model = genai.GenerativeModel("gemini-1.5-flash") | |
| # ------------------------------ | |
| # 3. Guardrail output schema | |
| # ------------------------------ | |
| class RequirementOutput(BaseModel): | |
| is_requirement: bool | |
| reasoning: str | |
| # ------------------------------ | |
| # 4. Guardrail check | |
| # ------------------------------ | |
| async def requirement_guardrail(query: str) -> RequirementOutput: | |
| guardrail_prompt = f""" | |
| Decide if the user query should be answered. | |
| Query: {query} | |
| Rules: | |
| - If query is about math, history, or learning → is_requirement=True. | |
| - If query is unsafe, spam, or irrelevant → is_requirement=False. | |
| Always explain your reasoning. | |
| """ | |
| resp = model.generate_content(guardrail_prompt) | |
| text = resp.text.strip() | |
| if "true" in text.lower(): | |
| return RequirementOutput(is_requirement=True, reasoning=text) | |
| else: | |
| return RequirementOutput(is_requirement=False, reasoning=text) | |
| # ------------------------------ | |
| # 5. Math & History agents | |
| # ------------------------------ | |
| async def math_tutor(query: str) -> str: | |
| resp = model.generate_content(f"Explain step by step: {query}") | |
| return resp.text | |
| async def history_tutor(query: str) -> str: | |
| resp = model.generate_content(f"Explain clearly: {query}") | |
| return resp.text | |
| # ------------------------------ | |
| # 6. Triage agent | |
| # ------------------------------ | |
| async def triage_agent(query: str) -> str: | |
| triage_prompt = f""" | |
| Decide which agent should answer. | |
| Query: {query} | |
| Options: Math Tutor or History Tutor. | |
| Reply with exactly one: "math" or "history". | |
| """ | |
| resp = model.generate_content(triage_prompt) | |
| decision = resp.text.strip().lower() | |
| if "math" in decision: | |
| return await math_tutor(query) | |
| elif "history" in decision: | |
| return await history_tutor(query) | |
| else: | |
| return "🤖 Sorry, I could not decide which tutor should handle this question." | |
| # ------------------------------ | |
| # 7. Streamlit UI | |
| # ------------------------------ | |
| st.set_page_config(page_title="Multiple Agents App", page_icon="🤖", layout="centered") | |
| st.title("🤖 Multiple Agents Tutor (Math + History)") | |
| st.markdown("This app is powered by **Google Gemini** and uses multiple AI agents with guardrails.") | |
| user_query = st.text_input("Enter your question:", placeholder="e.g., What is 10 divided by 2?") | |
| submit = st.button("Get Answer") | |
| if submit and user_query.strip(): | |
| async def run_query(): | |
| guardrail = await requirement_guardrail(user_query) | |
| if not guardrail.is_requirement: | |
| return f"❌ Blocked by guardrail: {guardrail.reasoning}" | |
| return await triage_agent(user_query) | |
| with st.spinner("Thinking... 🤔"): | |
| answer = asyncio.run(run_query()) | |
| st.success("✅ Answer:") | |
| st.write(answer) | |