cooney
06eb3c57ab
- Chart.js - pdf, csv 파일 업로드 후 데이터 정제하여 llm으로 처리 후 결과 도출 - sqlite로 데이터 저장 - ORM - SQLAlchemy
90 lines
2.6 KiB
Python
90 lines
2.6 KiB
Python
import os
|
|
|
|
from langchain_community.document_loaders import PyPDFLoader
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from langchain_community.vectorstores import FAISS
|
|
from langchain_core.runnables import RunnablePassthrough
|
|
from langchain_core.output_parsers import StrOutputParser
|
|
from langchain_core.prompts import ChatPromptTemplate
|
|
|
|
from backend.ai.embedding import watson_embedding
|
|
from backend.ai.llm import watson_llm
|
|
|
|
UPLOAD_PATH = "uploads"
|
|
|
|
def upload_document(file):
|
|
# file 저장
|
|
file_path = os.path.join(UPLOAD_PATH, file.filename)
|
|
|
|
with open(file_path, "wb") as f:
|
|
f.write(file.file.read())
|
|
|
|
# pdf 업로드 => 분할 => 인덱스 생성
|
|
# pdf 로드
|
|
loader = PyPDFLoader(file_path)
|
|
docs = loader.load()
|
|
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=30)
|
|
chunks = splitter.split_documents(docs)
|
|
|
|
faiss_store = FAISS.from_documents(documents=chunks, embedding=watson_embedding)
|
|
faiss_store.save_local("./db/vectorstore")
|
|
|
|
return {"message": "업로드 성공"}
|
|
|
|
# 질문 => 유사도 검색 => 문서 => llm 답변 생성
|
|
def rag_chat(question: str):
|
|
|
|
faiss_store = FAISS.load_local("./db/vectorstore", watson_embedding, allow_dangerous_deserialization=True)
|
|
|
|
retriever = faiss_store.as_retriever(search_kwargs={"k": 3})
|
|
|
|
### LLM
|
|
# 1. prompt
|
|
message = """\
|
|
당신은 PDF 기반 RAG AI 입니다.
|
|
다음 문서를 참고해서 질문에 답변하세요.
|
|
|
|
문서:
|
|
{context}
|
|
|
|
질문:
|
|
{question}
|
|
"""
|
|
|
|
rag_prompt = ChatPromptTemplate.from_template(message)
|
|
# 2. chain
|
|
chain = {"context": retriever, "question": RunnablePassthrough()} | rag_prompt | watson_llm | StrOutputParser()
|
|
|
|
# 3. answer
|
|
answer = chain.invoke(question)
|
|
return answer
|
|
|
|
async def rag_chat_stream(question: str):
|
|
|
|
faiss_store = FAISS.load_local("./db/vectorstore", watson_embedding, allow_dangerous_deserialization=True)
|
|
|
|
retriever = faiss_store.as_retriever(search_kwargs={"k": 3})
|
|
|
|
### LLM
|
|
# 1. prompt
|
|
message = """\
|
|
당신은 PDF 기반 RAG AI 입니다.
|
|
다음 문서를 참고해서 질문에 답변하세요.
|
|
|
|
문서:
|
|
{context}
|
|
|
|
질문:
|
|
{question}
|
|
"""
|
|
|
|
rag_prompt = ChatPromptTemplate.from_template(message)
|
|
# 2. chain
|
|
chain = {"context": retriever, "question": RunnablePassthrough()} | rag_prompt | watson_llm | StrOutputParser()
|
|
|
|
# 3. answer
|
|
# chain.stream() : 동기방식(요청 => 응답을 할 떄까지 기다리는 방식)
|
|
# chain.astream() : 비동기방식(다른 일 처리 가능)
|
|
async for chunk in chain.astream(question):
|
|
yield chunk |