Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closing database sessions #1216

Merged
merged 9 commits into from
Jun 20, 2024
Merged

Closing database sessions #1216

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
dc1b648
Closing database sessions
Josh-XT Jun 20, 2024
ca716b0
fix refs in all
Josh-XT Jun 20, 2024
23cfefb
move refs
Josh-XT Jun 20, 2024
936289b
fix imports
Josh-XT Jun 20, 2024
6c91b67
move webhook
Josh-XT Jun 20, 2024
5a930e0
fix update agent config
Josh-XT Jun 20, 2024
dd981a1
improve logging
Josh-XT Jun 20, 2024
c59ab04
use link for image
Josh-XT Jun 20, 2024
2ec9121
fix audio path
Josh-XT Jun 20, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
153 changes: 106 additions & 47 deletions agixt/Agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,9 @@
from Providers import Providers
from Extensions import Extensions
from Globals import getenv, DEFAULT_SETTINGS, DEFAULT_USER
from MagicalAuth import get_user_id, is_agixt_admin
from agixtsdk import AGiXTSDK
from fastapi import HTTPException
from datetime import datetime, timezone, timedelta
import logging
import json
Expand All @@ -29,23 +32,25 @@


def add_agent(agent_name, provider_settings=None, commands=None, user=DEFAULT_USER):
session = get_session()
if not agent_name:
return {"message": "Agent name cannot be empty."}
session = get_session()
# Check if agent already exists
agent = (
session.query(AgentModel)
.filter(AgentModel.name == agent_name, AgentModel.user.has(email=user))
.first()
)
if agent:
session.close()
return {"message": f"Agent {agent_name} already exists."}
agent = (
session.query(AgentModel)
.filter(AgentModel.name == agent_name, AgentModel.user.has(email=DEFAULT_USER))
.first()
)
if agent:
session.close()
return {"message": f"Agent {agent_name} already exists."}
user_data = session.query(User).filter(User.email == user).first()
user_id = user_data.id
Expand Down Expand Up @@ -80,7 +85,7 @@ def add_agent(agent_name, provider_settings=None, commands=None, user=DEFAULT_US
)
session.add(agent_command)
session.commit()

session.close()
return {"message": f"Agent {agent_name} created."}


Expand All @@ -94,6 +99,7 @@ def delete_agent(agent_name, user=DEFAULT_USER):
.first()
)
if not agent:
session.close()
return {"message": f"Agent {agent_name} not found."}, 404

# Delete associated chain steps
Expand Down Expand Up @@ -125,7 +131,7 @@ def delete_agent(agent_name, user=DEFAULT_USER):
# Delete the agent
session.delete(agent)
session.commit()

session.close()
return {"message": f"Agent {agent_name} deleted."}, 200


Expand All @@ -139,11 +145,11 @@ def rename_agent(agent_name, new_name, user=DEFAULT_USER):
.first()
)
if not agent:
session.close()
return {"message": f"Agent {agent_name} not found."}, 404

agent.name = new_name
session.commit()

session.close()
return {"message": f"Agent {agent_name} renamed to {new_name}."}, 200


Expand All @@ -162,21 +168,16 @@ def get_agents(user=DEFAULT_USER):
if agent.name in [a["name"] for a in output]:
continue
output.append({"name": agent.name, "id": agent.id, "status": False})
session.close()
return output


class Agent:
def __init__(self, agent_name=None, user=DEFAULT_USER, ApiClient=None):
def __init__(self, agent_name=None, user=DEFAULT_USER, ApiClient: AGiXTSDK = None):
self.agent_name = agent_name if agent_name is not None else "AGiXT"
self.session = get_session()
user = user if user is not None else DEFAULT_USER
self.user = user.lower()
try:
user_data = self.session.query(User).filter(User.email == self.user).first()
self.user_id = user_data.id
except Exception as e:
logging.error(f"User {self.user} not found.")
raise
self.user_id = get_user_id(user=self.user)
self.AGENT_CONFIG = self.get_agent_config()
self.load_config_keys()
if "settings" not in self.AGENT_CONFIG:
Expand Down Expand Up @@ -286,35 +287,33 @@ def load_config_keys(self):
setattr(self, key, self.AGENT_CONFIG[key])

