Spaces:
Running
Running
File size: 5,885 Bytes
576227b 434b328 576227b 434b328 576227b 434b328 576227b 434b328 576227b 434b328 576227b 434b328 576227b 434b328 576227b 434b328 576227b 434b328 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 130 131 132 133 134 |
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
from src.utils.streamlit_interface import get_user_message, output_assistant_response
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()
self.messages = []
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 = get_user_message("An error occurred. Do you want to retry? (y/n): ")
if shouldRetry and shouldRetry.lower() == "y":
return self.run(messages)
else:
output_assistant_response("Ending the conversation.")
return messages
logger.debug(f"Response: {response}")
if (not response.text and not response.function_calls):
output_assistant_response("No response from the model.")
# Attach the llm response to the messages
if response.text is not None:
output_assistant_response("CEO: " + response.text)
# 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 = get_user_message("Should I continue? (y/n): ")
if shouldContinue.lower() == "y":
return self.run(messages)
else:
output_assistant_response("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 = get_user_message("User: ")
# question = input("User: ")
if question and ("exit" in question.lower() or "quit" in question.lower()):
output_assistant_response("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) |