import gradio as gr
import os
import json
import requests
import openai

# Set up API keys
openai.api_key = os.getenv('OPENAI_API_KEY')
ragie_api_key = os.getenv('RAGIE_API_KEY')

def get_response(query):
    """Combined Ragie and OpenAI following Step 4"""
    try:
        # First, get chunks from Ragie using our working method
        response = requests.post(
            "https://api.ragie.ai/retrievals",
            headers={
                'accept': 'application/json',
                'authorization': f'Bearer {ragie_api_key}',
                'content-type': 'application/json'
            },
            json={
                "query": query,
                "top_k": 8,
                "rerank": True
            }
        )
        
        print(f"\nRagie API Response Status: {response.status_code}")
        
        if not response.ok:
            print(f"Failed to retrieve data from Ragie API: {response.status} {response.reason}")
            return "Failed to retrieve data from Ragie API"
            
        # Get chunks exactly as shown in Step 4
        data = response.json()
        chunk_text = [chunk["text"] for chunk in data.get("scored_chunks", [])]
        
        # Use exact system prompt from Step 4
        system_prompt = f"""These are very important to follow:

You are "Ragie AI", a professional but friendly AI chatbot working as an assitant to the user.

Your current task is to help the user based on all of the information available to you shown below.
Answer informally, directly, and concisely without a heading or greeting but include everything relevant.
Use richtext Markdown when appropriate including bold, italic, paragraphs, and lists when helpful.
If using LaTeX, use double $$ as delimiter instead of single $. Use $$...$$ instead of parentheses.
Organize information into multiple sections or points when appropriate.
Don't include raw item IDs or other raw fields from the source.
Don't use XML or other markup unless requested by the user.

Here is all of the information available to answer the user:
===
{chr(10).join(chunk_text)}
===

If the user asked for a search and there are no results, make sure to let the user know that you couldn't find anything,
and what they might be able to do to find the information they need.

END SYSTEM INSTRUCTIONS"""

        # OpenAI chat completion using older package format
        chat_completion = openai.ChatCompletion.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": query}
            ]
        )
        
        return chat_completion.choices[0].message['content']
        
    except Exception as e:
        return f"Error: {str(e)}"

# Create interface
demo = gr.Interface(
    fn=get_response,
    inputs=gr.Textbox(label="Enter your query"),
    outputs=gr.Textbox(label="Response", lines=15),
    title="Ragie + OpenAI Chatbot",
    description="Ask questions about your documents. The system will retrieve relevant information and generate a response."
)

if __name__ == "__main__":
    demo.launch()