langchain_backend/main.py

import os
from langchain_backend.utils import create_prompt_llm_chain, create_retriever, getPDF
from langchain_backend import utils
from langchain.chains import create_retrieval_chain
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings

os.environ.get("OPENAI_API_KEY")

def get_llm_answer(system_prompt, user_prompt, pdf_url, model, embedding):
  if embedding == "gpt":
    embedding_object = OpenAIEmbeddings()
  else:
    embedding_object = HuggingFaceEmbeddings(model_name=embedding)

  vectorstore = Chroma(
      collection_name="documents",
      embedding_function=embedding_object
    )

  print('model: ', model)
  print('embedding: ', embedding)
  pages = []
  if pdf_url:
    pages = getPDF(pdf_url)
  else:
    pages = getPDF()
  retriever = create_retriever(pages, vectorstore)
  rag_chain = create_retrieval_chain(retriever, create_prompt_llm_chain(system_prompt, model))
  results = rag_chain.invoke({"input": user_prompt})
  print('allIds ARQUIVO MAIN: ',  utils.allIds)
  vectorstore.delete( utils.allIds)
  vectorstore.delete_collection()
  utils.allIds = []
  print('utils.allIds: ', utils.allIds)
  return results