mgbam commited on
Commit
8bfd9b5
·
verified ·
1 Parent(s): e0a9fc4

Update tools/gemini_tool.py

Browse files
Files changed (1) hide show
  1. 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)}"