Spaces:
Sleeping
Sleeping
File size: 2,402 Bytes
7a3bf4a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
from typing import Any, Optional
from langchain_core.callbacks import CallbackManagerForToolRun
from langchain_core.tools import BaseTool
from langchain_core.tools.base import ArgsSchema
from pydantic import BaseModel, Field
from google.genai.types import Tool, GoogleSearch, GenerateContentConfig, Content, Part
from google import genai
class GoogleSearchInput(BaseModel):
query: str = Field(description="The query to search for")
class GoogleSearchTool(BaseTool):
name: str = "google_search"
description: str = "Useful for searching Google to find information about current events, data, or public information."
args_schema: Optional[ArgsSchema] = GoogleSearchInput
client: Any = None
model_id: str = "gemini-2.0-flash"
return_direct: bool = False
google_search_tool: Any = None
def __init__(self, api_key: Optional[str] = None, **kwargs):
super().__init__(**kwargs)
self.client = genai.Client(api_key=api_key)
self.google_search_tool = Tool(google_search=GoogleSearch())
def _run(self, query: str, run_manager: Optional[CallbackManagerForToolRun]=None) -> str:
"""Run the Google search with the query."""
try:
response = self.client.models.generate_content(
model=self.model_id,
contents=Content(
parts=[Part.from_text(text=f"Search for information about: {query}")],
role="user"
),
config=GenerateContentConfig(
tools=[self.google_search_tool],
response_modalities=["TEXT"]
)
)
# Extract search results
return response.text
except Exception as e:
return f"Error performing Google search: {str(e)}"
async def _arun(self, query: str, run_manager: Optional[CallbackManagerForToolRun]=None) -> str:
"""Run the Google search with the query asynchronously."""
return self._run(query, run_manager=run_manager.get_sync())
# Example usage
if __name__ == "__main__":
from dotenv import load_dotenv
# Load API key from environment variables
load_dotenv()
# Create the Google Search runnable
google_search = GoogleSearchTool()
# Run a search query
result = google_search.invoke({"query": "What is the latest news about AI?"})
print(result)
|