import requests

API_ENDPOINT = "http://54.254.230.28:8888"

# function to call api
def call_api_stream(api_path, api_params):
    session = requests.Session()
    url = f"{API_ENDPOINT}/{api_path}"
    response = session.post(
        url, json=api_params, headers={"Content-Type": "application/json"},
        stream=True
    )
    return response

def call_api(api_path, api_params):
    session = requests.Session()
    url = f"{API_ENDPOINT}/{api_path}"
    response = session.post(
        url, json=api_params, headers={"Content-Type": "application/json"}
    )
    return response.json()

def api_rag_qa_chain_demo(openai_model_name, query, year, company_name):
    api_path = "qa/demo"
    api_params = {
        "openai_model_name": openai_model_name,
        "query": query,
        "year": year,
        "company_name": company_name,
    }
    return call_api_stream(api_path, api_params)

def api_rag_summ_chain_demo(openai_model_name, query, year, company_name, tone):
    api_path = "qa/waterfee"
    api_params = {
        "openai_model_name": openai_model_name,
        "query": query,
        "year": year,
        "company_name": company_name,
        "tone": tone,
    }
    return call_api_stream(api_path, api_params)