Update app.py

app.py

@@ -1,61 +1,76 @@
-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
-
-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))
+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=8000)