Spaces:

Eng-Musa
/

MusAI

Sleeping

MusAI / clients /groq_client.py

remove groq/compound-mini

69a7ed3 29 days ago

11.6 kB

import os
import json
from uuid import uuid4
from groq import Groq
from langchain_core.documents import Document

# CHANGED: Using local HuggingFaceEmbeddings instead of online API
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from dotenv import load_dotenv
import random
import shutil
from optimized_quiz import OPTIMIZED_QUESTIONS
from chat_resources import *
from datetime import datetime, timedelta

# Load main env
load_dotenv(dotenv_path=".env")

# Load secrets env (this can override duplicates from .env)
load_dotenv(dotenv_path=".secrets.env", override=True)

# Config
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
HF_API_KEY = os.getenv("HF_API_KEY")
DATA_PATH = "data.json"
CHROMA_PATH = "chroma_db"

TEMPERATURE = float(os.getenv("G_TEMPERATURE", 0.7))
MAX_TOKENS = int(os.getenv("G_MAX_TOKENS", 400))
RETRIEVE_K = int(os.getenv("G_RETRIEVE_K", 3))
TOP_P = float(os.getenv("G_TOP_P", 1.0))
MAX_CONVERSATION_HISTORY = int(os.getenv("G_MAX_CONVERSATION_HISTORY", 5))
MMR = str(os.getenv("MMR", "mmr"))
G_FETCH_K = int(os.getenv("G_FETCH_K", 20))
LAMBDA_MULT = float(os.getenv("LAMBDA_MULT", 0.5))


