import os
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
from langchain_community.vectorstores import Chroma
from langchain.prompts import ChatPromptTemplate
from get_embedding_function import get_embedding_function
from langchain_groq import ChatGroq
import chainlit as cl

app = FastAPI()

# Configurar variáveis de ambiente
os.environ["OPENAI_API_BASE"] = 'https://api.groq.com/openai/v1'
os.environ["OPENAI_MODEL_NAME"] = 'llama3-8b-8192'
os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")

CHROMA_PATH = "chroma"

PROMPT_TEMPLATE = """
You are 'Vasu', an experienced professor with extensive knowledge in Cryptocurrency, Artificial Intelligence, and related projects.
Provide relevant 'Links' "http://", but include links only when they are particularly useful for understanding the response.
Answer the question based solely on the following context: {context}

Based on the above context, answer the question: {question}.
"""

class QueryRequest(BaseModel):
    query: str

class QueryResponse(BaseModel):
    response: str
    sources: List[str]

def query_rag(query_text: str):
    # Configurar o modelo Groq
    chat_groq = ChatGroq(temperature=0, model_name="llama3-8b-8192")

    # Preparar o DB
    embedding_function = get_embedding_function()
    db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)

    # Buscar no DB
    results = db.similarity_search_with_score(query_text, k=10)

    context_text = "\n\n---\n\n".join([doc.page_content for doc, _score in results])
    prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
    prompt = prompt_template.format(context=context_text, question=query_text)

    # Obter a resposta usando Groq
    response_text = chat_groq.invoke(prompt).content

    sources = [doc.metadata.get("id", None) for doc, _score in results]
    return response_text, sources

@app.post("/query", response_model=QueryResponse)
async def query_api(request: QueryRequest):
    try:
        response_text, sources = query_rag(request.query)
        return QueryResponse(response=response_text, sources=sources)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@cl.on_message
async def chainlit_main(message: cl.Message):
    query_text = message.content  # Obter a mensagem do usuário a partir do Chainlit
    response_text = query_rag(query_text)
    
    # Enviar a resposta de volta para o Chainlit
    await cl.Message(
        content=f"{response_text}",
    ).send()

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)