from google import genai
import gradio as gr
import os

API_KEY = os.getenv("GEMINI_API_KEY")

client = genai.Client(
    api_key=API_KEY
)

# RAG用知識ベース
with open("kosen.txt", encoding="utf-8") as f:
    knowledge = f.read()


def rag_chat(message, history):
    # 過去の会話を文字列化
    conversation = ""
    if history is None:
        history = []
    
    for msg in history:
        role = msg["role"]
        content = msg["content"]
        conversation += f"{role}: {content}\n"
    # Geminiへ渡すプロンプト
    prompt = f"""
あなたは富山高等専門学校の案内AIです。
以下の資料を参考に回答してください。
ただし、資料にない内容に関しては、「参考資料に情報がないため、geminiによる回答を行います。
ただし、不正確である場合がありますので予めご了承ください。」と伝えた上で、geminiによる回答を行ってください。
=====資料=====
{knowledge}
============
=====過去の会話=====
{conversation}
===================
ユーザー:
{message}
"""
    response = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=prompt
    )
    return response.text
app = gr.Interface(
    fn=rag_chat,
    inputs=gr.Textbox(
        label="質問",
        lines=3
    ),
    outputs=gr.Textbox(
        label="回答",
        lines=10
    ),
    title="富山高専RAGアプリ"
)

app.launch()