class GroqClient:
    def __init__(self):
        self._sessions = {}  # {ip: {'shown': set(), 'last_activity': datetime}}
        self.SESSION_TIMEOUT = timedelta(minutes=30)

        self.documents = self.load_json_data(DATA_PATH)
        if not self.documents:
            raise RuntimeError("No data loaded")

        self.vector_store = self.init_vector_store(self.documents)

        self.retriever = self.vector_store.as_retriever(
            search_type=MMR,  # Use Maximal Marginal Relevance
            search_kwargs={
                "k": RETRIEVE_K,  # Final number of docs to return
                "fetch_k": G_FETCH_K,  # Number of docs to initially fetch before filtering for diversity
                "lambda_mult": LAMBDA_MULT,  # Balance between relevance (1.0) and diversity (0.0)
            },
        )

        if not GROQ_API_KEY:
            raise RuntimeError("GROQ_API_KEY not found in environment")

        self.client = Groq(api_key=GROQ_API_KEY)

        self.SYSTEM_MESSSAGE = SYSTEM_MESSSAGE

        self.PROMPT_TEMPLATE = PROMPT_TEMPLATE

        self.BLACKLIST = BLACKLIST

    def load_json_data(self, path):
        try:
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)

            documents = []

            if "qa" in data:
                for item in data["qa"]:
                    text = f"Q: {item['question']}\nA: {item['answer']}"
                    documents.append(
                        Document(
                            page_content=text,
                            metadata={
                                "id": item.get("id", str(uuid4())),
                                "category": item.get("category", "QA"),
                            },
                        )
                    )

            if "chunks" in data:
                for item in data["chunks"]:
                    documents.append(
                        Document(
                            page_content=item["chunk"],
                            metadata={
                                "id": item.get("id", str(uuid4())),
                                "category": "Chunk",
                            },
                        )
                    )

            return documents

        except Exception as e:
            print(f"Error loading JSON data: {e}")
            return []

    def init_vector_store(self, documents):
        # CHANGED: Using local HuggingFaceEmbeddings instead of online API
        embeddings_model = HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-MiniLM-L6-v2",
            model_kwargs={"device": "cpu"},
            encode_kwargs={"normalize_embeddings": True},
        )

        # Clear old data to avoid duplicates
        if os.path.exists(CHROMA_PATH):
            shutil.rmtree(CHROMA_PATH, ignore_errors=True)

        uuids = [str(uuid4()) for _ in documents]

        vector_store = Chroma(
            collection_name="user_data",
            embedding_function=embeddings_model,
            persist_directory=CHROMA_PATH,
        )

        vector_store.add_documents(documents=documents, ids=uuids)
        return vector_store

    def handle_unknown_query(self):
        return random.choice(FALLBACK_RESPONSES)

    # def get_next_questions(self):
    #     return random.sample(OPTIMIZED_QUESTIONS, 3)
    def cleanup_expired_sessions(self):
        """Remove expired sessions to avoid memory overload."""
        current_time = datetime.now()
        expired_ips = [
            ip
            for ip, data in self._sessions.items()
            if current_time - data["last_activity"] > self.SESSION_TIMEOUT
        ]
        for ip in expired_ips:
            del self._sessions[ip]

    def get_next_questions(self, ip: str):
        """Return 3 non-repeated random questions within the session."""
        self.cleanup_expired_sessions()

        current_time = datetime.now()
        if ip not in self._sessions:
            self._sessions[ip] = {"shown": set(), "last_activity": current_time}
        else:
            if (
                current_time - self._sessions[ip]["last_activity"]
                > self.SESSION_TIMEOUT
            ):
                self._sessions[ip]["shown"].clear()
            self._sessions[ip]["last_activity"] = current_time

        shown = self._sessions[ip]["shown"]
        remaining = [q for q in OPTIMIZED_QUESTIONS if q not in shown]

        if len(remaining) < 3:
            shown.clear()
            remaining = OPTIMIZED_QUESTIONS[:]

        selected = random.sample(remaining, 3)
        shown.update(selected)
        return selected

    # ---------------MAIN-----------------

    # Non-streaming ask method for backwards compatibility
    # ============================================
    def ask(self, raw_query: str) -> tuple[str, str | None]:
        if (
            not raw_query
            or raw_query is None
            or raw_query == ""
            or len(raw_query) > 1000
        ):
            return "Please provide a valid query under 1,000 characters.", None

        if raw_query.lower() in GREETINGS_TRIGGERS:
            return random.choice(GREETINGS), None

        try:
            docs = self.retriever.invoke(raw_query)
        except Exception as e:
            return f"Error retrieving documents: {e}", None

        if not docs:
            return self.handle_unknown_query(), None

        context = "\n".join([d.page_content for d in docs])
        fallback = self.handle_unknown_query()
        prompt = self.PROMPT_TEMPLATE.format(
            context=context, question=raw_query, fallback_response=fallback
        )

        messages = [
            {
                "role": "system",
                "content": self.SYSTEM_MESSSAGE,
            },
        ] + [
            {"role": "user", "content": prompt},
        ]

        # Try multiple models with fallback
        models_to_try = [
            "openai/gpt-oss-120b",
            "openai/gpt-oss-20b",
            "compound-beta-mini",
            "openai/gpt-oss-safeguard-20b",
            
        ]

        random.shuffle(models_to_try)

        for model in models_to_try:
            try:
                print(model)
                completion = self.client.chat.completions.create(
                    model=model,
                    messages=messages,
                    temperature=TEMPERATURE,
                    max_completion_tokens=MAX_TOKENS,
                    top_p=TOP_P,
                    stream=False,
                )
                response = completion.choices[0].message.content
                if response and response.strip():
                    return response.strip(), model
                else:
                    continue  # Try next model

            except Exception as e:
                # Check if it's a rate limit error
                if "rate_limit_exceeded" in str(e) or "429" in str(e):
                    # print(f"Rate limit hit for model {model}, trying fallback...")
                    continue
                else:
                    # For other errors, return immediately
                    return f"Error while calling LLM: {e}", None

        # If all models fail
        return (
            "I'm temporarily experiencing high demand. Please try again in a few minutes or rephrase your question.",
            None,
        )

    # Streaming ask method (Generator)
    # ============================================
    def ask_stream(self, raw_query: str):
        """Generator function that yields response chunks in real-time"""

        # Validation
        if (
            not raw_query
            or raw_query is None
            or raw_query == ""
            or len(raw_query) > 1000
        ):
            yield "Please provide a valid query under 1,000 characters."
            return

        if raw_query.lower() in GREETINGS_TRIGGERS:
            yield random.choice(GREETINGS)
            return

        # Retrieve documents
        try:
            docs = self.retriever.invoke(raw_query)
        except Exception as e:
            yield f"Error retrieving documents: {e}"
            return

        if not docs:
            yield self.handle_unknown_query()
            return

        # Prepare context and prompt
        context = "\n".join([d.page_content for d in docs])
        fallback = self.handle_unknown_query()
        prompt = self.PROMPT_TEMPLATE.format(
            context=context, question=raw_query, fallback_response=fallback
        )

        messages = [
            {
                "role": "system",
                "content": self.SYSTEM_MESSSAGE,
            },
            {"role": "user", "content": prompt},
        ]

        # Try multiple models with fallback
        models_to_try = [
            "openai/gpt-oss-120b",
            "openai/gpt-oss-20b",
            "compound-beta-mini",
            "llama-3.1-8b-instant",
            "llama-3.3-70b-versatile",
        ]

        random.shuffle(models_to_try)

        for model in models_to_try:
            try:
                completion = self.client.chat.completions.create(
                    model=model,
                    messages=messages,
                    temperature=TEMPERATURE,
                    max_completion_tokens=MAX_TOKENS,
                    top_p=TOP_P,
                    stream=True,  # Enable streaming
                )

                # Stream chunks as they arrive
                has_content = False
                for chunk in completion:
                    if chunk.choices[0].delta.content:
                        has_content = True
                        yield chunk.choices[0].delta.content

                    # Check if stream finished
                    if chunk.choices[0].finish_reason:
                        break

                # If we got content, we're done
                if has_content:
                    return

            except Exception as e:
                # Check if it's a rate limit error
                if "rate_limit_exceeded" in str(e) or "429" in str(e):
                    continue  # Try next model
                else:
                    yield f"Error while calling LLM: {e}"
                    return

        # If all models fail
        yield "I'm temporarily experiencing high demand. Please try again in a few minutes or rephrase your question."

	import os
	import json
	from uuid import uuid4
	from groq import Groq
	from langchain_core.documents import Document

	# CHANGED: Using local HuggingFaceEmbeddings instead of online API
	from langchain_huggingface import HuggingFaceEmbeddings
	from langchain_chroma import Chroma
	from dotenv import load_dotenv
	import random
	import shutil
	from optimized_quiz import OPTIMIZED_QUESTIONS
	from chat_resources import *
	from datetime import datetime, timedelta

	# Load main env
	load_dotenv(dotenv_path=".env")

	# Load secrets env (this can override duplicates from .env)
	load_dotenv(dotenv_path=".secrets.env", override=True)

	# Config
	GROQ_API_KEY = os.getenv("GROQ_API_KEY")
	HF_API_KEY = os.getenv("HF_API_KEY")
	DATA_PATH = "data.json"
	CHROMA_PATH = "chroma_db"

	TEMPERATURE = float(os.getenv("G_TEMPERATURE", 0.7))
	MAX_TOKENS = int(os.getenv("G_MAX_TOKENS", 400))
	RETRIEVE_K = int(os.getenv("G_RETRIEVE_K", 3))
	TOP_P = float(os.getenv("G_TOP_P", 1.0))
	MAX_CONVERSATION_HISTORY = int(os.getenv("G_MAX_CONVERSATION_HISTORY", 5))
	MMR = str(os.getenv("MMR", "mmr"))
	G_FETCH_K = int(os.getenv("G_FETCH_K", 20))
	LAMBDA_MULT = float(os.getenv("LAMBDA_MULT", 0.5))


	class GroqClient:
	def __init__(self):
	self._sessions = {} # {ip: {'shown': set(), 'last_activity': datetime}}
	self.SESSION_TIMEOUT = timedelta(minutes=30)

	self.documents = self.load_json_data(DATA_PATH)
	if not self.documents:
	raise RuntimeError("No data loaded")

	self.vector_store = self.init_vector_store(self.documents)

	self.retriever = self.vector_store.as_retriever(
	search_type=MMR, # Use Maximal Marginal Relevance
	search_kwargs={
	"k": RETRIEVE_K, # Final number of docs to return
	"fetch_k": G_FETCH_K, # Number of docs to initially fetch before filtering for diversity
	"lambda_mult": LAMBDA_MULT, # Balance between relevance (1.0) and diversity (0.0)
	},
	)

	if not GROQ_API_KEY:
	raise RuntimeError("GROQ_API_KEY not found in environment")

	self.client = Groq(api_key=GROQ_API_KEY)

	self.SYSTEM_MESSSAGE = SYSTEM_MESSSAGE

	self.PROMPT_TEMPLATE = PROMPT_TEMPLATE

	self.BLACKLIST = BLACKLIST

	def load_json_data(self, path):
	try:
	with open(path, "r", encoding="utf-8") as f:
	data = json.load(f)

	documents = []

	if "qa" in data:
	for item in data["qa"]:
	text = f"Q: {item['question']}\nA: {item['answer']}"
	documents.append(
	Document(
	page_content=text,
	metadata={
	"id": item.get("id", str(uuid4())),
	"category": item.get("category", "QA"),
	},
	)
	)

	if "chunks" in data:
	for item in data["chunks"]:
	documents.append(
	Document(
	page_content=item["chunk"],
	metadata={
	"id": item.get("id", str(uuid4())),
	"category": "Chunk",
	},
	)
	)

	return documents

	except Exception as e:
	print(f"Error loading JSON data: {e}")
	return []

	def init_vector_store(self, documents):
	# CHANGED: Using local HuggingFaceEmbeddings instead of online API
	embeddings_model = HuggingFaceEmbeddings(
	model_name="sentence-transformers/all-MiniLM-L6-v2",
	model_kwargs={"device": "cpu"},
	encode_kwargs={"normalize_embeddings": True},
	)

	# Clear old data to avoid duplicates
	if os.path.exists(CHROMA_PATH):
	shutil.rmtree(CHROMA_PATH, ignore_errors=True)

	uuids = [str(uuid4()) for _ in documents]

	vector_store = Chroma(
	collection_name="user_data",
	embedding_function=embeddings_model,
	persist_directory=CHROMA_PATH,
	)

	vector_store.add_documents(documents=documents, ids=uuids)
	return vector_store

	def handle_unknown_query(self):
	return random.choice(FALLBACK_RESPONSES)

	# def get_next_questions(self):
	# return random.sample(OPTIMIZED_QUESTIONS, 3)
	def cleanup_expired_sessions(self):
	"""Remove expired sessions to avoid memory overload."""
	current_time = datetime.now()
	expired_ips = [
	ip
	for ip, data in self._sessions.items()
	if current_time - data["last_activity"] > self.SESSION_TIMEOUT
	]
	for ip in expired_ips:
	del self._sessions[ip]

	def get_next_questions(self, ip: str):
	"""Return 3 non-repeated random questions within the session."""
	self.cleanup_expired_sessions()

	current_time = datetime.now()
	if ip not in self._sessions:
	self._sessions[ip] = {"shown": set(), "last_activity": current_time}
	else:
	if (
	current_time - self._sessions[ip]["last_activity"]
	> self.SESSION_TIMEOUT
	):
	self._sessions[ip]["shown"].clear()
	self._sessions[ip]["last_activity"] = current_time

	shown = self._sessions[ip]["shown"]
	remaining = [q for q in OPTIMIZED_QUESTIONS if q not in shown]

	if len(remaining) < 3:
	shown.clear()
	remaining = OPTIMIZED_QUESTIONS[:]

	selected = random.sample(remaining, 3)
	shown.update(selected)
	return selected

	# ---------------MAIN-----------------

	# Non-streaming ask method for backwards compatibility
	# ============================================
	def ask(self, raw_query: str) -> tuple[str, str \| None]:
	if (
	not raw_query
	or raw_query is None
	or raw_query == ""
	or len(raw_query) > 1000
	):
	return "Please provide a valid query under 1,000 characters.", None

	if raw_query.lower() in GREETINGS_TRIGGERS:
	return random.choice(GREETINGS), None

	try:
	docs = self.retriever.invoke(raw_query)
	except Exception as e:
	return f"Error retrieving documents: {e}", None

	if not docs:
	return self.handle_unknown_query(), None

	context = "\n".join([d.page_content for d in docs])
	fallback = self.handle_unknown_query()
	prompt = self.PROMPT_TEMPLATE.format(
	context=context, question=raw_query, fallback_response=fallback
	)

	messages = [
	{
	"role": "system",
	"content": self.SYSTEM_MESSSAGE,
	},
	] + [
	{"role": "user", "content": prompt},
	]

	# Try multiple models with fallback
	models_to_try = [
	"openai/gpt-oss-120b",
	"openai/gpt-oss-20b",
	"compound-beta-mini",
	"openai/gpt-oss-safeguard-20b",

	]

	random.shuffle(models_to_try)

	for model in models_to_try:
	try:
	print(model)
	completion = self.client.chat.completions.create(
	model=model,
	messages=messages,
	temperature=TEMPERATURE,
	max_completion_tokens=MAX_TOKENS,
	top_p=TOP_P,
	stream=False,
	)
	response = completion.choices[0].message.content
	if response and response.strip():
	return response.strip(), model
	else:
	continue # Try next model

	except Exception as e:
	# Check if it's a rate limit error
	if "rate_limit_exceeded" in str(e) or "429" in str(e):
	# print(f"Rate limit hit for model {model}, trying fallback...")
	continue
	else:
	# For other errors, return immediately
	return f"Error while calling LLM: {e}", None

	# If all models fail
	return (
	"I'm temporarily experiencing high demand. Please try again in a few minutes or rephrase your question.",
	None,
	)

	# Streaming ask method (Generator)
	# ============================================
	def ask_stream(self, raw_query: str):
	"""Generator function that yields response chunks in real-time"""

	# Validation
	if (
	not raw_query
	or raw_query is None
	or raw_query == ""
	or len(raw_query) > 1000
	):
	yield "Please provide a valid query under 1,000 characters."
	return

	if raw_query.lower() in GREETINGS_TRIGGERS:
	yield random.choice(GREETINGS)
	return

	# Retrieve documents
	try:
	docs = self.retriever.invoke(raw_query)
	except Exception as e:
	yield f"Error retrieving documents: {e}"
	return

	if not docs:
	yield self.handle_unknown_query()
	return

	# Prepare context and prompt
	context = "\n".join([d.page_content for d in docs])
	fallback = self.handle_unknown_query()
	prompt = self.PROMPT_TEMPLATE.format(
	context=context, question=raw_query, fallback_response=fallback
	)

	messages = [
	{
	"role": "system",
	"content": self.SYSTEM_MESSSAGE,
	},
	{"role": "user", "content": prompt},
	]

	# Try multiple models with fallback
	models_to_try = [
	"openai/gpt-oss-120b",
	"openai/gpt-oss-20b",
	"compound-beta-mini",
	"llama-3.1-8b-instant",
	"llama-3.3-70b-versatile",
	]

	random.shuffle(models_to_try)

	for model in models_to_try:
	try:
	completion = self.client.chat.completions.create(
	model=model,
	messages=messages,
	temperature=TEMPERATURE,
	max_completion_tokens=MAX_TOKENS,
	top_p=TOP_P,
	stream=True, # Enable streaming
	)

	# Stream chunks as they arrive
	has_content = False
	for chunk in completion:
	if chunk.choices[0].delta.content:
	has_content = True
	yield chunk.choices[0].delta.content

	# Check if stream finished
	if chunk.choices[0].finish_reason:
	break

	# If we got content, we're done
	if has_content:
	return

	except Exception as e:
	# Check if it's a rate limit error
	if "rate_limit_exceeded" in str(e) or "429" in str(e):
	continue # Try next model
	else:
	yield f"Error while calling LLM: {e}"
	return

	# If all models fail
	yield "I'm temporarily experiencing high demand. Please try again in a few minutes or rephrase your question."