Google

pip install google-generativeai requests

export GOOGLE_API_KEY="your-google-api-key"
export VALYU_API_KEY="your-valyu-api-key"

import google.generativeai as genai
import requests
import json
import os
from typing import Literal

# Configure Gemini
genai.configure(api_key=os.environ['GOOGLE_API_KEY'])

def valyu_search(
    query: str,
    search_type: Literal["all", "web", "proprietary", "news"] = "all",
    max_num_results: int = 5,
    relevance_threshold: float = 0.5,
    max_price: float = 30.0,
    category: str = None
) -> str:
    """
    Search for information using Valyu's comprehensive knowledge base.

    Args:
        query: Natural language search query
        search_type: Type of search - "all", "web", "proprietary", or "news"
        max_num_results: Number of results to return (1-20 for standard API keys, up to 100 with a [special API key](http://platform.valyu.ai/user/account/apikeys?req=increase_results))
        relevance_threshold: Minimum relevance score (0.0-1.0)
        max_price: Maximum cost in dollars. Only applies when provided. If not provided, adjusts automatically based on search type and max number of results.
        category: Natural language category to guide search

    Returns:
        JSON string with search results
    """
    url = "https://api.valyu.ai/v1/search"

    payload = {
        "query": query,
        "search_type": search_type,
        "max_num_results": max_num_results,
        "relevance_threshold": relevance_threshold,
        "max_price": max_price,
        "is_tool_call": True
    }

    if category:
        payload["category"] = category

    headers = {
        "x-api-key": os.environ["VALYU_API_KEY"],
        "Content-Type": "application/json"
    }

    try:
        response = requests.post(url, json=payload, headers=headers)
        response.raise_for_status()
        return json.dumps(response.json(), indent=2)
    except Exception as e:
        return f"Search error: {str(e)}"

# Define the function declaration for Gemini
valyu_function_declaration = genai.protos.FunctionDeclaration(
    name="valyu_search",
    description="Search for real-time information, academic papers, and comprehensive knowledge using Valyu's database",
    parameters=genai.protos.Schema(
        type=genai.protos.Type.OBJECT,
        properties={
            "query": genai.protos.Schema(
                type=genai.protos.Type.STRING,
                description="Natural language search query"
            ),
            "search_type": genai.protos.Schema(
                type=genai.protos.Type.STRING,
                enum=["all", "web", "proprietary", "news"],
                description="Type of search: 'all' for comprehensive, 'web' for current events, 'proprietary' for academic, 'news' for news articles only"
            ),
            "max_num_results": genai.protos.Schema(
                type=genai.protos.Type.INTEGER,
                description="Number of results to return (1-20 for standard API keys, up to 100 with a [special API key](http://platform.valyu.ai/user/account/apikeys?req=increase_results))"
            ),
            "relevance_threshold": genai.protos.Schema(
                type=genai.protos.Type.NUMBER,
                description="Minimum relevance score for results (0.0-1.0)"
            ),
            "max_price": genai.protos.Schema(
                type=genai.protos.Type.NUMBER,
                description="Maximum cost in dollars for this search"
            ),
            "category": genai.protos.Schema(
                type=genai.protos.Type.STRING,
                description="Natural language category to guide search context"
            )
        },
        required=["query"]
    )
)

# Create the tool
valyu_tool = genai.protos.Tool(
    function_declarations=[valyu_function_declaration]
)

# Initialize the model
model = genai.GenerativeModel(
    model_name="gemini-2.0-flash-exp",
    tools=[valyu_tool]
)

def chat_with_search(user_message: str):
    # Start a chat session
    chat = model.start_chat()

    # Send the user message
    response = chat.send_message(
        f"You are a helpful assistant with access to real-time search. "
        f"Use the valyu_search function to find current information when needed. "
        f"User query: {user_message}"
    )

    # Check if Gemini wants to call a function
    if response.candidates[0].content.parts:
        for part in response.candidates[0].content.parts:
            if hasattr(part, 'function_call') and part.function_call:
                function_call = part.function_call

                if function_call.name == "valyu_search":
                    # Extract function arguments
                    function_args = {}
                    for key, value in function_call.args.items():
                        function_args[key] = value

                    # Call the function
                    search_results = valyu_search(**function_args)

                    # Send function response back to Gemini
                    function_response = genai.protos.Part(
                        function_response=genai.protos.FunctionResponse(
                            name="valyu_search",
                            response={"result": search_results}
                        )
                    )

                    # Get final response with search results
                    final_response = chat.send_message(function_response)
                    return final_response.text

    # Return direct response if no function call
    return response.text

# Example usage
result = chat_with_search("What are the latest developments in quantum computing?")
print(result)