""" LLM Service — Gemini (primary) + Groq (fallback) orchestrator. Phase 2: Gemini Integration. - When USE_GEMINI=true (default), queries go through GeminiService. - When USE_GEMINI=false or GEMINI_API_KEY is missing, falls back to Groq. - _get_llm: Groq LLM factory (preserved for fallback) - MathAIEngine: Orchestrates retrieval, symbolic math, and LLM calls. """ import re import time import logging from typing import Dict, Any, Optional, Tuple from backend.src.config import settings, SYSTEM_TEMPLATE from backend.src.math.symbolic_engine import SymbolicMathEngine from backend.src.services.memory_service import MongoDBChatMemory try: from langchain_core.messages import HumanMessage, AIMessage, SystemMessage except ImportError: from langchain.schema import HumanMessage, AIMessage, SystemMessage logger = logging.getLogger("math_assistant.llm") _LLM_CACHE = {} class GroqLLMWrapper: """Wrapper around native Groq client to mimic Langchain's invoke() and astream() methods.""" def __init__(self, api_key: str, model_name: str): from groq import Groq, AsyncGroq self.client = Groq(api_key=api_key) self.async_client = AsyncGroq(api_key=api_key) self.model_name = model_name def invoke(self, messages): from langchain_core.messages import AIMessage groq_msgs = [] # If it's just a single string prompt if isinstance(messages, str): groq_msgs.append({"role": "user", "content": messages}) elif isinstance(messages, list): for m in messages: if isinstance(m, str): groq_msgs.append({"role": "user", "content": m}) elif isinstance(m, dict): # In case it's a raw dict groq_msgs.append(m) else: role = "user" if getattr(m, "type", "user") == "human" else getattr(m, "type", "user") content = m.content if isinstance(m.content, str) else str(m.content) groq_msgs.append({"role": role, "content": content}) chat_completion = self.client.chat.completions.create( messages=groq_msgs, model=self.model_name, temperature=0.1, max_tokens=2048, ) return AIMessage(content=chat_completion.choices[0].message.content) async def astream(self, messages): from langchain_core.messages import AIMessageChunk groq_msgs = [] if isinstance(messages, str): groq_msgs.append({"role": "user", "content": messages}) elif isinstance(messages, list): for m in messages: if isinstance(m, str): groq_msgs.append({"role": "user", "content": m}) elif isinstance(m, dict): groq_msgs.append(m) else: role = "user" if getattr(m, "type", "user") == "human" else getattr(m, "type", "user") content = m.content if isinstance(m.content, str) else str(m.content) groq_msgs.append({"role": role, "content": content}) stream = await self.async_client.chat.completions.create( messages=groq_msgs, model=self.model_name, temperature=0.1, max_tokens=2048, stream=True ) async for chunk in stream: content = chunk.choices[0].delta.content if content is not None: yield AIMessageChunk(content=content) def _get_llm(model=None): """Get or create a cached Groq LLM instance for a given model name.""" key = model or settings.groq_model_fallbacks[0] if key not in _LLM_CACHE: api_key = settings.groq_api_key if not api_key: raise ValueError("GROQ_API_KEY not set. Add it to your .env file.") logger.info(f"Initializing Native Groq LLM: {key}") _LLM_CACHE[key] = GroqLLMWrapper(api_key=api_key, model_name=key) return _LLM_CACHE[key] class MathAIEngine: """Orchestrates RAG retrieval, symbolic math, and LLM calls for math tutoring.""" def __init__(self, vector_store=None, session_id: str = "default"): try: self.llm = self._init_llm() if not settings.use_gemini else None except Exception as e: logger.warning(f"Failed to initialize Groq LLM: {e}") self.llm = None self.vector_store = vector_store self.memory = MongoDBChatMemory(session_id=session_id) self.symbolic = SymbolicMathEngine() self.session_id = session_id def _init_llm(self): """Initialize the fallback Groq LLM model.""" return _get_llm(settings.groq_model_fallbacks[0]) def _retrieve_context(self, query: str) -> Tuple[list, str]: """Retrieve relevant documents from the vector store.""" if not self.vector_store or not self.vector_store.is_ready(): return [], "No knowledge base available. Using general mathematical knowledge." docs = self.vector_store.similarity_search(query, k=3) if not docs: return [], "No specific context found." parts = [f"[Reference {i+1} - {d.metadata.get('topic','math')}]\n{d.page_content}" for i, d in enumerate(docs)] return docs, "\n\n---\n\n".join(parts) def _symbolic_hint(self, query: str) -> Optional[str]: """Try to get a symbolic verification hint for the query.""" ql = query.lower() for pattern, action in [ (r"(?:differentiate|derivative of|d/dx)\s+(.+?)(?:\s+with respect|\s*$)", "diff"), (r"(?:integrate|integral of)\s+(.+?)(?:\s+with respect|\s+dx|\s*$)", "int"), (r"solve\s+(.+?)\s+(?:for|=)", "solve"), ]: m = re.search(pattern, ql) if m: expr = m.group(1).strip() result = (self.symbolic.differentiate(expr) if action == "diff" else self.symbolic.integrate(expr) if action == "int" else self.symbolic.solve_equation(expr)) if result: return f"[Symbolic verification: {result}]" return None def query(self, user_input: str) -> Dict[str, Any]: """Process a user math query through the full RAG pipeline. Flow: symbolic hint → vector retrieval → chat history → Gemini (primary) or Groq (fallback) → store response. """ hint = self._symbolic_hint(user_input) source_docs, context = self._retrieve_context(user_input) chat_history = self.memory.get_langchain_messages(limit=4) # Store human message BEFORE LLM call so history order is correct self.memory.add_message("human", user_input) # ── Try Gemini first (Phase 2) ───────────────────────────────── if settings.use_gemini and settings.gemini_api_key: try: from backend.src.services.gemini_service import GeminiService gemini = GeminiService() answer = gemini.query(user_input, context=context, chat_history=chat_history) if answer and not answer.startswith("⚠️"): self.memory.add_message("assistant", answer) sources = [{"topic": d.metadata.get("topic", "unknown"), "source": d.metadata.get("source", "kb"), "difficulty": d.metadata.get("difficulty", "unknown")} for d in source_docs] return {"answer": answer, "sources": sources, "symbolic_hint": hint, "session_id": self.session_id, "context_docs": len(source_docs)} except Exception as e: logger.warning(f"Gemini failed, falling back to Groq: {e}") # ── Groq Fallback ────────────────────────────────────────────── # Build messages list manually — NEVER pass math content through # LangChain format_messages(), because curly braces in math break templates system_text = SYSTEM_TEMPLATE.replace("{context}", context) llm_messages = [SystemMessage(content=system_text)] for msg in chat_history: llm_messages.append(msg) llm_messages.append(HumanMessage(content=user_input)) # Auto-retry with model fallback answer = None last_error = None used_models = [] for model_name in settings.groq_model_fallbacks: if model_name in used_models: continue used_models.append(model_name) llm = _get_llm(model_name) for _attempt in range(2): try: raw = llm.invoke(llm_messages).content if raw and not any(p in raw.lower() for p in ["rate limit", "too many requests", "service unavailable"]): answer = raw if model_name != settings.groq_model_fallbacks[0]: answer = f"*(Using fallback model: {model_name})*\n\n" + answer break else: raise Exception(raw or "Empty response") except Exception as e: last_error = e err = str(e).lower() full_err = str(e) logger.warning(f"Model {model_name} attempt {_attempt+1} failed: {e}") if "per day" in full_err or "tokens per day" in full_err: logger.info(f"Daily limit on {model_name}, trying next model...") break elif "429" in full_err or "rate_limit" in err: time.sleep(2 ** _attempt) continue elif "timeout" in err or "connect" in err or "503" in err: time.sleep(1) continue else: break if answer: break if answer is None: full_err = str(last_error) err = full_err.lower() logger.error(f"LLM failed: [{type(last_error).__name__}] {full_err}") if "401" in full_err or "invalid_api_key" in err: answer = "⚠️ Invalid API key. Check GROQ_API_KEY in your .env file." elif "429" in full_err or "rate_limit_exceeded" in err: retry_match = re.search(r'try again in (.+?)\.', full_err) retry_info = f" Groq says: try again in **{retry_match.group(1)}**." if retry_match else "" if "per day" in full_err or "tokens per day" in full_err or "TPD" in full_err: answer = (f"⚠️ **Daily token limit reached** (Groq free tier: 100,000 tokens/day).{retry_info}\n\n" f"To keep using the app now, change your `.env`:\n```\nLLM_MODEL=llama3-8b-8192\n```\n" f"Then restart Streamlit. The 8B model has a separate 500k/day quota.") else: answer = f"⚠️ **Rate limit hit** (too many requests per minute).{retry_info} Wait 20–30 seconds and try again." elif "context_length" in err or ("context" in err and "length" in err): answer = "⚠️ Question + context too long. Try a shorter question." elif "connect" in err or "connection" in err: answer = "⚠️ Cannot reach Groq API. Check your internet connection." elif "timeout" in err: answer = "⚠️ Request timed out. Try again." else: answer = f"⚠️ Error ({type(last_error).__name__}): {full_err}" self.memory.add_message("assistant", answer) sources = [{"topic": d.metadata.get("topic", "unknown"), "source": d.metadata.get("source", "kb"), "difficulty": d.metadata.get("difficulty", "unknown")} for d in source_docs] return {"answer": answer, "sources": sources, "symbolic_hint": hint, "session_id": self.session_id, "context_docs": len(source_docs)} def clear_memory(self): """Clear chat history for this session.""" self.memory.clear_history() def get_history(self): """Get full chat history for this session.""" return self.memory.get_history(limit=50)