Update tools/gemini_tool.py
Browse files- tools/gemini_tool.py +49 -0
tools/gemini_tool.py
CHANGED
@@ -0,0 +1,49 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from langchain.tools import BaseTool
|
2 |
+
from typing import Type, Optional, Any
|
3 |
+
from pydantic import BaseModel, Field
|
4 |
+
import google.generativeai as genai
|
5 |
+
from config.settings import settings
|
6 |
+
from services.logger import app_logger
|
7 |
+
|
8 |
+
class GeminiInput(BaseModel):
|
9 |
+
query: str = Field(description="The query or prompt to send to Google Gemini.")
|
10 |
+
|
11 |
+
class GeminiTool(BaseTool):
|
12 |
+
name: str = "google_gemini_chat"
|
13 |
+
description: str = (
|
14 |
+
"Useful for when you need to answer questions or generate text using Google Gemini. "
|
15 |
+
"Use this for general knowledge, creative text generation, or complex reasoning tasks "
|
16 |
+
"that might benefit from a powerful large language model."
|
17 |
+
)
|
18 |
+
args_schema: Type[BaseModel] = GeminiInput
|
19 |
+
# return_direct: bool = True # If you want the agent to return Gemini's output directly
|
20 |
+
|
21 |
+
def _run(self, query: str) -> str:
|
22 |
+
if not settings.GEMINI_API_KEY:
|
23 |
+
app_logger.error("GEMINI_API_KEY not configured.")
|
24 |
+
return "Error: Gemini API key not configured."
|
25 |
+
try:
|
26 |
+
genai.configure(api_key=settings.GEMINI_API_KEY)
|
27 |
+
model = genai.GenerativeModel('gemini-pro')
|
28 |
+
response = model.generate_content(query)
|
29 |
+
return response.text
|
30 |
+
except Exception as e:
|
31 |
+
app_logger.error(f"Error calling Gemini API: {e}")
|
32 |
+
return f"Error interacting with Gemini: {str(e)}"
|
33 |
+
|
34 |
+
async def _arun(self, query: str) -> str:
|
35 |
+
# Asynchronous version (optional, implement if needed)
|
36 |
+
# For simplicity, using the synchronous version for now.
|
37 |
+
# You might need to use an async client for genai if available or run sync in thread.
|
38 |
+
if not settings.GEMINI_API_KEY:
|
39 |
+
app_logger.error("GEMINI_API_KEY not configured.")
|
40 |
+
return "Error: Gemini API key not configured."
|
41 |
+
try:
|
42 |
+
genai.configure(api_key=settings.GEMINI_API_KEY)
|
43 |
+
model = genai.GenerativeModel('gemini-pro')
|
44 |
+
# For async, genai might have an async client or you'd use `loop.run_in_executor`
|
45 |
+
response = await model.generate_content_async(query) # Assuming an async method
|
46 |
+
return response.text
|
47 |
+
except Exception as e:
|
48 |
+
app_logger.error(f"Error calling Gemini API asynchronously: {e}")
|
49 |
+
return f"Error interacting with Gemini: {str(e)}"
|