diff --git a/src/crewai_tools/tools/ai_mind_tool/ai_mind_tool.py b/src/crewai_tools/tools/ai_mind_tool/ai_mind_tool.py index 1059d0053..c36400d0b 100644 --- a/src/crewai_tools/tools/ai_mind_tool/ai_mind_tool.py +++ b/src/crewai_tools/tools/ai_mind_tool/ai_mind_tool.py @@ -7,7 +7,13 @@ from openai import OpenAI from pydantic import BaseModel -class AIMindInputSchema(BaseModel): +class AIMindToolConstants: + MINDS_API_BASE_URL = "https://mdb.ai/" + MIND_NAME_PREFIX = "crwai_mind_" + DATASOURCE_NAME_PREFIX = "crwai_ds_" + + +class AIMindToolInputSchema(BaseModel): """Input for AIMind Tool.""" query: str = "Question in natural language to ask the AI-Mind" @@ -22,7 +28,7 @@ class AIMindTool(BaseTool): "and Google BigQuery. " "Input should be a question in natural language." ) - args_schema: Type[BaseModel] = AIMindInputSchema + args_schema: Type[BaseModel] = AIMindToolInputSchema api_key: Optional[str] = None datasources: Optional[List[Dict[str, Any]]] = None mind_name: Optional[Text] = None @@ -49,7 +55,7 @@ class AIMindTool(BaseTool): datasources = [] for datasource in self.datasources: config = DatabaseConfig( - name=f"cai_ds_{secrets.token_hex(5)}", + name=f"{AIMindToolConstants.DATASOURCE_NAME_PREFIX}_{secrets.token_hex(5)}", engine=datasource["engine"], description=datasource["description"], connection_data=datasource["connection_data"], @@ -58,7 +64,7 @@ class AIMindTool(BaseTool): datasources.append(config) # Generate a random name for the Mind. - name = f"cai_mind_{secrets.token_hex(5)}" + name = f"{AIMindToolConstants.MIND_NAME_PREFIX}_{secrets.token_hex(5)}" mind = minds_client.minds.create( name=name, datasources=datasources, replace=True @@ -72,7 +78,7 @@ class AIMindTool(BaseTool): ): # Run the query on the AI-Mind. # The Minds API is OpenAI compatible and therefore, the OpenAI client can be used. - openai_client = OpenAI(base_url="https://mdb.ai/", api_key=self.api_key) + openai_client = OpenAI(base_url=AIMindToolConstants.MINDS_API_BASE_URL, api_key=self.api_key) completion = openai_client.chat.completions.create( model=self.mind_name,