|
@@ -5,6 +5,7 @@ from langchain_community.vectorstores import Chroma
|
|
|
from langchain_community.embeddings import OllamaEmbeddings
|
|
from langchain_community.embeddings import OllamaEmbeddings
|
|
|
from langchain_community.document_loaders import TextLoader
|
|
from langchain_community.document_loaders import TextLoader
|
|
|
from langchain_text_splitters import CharacterTextSplitter
|
|
from langchain_text_splitters import CharacterTextSplitter
|
|
|
|
|
+from langchain.memory import ConversationBufferWindowMemory
|
|
|
import os
|
|
import os
|
|
|
from dotenv import load_dotenv
|
|
from dotenv import load_dotenv
|
|
|
from tempfile import NamedTemporaryFile
|
|
from tempfile import NamedTemporaryFile
|
|
@@ -12,9 +13,6 @@ from tempfile import NamedTemporaryFile
|
|
|
# Load environment variables
|
|
# Load environment variables
|
|
|
load_dotenv()
|
|
load_dotenv()
|
|
|
|
|
|
|
|
-# Global dictionary for chat history
|
|
|
|
|
-chat_history = {}
|
|
|
|
|
-
|
|
|
|
|
# Initialize Ollama LLM and Embeddings
|
|
# Initialize Ollama LLM and Embeddings
|
|
|
llm = Ollama(model="tinyllama", temperature=0.7)
|
|
llm = Ollama(model="tinyllama", temperature=0.7)
|
|
|
embeddings = OllamaEmbeddings(model="tinyllama")
|
|
embeddings = OllamaEmbeddings(model="tinyllama")
|
|
@@ -42,59 +40,53 @@ def index_file(file_content: bytes, file_name: str):
|
|
|
os.unlink(temp_file_path)
|
|
os.unlink(temp_file_path)
|
|
|
|
|
|
|
|
# Define prompt templates
|
|
# Define prompt templates
|
|
|
-def get_prompt_with_history(session_id):
|
|
|
|
|
- history = chat_history.get(session_id, [])
|
|
|
|
|
- history_text = "\n".join([f"User: {msg['question']}\nAI: {msg['answer']}" for msg in history]) if history else "No previous conversation."
|
|
|
|
|
|
|
+def get_prompt_with_history(memory):
|
|
|
return PromptTemplate(
|
|
return PromptTemplate(
|
|
|
- input_variables=["question"],
|
|
|
|
|
- template=f"Previous conversation:\n{history_text}\n\nResponda à seguinte pergunta: {{question}}"
|
|
|
|
|
|
|
+ input_variables=["history", "question"],
|
|
|
|
|
+ template=f"Previous conversation:\n{{history}}\n\nResponda à seguinte pergunta: {{question}}"
|
|
|
)
|
|
)
|
|
|
|
|
|
|
|
-def get_prompt_with_history_and_docs(session_id, docs):
|
|
|
|
|
- history = chat_history.get(session_id, [])
|
|
|
|
|
- history_text = "\n".join([f"User: {msg['question']}\nAI: {msg['answer']}" for msg in history]) if history else "No previous conversation."
|
|
|
|
|
|
|
+def get_prompt_with_history_and_docs(memory, docs):
|
|
|
docs_text = "\n".join([f"Source: {doc.page_content}" for doc in docs]) if docs else "No relevant documents found."
|
|
docs_text = "\n".join([f"Source: {doc.page_content}" for doc in docs]) if docs else "No relevant documents found."
|
|
|
return PromptTemplate(
|
|
return PromptTemplate(
|
|
|
- input_variables=["question"],
|
|
|
|
|
- template=f"Previous conversation:\n{history_text}\n\nRelevant documents:\n{docs_text}\n\nResponda à seguinte pergunta usando as fontes relevantes e citando trechos como fontes: {{question}}"
|
|
|
|
|
|
|
+ input_variables=["history", "question"],
|
|
|
|
|
+ template=f"Previous conversation:\n{{history}}\n\nRelevant documents:\n{docs_text}\n\nResponda à seguinte pergunta usando as fontes relevantes e citando trechos como fontes: {{question}}"
|
|
|
)
|
|
)
|
|
|
|
|
|
|
|
def get_answer(session_id: str, question: str) -> str:
|
|
def get_answer(session_id: str, question: str) -> str:
|
|
|
- # Get or initialize chat history for this session
|
|
|
|
|
- if session_id not in chat_history:
|
|
|
|
|
- chat_history[session_id] = []
|
|
|
|
|
|
|
+ # Get or initialize memory for this session
|
|
|
|
|
+ memory = ConversationBufferWindowMemory(memory_key="history", input_key="question", k=3, session_id=session_id)
|
|
|
|
|
|
|
|
# Create chain with dynamic prompt including history
|
|
# Create chain with dynamic prompt including history
|
|
|
- prompt = get_prompt_with_history(session_id)
|
|
|
|
|
- chain = LLMChain(llm=llm, prompt=prompt)
|
|
|
|
|
|
|
+ prompt = get_prompt_with_history(memory)
|
|
|
|
|
+ chain = LLMChain(llm=llm, prompt=prompt, memory=memory)
|
|
|
|
|
|
|
|
# Get response
|
|
# Get response
|
|
|
response = chain.run(question=question)
|
|
response = chain.run(question=question)
|
|
|
response = response[:100] if len(response) > 100 else response # Truncate if needed
|
|
response = response[:100] if len(response) > 100 else response # Truncate if needed
|
|
|
|
|
|
|
|
- # Store the interaction in history
|
|
|
|
|
- chat_history[session_id].append({"question": question, "answer": response})
|
|
|
|
|
-
|
|
|
|
|
return response
|
|
return response
|
|
|
|
|
|
|
|
# RAG function for /ask endpoint
|
|
# RAG function for /ask endpoint
|
|
|
def ask_rag(session_id: str, question: str, file_content: bytes = None, file_name: str = None) -> dict:
|
|
def ask_rag(session_id: str, question: str, file_content: bytes = None, file_name: str = None) -> dict:
|
|
|
|
|
+ # Get or initialize memory for this session
|
|
|
|
|
+ memory = ConversationBufferWindowMemory(memory_key="history", input_key="question", k=3, session_id=session_id)
|
|
|
|
|
+
|
|
|
if file_content and file_name:
|
|
if file_content and file_name:
|
|
|
index_file(file_content, file_name)
|
|
index_file(file_content, file_name)
|
|
|
|
|
|
|
|
- if session_id not in chat_history:
|
|
|
|
|
- chat_history[session_id] = []
|
|
|
|
|
-
|
|
|
|
|
|
|
+ # Retrieve relevant documents
|
|
|
docs = vector_store.similarity_search(question, k=3)
|
|
docs = vector_store.similarity_search(question, k=3)
|
|
|
|
|
|
|
|
- prompt = get_prompt_with_history_and_docs(session_id, docs)
|
|
|
|
|
- chain = LLMChain(llm=llm, prompt=prompt)
|
|
|
|
|
|
|
+ # Create chain with dynamic prompt including history and docs
|
|
|
|
|
+ prompt = get_prompt_with_history_and_docs(memory, docs)
|
|
|
|
|
+ chain = LLMChain(llm=llm, prompt=prompt, memory=memory)
|
|
|
|
|
|
|
|
|
|
+ # Get response
|
|
|
response = chain.run(question=question)
|
|
response = chain.run(question=question)
|
|
|
response = response[:100] if len(response) > 100 else response
|
|
response = response[:100] if len(response) > 100 else response
|
|
|
|
|
|
|
|
- chat_history[session_id].append({"question": question, "answer": response})
|
|
|
|
|
-
|
|
|
|
|
|
|
+ # Prepare sources
|
|
|
sources = [doc.page_content for doc in docs]
|
|
sources = [doc.page_content for doc in docs]
|
|
|
|
|
|
|
|
return {"answer": response, "sources": sources}
|
|
return {"answer": response, "sources": sources}
|