{openNewTab(-3, "Settings", "Settings", false); getUserClick('Settings Viewed', {})}} className="top_right_icon" width={16} height={16} src="/images/settings.svg" alt="dropdown-icon"/>
setDropdown(true)}
onMouseLeave={() => setDropdown(false)}>
diff --git a/gui/pages/_app.css b/gui/pages/_app.css
index 8698f03a3f..479647a478 100644
--- a/gui/pages/_app.css
+++ b/gui/pages/_app.css
@@ -907,6 +907,14 @@ p {
.fw_500{font-weight: 500;}
+.br_4{border-radius: 4px}
+.br_5{border-radius: 5px}
+.br_6{border-radius: 6px}
+.br_8{border-radius: 8px}
+
+.fs_10{font-size: 10px}
+.fs_20{font-size: 20px}
+
.text_9{
color: #FFF;
font-family: Inter;
@@ -978,6 +986,14 @@ p {
line-height: normal;
}
+.text_20 {
+ color: #FFF;
+ font-size: 20px;
+ font-style: normal;
+ font-weight: 400;
+ line-height: normal;
+}
+
.text_20_bold{
color: #FFF;
font-size: 20px;
@@ -1138,13 +1154,12 @@ p {
.gap_16{gap:16px;}
.gap_20{gap:20px;}
-.border_gray{border: 1px solid rgba(255, 255, 255, 0.08)}
-.border_left_none{border-left: none;}
+.br_left_grey{border-left: 1px solid rgba(255, 255, 255, 0.08)}
.border_top_none{border-top: none;}
.border_bottom_none{border-bottom: none;}
.border_bottom_grey{border-bottom: 1px solid rgba(255, 255, 255, 0.08)}
-.border_radius_8{border-radius: 8px;}
-.border_radius_25{border-radius: 25px;}
+
+.bt_white{border-top: 1px solid rgba(255, 255, 255, 0.08);}
.bt_white{border-top: 1px solid rgba(255, 255, 255, 0.08);}
@@ -1154,13 +1169,16 @@ p {
.lh_16{line-height: 16px;}
.lh_17{line-height: 17px;}
.lh_18{line-height: 18px;}
+.lh_24{line-height: 24px;}
.padding_0{padding: 0}
.padding_5{padding: 5px;}
+.padding_6{padding: 6px;}
.padding_8{padding: 8px;}
.padding_10{padding: 10px;}
.padding_12{padding: 12px;}
.padding_16{padding: 16px;}
+.padding_20{padding: 20px;}
.padding_8_6{padding: 8px 6px;}
.padding_2_8{padding: 2px 8px;}
@@ -1385,19 +1403,6 @@ tr{
margin: 2px;
}
-.tools_used_tooltip{
- position: absolute;
- transform: translateX(-50%);
- padding: 6px;
- color: #fff;
- font-size: 12px;
- white-space: pre;
- z-index: 100;
- border-radius: 8px;
- background: #2E293F;
- margin-top: -40px;
-}
-
.image_class{
background: #FFFFFF80;
border-radius: 20px;
@@ -1478,8 +1483,11 @@ tr{
}
}
-.bg_black{background: black}
-.bg_white{background: white}
+.bg_black{background: black;}
+.bg_white{background: white;}
+.bg_none{background: none;}
+.bg_primary{background: #2E293F;}
+.bg_secondary{background: #272335;}
.bg_none{background: none}
.container {
@@ -1796,7 +1804,6 @@ tr{
.history_box, .history_box_selected {
width: 100%;
- padding: 10px;
color: white;
font-style: normal;
font-weight: 400;
@@ -1894,3 +1901,10 @@ tr{
flex-direction: row;
justify-content: center;
}
+
+.tooltip-class {
+ background-color: green;
+ border-radius: 6px;
+}
+
+
diff --git a/gui/pages/_app.js b/gui/pages/_app.js
index 6fe6ec6a4c..ad98f09bc0 100644
--- a/gui/pages/_app.js
+++ b/gui/pages/_app.js
@@ -14,17 +14,19 @@ import {
validateAccessToken,
checkEnvironment,
addUser,
- installToolkitTemplate, installAgentTemplate, installKnowledgeTemplate
+ installToolkitTemplate, installAgentTemplate, installKnowledgeTemplate, getFirstSignup
} from "@/pages/api/DashboardService";
-import {githubClientId} from "@/pages/api/apiConfig";
+import {githubClientId, mixpanelId} from "@/pages/api/apiConfig";
import {
getGithubClientId
} from "@/pages/api/DashboardService";
import {useRouter} from 'next/router';
import querystring from 'querystring';
-import {refreshUrl, loadingTextEffect} from "@/utils/utils";
+import {refreshUrl, loadingTextEffect, getUTMParametersFromURL, setLocalStorageValue, getUserClick} from "@/utils/utils";
import MarketplacePublic from "./Content/Marketplace/MarketplacePublic"
import {toast} from "react-toastify";
+import mixpanel from 'mixpanel-browser';
+
export default function App() {
const [selectedView, setSelectedView] = useState('');
@@ -101,12 +103,7 @@ export default function App() {
}
useEffect(() => {
- if (window.location.href.toLowerCase().includes('marketplace')) {
- setShowMarketplace(true);
- } else {
- installFromMarketplace();
- }
-
+ handleMarketplace()
loadingTextEffect('Initializing SuperAGI', setLoadingText, 500);
checkEnvironment()
@@ -115,6 +112,8 @@ export default function App() {
setEnv(env);
if (typeof window !== 'undefined') {
+ if(response.data.env === 'PROD' && mixpanelId())
+ mixpanel.init(mixpanelId(), { debug: false, track_pageview: true, persistence: 'localStorage' });
localStorage.setItem('applicationEnvironment', env);
}
@@ -123,35 +122,37 @@ export default function App() {
const queryParams = router.asPath.split('?')[1];
const parsedParams = querystring.parse(queryParams);
let access_token = parsedParams.access_token || null;
+ let first_login = parsedParams.first_time_login || false
+
+ const utmParams = getUTMParametersFromURL();
+ if (utmParams)
+ sessionStorage.setItem('utm_source', utmParams.utm_source);
+ const signupSource = sessionStorage.getItem('utm_source');
if (typeof window !== 'undefined' && access_token) {
localStorage.setItem('accessToken', access_token);
refreshUrl();
}
-
validateAccessToken()
.then((response) => {
setUserName(response.data.name || '');
+ if(mixpanelId())
+ mixpanel.identify(response.data.email)
+ if(first_login)
+ getUserClick('New Sign Up', {})
+ else
+ getUserClick('User Logged In', {})
+
+ if(signupSource) {
+ handleSignUpSource(signupSource)
+ }
fetchOrganisation(response.data.id);
})
.catch((error) => {
console.error('Error validating access token:', error);
});
} else {
- const userData = {
- "name": "SuperAGI User",
- "email": "super6@agi.com",
- "password": "pass@123",
- }
-
- addUser(userData)
- .then((response) => {
- setUserName(response.data.name);
- fetchOrganisation(response.data.id);
- })
- .catch((error) => {
- console.error('Error adding user:', error);
- });
+ handleLocalEnviroment()
}
})
.catch((error) => {
@@ -197,6 +198,40 @@ export default function App() {
}
}
+ const handleLocalEnviroment = () => {
+ const userData = {
+ "name": "SuperAGI User",
+ "email": "super6@agi.com",
+ "password": "pass@123",
+ }
+
+ addUser(userData)
+ .then((response) => {
+ setUserName(response.data.name);
+ fetchOrganisation(response.data.id);
+ })
+ .catch((error) => {
+ console.error('Error adding user:', error);
+ });
+ };
+ const handleSignUpSource = (signup) => {
+ getFirstSignup(signup)
+ .then((response) => {
+ sessionStorage.removeItem('utm_source');
+ })
+ .catch((error) => {
+ console.error('Error validating source:', error);
+ })
+ };
+
+ const handleMarketplace = () => {
+ if (window.location.href.toLowerCase().includes('marketplace')) {
+ setShowMarketplace(true);
+ } else {
+ installFromMarketplace();
+ }
+ };
+
useEffect(() => {
const clearLocalStorage = () => {
Object.keys(localStorage).forEach((key) => {
diff --git a/gui/pages/api/DashboardService.js b/gui/pages/api/DashboardService.js
index c26330d254..38e13c698f 100644
--- a/gui/pages/api/DashboardService.js
+++ b/gui/pages/api/DashboardService.js
@@ -389,10 +389,15 @@ export const getToolLogs = (toolName) => {
export const publishTemplateToMarketplace = (agentData) => {
return api.post(`/agent_templates/publish_template`, agentData);
};
+
export const getKnowledgeMetrics = (knowledgeName) => {
return api.get(`analytics/knowledge/${knowledgeName}/usage`)
}
export const getKnowledgeLogs = (knowledgeName) => {
return api.get(`analytics/knowledge/${knowledgeName}/logs`)
-}
\ No newline at end of file
+}
+
+export const getFirstSignup = (source) => {
+ return api.post(`/users/first_login_source/${source}`,);
+};
\ No newline at end of file
diff --git a/gui/pages/api/apiConfig.js b/gui/pages/api/apiConfig.js
index b7b3cdf20a..8f6338e0c8 100644
--- a/gui/pages/api/apiConfig.js
+++ b/gui/pages/api/apiConfig.js
@@ -2,6 +2,7 @@ import axios from 'axios';
const GITHUB_CLIENT_ID = process.env.GITHUB_CLIENT_ID;
const API_BASE_URL = process.env.NEXT_PUBLIC_API_BASE_URL || 'http://localhost:8001';
+const MIXPANEL_AUTH_ID = process.env.MIXPANEL_AUTH_ID
export const baseUrl = () => {
return API_BASE_URL;
@@ -11,6 +12,10 @@ export const githubClientId = () => {
return GITHUB_CLIENT_ID;
};
+export const mixpanelId = () => {
+ return MIXPANEL_AUTH_ID;
+};
+
const api = axios.create({
baseURL: API_BASE_URL,
headers: {
diff --git a/gui/utils/utils.js b/gui/utils/utils.js
index f0e4d4cc6b..ee46fb6bdd 100644
--- a/gui/utils/utils.js
+++ b/gui/utils/utils.js
@@ -1,9 +1,10 @@
-import {formatDistanceToNow} from 'date-fns';
+import {formatDistanceToNow, format, addMinutes} from 'date-fns';
import {utcToZonedTime} from 'date-fns-tz';
-import {baseUrl} from "@/pages/api/apiConfig";
+import {baseUrl, mixpanelId} from "@/pages/api/apiConfig";
import {EventBus} from "@/utils/eventBus";
import JSZip from "jszip";
import moment from 'moment';
+import mixpanel from 'mixpanel-browser'
const toolkitData = {
'Jira Toolkit': '/images/jira_icon.svg',
@@ -481,4 +482,67 @@ export const modelIcon = (model) => {
}
return icons[model];
+}
+
+export const modelGetAuth = (modelProvider) => {
+ const externalLinks = {
+ 'Replicate': 'https://replicate.com/account/api-tokens',
+ 'Hugging Face': 'https://huggingface.co/settings/tokens',
+ 'OpenAI': 'https://platform.openai.com/account/api-keys',
+ 'Google Palm': 'https://developers.generativeai.google/products/palm',
+ }
+
+ return externalLinks[modelProvider]
+}
+
+export const formatDateTime = (dateTimeString) => {
+ const date = new Date(dateTimeString);
+ const adjustedDate = addMinutes(addMinutes(date, 5 * 60), 30);
+ const formattedDate = format(adjustedDate, 'd MMM yyyy HH:mm');
+
+ return formattedDate;
+};
+
+export const convertWaitingPeriod = (waitingPeriod) => {
+ let convertedValue = waitingPeriod;
+ let unit = 'seconds';
+
+ if (convertedValue >= 60 && convertedValue < 3600) {
+ convertedValue = Math.floor(convertedValue / 60);
+ unit = 'minutes';
+ } else if (convertedValue >= 3600 && convertedValue < 86400) {
+ convertedValue = Math.floor(convertedValue / 3600);
+ unit = 'hours';
+ } else if (convertedValue >= 86400 && convertedValue < 604800) {
+ convertedValue = Math.floor(convertedValue / 86400);
+ unit = 'days';
+ } else if (convertedValue >= 604800) {
+ convertedValue = Math.floor(convertedValue / 604800);
+ unit = 'weeks';
+ }
+
+ return convertedValue + ' ' + unit;
+}
+
+export const getUTMParametersFromURL = () => {
+ const params = new URLSearchParams(window.location.search);
+
+ const utmParams = {
+ utm_source: params.get('utm_source') || '',
+ utm_medium: params.get('utm_medium') || '',
+ utm_campaign: params.get('utm_campaign') || '',
+ };
+
+ if (!utmParams.utm_source && !utmParams.utm_medium && !utmParams.utm_campaign) {
+ return null;
+ }
+
+ return utmParams;
+}
+
+export const getUserClick = (event, props) => {
+ const env = localStorage.getItem('applicationEnvironment');
+ if(env === 'PROD' && mixpanelId()){
+ mixpanel.track(event, props)
+ }
}
\ No newline at end of file
diff --git a/main.py b/main.py
index a0cd36223f..2e4b118525 100644
--- a/main.py
+++ b/main.py
@@ -307,14 +307,14 @@ def github_auth_handler(code: str = Query(...), Authorize: AuthJWT = Depends()):
db_user: User = db.session.query(User).filter(User.email == user_email).first()
if db_user is not None:
jwt_token = create_access_token(user_email, Authorize)
- redirect_url_success = f"{frontend_url}?access_token={jwt_token}"
+ redirect_url_success = f"{frontend_url}?access_token={jwt_token}&first_time_login={False}"
return RedirectResponse(url=redirect_url_success)
user = User(name=user_data["name"], email=user_email)
db.session.add(user)
db.session.commit()
jwt_token = create_access_token(user_email, Authorize)
- redirect_url_success = f"{frontend_url}?access_token={jwt_token}"
+ redirect_url_success = f"{frontend_url}?access_token={jwt_token}&first_time_login={True}"
return RedirectResponse(url=redirect_url_success)
else:
redirect_url_failure = "https://superagi.com/"
diff --git a/migrations/versions/3867bb00a495_added_first_login_source.py b/migrations/versions/3867bb00a495_added_first_login_source.py
new file mode 100644
index 0000000000..eaf28f0efe
--- /dev/null
+++ b/migrations/versions/3867bb00a495_added_first_login_source.py
@@ -0,0 +1,28 @@
+"""added_first_login_source
+
+Revision ID: 3867bb00a495
+Revises: 661ec8a4c32e
+Create Date: 2023-09-15 02:06:24.006555
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '3867bb00a495'
+down_revision = '661ec8a4c32e'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+ # ### commands auto generated by Alembic - please adjust! ###
+ op.add_column('users', sa.Column('first_login_source', sa.String(), nullable=True))
+ # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+ # ### commands auto generated by Alembic - please adjust! ###
+ op.drop_column('users', 'first_login_source')
+ # ### end Alembic commands ###
diff --git a/migrations/versions/661ec8a4c32e_open_ai_error_handling.py b/migrations/versions/661ec8a4c32e_open_ai_error_handling.py
new file mode 100644
index 0000000000..210c1a51bf
--- /dev/null
+++ b/migrations/versions/661ec8a4c32e_open_ai_error_handling.py
@@ -0,0 +1,30 @@
+"""open_ai_error_handling
+
+Revision ID: 661ec8a4c32e
+Revises: 40affbf3022b
+Create Date: 2023-09-07 10:41:07.462436
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '661ec8a4c32e'
+down_revision = 'c4f2f6ba602a'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+ # ### commands auto generated by Alembic - please adjust! ###
+ op.add_column('agent_execution_feeds', sa.Column('error_message', sa.String(), nullable=True))
+ op.add_column('agent_executions', sa.Column('last_shown_error_id', sa.Integer(), nullable=True))
+ # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+ # ### commands auto generated by Alembic - please adjust! ###
+ op.drop_column('agent_executions', 'last_shown_error_id')
+ op.drop_column('agent_execution_feeds', 'error_message')
+ # ### end Alembic commands ###
diff --git a/migrations/versions/c4f2f6ba602a_agent_workflow_wait_step.py b/migrations/versions/c4f2f6ba602a_agent_workflow_wait_step.py
new file mode 100644
index 0000000000..c37515c2bf
--- /dev/null
+++ b/migrations/versions/c4f2f6ba602a_agent_workflow_wait_step.py
@@ -0,0 +1,39 @@
+"""agent_workflow_wait_step
+
+Revision ID: c4f2f6ba602a
+Revises: 40affbf3022b
+Create Date: 2023-09-04 05:34:10.195248
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = 'c4f2f6ba602a'
+down_revision = '40affbf3022b'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+ # ### commands auto generated by Alembic - please adjust! ###
+ op.create_table('agent_workflow_step_waits',
+ sa.Column('id', sa.Integer(), nullable=False),
+ sa.Column('name', sa.String(), nullable=True),
+ sa.Column('description', sa.String(), nullable=True),
+ sa.Column('unique_id', sa.String(), nullable=True),
+ sa.Column('delay', sa.Integer(), nullable=True),
+ sa.Column('wait_begin_time', sa.DateTime(), nullable=True),
+ sa.Column('status', sa.String(), nullable=True),
+ sa.Column('created_at', sa.DateTime(), nullable=True),
+ sa.Column('updated_at', sa.DateTime(), nullable=True),
+ sa.PrimaryKeyConstraint('id')
+ )
+ # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+ # ### commands auto generated by Alembic - please adjust! ###
+ op.drop_table('agent_workflow_step_waits')
+ # ### end Alembic commands ###
diff --git a/superagi/agent/agent_iteration_step_handler.py b/superagi/agent/agent_iteration_step_handler.py
index 5ab21d4d66..0e2c4fcecf 100644
--- a/superagi/agent/agent_iteration_step_handler.py
+++ b/superagi/agent/agent_iteration_step_handler.py
@@ -11,6 +11,7 @@
from superagi.agent.tool_builder import ToolBuilder
from superagi.apm.event_handler import EventHandler
from superagi.config.config import get_config
+from superagi.helper.error_handler import ErrorHandler
from superagi.helper.token_counter import TokenCounter
from superagi.lib.logger import logger
from superagi.models.agent import Agent
@@ -73,6 +74,9 @@ def execute_step(self):
current_tokens = TokenCounter.count_message_tokens(messages = messages, model = self.llm.get_model())
response = self.llm.chat_completion(messages, TokenCounter(session=self.session, organisation_id=organisation.id).token_limit(self.llm.get_model()) - current_tokens)
+ if 'error' in response and response['message'] is not None:
+ ErrorHandler.handle_openai_errors(self.session, self.agent_id, self.agent_execution_id, response['message'])
+
if 'content' not in response or response['content'] is None:
raise RuntimeError(f"Failed to get response from llm")
@@ -82,11 +86,12 @@ def execute_step(self):
content = json.loads(response['content'])
tool = content.get('tool', {})
tool_name = tool.get('name', '') if tool else ''
- except Exception as e:
- logger.error(f"Decoding JSON has failed {e}")
+ except json.JSONDecodeError:
+ print("Decoding JSON has failed")
tool_name = ''
- CallLogHelper(session=self.session, organisation_id=organisation.id).create_call_log(execution.name,agent_config['agent_id'],total_tokens, tool_name,agent_config['model'])
+ CallLogHelper(session=self.session, organisation_id=organisation.id).create_call_log(execution.name,
+ agent_config['agent_id'], total_tokens, tool_name, agent_config['model'])
assistant_reply = response['content']
output_handler = get_output_handler(iteration_workflow_step.output_type,
diff --git a/superagi/agent/agent_message_builder.py b/superagi/agent/agent_message_builder.py
index 9008ef27c9..63698b8748 100644
--- a/superagi/agent/agent_message_builder.py
+++ b/superagi/agent/agent_message_builder.py
@@ -3,6 +3,7 @@
from sqlalchemy import asc
from superagi.config.config import get_config
+from superagi.helper.error_handler import ErrorHandler
from superagi.helper.prompt_reader import PromptReader
from superagi.helper.token_counter import TokenCounter
from superagi.models.agent_execution import AgentExecution
@@ -122,6 +123,9 @@ def _build_ltm_summary(self, past_messages, output_token_limit) -> str:
{"role": "assistant", "content": ltm_prompt}]
ltm_summary = self.llm.chat_completion(msgs)
+ if 'error' in ltm_summary and ltm_summary['message'] is not None:
+ ErrorHandler.handle_openai_errors(self.session, self.agent_id, self.agent_execution_id, ltm_summary['message'])
+
execution = AgentExecution(id=self.agent_execution_id)
agent_execution_configs = {"ltm_summary": ltm_summary["content"]}
AgentExecutionConfiguration.add_or_update_agent_execution_config(session=self.session, execution=execution,
diff --git a/superagi/agent/agent_tool_step_handler.py b/superagi/agent/agent_tool_step_handler.py
index 7aeb0d59b1..5b8c1c127c 100644
--- a/superagi/agent/agent_tool_step_handler.py
+++ b/superagi/agent/agent_tool_step_handler.py
@@ -7,6 +7,7 @@
from superagi.agent.output_parser import AgentSchemaToolOutputParser
from superagi.agent.queue_step_handler import QueueStepHandler
from superagi.agent.tool_builder import ToolBuilder
+from superagi.helper.error_handler import ErrorHandler
from superagi.helper.prompt_reader import PromptReader
from superagi.helper.token_counter import TokenCounter
from superagi.lib.logger import logger
@@ -105,6 +106,9 @@ def _process_input_instruction(self, agent_config, agent_execution_config, step_
# print(messages)
current_tokens = TokenCounter.count_message_tokens(messages, self.llm.get_model())
response = self.llm.chat_completion(messages, TokenCounter(session=self.session, organisation_id=self.organisation.id).token_limit(self.llm.get_model()) - current_tokens)
+
+ if 'error' in response and response['message'] is not None:
+ ErrorHandler.handle_openai_errors(self.session, self.agent_id, self.agent_execution_id, response['message'])
# ModelsHelper(session=self.session, organisation_id=organisation.id).create_call_log(execution.name,agent_config['agent_id'],response['response'].usage.total_tokens,json.loads(response['content'])['tool']['name'],agent_config['model'])
if 'content' not in response or response['content'] is None:
raise RuntimeError(f"Failed to get response from llm")
@@ -137,6 +141,10 @@ def _process_output_instruction(self, final_response: str, step_tool: AgentWorkf
current_tokens = TokenCounter.count_message_tokens(messages, self.llm.get_model())
response = self.llm.chat_completion(messages,
TokenCounter(session=self.session, organisation_id=self.organisation.id).token_limit(self.llm.get_model()) - current_tokens)
+
+ if 'error' in response and response['message'] is not None:
+ ErrorHandler.handle_openai_errors(self.session, self.agent_id, self.agent_execution_id, response['message'])
+
if 'content' not in response or response['content'] is None:
raise RuntimeError(f"ToolWorkflowStepHandler: Failed to get output response from llm")
total_tokens = current_tokens + TokenCounter.count_message_tokens(response, self.llm.get_model())
diff --git a/superagi/agent/agent_workflow_step_wait_handler.py b/superagi/agent/agent_workflow_step_wait_handler.py
new file mode 100644
index 0000000000..223859d31c
--- /dev/null
+++ b/superagi/agent/agent_workflow_step_wait_handler.py
@@ -0,0 +1,45 @@
+from datetime import datetime
+
+from superagi.agent.types.agent_execution_status import AgentExecutionStatus
+from superagi.lib.logger import logger
+from superagi.models.agent_execution import AgentExecution
+from superagi.models.workflows.agent_workflow_step import AgentWorkflowStep
+from superagi.models.workflows.agent_workflow_step_wait import AgentWorkflowStepWait
+from superagi.agent.types.wait_step_status import AgentWorkflowStepWaitStatus
+
+class AgentWaitStepHandler:
+ """Handle Agent Wait Step in the agent workflow."""
+
+ def __init__(self, session, agent_id, agent_execution_id):
+ self.session = session
+ self.agent_id = agent_id
+ self.agent_execution_id = agent_execution_id
+
+ def execute_step(self):
+ """Execute the agent wait step."""
+
+ logger.info("Executing Wait Step")
+ execution = AgentExecution.get_agent_execution_from_id(self.session, self.agent_execution_id)
+ workflow_step = AgentWorkflowStep.find_by_id(self.session, execution.current_agent_step_id)
+ step_wait = AgentWorkflowStepWait.find_by_id(self.session, workflow_step.action_reference_id)
+ if step_wait is not None:
+ step_wait.wait_begin_time = datetime.now()
+ step_wait.status = AgentWorkflowStepWaitStatus.WAITING.value
+ execution.status = AgentExecutionStatus.WAIT_STEP.value
+
+ self.session.commit()
+
+ def handle_next_step(self):
+ """Handle next step of agent workflow in case of wait step."""
+
+ execution = AgentExecution.get_agent_execution_from_id(self.session, self.agent_execution_id)
+ workflow_step = AgentWorkflowStep.find_by_id(self.session, execution.current_agent_step_id)
+ step_response = "default"
+ next_step = AgentWorkflowStep.fetch_next_step(self.session, workflow_step.id, step_response)
+ if str(next_step) == "COMPLETE":
+ agent_execution = AgentExecution.get_agent_execution_from_id(self.session, self.agent_execution_id)
+ agent_execution.current_agent_step_id = -1
+ agent_execution.status = "COMPLETED"
+ else:
+ AgentExecution.assign_next_step_id(self.session, self.agent_execution_id, next_step.id)
+ self.session.commit()
diff --git a/superagi/agent/queue_step_handler.py b/superagi/agent/queue_step_handler.py
index 1ce3d8313d..fcd9baf1f4 100644
--- a/superagi/agent/queue_step_handler.py
+++ b/superagi/agent/queue_step_handler.py
@@ -4,6 +4,7 @@
from superagi.agent.agent_message_builder import AgentLlmMessageBuilder
from superagi.agent.task_queue import TaskQueue
+from superagi.helper.error_handler import ErrorHandler
from superagi.helper.json_cleaner import JsonCleaner
from superagi.helper.prompt_reader import PromptReader
from superagi.helper.token_counter import TokenCounter
@@ -90,6 +91,10 @@ def _process_input_instruction(self, step_tool):
completion_prompt=step_tool.completion_prompt)
current_tokens = TokenCounter.count_message_tokens(messages, self.llm.get_model())
response = self.llm.chat_completion(messages, TokenCounter(session=self.session, organisation_id=self.organisation.id).token_limit(self.llm.get_model()) - current_tokens)
+
+ if 'error' in response and response['message'] is not None:
+ ErrorHandler.handle_openai_errors(self.session, self.agent_id, self.agent_execution_id, response['message'])
+
if 'content' not in response or response['content'] is None:
raise RuntimeError(f"Failed to get response from llm")
total_tokens = current_tokens + TokenCounter.count_message_tokens(response, self.llm.get_model())
diff --git a/superagi/agent/tool_executor.py b/superagi/agent/tool_executor.py
index 0170941644..303b8364a0 100644
--- a/superagi/agent/tool_executor.py
+++ b/superagi/agent/tool_executor.py
@@ -32,7 +32,7 @@ def execute(self, session, tool_name, tool_args):
status = "SUCCESS"
tool = tools[tool_name]
retry = False
- EventHandler(session=session).create_event('tool_used', {'tool_name': tool_name, 'agent_execution_id': self.agent_execution_id}, self.agent_id,
+ EventHandler(session=session).create_event('tool_used', {'tool_name': tool.name, 'agent_execution_id': self.agent_execution_id}, self.agent_id,
self.organisation_id),
try:
parsed_args = self.clean_tool_args(tool_args)
diff --git a/superagi/agent/types/__init__.py b/superagi/agent/types/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/superagi/agent/types/agent_execution_status.py b/superagi/agent/types/agent_execution_status.py
new file mode 100644
index 0000000000..c327c56cef
--- /dev/null
+++ b/superagi/agent/types/agent_execution_status.py
@@ -0,0 +1,18 @@
+from enum import Enum
+
+
+class AgentExecutionStatus(Enum):
+ RUNNING = 'RUNNING'
+ WAITING_FOR_PERMISSION = 'WAITING_FOR_PERMISSION'
+ ITERATION_LIMIT_EXCEEDED = 'ITERATION_LIMIT_EXCEEDED'
+ WAIT_STEP = 'WAIT_STEP'
+ COMPLETED = 'COMPLETED'
+
+ @classmethod
+ def get_agent_execution_status(cls, store):
+ if store is None:
+ raise ValueError("Storage type cannot be None.")
+ store = store.upper()
+ if store in cls.__members__:
+ return cls[store]
+ raise ValueError(f"{store} is not a valid storage name.")
diff --git a/superagi/agent/types/agent_workflow_step_action_types.py b/superagi/agent/types/agent_workflow_step_action_types.py
new file mode 100644
index 0000000000..dcce25fc13
--- /dev/null
+++ b/superagi/agent/types/agent_workflow_step_action_types.py
@@ -0,0 +1,17 @@
+from enum import Enum
+
+
+class AgentWorkflowStepAction(Enum):
+ ITERATION_WORKFLOW = 'ITERATION_WORKFLOW'
+ TOOL = 'TOOL'
+ WAIT_STEP = 'WAIT_STEP'
+
+
+ @classmethod
+ def get_agent_workflow_action_type(cls, store):
+ if store is None:
+ raise ValueError("Storage type cannot be None.")
+ store = store.upper()
+ if store in cls.__members__:
+ return cls[store]
+ raise ValueError(f"{store} is not a valid storage name.")
diff --git a/superagi/agent/types/wait_step_status.py b/superagi/agent/types/wait_step_status.py
new file mode 100644
index 0000000000..b81e72db1c
--- /dev/null
+++ b/superagi/agent/types/wait_step_status.py
@@ -0,0 +1,16 @@
+from enum import Enum
+
+
+class AgentWorkflowStepWaitStatus(Enum):
+ PENDING = 'PENDING'
+ WAITING = 'WAITING'
+ COMPLETED = 'COMPLETED'
+
+ @classmethod
+ def get_agent_workflow_step_wait_status(cls, store):
+ if store is None:
+ raise ValueError("Storage type cannot be None.")
+ store = store.upper()
+ if store in cls.__members__:
+ return cls[store]
+ raise ValueError(f"{store} is not a valid storage name.")
diff --git a/superagi/agent/workflow_seed.py b/superagi/agent/workflow_seed.py
index 7f95b1244a..2400faac86 100644
--- a/superagi/agent/workflow_seed.py
+++ b/superagi/agent/workflow_seed.py
@@ -37,16 +37,16 @@ def build_sales_workflow(cls, session):
# "Search for leads based on the given goals",
# step_type="TRIGGER")
#
- # step2 = AgentWorkflowStep.find_or_create_tool_workflow_step(session, agent_workflow.id,
- # str(agent_workflow.id) + "_step2",
- # WriteFileTool().name,
- # "Write the leads to a csv file")
+ step2 = AgentWorkflowStep.find_or_create_tool_workflow_step(session, agent_workflow.id,
+ str(agent_workflow.id) + "_step2",
+ ListFileTool().name,
+ "list the files",
+ step_type="TRIGGER")
step3 = AgentWorkflowStep.find_or_create_tool_workflow_step(session, agent_workflow.id,
str(agent_workflow.id) + "_step3",
ReadFileTool().name,
- "Read the leads from the file generated in the previous run",
- step_type="TRIGGER")
+ "Read the leads from the file")
# task queue ends when the elements gets over
step4 = AgentWorkflowStep.find_or_create_tool_workflow_step(session, agent_workflow.id,
@@ -67,7 +67,7 @@ def build_sales_workflow(cls, session):
step7 = AgentWorkflowStep.find_or_create_tool_workflow_step(session, agent_workflow.id,
str(agent_workflow.id) + "_step7",
- GoogleSearchTool().name,
+ SearxSearchTool().name,
"Search about the company given in the high-end goal only")
step8 = AgentWorkflowStep.find_or_create_tool_workflow_step(session, agent_workflow.id,
@@ -75,8 +75,23 @@ def build_sales_workflow(cls, session):
SendEmailTool().name,
"Customize the Email according to the company information in the mail")
+ step9 = AgentWorkflowStep.find_or_create_wait_workflow_step(session, agent_workflow.id,
+ str(agent_workflow.id) + "_step9",
+ "Wait for 2 minutes",
+ 2*60)
+
+ step10 = AgentWorkflowStep.find_or_create_tool_workflow_step(session, agent_workflow.id,
+ str(agent_workflow.id) + "_step10",
+ ReadEmailTool().name,
+ "Read the email from adarshdeepmurari@gmail.com")
+
+ step11 = AgentWorkflowStep.find_or_create_tool_workflow_step(session, agent_workflow.id,
+ str(agent_workflow.id) + "_step11",
+ SendEmailTool().name,
+ "Customize the Email according to the company information in the mail")
+
# AgentWorkflowStep.add_next_workflow_step(session, step1.id, step2.id)
- # AgentWorkflowStep.add_next_workflow_step(session, step2.id, step3.id)
+ AgentWorkflowStep.add_next_workflow_step(session, step2.id, step3.id)
AgentWorkflowStep.add_next_workflow_step(session, step3.id, step4.id)
AgentWorkflowStep.add_next_workflow_step(session, step4.id, -1, "COMPLETE")
AgentWorkflowStep.add_next_workflow_step(session, step4.id, step5.id)
@@ -84,7 +99,10 @@ def build_sales_workflow(cls, session):
AgentWorkflowStep.add_next_workflow_step(session, step6.id, step7.id, "YES")
AgentWorkflowStep.add_next_workflow_step(session, step6.id, step5.id, "NO")
AgentWorkflowStep.add_next_workflow_step(session, step7.id, step8.id)
- AgentWorkflowStep.add_next_workflow_step(session, step8.id, step4.id)
+ AgentWorkflowStep.add_next_workflow_step(session, step8.id, step9.id)
+ AgentWorkflowStep.add_next_workflow_step(session, step9.id, step10.id)
+ AgentWorkflowStep.add_next_workflow_step(session, step10.id, step11.id)
+ AgentWorkflowStep.add_next_workflow_step(session, step11.id, step4.id)
session.commit()
@classmethod
diff --git a/superagi/apm/call_log_helper.py b/superagi/apm/call_log_helper.py
index 85d359d4ef..03f183c0fe 100644
--- a/superagi/apm/call_log_helper.py
+++ b/superagi/apm/call_log_helper.py
@@ -50,13 +50,11 @@ def fetch_data(self, model: str):
'runs': []
}
- # Fetch all runs for this model
- runs = self.session.query(CallLogs).filter(CallLogs.model == model, CallLogs.org_id == self.organisation_id).all()
+ runs = self.session.query(CallLogs).filter(CallLogs.model == model,
+ CallLogs.org_id == self.organisation_id).all()
for run in runs:
- # Get agent's name using agent_id as a foreign key
agent = self.session.query(Agent).filter(Agent.id == run.agent_id).first()
- # Get toolkit's name using tool_used as a linking key
toolkit = None
tool = self.session.query(Tool).filter(Tool.name == run.tool_used).first()
if tool:
@@ -66,15 +64,17 @@ def fetch_data(self, model: str):
'id': run.id,
'agent_execution_name': run.agent_execution_name,
'agent_id': run.agent_id,
- 'agent_name': agent.name if agent is not None else None, # add agent_name to dictionary
+ 'agent_name': agent.name if agent is not None else None,
'tokens_consumed': run.tokens_consumed,
'tool_used': run.tool_used,
- 'toolkit_name': toolkit.name if toolkit is not None else None, # add toolkit_name to dictionary
+ 'toolkit_name': toolkit.name if toolkit is not None else None,
'org_id': run.org_id,
'created_at': run.created_at,
'updated_at': run.updated_at,
})
+ model_data['runs'] = model_data['runs'][::-1]
+
return model_data
except SQLAlchemyError as err:
diff --git a/superagi/apm/knowledge_handler.py b/superagi/apm/knowledge_handler.py
index 098791d4cf..30347b01b6 100644
--- a/superagi/apm/knowledge_handler.py
+++ b/superagi/apm/knowledge_handler.py
@@ -7,6 +7,7 @@
from sqlalchemy.sql import func
from sqlalchemy.orm import aliased
from superagi.models.agent_config import AgentConfiguration
+from superagi.models.agent_execution_config import AgentExecutionConfiguration
import pytz
from datetime import datetime
@@ -43,7 +44,7 @@ def get_knowledge_usage_by_name(self, knowledge_name: str) -> Dict[str, Dict[str
'knowledge_calls': self.session.query(
EventAlias
).filter(
- EventAlias.event_property['tool_name'].astext == 'knowledgesearch',
+ EventAlias.event_property['tool_name'].astext == 'Knowledge Search',
EventAlias.event_name == 'tool_used',
EventAlias.org_id == self.organisation_id,
EventAlias.agent_id.in_(self.session.query(Event.agent_id).filter(
@@ -89,6 +90,16 @@ def get_knowledge_events_by_name(self, knowledge_name: str) -> List[Dict[str, Un
event_run = next((er for er in event_runs if er.agent_id == knowledge_event.agent_id and er.event_property['agent_execution_id'] == agent_execution_id), None)
agent_created_event = next((ace for ace in agent_created_events if ace.agent_id == knowledge_event.agent_id), None)
+
+ model_query = self.session.query(AgentExecutionConfiguration).filter(
+ AgentExecutionConfiguration.agent_execution_id == agent_execution_id,
+ AgentExecutionConfiguration.key == 'model'
+ ).first()
+
+ if model_query and model_query.value != 'None':
+ model_value = model_query.value
+ else:
+ model_value = None
try:
user_timezone = AgentConfiguration.get_agent_config_by_key_and_agent_id(session=self.session, key='user_timezone', agent_id=knowledge_event.agent_id)
if user_timezone and user_timezone.value != 'None':
@@ -108,7 +119,7 @@ def get_knowledge_events_by_name(self, knowledge_name: str) -> List[Dict[str, Un
'calls': event_run.event_property['calls'],
'agent_execution_name': event_run.event_property['name'],
'agent_name': agent_created_event.event_property['agent_name'],
- 'model': agent_created_event.event_property['model']
+ 'model': model_value if model_value else agent_created_event.event_property['model']
}
if agent_execution_id not in [i['agent_execution_id'] for i in results]:
results.append(result_dict)
diff --git a/superagi/apm/tools_handler.py b/superagi/apm/tools_handler.py
index da3f97cc6e..0e583548dd 100644
--- a/superagi/apm/tools_handler.py
+++ b/superagi/apm/tools_handler.py
@@ -1,7 +1,7 @@
from typing import List, Dict, Union
from sqlalchemy import func, distinct, and_
from sqlalchemy.orm import Session
-from sqlalchemy import Integer
+from sqlalchemy import Integer, String
from fastapi import HTTPException
from superagi.models.events import Event
from superagi.models.tool import Tool
@@ -10,6 +10,7 @@
from sqlalchemy.sql import label
from datetime import datetime
from superagi.models.agent_config import AgentConfiguration
+from superagi.models.agent_execution_config import AgentExecutionConfiguration
import pytz
class ToolsHandler:
@@ -19,10 +20,10 @@ def __init__(self, session: Session, organisation_id: int):
def get_tool_and_toolkit(self):
tools_and_toolkits = self.session.query(
- Tool.name.label('tool_name'), Toolkit.name.label('toolkit_name')).join(
+ func.lower(Tool.name).label('tool_name'), Toolkit.name.label('toolkit_name')).join(
Toolkit, Tool.toolkit_id == Toolkit.id).all()
- return {item.tool_name: item.toolkit_name for item in tools_and_toolkits}
+ return {item.tool_name.lower(): item.toolkit_name for item in tools_and_toolkits}
def calculate_tool_usage(self) -> List[Dict[str, int]]:
tool_usage = []
@@ -55,9 +56,11 @@ def calculate_tool_usage(self) -> List[Dict[str, int]]:
'tool_name': row.tool_name,
'unique_agents': row.unique_agents,
'total_usage': row.total_usage,
- 'toolkit': tool_and_toolkit.get(row.tool_name, None)
+ 'toolkit': tool_and_toolkit.get(row.tool_name.lower(), None)
} for row in result]
+ tool_usage.sort(key=lambda tool: tool['total_usage'], reverse=True)
+
return tool_usage
def get_tool_usage_by_name(self, tool_name: str) -> Dict[str, Dict[str, int]]:
@@ -65,30 +68,34 @@ def get_tool_usage_by_name(self, tool_name: str) -> Dict[str, Dict[str, int]]:
if not is_tool_name_valid:
raise HTTPException(status_code=404, detail="Tool not found")
- formatted_tool_name = tool_name.lower().replace(" ", "")
- tool_used_event = self.session.query(
- Event.event_property['tool_name'].label('tool_name'),
+ tool_name_event = self.session.query(
+ Event.event_property['tool_name'].cast(String).label('tool_name'),
func.count(Event.id).label('tool_calls'),
func.count(distinct(Event.agent_id)).label('tool_unique_agents')
).filter(
- Event.event_name == 'tool_used',
+ Event.event_name == 'tool_used',
Event.org_id == self.organisation_id,
- Event.event_property['tool_name'].astext == formatted_tool_name
+ Event.event_property['tool_name'].astext == tool_name
).group_by(
- Event.event_property['tool_name']
+ Event.event_property['tool_name'].cast(String)
).first()
- if tool_used_event is None:
- return {}
+ tool_data = {}
+ tool_calls = 0
+ tool_unique_agents = 0
+
+ if tool_name_event:
+ tool_calls += tool_name_event.tool_calls
+ tool_unique_agents += tool_name_event.tool_unique_agents
tool_data = {
- 'tool_calls': tool_used_event.tool_calls,
- 'tool_unique_agents': tool_used_event.tool_unique_agents
- }
+ 'tool_calls': tool_calls,
+ 'tool_unique_agents': tool_unique_agents
+ }
return tool_data
-
+
def get_tool_events_by_name(self, tool_name: str) -> List[Dict[str, Union[str, int, List[str]]]]:
is_tool_name_valid = self.session.query(Tool).filter_by(name=tool_name).first()
@@ -96,12 +103,10 @@ def get_tool_events_by_name(self, tool_name: str) -> List[Dict[str, Union[str, i
if not is_tool_name_valid:
raise HTTPException(status_code=404, detail="Tool not found")
- formatted_tool_name = tool_name.lower().replace(" ", "")
-
tool_events = self.session.query(Event).filter(
Event.org_id == self.organisation_id,
Event.event_name == 'tool_used',
- Event.event_property['tool_name'].astext == formatted_tool_name
+ Event.event_property['tool_name'].astext == tool_name
).all()
tool_events = [te for te in tool_events if 'agent_execution_id' in te.event_property]
@@ -123,6 +128,16 @@ def get_tool_events_by_name(self, tool_name: str) -> List[Dict[str, Union[str, i
event_run = next((er for er in event_runs if er.agent_id == tool_event.agent_id and er.event_property['agent_execution_id'] == agent_execution_id), None)
agent_created_event = next((ace for ace in agent_created_events if ace.agent_id == tool_event.agent_id), None)
+
+ model_query = self.session.query(AgentExecutionConfiguration).filter(
+ AgentExecutionConfiguration.agent_execution_id == agent_execution_id,
+ AgentExecutionConfiguration.key == 'model'
+ ).first()
+
+ if model_query and model_query.value != 'None':
+ model_value = model_query.value
+ else:
+ model_value = None
try:
user_timezone = AgentConfiguration.get_agent_config_by_key_and_agent_id(session=self.session, key='user_timezone', agent_id=tool_event.agent_id)
if user_timezone and user_timezone.value != 'None':
@@ -139,7 +154,7 @@ def get_tool_events_by_name(self, tool_name: str) -> List[Dict[str, Union[str, i
).filter(
Event.org_id == self.organisation_id,
Event.event_name == 'tool_used',
- Event.event_property['tool_name'].astext != formatted_tool_name,
+ Event.event_property['tool_name'].astext != tool_name,
Event.agent_id == tool_event.agent_id,
Event.id.between(tool_event.id, event_run.id)
).all()
@@ -154,7 +169,7 @@ def get_tool_events_by_name(self, tool_name: str) -> List[Dict[str, Union[str, i
'agent_execution_name': event_run.event_property['name'],
'other_tools': other_tools,
'agent_name': agent_created_event.event_property['agent_name'],
- 'model': agent_created_event.event_property['model']
+ 'model': model_value if model_value else agent_created_event.event_property['model']
}
if agent_execution_id not in [i['agent_execution_id'] for i in results]:
diff --git a/superagi/controllers/agent_execution_feed.py b/superagi/controllers/agent_execution_feed.py
index e92f22face..93cf47bb77 100644
--- a/superagi/controllers/agent_execution_feed.py
+++ b/superagi/controllers/agent_execution_feed.py
@@ -1,7 +1,9 @@
+import asyncio
from datetime import datetime
+import time
from typing import Optional
-from fastapi import APIRouter
+from fastapi import APIRouter, BackgroundTasks
from fastapi import HTTPException, Depends
from fastapi_jwt_auth import AuthJWT
from fastapi_sqlalchemy import db
@@ -16,6 +18,10 @@
from superagi.helper.feed_parser import parse_feed
from superagi.models.agent_execution import AgentExecution
from superagi.models.agent_execution_feed import AgentExecutionFeed
+from superagi.lib.logger import logger
+from superagi.agent.types.agent_workflow_step_action_types import AgentWorkflowStepAction
+from superagi.models.workflows.agent_workflow_step import AgentWorkflowStep
+from superagi.models.workflows.agent_workflow_step_wait import AgentWorkflowStepWait
import re
# from superagi.types.db import AgentExecutionFeedOut, AgentExecutionFeedIn
@@ -144,7 +150,7 @@ def update_agent_execution_feed(agent_execution_feed_id: int,
@router.get("/get/execution/{agent_execution_id}")
def get_agent_execution_feed(agent_execution_id: int,
- Authorize: AuthJWT = Depends(check_auth)):
+ Authorize: AuthJWT = Depends(check_auth)):
"""
Get agent execution feed with other execution details.
@@ -165,7 +171,17 @@ def get_agent_execution_feed(agent_execution_id: int,
asc(AgentExecutionFeed.created_at)).all()
# # parse json
final_feeds = []
+ error = ""
for feed in feeds:
+ if feed.error_message:
+ if (agent_execution.last_shown_error_id is None) or (feed.id > agent_execution.last_shown_error_id):
+ #new error occured
+ error = feed.error_message
+ agent_execution.last_shown_error_id = feed.id
+ agent_execution.status = "ERROR_PAUSED"
+ db.session.commit()
+ if feed.id == agent_execution.last_shown_error_id and agent_execution.status == "ERROR_PAUSED":
+ error = feed.error_message
if feed.feed != "" and re.search(r"The current time and date is\s(\w{3}\s\w{3}\s\s?\d{1,2}\s\d{2}:\d{2}:\d{2}\s\d{4})",feed.feed) == None :
final_feeds.append(parse_feed(feed))
@@ -183,13 +199,22 @@ def get_agent_execution_feed(agent_execution_id: int,
"tool_name": permission.tool_name,
"question": permission.question,
"user_feedback": permission.user_feedback,
- "time_difference":get_time_difference(permission.created_at,str(datetime.now()))
+ "time_difference": get_time_difference(permission.created_at, str(datetime.now()))
} for permission in execution_permissions
]
+
+ waiting_period = None
+
+ if agent_execution.status == AgentWorkflowStepAction.WAIT_STEP.value:
+ workflow_step = AgentWorkflowStep.find_by_id(db.session, agent_execution.current_agent_step_id)
+ waiting_period = (AgentWorkflowStepWait.find_by_id(db.session, workflow_step.action_reference_id)).delay
+
return {
"status": agent_execution.status,
"feeds": final_feeds,
- "permissions": permissions
+ "permissions": permissions,
+ "waiting_period": waiting_period,
+ "errors": error
}
diff --git a/superagi/controllers/models_controller.py b/superagi/controllers/models_controller.py
index 5aab5d2076..2521f9abd1 100644
--- a/superagi/controllers/models_controller.py
+++ b/superagi/controllers/models_controller.py
@@ -103,7 +103,7 @@ async def fetch_data(request: ModelName, organisation=Depends(get_user_organisat
@router.get("/get/list", status_code=200)
-def get_knowledge_list(page: int = 0, organisation=Depends(get_user_organisation)):
+def get_models_list(page: int = 0, organisation=Depends(get_user_organisation)):
"""
Get Marketplace Model list.
@@ -122,7 +122,7 @@ def get_knowledge_list(page: int = 0, organisation=Depends(get_user_organisation
@router.get("/marketplace/list/{page}", status_code=200)
-def get_marketplace_knowledge_list(page: int = 0):
+def get_marketplace_models_list(page: int = 0):
organisation_id = get_config("MARKETPLACE_ORGANISATION_ID")
if organisation_id is not None:
organisation_id = int(organisation_id)
diff --git a/superagi/controllers/types/models_types.py b/superagi/controllers/types/models_types.py
index 0a3b99ce1d..4f6b9cc15b 100644
--- a/superagi/controllers/types/models_types.py
+++ b/superagi/controllers/types/models_types.py
@@ -11,4 +11,4 @@ def get_models_types(cls, model_type):
model_type = model_type.upper()
if model_type in cls.__members__:
return cls[model_type]
- raise ValueError(f"{model_type} is not a valid storage name.")
+ raise ValueError(f"{model_type} is not a valid storage name.")
\ No newline at end of file
diff --git a/superagi/controllers/user.py b/superagi/controllers/user.py
index 19bf2c9354..c550fd8891 100644
--- a/superagi/controllers/user.py
+++ b/superagi/controllers/user.py
@@ -11,12 +11,14 @@
from superagi.models.user import User
from fastapi import APIRouter
-from superagi.helper.auth import check_auth
+from superagi.helper.auth import check_auth, get_current_user
from superagi.lib.logger import logger
+
# from superagi.types.db import UserBase, UserIn, UserOut
router = APIRouter()
+
class UserBase(BaseModel):
name: str
email: str
@@ -42,6 +44,7 @@ class UserIn(UserBase):
class Config:
orm_mode = True
+
# CRUD Operations
@router.post("/add", response_model=UserOut, status_code=201)
def create_user(user: UserIn,
@@ -126,3 +129,16 @@ def update_user(user_id: int,
db.session.commit()
return db_user
+
+
+@router.post("/first_login_source/{source}")
+def update_first_login_source(source: str, Authorize: AuthJWT = Depends(check_auth)):
+ """ Update first login source of the user """
+ user = get_current_user(Authorize)
+ # valid_sources = ['google', 'github', 'email']
+ if user.first_login_source is None or user.first_login_source == '':
+ user.first_login_source = source
+ db.session.commit()
+ db.session.flush()
+ logger.info("User : ",user)
+ return user
diff --git a/superagi/helper/error_handler.py b/superagi/helper/error_handler.py
new file mode 100644
index 0000000000..839472574c
--- /dev/null
+++ b/superagi/helper/error_handler.py
@@ -0,0 +1,10 @@
+from superagi.models.agent_execution import AgentExecution
+from superagi.models.agent_execution_feed import AgentExecutionFeed
+
+class ErrorHandler:
+
+ def handle_openai_errors(session, agent_id, agent_execution_id, error_message):
+ execution = session.query(AgentExecution).filter(AgentExecution.id == agent_execution_id).first()
+ agent_feed = AgentExecutionFeed(agent_execution_id=agent_execution_id, agent_id=agent_id, role="system", feed="", error_message=error_message, feed_group_id=execution.current_feed_group_id)
+ session.add(agent_feed)
+ session.commit()
\ No newline at end of file
diff --git a/superagi/helper/feed_parser.py b/superagi/helper/feed_parser.py
index c05e217b53..c02040079d 100644
--- a/superagi/helper/feed_parser.py
+++ b/superagi/helper/feed_parser.py
@@ -2,6 +2,7 @@
from datetime import datetime
from superagi.helper.time_helper import get_time_difference
+from superagi.lib.logger import logger
def parse_feed(feed):
@@ -40,7 +41,8 @@ def parse_feed(feed):
return {"role": "assistant", "feed": final_output, "updated_at": feed.updated_at,
"time_difference": feed.time_difference}
except Exception:
- return feed
+ return {"role": "assistant", "feed": feed.feed, "updated_at": feed.updated_at,
+ "time_difference": feed.time_difference}
if feed.role == "system":
final_output = feed.feed
@@ -48,5 +50,9 @@ def parse_feed(feed):
final_output = feed.feed.split("TOOLS:")[0]
return {"role": "system", "feed": final_output, "updated_at": feed.updated_at,
"time_difference": feed.time_difference}
-
+
+ if feed.role == "user":
+ return {"role": "user", "feed": feed.feed, "updated_at": feed.updated_at,
+ "time_difference": feed.time_difference}
+
return feed
diff --git a/superagi/helper/s3_helper.py b/superagi/helper/s3_helper.py
index a454db588c..c084bd1090 100644
--- a/superagi/helper/s3_helper.py
+++ b/superagi/helper/s3_helper.py
@@ -144,6 +144,7 @@ def list_files_from_s3(self, file_path):
logger.info(f"Listing files from s3 with prefix: {file_path}")
response = self.s3.list_objects_v2(Bucket=get_config("BUCKET_NAME"), Prefix=file_path)
if 'Contents' in response:
+ logger.info(response['Contents'])
file_list = [obj['Key'] for obj in response['Contents']]
return file_list
else:
diff --git a/superagi/helper/validate_csv.py b/superagi/helper/validate_csv.py
index 5e1c6fd464..882679ccf1 100644
--- a/superagi/helper/validate_csv.py
+++ b/superagi/helper/validate_csv.py
@@ -11,7 +11,7 @@ def correct_csv_encoding(file_path):
if encoding != 'utf-8':
data = []
with open(file_path, 'r', encoding=encoding) as f:
- reader = csv.reader(f, delimiter=';', quotechar='"')
+ reader = csv.reader(f, delimiter=',', quotechar='"')
for row in reader:
try:
data.append(row)
@@ -20,7 +20,6 @@ def correct_csv_encoding(file_path):
continue
df = pd.DataFrame(data)
-
df.to_csv(file_path, encoding='utf-8', index=False)
logger.info("File is converted to utf-8 encoding.")
else:
diff --git a/superagi/jobs/agent_executor.py b/superagi/jobs/agent_executor.py
index c3e52b3ba2..4418dd86e7 100644
--- a/superagi/jobs/agent_executor.py
+++ b/superagi/jobs/agent_executor.py
@@ -5,21 +5,28 @@
import superagi.worker
from superagi.agent.agent_iteration_step_handler import AgentIterationStepHandler
from superagi.agent.agent_tool_step_handler import AgentToolStepHandler
+from superagi.agent.agent_workflow_step_wait_handler import AgentWaitStepHandler
+from superagi.agent.types.wait_step_status import AgentWorkflowStepWaitStatus
from superagi.apm.event_handler import EventHandler
+from superagi.config.config import get_config
from superagi.lib.logger import logger
from superagi.llms.google_palm import GooglePalm
from superagi.llms.hugging_face import HuggingFace
from superagi.llms.replicate import Replicate
from superagi.llms.llm_model_factory import get_model
+from superagi.llms.replicate import Replicate
from superagi.models.agent import Agent
from superagi.models.agent_config import AgentConfiguration
from superagi.models.agent_execution import AgentExecution
from superagi.models.db import connect_db
from superagi.models.workflows.agent_workflow_step import AgentWorkflowStep
-from superagi.types.model_source_types import ModelSourceType
+from superagi.models.workflows.agent_workflow_step_wait import AgentWorkflowStepWait
from superagi.types.vector_store_types import VectorStoreType
from superagi.vector_store.embedding.openai import OpenAiEmbedding
from superagi.vector_store.vector_factory import VectorFactory
+from superagi.worker import execute_agent
+from superagi.agent.types.agent_workflow_step_action_types import AgentWorkflowStepAction
+from superagi.agent.types.agent_execution_status import AgentExecutionStatus
from superagi.vector_store.redis import Redis
from superagi.config.config import get_config
@@ -46,7 +53,7 @@ def execute_next_step(self, agent_execution_id):
agent = session.query(Agent).filter(Agent.id == agent_execution.agent_id).first()
agent_config = Agent.fetch_configuration(session, agent.id)
if agent.is_deleted or (
- agent_execution.status != "RUNNING" and agent_execution.status != "WAITING_FOR_PERMISSION"):
+ agent_execution.status != AgentExecutionStatus.RUNNING.value and agent_execution.status != AgentExecutionStatus.WAITING_FOR_PERMISSION.value):
logger.error(f"Agent execution stopped. {agent.id}: {agent_execution.status}")
return
@@ -56,7 +63,8 @@ def execute_next_step(self, agent_execution_id):
return
try:
- model_config = AgentConfiguration.get_model_api_key(session, agent_execution.agent_id, agent_config["model"])
+ model_config = AgentConfiguration.get_model_api_key(session, agent_execution.agent_id,
+ agent_config["model"])
model_api_key = model_config['api_key']
model_llm_source = model_config['provider']
except Exception as e:
@@ -66,9 +74,10 @@ def execute_next_step(self, agent_execution_id):
try:
memory = None
if "OpenAI" in model_llm_source:
- vector_store_type = VectorStoreType.get_vector_store_type(get_config("LTM_DB","Redis"))
+ vector_store_type = VectorStoreType.get_vector_store_type(get_config("LTM_DB", "Redis"))
memory = VectorFactory.get_vector_storage(vector_store_type, "super-agent-index1",
- AgentExecutor.get_embedding(model_llm_source, model_api_key))
+ AgentExecutor.get_embedding(model_llm_source,
+ model_api_key))
except Exception as e:
logger.info(f"Unable to setup the connection...{e}")
memory = None
@@ -76,21 +85,9 @@ def execute_next_step(self, agent_execution_id):
agent_workflow_step = session.query(AgentWorkflowStep).filter(
AgentWorkflowStep.id == agent_execution.current_agent_step_id).first()
try:
- if agent_workflow_step.action_type == "TOOL":
- tool_step_handler = AgentToolStepHandler(session,
- llm=get_model(model=agent_config["model"], api_key=model_api_key, organisation_id=organisation.id)
- , agent_id=agent.id, agent_execution_id=agent_execution_id,
- memory=memory)
- tool_step_handler.execute_step()
- elif agent_workflow_step.action_type == "ITERATION_WORKFLOW":
- iteration_step_handler = AgentIterationStepHandler(session,
- llm=get_model(model=agent_config["model"],
- api_key=model_api_key,
- organisation_id=organisation.id)
- , agent_id=agent.id,
- agent_execution_id=agent_execution_id, memory=memory)
- print(get_model(model=agent_config["model"], api_key=model_api_key, organisation_id=organisation.id))
- iteration_step_handler.execute_step()
+ self.__execute_workflow_step(agent, agent_config, agent_execution_id, agent_workflow_step, memory,
+ model_api_key, organisation, session)
+
except Exception as e:
logger.info("Exception in executing the step: {}".format(e))
superagi.worker.execute_agent.apply_async((agent_execution_id, datetime.now()), countdown=15)
@@ -107,6 +104,30 @@ def execute_next_step(self, agent_execution_id):
session.close()
engine.dispose()
+ def __execute_workflow_step(self, agent, agent_config, agent_execution_id, agent_workflow_step, memory,
+ model_api_key, organisation, session):
+ logger.info("Executing Workflow step : ", agent_workflow_step.action_type)
+ if agent_workflow_step.action_type == AgentWorkflowStepAction.TOOL.value:
+ tool_step_handler = AgentToolStepHandler(session,
+ llm=get_model(model=agent_config["model"], api_key=model_api_key,
+ organisation_id=organisation.id)
+ , agent_id=agent.id, agent_execution_id=agent_execution_id,
+ memory=memory)
+ tool_step_handler.execute_step()
+ elif agent_workflow_step.action_type == AgentWorkflowStepAction.ITERATION_WORKFLOW.value:
+ iteration_step_handler = AgentIterationStepHandler(session,
+ llm=get_model(model=agent_config["model"],
+ api_key=model_api_key,
+ organisation_id=organisation.id)
+ , agent_id=agent.id,
+ agent_execution_id=agent_execution_id, memory=memory)
+ print(get_model(model=agent_config["model"], api_key=model_api_key, organisation_id=organisation.id))
+ iteration_step_handler.execute_step()
+ elif agent_workflow_step.action_type == AgentWorkflowStepAction.WAIT_STEP.value:
+ (AgentWaitStepHandler(session=session, agent_id=agent.id,
+ agent_execution_id=agent_execution_id)
+ .execute_step())
+
@classmethod
def get_embedding(cls, model_source, model_api_key):
if "OpenAI" in model_source:
@@ -122,7 +143,7 @@ def get_embedding(cls, model_source, model_api_key):
def _check_for_max_iterations(self, session, organisation_id, agent_config, agent_execution_id):
db_agent_execution = session.query(AgentExecution).filter(AgentExecution.id == agent_execution_id).first()
if agent_config["max_iterations"] <= db_agent_execution.num_of_calls:
- db_agent_execution.status = "ITERATION_LIMIT_EXCEEDED"
+ db_agent_execution.status = AgentExecutionStatus.ITERATION_LIMIT_EXCEEDED.value
EventHandler(session=session).create_event('run_iteration_limit_crossed',
{'agent_execution_id': db_agent_execution.id,
@@ -133,4 +154,33 @@ def _check_for_max_iterations(self, session, organisation_id, agent_config, agen
session.commit()
logger.info("ITERATION_LIMIT_CROSSED")
return True
- return False
\ No newline at end of file
+ return False
+
+ def execute_waiting_workflows(self):
+ """Check if wait time of wait workflow step is over and can be resumed."""
+
+ session = Session()
+ waiting_agent_executions = session.query(AgentExecution).filter(
+ AgentExecution.status == AgentExecutionStatus.WAIT_STEP.value,
+ ).all()
+ for agent_execution in waiting_agent_executions:
+ workflow_step = session.query(AgentWorkflowStep).filter(
+ AgentWorkflowStep.id == agent_execution.current_agent_step_id).first()
+ step_wait = AgentWorkflowStepWait.find_by_id(session, workflow_step.action_reference_id)
+ if step_wait is not None:
+ wait_time = step_wait.delay if not None else 0
+ logger.info(f"Agent Execution ID: {agent_execution.id}")
+ logger.info(f"Wait time: {wait_time}")
+ logger.info(f"Wait begin time: {step_wait.wait_begin_time}")
+ logger.info(f"Current time: {datetime.now()}")
+ logger.info(f"Wait Difference : {(datetime.now() - step_wait.wait_begin_time).total_seconds()}")
+ if ((datetime.now() - step_wait.wait_begin_time).total_seconds() > wait_time
+ and step_wait.status == AgentWorkflowStepWaitStatus.WAITING.value):
+ agent_execution.status = AgentExecutionStatus.RUNNING.value
+ step_wait.status = AgentWorkflowStepWaitStatus.COMPLETED.value
+ session.commit()
+ session.flush()
+ AgentWaitStepHandler(session=session, agent_id=agent_execution.agent_id,
+ agent_execution_id=agent_execution.id).handle_next_step()
+ execute_agent.delay(agent_execution.id, datetime.now())
+ session.close()
diff --git a/superagi/llms/openai.py b/superagi/llms/openai.py
index b8a97a44a8..454f7ce191 100644
--- a/superagi/llms/openai.py
+++ b/superagi/llms/openai.py
@@ -77,16 +77,16 @@ def chat_completion(self, messages, max_tokens=get_config("MAX_MODEL_TOKEN_LIMIT
return {"response": response, "content": content}
except AuthenticationError as auth_error:
logger.info("OpenAi AuthenticationError:", auth_error)
- return {"error": "ERROR_AUTHENTICATION", "message": "Authentication error please check the api keys.."}
+ return {"error": "ERROR_AUTHENTICATION", "message": "Authentication error please check the api keys: "+str(auth_error)}
except RateLimitError as api_error:
logger.info("OpenAi RateLimitError:", api_error)
- return {"error": "ERROR_RATE_LIMIT", "message": "Openai rate limit exceeded.."}
+ return {"error": "ERROR_RATE_LIMIT", "message": "Openai rate limit exceeded: "+str(api_error)}
except InvalidRequestError as invalid_request_error:
logger.info("OpenAi InvalidRequestError:", invalid_request_error)
- return {"error": "ERROR_INVALID_REQUEST", "message": "Openai invalid request error.."}
+ return {"error": "ERROR_INVALID_REQUEST", "message": "Openai invalid request error: "+str(invalid_request_error)}
except Exception as exception:
logger.info("OpenAi Exception:", exception)
- return {"error": "ERROR_OPENAI", "message": "Open ai exception"}
+ return {"error": "ERROR_OPENAI", "message": "Open ai exception: "+str(exception)}
def verify_access_key(self):
"""
diff --git a/superagi/models/agent_execution.py b/superagi/models/agent_execution.py
index 5cba5f5092..73e1f8f15c 100644
--- a/superagi/models/agent_execution.py
+++ b/superagi/models/agent_execution.py
@@ -37,6 +37,7 @@ class AgentExecution(DBBaseModel):
permission_id = Column(Integer)
iteration_workflow_step_id = Column(Integer)
current_feed_group_id = Column(String)
+ last_shown_error_id = Column(Integer)
def __repr__(self):
"""
diff --git a/superagi/models/agent_execution_feed.py b/superagi/models/agent_execution_feed.py
index 61f4495856..84cc4fbc1a 100644
--- a/superagi/models/agent_execution_feed.py
+++ b/superagi/models/agent_execution_feed.py
@@ -27,6 +27,7 @@ class AgentExecutionFeed(DBBaseModel):
role = Column(String)
extra_info = Column(String)
feed_group_id = Column(String)
+ error_message = Column(String)
def __repr__(self):
"""
diff --git a/superagi/models/user.py b/superagi/models/user.py
index 5190957566..6ef8bb6940 100644
--- a/superagi/models/user.py
+++ b/superagi/models/user.py
@@ -24,6 +24,7 @@ class User(DBBaseModel):
email = Column(String, unique=True)
password = Column(String)
organisation_id = Column(Integer)
+ first_login_source = Column(String)
def __repr__(self):
"""
@@ -34,4 +35,4 @@ def __repr__(self):
"""
return f"User(id={self.id}, name='{self.name}', email='{self.email}', password='{self.password}'," \
- f"organisation_id={self.organisation_id})"
+ f"organisation_id={self.organisation_id}, first_login_source={self.first_login_source})"
diff --git a/superagi/models/workflows/agent_workflow_step.py b/superagi/models/workflows/agent_workflow_step.py
index 49af055d2c..b7a3ada160 100644
--- a/superagi/models/workflows/agent_workflow_step.py
+++ b/superagi/models/workflows/agent_workflow_step.py
@@ -1,12 +1,12 @@
import json
-from sqlalchemy import Column, Integer, String, Text, Boolean
+from sqlalchemy import Column, Integer, String
from sqlalchemy.dialects.postgresql import JSONB
from superagi.lib.logger import logger
-from superagi.models import db
from superagi.models.base_model import DBBaseModel
from superagi.models.workflows.agent_workflow_step_tool import AgentWorkflowStepTool
+from superagi.models.workflows.agent_workflow_step_wait import AgentWorkflowStepWait
from superagi.models.workflows.iteration_workflow import IterationWorkflow
@@ -30,9 +30,9 @@ class AgentWorkflowStep(DBBaseModel):
agent_workflow_id = Column(Integer)
unique_id = Column(String)
step_type = Column(String) # TRIGGER, NORMAL
- action_type = Column(String) # TOOL, ITERATION_WORKFLOW, LLM
+ action_type = Column(String) # TOOL, ITERATION_WORKFLOW, LLM, WAIT_STEP
action_reference_id = Column(Integer) # id of the action
- next_steps = Column(JSONB) # edge_ref_id, response, step_id
+ next_steps = Column(JSONB) # edge_ref_id, response, step_id
def __repr__(self):
"""
@@ -45,7 +45,7 @@ def __repr__(self):
return f"AgentWorkflowStep(id={self.id}, status='{self.agent_workflow_id}', " \
f"prompt='{self.unique_id}', agent_id={self.step_type}, action_type={self.action_type}, " \
f"action_reference_id={self.action_reference_id}, next_steps={self.next_steps})"
-
+
def to_dict(self):
"""
Converts the AgentWorkflowStep object to a dictionary.
@@ -150,6 +150,29 @@ def find_or_create_tool_workflow_step(cls, session, agent_workflow_id: int, uniq
session.commit()
return workflow_step
+ @classmethod
+ def find_or_create_wait_workflow_step(cls, session, agent_workflow_id: int, unique_id: str,
+ wait_description: str, delay: int, step_type="NORMAL"):
+ """ Find or create a wait workflow step"""
+ logger.info("Finding or creating wait step")
+ workflow_step = session.query(AgentWorkflowStep).filter(
+ AgentWorkflowStep.agent_workflow_id == agent_workflow_id, AgentWorkflowStep.unique_id == unique_id).first()
+
+ step_wait = AgentWorkflowStepWait.find_or_create_wait(session=session, step_unique_id=unique_id,
+ description=wait_description, delay=delay)
+ if workflow_step is None:
+ workflow_step = AgentWorkflowStep(unique_id=unique_id, step_type=step_type,
+ agent_workflow_id=agent_workflow_id)
+ session.add(workflow_step)
+ session.commit()
+ workflow_step.step_type = step_type
+ workflow_step.agent_workflow_id = agent_workflow_id
+ workflow_step.action_reference_id = step_wait.id
+ workflow_step.action_type = "WAIT_STEP"
+ workflow_step.next_steps = []
+ session.commit()
+ return workflow_step
+
@classmethod
def find_or_create_iteration_workflow_step(cls, session, agent_workflow_id: int, unique_id: str,
iteration_workflow_name: str, step_type="NORMAL"):
@@ -184,7 +207,8 @@ def find_or_create_iteration_workflow_step(cls, session, agent_workflow_id: int,
return workflow_step
@classmethod
- def add_next_workflow_step(cls, session, current_agent_step_id: int, next_step_id: int, step_response: str = "default"):
+ def add_next_workflow_step(cls, session, current_agent_step_id: int, next_step_id: int,
+ step_response: str = "default"):
""" Add Next workflow steps in the next_steps column
Args:
@@ -251,7 +275,3 @@ def fetch_next_step(cls, session, current_agent_step_id: int, step_response: str
return "COMPLETE"
return AgentWorkflowStep.find_by_unique_id(session, default_steps[0]["step_id"])
return None
-
-
-
-
diff --git a/superagi/models/workflows/agent_workflow_step_wait.py b/superagi/models/workflows/agent_workflow_step_wait.py
new file mode 100644
index 0000000000..f24540b1ff
--- /dev/null
+++ b/superagi/models/workflows/agent_workflow_step_wait.py
@@ -0,0 +1,89 @@
+import json
+
+from sqlalchemy import Column, Integer, String, DateTime
+
+from superagi.models.base_model import DBBaseModel
+
+
+class AgentWorkflowStepWait(DBBaseModel):
+ """
+ Step for a Agent Workflow to wait
+
+ Attributes:
+ id (int): The unique identifier of the wait block step.
+ name (str): The name of the wait block step.
+ description (str): The description of the wait block step.
+ delay (int): The delay time in seconds.
+ wait_begin_time (DateTime): The start time of the wait block.
+ """
+
+ __tablename__ = 'agent_workflow_step_waits'
+
+ id = Column(Integer, primary_key=True)
+ name = Column(String)
+ description = Column(String)
+ unique_id = Column(String)
+ delay = Column(Integer) # Delay is stored in seconds
+ wait_begin_time = Column(DateTime)
+ status = Column(String) # 'PENDING', 'WAITING', 'COMPLETED'
+
+ def __repr__(self):
+ """
+ Returns a string representation of the WaitBlockStep object.
+
+ Returns:
+ str: String representation of the WaitBlockStep.
+ """
+
+ return f"WaitBlockStep(id={self.id}, name='{self.name}', delay='{self.delay}', " \
+ f"wait_begin_time='{self.wait_begin_time}'"
+
+ def to_dict(self):
+ """
+ Converts the WaitBlockStep object to a dictionary.
+
+ Returns:
+ dict: Dictionary representation of the WaitBlockStep.
+ """
+
+ return {
+ 'id': self.id,
+ 'name': self.name,
+ 'delay': self.delay,
+ 'wait_begin_time': self.wait_begin_time
+ }
+
+ def to_json(self):
+ """
+ Converts the WaitBlockStep object to a JSON string.
+
+ Returns:
+ str: JSON string representation of the WaitBlockStep.
+ """
+
+ return json.dumps(self.to_dict())
+
+ @classmethod
+ def find_by_id(cls, session, step_id: int):
+ return session.query(AgentWorkflowStepWait).filter(AgentWorkflowStepWait.id == step_id).first()
+
+ @classmethod
+ def find_or_create_wait(cls, session, step_unique_id: str, description: str, delay: int):
+ unique_id = f"{step_unique_id}_wait"
+ wait = session.query(AgentWorkflowStepWait).filter(AgentWorkflowStepWait.unique_id == unique_id).first()
+ if wait is None:
+ wait = AgentWorkflowStepWait(
+ unique_id=unique_id,
+ name=unique_id,
+ delay=delay,
+ description=description,
+ status='PENDING'
+ )
+ session.add(wait)
+ else:
+ wait.delay = delay
+ wait.description = description
+ wait.status = 'PENDING'
+ session.commit()
+ session.flush()
+ return wait
diff --git a/superagi/tools/apollo/apollo_search.py b/superagi/tools/apollo/apollo_search.py
index d9951a4e5e..974638cc76 100644
--- a/superagi/tools/apollo/apollo_search.py
+++ b/superagi/tools/apollo/apollo_search.py
@@ -73,7 +73,7 @@ def _execute(self, person_titles: list[str], page: int = 1, per_page: int = 25,
num_of_employees, person_location, organization_domains)
logger.info(people_data)
people_list = []
- if 'people' in people_data and len(people_data['people']) > 0:
+ if people_data and 'people' in people_data and len(people_data['people']) > 0:
for person in people_data['people']:
people_list.append({'first_name': person['first_name'],
'last_name': person['last_name'],
@@ -130,3 +130,4 @@ def apollo_search_results(self, page, per_page, person_titles, num_of_employees
return response.json()
else:
return None
+
diff --git a/superagi/tools/code/improve_code.py b/superagi/tools/code/improve_code.py
index c54856811a..a825ba47ae 100644
--- a/superagi/tools/code/improve_code.py
+++ b/superagi/tools/code/improve_code.py
@@ -4,10 +4,13 @@
from pydantic import BaseModel, Field
from superagi.agent.agent_prompt_builder import AgentPromptBuilder
+from superagi.helper.error_handler import ErrorHandler
from superagi.helper.prompt_reader import PromptReader
from superagi.helper.token_counter import TokenCounter
from superagi.lib.logger import logger
from superagi.llms.base_llm import BaseLlm
+from superagi.models.agent_execution import AgentExecution
+from superagi.models.agent_execution_feed import AgentExecutionFeed
from superagi.resource_manager.file_manager import FileManager
from superagi.tools.base_tool import BaseTool
from superagi.tools.tool_response_query_manager import ToolResponseQueryManager
@@ -29,6 +32,7 @@ class ImproveCodeTool(BaseTool):
"""
llm: Optional[BaseLlm] = None
agent_id: int = None
+ agent_execution_id: int = None
name = "ImproveCodeTool"
description = (
"This tool improves the generated code."
@@ -71,6 +75,9 @@ def _execute(self) -> str:
# Use LLM to generate improved code
result = self.llm.chat_completion([{'role': 'system', 'content': prompt}])
+
+ if result is not None and 'error' in result and result['message'] is not None:
+ ErrorHandler.handle_openai_errors(self.toolkit_config.session, self.agent_id, self.agent_execution_id, result['message'])
# Extract the response first
response = result.get('response')
diff --git a/superagi/tools/code/write_code.py b/superagi/tools/code/write_code.py
index 87fdca1184..a799fc656c 100644
--- a/superagi/tools/code/write_code.py
+++ b/superagi/tools/code/write_code.py
@@ -4,10 +4,13 @@
from pydantic import BaseModel, Field
from superagi.agent.agent_prompt_builder import AgentPromptBuilder
+from superagi.helper.error_handler import ErrorHandler
from superagi.helper.prompt_reader import PromptReader
from superagi.helper.token_counter import TokenCounter
from superagi.lib.logger import logger
from superagi.llms.base_llm import BaseLlm
+from superagi.models.agent_execution import AgentExecution
+from superagi.models.agent_execution_feed import AgentExecutionFeed
from superagi.resource_manager.file_manager import FileManager
from superagi.tools.base_tool import BaseTool
from superagi.tools.tool_response_query_manager import ToolResponseQueryManager
@@ -34,6 +37,7 @@ class CodingTool(BaseTool):
"""
llm: Optional[BaseLlm] = None
agent_id: int = None
+ agent_execution_id: int = None
name = "CodingTool"
description = (
"You will get instructions for code to write. You will write a very long answer. "
@@ -75,6 +79,9 @@ def _execute(self, code_description: str) -> str:
token_limit = TokenCounter(session=self.toolkit_config.session, organisation_id=organisation.id).token_limit(self.llm.get_model())
result = self.llm.chat_completion(messages, max_tokens=(token_limit - total_tokens - 100))
+
+ if 'error' in result and result['message'] is not None:
+ ErrorHandler.handle_openai_errors(self.toolkit_config.session, self.agent_id, self.agent_execution_id, result['message'])
# Get all filenames and corresponding code blocks
regex = r"(\S+?)\n```\S*\n(.+?)```"
diff --git a/superagi/tools/code/write_spec.py b/superagi/tools/code/write_spec.py
index 91b5bd8434..9793d726e7 100644
--- a/superagi/tools/code/write_spec.py
+++ b/superagi/tools/code/write_spec.py
@@ -3,10 +3,13 @@
from pydantic import BaseModel, Field
from superagi.agent.agent_prompt_builder import AgentPromptBuilder
+from superagi.helper.error_handler import ErrorHandler
from superagi.helper.prompt_reader import PromptReader
from superagi.helper.token_counter import TokenCounter
from superagi.lib.logger import logger
from superagi.llms.base_llm import BaseLlm
+from superagi.models.agent_execution import AgentExecution
+from superagi.models.agent_execution_feed import AgentExecutionFeed
from superagi.resource_manager.file_manager import FileManager
from superagi.tools.base_tool import BaseTool
from superagi.models.agent import Agent
@@ -37,6 +40,7 @@ class WriteSpecTool(BaseTool):
"""
llm: Optional[BaseLlm] = None
agent_id: int = None
+ agent_execution_id: int = None
name = "WriteSpecTool"
description = (
"A tool to write the spec of a program."
@@ -69,6 +73,9 @@ def _execute(self, task_description: str, spec_file_name: str) -> str:
token_limit = TokenCounter(session=self.toolkit_config.session, organisation_id=organisation.id).token_limit(self.llm.get_model())
result = self.llm.chat_completion(messages, max_tokens=(token_limit - total_tokens - 100))
+
+ if 'error' in result and result['message'] is not None:
+ ErrorHandler.handle_openai_errors(self.toolkit_config.session, self.agent_id, self.agent_execution_id, result['message'])
# Save the specification to a file
write_result = self.resource_manager.write_file(spec_file_name, result["content"])
diff --git a/superagi/tools/code/write_test.py b/superagi/tools/code/write_test.py
index de0ece5349..3b446fc36d 100644
--- a/superagi/tools/code/write_test.py
+++ b/superagi/tools/code/write_test.py
@@ -4,10 +4,13 @@
from pydantic import BaseModel, Field
from superagi.agent.agent_prompt_builder import AgentPromptBuilder
+from superagi.helper.error_handler import ErrorHandler
from superagi.helper.prompt_reader import PromptReader
from superagi.helper.token_counter import TokenCounter
from superagi.lib.logger import logger
from superagi.llms.base_llm import BaseLlm
+from superagi.models.agent_execution import AgentExecution
+from superagi.models.agent_execution_feed import AgentExecutionFeed
from superagi.resource_manager.file_manager import FileManager
from superagi.tools.base_tool import BaseTool
from superagi.tools.tool_response_query_manager import ToolResponseQueryManager
@@ -38,6 +41,7 @@ class WriteTestTool(BaseTool):
"""
llm: Optional[BaseLlm] = None
agent_id: int = None
+ agent_execution_id: int = None
name = "WriteTestTool"
description = (
"You are a super smart developer using Test Driven Development to write tests according to a specification.\n"
@@ -86,6 +90,9 @@ def _execute(self, test_description: str, test_file_name: str) -> str:
token_limit = TokenCounter(session=self.toolkit_config.session, organisation_id=organisation.id).token_limit(self.llm.get_model())
result = self.llm.chat_completion(messages, max_tokens=(token_limit - total_tokens - 100))
+
+ if 'error' in result and result['message'] is not None:
+ ErrorHandler.handle_openai_errors(self.toolkit_config.session, self.agent_id, self.agent_execution_id, result['message'])
regex = r"(\S+?)\n```\S*\n(.+?)```"
matches = re.finditer(regex, result["content"], re.DOTALL)
diff --git a/superagi/tools/duck_duck_go/duck_duck_go_search.py b/superagi/tools/duck_duck_go/duck_duck_go_search.py
index 47223e4cc4..43c0d2ec09 100644
--- a/superagi/tools/duck_duck_go/duck_duck_go_search.py
+++ b/superagi/tools/duck_duck_go/duck_duck_go_search.py
@@ -2,12 +2,15 @@
import requests
from typing import Type, Optional,Union
import time
+from superagi.helper.error_handler import ErrorHandler
from superagi.lib.logger import logger
from pydantic import BaseModel, Field
from duckduckgo_search import DDGS
from itertools import islice
from superagi.helper.token_counter import TokenCounter
from superagi.llms.base_llm import BaseLlm
+from superagi.models.agent_execution import AgentExecution
+from superagi.models.agent_execution_feed import AgentExecutionFeed
from superagi.tools.base_tool import BaseTool
from superagi.helper.webpage_extractor import WebpageExtractor
@@ -33,6 +36,8 @@ class DuckDuckGoSearchTool(BaseTool):
"""
llm: Optional[BaseLlm] = None
name = "DuckDuckGoSearch"
+ agent_id: int = None
+ agent_execution_id: int = None
description = (
"A tool for performing a DuckDuckGo search and extracting snippets and webpages."
"Input should be a search query."
@@ -169,4 +174,7 @@ def summarise_result(self, query, snippets):
messages = [{"role": "system", "content": summarize_prompt}]
result = self.llm.chat_completion(messages, max_tokens=self.max_token_limit)
+
+ if 'error' in result and result['message'] is not None:
+ ErrorHandler.handle_openai_errors(self.toolkit_config.session, self.agent_id, self.agent_execution_id, result['message'])
return result["content"]
diff --git a/superagi/tools/github/review_pull_request.py b/superagi/tools/github/review_pull_request.py
index a230c07017..c76a460b47 100644
--- a/superagi/tools/github/review_pull_request.py
+++ b/superagi/tools/github/review_pull_request.py
@@ -2,6 +2,7 @@
from typing import Type, Optional
from pydantic import BaseModel, Field
+from superagi.helper.error_handler import ErrorHandler
from superagi.helper.github_helper import GithubHelper
from superagi.helper.json_cleaner import JsonCleaner
@@ -9,6 +10,8 @@
from superagi.helper.token_counter import TokenCounter
from superagi.llms.base_llm import BaseLlm
from superagi.models.agent import Agent
+from superagi.models.agent_execution import AgentExecution
+from superagi.models.agent_execution_feed import AgentExecutionFeed
from superagi.tools.base_tool import BaseTool
@@ -87,6 +90,9 @@ def run_code_review(self, github_helper, content, latest_commit_id, organisation
token_limit = TokenCounter(session=self.toolkit_config.session,
organisation_id=organisation.id).token_limit(self.llm.get_model())
result = self.llm.chat_completion(messages, max_tokens=(token_limit - total_tokens - 100))
+
+ if 'error' in result and result['message'] is not None:
+ ErrorHandler.handle_openai_errors(self.toolkit_config.session, self.agent_id, self.agent_execution_id, result['message'])
response = result["content"]
if response.startswith("```") and response.endswith("```"):
response = "```".join(response.split("```")[1:-1])
diff --git a/superagi/tools/google_search/google_search.py b/superagi/tools/google_search/google_search.py
index c0c7e875b2..9880266da2 100644
--- a/superagi/tools/google_search/google_search.py
+++ b/superagi/tools/google_search/google_search.py
@@ -2,13 +2,15 @@
from typing import Type, Optional
from pydantic import BaseModel, Field
+from superagi.helper.error_handler import ErrorHandler
from superagi.helper.google_search import GoogleSearchWrap
from superagi.helper.token_counter import TokenCounter
from superagi.llms.base_llm import BaseLlm
+from superagi.models.agent_execution import AgentExecution
+from superagi.models.agent_execution_feed import AgentExecutionFeed
from superagi.tools.base_tool import BaseTool
-
class GoogleSearchSchema(BaseModel):
query: str = Field(
...,
@@ -26,6 +28,8 @@ class GoogleSearchTool(BaseTool):
"""
llm: Optional[BaseLlm] = None
name = "GoogleSearch"
+ agent_id: int = None
+ agent_execution_id: int = None
description = (
"A tool for performing a Google search and extracting snippets and webpages."
"Input should be a search query."
@@ -88,4 +92,7 @@ def summarise_result(self, query, snippets):
messages = [{"role": "system", "content": summarize_prompt}]
result = self.llm.chat_completion(messages, max_tokens=self.max_token_limit)
+
+ if 'error' in result and result['message'] is not None:
+ ErrorHandler.handle_openai_errors(self.toolkit_config.session, self.agent_id, self.agent_execution_id, result['message'])
return result["content"]
diff --git a/superagi/tools/google_serp_search/google_serp_search.py b/superagi/tools/google_serp_search/google_serp_search.py
index e3b0cb4384..4f9d46c43b 100644
--- a/superagi/tools/google_serp_search/google_serp_search.py
+++ b/superagi/tools/google_serp_search/google_serp_search.py
@@ -1,8 +1,11 @@
from typing import Type, Optional, Any
from pydantic import BaseModel, Field
import aiohttp
+from superagi.helper.error_handler import ErrorHandler
from superagi.helper.google_serp import GoogleSerpApiWrap
from superagi.llms.base_llm import BaseLlm
+from superagi.models.agent_execution import AgentExecution
+from superagi.models.agent_execution_feed import AgentExecutionFeed
from superagi.tools.base_tool import BaseTool
import os
@@ -28,6 +31,8 @@ class GoogleSerpTool(BaseTool):
"""
llm: Optional[BaseLlm] = None
name = "GoogleSerp"
+ agent_id: int = None
+ agent_execution_id: int = None
description = (
"A tool for performing a Google SERP search and extracting snippets and webpages."
"Input should be a search query."
@@ -66,4 +71,7 @@ def summarise_result(self, query, snippets):
messages = [{"role": "system", "content": summarize_prompt}]
result = self.llm.chat_completion(messages, max_tokens=self.max_token_limit)
+
+ if 'error' in result and result['message'] is not None:
+ ErrorHandler.handle_openai_errors(self.toolkit_config.session, self.agent_id, self.agent_execution_id, result['message'])
return result["content"]
diff --git a/superagi/tools/instagram_tool/instagram.py b/superagi/tools/instagram_tool/instagram.py
index bc423b6e90..5d9e68b357 100644
--- a/superagi/tools/instagram_tool/instagram.py
+++ b/superagi/tools/instagram_tool/instagram.py
@@ -3,11 +3,13 @@
import boto3
import os
from superagi.config.config import get_config
+from superagi.helper.error_handler import ErrorHandler
from superagi.helper.resource_helper import ResourceHelper
from typing import Type, Optional
from pydantic import BaseModel, Field
from superagi.helper.token_counter import TokenCounter
from superagi.llms.base_llm import BaseLlm
+from superagi.models.agent_execution_feed import AgentExecutionFeed
from superagi.tools.base_tool import BaseTool
import os
import requests
@@ -114,6 +116,8 @@ def create_caption(self, photo_description: str) -> str:
messages = [{"role": "system", "content": caption_prompt}]
result = self.llm.chat_completion(messages, max_tokens=self.max_token_limit)
+ if 'error' in result and result['message'] is not None:
+ ErrorHandler.handle_openai_errors(self.toolkit_config.session, self.agent_id, self.agent_execution_id, result['message'])
caption=result["content"]
encoded_caption=urllib. parse. quote(caption)
diff --git a/superagi/tools/searx/searx.py b/superagi/tools/searx/searx.py
index 49e09a072b..6ba4555116 100644
--- a/superagi/tools/searx/searx.py
+++ b/superagi/tools/searx/searx.py
@@ -1,6 +1,9 @@
from typing import Type, Optional
from pydantic import BaseModel, Field
+from superagi.helper.error_handler import ErrorHandler
from superagi.llms.base_llm import BaseLlm
+from superagi.models.agent_execution import AgentExecution
+from superagi.models.agent_execution_feed import AgentExecutionFeed
from superagi.tools.base_tool import BaseTool
from superagi.tools.searx.search_scraper import search_results
@@ -22,6 +25,8 @@ class SearxSearchTool(BaseTool):
"""
llm: Optional[BaseLlm] = None
name = "SearxSearch"
+ agent_id:int =None
+ agent_execution_id:int =None
description = (
"A tool for performing a Searx search and extracting snippets and webpages."
"Input should be a search query."
@@ -67,4 +72,7 @@ def summarise_result(self, query, snippets):
messages = [{"role": "system", "content": summarize_prompt}]
result = self.llm.chat_completion(messages, max_tokens=self.max_token_limit)
+
+ if 'error' in result and result['message'] is not None:
+ ErrorHandler.handle_openai_errors(self.toolkit_config.session, self.agent_id, self.agent_execution_id, result['message'])
return result["content"]
diff --git a/superagi/tools/thinking/tools.py b/superagi/tools/thinking/tools.py
index 8bdc38ca56..370eeee4f6 100644
--- a/superagi/tools/thinking/tools.py
+++ b/superagi/tools/thinking/tools.py
@@ -3,9 +3,12 @@
from pydantic import BaseModel, Field
from superagi.agent.agent_prompt_builder import AgentPromptBuilder
+from superagi.helper.error_handler import ErrorHandler
from superagi.helper.prompt_reader import PromptReader
from superagi.lib.logger import logger
from superagi.llms.base_llm import BaseLlm
+from superagi.models.agent_execution import AgentExecution
+from superagi.models.agent_execution_feed import AgentExecutionFeed
from superagi.tools.base_tool import BaseTool
from superagi.tools.tool_response_query_manager import ToolResponseQueryManager
@@ -34,6 +37,7 @@ class ThinkingTool(BaseTool):
args_schema: Type[ThinkingSchema] = ThinkingSchema
goals: List[str] = []
agent_execution_id:int=None
+ agent_id:int = None
permission_required: bool = False
tool_response_manager: Optional[ToolResponseQueryManager] = None
@@ -62,6 +66,9 @@ def _execute(self, task_description: str):
prompt = prompt.replace("{relevant_tool_response}",relevant_tool_response)
messages = [{"role": "system", "content": prompt}]
result = self.llm.chat_completion(messages, max_tokens=self.max_token_limit)
+
+ if 'error' in result and result['message'] is not None:
+ ErrorHandler.handle_openai_errors(self.toolkit_config.session, self.agent_id, self.agent_execution_id, result['message'])
return result["content"]
except Exception as e:
logger.error(e)
diff --git a/superagi/tools/tool_response_query_manager.py b/superagi/tools/tool_response_query_manager.py
index 0bf708b396..2429945b15 100644
--- a/superagi/tools/tool_response_query_manager.py
+++ b/superagi/tools/tool_response_query_manager.py
@@ -15,8 +15,7 @@ def get_last_response(self, tool_name: str = None):
def get_relevant_response(self, query: str,metadata:dict, top_k: int = 5):
if self.memory is None:
return ""
- documents = self.memory.get_matching_text(query, metadata=metadata,
- top_k=top_k)
+ documents = self.memory.get_matching_text(query, metadata=metadata)
relevant_responses = ""
for document in documents["documents"]:
relevant_responses += document.text_content
diff --git a/superagi/worker.py b/superagi/worker.py
index f31e5fc96f..b6c2cace43 100644
--- a/superagi/worker.py
+++ b/superagi/worker.py
@@ -34,6 +34,10 @@
'task': 'initialize-schedule-agent',
'schedule': timedelta(minutes=5),
},
+ 'execute_waiting_workflows': {
+ 'task': 'execute_waiting_workflows',
+ 'schedule': timedelta(minutes=2),
+ },
}
app.conf.beat_schedule = beat_schedule
@@ -41,11 +45,19 @@
def agent_status_change(target, val,old_val,initiator):
if not hasattr(sys, '_called_from_test'):
webhook_callback.delay(target.id,val,old_val)
-
+
+@app.task(name="execute_waiting_workflows", autoretry_for=(Exception,), retry_backoff=2, max_retries=5)
+def execute_waiting_workflows():
+ """Check if wait time of wait workflow step is over and can be resumed."""
+
+ from superagi.jobs.agent_executor import AgentExecutor
+ logger.info("Executing waiting workflows job")
+ AgentExecutor().execute_waiting_workflows()
+
@app.task(name="initialize-schedule-agent", autoretry_for=(Exception,), retry_backoff=2, max_retries=5)
def initialize_schedule_agent_task():
"""Executing agent scheduling in the background."""
-
+
schedule_helper = AgentScheduleHelper()
schedule_helper.update_next_scheduled_time()
schedule_helper.run_scheduled_agents()
diff --git a/tests/unit_tests/agent/test_agent_workflow_step_wait_handler.py b/tests/unit_tests/agent/test_agent_workflow_step_wait_handler.py
new file mode 100644
index 0000000000..c0fb554404
--- /dev/null
+++ b/tests/unit_tests/agent/test_agent_workflow_step_wait_handler.py
@@ -0,0 +1,67 @@
+from datetime import datetime
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from superagi.models.agent_execution import AgentExecution
+from superagi.models.workflows.agent_workflow_step import AgentWorkflowStep
+from superagi.agent.agent_workflow_step_wait_handler import AgentWaitStepHandler
+
+
+# Mock datetime.now() for testing
+@pytest.fixture
+def mock_datetime_now():
+ return datetime(2023, 9, 6, 12, 0, 0)
+
+
+@pytest.fixture(autouse=True)
+def mock_datetime_now_fixture(monkeypatch, mock_datetime_now):
+ monkeypatch.setattr("superagi.agent.agent_workflow_step_wait_handler.datetime",
+ MagicMock(now=MagicMock(return_value=mock_datetime_now)))
+
+# Test cases
+@patch.object(AgentExecution, 'get_agent_execution_from_id')
+@patch.object(AgentWorkflowStep, 'find_by_id')
+@patch.object(AgentWorkflowStep, 'fetch_next_step')
+def test_handle_next_step_complete(mock_fetch_next_step, mock_find_by_id, mock_get_agent_execution_from_id, mock_datetime_now_fixture):
+ mock_session = MagicMock()
+ mock_agent_execution = MagicMock(current_agent_step_id=1, status="WAIT_STEP")
+
+ mock_get_agent_execution_from_id.return_value = mock_agent_execution
+ mock_find_by_id.return_value = MagicMock()
+
+ mock_next_step = MagicMock(id=2)
+ mock_next_step.__str__.return_value = "COMPLETE"
+ mock_fetch_next_step.return_value = mock_next_step
+
+ handler = AgentWaitStepHandler(mock_session, 1, 2)
+
+ handler.handle_next_step()
+
+ # Assertions
+ assert mock_agent_execution.current_agent_step_id == -1
+ assert mock_agent_execution.status == "COMPLETED"
+ mock_session.commit.assert_called_once()
+
+
+# Test cases
+@patch.object(AgentExecution, 'get_agent_execution_from_id')
+@patch.object(AgentWorkflowStep, 'find_by_id')
+@patch.object(AgentWorkflowStep, 'fetch_next_step')
+def test_execute_step(mock_fetch_next_step, mock_find_by_id, mock_get_agent_execution_from_id):
+ mock_session = MagicMock()
+ mock_agent_execution = MagicMock(current_agent_step_id=1, status="WAIT_STEP")
+ mock_step_wait = MagicMock(status="WAITING")
+
+ mock_get_agent_execution_from_id.return_value = mock_agent_execution
+ mock_find_by_id.return_value = mock_step_wait
+ mock_fetch_next_step.return_value = MagicMock()
+
+ handler = AgentWaitStepHandler(mock_session, 1, 2)
+
+ handler.execute_step()
+
+ # Assertions
+ assert mock_step_wait.status == "WAITING"
+ assert mock_agent_execution.status == "WAIT_STEP"
+ mock_session.commit.assert_called_once()
\ No newline at end of file
diff --git a/tests/unit_tests/apm/test_tools_handler.py b/tests/unit_tests/apm/test_tools_handler.py
index f805bbdf10..da13bca1ef 100644
--- a/tests/unit_tests/apm/test_tools_handler.py
+++ b/tests/unit_tests/apm/test_tools_handler.py
@@ -27,31 +27,29 @@ def test_calculate_tool_usage(tools_handler, mock_session):
tool_used_subquery.c.tool_name = 'Tool1'
tool_used_subquery.c.agent_id = 1
-
agent_count_subquery.c.tool_name = 'Tool1'
agent_count_subquery.c.unique_agents = 1
-
total_usage_subquery.c.tool_name = 'Tool1'
total_usage_subquery.c.total_usage = 5
- tools_handler.get_tool_and_toolkit = MagicMock()
- tools_handler.get_tool_and_toolkit.return_value = {'Tool1': 'Toolkit1'}
-
- mock_session.query().filter_by().subquery.return_value = tool_used_subquery
- mock_session.query().group_by().subquery.return_value = agent_count_subquery
- mock_session.query().group_by().subquery.return_value = total_usage_subquery
+ mock_session.query.return_value.filter_by.return_value.subquery.return_value = tool_used_subquery
+ mock_session.query.return_value.group_by.return_value.subquery.side_effect = [agent_count_subquery, total_usage_subquery]
result_obj = MagicMock()
result_obj.tool_name = 'Tool1'
result_obj.unique_agents = 1
result_obj.total_usage = 5
- mock_session.query().join().all.return_value = [result_obj]
+
+ mock_session.query.return_value.join.return_value.all.return_value = [result_obj]
+
+ tools_handler.get_tool_and_toolkit = MagicMock(return_value={'tool1': 'Toolkit1'})
result = tools_handler.calculate_tool_usage()
assert isinstance(result, list)
expected_output = [{'tool_name': 'Tool1', 'unique_agents': 1, 'total_usage': 5, 'toolkit': 'Toolkit1'}]
+
assert result == expected_output
def test_get_tool_and_toolkit(tools_handler, mock_session):
diff --git a/tests/unit_tests/controllers/test_agent_execution_feeds.py b/tests/unit_tests/controllers/test_agent_execution_feeds.py
new file mode 100644
index 0000000000..bb0322617c
--- /dev/null
+++ b/tests/unit_tests/controllers/test_agent_execution_feeds.py
@@ -0,0 +1,43 @@
+from unittest.mock import MagicMock, Mock, create_autospec, patch
+import pytest
+from fastapi.testclient import TestClient
+from fastapi import HTTPException
+from main import app
+from fastapi_sqlalchemy import db
+from superagi.controllers.agent_execution_feed import get_agent_execution_feed
+
+@patch('superagi.controllers.agent_execution_feed.db')
+def test_get_agent_execution_feed(mock_query):
+ mock_session = create_autospec(pytest.Session)
+
+ AgentExecution = MagicMock()
+ agent_execution = AgentExecution()
+ agent_execution.status = "PAUSED"
+ agent_execution.last_shown_error_id = None
+
+ AgentExecutionFeed = MagicMock()
+ agent_execution_feed = AgentExecutionFeed()
+ agent_execution_feed.error_message = None
+
+ feeds = [agent_execution_feed]
+
+ check_auth = MagicMock()
+ AuthJWT = MagicMock()
+ check_auth.return_value = AuthJWT
+ asc = MagicMock()
+
+ AgentExecutionPermission = MagicMock()
+ agent_execution_permission = AgentExecutionPermission()
+ agent_execution_permission.id = 1
+ agent_execution_permission.created_at = "2021-12-13T00:00:00"
+ agent_execution_permission.response = "Yes"
+ agent_execution_permission.status = "Completed"
+ agent_execution_permission.tool_name = "Tool1"
+ agent_execution_permission.question = "Question1"
+ agent_execution_permission.user_feedback = "Feedback1"
+
+ permissions = [agent_execution_permission]
+ mock_agent_execution = Mock()
+ mock_query.return_value.filter.return_value.first.return_value = mock_agent_execution
+ mock_agent_execution_id = 1
+ assert get_agent_execution_feed(mock_agent_execution_id)
\ No newline at end of file
diff --git a/tests/unit_tests/controllers/test_user.py b/tests/unit_tests/controllers/test_user.py
new file mode 100644
index 0000000000..0ab124dd75
--- /dev/null
+++ b/tests/unit_tests/controllers/test_user.py
@@ -0,0 +1,37 @@
+from unittest.mock import patch
+
+import pytest
+from fastapi.testclient import TestClient
+
+from main import app
+from superagi.models.user import User
+
+client = TestClient(app)
+
+# Define a fixture for an authenticated user
+@pytest.fixture
+def authenticated_user():
+ # Create a mock user object with necessary attributes
+ user = User()
+
+ # Set user attributes
+ user.id = 1 # User ID
+ user.username = "testuser" # User's username
+ user.email = "super6@agi.com" # User's email
+ user.first_login_source = None # User's first login source
+ user.token = "mock-jwt-token"
+
+ return user
+
+# Test case for updating first login source when it's not set
+def test_update_first_login_source(authenticated_user):
+ with patch('superagi.helper.auth.db') as mock_auth_db:
+ source = "github" # Specify the source you want to set
+
+ mock_auth_db.session.query.return_value.filter.return_value.first.return_value = authenticated_user
+ response = client.post(f"users/first_login_source/{source}", headers={"Authorization": f"Bearer {authenticated_user.token}"})
+
+ # Verify the HTTP response
+ assert response.status_code == 200
+ assert "first_login_source" in response.json() # Check if the "first_login_source" field is in the response
+ assert response.json()["first_login_source"] == "github" # Check if the "source" field equals "github"
\ No newline at end of file
diff --git a/tests/unit_tests/helper/test_error_handling.py b/tests/unit_tests/helper/test_error_handling.py
new file mode 100644
index 0000000000..924b227c7b
--- /dev/null
+++ b/tests/unit_tests/helper/test_error_handling.py
@@ -0,0 +1,19 @@
+import pytest
+from unittest.mock import Mock, patch
+from superagi.models.agent_execution import AgentExecution
+from superagi.models.agent_execution_feed import AgentExecutionFeed
+from superagi.helper.error_handler import ErrorHandler
+
+def test_handle_error():
+ session = Mock()
+ agent_id = 1
+ agent_execution_id = 2
+ error_message = 'Test error'
+
+ mock_query = Mock()
+ mock_query.filter().first.return_value = AgentExecution(id=agent_execution_id)
+ session.query.return_value = mock_query
+
+ ErrorHandler.handle_openai_errors(session, agent_id, agent_execution_id, error_message)
+
+ session.query.assert_called_once_with(AgentExecution)
diff --git a/tests/unit_tests/helper/test_feed_parser.py b/tests/unit_tests/helper/test_feed_parser.py
index 0f321a0544..2915949ebd 100644
--- a/tests/unit_tests/helper/test_feed_parser.py
+++ b/tests/unit_tests/helper/test_feed_parser.py
@@ -1,22 +1,16 @@
import unittest
from datetime import datetime
-
from superagi.helper.feed_parser import parse_feed
from superagi.models.agent_execution_feed import AgentExecutionFeed
-
-
class TestParseFeed(unittest.TestCase):
-
def test_parse_feed_system(self):
current_time = datetime.now()
+ sample_feed = AgentExecutionFeed(
+ id=2, agent_execution_id=100, agent_id=200, role="user",
+ feed='System message', updated_at=current_time
+ )
- # Create a sample AgentExecutionFeed object with a system role
- sample_feed = AgentExecutionFeed(id=2, agent_execution_id=100, agent_id=200, role="assistant",
- feed='System message',
- updated_at=current_time)
-
- # Call the parse_feed function with the sample_feed object
result = parse_feed(sample_feed)
-
- # In this test case, we only ensure that the parse_feed function doesn't modify the given feed
- self.assertEqual(result, sample_feed, "Incorrect output from parse_feed function for system role")
+
+ self.assertEqual(result['feed'], sample_feed.feed, "Incorrect output from parse_feed function for system role")
+ self.assertEqual(result['role'], sample_feed.role, "Incorrect output from parse_feed function for system role")
\ No newline at end of file
diff --git a/tests/unit_tests/tools/file/test_list_files.py b/tests/unit_tests/tools/file/test_list_files.py
index 3177ed2b91..f895e6d912 100644
--- a/tests/unit_tests/tools/file/test_list_files.py
+++ b/tests/unit_tests/tools/file/test_list_files.py
@@ -1,4 +1,4 @@
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
import pytest
@@ -9,6 +9,9 @@
def list_file_tool():
list_file_tool = ListFileTool()
list_file_tool.agent_id = 1 # Set a dummy agent ID for testing.
+ # list_file_tool = "test_agent"
+ mock_session = MagicMock(name="session")
+ list_file_tool.toolkit_config.session = mock_session
yield list_file_tool
@@ -23,9 +26,12 @@ def test_list_files(list_file_tool):
assert files == ['file1.txt', 'file3.txt', 'file4.txt']
-
def test_execute(list_file_tool):
- with patch.object(ListFileTool, 'list_files', return_value=['file1.txt', 'file2.txt']):
+ mock_get_formatted_agent_level_path = MagicMock(return_value="SuperAGI/workspace/input/{agent_id}/")
+
+ with patch.object(ListFileTool, 'list_files', return_value=['file1.txt', 'file2.txt']), \
+ patch('superagi.helper.resource_helper.ResourceHelper.get_formatted_agent_level_path', new=mock_get_formatted_agent_level_path):
+
files = list_file_tool._execute()
assert files == ['file1.txt', 'file2.txt']
\ No newline at end of file