Spaces:
Sleeping
Sleeping
import os | |
import getpass | |
import spacy | |
import pandas as pd | |
from typing import Optional, List, Dict, Any | |
import subprocess | |
from langchain.llms.base import LLM | |
from langchain.docstore.document import Document | |
from langchain.embeddings import HuggingFaceEmbeddings | |
from langchain.vectorstores import FAISS | |
from langchain.chains import RetrievalQA | |
from smolagents import CodeAgent, DuckDuckGoSearchTool, ManagedAgent, LiteLLMModel | |
from pydantic import BaseModel, Field, ValidationError, validator | |
from mistralai import Mistral | |
from langchain.prompts import PromptTemplate | |
# Import chains and tools | |
from classification_chain import get_classification_chain | |
from cleaner_chain import get_cleaner_chain | |
from refusal_chain import get_refusal_chain | |
from tailor_chain import get_tailor_chain | |
from prompts import classification_prompt, refusal_prompt, tailor_prompt | |
# Initialize Mistral API client | |
mistral_api_key = os.environ.get("MISTRAL_API_KEY") | |
client = Mistral(api_key=mistral_api_key) | |
# Initialize LiteLLM model for web search | |
pydantic_agent = LiteLLMModel(model_id="gemini/gemini-pro", api_key=os.environ.get("GEMINI_API_KEY")) | |
# Pydantic models for validation and type safety | |
class QueryInput(BaseModel): | |
query: str = Field(..., min_length=1, description="The input query string") | |
def check_query_is_string(cls, v): | |
if not isinstance(v, str): | |
raise ValueError("Query must be a valid string") | |
if v.strip() == "": | |
raise ValueError("Query cannot be empty or just whitespace") | |
return v.strip() | |
class ModerationResult(BaseModel): | |
is_safe: bool = Field(..., description="Whether the content is safe") | |
categories: Dict[str, bool] = Field(default_factory=dict, description="Detected content categories") | |
original_text: str = Field(..., description="The original input text") | |
# Load spaCy model for NER | |
def install_spacy_model(): | |
try: | |
spacy.load("en_core_web_sm") | |
print("spaCy model 'en_core_web_sm' is already installed.") | |
except OSError: | |
print("Downloading spaCy model 'en_core_web_sm'...") | |
subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"], check=True) | |
print("spaCy model 'en_core_web_sm' downloaded successfully.") | |
install_spacy_model() | |
nlp = spacy.load("en_core_web_sm") | |
def extract_main_topic(query: str) -> str: | |
try: | |
query_input = QueryInput(query=query) | |
doc = nlp(query_input.query) | |
main_topic = None | |
for ent in doc.ents: | |
if ent.label_ in ["ORG", "PRODUCT", "PERSON", "GPE", "TIME"]: | |
main_topic = ent.text | |
break | |
if not main_topic: | |
for token in doc: | |
if token.pos_ in ["NOUN", "PROPN"]: | |
main_topic = token.text | |
break | |
return main_topic if main_topic else "this topic" | |
except Exception as e: | |
print(f"Error extracting main topic: {e}") | |
return "this topic" | |
def moderate_text(query: str) -> ModerationResult: | |
try: | |
query_input = QueryInput(query=query) | |
response = client.classifiers.moderate_chat( | |
model="mistral-moderation-latest", | |
inputs=[{"role": "user", "content": query_input.query}] | |
) | |
is_safe = True | |
categories = {} | |
if hasattr(response, 'results') and response.results: | |
categories = { | |
"violence": response.results[0].categories.get("violence_and_threats", False), | |
"hate": response.results[0].categories.get("hate_and_discrimination", False), | |
"dangerous": response.results[0].categories.get("dangerous_and_criminal_content", False), | |
"selfharm": response.results[0].categories.get("selfharm", False) | |
} | |
is_safe = not any(categories.values()) | |
return ModerationResult( | |
is_safe=is_safe, | |
categories=categories, | |
original_text=query_input.query | |
) | |
except ValidationError as e: | |
raise ValueError(f"Input validation failed: {str(e)}") | |
except Exception as e: | |
raise RuntimeError(f"Moderation failed: {str(e)}") | |
def classify_query(query: str) -> str: | |
try: | |
query_input = QueryInput(query=query) | |
wellness_keywords = ["box breathing", "meditation", "yoga", "mindfulness", "breathing exercises"] | |
if any(keyword in query_input.query.lower() for keyword in wellness_keywords): | |
return "Wellness" | |
class_result = classification_chain.invoke({"query": query_input.query}) | |
classification = class_result.get("text", "").strip() | |
return classification if classification != "" else "OutOfScope" | |
except ValidationError as e: | |
raise ValueError(f"Classification input validation failed: {str(e)}") | |
except Exception as e: | |
raise RuntimeError(f"Classification failed: {str(e)}") | |
def build_or_load_vectorstore(csv_path: str, store_dir: str) -> FAISS: | |
try: | |
if os.path.exists(store_dir): | |
print(f"Loading existing FAISS store from '{store_dir}'") | |
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/multi-qa-mpnet-base-dot-v1") | |
return FAISS.load_local(store_dir, embeddings) | |
print(f"Building new FAISS store from CSV: {csv_path}") | |
df = pd.read_csv(csv_path) | |
df = df.loc[:, ~df.columns.str.contains('^Unnamed')] | |
df.columns = df.columns.str.strip() | |
docs = [ | |
Document(page_content=str(row["Answers"]), metadata={"question": str(row["Question"])}) | |
for _, row in df.iterrows() | |
] | |
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/multi-qa-mpnet-base-dot-v1") | |
vectorstore = FAISS.from_documents(docs, embedding=embeddings) | |
vectorstore.save_local(store_dir) | |
return vectorstore | |
except Exception as e: | |
raise RuntimeError(f"Error building/loading vector store: {str(e)}") | |
def build_rag_chain(llm_model: LiteLLMModel, vectorstore: FAISS) -> RetrievalQA: | |
class GeminiLangChainLLM(LLM): | |
def _call(self, prompt: str, stop: Optional[list] = None, **kwargs) -> str: | |
messages = [{"role": "user", "content": prompt}] | |
return llm_model(messages, stop_sequences=stop) | |
def _llm_type(self) -> str: | |
return "custom_gemini" | |
try: | |
retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3}) | |
gemini_as_llm = GeminiLangChainLLM() | |
return RetrievalQA.from_chain_type( | |
llm=gemini_as_llm, | |
chain_type="stuff", | |
retriever=retriever, | |
return_source_documents=True | |
) | |
except Exception as e: | |
raise RuntimeError(f"Error building RAG chain: {str(e)}") | |
def run_pipeline(query: str) -> str: | |
try: | |
query = sanitize_message(query) | |
moderation_result = moderate_text(query) | |
if not moderation_result.is_safe: | |
return "Sorry, this query contains harmful or inappropriate content." | |
classification = classify_query(moderation_result.original_text) | |
if classification == "OutOfScope": | |
refusal_text = refusal_chain.run({"topic": "this topic"}) | |
return tailor_chain.run({"response": refusal_text}).strip() | |
if classification == "Wellness": | |
rag_result = wellness_rag_chain({"query": moderation_result.original_text}) | |
csv_answer = rag_result["result"].strip() | |
web_answer = "" if csv_answer else do_web_search(moderation_result.original_text) | |
final_merged = merge_responses(csv_answer, web_answer) | |
return tailor_chain.run({"response": final_merged}).strip() | |
if classification == "Brand": | |
rag_result = brand_rag_chain({"query": moderation_result.original_text}) | |
csv_answer = rag_result["result"].strip() | |
final_merged = merge_responses(csv_answer, "") | |
return tailor_chain.run({"response": final_merged}).strip() | |
refusal_text = refusal_chain.run({"topic": "this topic"}) | |
return tailor_chain.run({"response": refusal_text}).strip() | |
# Initialize chains and vectorstores | |
# Initialize chains and vectorstores | |
try: | |
classification_chain = get_classification_chain() | |
refusal_chain = get_refusal_chain() | |
tailor_chain = get_tailor_chain() | |
cleaner_chain = get_cleaner_chain() | |
wellness_csv = "AIChatbot.csv" | |
brand_csv = "BrandAI.csv" | |
wellness_store_dir = "faiss_wellness_store" | |
brand_store_dir = "faiss_brand_store" | |
wellness_vectorstore = build_or_load_vectorstore(wellness_csv, wellness_store_dir) | |
brand_vectorstore = build_or_load_vectorstore(brand_csv, brand_store_dir) | |
gemini_llm = LiteLLMModel(model_id="gemini/gemini-pro", api_key=os.environ.get("GEMINI_API_KEY")) | |
wellness_rag_chain = build_rag_chain(gemini_llm, wellness_vectorstore) | |
brand_rag_chain = build_rag_chain(gemini_llm, brand_vectorstore) | |
print("Pipeline initialized successfully!") | |
except Exception as e: | |
print(f"Error initializing pipeline: {str(e)}") | |
def run_with_chain(query: str) -> str: | |
return run_pipeline(query) | |