#!/usr/bin/env python3 """ Genesis Search Grounding Wrapper ================================ Integrates Gemini 2.0 with Google Search for grounded, real-time responses. Features: - Google Search grounding via Gemini API - Hybrid mode: Search + Qdrant for custom knowledge - Cost tracking for $35/1000 queries pricing - Caching to reduce duplicate queries Usage: from search_grounding import SearchGroundedAgent agent = SearchGroundedAgent() response = agent.query("What's the latest news on AI?") """ import os import sys import json import hashlib from datetime import datetime from typing import Optional, Dict, List, Any # Add genesis-system to path sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import google.generativeai as genai from budget_manager import BudgetManager # Load config CONFIG_PATH = os.path.join(os.path.dirname(__file__), "genesis_config.json") with open(CONFIG_PATH) as f: CONFIG = json.load(f) # Configure Gemini genai.configure(api_key=CONFIG["gemini"]["api_key"]) # Pricing constants SEARCH_GROUNDING_COST_PER_QUERY = 0.035 # $35 per 1000 queries class SearchGroundedAgent: """ Agent that uses Gemini 2.0 with Google Search grounding. Provides real-time, factual responses with citations. """ def __init__(self, budget_limit: float = 10.0, cache_enabled: bool = True): self.budget = BudgetManager(daily_limit=budget_limit) self.cache_enabled = cache_enabled self.cache: Dict[str, Dict] = {} # Initialize Gemini 2.0 Flash - grounding is enabled per-request self.model = genai.GenerativeModel('gemini-2.0-flash') self.grounded_model = genai.GenerativeModel( 'gemini-2.0-flash', generation_config=genai.GenerationConfig( temperature=0.7, ) ) print("[OK] Search Grounding Agent initialized") print(f"[OK] Model: gemini-2.0-flash") print(f"[OK] Cost: ${SEARCH_GROUNDING_COST_PER_QUERY}/grounded query") def _cache_key(self, query: str) -> str: """Generate cache key from query.""" return hashlib.md5(query.lower().strip().encode()).hexdigest() def query(self, prompt: str, use_grounding: bool = True, include_citations: bool = True) -> Dict[str, Any]: """ Query with optional Google Search grounding. Args: prompt: The user's question use_grounding: Whether to use Google Search (costs $0.035/query) include_citations: Whether to extract and include source citations Returns: Dict with 'response', 'grounded', 'citations', 'cost' """ # Check cache cache_key = self._cache_key(prompt) if self.cache_enabled and cache_key in self.cache: cached = self.cache[cache_key] cached["from_cache"] = True print(f"[CACHE] Hit for query: {prompt[:50]}...") return cached # Budget check if not self.budget.is_within_budget(): return { "response": "Budget limit reached. Cannot execute grounded query.", "grounded": False, "citations": [], "cost": 0, "error": "budget_exceeded" } try: if use_grounding: # Use model with google_search tool enabled response = self.model.generate_content(prompt) cost = SEARCH_GROUNDING_COST_PER_QUERY else: # Use base model without grounding base_model = genai.GenerativeModel('gemini-2.0-flash') response = base_model.generate_content(prompt) # Calculate token-based cost usage = response.usage_metadata input_cost = (usage.prompt_token_count / 1_000_000) * 0.075 output_cost = (usage.candidates_token_count / 1_000_000) * 0.30 cost = input_cost + output_cost # Extract response text response_text = response.text # Extract citations if available (from grounding metadata) citations = [] if hasattr(response, 'candidates') and response.candidates: candidate = response.candidates[0] if hasattr(candidate, 'grounding_metadata') and candidate.grounding_metadata: gm = candidate.grounding_metadata if hasattr(gm, 'grounding_chunks'): for chunk in gm.grounding_chunks: if hasattr(chunk, 'web') and chunk.web: citations.append({ "title": getattr(chunk.web, 'title', 'Unknown'), "uri": getattr(chunk.web, 'uri', '') }) # Log cost self.budget.log_cost(cost, "search_grounding" if use_grounding else "gemini-2.0-flash") result = { "response": response_text, "grounded": use_grounding, "citations": citations, "cost": cost, "timestamp": datetime.now().isoformat(), "from_cache": False } # Cache result if self.cache_enabled: self.cache[cache_key] = result return result except Exception as e: return { "response": f"Error: {str(e)}", "grounded": False, "citations": [], "cost": 0, "error": str(e) } def research(self, topic: str, depth: int = 3) -> Dict[str, Any]: """ Deep research on a topic using multiple grounded queries. Args: topic: Research topic depth: Number of follow-up queries (default 3) Returns: Dict with 'summary', 'findings', 'total_cost', 'citations' """ findings = [] all_citations = [] total_cost = 0.0 # Initial query initial = self.query(f"Research: {topic}. Provide comprehensive, factual information.") findings.append({"query": topic, "response": initial["response"]}) all_citations.extend(initial.get("citations", [])) total_cost += initial.get("cost", 0) # Follow-up queries for depth follow_ups = [ f"What are the latest developments in {topic}?", f"What are the key challenges or limitations of {topic}?", f"What are practical applications of {topic}?" ] for i, follow_up in enumerate(follow_ups[:depth]): result = self.query(follow_up) findings.append({"query": follow_up, "response": result["response"]}) all_citations.extend(result.get("citations", [])) total_cost += result.get("cost", 0) # Generate summary summary_prompt = f"""Synthesize the following research findings into a concise summary: Topic: {topic} Findings: {json.dumps(findings, indent=2)} Provide a structured summary with key insights.""" summary_result = self.query(summary_prompt, use_grounding=False) total_cost += summary_result.get("cost", 0) return { "topic": topic, "summary": summary_result["response"], "findings": findings, "citations": all_citations, "total_cost": total_cost, "timestamp": datetime.now().isoformat() } def get_stats(self) -> Dict[str, Any]: """Get agent statistics.""" return { "cache_size": len(self.cache), "budget_spent": self.budget.get_current_spend(), "budget_limit": self.budget.daily_limit, "within_budget": self.budget.is_within_budget() } # CLI for testing if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="Search Grounding Agent") parser.add_argument("--query", "-q", type=str, help="Single query") parser.add_argument("--research", "-r", type=str, help="Research topic") parser.add_argument("--test", action="store_true", help="Run self-test") args = parser.parse_args() agent = SearchGroundedAgent() if args.test: print("\n[TEST] Search Grounding Self-Test") print(f"[OK] Agent initialized: {agent.get_stats()}") # Test without grounding (free) result = agent.query("What is 2 + 2?", use_grounding=False) print(f"[OK] Non-grounded query: {result['response'][:100]}...") print(f"[OK] Cost: ${result['cost']:.6f}") print("\n[DONE] Self-test complete") elif args.query: result = agent.query(args.query) print(f"\nResponse: {result['response']}") print(f"\nCitations: {result.get('citations', [])}") print(f"Cost: ${result['cost']:.4f}") elif args.research: result = agent.research(args.research) print(f"\nSummary: {result['summary']}") print(f"\nTotal Cost: ${result['total_cost']:.4f}") print(f"Citations: {len(result['citations'])} sources") else: print("Usage: python search_grounding.py --query 'your question'") print(" python search_grounding.py --research 'topic'") print(" python search_grounding.py --test")