FernAI / app.py
redfernstech's picture
Update app.py
6b539de verified
# import os
# import time
# from fastapi import FastAPI,Request
# from fastapi.responses import HTMLResponse
# from fastapi.staticfiles import StaticFiles
# from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate, Settings
# from llama_index.embeddings.huggingface import HuggingFaceEmbedding
# from pydantic import BaseModel
# from fastapi.responses import JSONResponse
# import uuid # for generating unique IDs
# import datetime
# from fastapi.middleware.cors import CORSMiddleware
# from fastapi.templating import Jinja2Templates
# from huggingface_hub import InferenceClient
# import json
# import re
# from gradio_client import Client
# from simple_salesforce import Salesforce, SalesforceLogin
# from llama_index.llms.huggingface import HuggingFaceLLM
# # from llama_index.llms.huggingface import HuggingFaceInferenceAPI
# # Define Pydantic model for incoming request body
# class MessageRequest(BaseModel):
# message: str
# repo_id = "meta-llama/Meta-Llama-3-8B-Instruct"
# llm_client = InferenceClient(
# model=repo_id,
# token=os.getenv("HF_TOKEN"),
# )
# os.environ["HF_TOKEN"] = os.getenv("HF_TOKEN")
# username = os.getenv("username")
# password = os.getenv("password")
# security_token = os.getenv("security_token")
# domain = os.getenv("domain")# Using sandbox environment
# session_id, sf_instance = SalesforceLogin(username=username, password=password, security_token=security_token, domain=domain)
# # Create Salesforce object
# sf = Salesforce(instance=sf_instance, session_id=session_id)
# app = FastAPI()
# @app.middleware("http")
# async def add_security_headers(request: Request, call_next):
# response = await call_next(request)
# response.headers["Content-Security-Policy"] = "frame-ancestors *; frame-src *; object-src *;"
# response.headers["X-Frame-Options"] = "ALLOWALL"
# return response
# # Allow CORS requests from any domain
# app.add_middleware(
# CORSMiddleware,
# allow_origins=["*"],
# allow_credentials=True,
# allow_methods=["*"],
# allow_headers=["*"],
# )
# @app.get("/favicon.ico")
# async def favicon():
# return HTMLResponse("") # or serve a real favicon if you have one
# app.mount("/static", StaticFiles(directory="static"), name="static")
# templates = Jinja2Templates(directory="static")
# # Configure Llama index settings
# Settings.llm = HuggingFaceLLM(
# model_name="meta-llama/Meta-Llama-3-8B-Instruct",
# tokenizer_name="meta-llama/Meta-Llama-3-8B-Instruct",
# context_window=3000,
# token=os.getenv("HF_TOKEN"),
# max_new_tokens=512,
# generate_kwargs={"temperature": 0.1},
# )
# Settings.embed_model = HuggingFaceEmbedding(
# model_name="BAAI/bge-small-en-v1.5"
# )
# PERSIST_DIR = "db"
# PDF_DIRECTORY = 'data'
# # Ensure directories exist
# os.makedirs(PDF_DIRECTORY, exist_ok=True)
# os.makedirs(PERSIST_DIR, exist_ok=True)
# chat_history = []
# current_chat_history = []
# def data_ingestion_from_directory():
# documents = SimpleDirectoryReader(PDF_DIRECTORY).load_data()
# storage_context = StorageContext.from_defaults()
# index = VectorStoreIndex.from_documents(documents)
# index.storage_context.persist(persist_dir=PERSIST_DIR)
# def initialize():
# start_time = time.time()
# data_ingestion_from_directory() # Process PDF ingestion at startup
# print(f"Data ingestion time: {time.time() - start_time} seconds")
# def split_name(full_name):
# # Split the name by spaces
# words = full_name.strip().split()
# # Logic for determining first name and last name
# if len(words) == 1:
# first_name = ''
# last_name = words[0]
# elif len(words) == 2:
# first_name = words[0]
# last_name = words[1]
# else:
# first_name = words[0]
# last_name = ' '.join(words[1:])
# return first_name, last_name
# initialize() # Run initialization tasks
# def handle_query(query):
# chat_text_qa_msgs = [
# (
# "user",
# """
# You are the Clara Redfernstech chatbot. Your goal is to provide accurate, professional, and helpful answers to user queries based on the company's data. Always ensure your responses are clear and concise. Give response within 10-15 words only
# {context_str}
# Question:
# {query_str}
# """
# )
# ]
# text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
# storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
# index = load_index_from_storage(storage_context)
# context_str = ""
# for past_query, response in reversed(current_chat_history):
# if past_query.strip():
# context_str += f"User asked: '{past_query}'\nBot answered: '{response}'\n"
# query_engine = index.as_query_engine(text_qa_template=text_qa_template, context_str=context_str)
# answer = query_engine.query(query)
# if hasattr(answer, 'response'):
# response=answer.response
# elif isinstance(answer, dict) and 'response' in answer:
# response =answer['response']
# else:
# response ="Sorry, I couldn't find an answer."
# current_chat_history.append((query, response))
# return response
# @app.get("/ch/{id}", response_class=HTMLResponse)
# async def load_chat(request: Request, id: str):
# return templates.TemplateResponse("index.html", {"request": request, "user_id": id})
# # Route to save chat history
# @app.post("/hist/")
# async def save_chat_history(history: dict):
# # Check if 'userId' is present in the incoming dictionary
# user_id = history.get('userId')
# print(user_id)
# # Ensure user_id is defined before proceeding
# if user_id is None:
# return {"error": "userId is required"}, 400
# # Construct the chat history string
# hist = ''.join([f"'{entry['sender']}: {entry['message']}'\n" for entry in history['history']])
# hist = "You are a Redfernstech summarize model. Your aim is to use this conversation to identify user interests solely based on that conversation: " + hist
# print(hist)
# # Get the summarized result from the client model
# result = hist
# try:
# sf.Lead.update(user_id, {'Description': result})
# except Exception as e:
# return {"error": f"Failed to update lead: {str(e)}"}, 500
# return {"summary": result, "message": "Chat history saved"}
# @app.post("/webhook")
# async def receive_form_data(request: Request):
# form_data = await request.json()
# # Log in to Salesforce
# first_name, last_name = split_name(form_data['name'])
# data = {
# 'FirstName': first_name,
# 'LastName': last_name,
# 'Description': 'hii', # Static description
# 'Company': form_data['company'], # Assuming company is available in form_data
# 'Phone': form_data['phone'].strip(), # Phone from form data
# 'Email': form_data['email'], # Email from form data
# }
# a=sf.Lead.create(data)
# # Generate a unique ID (for tracking user)
# unique_id = a['id']
# # Here you can do something with form_data like saving it to a database
# print("Received form data:", form_data)
# # Send back the unique id to the frontend
# return JSONResponse({"id": unique_id})
# @app.post("/chat/")
# async def chat(request: MessageRequest):
# message = request.message # Access the message from the request body
# response = handle_query(message) # Process the message
# message_data = {
# "sender": "User",
# "message": message,
# "response": response,
# "timestamp": datetime.datetime.now().isoformat()
# }
# chat_history.append(message_data)
# return {"response": response}
# @app.get("/")
# def read_root():
# return {"message": "Welcome to the API"}
import os
import datetime
import json
import logging
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from fastapi.templating import Jinja2Templates
from simple_salesforce import Salesforce, SalesforceLogin
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from llama_index.core import StorageContext, VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.core import load_index_from_storage
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.langchain import LangChainLLM # Added for Groq integration
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Define Pydantic model for incoming request body
class MessageRequest(BaseModel):
message: str
# Initialize FastAPI app
app = FastAPI()
# Allow CORS (restrict origins in production)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # TODO: Restrict to specific origins in production
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Mount static files and templates
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="static")
# Validate environment variables
required_env_vars = ["CHATGROQ_API_KEY", "username", "password", "security_token", "domain", "HF_TOKEN"]
for var in required_env_vars:
if not os.getenv(var):
logger.error(f"Environment variable {var} is not set")
raise ValueError(f"Environment variable {var} is not set")
# Initialize Groq model
GROQ_API_KEY = os.getenv("CHATGROQ_API_KEY")
GROQ_MODEL = "llama3-8b-8192"
try:
groq_llm = ChatGroq(
model_name=GROQ_MODEL,
api_key=GROQ_API_KEY,
temperature=0.1,
max_tokens=50
)
except Exception as e:
logger.error(f"Failed to initialize Groq model: {e}")
raise HTTPException(status_code=500, detail="Failed to initialize Groq model")
# Configure LlamaIndex settings
Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
Settings.llm = LangChainLLM(llm=groq_llm) # Use Groq LLM for LlamaIndex
# Salesforce credentials
username = os.getenv("username")
password = os.getenv("password")
security_token = os.getenv("security_token")
domain = os.getenv("domain") # e.g., 'test' for sandbox
# Initialize Salesforce connection
sf = None
try:
session_id, sf_instance = SalesforceLogin(
username=username, password=password, security_token=security_token, domain=domain
)
sf = Salesforce(instance=sf_instance, session_id=session_id)
logger.info("Salesforce connection established")
except Exception as e:
logger.warning(f"Failed to connect to Salesforce: {e}. Continuing without Salesforce integration.")
# Chat history
chat_history = []
current_chat_history = []
MAX_HISTORY_SIZE = 100
# Directories for data ingestion and storage
PDF_DIRECTORY = "data"
PERSIST_DIR = "db"
# Ensure directories exist
os.makedirs(PDF_DIRECTORY, exist_ok=True)
os.makedirs(PERSIST_DIR, exist_ok=True)
def data_ingestion_from_directory():
"""Ingest documents from PDF_DIRECTORY and store embeddings in PERSIST_DIR."""
try:
documents = SimpleDirectoryReader(PDF_DIRECTORY).load_data()
if not documents:
logger.warning("No documents found in PDF_DIRECTORY")
return False
storage_context = StorageContext.from_defaults()
index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
index.storage_context.persist(persist_dir=PERSIST_DIR)
logger.info("Data ingestion and embedding storage completed successfully")
return True
except Exception as e:
logger.error(f"Error during data ingestion: {e}")
raise HTTPException(status_code=500, detail=f"Data ingestion failed: {str(e)}")
def initialize():
"""Initialize the application by ingesting data and setting up embeddings."""
try:
if not data_ingestion_from_directory():
logger.info("No documents to ingest, proceeding with empty index")
except Exception as e:
logger.error(f"Initialization failed: {e}")
raise HTTPException(status_code=500, detail="Initialization failed")
# Run initialization
initialize()
def handle_query(query: str) -> str:
"""Handle user query by retrieving relevant documents and querying Groq LLM."""
# Prepare context from chat history
chat_context = ""
for past_query, response in reversed(current_chat_history[-10:]):
if past_query.strip():
chat_context += f"User: {past_query}\nBot: {response}\n"
# Load vector index and retrieve relevant documents
try:
storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
index = load_index_from_storage(storage_context)
query_engine = index.as_query_engine(similarity_top_k=2)
retrieved = query_engine.query(query)
doc_context = retrieved.response if hasattr(retrieved, 'response') else "No relevant information found."
logger.info(f"Retrieved context for query '{query}': {doc_context[:100]}...")
except Exception as e:
logger.error(f"Error retrieving documents: {e}")
doc_context = "Failed to retrieve relevant information."
# Construct prompt for Redferns Tech chatbot
prompt_template = ChatPromptTemplate.from_messages([
("system", """
You are Clara, a chatbot for Redferns Tech. Provide accurate, professional answers in 10-15 words.
Use the provided document context and chat history to inform your response.
If you don't know the answer, politely say: "I'm sorry, I don't have the information to answer that."
Document Context:
{doc_context}
Chat History:
{chat_context}
Question:
{query}
"""),
])
prompt = prompt_template.format(doc_context=doc_context, chat_context=chat_context, query=query)
# Query Groq model
try:
response = groq_llm.invoke(prompt)
response_text = response.content.strip()
if not response_text or response_text.lower() == "unknown":
response_text = "I'm sorry, I don't have the information to answer that."
except Exception as e:
logger.error(f"Error querying Groq API: {e}")
response_text = "I'm sorry, I don't have the information to answer that."
# Update chat history
if len(current_chat_history) >= MAX_HISTORY_SIZE:
current_chat_history.pop(0)
current_chat_history.append((query, response_text))
return response_text
@app.get("/ch/{id}", response_class=HTMLResponse)
async def load_chat(request: Request, id: str):
"""Serve the chat interface for a specific user ID."""
return templates.TemplateResponse("index.html", {"request": request, "user_id": id})
@app.post("/hist/")
async def save_chat_history(history: dict):
"""Save chat history to Salesforce."""
if not sf:
logger.error("Salesforce integration is disabled")
return JSONResponse({"error": "Salesforce integration is unavailable"}, status_code=503)
user_id = history.get('userId')
if not user_id:
logger.error("userId is missing in history request")
return JSONResponse({"error": "userId is required"}, status_code=400)
hist = ''.join([f"{entry['sender']}: {entry['message']}\n" for entry in history['history']])
summary_prompt = f"Summarize user interests from this conversation:\n{hist}"
try:
sf.Lead.update(user_id, {'Description': summary_prompt})
logger.info(f"Chat history updated for user {user_id}")
return {"summary": summary_prompt, "message": "Chat history saved"}
except Exception as e:
logger.error(f"Failed to update lead: {e}")
return JSONResponse({"error": f"Failed to update lead: {str(e)}"}, status_code=500)
@app.post("/webhook")
async def receive_form_data(request: Request):
"""Create a Salesforce lead from form data."""
if not sf:
logger.error("Salesforce integration is disabled")
return JSONResponse({"error": "Salesforce integration is unavailable"}, status_code=503)
try:
form_data = await request.json()
except json.JSONDecodeError:
logger.error("Invalid JSON in webhook request")
return JSONResponse({"error": "Invalid JSON"}, status_code=400)
first_name, last_name = split_name(form_data.get('name', ''))
data = {
'FirstName': first_name,
'LastName': last_name,
'Description': 'Lead created via webhook',
'Company': form_data.get('company', ''),
'Phone': form_data.get('phone', '').strip(),
'Email': form_data.get('email', ''),
}
try:
result = sf.Lead.create(data)
unique_id = result['id']
logger.info(f"Lead created with ID {unique_id}")
return JSONResponse({"id": unique_id})
except Exception as e:
logger.error(f"Failed to create lead: {e}")
return JSONResponse({"error": f"Failed to create lead: {str(e)}"}, status_code=500)
@app.post("/chat/")
async def chat(request: MessageRequest):
"""Handle chat messages and return responses."""
message = request.message.strip()
if not message:
return JSONResponse({"error": "Message cannot be empty"}, status_code=400)
response = handle_query(message)
message_data = {
"sender": "User",
"message": message,
"response": response,
"timestamp": datetime.datetime.now().isoformat()
}
if len(chat_history) >= MAX_HISTORY_SIZE:
chat_history.pop(0)
chat_history.append(message_data)
logger.info(f"Chat message processed: {message}")
return {"response": response}
@app.get("/health")
async def health_check():
"""Check the health of the application."""
try:
storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
index = load_index_from_storage(storage_context)
logger.info("Vector index loaded successfully")
return {"status": "healthy", "pdf_ingestion": "successful"}
except Exception as e:
logger.error(f"Health check failed: {e}")
return {"status": "unhealthy", "error": str(e)}
@app.get("/")
async def read_root():
"""Root endpoint for the API."""
return {"message": "Welcome to the Redferns Tech Chatbot API"}
def split_name(full_name: str) -> tuple:
"""Split a full name into first and last names."""
words = full_name.strip().split()
return (words[0], ' '.join(words[1:])) if words else ('', '')