def get_agent_config(self):
session = get_session()
agent = (
self.session.query(AgentModel)
session.query(AgentModel)
.filter(
AgentModel.name == self.agent_name, AgentModel.user_id == self.user_id
)
.first()
)
if not agent:
# Check if it is a global agent
global_user = (
self.session.query(User).filter(User.email == DEFAULT_USER).first()
)
global_user = session.query(User).filter(User.email == DEFAULT_USER).first()
agent = (
self.session.query(AgentModel)
session.query(AgentModel)
.filter(
AgentModel.name == self.agent_name,
AgentModel.user_id == global_user.id,
)
.first()
)

config = {"settings": {}, "commands": {}}
if agent:
all_commands = self.session.query(Command).all()
all_commands = session.query(Command).all()
agent_settings = (
self.session.query(AgentSettingModel).filter_by(agent_id=agent.id).all()
session.query(AgentSettingModel).filter_by(agent_id=agent.id).all()
)
agent_commands = (
self.session.query(AgentCommand)
session.query(AgentCommand)
.join(Command)
.filter(
AgentCommand.agent_id == agent.id,
Expand All @@ -331,7 +330,10 @@ def get_agent_config(self):
)
for setting in agent_settings:
config["settings"][setting.name] = setting.value
session.commit()
session.close()
return config
session.close()
return {"settings": DEFAULT_SETTINGS, "commands": {}}

async def inference(self, prompt: str, tokens: int = 0, images: list = []):
Expand Down Expand Up @@ -387,24 +389,68 @@ def get_commands_string(self):
return verbose_commands

def update_agent_config(self, new_config, config_key):
session = get_session()
agent = (
self.session.query(AgentModel)
session.query(AgentModel)
.filter(
AgentModel.name == self.agent_name, AgentModel.user_id == self.user_id
)
.first()
)
if not agent:
logging.error(f"Agent '{self.agent_name}' not found in the database.")
return
if self.user == DEFAULT_USER:
return f"Agent {self.agent_name} not found."
# Check if it is a global agent.
global_user = session.query(User).filter(User.email == DEFAULT_USER).first()
global_agent = (
session.query(AgentModel)
.filter(
AgentModel.name == self.agent_name,
AgentModel.user_id == global_user.id,
)
.first()
)
# if it is a global agent, copy it to the user's agents.
if global_agent:
agent = AgentModel(
name=self.agent_name,
user_id=self.user_id,
provider_id=global_agent.provider_id,
)
session.add(agent)
agent_settings = (
session.query(AgentSettingModel)
.filter_by(agent_id=global_agent.id)
.all()
)
for setting in agent_settings:
agent_setting = AgentSettingModel(
agent_id=agent.id,
name=setting.name,
value=setting.value,
)
session.add(agent_setting)
agent_commands = (
session.query(AgentCommand)
.filter_by(agent_id=global_agent.id)
.all()
)
for agent_command in agent_commands:
agent_command = AgentCommand(
agent_id=agent.id,
command_id=agent_command.command_id,
state=agent_command.state,
)
session.add(agent_command)
session.commit()
session.close()
return f"Agent {self.agent_name} configuration updated successfully."
if config_key == "commands":
for command_name, enabled in new_config.items():
command = (
self.session.query(Command).filter_by(name=command_name).first()
)
command = session.query(Command).filter_by(name=command_name).first()
if command:
agent_command = (
self.session.query(AgentCommand)
session.query(AgentCommand)
.filter_by(agent_id=agent.id, command_id=command.id)
.first()
)
Expand All @@ -414,12 +460,12 @@ def update_agent_config(self, new_config, config_key):
agent_command = AgentCommand(
agent_id=agent.id, command_id=command.id, state=enabled
)
self.session.add(agent_command)
session.add(agent_command)
else:
for setting_name, setting_value in new_config.items():
logging.info(f"Setting {setting_name} to {setting_value}.")
agent_setting = (
self.session.query(AgentSettingModel)
session.query(AgentSettingModel)
.filter_by(agent_id=agent.id, name=setting_name)
.first()
)
Expand All @@ -429,15 +475,18 @@ def update_agent_config(self, new_config, config_key):
agent_setting = AgentSettingModel(
agent_id=agent.id, name=setting_name, value=str(setting_value)
)
self.session.add(agent_setting)
session.add(agent_setting)
try:
self.session.commit()
session.commit()
session.close()
logging.info(f"Agent {self.agent_name} configuration updated successfully.")
except Exception as e:
self.session.rollback()
session.rollback()
session.close()
logging.error(f"Error updating agent configuration: {str(e)}")
raise

raise HTTPException(
status_code=500, detail=f"Error updating agent configuration: {str(e)}"
)
return f"Agent {self.agent_name} configuration updated."

def get_browsed_links(self, conversation_id=None):
Expand All @@ -447,21 +496,24 @@ def get_browsed_links(self, conversation_id=None):
Returns:
list: The list of URLs that have been browsed by the agent.
"""
session = get_session()
agent = (
self.session.query(AgentModel)
session.query(AgentModel)
.filter(
AgentModel.name == self.agent_name, AgentModel.user_id == self.user_id
)
.first()
)
if not agent:
session.close()
return []
browsed_links = (
self.session.query(AgentBrowsedLink)
session.query(AgentBrowsedLink)
.filter_by(agent_id=agent.id, conversation_id=conversation_id)
.order_by(AgentBrowsedLink.id.desc())
.all()
)
session.close()
if not browsed_links:
return []
return browsed_links
Expand Down Expand Up @@ -495,8 +547,9 @@ def add_browsed_link(self, url, conversation_id=None):
Returns:
str: The response message.
"""
session = get_session()
agent = (
self.session.query(AgentModel)
session.query(AgentModel)
.filter(
AgentModel.name == self.agent_name, AgentModel.user_id == self.user_id
)
Expand All @@ -507,8 +560,9 @@ def add_browsed_link(self, url, conversation_id=None):
browsed_link = AgentBrowsedLink(
agent_id=agent.id, url=url, conversation_id=conversation_id
)
self.session.add(browsed_link)
self.session.commit()
session.add(browsed_link)
session.commit()
session.close()
return f"Link {url} added to browsed links."

def delete_browsed_link(self, url, conversation_id=None):
Expand All @@ -521,8 +575,9 @@ def delete_browsed_link(self, url, conversation_id=None):
Returns:
str: The response message.
"""
session = get_session()
agent = (
self.session.query(AgentModel)
session.query(AgentModel)
.filter(
AgentModel.name == self.agent_name,
AgentModel.user_id == self.user_id,
Expand All @@ -532,33 +587,37 @@ def delete_browsed_link(self, url, conversation_id=None):
if not agent:
return f"Agent {self.agent_name} not found."
browsed_link = (
self.session.query(AgentBrowsedLink)
session.query(AgentBrowsedLink)
.filter_by(agent_id=agent.id, url=url, conversation_id=conversation_id)
.first()
)
if not browsed_link:
return f"Link {url} not found."
self.session.delete(browsed_link)
self.session.commit()
session.delete(browsed_link)
session.commit()
session.close()
return f"Link {url} deleted from browsed links."

def get_agent_id(self):
session = get_session()
agent = (
self.session.query(AgentModel)
session.query(AgentModel)
.filter(
AgentModel.name == self.agent_name, AgentModel.user_id == self.user_id
)
.first()
)
if not agent:
agent = (
self.session.query(AgentModel)
session.query(AgentModel)
.filter(
AgentModel.name == self.agent_name,
AgentModel.user.has(email=DEFAULT_USER),
)
.first()
)
session.close()
if not agent:
return None
session.close()
return agent.id
Loading
Loading