Update app.py

app.py

@@ -1,114 +1,92 @@
 import os
 import asyncio
 import streamlit as st
-from pydantic import BaseModel
 from dotenv import load_dotenv
-from agents import (
-    Agent,
-    InputGuardrail,
-    GuardrailFunctionOutput,
-    Runner,
-    OpenAIChatCompletionsModel,
-    RunConfig,
-    AsyncOpenAI,
-)
+import google.generativeai as genai
+from pydantic import BaseModel
 
 # ------------------------------
 # 1. Load environment variables
 # ------------------------------
 load_dotenv()
-GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
 
 if not GEMINI_API_KEY:
-    st.error("❌ GEMINI_API_KEY is not set in the .env file.")
+    st.error("❌ GEMINI_API_KEY is not set in Secrets.")
     st.stop()
 
 # ------------------------------
-# 2. Gemini client
-# ------------------------------
-external_client = AsyncOpenAI(
-    api_key=GEMINI_API_KEY,
-    base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
-)
-
-# ------------------------------
-# 3. Gemini model
+# 2. Configure Gemini client
 # ------------------------------
-model = OpenAIChatCompletionsModel(
-    model="gemini-2.0-flash",  # can be gemini-1.5-flash or gemini-1.5-pro
-    openai_client=external_client,
-)
+genai.configure(api_key=GEMINI_API_KEY)
+model = genai.GenerativeModel("gemini-1.5-flash")
 
 # ------------------------------
-# 4. RunConfig
-# ------------------------------
-config = RunConfig(
-    model=model,
-    model_provider=external_client,
-    tracing_disabled=True
-)
-
-# ------------------------------
-# 5. Guardrail output schema
+# 3. Guardrail output schema
 # ------------------------------
 class RequirementOutput(BaseModel):
     is_requirement: bool
     reasoning: str
 
 # ------------------------------
-# 6. Guardrail agent
+# 4. Guardrail check
 # ------------------------------
-guardrail_agent = Agent(
-    name="Guardrail check",
-    instructions=(
-        "You are a filter agent. "
-        "Your job is to decide if the user query is safe, relevant, and should be answered. "
-        "- If the query is about math, history, or general learning → set is_requirement=True. "
-        "- If the query is harmful, unsafe, spam, or irrelevant → set is_requirement=False. "
-        "Always explain your reasoning briefly."
-    ),
-    output_type=RequirementOutput,
-)
+async def requirement_guardrail(query: str) -> RequirementOutput:
+    guardrail_prompt = f"""
+    Decide if the user query should be answered.
 
-# ------------------------------
-# 7. Math & History tutor agents
-# ------------------------------
-math_tutor_agent = Agent(
-    name="Math Tutor",
-    handoff_description="Specialist agent for Math Tutoring",
-    instructions="You provide help with math-related queries. Explain step by step with reasoning and examples.",
-)
-
-history_tutor_agent = Agent(
-    name="History Tutor",
-    handoff_description="Specialist agent for History Tutoring",
-    instructions="You provide assistance with history-related queries. Explain important events and context clearly.",
-)
+    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)
 
 # ------------------------------
-# 8. Guardrail function
+# 5. Math & History agents
 # ------------------------------
-async def requirement_guardrail(ctx, agent, input_data):
-    result = await Runner.run(guardrail_agent, input_data, context=ctx.context, run_config=config)
-    final_output = result.final_output_as(RequirementOutput)
+async def math_tutor(query: str) -> str:
+    resp = model.generate_content(f"Explain step by step: {query}")
+    return resp.text
 
-    return GuardrailFunctionOutput(
-        output_info=final_output,
-        tripwire_triggered=not final_output.is_requirement
-    )
+async def history_tutor(query: str) -> str:
+    resp = model.generate_content(f"Explain clearly: {query}")
+    return resp.text
 
 # ------------------------------
-# 9. Triage agent
+# 6. Triage agent
 # ------------------------------
-triage_agent = Agent(
-    name="Triage Agent",
-    instructions="You decide which agent (Math or History) should handle the query.",
-    handoffs=[math_tutor_agent, history_tutor_agent],
-    input_guardrails=[InputGuardrail(guardrail_function=requirement_guardrail)],
-)
+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."
 
 # ------------------------------
-# 10. Streamlit UI
+# 7. Streamlit UI
 # ------------------------------
 st.set_page_config(page_title="Multiple Agents App", page_icon="🤖", layout="centered")
 st.title("🤖 Multiple Agents Tutor (Math + History)")
@@ -118,11 +96,13 @@ user_query = st.text_input("Enter your question:", placeholder="e.g., What is 10
 submit = st.button("Get Answer")
 
 if submit and user_query.strip():
-    async def run_agent():
-        result = await Runner.run(triage_agent, user_query, run_config=config)
-        return result.final_output
+    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_agent())
+        answer = asyncio.run(run_query())
         st.success("✅ Answer:")
         st.write(answer)