Merge pull request #1168 from TransformerOptimus/knowledge_model_fix

fix for knowledge search tool
TransformerOptimus · Aug 31, 2023 · 7254014 · 7254014
2 parents 5378cbd + f71bba4
commit 7254014
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/superagi/tools/knowledge_search/knowledge_search.py b/superagi/tools/knowledge_search/knowledge_search.py
@@ -44,8 +44,8 @@ def _execute(self, query: str):
         vector_db_index = VectordbIndices.get_vector_index_from_id(session, knowledge.vector_db_index_id)
         vector_db = Vectordbs.get_vector_db_from_id(session, vector_db_index.vector_db_id)
         db_creds = VectordbConfigs.get_vector_db_config_from_db_id(session, vector_db.id)
-        model_api_key = Configuration.fetch_configuration(session, organisation_id, "model_api_key")
-        model_source = Configuration.fetch_configuration(session, organisation_id, "model_source")
+        model_api_key = self.get_tool_config('OPENAI_API_KEY')
+        model_source = 'OpenAI'
         embedding_model = AgentExecutor.get_embedding(model_source, model_api_key)
         try:
             if vector_db_index.state == "Custom":

diff --git a/superagi/tools/knowledge_search/knowledge_search_toolkit.py b/superagi/tools/knowledge_search/knowledge_search_toolkit.py
@@ -2,7 +2,7 @@
 from typing import List
 from superagi.tools.base_tool import BaseTool, BaseToolkit, ToolConfiguration
 from superagi.tools.knowledge_search.knowledge_search import KnowledgeSearchTool
-
+from superagi.types.key_type import ToolConfigKeyType
 
 class KnowledgeSearchToolkit(BaseToolkit, ABC):
     name: str = "Knowledge Search Toolkit"
@@ -12,4 +12,6 @@ def get_tools(self) -> List[BaseTool]:
         return [KnowledgeSearchTool()]
 
     def get_env_keys(self) -> List[ToolConfiguration]:
-        return []
+        return [
+            ToolConfiguration(key="OPENAI_API_KEY", key_type=ToolConfigKeyType.STRING, is_required=False, is_secret=True)
+        ]