MCP with Gemini Integration

""" Enhanced Knowledge Base Tool for Gemini LLM Integration. Provides intelligent semantic search through company knowledge base. """ import json import logging import os from typing import Dict, Any, List from pathlib import Path try: from src.tools import ToolDefinition except ImportError: from tools import ToolDefinition logger = logging.getLogger(__name__) class KnowledgeBaseTool: """Enhanced tool for searching company knowledge base with semantic understanding.""" def __init__(self): logger.info("Enhanced knowledge base tool initialized") self.knowledge_base = self._load_knowledge_base() self._validate_knowledge_base() # Initialize rate-limited client instead of direct client self._initialize_rate_limited_client() def _load_knowledge_base(self) -> Dict[str, Any]: """Load and cache the knowledge base with enhanced error handling.""" try: current_dir = os.path.dirname(os.path.abspath(__file__)) kb_path = os.path.join(current_dir, "..", "..", "data", "knowledge_base.json") logger.info(f"Loading knowledge base from: {kb_path}") if not os.path.exists(kb_path): # Create a comprehensive sample knowledge base sample_kb = { "qa_pairs": [ { "question": "What is the company's vacation policy?", "answer": "Employees receive 15 days of paid vacation per year for the first 2 years, 20 days after 2 years of service, and 25 days after 5 years. Vacation days must be approved by your manager in advance. Unused vacation days can be carried over up to 10 days to the next year." }, { "question": "What is the remote work policy?", "answer": "Employees can work remotely up to 3 days per week with manager approval. Full remote work is available for senior employees and special circumstances. All remote workers must maintain core hours of 10 AM - 3 PM EST and be available for meetings. Remote work requires a stable internet connection and appropriate workspace." }, { "question": "What benefits does the company offer?", "answer": "We offer comprehensive health insurance (medical, dental, vision), 401k with 4% company match, life insurance, disability insurance, flexible spending accounts, employee assistance program, professional development stipend of $2000 per year, gym membership reimbursement, and commuter benefits." }, { "question": "How many sick days do employees get?", "answer": "Employees receive 10 paid sick days per year. Unused sick days can be carried over up to a maximum of 20 days. Extended sick leave may be available under FMLA. Sick days can be used for personal illness, medical appointments, or caring for sick family members." }, { "question": "What is the dress code?", "answer": "We have a business casual dress code. On Fridays, casual dress is acceptable. For client meetings, business professional attire is required. The dress code applies to both in-office and remote work when on video calls. Jeans are acceptable on casual Fridays." }, { "question": "What are the working hours?", "answer": "Standard working hours are 9 AM - 5 PM EST, Monday through Friday. Flexible hours are available with manager approval. Core hours of 10 AM - 3 PM EST must be maintained for team collaboration. Overtime may be required for special projects." }, { "question": "How do I request time off?", "answer": "Submit time-off requests through the HR portal at least 2 weeks in advance for vacation and 1 week for sick days. Include the dates, reason, and any relevant details. Your manager will review and approve or request changes. Emergency time off can be requested with manager approval." }, { "question": "What is the performance review process?", "answer": "Performance reviews are conducted annually in March. The process includes self-assessment, manager evaluation, and goal setting for the next year. Reviews cover job performance, goals achievement, and professional development. Feedback is provided throughout the year, not just during formal reviews." }, { "question": "What training and development opportunities are available?", "answer": "We offer $2000 annual professional development stipend, internal training programs, conference attendance, online course reimbursement, mentorship programs, and skill-building workshops. Employees are encouraged to pursue certifications and advanced degrees with company support." }, { "question": "What is the company's policy on overtime?", "answer": "Overtime is paid at 1.5x rate for hours worked beyond 40 hours per week. Overtime must be pre-approved by your manager. Emergency overtime may be required for critical projects. Comp time is available as an alternative to overtime pay with manager approval." } ] } # Ensure the data directory exists os.makedirs(os.path.dirname(kb_path), exist_ok=True) with open(kb_path, 'w') as f: json.dump(sample_kb, f, indent=2) logger.info("Created comprehensive sample knowledge base") with open(kb_path, 'r') as f: kb_data = json.load(f) if not isinstance(kb_data, dict) or 'qa_pairs' not in kb_data: raise ValueError("Invalid knowledge base format: missing 'qa_pairs' key") return kb_data except Exception as e: logger.error(f"Error loading knowledge base: {e}") raise RuntimeError(f"Failed to load knowledge base: {str(e)}") def _validate_knowledge_base(self): """Validate the knowledge base structure and content.""" if not self.knowledge_base.get('qa_pairs'): raise ValueError("Knowledge base is empty or missing qa_pairs") qa_pairs = self.knowledge_base['qa_pairs'] if not isinstance(qa_pairs, list): raise ValueError("qa_pairs must be a list") for i, qa in enumerate(qa_pairs): if not isinstance(qa, dict): raise ValueError(f"QA pair {i} must be a dictionary") if 'question' not in qa or 'answer' not in qa: raise ValueError(f"QA pair {i} missing 'question' or 'answer' key") if not qa['question'] or not qa['answer']: raise ValueError(f"QA pair {i} has empty question or answer") logger.info(f"Knowledge base validated successfully with {len(qa_pairs)} Q&A pairs") def _initialize_rate_limited_client(self): """Initialize rate-limited Gemini client for semantic search.""" from dotenv import load_dotenv try: from src.utils.gemini_client import get_rate_limited_client except ImportError: from utils.gemini_client import get_rate_limited_client load_dotenv("../../.env") gemini_api_key = os.getenv('GEMINI_API_KEY') if not gemini_api_key: raise ValueError("GEMINI_API_KEY not found in environment variables. Please set your API key in the .env file.") try: self.rate_limited_client = get_rate_limited_client(gemini_api_key) logger.info("Rate-limited Gemini client initialized successfully for semantic search") except Exception as e: logger.error(f"Failed to initialize rate-limited Gemini client: {e}") raise RuntimeError(f"Rate-limited Gemini client initialization failed: {str(e)}") async def search_knowledge_base(self, query: str) -> str: """Search and retrieve information from the company knowledge base using advanced LLM-powered semantic search. This tool uses advanced AI to understand your question and find the most relevant information from company policies, benefits, and procedures stored in the knowledge base. Use this for any questions about: - Vacation and time-off policies - Remote work policies - Employee benefits - Sick leave policies - Dress codes - Working hours - Training and development - HR procedures - Performance reviews - Overtime policies - Any other company policies or procedures Args: query: The user's question about company policies or information Returns: The most relevant information from the knowledge base based on semantic understanding """ logger.info(f"search_knowledge_base called with query: {query}") try: if not query or not query.strip(): raise ValueError("Query cannot be empty") query = query.strip() if not self.knowledge_base.get('qa_pairs'): raise ValueError("Knowledge base is empty or not available. Please check the data/knowledge_base.json file.") # Use enhanced LLM-powered semantic search return await self._enhanced_semantic_search(query) except Exception as e: logger.error(f"Error in knowledge base search: {e}", exc_info=True) raise RuntimeError(f"Failed to search knowledge base: {str(e)}") async def _enhanced_semantic_search(self, query: str) -> str: """Use Gemini LLM for enhanced semantic search through the knowledge base.""" try: # Prepare the knowledge base content for the LLM kb_content = [] for i, qa in enumerate(self.knowledge_base.get('qa_pairs', []), 1): kb_content.append(f"{i}. Q: {qa.get('question', '')}\n A: {qa.get('answer', '')}") kb_text = "\n\n".join(kb_content) # Create an enhanced semantic search prompt search_prompt = f"""You are an intelligent company knowledge base assistant. A user has asked a question, and you need to find the most relevant information from our comprehensive company knowledge base. User Question: "{query}" Company Knowledge Base: {kb_text} Instructions: 1. Analyze the user's question to understand what they're looking for 2. Consider synonyms, related concepts, and different ways to ask the same question 3. Find the most relevant Q&A pair(s) that answer their question 4. If you find relevant information, return it in this format: "Here's what I found in the company knowledge base: Q: [Question] A: [Answer]" 5. If you find multiple relevant items, include all of them 6. If no information directly answers their question, say "I couldn't find specific information about that in our knowledge base. Here are some related topics that might help:" and list the closest matches 7. Always be helpful, professional, and accurate 8. Don't make up information that's not in the knowledge base 9. If the user's question is unclear, ask for clarification Please provide your response now.""" # Call Gemini for enhanced semantic search using rate-limited client from google.generativeai import types response = await self.rate_limited_client.generate_content( contents=search_prompt, generation_config=types.GenerationConfig( temperature=0.1, # Low temperature for consistent, factual responses max_output_tokens=1024 ) ) if response and hasattr(response, 'text') and response.text: logger.info("Enhanced semantic search completed successfully") return response.text.strip() else: raise RuntimeError("No response received from Gemini API. Please check your API key and network connection.") except Exception as e: logger.error(f"Error in enhanced semantic search: {e}") raise RuntimeError(f"Enhanced semantic search failed: {str(e)}") # Global instance knowledge_base_tool = KnowledgeBaseTool() def register_tool() -> ToolDefinition: """Register the enhanced knowledge base tool.""" return ToolDefinition( name="get_knowledge_base", description="Search and retrieve information from the company knowledge base using advanced LLM-powered semantic search. Use for questions about company policies, benefits, procedures, HR information, and employee guidelines.", handler=knowledge_base_tool.search_knowledge_base, input_schema={ "type": "object", "properties": { "query": { "type": "string", "description": "The user's question about company policies, benefits, or information" } }, "required": ["query"] }, examples=[ "What is the company's vacation policy?", "How many sick days do employees get?", "What is the remote work policy?", "What benefits does the company offer?", "What is the dress code?", "How do I request time off?", "What are the working hours?", "What is the performance review process?" ] )