Models superagi (#1117)

TransformerOptimus · Aug 25, 2023 · 1d5d066 · 1d5d066
1 parent fc3c616
commit 1d5d066
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 5 deletions.
diff --git a/superagi/agent/agent_iteration_step_handler.py b/superagi/agent/agent_iteration_step_handler.py
@@ -77,7 +77,7 @@ def execute_step(self):
         if 'content' not in response or response['content'] is None:
             raise RuntimeError(f"Failed to get response from llm")
 
-        total_tokens = current_tokens + TokenCounter(session=self.session, organisation_id=organisation.id).count_message_tokens(response['content'], self.llm.get_model())
+        total_tokens = current_tokens + TokenCounter.count_message_tokens(response['content'], self.llm.get_model())
         AgentExecution.update_tokens(self.session, self.agent_execution_id, total_tokens)
         try:
             content = json.loads(response['content'])

diff --git a/superagi/agent/agent_tool_step_handler.py b/superagi/agent/agent_tool_step_handler.py
@@ -121,7 +121,7 @@ def _build_tool_obj(self, agent_config, agent_execution_config, tool_name: str):
         if tool_name == "QueryResourceTool":
             resource_summary = ResourceSummarizer(session=self.session,
                                                   agent_id=self.agent_id,
-                                                  model= agent_config["model"]).fetch_or_create_agent_resource_summary(
+                                                  model=agent_config["model"]).fetch_or_create_agent_resource_summary(
                 default_summary=agent_config.get("resource_summary"))
 
         organisation = Agent.find_org_by_agent_id(self.session, self.agent_id)

diff --git a/superagi/agent/queue_step_handler.py b/superagi/agent/queue_step_handler.py
@@ -12,6 +12,7 @@
 from superagi.models.agent_execution_feed import AgentExecutionFeed
 from superagi.models.workflows.agent_workflow_step import AgentWorkflowStep
 from superagi.models.workflows.agent_workflow_step_tool import AgentWorkflowStepTool
+from superagi.models.agent import Agent
 from superagi.types.queue_status import QueueStatus
 
 
@@ -22,6 +23,7 @@ def __init__(self, session, llm, agent_id: int, agent_execution_id: int):
         self.llm = llm
         self.agent_execution_id = agent_execution_id
         self.agent_id = agent_id
+        self.organisation = Agent.find_org_by_agent_id(self.session, agent_id=self.agent_id)
 
     def _queue_identifier(self, step_tool):
         return step_tool.unique_id + "_" + str(self.agent_execution_id)
@@ -87,7 +89,7 @@ def _process_input_instruction(self, step_tool):
             .build_agent_messages(prompt, agent_feeds, history_enabled=step_tool.history_enabled,
                                   completion_prompt=step_tool.completion_prompt)
         current_tokens = TokenCounter.count_message_tokens(messages, self.llm.get_model())
-        response = self.llm.chat_completion(messages, TokenCounter.token_limit(self.llm.get_model()) - current_tokens)
+        response = self.llm.chat_completion(messages, TokenCounter(session=self.session, organisation_id=self.organisation.id).token_limit(self.llm.get_model()) - current_tokens)
         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/worker.py b/superagi/worker.py
@@ -11,7 +11,7 @@
 from superagi.config.config import get_config
 from superagi.helper.agent_schedule_helper import AgentScheduleHelper
 from superagi.models.configuration import Configuration
-
+from superagi.models.agent import Agent
 from superagi.models.db import connect_db
 from superagi.types.model_source_types import ModelSourceType
 
@@ -71,7 +71,9 @@ def summarize_resource(agent_id: int, resource_id: int):
     engine = connect_db()
     Session = sessionmaker(bind=engine)
     session = Session()
-    model_source = Configuration.fetch_value_by_agent_id(session, agent_id, "model_source") or "OpenAi"
+    agent_config = Agent.fetch_configuration(session, agent_id)
+    organisation = Agent.find_org_by_agent_id(session, agent_id)
+    model_source = Configuration.fetch_configurations(session, organisation.id, "model_source", agent_config["model"]) or "OpenAi"
     if ModelSourceType.GooglePalm.value in model_source or ModelSourceType.Replicate.value in model_source:
         return