Spaces:
Running
Running
File size: 5,527 Bytes
576227b |
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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
from google import genai
from google.genai import types
import os
from dotenv import load_dotenv
import sys
from src.tool_loader import ToolLoader
from src.utils.suppress_outputs import suppress_output
import logging
logger = logging.getLogger(__name__)
handler = logging.StreamHandler(sys.stdout)
handler.setLevel(logging.INFO)
logger.addHandler(handler)
class GeminiManager:
def __init__(self, toolsLoader: ToolLoader, system_prompt_file="./models/system3.prompt", gemini_model="gemini-2.5-pro-exp-03-25"):
load_dotenv()
self.API_KEY = os.getenv("GEMINI_KEY")
self.client = genai.Client(api_key=self.API_KEY)
self.toolsLoader: ToolLoader = toolsLoader
self.toolsLoader.load_tools()
self.model_name = gemini_model
with open(system_prompt_file, 'r', encoding="utf8") as f:
self.system_prompt = f.read()
def generate_response(self, messages):
return self.client.models.generate_content(
#model='gemini-2.5-pro-preview-03-25',
model=self.model_name,
#model='gemini-2.5-pro-exp-03-25',
#model='gemini-2.0-flash',
contents=messages,
config=types.GenerateContentConfig(
system_instruction=self.system_prompt,
temperature=0.2,
tools=self.toolsLoader.getTools(),
),
)
def handle_tool_calls(self, response):
parts = []
for function_call in response.function_calls:
toolResponse = None
logger.info(f"Function Name: {function_call.name}, Arguments: {function_call.args}")
try:
toolResponse = self.toolsLoader.runTool(function_call.name, function_call.args)
except Exception as e:
logger.warning(f"Error running tool: {e}")
toolResponse = {
"status": "error",
"message": f"Tool {function_call.name} failed to run.",
"output": str(e),
}
logger.debug(f"Tool Response: {toolResponse}")
tool_content = types.Part.from_function_response(
name=function_call.name,
response = {"result":toolResponse})
try:
self.toolsLoader.load_tools()
except Exception as e:
logger.info(f"Error loading tools: {e}. Deleting the tool.")
# delete the created tool
self.toolsLoader.delete_tool(toolResponse['output']['tool_name'], toolResponse['output']['tool_file_path'])
tool_content = types.Part.from_function_response(
name=function_call.name,
response={"result":f"{function_call.name} with {function_call.args} doesn't follow the required format, please read the other tool implementations for reference." + str(e)})
parts.append(tool_content)
return types.Content(
role='model' if self.model_name == "gemini-2.5-pro-exp-03-25" else 'tool',
parts=parts
)
def run(self, messages):
try:
response = suppress_output(self.generate_response)(messages)
except Exception as e:
logger.debug(f"Error generating response: {e}")
shouldRetry = input("An error occurred. Do you want to retry? (y/n): ")
if shouldRetry.lower() == "y":
return self.run(messages)
else:
print("Ending the conversation.")
return messages
logger.debug(f"Response: {response}")
if (not response.text and not response.function_calls):
print("No response from the model.")
# Attach the llm response to the messages
if response.text is not None:
print("CEO:", response.text)
assistant_content = types.Content(
role='model' if self.model_name == "gemini-2.5-pro-exp-03-25" else 'assistant',
parts=[types.Part.from_text(text=response.text)],
)
messages.append(assistant_content)
# Attach the function call response to the messages
if response.candidates[0].content and response.candidates[0].content.parts:
messages.append(response.candidates[0].content)
# Invoke the function calls if any and attach the response to the messages
if response.function_calls:
messages.append(self.handle_tool_calls(response))
shouldContinue = input("Should I continue? (y/n): ")
if shouldContinue.lower() == "y":
return self.run(messages)
else:
print("Ending the conversation.")
return messages
else:
logger.debug("No tool calls found in the response.")
# Start the loop again
return self.start_conversation(messages)
def start_conversation(self, messages=[]):
question = input("User: ")
if ("exit" in question.lower() or "quit" in question.lower()):
print("Ending the conversation.")
return messages
user_content = types.Content(
role='user',
parts=[types.Part.from_text(text=question)],
)
messages.append(user_content)
# Start the conversation loop
return self.run(messages) |