Spaces:
Sleeping
Sleeping
from flask import Flask, request | |
import os | |
from langchain.vectorstores import Chroma | |
from langchain.document_loaders import PyPDFLoader | |
from langchain.embeddings import HuggingFaceEmbeddings | |
from langchain.text_splitter import RecursiveCharacterTextSplitter | |
import requests | |
from twilio.rest import Client | |
# Flask app | |
app = Flask(__name__) | |
# ChromaDB path | |
CHROMA_PATH = '/code/chroma_db' | |
if not os.path.exists(CHROMA_PATH): | |
os.makedirs(CHROMA_PATH) | |
from ai71 import AI71 | |
def generate_response(query, chat_history): | |
response = '' | |
try: | |
ai71_client = AI71(api_key=AI71_API_KEY) | |
chat_completion = ai71_client.chat.completions.create( | |
model="tiiuae/falcon-180b-chat", | |
messages=[ | |
{"role": "system", "content": "You are the best agricultural assistant. Remember to give a response in not more than 2 sentences."}, | |
{"role": "user", "content": f"Answer the query based on history {chat_history}: {query}"} | |
], | |
stream=True | |
) | |
for chunk in chat_completion: | |
if chunk.choices[0].delta.content: | |
response += chunk.choices[0].delta.content | |
# Clean up response text | |
response = response.replace("###", '').replace('\nUser:', '') | |
except Exception as e: | |
print(f"Error generating response: {e}") | |
response = "An error occurred while generating the response." | |
return response | |
# Initialize ChromaDB | |
def initialize_chroma(): | |
try: | |
embedding_function = HuggingFaceEmbeddings() | |
db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function) | |
# Perform an initial operation to ensure the database is correctly initialized | |
db.similarity_search_with_score("test query", k=1) | |
print("Chroma initialized successfully.") | |
except Exception as e: | |
print(f"Error initializing Chroma: {e}") | |
initialize_chroma() | |
# Set AI71 API key | |
AI71_API_KEY = os.environ.get('AI71_API_KEY') | |
account_sid = os.environ.get('TWILIO_ACCOUNT_SID') | |
auth_token = os.environ.get('TWILIO_AUTH_TOKEN') | |
client = Client(account_sid, auth_token) | |
from_whatsapp_number = 'whatsapp:+14155238886' | |
# Download file utility | |
def download_file(url, ext): | |
local_filename = f'/code/uploaded_file{ext}' | |
with requests.get(url, stream=True) as r: | |
with open(local_filename, 'wb') as f: | |
for chunk in r.iter_content(chunk_size=8192): | |
f.write(chunk) | |
return local_filename | |
# Process PDF and return text | |
import fitz # PyMuPDF | |
def extract_text_from_pdf(pdf_filepath): | |
text = '' | |
try: | |
pdf_document = fitz.open(pdf_filepath) | |
for page_num in range(len(pdf_document)): | |
page = pdf_document.load_page(page_num) | |
text += page.get_text() | |
pdf_document.close() | |
except Exception as e: | |
print(f"Error extracting text from PDF: {e}") | |
return None | |
return text | |
def query_rag(query_text: str, chat_history): | |
try: | |
embedding_function = HuggingFaceEmbeddings() | |
db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function) | |
results = db.similarity_search_with_score(query_text, k=5) | |
if not results: | |
return "Sorry, I couldn't find any relevant information." | |
context_text = "\n\n---\n\n".join([doc.page_content for doc, _score in results]) | |
prompt = f"Context:\n{context_text}\n\nQuestion:\n{query_text}" | |
response = generate_response(prompt, chat_history) | |
return response | |
except Exception as e: | |
print(f"Error querying RAG system: {e}") | |
return "An error occurred while querying the RAG system." | |
def save_pdf_and_update_database(pdf_filepath): | |
try: | |
text = extract_text_from_pdf(pdf_filepath) | |
if not text: | |
print("Error extracting text from PDF.") | |
return | |
text_splitter = RecursiveCharacterTextSplitter( | |
chunk_size=800, | |
chunk_overlap=80, | |
length_function=len, | |
is_separator_regex=False, | |
) | |
chunks = text_splitter.split_text(text) | |
embedding_function = HuggingFaceEmbeddings() | |
db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function) | |
db.add_documents(chunks) | |
db.persist() | |
print("PDF processed and data updated in Chroma.") | |
except Exception as e: | |
print(f"Error processing PDF: {e}") | |
# Flask route to handle WhatsApp webhook | |
def whatsapp_webhook(): | |
incoming_msg = request.values.get('Body', '').lower() | |
sender = request.values.get('From') | |
num_media = int(request.values.get('NumMedia', 0)) | |
chat_history = [] # You need to handle chat history appropriately | |
if num_media > 0: | |
media_url = request.values.get('MediaUrl0') | |
content_type = request.values.get('MediaContentType0') | |
if content_type == 'application/pdf': | |
filepath = download_file(media_url, ".pdf") | |
save_pdf_and_update_database(filepath) | |
response_text = "PDF has been processed. You can now ask questions related to its content." | |
else: | |
response_text = "Unsupported file type. Please upload a PDF document." | |
else: | |
# Use RAG to generate a response based on the query | |
response_text = query_rag(incoming_msg, chat_history) | |
# Send the response back to the sender | |
send_message(sender, response_text) | |
return '', 204 | |
# Function to send message | |
def send_message(to, body): | |
try: | |
message = client.messages.create( | |
from_=from_whatsapp_number, | |
body=body, | |
to=to | |
) | |
print(f"Message sent with SID: {message.sid}") | |
except Exception as e: | |
print(f"Error sending message: {e}") | |
def send_initial_message(to_number): | |
send_message( | |
f'whatsapp:{to_number}', | |
'Welcome to the Agri AI Chatbot! How can I assist you today? You can send an image with "pest" or "disease" to classify it.' | |
) | |
if __name__ == "__main__": | |
send_initial_message('919080522395') | |
send_initial_message('916382792828') | |
app.run(host='0.0.0.0', port=7860) |