change llama_index_agent to knowledgebase agent

bolna/agent_manager/task_manager.py

@@ -561,8 +561,8 @@ def __get_agent_object(self, llm, agent_type, assistant_config = None ):
     def __setup_tasks(self, llm = None, agent_type = None, assistant_config= None):
         if self.task_config["task_type"] == "conversation" and not self.__is_multiagent():
             self.tools["llm_agent"] = self.__get_agent_object(llm, agent_type, assistant_config)
-            if agent_type == "llama-index-rag":
-                logger.info("#### Setting up llama-index-rag agent ####")
+            if agent_type == "knowledgebase_agent":
+                logger.info("#### Setting up knowledgebase_agent agent ####")
                 extra_config = self.task_config["tools_config"]["llm_agent"].get("extra_config", {})
                 vector_store_config = extra_config.get("vector_store", {})
                 self.tools["llm_agent"] = LlamaIndexRag(
@@ -608,7 +608,7 @@ def __get_final_prompt(self, prompt, today):
     async def load_prompt(self, assistant_name, task_id, local, **kwargs):
         logger.info("prompt and config setup started")
         agent_type = self.task_config["tools_config"]["llm_agent"].get("agent_type", "simple_llm_agent")
-        if self.task_config["task_type"] == "webhook" or agent_type in ["openai_assistant", "llamaindex_rag_agent"]:
+        if self.task_config["task_type"] == "webhook" or agent_type in ["openai_assistant", "knowledgebase_agent"]:
             return
         self.is_local = local
         today = datetime.now().strftime("%A, %B %d, %Y")

bolna/models.py

@@ -156,8 +156,6 @@ class VectorStore(BaseModel):
     provider: str
     provider_config: Union[LanceDBProviderConfig, MongoDBProviderConfig]
 
-class ExtraConfig(BaseModel):
-    vector_store : VectorStore
 
 class LLM(BaseModel):
     model: Optional[str] = "gpt-3.5-turbo"
@@ -181,6 +179,7 @@ class SIMPLE_LLM_AGENT(LLM):
     extraction_details: Optional[str] = None
     summarization_details: Optional[str] = None
 
+
 class Node(BaseModel):
     id: str
     type: str #Can be router or conversation for now
@@ -209,12 +208,15 @@ class MultiAgent(BaseModel):
     default_agent: str
     embedding_model: Optional[str] = "Snowflake/snowflake-arctic-embed-l"
 
+class KnowledgebaseAgent(LLM):
+    vector_store: VectorStore
+
 class LLM_AGENT(BaseModel):
     agent_flow_type: str
     agent_type: str #can be llamaindex_rag, simple_llm_agent, router_agent, dag_agent, openai_assistant, custom, etc 
     #extra_config: Union[OpenaiAssistants, LLM_AGENT_GRAPH, MultiAgent, LLM, SIMPLE_LLM_AGENT]
     guardrails: Optional[Routes] = None #Just to reduce confusion
-    extra_config: Union[OpenaiAssistants, LLM_AGENT_GRAPH, MultiAgent, LLM]
+    extra_config: Union[OpenaiAssistants, LLM_AGENT_GRAPH, MultiAgent, KnowledgebaseAgent, SIMPLE_LLM_AGENT, LLM]
 
 
 class MessagingModel(BaseModel):